diff --git a/src/astra/pipelines/aspcap/__init__.py b/src/astra/pipelines/aspcap/__init__.py index c2f5bea..05731a4 100644 --- a/src/astra/pipelines/aspcap/__init__.py +++ b/src/astra/pipelines/aspcap/__init__.py @@ -103,16 +103,24 @@ def aspcap( for spectrum in spectra_with_no_initial_guess: yield ASPCAP.from_spectrum(spectrum, flag_no_suitable_initial_guess=True) - coarse_results = _aspcap_stage("coarse", coarse_plans, parent_dir, max_processes, max_threads, max_concurrent_loading, soft_thread_ratio) + all_coarse_results = _aspcap_stage("coarse", coarse_plans, parent_dir, max_processes, max_threads, max_concurrent_loading, soft_thread_ratio) - stellar_parameter_plans, best_coarse_results = plan_stellar_parameters_stage(spectra, coarse_results) - - stellar_parameter_results = _aspcap_stage("params", stellar_parameter_plans, parent_dir, max_processes, max_threads, max_concurrent_loading, soft_thread_ratio) + stellar_parameter_plans, coarse_results = plan_stellar_parameters_stage(spectra, all_coarse_results) + results = _aspcap_stage("params", stellar_parameter_plans, parent_dir, max_processes, max_threads, max_concurrent_loading, soft_thread_ratio) - # yeet back some ASPCAP results - for r in stellar_parameter_results: - coarse = best_coarse_results[r["spectrum_pk"]] + for r in results: + coarse = coarse_results[r["spectrum_pk"]] + + v_sini = 10**(r.get("log10_v_sini", np.nan)) + e_v_sini = r.get("e_log10_v_sini", np.nan) * v_sini * np.log(10) + v_micro = 10**(r.get("log10_v_micro", np.nan)) + e_v_micro = r.get("e_log10_v_micro", np.nan) * v_micro * np.log(10) + r.update( + v_sini=v_sini, + e_v_sini=e_v_sini, + v_micro=v_micro, + e_v_micro=e_v_micro, coarse_teff=coarse.teff, coarse_logg=coarse.logg, coarse_v_micro=10**(coarse.log10_v_micro or np.nan), @@ -136,21 +144,11 @@ def aspcap( ferre_time_coarse=coarse.t_elapsed, ferre_time_params=r["t_elapsed"], ) - yield ASPCAP(**r) - - - -# TODO: move elsewhere? -def _pre_compute_continuum(coarse_result, spectrum, pre_continuum): - try: - # Apply continuum normalization. - pre_computed_continuum = pre_continuum.fit(spectrum, coarse_result) - except: - log.exception(f"Exception when computing continuum for spectrum {spectrum} from coarse result {coarse_result}:") - return (spectrum.spectrum_pk, None) - else: - return (spectrum.spectrum_pk, pre_computed_continuum) + + # TODO: Chemical abundances. + for r in results: + yield ASPCAP(**r) diff --git a/src/astra/pipelines/aspcap/coarse.py b/src/astra/pipelines/aspcap/coarse.py index 03078cf..0f516fb 100644 --- a/src/astra/pipelines/aspcap/coarse.py +++ b/src/astra/pipelines/aspcap/coarse.py @@ -1,7 +1,6 @@ import os import numpy as np from glob import glob -from tqdm import tqdm from astra import task from astra.models.spectrum import Spectrum from astra.models.aspcap import FerreCoarse @@ -81,7 +80,7 @@ def plan_coarse_stellar_parameters( all_kwds = [] spectrum_primary_keys_with_at_least_one_initial_guess = set() - for spectrum, input_initial_guess in tqdm(initial_guess_callable(spectra), total=0, desc="Initial guesses"): + for spectrum, input_initial_guess in initial_guess_callable(spectra): n_initial_guesses = 0 for strict in (True, False): @@ -234,7 +233,6 @@ def plan_coarse_stellar_parameters( # Bundle them together into executables based on common header paths. header_paths = list(set([ea["header_path"] for ea in all_kwds])) - log.info(f"Found {len(header_paths)} unique header paths") grouped_task_kwds = { header_path: [] for header_path in header_paths } for kwds in all_kwds: diff --git a/src/astra/pipelines/ferre/pre_process.py b/src/astra/pipelines/ferre/pre_process.py index ddeda44..3164089 100644 --- a/src/astra/pipelines/ferre/pre_process.py +++ b/src/astra/pipelines/ferre/pre_process.py @@ -99,12 +99,7 @@ def pre_process_ferre( control_kwds[key] = prefix + control_kwds[key] absolute_pwd = expand_path(pwd) - #log.info(f"FERRE working directory: {absolute_pwd}") - - # Construct mask to match FERRE model grid. - #chip_wavelengths = tuple(map(utils.wavelength_array, segment_headers)) - values_or_cycle_none = lambda x: x if (x is not None and len(x) > 0) else cycle([None]) all_initial_parameters = dict_to_list(dict( teff=values_or_cycle_none(initial_teff), @@ -119,13 +114,8 @@ def pre_process_ferre( upstream_pk=values_or_cycle_none(upstream_pk) )) - # Retrict to the pixels within the model wavelength grid. - # TODO: Assuming all spectra are the same. - #mask = _get_ferre_chip_mask(spectra[0].wavelength, chip_wavelengths) - - # TODO: use mask2 mask = utils.get_apogee_pixel_mask() - #assert np.all(mask == mask2) + has_warned_on_bad_pixels = False index, skipped, batch_names, batch_initial_parameters, batch_flux, batch_e_flux = (0, [], [], [], [], []) @@ -145,7 +135,7 @@ def pre_process_ferre( flux = np.copy(spectrum.flux) e_flux = np.copy(spectrum.ivar)**-0.5 except: - log.warning(f"Exception accessing pixel arrays for spectrum {spectrum}") + #log.warning(f"Exception accessing pixel arrays for spectrum {spectrum}") skipped.append((spectrum, {"flag_spectrum_io_error": True})) continue @@ -202,10 +192,8 @@ def pre_process_ferre( return (pwd, 0, skipped) control_kwds_formatted = utils.format_ferre_control_keywords(control_kwds, n_obj=1 + index) - #log.info(f"FERRE control keywords:\n{control_kwds_formatted}") # Convert list of dicts of initial parameters to array. - #log.info(f"Validating initial and frozen parameters") batch_initial_parameters_array = utils.validate_initial_and_frozen_parameters( headers, batch_initial_parameters, @@ -215,18 +203,15 @@ def pre_process_ferre( ) # Create directory and write the control file os.makedirs(absolute_pwd, exist_ok=True) - #log.info(f"Writing control file") with open(os.path.join(absolute_pwd, "input.nml"), "w") as fp: fp.write(control_kwds_formatted) # hack: we do basename here in case we wrote the prefix to PFILE for the abundances run - #log.info(f"Writing input parameters") with open(os.path.join(absolute_pwd, os.path.basename(control_kwds["pfile"])), "w") as fp: for name, point in zip(batch_names, batch_initial_parameters_array): fp.write(utils.format_ferre_input_parameters(*point, name=name)) if write_input_pixel_arrays: - #log.info(f"Writing input pixel arrays") LARGE = 1e10 batch_flux = np.array(batch_flux) @@ -242,21 +227,17 @@ def pre_process_ferre( non_finite_flux = ~np.isfinite(batch_flux) batch_flux[non_finite_flux] = 0.0 batch_e_flux[non_finite_flux] = LARGE - #if np.any(non_finite_flux): - # log.warning(f"Non-finite fluxes found. Setting them to zero and setting flux error to {LARGE:.1e}") finite_e_flux = np.isfinite(batch_e_flux) batch_e_flux[~finite_e_flux] = LARGE if not np.any(finite_e_flux): log.warning(f"ALL flux errors are non-finite!") - # Write data arrays. savetxt_kwds = dict(fmt="%.4e")#footer="\n") np.savetxt(flux_path, batch_flux, **savetxt_kwds) np.savetxt(e_flux_path, batch_e_flux, **savetxt_kwds) n_obj = len(batch_names) - log.info(f"FERRE pre-processing complete for {pwd} ({n_obj} objects)") return (pwd, n_obj, skipped) '''