{
"cells": [
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"# High-Resolution Transmission Retrievals\n",
"\n",
"This tutorial covers how to run a retrieval with high-resolution ground-based transmission spectrum data using POSEIDON. \n",
"\n",
"Before you run this notebook, you should first run the [\\\"Ground-Based High-Resolution Transmission Spectroscopy (Cross Correlation)\\\"](transmission_high_res_cross_correlate.html) tutorial to preprocess the WASP-121b data. If you have data_processed.hdf5 saved in your planet directory, you are all set!"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Loading WASP-121b Transmission Data\n",
"\n",
"First, we will load the processed data for your planet (here, WASP-121b). For more information about this dataset and to learn the basics of high-resolution cross correlation spectroscopy, see the [\\\"Ground-Based High-Resolution Transmission Spectroscopy (Cross Correlation)\\\"](transmission_high_res_cross_correlate.html) tutorial."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from POSEIDON.high_res import read_high_res_data\n",
"\n",
"planet_name = 'WASP-121b'\n",
"\n",
"data_dir = '../../../POSEIDON/reference_data/observations/' + planet_name # The directory where you've put the data\n",
"\n",
"data = read_high_res_data(data_dir, names=[\"blue\"]) # only use blue arm for faster retrieval\n",
"data = read_high_res_data(data_dir, names=[\"blue\", \"redl\", \"redu\"])"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Creating a Retrieval Model\n",
"\n",
"Now, let's provide the wavelength grid and properties of the host star and your planet. The wavelength range should match the range of your data. The blue arm spans 0.37 microns to 0.51 microns. If you decide to use both blue and red arms, you should increase the range to 0.37 microns to 0.87 microns.\n",
"\n",
"We use R=250,000 as a tradeoff between computational speed and accuracy. For more discussion, see the previous tutorial."
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from POSEIDON.core import define_model, wl_grid_constant_R\n",
"from POSEIDON.core import create_star, create_planet\n",
"from POSEIDON.constants import R_Sun, R_J, M_J\n",
"\n",
"# ***** Wavelength grid *****#\n",
"\n",
"wl_min = 0.37 # Minimum wavelength (um)\n",
"wl_max = 0.51 # Maximum wavelength (um) for blue arm\n",
"wl_max = 0.87 # change to include red arm\n",
"R = 250000 # Spectral resolution of grid \n",
"\n",
"# Create a wavelength grid with constant R\n",
"wl = wl_grid_constant_R(wl_min, wl_max, R)\n",
"\n",
"# ***** Define stellar properties *****#\n",
"\n",
"R_s = 1.458 * R_Sun # Stellar radius (m)\n",
"T_s = 6776 # Stellar effective temperature (K) \n",
"Met_s = 0.13 # Stellar metallicity [log10(Fe/H_star / Fe/H_solar)] \n",
"log_g_s = 4.24 # Stellar log surface gravity (log10(cm/s^2) by convention) \n",
"\n",
"star = create_star(R_s, T_s, log_g_s, Met_s, wl=wl, stellar_grid=\"phoenix\")\n",
"\n",
"# ***** Define planet properties *****#\n",
"\n",
"planet_name = \"WASP-121b\" # Planet name used for plots, output files etc.\n",
"\n",
"R_p = 1.753 * R_J # Planetary radius (m) \n",
"M_p = 1.157 * M_J # Mass of planet (kg) \n",
"\n",
"# Create the planet object\n",
"planet = create_planet(planet_name, R_p, mass=M_p)\n",
"\n",
"# If distance not specified, use fiducial value\n",
"if planet[\"system_distance\"] is None:\n",
" planet[\"system_distance\"] = 1 # This value only used for flux ratios, so it cancels\n",
"d = planet[\"system_distance\"]"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"Existing literature have shown detection of $\\rm{Fe}$ in the atmosphere of WASP-121b. There are strong $\\rm{Fe}$ absorption features in the wavelength range as well.\n",
"\n",
"So for a first attempt, we consider a model with $\\rm{Fe}$, an isothermal temperature profile, and no clouds.\n",
"\n",
"For additional parameters used in high resolution retrieval, we include: $a$ (the scale parameter), $b$ (the scale parameter for noise), $K_p$ (the Keplerian orbital velocity), $V_{sys}$ (the systematic velocity), and $W_{conv}$ (width of the gaussian convolution kernel used for line broadening). You can opt to use the MLE estimator of $\\beta$ and not include it as a free parameter, which we are going to do here. [Gibson et al. 2022](https://ui.adsabs.harvard.edu/abs/2022MNRAS.512.4618G/abstract) contains a discussion on this choice. An additional parameter available is $\\Delta \\phi$, which offsets the ephemeris. However $\\Delta \\phi$ is very degenerate with $V_{sys}$ if the range of covered orbital phase is small.\n",
"\n",
"Be sure to reference [Gibson et al. 2022](https://ui.adsabs.harvard.edu/abs/2022MNRAS.512.4618G/abstract) if you want a more detailed description of these parameters."
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Free parameters: ['R_p_ref' 'T' 'log_Fe' 'log_Cr' 'log_Mg' 'log_V' 'log_Ti' 'K_p' 'V_sys'\n",
" 'W_conv' 'log_alpha_HR']\n"
]
}
],
"source": [
"# ***** Define model *****#\n",
"\n",
"model_name = \"High-res retrieval\" # Model name used for plots, output files etc.\n",
"\n",
"bulk_species = [\"H2\", \"He\"] # H2 + He comprises the bulk atmosphere\n",
"param_species = [\"Fe\", \"Cr\", \"Mg\", \"V\", \"Ti\"] # Add more chemical species to the model here\n",
"\n",
"method = \"sysrem\"\n",
"\n",
"# Create the model object\n",
"model = define_model(model_name, bulk_species, param_species,\n",
" PT_profile = \"isotherm\", reference_parameter = \"R_p_ref\",\n",
" high_res_method = \"sysrem\", # Important! Should be the same as the method used to preprocess the data\n",
" alpha_high_res_option = 'log', \n",
" fix_alpha_high_res = False, fix_W_conv_high_res = False, \n",
" fix_beta_high_res = True, fix_Delta_phi_high_res = True,\n",
" )\n",
"\n",
"# Check the free parameters defining this model\n",
"print(\"Free parameters: \" + str(model[\"param_names\"]))"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Setting Retrieval Priors\n",
"\n",
"One of the most important aspects in any Bayesian analysis is deciding what priors to use for the free parameters. Specifying a prior has two steps: (i) choosing the type of probability distribution; and (ii) choosing the allowable range.\n",
"\n",
"Most free parameters in atmospheric retrievals with POSEIDON use the following prior types:\n",
"\n",
"- Uniform: you provide the minimum and maximum values for the parameter.\n",
"- Gaussian: you provide the mean and standard deviation for the parameter.\n",
"\n",
"\n",
"\n",
" **Note:**\n",
"\n",
" If you do not specify a prior type or range for a given parameter, POSEIDON will ascribe a default prior type (generally uniform) and a 'generous' range.\n",
"\n",
"
\n",
"\n",
"\n",
"Your first retrieval is defined by seven free parameters printed above: (1) the radius at the (fixed) reference pressure; (2) the isothermal atmospheric temperature; (3) the log-mixing ratio of $\\rm{Fe}$; and (4) the four high resolution parameters. \n",
"\n",
"Since we are assuming no *a priori* information on WASP-121b's atmosphere, we will use uniform priors for all the parameters."
]
},