Skip to content

Commit

Permalink
Finished implementing clouds with free abundances
Browse files Browse the repository at this point in the history
  • Loading branch information
tomasstolker committed Mar 7, 2021
1 parent d3c85d1 commit 14ceec0
Show file tree
Hide file tree
Showing 8 changed files with 278 additions and 152 deletions.
94 changes: 44 additions & 50 deletions species/analysis/retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -357,9 +357,6 @@ def set_parameters(self,
for item in self.cloud_species:
self.parameters.append(item[:-3])

if 'c_o_ratio' in bounds:
self.parameters.append('c_o_ratio')

if quenching:
self.parameters.append('log_p_quench')

Expand Down Expand Up @@ -1185,10 +1182,53 @@ def loglike(cube,
else:
pt_smooth = self.pt_smooth

# C/O and [Fe/H]

if chemistry == 'equilibrium':
metallicity = cube[cube_index['metallicity']]
c_o_ratio = cube[cube_index['c_o_ratio']]

elif chemistry == 'free':
# TODO Set [Fe/H] = 0 for P-T profile
metallicity = 0.

# Create a dictionary with the mass fractions

log_x_abund = {}

for item in self.line_species:
log_x_abund[item] = cube[cube_index[item]]

# Check if the sum of fractional abundances is smaller than unity

if np.sum(10.**np.asarray(list(log_x_abund.values()))) > 1.:
return -np.inf

# Check if the C/H and O/H ratios are within the prior boundaries

c_h_ratio, o_h_ratio, c_o_ratio = \
retrieval_util.calc_metal_ratio(log_x_abund)

if 'c_h_ratio' in bounds and (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

# Create the P-T profile

temp, knot_temp = retrieval_util.create_pt_profile(
cube, cube_index, pt_profile, self.pressure, knot_press, pt_smooth)
cube, cube_index, pt_profile, self.pressure, knot_press,
metallicity, c_o_ratio, pt_smooth)

# Prepare the scaling based on the cloud optical depth

Expand Down Expand Up @@ -1248,42 +1288,6 @@ def loglike(cube,
else:
log_p_quench = -10.

# Free chemistry mass fraction and abundance ratios

if chemistry == 'free':
# Create a dictionary with the mass fractions

log_x_abund = {}

for item in self.line_species:
log_x_abund[item] = cube[cube_index[item]]

# Check if the sum of fractional abundances is smaller than unity

if np.sum(10.**np.asarray(list(log_x_abund.values()))) > 1.:
return -np.inf

# Check if the C/H and O/H ratios are within the prior boundaries

if 'c_h_ratio' in bounds or 'o_h_ratio' in bounds or 'c_o_ratio' in bounds:
c_h_ratio, o_h_ratio, c_o_ratio = \
retrieval_util.calc_metal_ratio(log_x_abund)

if 'c_h_ratio' in bounds and (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

# Calculate the emission spectrum

start = time.time()
Expand Down Expand Up @@ -1325,9 +1329,6 @@ def loglike(cube,
cube[cube_index['metallicity']],
cloud_fractions)

c_o_ratio = cube[cube_index['c_o_ratio']]
metallicity = cube[cube_index['metallicity']]

log_x_abund = None

elif chemistry == 'free':
Expand All @@ -1338,13 +1339,6 @@ def loglike(cube,
for item in self.cloud_species:
log_x_base[item[:-3]] = cube[cube_index[item]]

# Calculate C/O ratio from abundances

c_h_ratio, o_h_ratio, c_o_ratio = retrieval_util.calc_metal_ratio(log_x_abund)

# TODO Force [Fe/H] = 0 for calculating the cloud base
metallicity = 0.

# The try-except is required to catch numerical precision errors with the clouds
# try:
wlen_micron, flux_lambda, _ = retrieval_util.calc_spectrum_clouds(
Expand Down
6 changes: 3 additions & 3 deletions species/data/database.py
Original file line number Diff line number Diff line change
Expand Up @@ -2374,9 +2374,9 @@ def get_retrieval_teff(self,
flux_int = simps(sample_scale*box_item.flux, box_item.wavelength)
teff[i] = (flux_int/constants.SIGMA_SB)**0.25

np.savetxt(f'output/spectrum/spectrum{i:04d}.dat',
np.column_stack([box_item.wavelength, sample_scale*box_item.flux]),
header='Wavelength (um) - Flux (W m-2 um-1)')
# np.savetxt(f'output/spectrum/spectrum{i:04d}.dat',
# np.column_stack([box_item.wavelength, sample_scale*box_item.flux]),
# header='Wavelength (um) - Flux (W m-2 um-1)')

q_16, q_50, q_84 = np.percentile(teff, [16., 50., 84.])

Expand Down
63 changes: 53 additions & 10 deletions species/plot/plot_mcmc.py
Original file line number Diff line number Diff line change
Expand Up @@ -238,10 +238,16 @@ def plot_posterior(tag: str,
if 'H2O' in samples_box.parameters:
samples_box.parameters.append('c_h_ratio')
samples_box.parameters.append('o_h_ratio')
# samples_box.parameters.append('c_o_ratio')

ndim += 2

abund_index = {}
for i, item in enumerate(samples_box.parameters):
if item == 'CO':
if item == 'CH4':
abund_index['CH4'] = i

elif item == 'CO':
abund_index['CO'] = i

elif item == 'CO_all_iso':
Expand All @@ -250,24 +256,43 @@ def plot_posterior(tag: str,
elif item == 'CO2':
abund_index['CO2'] = i

elif item == 'CH4':
abund_index['CH4'] = i
elif item == 'FeH':
abund_index['FeH'] = i

elif item == 'H2O':
abund_index['H2O'] = i

elif item == 'H2S':
abund_index['H2S'] = i

elif item == 'Na':
abund_index['Na'] = i

elif item == 'NH3':
abund_index['NH3'] = i

elif item == 'H2S':
abund_index['H2S'] = i
elif item == 'K':
abund_index['K'] = i

elif item == 'PH3':
abund_index['PH3'] = i

elif item == 'TiO':
abund_index['TiO'] = i

elif item == 'VO':
abund_index['VO'] = 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 samples_box.parameters:
abund['CH4'] = item[abund_index['CH4']]

if 'CO' in samples_box.parameters:
abund['CO'] = item[abund_index['CO']]

Expand All @@ -277,19 +302,37 @@ def plot_posterior(tag: str,
if 'CO2' in samples_box.parameters:
abund['CO2'] = item[abund_index['CO2']]

if 'CH4' in samples_box.parameters:
abund['CH4'] = item[abund_index['CH4']]
if 'FeH' in samples_box.parameters:
abund['FeH'] = item[abund_index['FeH']]

if 'H2O' in samples_box.parameters:
abund['H2O'] = item[abund_index['H2O']]

if 'H2S' in samples_box.parameters:
abund['H2S'] = item[abund_index['H2S']]

if 'Na' in samples_box.parameters:
abund['Na'] = item[abund_index['Na']]

if 'NH3' in samples_box.parameters:
abund['NH3'] = item[abund_index['NH3']]

if 'H2S' in samples_box.parameters:
abund['H2S'] = item[abund_index['H2S']]
if 'K' in samples_box.parameters:
abund['K'] = item[abund_index['K']]

if 'PH3' in samples_box.parameters:
abund['PH3'] = item[abund_index['PH3']]

if 'TiO' in samples_box.parameters:
abund['TiO'] = item[abund_index['TiO']]

if 'VO' in samples_box.parameters:
abund['VO'] = item[abund_index['VO']]

if 'VO' in samples_box.parameters:
abund['VO'] = item[abund_index['VO']]

c_h_ratio[i], o_h_ratio[i] = retrieval_util.calc_metal_ratio(abund)
c_h_ratio[i], o_h_ratio[i], c_o_ratio[i] = retrieval_util.calc_metal_ratio(abund)

if vmr and samples_box.spectrum == 'petitradtrans' and \
'metallicity' not in samples_box.parameters:
Expand Down
61 changes: 46 additions & 15 deletions species/plot/plot_retrieval.py
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ def plot_pt_profile(tag: str,
print(f'Plotting the P-T profiles: {output}...', end='', flush=True)

cloud_species = ['Fe(c)', 'MgSiO3(c)', 'Al2O3(c)', 'Na2S(c)', 'KCl(c)']
cloud_check = ['fe', 'mgsio3', 'al2o3', 'na2s', 'kcl']
# cloud_check = ['fe', 'mgsio3', 'al2o3', 'na2s', 'kcl']

cloud_color = {'Fe(c)': 'tab:blue', 'MgSiO3(c)': 'tab:orange',
'Al2O3(c)': 'tab:green', 'Na2S(c)': 'tab:cyan',
Expand Down Expand Up @@ -95,21 +95,21 @@ def plot_pt_profile(tag: str,

ax = plt.subplot(gridsp[0, 0])

top = True

for item in cloud_check:
if f'{item}_fraction' in median:
top = False

elif f'{item}_tau' in median:
top = False
# top = True
#
# for item in cloud_check:
# if f'{item}_fraction' in median:
# top = False
#
# elif f'{item}_tau' in median:
# top = False

ax.tick_params(axis='both', which='major', colors='black', labelcolor='black',
direction='in', width=1, length=5, labelsize=12, top=top,
direction='in', width=1, length=5, labelsize=12, top=True,
bottom=True, left=True, right=True)

ax.tick_params(axis='both', which='minor', colors='black', labelcolor='black',
direction='in', width=1, length=3, labelsize=12, top=top,
direction='in', width=1, length=3, labelsize=12, top=True,
bottom=True, left=True, right=True)

ax.set_xlabel('Temperature (K)', fontsize=13)
Expand Down Expand Up @@ -151,15 +151,36 @@ def plot_pt_profile(tag: str,
knot_press = np.logspace(np.log10(pressure[0]), np.log10(pressure[-1]), 15)

for i, item in enumerate(samples):
# C/O and [Fe/H]

if box.attributes['chemistry'] == 'equilibrium':
metallicity = item[param_index['metallicity']]
c_o_ratio = item[param_index['c_o_ratio']]

elif box.attributes['chemistry'] == 'free':
# TODO Set [Fe/H] = 0
metallicity = 0.

# Create a dictionary with the mass fractions

log_x_abund = {}

for i in range(box.attributes['n_line_species']):
line_item = box.attributes[f'line_species{i}']
log_x_abund[line_item] = item[param_index[line_item]]

# 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)

if pt_profile == 'molliere':
t3_param = np.array([item[param_index['t1']],
item[param_index['t2']],
item[param_index['t3']]])

temp, _ = retrieval_util.pt_ret_model(
t3_param, 10.**item[param_index['log_delta']], item[param_index['alpha']],
item[param_index['tint']], pressure, item[param_index['metallicity']],
item[param_index['c_o_ratio']])
item[param_index['tint']], pressure, metallicity, c_o_ratio)

elif pt_profile == 'free':
knot_temp = []
Expand All @@ -182,6 +203,11 @@ def plot_pt_profile(tag: str,
# np.column_stack([pressure, temp]),
# header='Pressure (bar) - Temperature (K)')

if box.attributes['chemistry'] == 'free':
# TODO Set [Fe/H] = 0
median['metallicity'] = metallicity
median['c_o_ratio'] = c_o_ratio

if pt_profile == 'molliere':
temp, _ = retrieval_util.pt_ret_model(
np.array([median['t1'], median['t2'], median['t3']]), 10.**median['log_delta'],
Expand All @@ -206,7 +232,10 @@ def plot_pt_profile(tag: str,

ax.plot(temp, pressure, '-', lw=1, color='black', zorder=2)

if 'metallicity' in parameters and 'c_o_ratio' in parameters:
# Add cloud condensation profiles

if extra_axis == 'grains' and 'metallicity' in median and 'c_o_ratio' in median:

if 'log_p_quench' in median:
quench_press = 10.**median['log_p_quench']
else:
Expand Down Expand Up @@ -589,7 +618,9 @@ def plot_clouds(tag: str,
box = species_db.get_samples(tag)
median = box.median_sample

if f'{composition.lower()}_fraction' not in median and 'log_tau_cloud' not in median:
if f'{composition.lower()}_fraction' not in median and \
'log_tau_cloud' not in median and f'{composition}(c)' not in median:

raise ValueError(f'The mass fraction of the {composition} clouds is not found. The median '
f'sample contains the following parameters: {list(median.keys())}')

Expand Down
Loading

0 comments on commit 14ceec0

Please sign in to comment.