Skip to content

Commit

Permalink
Minor improvements
Browse files Browse the repository at this point in the history
  • Loading branch information
@tomasstolker
tomasstolker committed Jan 27, 2025
1 parent ea0293b commit 6a74327
Show file tree
Hide file tree
Showing 4 changed files with 49 additions and 128 deletions.
8 changes: 4 additions & 4 deletions species/data/database.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4019,8 +4019,8 @@ def add_retrieval(
quench_press = None

abund_in = eq_chem.interpolate_mass_fractions(
np.full(pressure.shape, sample_dict["c_o_ratio"]),
np.full(pressure.shape, sample_dict["metallicity"]),
np.full(pressure.shape[0], sample_dict["c_o_ratio"]),
np.full(pressure.shape[0], sample_dict["metallicity"]),
temp,
pressure,
carbon_pressure_quench=quench_press,
Expand Down Expand Up @@ -4740,8 +4740,8 @@ def petitcode_param(
# Interpolate the abundances, following chemical equilibrium

abund_in, mmw, _ = eq_chem.interpolate_mass_fractions(
np.full(pressure.shape, model_param["c_o_ratio"]),
np.full(pressure.shape, model_param["metallicity"]),
np.full(pressure.shape[0], model_param["c_o_ratio"]),
np.full(pressure.shape[0], model_param["metallicity"]),
temperature,
pressure,
carbon_pressure_quench=p_quench,
Expand Down
151 changes: 36 additions & 115 deletions species/fit/retrieval.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -291,10 +291,6 @@ def __init__(
object_name, inc_phot=True, inc_spec=True, verbose=False
)

# Copy the cloud species into a new list because the values will be adjusted by Radtrans

# self.cloud_species_full = self.cloud_species.copy()

# Get photometric data

self.objphot = []
Expand Down Expand Up @@ -634,19 +630,6 @@ def _set_parameters(
# Cloud parameters

if "log_kappa_0" in bounds:
inspect_prt = inspect.getfullargspec(rt_object.calculate_flux)

if "give_absorption_opacity" not in inspect_prt.args:
raise RuntimeError(
"The Radtrans.calculate_flux method "
"from petitRADTRANS does not have "
"the give_absorption_opacity "
"parameter. Probably you are "
"using an outdated version so "
"please update petitRADTRANS "
"to the latest version."
)

if "fsed_1" in bounds and "fsed_2" in bounds:
self.parameters.append("fsed_1")
self.parameters.append("fsed_2")
Expand All @@ -672,19 +655,6 @@ def _set_parameters(
self.parameters.append("lambda_ray")

elif "log_kappa_gray" in bounds:
inspect_prt = inspect.getfullargspec(rt_object.calculate_flux)

if "give_absorption_opacity" not in inspect_prt.args:
raise RuntimeError(
"The Radtrans.calculate_flux method "
"from petitRADTRANS does not have "
"the give_absorption_opacity "
"parameter. Probably you are "
"using an outdated version so "
"please update petitRADTRANS "
"to the latest version."
)

self.parameters.append("log_kappa_gray")
self.parameters.append("log_cloud_top")

Expand Down Expand Up @@ -722,10 +692,7 @@ def _set_parameters(

if len(self.cloud_species) > 1:
for item in self.cloud_species[1:]:
cloud_1 = item
cloud_2 = self.cloud_species[0]

self.parameters.append(f"{cloud_1}_{cloud_2}_ratio")
self.parameters.append(f"{item}_{self.cloud_species[0]}_ratio")

# Add the flux scaling parameters

Expand Down Expand Up @@ -1199,19 +1166,18 @@ def _prior_transform(

if "log_tau_cloud" in bounds:
for item in self.cloud_species[1:]:
cloud_1 = item
cloud_2 = self.cloud_species[0]

mass_ratio = (
bounds[f"{cloud_1}_{cloud_2}_ratio"][0]
bounds[f"{item}_{self.cloud_species[0]}_ratio"][0]
+ (
bounds[f"{cloud_1}_{cloud_2}_ratio"][1]
- bounds[f"{cloud_1}_{cloud_2}_ratio"][0]
bounds[f"{item}_{self.cloud_species[0]}_ratio"][1]
- bounds[f"{item}_{self.cloud_species[0]}_ratio"][0]
)
* cube[cube_index[f"{cloud_1}_{cloud_2}_ratio"]]
* cube[cube_index[f"{item}_{self.cloud_species[0]}_ratio"]]
)

cube[cube_index[f"{cloud_1}_{cloud_2}_ratio"]] = mass_ratio
cube[cube_index[f"{item}_{self.cloud_species[0]}_ratio"]] = (
mass_ratio
)

else:
for item in self.cloud_species:
Expand Down Expand Up @@ -1558,19 +1524,18 @@ def _prior_transform(

if len(self.cloud_species) > 1:
for item in self.cloud_species[1:]:
cloud_1 = item
cloud_2 = self.cloud_species[0]

mass_ratio = (
bounds[f"{cloud_1}_{cloud_2}_ratio"][0]
bounds[f"{item}_{self.cloud_species[0]}_ratio"][0]
+ (
bounds[f"{cloud_1}_{cloud_2}_ratio"][1]
- bounds[f"{cloud_1}_{cloud_2}_ratio"][0]
bounds[f"{item}_{self.cloud_species[0]}_ratio"][1]
- bounds[f"{item}_{self.cloud_species[0]}_ratio"][0]
)
* cube[cube_index[f"{cloud_1}_{cloud_2}_ratio"]]
* cube[cube_index[f"{item}_{self.cloud_species[0]}_ratio"]]
)

cube[cube_index[f"{cloud_1}_{cloud_2}_ratio"]] = mass_ratio
cube[cube_index[f"{item}_{self.cloud_species[0]}_ratio"]] = (
mass_ratio
)

elif self.chemistry == "equilibrium":
# Cloud mass fractions at the cloud base,
Expand Down Expand Up @@ -2079,11 +2044,8 @@ def _lnlike(
cloud_fractions[item] = 0.0

else:
cloud_1 = item
cloud_2 = self.cloud_species[0]

cloud_fractions[item] = cube[
cube_index[f"{cloud_1}_{cloud_2}_ratio"]
cube_index[f"{item}_{self.cloud_species[0]}_ratio"]
]

log_x_base = log_x_cloud_base(
Expand All @@ -2105,11 +2067,8 @@ def _lnlike(
log_x_base[item] = 0.0

else:
cloud_1 = item
cloud_2 = self.cloud_species[0]

log_x_base[item] = cube[
cube_index[f"{cloud_1}_{cloud_2}_ratio"]
cube_index[f"{item}_{self.cloud_species[0]}_ratio"]
]

else:
Expand Down Expand Up @@ -2200,7 +2159,7 @@ def _lnlike(
if self.check_flux is not None:
# Pressure index at the radiative-convective boundary
# if conv_press is None:
# i_conv = lowres_radtrans.press.shape[0]
# i_conv = lowres_radtrans.press.size]
# else:
# i_conv = np.argmax(conv_press < 1e-6 * lowres_radtrans.press)

Expand Down Expand Up @@ -2359,7 +2318,7 @@ def _lnlike(
ln_prior += -0.5 * f_bol.size * np.log(2.0 * np.pi * sigma_fbol**2)

# for i in range(i_conv):
# for i in range(lowres_radtrans.press.shape[0]):
# for i in range(lowres_radtrans.press.size]):
# if not isclose(
# f_bol_spec,
# f_bol,
Expand Down Expand Up @@ -2880,7 +2839,7 @@ def _lnlike(
/ (2.0 * corr_len[spec_item] ** 2)
)
+ (1.0 - corr_amp[spec_item] ** 2)
* np.eye(data_wavel.shape[0])
* np.eye(data_wavel.size)
* error_i**2
)

Expand All @@ -2907,7 +2866,7 @@ def _lnlike(
# See Eq. 9 in Brogi & Line (2019)

# Number of wavelengths
n_wavel = float(data_flux.shape[0])
n_wavel = float(data_flux.size)

# Apply the optional flux scaling to the data
data_flux_scaled = scaling[spec_item] * data_flux
Expand Down Expand Up @@ -3522,12 +3481,27 @@ def setup_retrieval(
# Pressure array for Radtrans

if self.pressure_grid in ["standard", "manual"]:
print(
f"\nNumber of pressure levels used with the "
f"radiative transfer: {self.pressure.size}"
)

radtrans_press = np.copy(self.pressure)

elif self.pressure_grid == "smaller":
print(
f"\nNumber of pressure levels used with the "
f"radiative transfer: {self.pressure[::3].size}"
)

radtrans_press = self.pressure[::3]

elif self.pressure_grid == "clouds":
print(
"\nNumber of pressure levels used with the "
"radiative transfer: adaptive refinement"
)

radtrans_press = self.pressure[::24]

# Create an instance of Ratrans
Expand Down Expand Up @@ -3642,59 +3616,6 @@ def setup_retrieval(
scattering_in_emission=self.scattering,
)

# Create the RT arrays

# if self.pressure_grid == "standard":
# print(
# f"\nNumber of pressure levels used with the "
# f"radiative transfer: {self.pressure.size}"
# )
#
# rt_object.setup_opa_structure(self.pressure)
#
# for item in self.ccf_radtrans.values():
# item.setup_opa_structure(self.pressure)
#
# if self.check_flux is not None:
# lowres_radtrans.setup_opa_structure(self.pressure)
#
# elif self.pressure_grid == "smaller":
# print(
# f"\nNumber of pressure levels used with the "
# f"radiative transfer: {self.pressure[::3].size}"
# )
#
# rt_object.setup_opa_structure(self.pressure[::3])
#
# for item in self.ccf_radtrans.values():
# item.setup_opa_structure(self.pressure[::3])
#
# if self.check_flux is not None:
# lowres_radtrans.setup_opa_structure(self.pressure[::3])
#
# elif self.pressure_grid == "clouds":
# if len(self.cloud_species) == 0:
# raise ValueError(
# "Please select a different pressure_grid. Setting the argument "
# "to 'clouds' is only possible with the use of cloud species."
# )
#
# # The pressure structure is reinitiated after the
# # refinement around the cloud deck so the current
# # initializiation to 60 pressure points is not used
# print(
# "\nNumber of pressure levels used with the "
# "radiative transfer: adaptive refinement"
# )
#
# rt_object.setup_opa_structure(self.pressure[::24])
#
# for item in self.ccf_radtrans.values():
# item.setup_opa_structure(self.pressure[::24])
#
# if self.check_flux is not None:
# lowres_radtrans.setup_opa_structure(self.pressure[::24])

# Create the knot pressures for temperature profile

if self.pt_profile in ["free", "monotonic"]:
Expand Down
4 changes: 2 additions & 2 deletions species/plot/plot_retrieval.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -811,7 +811,7 @@ def plot_pt_profile(
)

if extra_axis is not None:
ax2.legend(loc="upper right", frameon=False, fontsize=12.0)
ax2.legend(loc="upper right", frameon=False, fontsize=10.0)

else:
if extra_axis is not None:
Expand Down Expand Up @@ -1493,7 +1493,7 @@ def plot_clouds(
ha="left",
va="bottom",
transform=ax1.transAxes,
color="black",
color="white",
fontsize=13.0,
)

Expand Down
14 changes: 7 additions & 7 deletions species/util/retrieval_util.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -132,8 +132,8 @@ def pt_ret_model(
# Interpolate the abundances, following chemical equilibrium

_, _, nabla_ad = eq_chem.interpolate_mass_fractions(
np.full(tedd.shape[0], c_o_ratio),
np.full(tedd.shape[0], metallicity),
np.full(tedd.size, c_o_ratio),
np.full(tedd.size, metallicity),
tedd,
press,
carbon_pressure_quench=None,
Expand Down Expand Up @@ -170,8 +170,8 @@ def pt_ret_model(
t_take = copy.copy(tfinal)

_, _, nabla_ad = eq_chem.interpolate_mass_fractions(
np.full(tedd.shape[0], c_o_ratio),
np.full(tedd.shape[0], metallicity),
np.full(t_take.size, c_o_ratio),
np.full(t_take.size, metallicity),
t_take,
press,
carbon_pressure_quench=None,
Expand Down Expand Up @@ -514,7 +514,7 @@ def create_pt_profile(

elif pt_profile in ["free", "monotonic"]:
knot_temp = []
for i in range(knot_press.shape[0]):
for i in range(knot_press.size):
knot_temp.append(cube[cube_index[f"t{i}"]])

knot_temp = np.asarray(knot_temp)
Expand Down Expand Up @@ -2741,8 +2741,8 @@ def quench_pressure(

# Interpolate the equilibbrium abundances

co_array = np.full(pressure.shape[0], c_o_ratio)
feh_array = np.full(pressure.shape[0], metallicity)
co_array = np.full(pressure.size, c_o_ratio)
feh_array = np.full(pressure.size, metallicity)

from petitRADTRANS.chemistry.pre_calculated_chemistry import (
PreCalculatedEquilibriumChemistryTable,
Expand Down

0 comments on commit 6a74327

Please sign in to comment.