
{
"cell_type": "code",
"execution_count": 4,
"metadata": {},
"outputs": [],
"source": [
"from POSEIDON.core import set_priors\n",
"\n",
"# ***** Set priors for retrieval *****#\n",
"\n",
"# Initialise prior type dictionary\n",
"prior_types = {}\n",
"\n",
"# Specify whether priors are linear, Gaussian, etc.\n",
"prior_types[\"T\"] = \"uniform\"\n",
"prior_types[\"R_p_ref\"] = \"gaussian\"\n",
"prior_types[\"log_X\"] = \"uniform\"\n",
"prior_types[\"K_p\"] = \"uniform\"\n",
"prior_types[\"V_sys\"] = \"uniform\"\n",
"prior_types[\"log_alph_HR\"] = \"uniform\"\n",
"prior_types[\"beta_HR\"] = \"uniform\"\n",
"prior_types[\"W_conv\"] = \"uniform\"\n",
"\n",
"# Initialise prior range dictionary\n",
"prior_ranges = {}\n",
"\n",
"# Specify prior ranges for each free parameter\n",
"prior_ranges[\"T\"] = [1000, 4000]\n",
"prior_ranges[\"R_p_ref\"] = [R_p, 0.05*R_J]\n",
"prior_ranges[\"log_X\"] = [-15, 0]\n",
"prior_ranges[\"K_p\"] = [170, 230]\n",
"prior_ranges[\"V_sys\"] = [-10, 10]\n",
"prior_ranges[\"log_alpha_HR\"] = [-1, 2]\n",
"prior_ranges[\"beta_HR\"] = [0.1, 10]\n",
"prior_ranges[\"W_conv\"] = [1, 50]\n",
"\n",
"# Create prior object for retrieval\n",
"priors = set_priors(planet, star, model, data, prior_types, prior_ranges)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Pre-load Opacities\n",
"\n",
"The last step before running a retrieval is to pre-interpolate the cross sections for our model and store them in memory. For more details on this process, see the forward model tutorial.\n",
"\n",
"\n",
"\n",
" **Warning:**\n",
"\n",
" Ensure the range of $T_{\\rm{fine}}$ used for opacity pre-interpolation is at least as large as the desired prior range for temperatures to be explored in the retrieval. Any models with layer temperatures falling outside the range of $T_{\\rm{fine}}$ will be automatically rejected (for retrievals with non-isothermal P-T profiles, this prevents unphysical profiles with negative temperatures etc.)\n",
"\n",
"
"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from POSEIDON.core import read_opacities\n",
"import numpy as np\n",
"\n",
"# ***** Read opacity data *****#\n",
"\n",
"opacity_treatment = \"opacity_sampling\"\n",
"\n",
"# Define fine temperature grid (K)\n",
"T_fine_min = 1000 # 400 K lower limit suffices for a typical hot Jupiter\n",
"T_fine_max = 4000 # 2000 K upper limit suffices for a typical hot Jupiter\n",
"T_fine_step = 50 # 20 K steps are a good tradeoff between accuracy and RAM\n",
"\n",
"T_fine = np.arange(T_fine_min, (T_fine_max + T_fine_step), T_fine_step)\n",
"\n",
"# Define fine pressure grid (log10(P/bar))\n",
"log_P_fine_min = -12.0 # 1 ubar is the lowest pressure in the opacity database\n",
"log_P_fine_max = 2 # 100 bar is the highest pressure in the opacity database\n",
"log_P_fine_step = 0.2 # 0.2 dex steps are a good tradeoff between accuracy and RAM\n",
"\n",
"log_P_fine = np.arange(log_P_fine_min, (log_P_fine_max + log_P_fine_step), log_P_fine_step)\n",
"\n",
"# Now we can pre-interpolate the sampled opacities (may take up to a minute)\n",
"opac = read_opacities(model, wl, opacity_treatment, T_fine, log_P_fine)"
]
},
{
"attachments": {},
"cell_type": "markdown",
"metadata": {},
"source": [
"## Run Retrieval\n",
"\n",
"You are now ready to run your high resolution atmospheric retrieval on this dataset!\n",
"\n",
"\n",
"\n",
" **Tip:**\n",
"\n",
" Retrievals run faster on multiple cores. When running the cells in this Jupyter notebook, only a single core will be used. You can run a multi-core retrieval on 24 cores by converting this Jupyter notebook into a python script, then calling mpirun on the .py file:\n",
"\n",
" ```\n",
" mpirun -n 24 python -u YOUR_RETRIEVAL_SCRIPT.py\n",
" ```\n",
" \n",
"
\n",
"\n",
"\n",
"\n",
"\n",
" **Important Note:**\n",
" A high resolution forward model is computationally expensive (~1 second per model). With 400 live points, it takes ~100,000 evaluations for the model to converge. With 36 cores, this amounts to ~1 hour. \n",
" \n",
" Instead of waiting until the end of time for the next cell to finish, you could run the 'transmission_high_res_retrieval.py' file in this folder, which is the same code converted from this notebook, and parallelise with multiple cores in command line. \n",
" \n",
" To check the code is working before launching a high-res retrieval, you can run the cell below and wait for a couple of minutes. Once it says \"live points generated\" and still no error, you are good to run it on multiple cores!\n",
" \n",
"
"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from POSEIDON.retrieval import run_retrieval\n",
"\n",
"# ***** Specify fixed atmospheric settings for retrieval *****#\n",
"\n",
"# Atmospheric pressure grid\n",
"P_min = 1e-12 # 1 pbar\n",
"P_max = 100 # 100 bar\n",
"N_layers = 100 # 100 layers\n",
"\n",
"# Let's space the layers uniformly in log-pressure\n",
"P = np.logspace(np.log10(P_max), np.log10(P_min), N_layers)\n",
"\n",
"# Specify the reference pressure and radius\n",
"P_ref = 1e-2 # Reference pressure (bar)\n",
"\n",
"# ***** Run atmospheric retrieval *****#\n",
"\n",
"run_retrieval(planet, star, model, opac, data, priors, wl, P, P_ref, R_p_ref = R_p,\n",
" R = R, spectrum_type = \"transmission\", sampling_algorithm = \"MultiNest\",\n",
" N_live = 400, verbose = True, N_output_samples = 1000, \n",
" resume = False)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Plot Retrieval Results"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Generate a corner plot after the retrieval is finished\n",
"from POSEIDON.corner import generate_cornerplot\n",
"\n",
"fig_corner = generate_cornerplot(planet, model)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"# Read retrieved PT profile and plot it\n",
"from POSEIDON.utility import read_retrieved_PT\n",
"from POSEIDON.visuals import plot_PT_retrieved\n",
"\n",
"# Read the retrieved PT profile\n",
"P, T_low2, T_low1, T_median, \\\n",
"T_high1, T_high2 = read_retrieved_PT(planet_name, model_name)\n",
"\n",
"PT_median = [(T_median, P)]\n",
"PT_low2 = [(T_low2, P)]\n",
"PT_low1 = [(T_low1, P)]\n",
"PT_high1 = [(T_high1, P)]\n",
"PT_high2 = [(T_high2, P)]\n",
"\n",
"# Plot the retrieved PT profile\n",
"plot_PT_retrieved(planet_name, PT_median, PT_low2, PT_low1, PT_high1, PT_high2,\n",
" # T_true=None, # Uncomment this line if you have a PT profile to compare to\n",
" # # colour_list=[], # Uncomment this line if you want to specify colors\n",
" T_min=2000, T_max=4000,\n",
" legend_location=\"lower left\"\n",
" )"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Below is the corner plot and retrieved PT profile from a retrieval on this dataset.\n",
"\n",
""
]
}
],
"metadata": {
"kernelspec": {
"display_name": "POSEIDON_python_3.11",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.9"
},
"orig_nbformat": 4
},
"nbformat": 4,
"nbformat_minor": 2
}