From 9695761f73f38ccfbefbdc54999184d6a4b341e4 Mon Sep 17 00:00:00 2001 From: Tomas Stolker Date: Fri, 14 Jul 2023 15:40:18 +0200 Subject: [PATCH] Finished implementing support for the abund_nodes parameter with free retrievals, added inc_abund parameter to plot_posterior, capital sensitive units in convert_units function --- species/analysis/retrieval.py | 216 ++++++++++++------- species/data/database.py | 64 ++++-- species/plot/plot_mcmc.py | 196 +++++++---------- species/plot/plot_retrieval.py | 41 +++- species/read/read_radtrans.py | 105 ++++++++-- species/util/data_util.py | 66 +++--- species/util/plot_util.py | 4 +- species/util/retrieval_util.py | 372 ++++++++++++++++++--------------- 8 files changed, 631 insertions(+), 433 deletions(-) diff --git a/species/analysis/retrieval.py b/species/analysis/retrieval.py index 0006ef57..69094fdc 100644 --- a/species/analysis/retrieval.py +++ b/species/analysis/retrieval.py @@ -18,8 +18,6 @@ import matplotlib.pyplot as plt import numpy as np -from scipy import stats - try: import pymultinest except: @@ -32,7 +30,7 @@ from molmass import Formula from scipy.integrate import simps -from scipy.stats import invgamma +from scipy.stats import invgamma, norm from typeguard import typechecked from species.analysis import photometry @@ -347,10 +345,11 @@ def __init__( self.parameters = [] - # Initiate the optional P-T parameters + # Initiate the optional P-T and abundance parameters self.pt_smooth = None self.temp_nodes = None + self.abund_nodes = None # Weighting of the photometric and spectroscopic data @@ -456,15 +455,18 @@ def set_parameters( # Abundance parameters if chemistry == "equilibrium": - self.parameters.append("metallicity") self.parameters.append("c_o_ratio") elif chemistry == "free": + if self.abund_nodes is None: + for line_item in self.line_species: + self.parameters.append(line_item) - for i in range(self.abund_nodes): - for item in self.line_species: - self.parameters.append(f"{item}_{i}") + else: + for node_idx in range(self.abund_nodes): + for line_item in self.line_species: + self.parameters.append(f"{line_item}_{node_idx}") # Non-equilibrium chemistry @@ -829,7 +831,7 @@ def run_multinest( can be included when ``pt_profile='monotonic'`` or ``pt_profile='free'``. A prior will be applied that penalizes wiggles in the P-T profile through - the second derivative of the temperature structure + the second derivative of the temperature structure (see `Line et al. (2015) `_ for details). @@ -901,9 +903,8 @@ def run_multinest( ``pt_profile='monotonic'`` or ``pt_profile='free'``. abund_nodes : int, None Number of free abundances nodes that are used with - ``chemistry='free'``. The default value of 1 is used, - that is, constant abundances with altitude if the - argument is set to ``None``. + ``chemistry='free'``. Constant abundances with + altitude are used if the argument is set to ``None``. prior : dict(str, tuple(float, float)), None Dictionary with Gaussian priors for one or multiple parameters. The prior can be set for any of the @@ -939,7 +940,7 @@ def run_multinest( warnings.warn( "The 'quenching' parameter can only be used in " "combination with chemistry='equilibrium'. The " - "argument of \'quenching\' will be set to None." + "argument of 'quenching' will be set to None." ) quenching = None @@ -966,8 +967,8 @@ def run_multinest( # Set number of free abundance nodes if chemistry == "free": - if abund_nodes is None: - self.abund_nodes = 1 + if abund_nodes is None or abund_nodes == 1: + self.abund_nodes = None else: self.abund_nodes = abund_nodes @@ -1009,7 +1010,9 @@ def run_multinest( if "poor_mans_nonequ_chem" in sys.modules: from poor_mans_nonequ_chem.poor_mans_nonequ_chem import interpol_abundances else: - from petitRADTRANS.poor_mans_nonequ_chem.poor_mans_nonequ_chem import interpol_abundances + from petitRADTRANS.poor_mans_nonequ_chem.poor_mans_nonequ_chem import ( + interpol_abundances, + ) print(" [DONE]") @@ -1079,17 +1082,18 @@ def run_multinest( if "o_h_ratio" in bounds: del bounds["o_h_ratio"] - if chemistry == "free" and self.abund_nodes > 1: + if chemistry == "free" and self.abund_nodes is not None: for param_item in ["c_h_ratio", "o_h_ratio", "c_o_ratio"]: if param_item in bounds: - warnings.warn(f"The \'{param_item}\' parameter " - "can not be used if the " - "\'abund_nodes\' argument is set " - "to a value larger than 1. The " - "prior boundaries of " - f"\'{param_item}\' will therefore " - "be removed from the \'bounds\' " - "dictionary.") + warnings.warn( + f"The '{param_item}' parameter " + "can not be used if the " + "'abund_nodes' argument is set. " + "The prior boundaries of " + f"'{param_item}' will therefore " + "be removed from the 'bounds' " + "dictionary." + ) del bounds[param_item] @@ -1188,8 +1192,10 @@ def run_multinest( # The pressure structure is reinitiated after the # refinement around the cloud deck so the current # initializiation to 60 pressure points is not used - print("Number of pressure levels used with the " - "radiative transfer: adaptive refinement") + print( + "Number of pressure levels used with the " + "radiative transfer: adaptive refinement" + ) rt_object.setup_opa_structure(self.pressure[::24]) @@ -1203,9 +1209,7 @@ def run_multinest( if pt_profile in ["free", "monotonic"]: knot_press = np.logspace( - np.log10(self.pressure[0]), - np.log10(self.pressure[-1]), - self.temp_nodes + np.log10(self.pressure[0]), np.log10(self.pressure[-1]), self.temp_nodes ) else: @@ -1213,11 +1217,11 @@ def run_multinest( # Create the knot pressures for abundance profile - if chemistry == "free" and self.abund_nodes > 1: + if chemistry == "free" and self.abund_nodes is not None: knot_press_abund = np.logspace( np.log10(self.pressure[0]), np.log10(self.pressure[-1]), - self.abund_nodes + self.abund_nodes, ) else: @@ -1273,16 +1277,15 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: # Parallax (mas), Gaussian prior - cube[cube_index["parallax"]] = stats.norm.ppf( + cube[cube_index["parallax"]] = norm.ppf( cube[cube_index["parallax"]], loc=self.parallax[0], - scale=self.parallax[1] + scale=self.parallax[1], ) # Pressure-temperature profile if pt_profile in ["molliere", "mod-molliere"]: - # Internal temperature (K) of the Eddington # approximation (middle altitudes) # see Eq. 2 in Mollière et al. (2020) @@ -1300,7 +1303,7 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: if pt_profile == "molliere": # Connection temperature (K) - t_connect = (3.0 / 4.0 * tint ** 4.0 * (0.1 + 2.0 / 3.0)) ** 0.25 + t_connect = (3.0 / 4.0 * tint**4.0 * (0.1 + 2.0 / 3.0)) ** 0.25 # The temperature (K) at temp_3 is scaled down from t_connect temp_3 = t_connect * (1 - cube[cube_index["t3"]]) @@ -1416,7 +1419,6 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: ) if pt_profile == "eddington": - # Internal temperature (K) for the # Eddington approximation if "tint" in bounds: @@ -1461,7 +1463,7 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: # Input log_gamma_r is sampled between 0 and 1 gamma_r = invgamma.ppf( - cube[cube_index["log_gamma_r"]], a=1.0, scale=10.0 ** log_beta_r + cube[cube_index["log_gamma_r"]], a=1.0, scale=10.0**log_beta_r ) cube[cube_index["log_gamma_r"]] = np.log10(gamma_r) @@ -1499,23 +1501,52 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: log_x_abund = {} - for node_idx in range(self.abund_nodes): + if self.abund_nodes is None: for line_item in self.line_species: - item = f"{line_item}_{node_idx}" - if line_item in bounds: - cube[cube_index[item]] = ( + cube[cube_index[line_item]] = ( bounds[line_item][0] + (bounds[line_item][1] - bounds[line_item][0]) - * cube[cube_index[item]] + * cube[cube_index[line_item]] ) - elif item not in ["K", "K_lor_cut", "K_burrows", "K_allard"]: + elif line_item not in [ + "K", + "K_lor_cut", + "K_burrows", + "K_allard", + ]: # Default: -10. - 0. dex - cube[cube_index[item]] = -10.0 * cube[cube_index[item]] + cube[cube_index[line_item]] = ( + -10.0 * cube[cube_index[line_item]] + ) # Add the log10 of the mass fraction to the abundace dictionary - log_x_abund[item] = cube[cube_index[item]] + log_x_abund[line_item] = cube[cube_index[line_item]] + + else: + for node_idx in range(self.abund_nodes): + for line_item in self.line_species: + item = f"{line_item}_{node_idx}" + + if line_item in bounds: + cube[cube_index[item]] = ( + bounds[line_item][0] + + (bounds[line_item][1] - bounds[line_item][0]) + * cube[cube_index[item]] + ) + + elif item not in [ + "K", + "K_lor_cut", + "K_burrows", + "K_allard", + ]: + # Default: -10. - 0. dex + cube[cube_index[item]] = -10.0 * cube[cube_index[item]] + + # Add the log10 of the mass fraction to the abundace dictionary + log_x_abund[item] = cube[cube_index[item]] if ( "Na" in self.line_species @@ -1523,21 +1554,48 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: or "Na_burrows" in self.line_species or "Na_allard" in self.line_species ): + if self.abund_nodes is None: + log_x_k_abund = retrieval_util.potassium_abundance( + log_x_abund, self.line_species, abund_nodes + ) - log_x_k_abund = retrieval_util.potassium_abundance(log_x_abund, self.line_species, knot_press_abund) - - for node_idx in range(self.abund_nodes): if "K" in self.line_species: - cube[cube_index[f"K_{node_idx}"]] = log_x_k_abund[node_idx] + cube[cube_index["K"]] = log_x_k_abund elif "K_lor_cut" in self.line_species: - cube[cube_index[f"K_lor_cut_{node_idx}"]] = log_x_k_abund[node_idx] + cube[cube_index["K_lor_cut"]] = log_x_k_abund elif "K_burrows" in self.line_species: - cube[cube_index[f"K_burrows_{node_idx}"]] = log_x_k_abund[node_idx] + cube[cube_index["K_burrows"]] = log_x_k_abund elif "K_allard" in self.line_species: - cube[cube_index[f"K_allard_{node_idx}"]] = log_x_k_abund[node_idx] + cube[cube_index["K_allard"]] = log_x_k_abund + + else: + log_x_k_abund = retrieval_util.potassium_abundance( + log_x_abund, self.line_species, abund_nodes + ) + + for node_idx in range(self.abund_nodes): + if "K" in self.line_species: + cube[cube_index[f"K_{node_idx}"]] = log_x_k_abund[ + node_idx + ] + + elif "K_lor_cut" in self.line_species: + cube[ + cube_index[f"K_lor_cut_{node_idx}"] + ] = log_x_k_abund[node_idx] + + elif "K_burrows" in self.line_species: + cube[ + cube_index[f"K_burrows_{node_idx}"] + ] = log_x_k_abund[node_idx] + + elif "K_allard" in self.line_species: + cube[ + cube_index[f"K_allard_{node_idx}"] + ] = log_x_k_abund[node_idx] # log10 abundances of the cloud species @@ -2002,10 +2060,10 @@ def prior_func(cube, n_dim: int, n_param: int) -> None: # Standard deviation of the Gaussian kernel for smoothing the P-T profile if "pt_smooth" in bounds: - cube[cube_index[f"pt_smooth"]] = ( + cube[cube_index["pt_smooth"]] = ( bounds["pt_smooth"][0] + (bounds["pt_smooth"][1] - bounds["pt_smooth"][0]) - * cube[cube_index[f"pt_smooth"]] + * cube[cube_index["pt_smooth"]] ) # Mixing-length for convective flux @@ -2140,11 +2198,18 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: # Create a dictionary with the mass fractions - log_x_abund = {} + if self.abund_nodes is None: + log_x_abund = {} + for line_item in self.line_species: + log_x_abund[line_item] = cube[cube_index[line_item]] - for node_idx in range(self.abund_nodes): - for item in self.line_species: - log_x_abund[f"{item}_{node_idx}"] = cube[cube_index[f"{item}_{node_idx}"]] + else: + log_x_abund = {} + for node_idx in range(self.abund_nodes): + for line_item in self.line_species: + log_x_abund[f"{line_item}_{node_idx}"] = cube[ + cube_index[f"{line_item}_{node_idx}"] + ] # Check if the sum of fractional abundances is smaller than unity @@ -2153,7 +2218,7 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: # Check if the C/H and O/H ratios are within the prior boundaries - if self.abund_nodes == 1: + if self.abund_nodes is None: c_h_ratio, o_h_ratio, c_o_ratio = retrieval_util.calc_metal_ratio( log_x_abund, self.line_species ) @@ -2162,22 +2227,19 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: c_h_ratio < bounds["c_h_ratio"][0] or c_h_ratio > bounds["c_h_ratio"][1] ): - return -np.inf if "o_h_ratio" in bounds and ( o_h_ratio < bounds["o_h_ratio"][0] or o_h_ratio > bounds["o_h_ratio"][1] ): - return -np.inf if "c_o_ratio" in bounds and ( c_o_ratio < bounds["c_o_ratio"][0] or c_o_ratio > bounds["c_o_ratio"][1] ): - - return -np.inf + return -np.inf else: c_o_ratio = 0.55 @@ -2372,7 +2434,11 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: tau_cloud = None - if "log_kappa_0" in bounds or "log_kappa_gray" in bounds or "log_kappa_abs" in bounds: + if ( + "log_kappa_0" in bounds + or "log_kappa_gray" in bounds + or "log_kappa_abs" in bounds + ): if "log_tau_cloud" in self.parameters: tau_cloud = 10.0 ** cube[cube_index["log_tau_cloud"]] @@ -2686,11 +2752,11 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: # bolometric flux at each pressure ln_prior += np.sum( - -0.5 * (f_bol - f_bol_spec) ** 2 / sigma_fbol ** 2 + -0.5 * (f_bol - f_bol_spec) ** 2 / sigma_fbol**2 ) ln_prior += ( - -0.5 * f_bol.size * np.log(2.0 * np.pi * sigma_fbol ** 2) + -0.5 * f_bol.size * np.log(2.0 * np.pi * sigma_fbol**2) ) # for i in range(i_conv): @@ -2834,9 +2900,9 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: np.argmin(np.abs(rt_object.tau_rosse - 1.0)) ] - index_tp = (press_tmp > rosse_pphot / 10.0) & ( - press_tmp < rosse_pphot * 10.0 - ) + # index_tp = (press_tmp > rosse_pphot / 10.0) & ( + # press_tmp < rosse_pphot * 10.0 + # ) # tau_pow = np.mean( # np.diff(np.log(rt_object.tau_rosse[index_tp])) @@ -2871,8 +2937,8 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: if hasattr(rt_object, "tau_pow"): ln_like += -0.5 * ( cube[cube_index["alpha"]] - rt_object.tau_pow - ) ** 2.0 / sigma_alpha ** 2.0 - 0.5 * np.log( - 2.0 * np.pi * sigma_alpha ** 2.0 + ) ** 2.0 / sigma_alpha**2.0 - 0.5 * np.log( + 2.0 * np.pi * sigma_alpha**2.0 ) else: @@ -2896,7 +2962,6 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: lbl_flux = {} for item in cross_corr: - ( lbl_wavel[item], lbl_flux[item], @@ -3170,7 +3235,7 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: ) + (1.0 - corr_amp[item] ** 2) * np.eye(data_wavel.shape[0]) - * error_i ** 2 + * error_i**2 ) dot_tmp = np.dot( @@ -3190,7 +3255,7 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: -0.5 * weight * np.sum( - flux_diff ** 2 / data_var + flux_diff**2 / data_var + np.log(2.0 * np.pi * data_var) ) ) @@ -3351,6 +3416,7 @@ def loglike_func(cube, n_dim: int, n_param: int) -> float: radtrans_dict["pressure_grid"] = self.pressure_grid radtrans_dict["wavel_range"] = self.wavel_range radtrans_dict["temp_nodes"] = self.temp_nodes + radtrans_dict["abund_nodes"] = self.abund_nodes radtrans_dict["max_press"] = self.max_pressure if "pt_smooth" not in bounds: diff --git a/species/data/database.py b/species/data/database.py index 9bb95f97..2a6ad0f9 100644 --- a/species/data/database.py +++ b/species/data/database.py @@ -1081,13 +1081,13 @@ def add_object( if verbose: print(f" - {mag_item}:") + print(f" - Mean wavelength (um) = {mean_wavel:.4e}") + print( f" - Apparent magnitude = {app_mag[mag_item][0]:.2f} +/- " f"{app_mag[mag_item][1]:.2f}" ) - print(f" - Mean wavelength (um) = {mean_wavel:.4e}") - print( f" - Flux (W m-2 um-1) = {flux[mag_item]:.2e} +/- " f"{error[mag_item]:.2e}" @@ -1124,13 +1124,13 @@ def add_object( app_mag_item = (dered_mag, app_mag[mag_item][i][1]) if verbose: + print(f" - Mean wavelength (um) = {mean_wavel:.4e}") + print( f" - Apparent magnitude = {app_mag_item[0]:.2f} +/- " f"{app_mag_item[1]:.2f}" ) - print(f" - Mean wavelength (um) = {mean_wavel:.4e}") - print( f" - Flux (W m-2 um-1) = {flux[mag_item][i]:.2e} +/- " f"{error[mag_item][i]:.2e}" @@ -1184,14 +1184,18 @@ def add_object( if flux_item in units: flux_in = np.array([[mean_wavel, data[2], data[3]]]) - flux_out = data_util.convert_units(flux_in, ('um', units[flux_item])) + flux_out = data_util.convert_units( + flux_in, ("um", units[flux_item]) + ) data = [np.nan, np.nan, flux_out[0, 1], flux_out[0, 2]] if verbose: print(f" - {flux_item}:") print(f" - Mean wavelength (um) = {mean_wavel:.4e}") - print(f" - Flux (W m-2 um-1) = {data[2]:.2e} +/- {data[3]:.2e}") + print( + f" - Flux (W m-2 um-1) = {data[2]:.2e} +/- {data[3]:.2e}" + ) # None, None, (W m-2 um-1), (W m-2 um-1) dset = h5_file.create_dataset( @@ -1230,7 +1234,9 @@ def add_object( covariance = hdulist[1].data["COVARIANCE"] # (W m-2 um-1)^2 error = np.sqrt(np.diag(covariance)) # (W m-2 um-1) - read_spec[spec_item] = np.column_stack([wavelength, flux, error]) + read_spec[spec_item] = np.column_stack( + [wavelength, flux, error] + ) else: # Otherwise try to read a 2D dataset with 3 columns @@ -1246,7 +1252,9 @@ def add_object( and spec_item not in read_spec ): if spec_item in units: - data = data_util.convert_units(data, units[spec_item]) + data = data_util.convert_units( + data, units[spec_item] + ) read_spec[spec_item] = data @@ -1294,7 +1302,10 @@ def add_object( ) read_spec[spec_item][:, 1] *= 10.0 ** (0.4 * ext_mag) - if read_spec[spec_item].shape[0] == 3 and read_spec[spec_item].shape[1] != 3: + if ( + read_spec[spec_item].shape[0] == 3 + and read_spec[spec_item].shape[1] != 3 + ): warnings.warn( f"Transposing the data of {spec_item} because " f"the first instead of the second axis " @@ -1363,12 +1374,14 @@ def add_object( else: if spec_item in units: - warnings.warn("The unit conversion has not been " - "implemented for covariance matrices. " - "Please open an issue on the Github " - "page if such functionality is required " - "or provide the file with covariances " - "in (W m-2 um-1)^2.") + warnings.warn( + "The unit conversion has not been " + "implemented for covariance matrices. " + "Please open an issue on the Github " + "page if such functionality is required " + "or provide the file with covariances " + "in (W m-2 um-1)^2." + ) # Otherwise try to read a square, 2D dataset if verbose: @@ -1385,7 +1398,10 @@ def add_object( if data.ndim == 2 and data.shape[0] == data.shape[1]: if spec_item not in read_cov: - if data.shape[0] == read_spec[spec_item].shape[0]: + if ( + data.shape[0] + == read_spec[spec_item].shape[0] + ): if np.all(np.diag(data) == 1.0): warnings.warn(corr_warn) @@ -3309,6 +3325,11 @@ def add_retrieval( if "max_press" in radtrans: dset.attrs["max_press"] = radtrans["max_press"] + if "abund_nodes" in radtrans: + dset.attrs["abund_nodes"] = radtrans["abund_nodes"] + else: + dset.attrs["abund_nodes"] = "None" + print(" [DONE]") # Set number of pressures @@ -3675,6 +3696,16 @@ def get_retrieval_spectra( n_line_species = dset.attrs["n_line_species"] n_cloud_species = dset.attrs["n_cloud_species"] + # Get number of abundance nodes + + if "abund_nodes" in dset.attrs: + if dset.attrs["abund_nodes"] == "None": + abund_nodes = None + else: + abund_nodes = dset.attrs["abund_nodes"] + else: + abund_nodes = None + # Convert numpy boolean to regular boolean scattering = bool(dset.attrs["scattering"]) @@ -3816,6 +3847,7 @@ def get_retrieval_spectra( distance=distance, pt_smooth=pt_smooth, temp_nodes=temp_nodes, + abund_nodes=abund_nodes, read_rad=read_rad, sample=item, ) diff --git a/species/plot/plot_mcmc.py b/species/plot/plot_mcmc.py index cfc4a139..9f0528fc 100644 --- a/species/plot/plot_mcmc.py +++ b/species/plot/plot_mcmc.py @@ -196,6 +196,7 @@ def plot_posterior( inc_log_mass: bool = False, inc_pt_param: bool = False, inc_loglike: bool = False, + inc_abund: bool = True, output: Optional[str] = None, object_type: str = "planet", param_inc: Optional[List[str]] = None, @@ -249,6 +250,9 @@ def plot_posterior( inc_loglike : bool Include the log10 of the likelihood as additional parameter in the corner plot. + inc_abund : bool + Include the abundances when retrieving free abundances with + :class:`~species.analysis.retrieval.AtmosphericRetrieval`. output : str, None Output filename for the plot. The plot is shown in an interface window if the argument is set to ``None``. @@ -324,127 +328,46 @@ def plot_posterior( for item in item_del: box.parameters.remove(item) - if box.spectrum == "petitradtrans" and box.attributes["chemistry"] == "free": - box.parameters.append("c_h_ratio") - box.parameters.append("o_h_ratio") - box.parameters.append("c_o_ratio") - - ndim += 3 - - abund_index = {} - for i, item in enumerate(box.parameters): - if item == "CH4": - abund_index["CH4"] = i - - elif item == "CO": - abund_index["CO"] = i - - elif item == "CO_all_iso": - abund_index["CO_all_iso"] = i - - elif item == "CO_all_iso_HITEMP": - abund_index["CO_all_iso_HITEMP"] = i - - elif item == "CO2": - abund_index["CO2"] = i - - elif item == "FeH": - abund_index["FeH"] = i - - elif item == "H2O": - abund_index["H2O"] = i - - elif item == "H2O_HITEMP": - abund_index["H2O_HITEMP"] = i - - elif item == "H2S": - abund_index["H2S"] = i - - elif item == "Na": - abund_index["Na"] = i - - elif item == "NH3": - abund_index["NH3"] = i - - elif item == "K": - abund_index["K"] = i - - elif item == "PH3": - abund_index["PH3"] = i - - elif item == "TiO": - abund_index["TiO"] = i - - elif item == "TiO_all_Exomol": - abund_index["TiO_all_Exomol"] = i - - elif item == "VO": - abund_index["VO"] = i - - elif item == "VO_Plez": - abund_index["VO_Plez"] = i - - c_h_ratio = np.zeros(samples.shape[0]) - o_h_ratio = np.zeros(samples.shape[0]) - c_o_ratio = np.zeros(samples.shape[0]) - - for i, item in enumerate(samples): - abund = {} - - if "CH4" in box.parameters: - abund["CH4"] = item[abund_index["CH4"]] - - if "CO" in box.parameters: - abund["CO"] = item[abund_index["CO"]] - - if "CO_all_iso" in box.parameters: - abund["CO_all_iso"] = item[abund_index["CO"]] - - if "CO_all_iso_HITEMP" in box.parameters: - abund["CO_all_iso_HITEMP"] = item[abund_index["CO_all_iso_HITEMP"]] - - if "CO2" in box.parameters: - abund["CO2"] = item[abund_index["CO2"]] - - if "FeH" in box.parameters: - abund["FeH"] = item[abund_index["FeH"]] + if box.spectrum == "petitradtrans": + n_line_species = box.attributes["n_line_species"] - if "H2O" in box.parameters: - abund["H2O"] = item[abund_index["H2O"]] + line_species = [] + for i in range(n_line_species): + line_species.append(box.attributes[f"line_species{i}"]) - if "H2O_HITEMP" in box.parameters: - abund["H2O_HITEMP"] = item[abund_index["H2O_HITEMP"]] + if "abund_nodes" not in box.attributes: + box.attributes["abund_nodes"] = "None" - if "H2S" in box.parameters: - abund["H2S"] = item[abund_index["H2S"]] - - if "Na" in box.parameters: - abund["Na"] = item[abund_index["Na"]] - - if "K" in box.parameters: - abund["K"] = item[abund_index["K"]] - - if "NH3" in box.parameters: - abund["NH3"] = item[abund_index["NH3"]] - - if "PH3" in box.parameters: - abund["PH3"] = item[abund_index["PH3"]] - - if "TiO" in box.parameters: - abund["TiO"] = item[abund_index["TiO"]] - - if "TiO_all_Exomol" in box.parameters: - abund["TiO_all_Exomol"] = item[abund_index["TiO_all_Exomol"]] - - if "VO" in box.parameters: - abund["VO"] = item[abund_index["VO"]] - - if "VO_Plez" in box.parameters: - abund["VO_Plez"] = item[abund_index["VO_Plez"]] - - c_h_ratio[i], o_h_ratio[i], c_o_ratio[i] = retrieval_util.calc_metal_ratio( - abund - ) + if box.spectrum == "petitradtrans" and box.attributes["chemistry"] == "free": + if box.attributes["abund_nodes"] == "None": + box.parameters.append("c_h_ratio") + box.parameters.append("o_h_ratio") + box.parameters.append("c_o_ratio") + + ndim += 3 + + abund_index = {} + + for line_item in line_species: + abund_index[line_item] = box.parameters.index(line_item) + + c_h_ratio = np.zeros(samples.shape[0]) + o_h_ratio = np.zeros(samples.shape[0]) + c_o_ratio = np.zeros(samples.shape[0]) + + for sample_item in samples: + abund_dict = {} + for line_item in line_species: + abund_dict[line_item] = sample_item[abund_index[line_item]] + + ( + c_h_ratio[i], + o_h_ratio[i], + c_o_ratio[i], + ) = retrieval_util.calc_metal_ratio( + abund_dict, + line_species, + ) if ( vmr @@ -454,7 +377,7 @@ def plot_posterior( print("Changing mass fractions to number fractions...", end="", flush=True) # Get all available line species - line_species = retrieval_util.get_line_species() + all_line_species = retrieval_util.get_line_species() # Get the atomic and molecular masses masses = retrieval_util.atomic_masses() @@ -467,7 +390,7 @@ def plot_posterior( log_x_abund = {} for param_item in box.parameters: - if param_item in line_species: + if param_item in all_line_species: # Get the index of the parameter param_index = box.parameters.index(param_item) @@ -481,7 +404,7 @@ def plot_posterior( mmw = retrieval_util.mean_molecular_weight(x_abund) for param_item in box.parameters: - if param_item in line_species: + if param_item in all_line_species: # Get the index of the parameter param_index = box.parameters.index(param_item) @@ -715,6 +638,8 @@ def plot_posterior( 10.0 ** samples[:, mass_index] * constants.M_JUP / constants.M_SUN ) + # Include the log-likelihood + if inc_loglike: # Get ln(L) of the samples ln_prob = box.ln_prob[..., np.newaxis] @@ -735,6 +660,35 @@ def plot_posterior( box.parameters.append("log_prob") ndim += 1 + # Remove abundances + + if not inc_abund and box.attributes["chemistry"] == "free": + index_del = [] + item_del = [] + + if box.attributes["abund_nodes"] == "None": + for line_item in line_species: + param_index = np.argwhere(np.array(box.parameters) == line_item)[0] + index_del.append(param_index) + item_del.append(line_item) + + else: + for line_item in line_species: + for node_idx in range(box.attributes["abund_nodes"]): + param_index = np.argwhere( + np.array(box.parameters) == f"{line_item}_{node_idx}" + )[0] + index_del.append(param_index) + item_del.append(f"{line_item}_{node_idx}") + + samples = np.delete(samples, index_del, axis=1) + ndim -= len(index_del) + + for item in item_del: + box.parameters.remove(item) + + # Update axes labels + labels = plot_util.update_labels(box.parameters, object_type=object_type) # Check if parameter values were fixed diff --git a/species/plot/plot_retrieval.py b/species/plot/plot_retrieval.py index 1daa14d0..e57d9c4d 100644 --- a/species/plot/plot_retrieval.py +++ b/species/plot/plot_retrieval.py @@ -172,6 +172,15 @@ def plot_pt_profile( # For backward compatibility temp_nodes = 15 + if "abund_nodes" in box.attributes: + if box.attributes["abund_nodes"] == "None": + abund_nodes = None + else: + abund_nodes = box.attributes["abund_nodes"] + else: + # For backward compatibility + abund_nodes = None + if "max_press" in box.attributes: max_press = box.attributes["max_press"] else: @@ -231,15 +240,29 @@ def plot_pt_profile( # Create a dictionary with the mass fractions - log_x_abund = {} + if abund_nodes is None: + log_x_abund = {} + + for line_idx in range(box.attributes["n_line_species"]): + line_item = box.attributes[f"line_species{line_idx}"] + log_x_abund[line_item] = item[param_index[line_item]] - for j in range(box.attributes["n_line_species"]): - line_item = box.attributes[f"line_species{j}"] - log_x_abund[line_item] = item[param_index[line_item]] + # Check if the C/H and O/H ratios are within the prior boundaries - # Check if the C/H and O/H ratios are within the prior boundaries + _, _, c_o_ratio = retrieval_util.calc_metal_ratio(log_x_abund) - _, _, c_o_ratio = retrieval_util.calc_metal_ratio(log_x_abund) + else: + log_x_abund = {} + for line_idx in range(box.attributes["n_line_species"]): + line_item = box.attributes[f"line_species{line_idx}"] + for node_idx in range(abund_nodes): + log_x_abund[f"{line_item}_{node_idx}"] = model_param[ + f"{line_item}_{node_idx}" + ] + + # TODO Set C/O = 0.55 for Molliere P-T profile + # and cloud condensation profiles + c_o_ratio = 0.55 if pt_profile == "molliere": t3_param = np.array( @@ -716,9 +739,9 @@ def plot_pt_profile( else: raise ValueError( - "The Radtrans object does not contain any cloud species. Please " - "set the argument of 'extra_axis' either to 'photosphere' or " - "None." + "The Radtrans object does not contain any cloud " + "species. Please set the argument of 'extra_axis' " + "either to 'photosphere' or None." ) color_iter = iter(cloud_colors) diff --git a/species/read/read_radtrans.py b/species/read/read_radtrans.py index 558e48fb..c4f581e9 100644 --- a/species/read/read_radtrans.py +++ b/species/read/read_radtrans.py @@ -229,7 +229,7 @@ def get_model( spec_res: Optional[float] = None, wavel_resample: Optional[np.ndarray] = None, plot_contribution: Optional[Union[bool, str]] = False, - temp_nodes: Optional[int] = None, + temp_nodes: Optional[Union[int, np.integer]] = None, ) -> box.ModelBox: """ Function for calculating a model spectrum with @@ -458,17 +458,53 @@ def get_model( else: chemistry = "free" - # Check if all line species from the Radtrans object - # are also present in the model_param dictionary + check_nodes = {} - for item in self.line_species: - if item not in model_param: - raise RuntimeError( - f"The abundance of {item} is not found " - f"in the dictionary with parameters of " - f"'model_param'. Please add the log10 " - f"mass fraction of {item}." - ) + for line_item in self.line_species: + abund_count = 0 + + for node_idx in range(100): + if f"{line_item}_{node_idx}" in model_param: + abund_count += 1 + else: + break + + check_nodes[line_item] = abund_count + + # Check if there are an equal number of + # abundance nodes for all the line species + + nodes_list = list(check_nodes.values()) + + if not all(value == nodes_list[0] for value in nodes_list): + raise ValueError("The number of abundance nodes is " + "not equal for all the lines " + f"species: {check_nodes}") + + if all(value == 0 for value in nodes_list): + abund_nodes = None + else: + abund_nodes = nodes_list[0] + + for line_item in self.line_species: + if abund_nodes is None: + if line_item not in model_param: + raise RuntimeError( + f"The abundance of {line_item} is not " + "found in the dictionary with parameters " + "of 'model_param'. Please add the log10 " + f"mass fraction of {line_item}." + ) + + else: + for node_idx in range(abund_nodes): + if f"{line_item}_{node_idx}" not in model_param: + raise RuntimeError( + f"The abundance of {line_item} is not " + "found in the dictionary with parameters " + "of 'model_param'. Please add the log10 " + f"mass fraction of {line_item}." + ) # Check quenching parameter @@ -490,6 +526,16 @@ def get_model( "'pressure', 'diffusion', or None." ) + # Abundance nodes + + if chemistry == "free" and abund_nodes is not None: + knot_press_abund = np.logspace( + np.log10(self.pressure[0]), np.log10(self.pressure[-1]), abund_nodes + ) + + else: + knot_press_abund = None + # C/O and [Fe/H] if chemistry == "equilibrium": @@ -508,12 +554,26 @@ def get_model( # Create a dictionary with the mass fractions - log_x_abund = {} + if abund_nodes is None: + log_x_abund = {} + for line_item in self.line_species: + log_x_abund[line_item] = model_param[line_item] - for item in self.line_species: - log_x_abund[item] = model_param[item] + _, _, c_o_ratio = retrieval_util.calc_metal_ratio( + log_x_abund, self.line_species + ) + + else: + log_x_abund = {} + for line_item in self.line_species: + for node_idx in range(abund_nodes): + log_x_abund[f"{line_item}_{node_idx}"] = model_param[ + f"{line_item}_{node_idx}" + ] - _, _, c_o_ratio = retrieval_util.calc_metal_ratio(log_x_abund) + # TODO Set C/O = 0.55 for Molliere P-T profile + # and cloud condensation profiles + c_o_ratio = 0.55 # Create the P-T profile @@ -543,8 +603,8 @@ def get_model( if temp_nodes is None: temp_nodes = 0 - for i in range(100): - if f"t{i}" in model_param: + for temp_idx in range(100): + if f"t{temp_idx}" in model_param: temp_nodes += 1 else: break @@ -554,8 +614,8 @@ def get_model( ) knot_temp = [] - for i in range(temp_nodes): - knot_temp.append(model_param[f"t{i}"]) + for temp_idx in range(temp_nodes): + knot_temp.append(model_param[f"t{temp_idx}"]) knot_temp = np.asarray(knot_temp) @@ -604,7 +664,6 @@ def get_model( or "log_kappa_gray" in model_param or "log_kappa_abs" in model_param ): - tau_cloud = None log_x_base = None @@ -627,7 +686,6 @@ def get_model( cloud_fractions = {} for item in self.cloud_species: - if f"{item[:-3].lower()}_fraction" in model_param: cloud_fractions[item] = model_param[ f"{item[:-3].lower()}_fraction" @@ -769,6 +827,7 @@ def get_model( cloud_dict, model_param["logg"], chemistry=chemistry, + knot_press_abund=knot_press_abund, pressure_grid=self.pressure_grid, plotting=False, contribution=True, @@ -796,6 +855,7 @@ def get_model( cloud_dict, model_param["logg"], chemistry=chemistry, + knot_press_abund=knot_press_abund, pressure_grid=self.pressure_grid, plotting=False, contribution=True, @@ -831,6 +891,7 @@ def get_model( model_param, model_param["logg"], chemistry=chemistry, + knot_press_abund=knot_press_abund, pressure_grid=self.pressure_grid, plotting=False, contribution=True, @@ -852,6 +913,7 @@ def get_model( None, pressure_grid=self.pressure_grid, chemistry=chemistry, + knot_press_abund=knot_press_abund, contribution=True, ) @@ -871,6 +933,7 @@ def get_model( None, log_x_abund, chemistry=chemistry, + knot_press_abund=knot_press_abund, pressure_grid=self.pressure_grid, contribution=True, ) diff --git a/species/util/data_util.py b/species/util/data_util.py index e61e44e4..a81b542b 100644 --- a/species/util/data_util.py +++ b/species/util/data_util.py @@ -843,6 +843,7 @@ def retrieval_spectrum( distance: Optional[float], pt_smooth: Optional[float], temp_nodes: Optional[np.integer], + abund_nodes: Optional[np.integer], read_rad: read_radtrans.ReadRadtrans, sample: np.ndarray, ) -> box.ModelBox: @@ -882,6 +883,9 @@ def retrieval_spectrum( temp_nodes : int, None Number of free temperature nodes that are used when ``pt_profile='monotonic'`` or ``pt_profile='free'``. + abund_nodes : int, None + Number of free abundance nodes that are used when + ``chemistry='free'``. read_rad : read_radtrans.ReadRadtrans Instance of :class:`~species.read.read_radtrans.ReadRadtrans`. sample : np.ndarray @@ -939,8 +943,16 @@ def retrieval_spectrum( model_param["metallicity"] = sample[indices["metallicity"]] elif chemistry == "free": - for species_item in line_species: - model_param[species_item] = sample[indices[species_item]] + if abund_nodes is None: + for line_item in line_species: + model_param[line_item] = sample[indices[line_item]] + + else: + for line_item in line_species: + for node_idx in range(abund_nodes): + model_param[f"{line_item}_{node_idx}"] = sample[ + indices[f"{line_item}_{node_idx}"] + ] if quenching == "pressure": model_param["log_p_quench"] = sample[indices["log_p_quench"]] @@ -1054,10 +1066,10 @@ def convert_units(flux_in: np.ndarray, units_in: Tuple[str, str]) -> np.ndarray: photometric fluxes, the array should also be 2D but with a single row/wavelength. units_in : tuple(str, str) - Tuple with the units of the wavelength ("um", "angstrom", "nm", - "mm", "cm", "m") and the units of the flux density - ("w m-2 um-1", "w m-2 m-1", "w m-2 hz-1", "erg s-1 cm-2 hz-1", - "jy", "mjy"). + Tuple with the units of the wavelength ("um", "angstrom", "A", + "nm", "mm", "cm", "m") and the units of the flux density + ("W m-2 um-1", "W m-2 m-1", "W m-2 Hz-1", "erg s-1 cm-2 Hz-1", + "mJy", "Jy", "MJy"). Returns ------- @@ -1071,12 +1083,12 @@ def convert_units(flux_in: np.ndarray, units_in: Tuple[str, str]) -> np.ndarray: # Convert wavelengths to micrometer (um) - wavel_units = ["um", "angstrom", "nm", "mm", "cm", "m"] + wavel_units = ["um", "angstrom", "A", "nm", "mm", "cm", "m"] if units_in[0] == "um": - pass + flux_out[:, 0] = flux_in[:, 0].copy() - elif units_in[0] == "angstrom": + elif units_in[0] in ["angstrom", "A"]: flux_out[:, 0] = flux_in[:, 0] * 1e-4 elif units_in[0] == "nm": @@ -1100,36 +1112,42 @@ def convert_units(flux_in: np.ndarray, units_in: Tuple[str, str]) -> np.ndarray: # Convert flux density to W m-2 um-1 flux_units = [ - "w m-2 um-1", - "w m-2 m-1", - "w m-2 hz-1", - "erg s-1 cm-2 hz-1", - "jy", - "mjy", + "W m-2 um-1", + "W m-2 m-1", + "W m-2 Hz-1", + "erg s-1 cm-2 Hz-1", + "mJy", + "Jy", + "MJy", ] - if units_in[1] == "w m-2 um-1": - pass + if units_in[1] == "W m-2 um-1": + flux_out[:, 1] = flux_in[:, 1].copy() + flux_out[:, 2] = flux_in[:, 2].copy() - elif units_in[1] == "w m-2 m-1": + elif units_in[1] == "W m-2 m-1": flux_out[:, 1] = flux_in[:, 1] * 1e-6 flux_out[:, 2] = flux_in[:, 2] * 1e-6 - elif units_in[1] == "w m-2 hz-1": + elif units_in[1] == "W m-2 Hz-1": flux_out[:, 1] = flux_in[:, 1] * speed_light / flux_out[:, 0] ** 2 flux_out[:, 2] = flux_in[:, 2] * speed_light / flux_out[:, 0] ** 2 - elif units_in[1] == "erg s-1 cm-2 hz-1": + elif units_in[1] == "erg s-1 cm-2 Hz-1": flux_out[:, 1] = flux_in[:, 1] * 1e-3 * speed_light / flux_out[:, 0] ** 2 flux_out[:, 2] = flux_in[:, 2] * 1e-3 * speed_light / flux_out[:, 0] ** 2 - elif units_in[1] == "jy": + elif units_in[1] == "mJy": + flux_out[:, 1] = flux_in[:, 1] * 1e-29 * speed_light / flux_out[:, 0] ** 2 + flux_out[:, 2] = flux_in[:, 2] * 1e-29 * speed_light / flux_out[:, 0] ** 2 + + elif units_in[1] == "Jy": flux_out[:, 1] = flux_in[:, 1] * 1e-26 * speed_light / flux_out[:, 0] ** 2 flux_out[:, 2] = flux_in[:, 2] * 1e-26 * speed_light / flux_out[:, 0] ** 2 - elif units_in[1] == "mjy": - flux_out[:, 1] = flux_in[:, 1] * 1e-29 * speed_light / flux_out[:, 0] ** 2 - flux_out[:, 2] = flux_in[:, 2] * 1e-29 * speed_light / flux_out[:, 0] ** 2 + elif units_in[1] == "MJy": + flux_out[:, 1] = flux_in[:, 1] * 1e-20 * speed_light / flux_out[:, 0] ** 2 + flux_out[:, 2] = flux_in[:, 2] * 1e-20 * speed_light / flux_out[:, 0] ** 2 else: raise ValueError( diff --git a/species/util/plot_util.py b/species/util/plot_util.py index 4bd0a5ce..4f17488b 100644 --- a/species/util/plot_util.py +++ b/species/util/plot_util.py @@ -1379,7 +1379,7 @@ def create_model_label( if par_unit[i] is None: if len(label) > 0: label += ", " - elif model_name is not None: + elif inc_model_name and model_name is not None: label += f"{model_name_new}: " label += f"{par_label[i]} = {value}" @@ -1387,7 +1387,7 @@ def create_model_label( else: if len(label) > 0: label += ", " - elif model_name is not None: + elif inc_model_name and model_name is not None: label += f"{model_name_new}: " label += f"{par_label[i]} = {value} {par_unit[i]}" diff --git a/species/util/retrieval_util.py b/species/util/retrieval_util.py index 3a2963ee..e6eb6744 100644 --- a/species/util/retrieval_util.py +++ b/species/util/retrieval_util.py @@ -579,7 +579,6 @@ def create_abund_dict( chemistry: str, pressure_grid: str = "smaller", indices: Optional[np.ndarray] = None, - knot_press_abund: Optional[np.ndarray] = None ) -> Dict[str, np.ndarray]: """ Function to update the names in the abundance dictionary. @@ -609,9 +608,6 @@ def create_abund_dict( Pressure indices from the adaptive refinement in a cloudy atmosphere. Only required with ``pressure_grid='clouds'``. Otherwise, the argument can be set to ``None``. - knot_press_abund : np.ndarray, None - Pressure nodes at which the abundances are sampled. Only - required when ``chemistry='free'``. Returns ------- @@ -619,11 +615,6 @@ def create_abund_dict( Dictionary with the updated names of the abundances. """ - if knot_press_abund is None: - abund_nodes = 1 - else: - abund_nodes = knot_press_abund.size - # create a dictionary with the updated abundance names abund_out = {} @@ -777,7 +768,6 @@ def calc_spectrum_clear( p_quench: Optional[float], log_x_abund: Optional[dict], chemistry: str, - line_species: List[str], knot_press_abund: Optional[np.ndarray], pressure_grid: str = "smaller", contribution: bool = False, @@ -837,19 +827,23 @@ def calc_spectrum_clear( Emission contribution. """ - # Import interpol_abundances here because it slows down importing - # species otherwise. Importing interpol_abundances is only slow the - # first time, which occurs at the start of the run_multinest method - # of AtmosphericRetrieval - - if "poor_mans_nonequ_chem" in sys.modules: - from poor_mans_nonequ_chem.poor_mans_nonequ_chem import interpol_abundances + if knot_press_abund is None: + abund_nodes = None else: - from petitRADTRANS.poor_mans_nonequ_chem.poor_mans_nonequ_chem import ( - interpol_abundances, - ) + abund_nodes = knot_press_abund.size if chemistry == "equilibrium": + # Import interpol_abundances here because it slows down + # importing species otherwise. Importing interpol_abundances + # is only slow the first time, which occurs at the start of + # the run_multinest method of AtmosphericRetrieval + if "poor_mans_nonequ_chem" in sys.modules: + from poor_mans_nonequ_chem.poor_mans_nonequ_chem import interpol_abundances + else: + from petitRADTRANS.poor_mans_nonequ_chem.poor_mans_nonequ_chem import ( + interpol_abundances, + ) + # Chemical equilibrium abund_in = interpol_abundances( np.full(pressure.shape, c_o_ratio), @@ -867,7 +861,7 @@ def calc_spectrum_clear( # Create a dictionary with all mass fractions abund_in = mass_fractions( - log_x_abund, rt_object, line_species, knot_press_abund + log_x_abund, rt_object, rt_object.line_species, abund_nodes ) # Mean molecular weight @@ -893,7 +887,6 @@ def calc_spectrum_clear( chemistry, pressure_grid=pressure_grid, indices=None, - knot_press_abund=knot_press_abund, ) # calculate the emission spectrum @@ -1017,6 +1010,11 @@ def calc_spectrum_clouds( Array with mean molecular weight. """ + if knot_press_abund is None: + abund_nodes = None + else: + abund_nodes = knot_press_abund.size + if chemistry == "equilibrium": # Import interpol_abundances here because it slows down # importing species otherwise. Importing interpol_abundances @@ -1046,50 +1044,55 @@ def calc_spectrum_clouds( # Free abundances # Create a dictionary with all mass fractions - abund_in = mass_fractions(log_x_abund, rt_object.line_species, knot_press_abund) - - # Mean molecular weight - if knot_press_abund is None: - mmw = mean_molecular_weight(abund_in) - mmw *= np.ones_like(pressure) + abund_in = mass_fractions(log_x_abund, rt_object.line_species, abund_nodes) # Create list of all species abund_species = rt_object.line_species.copy() abund_species.append("H2") abund_species.append("He") - # Create arrays of constant atmosphere abundance - if knot_press_abund is None: + if abund_nodes is None: + # Mean molecular weight + + mmw_float = mean_molecular_weight(abund_in) + mmw = np.full(pressure.size, mmw_float) + + # Create arrays of constant abundances with pressure + for abund_item in abund_species: - abund_in[abund_item] = np.full(pressure.size, abund_in[f"{abund_item}_0"]) - del abund_in[f"{abund_item}_0"] + abund_in[abund_item] = np.full(pressure.size, abund_in[abund_item]) else: + # Create arrays of pressure-dependent abundances + for abund_item in abund_species: knot_abund = np.zeros(knot_press_abund.shape) - for nod_idx in range(knot_press_abund.size): - knot_abund[nod_idx] = abund_in[f"{abund_item}_{nod_idx}"] - del abund_in[f"{abund_item}_{nod_idx}"] + for node_idx in range(abund_nodes): + knot_abund[node_idx] = abund_in[f"{abund_item}_{node_idx}"] + del abund_in[f"{abund_item}_{node_idx}"] abund_in[abund_item] = pt_spline_interp( knot_press_abund, knot_abund, pressure, pt_smooth=0.1 ) - mmw = np.zeros(pressure.shape[0]) + # Mean molecular weight + + mmw = np.zeros(pressure.size) - for i in range(pressure.size): + for press_idx in range(pressure.size): abund_dict = {} - for key, value in abund_in.items(): - abund_dict[f"{key}_0"] = value[i] - mmw[i] = mean_molecular_weight(abund_dict) + for abund_item, abund_value in abund_in.items(): + abund_dict[abund_item] = abund_value[press_idx] + + mmw[press_idx] = mean_molecular_weight(abund_dict) if log_x_base is not None: p_base = {} - for item in log_x_base: + for cloud_item in log_x_base: p_base_item = find_cloud_deck( - item, + cloud_item, pressure, temperature, metallicity, @@ -1098,15 +1101,15 @@ def calc_spectrum_clouds( plotting=plotting, ) - abund_in[f"{item}(c)"] = np.zeros_like(temperature) + abund_in[f"{cloud_item}(c)"] = np.zeros_like(temperature) - abund_in[f"{item}(c)"][pressure < p_base_item] = ( - 10.0 ** log_x_base[item] + abund_in[f"{cloud_item}(c)"][pressure < p_base_item] = ( + 10.0 ** log_x_base[cloud_item] * (pressure[pressure <= p_base_item] / p_base_item) ** cloud_dict["fsed"] ) - p_base[f"{item}(c)"] = p_base_item + p_base[f"{cloud_item}(c)"] = p_base_item # Adaptive pressure refinement around the cloud base if pressure_grid == "clouds": @@ -1120,7 +1123,6 @@ def calc_spectrum_clouds( chemistry, pressure_grid=pressure_grid, indices=indices, - knot_press_abund=knot_press_abund, ) # Create dictionary with sedimentation parameters @@ -1496,7 +1498,7 @@ def kappa_scat( def mass_fractions( log_x_abund: Dict[str, float], line_species: List[str], - knot_press_abund: Optional[np.ndarray], + abund_nodes: Optional[int] = None, ) -> Dict[str, float]: """ Function to return a dictionary with the mass fractions of @@ -1508,10 +1510,9 @@ def mass_fractions( Dictionary with the log10 of the mass fractions of metals. line_species : list(str) List with the line species. - knot_press_abund : np.ndarray, None - Array with the pressures (bar) for the abundances. The - argument can be set to ``None`` if abundances are constant - with pressure. + abund_nodes : int, None + Number of abundance nodes. The argument can be set to + ``None`` if abundances are constant with pressure. Returns ------- @@ -1519,38 +1520,60 @@ def mass_fractions( Dictionary with the mass fractions of all species. """ - if knot_press_abund is None: - abund_nodes = 1 - else: - abund_nodes = knot_press_abund.size - - # Initiate abundance dictionary - abund = {} + if abund_nodes is None: + # Initiate abundance dictionary + abund = {} - for nod_idx in range(abund_nodes): # Initiate the total mass fraction of the metals metal_sum = 0.0 - for item in line_species: - # Add the mass fraction to the dictionary - abund[f"{item}_{nod_idx}"] = 10.0 ** log_x_abund[f"{item}_{nod_idx}"] + for line_item in line_species: + if line_item in log_x_abund: + # Add the mass fraction to the dictionary + abund[line_item] = 10.0 ** log_x_abund[line_item] - # Update the total mass fraction of the metals - metal_sum += abund[f"{item}_{nod_idx}"] + # Update the total mass fraction of the metals + metal_sum += abund[line_item] # Mass fraction of H2 and He ab_h2_he = 1.0 - metal_sum # Add H2 and He mass fraction to the dictionary - abund[f"H2_{nod_idx}"] = ab_h2_he * 0.75 - abund[f"He_{nod_idx}"] = ab_h2_he * 0.25 + abund["H2"] = ab_h2_he * 0.75 + abund["He"] = ab_h2_he * 0.25 + + else: + # Initiate abundance dictionary + abund = {} + + for node_idx in range(abund_nodes): + # Initiate the total mass fraction of the metals + metal_sum = 0.0 + + for line_item in line_species: + if f"{line_item}_{node_idx}" in log_x_abund: + # Add the mass fraction to the dictionary + abund[f"{line_item}_{node_idx}"] = ( + 10.0 ** log_x_abund[f"{line_item}_{node_idx}"] + ) + + # Update the total mass fraction of the metals + metal_sum += abund[f"{line_item}_{node_idx}"] + + # Mass fraction of H2 and He + ab_h2_he = 1.0 - metal_sum + + # Add H2 and He mass fraction to the dictionary + abund[f"H2_{node_idx}"] = ab_h2_he * 0.75 + abund[f"He_{node_idx}"] = ab_h2_he * 0.25 return abund @typechecked def calc_metal_ratio( - log_x_abund: Dict[str, float], line_species: List[str] + log_x_abund: Dict[str, float], + line_species: List[str], ) -> Tuple[float, float, float]: """ Function for calculating [C/H], [O/H], and C/O for a given set @@ -1583,9 +1606,8 @@ def calc_metal_ratio( masses = atomic_masses() # Create a dictionary with all mass fractions - abund = mass_fractions( - log_x_abund=log_x_abund, line_species=line_species, knot_press_abund=None - ) + + abund = mass_fractions(log_x_abund, line_species, abund_nodes=None) # Calculate the mean molecular weight from the input mass fractions mmw = mean_molecular_weight(abund) @@ -1597,89 +1619,89 @@ def calc_metal_ratio( # Calculate the total C abundance - if "CO_0" in abund: - c_abund += abund["CO_0"] * mmw / masses["CO"] + if "CO" in abund: + c_abund += abund["CO"] * mmw / masses["CO"] - if "CO_all_iso_0" in abund: - c_abund += abund["CO_all_iso_0"] * mmw / masses["CO"] + if "CO_all_iso" in abund: + c_abund += abund["CO_all_iso"] * mmw / masses["CO"] - if "CO_all_iso_HITEMP_0" in abund: - c_abund += abund["CO_all_iso_HITEMP_0"] * mmw / masses["CO"] + if "CO_all_iso_HITEMP" in abund: + c_abund += abund["CO_all_iso_HITEMP"] * mmw / masses["CO"] - if "CO_all_iso_Chubb_0" in abund: - c_abund += abund["CO_all_iso_Chubb_0"] * mmw / masses["CO"] + if "CO_all_iso_Chubb" in abund: + c_abund += abund["CO_all_iso_Chubb"] * mmw / masses["CO"] - if "CO2_0" in abund: - c_abund += abund["CO2_0"] * mmw / masses["CO2"] + if "CO2" in abund: + c_abund += abund["CO2"] * mmw / masses["CO2"] - if "CO2_main_iso_0" in abund: - c_abund += abund["CO2_main_iso_0"] * mmw / masses["CO2"] + if "CO2_main_iso" in abund: + c_abund += abund["CO2_main_iso"] * mmw / masses["CO2"] - if "CH4_0" in abund: - c_abund += abund["CH4_0"] * mmw / masses["CH4"] + if "CH4" in abund: + c_abund += abund["CH4"] * mmw / masses["CH4"] - if "CH4_main_iso_0" in abund: - c_abund += abund["CH4_main_iso_0"] * mmw / masses["CH4"] + if "CH4_main_iso" in abund: + c_abund += abund["CH4_main_iso"] * mmw / masses["CH4"] # Calculate the total O abundance - if "CO_0" in abund: - o_abund += abund["CO_0"] * mmw / masses["CO"] + if "CO" in abund: + o_abund += abund["CO"] * mmw / masses["CO"] - if "CO_all_iso_0" in abund: - o_abund += abund["CO_all_iso_0"] * mmw / masses["CO"] + if "CO_all_iso" in abund: + o_abund += abund["CO_all_iso"] * mmw / masses["CO"] - if "CO_all_iso_HITEMP_0" in abund: - o_abund += abund["CO_all_iso_HITEMP_0"] * mmw / masses["CO"] + if "CO_all_iso_HITEMP" in abund: + o_abund += abund["CO_all_iso_HITEMP"] * mmw / masses["CO"] - if "CO_all_iso_Chubb_0" in abund: - o_abund += abund["CO_all_iso_Chubb_0"] * mmw / masses["CO"] + if "CO_all_iso_Chubb" in abund: + o_abund += abund["CO_all_iso_Chubb"] * mmw / masses["CO"] - if "CO2_0" in abund: - o_abund += 2.0 * abund["CO2_0"] * mmw / masses["CO2"] + if "CO2" in abund: + o_abund += 2.0 * abund["CO2"] * mmw / masses["CO2"] - if "CO2_main_iso_0" in abund: - o_abund += 2.0 * abund["CO2_main_iso_0"] * mmw / masses["CO2"] + if "CO2_main_iso" in abund: + o_abund += 2.0 * abund["CO2_main_iso"] * mmw / masses["CO2"] - if "H2O_0" in abund: - o_abund += abund["H2O_0"] * mmw / masses["H2O"] + if "H2O" in abund: + o_abund += abund["H2O"] * mmw / masses["H2O"] - if "H2O_HITEMP_0" in abund: - o_abund += abund["H2O_HITEMP_0"] * mmw / masses["H2O"] + if "H2O_HITEMP" in abund: + o_abund += abund["H2O_HITEMP"] * mmw / masses["H2O"] - if "H2O_main_iso_0" in abund: - o_abund += abund["H2O_main_iso_0"] * mmw / masses["H2O"] + if "H2O_main_iso" in abund: + o_abund += abund["H2O_main_iso"] * mmw / masses["H2O"] # Calculate the total H abundance - h_abund += 2.0 * abund["H2_0"] * mmw / masses["H2"] + h_abund += 2.0 * abund["H2"] * mmw / masses["H2"] - if "CH4_0" in abund: - h_abund += 4.0 * abund["CH4_0"] * mmw / masses["CH4"] + if "CH4" in abund: + h_abund += 4.0 * abund["CH4"] * mmw / masses["CH4"] - if "CH4_main_iso_0" in abund: - h_abund += 4.0 * abund["CH4_main_iso_0"] * mmw / masses["CH4"] + if "CH4_main_iso" in abund: + h_abund += 4.0 * abund["CH4_main_iso"] * mmw / masses["CH4"] - if "H2O_0" in abund: - h_abund += 2.0 * abund["H2O_0"] * mmw / masses["H2O"] + if "H2O" in abund: + h_abund += 2.0 * abund["H2O"] * mmw / masses["H2O"] - if "H2O_HITEMP_0" in abund: - h_abund += 2.0 * abund["H2O_HITEMP_0"] * mmw / masses["H2O"] + if "H2O_HITEMP" in abund: + h_abund += 2.0 * abund["H2O_HITEMP"] * mmw / masses["H2O"] - if "H2O_main_iso_0" in abund: - h_abund += 2.0 * abund["H2O_main_iso_0"] * mmw / masses["H2O"] + if "H2O_main_iso" in abund: + h_abund += 2.0 * abund["H2O_main_iso"] * mmw / masses["H2O"] - if "NH3_0" in abund: - h_abund += 3.0 * abund["NH3_0"] * mmw / masses["NH3"] + if "NH3" in abund: + h_abund += 3.0 * abund["NH3"] * mmw / masses["NH3"] - if "NH3_main_iso_0" in abund: - h_abund += 3.0 * abund["NH3_main_iso_0"] * mmw / masses["NH3"] + if "NH3_main_iso" in abund: + h_abund += 3.0 * abund["NH3_main_iso"] * mmw / masses["NH3"] - if "H2S_0" in abund: - h_abund += 2.0 * abund["H2S_0"] * mmw / masses["H2S"] + if "H2S" in abund: + h_abund += 2.0 * abund["H2S"] * mmw / masses["H2S"] - if "H2S_main_iso_0" in abund: - h_abund += 2.0 * abund["H2S_main_iso_0"] * mmw / masses["H2S"] + if "H2S_main_iso" in abund: + h_abund += 2.0 * abund["H2S_main_iso"] * mmw / masses["H2S"] return ( np.log10(c_abund / h_abund / c_h_solar), @@ -1710,25 +1732,23 @@ def mean_molecular_weight(abundances: Dict[str, float]) -> float: mmw_sum = 0.0 for abund_item in abundances: - key = abund_item[:-2] - - if key in ["CO_all_iso", "CO_all_iso_HITEMP", "CO_all_iso_Chubb"]: + if abund_item in ["CO_all_iso", "CO_all_iso_HITEMP", "CO_all_iso_Chubb"]: mmw_sum += abundances[abund_item] / masses["CO"] - elif key in ["Na_lor_cut", "Na_allard", "Na_burrows"]: + elif abund_item in ["Na_lor_cut", "Na_allard", "Na_burrows"]: mmw_sum += abundances[abund_item] / masses["Na"] - elif key in ["K_lor_cut", "K_allard", "K_burrows"]: + elif abund_item in ["K_lor_cut", "K_allard", "K_burrows"]: mmw_sum += abundances[abund_item] / masses["K"] - elif key == "CH4_main_iso": + elif abund_item == "CH4_main_iso": mmw_sum += abundances[abund_item] / masses["CH4"] - elif key in ["H2O_main_iso", "H2O_HITEMP"]: + elif abund_item in ["H2O_main_iso", "H2O_HITEMP"]: mmw_sum += abundances[abund_item] / masses["H2O"] else: - mmw_sum += abundances[abund_item] / masses[key] + mmw_sum += abundances[abund_item] / masses[abund_item] return 1.0 / mmw_sum @@ -1737,8 +1757,8 @@ def mean_molecular_weight(abundances: Dict[str, float]) -> float: def potassium_abundance( log_x_abund: Dict[str, float], line_species: List[str], - knot_press_abund: Optional[np.ndarray], -) -> List[float]: + abund_nodes: Optional[int] = None, +) -> Union[float, List[float]]: """ Function to calculate the mass fraction of potassium at a solar ratio of the sodium and potassium abundances. @@ -1749,23 +1769,19 @@ def potassium_abundance( Dictionary with the log10 of the mass fractions. line_species : list(str) List with the line species. - knot_press_abund : np.ndarray, None - Array with the pressures (bar) for the abundances. The - argument can be set to ``None`` if abundances are constant - with pressure. + abund_nodes : int, None + Number of abundance nodes. The argument can be set to + ``None`` if abundances are constant with pressure. Returns ------- - list(float) - Log10 of the mass fractions of potassium in the order - of ``knot_press_abund``. + float, list(float) + Log10 of the mass fraction of potassium or a list with the + log10 mass fraction of potassium in case ``abund_nodes`` is + not ``None``. The length of the list is equal to the argument + of ``abund_nodes`` in that case. """ - if knot_press_abund is None: - abund_nodes = 1 - else: - abund_nodes = knot_press_abund.size - # Solar volume mixing ratios of Na and K (Asplund et al. 2009) n_na_solar = 1.60008694353205e-06 n_k_solar = 9.86605611925677e-08 @@ -1774,39 +1790,66 @@ def potassium_abundance( masses = atomic_masses() # Create a dictionary with all mass fractions - x_abund = mass_fractions(log_x_abund, line_species, knot_press_abund) - - log_x_k = [] + x_abund = mass_fractions(log_x_abund, line_species, abund_nodes) - for node_idx in range(abund_nodes): + if abund_nodes is None: # Calculate mean molecular weight from the mass fractions - abund_dict = {} - for abund_item in line_species: - abund_dict[f"{abund_item}_0"] = x_abund[f"{abund_item}_{node_idx}"] - - mmw = mean_molecular_weight(abund_dict) + mmw = mean_molecular_weight(x_abund) # Volume mixing ratio of sodium - if f"Na_{node_idx}" in log_x_abund: - n_na_abund = x_abund[f"Na_{node_idx}"] * mmw / masses["Na"] + if "Na" in log_x_abund: + n_na_abund = x_abund["Na"] * mmw / masses["Na"] - elif f"Na_lor_cut_{node_idx}" in log_x_abund: - n_na_abund = x_abund[f"Na_lor_cut_{node_idx}"] * mmw / masses["Na"] + elif "Na_lor_cut" in log_x_abund: + n_na_abund = x_abund["Na_lor_cut"] * mmw / masses["Na"] - elif f"Na_allard_{node_idx}" in log_x_abund: - n_na_abund = x_abund[f"Na_allard_{node_idx}"] * mmw / masses["Na"] + elif "Na_allard" in log_x_abund: + n_na_abund = x_abund["Na_allard"] * mmw / masses["Na"] - elif f"Na_burrows_{node_idx}" in log_x_abund: - n_na_abund = x_abund[f"Na_burrows_{node_idx}"] * mmw / masses["Na"] + elif "Na_burrows" in log_x_abund: + n_na_abund = x_abund["Na_burrows"] * mmw / masses["Na"] # Volume mixing ratio of potassium n_k_abund = n_na_abund * n_k_solar / n_na_solar # Mass fraction of potassium - log_x_k.append(np.log10(n_k_abund * masses["K"] / mmw)) + log_x_k = np.log10(n_k_abund * masses["K"] / mmw) + + else: + log_x_k = [] + + for node_idx in range(abund_nodes): + # Calculate mean molecular weight from the mass fractions + + abund_dict = {} + for abund_item in line_species: + abund_dict[abund_item] = x_abund[f"{abund_item}_{node_idx}"] + + mmw = mean_molecular_weight(abund_dict) + + # Volume mixing ratio of sodium + + if f"Na_{node_idx}" in log_x_abund: + n_na_abund = x_abund[f"Na_{node_idx}"] * mmw / masses["Na"] + + elif f"Na_lor_cut_{node_idx}" in log_x_abund: + n_na_abund = x_abund[f"Na_lor_cut_{node_idx}"] * mmw / masses["Na"] + + elif f"Na_allard_{node_idx}" in log_x_abund: + n_na_abund = x_abund[f"Na_allard_{node_idx}"] * mmw / masses["Na"] + + elif f"Na_burrows_{node_idx}" in log_x_abund: + n_na_abund = x_abund[f"Na_burrows_{node_idx}"] * mmw / masses["Na"] + + # Volume mixing ratio of potassium + + n_k_abund = n_na_abund * n_k_solar / n_na_solar + + # Mass fraction of potassium + log_x_k.append(np.log10(n_k_abund * masses["K"] / mmw)) return log_x_k @@ -2287,7 +2330,6 @@ def scale_cloud_abund( chemistry, pressure_grid=pressure_grid, indices=indices, - knot_press_abund=None, ) # Interpolate the line opacities to the temperature structure