Retrievals: Dealing with multiple datasets
An advanced retrieval tutorial.
Written by Evert Nasedkin and Paul Mollière
This tutorial will use JWST observations of WASP39 b to demonstrate how to incorporate multiple datasets into the pRT retrieval framework.
Getting started
You should already have an installation of pRT on your machine, please see the installation manual if you have not installed it yet. pyMultiNest is also required. Ultranest is required if you want to use that as your sampler rather than pyMultiNest. See the Ultranest documentation for why you might want to choose this method. Using nested sampling rather than MCMC is faster, handles multimodal cases better, and directly provides the Bayesian evidence, which allows for easier model comparison.
In this tutorial, we will outline the process of setting up a RetrievalConfig object, which is the class used to set up a pRT retrieval. The basic process will always be to set up the configuration, and then pass it to the Retrieval class to run the retrieval using pyMultiNest. Several standard plotting outputs will also be produced by the retrieval class. Most of the classes and functions used in this tutorial have more advanced features than what will be explained here, so it’s highly recommended to take a look at the code and API documentation. There should be enough flexibility built in to cover most typical retrieval studies, but if you have feature requests please get in touch, or open an issue on gitlab.
[1]:
# Let's start by importing everything we need.
import os
# To not have numpy start parallelizing on its own
os.environ["OMP_NUM_THREADS"] = "1"
#os.environ["pRT_input_data_path"] = "/path/to/petitRADTRANS/petitRADTRANS/input_data/"
import numpy as np
import matplotlib.pyplot as plt
from petitRADTRANS import Radtrans
from petitRADTRANS import nat_cst as nc
# Import the class used to set up the retrieval.
from petitRADTRANS.retrieval import Retrieval,RetrievalConfig
# Import Prior functions, if necessary.
from petitRADTRANS.retrieval.util import gaussian_prior
# Import atmospheric model function
from petitRADTRANS.retrieval.models import guillot_patchy_transmission
Lets start out by setting up a simple run definition. We’ll add the data AFTER we define the model function below Full details of the parameters can be found in the API documentation.
[2]:
# Lets start out by setting up a simple run definition
# Full details of the parameters can be found in retrieval_config.py
# Since our retrieval has already ran before, we'll set the mode to 'evaluate' so we can make some plots.
RunDefinition = RetrievalConfig(retrieval_name = "WASP39b_Guillot_FreeChem_PatchyGreyHaze",
run_mode = "retrieval", # This must be 'retrieval' to run PyMultiNest
pressures = np.logspace(-8,3,100), # Extend up to 10^-8 bar
AMR = False, # We won't be using adaptive mesh refinement for the pressure grid
scattering = False) # This would turn on scattering when calculating emission spectra.
# Scattering is automatically included for transmission spectra.
Data
Let’s start with reading in the data. The data must be a 1D spectrum with error bars or a covariance matrix.
As in the basic tutorial, we’re reading in text files, but this time we also include the column that describes the wavelength bins:
# Wavelength [micron], Bins [micron], Flux [W/m2/micron or (Rp/Rstar)^2], Flux Error [W/m2/micron or (Rp/Rstar)^2]
The Data class is arguably the most important part of setting up the retrieval. Not only do you input your data here, but you also choose your model function and resolution. This means that you can design a retrieval around different datatypes and retrieve simultaneously on both - for example, if you want the day and nightside of a planet, or want to combine the eastward and westward limbs of a transmission spectrum with different models.
You can also set a distance to your object, which will allow you to automatically scale the flux and error of your data using the scale_to_distance() method - useful if you have data normalized to 10pc! Finally, there’s also a arguments scale, scale_err and offset_bool, which tells the retrieval that the flux or uncertaines should be scaled by an arbitrary multiplicative factor or have an additive offset, both which is set up as a normal retrieval parameter using the
RetrivalConfig.add_parameter() method. The name must be of the format DATANAME_scale_factor or DATANAME_offset. This is useful if two datasets are incompatible in absolute photometry, but you still want to use the spectral shape to inform the retrieval.
In this retrieval we’re going to include several datasets from different JWST instruments, starting with NIRISS SOSS orders 1 and 2. To include both of them , we simply add more than one dataset to our RetrievalConfig object. Notice that we’re also telling the retrieval that we want the order 2 data to have an additive offset: this will allow the data to float relative to order one, which remains fixed. This can be used to compensate for differences in transit depth between different
instruments.
We’re also using the buillt in guillot_patchy_transmission model for this retrieval, rather than writing our own model function from scratch.
[3]:
import petitRADTRANS # need to get the name for the example data
path_to_data = ""
path_to_data = petitRADTRANS.__file__.split('__init__.py')[0] # Default location for the example
transmission_directory = "retrieval/examples/transmission/"
RunDefinition.add_data('JWST/NIRISSSOSS/O1',
f"{path_to_data}{transmission_directory}observations/JWST/WASP39b_niriss_soss1.txt",
data_resolution=700,
model_resolution=300,
model_generating_function = guillot_patchy_transmission,
external_pRT_reference=None)
RunDefinition.add_data('JWST/NIRISSSOSS/O2',
f"{path_to_data}{transmission_directory}observations/JWST/WASP39b_niriss_soss2.txt",
data_resolution=700,
model_resolution=300,
offset_bool=True,
model_generating_function = guillot_patchy_transmission,
external_pRT_reference = None)
External references
Sometimes the datasets will include regions that overlap in wavelength, or where one dataset falls entirely within the wavelength range of another dataset. In that case we can use an external reference: the dataset that falls within the wavelength range of the other will use the spectrum computed using the other dataset’s Radtrans to compute the likelihood. This way we only need to compute the spectrum once in the same wavelength range, rather than initialising two Radtrans and
calculating a spectrum for each. However, be careful here: the model resolution in the reference object should be high enough to properly sample any datasets that reference it!
In this example, the NIRSpec PRISM data covers the entire NIRISS SOSS wavelength range, so we can use it as a reference for both NIRISS SOSS orders.
[4]:
RunDefinition.data = {} # Remove the previous data that was added above.
RunDefinition.add_data('JWST/NIRSPEC/PRISM',
f"{path_to_data}{transmission_directory}observations/JWST/WASP39b_nirspec_prism.txt",
data_resolution=100,
model_resolution=300,
offset_bool=True,
model_generating_function = guillot_patchy_transmission)
RunDefinition.add_data('JWST/NIRISSSOSS/O1',
f"{path_to_data}{transmission_directory}observations/JWST/WASP39b_niriss_soss1.txt",
data_resolution=700,
model_resolution=300,
model_generating_function = guillot_patchy_transmission,
external_pRT_reference='JWST/NIRSPEC/PRISM')
RunDefinition.add_data('JWST/NIRISSSOSS/O2',
f"{path_to_data}{transmission_directory}observations/JWST/WASP39b_niriss_soss2.txt",
data_resolution=700,
model_resolution=300,
offset_bool=True,
model_generating_function = guillot_patchy_transmission,
external_pRT_reference = 'JWST/NIRSPEC/PRISM')
Model parameters
Here we’re using a more complicated atmospheric model to fit the JWST data. The temperature profile is taken from Guillot 2010, and includes four parameters to describe the shape. We’re freely retrieving the chemical abundances, and include both patchy grey clouds and an enhanced power law slope as a proxy for hazes.
[5]:
# WASP 39 parameters
RunDefinition.add_parameter(name = 'D_pl', free = False, value = 230.0*nc.pc)
RunDefinition.add_parameter(name = 'Rstar', free = False, value = 0.9324 * nc.r_sun)
# Fix the reference pressure in bar
RunDefinition.add_parameter('reference_pressure',False,value = 0.01)
# Choose two of log_g, radius and mass priors
RunDefinition.add_parameter('log_g',True,
transform_prior_cube_coordinate = \
lambda x : 2.0+3.5*x)
RunDefinition.add_parameter('R_pl', True,
transform_prior_cube_coordinate = \
lambda x : 0.8*nc.r_jup_mean+ (x*0.8*nc.r_jup_mean))
# Priors for Guillot 2010 Temperature Profile
RunDefinition.add_parameter("T_int", True,\
transform_prior_cube_coordinate = \
lambda x : 100 + 3500*x)
RunDefinition.add_parameter("T_equ", True,\
transform_prior_cube_coordinate = \
lambda x : 100 + 3500*x)
RunDefinition.add_parameter("gamma", True,\
transform_prior_cube_coordinate = \
lambda x : 10**(-(x/2.)**2./2.))
RunDefinition.add_parameter("log_kappa_IR", True,\
transform_prior_cube_coordinate = \
lambda x : -4.0 + 6.0*x)
# Grey cloud top pressure
RunDefinition.add_parameter('log_Pcloud',True,
transform_prior_cube_coordinate = \
lambda x : -8.+11.*x)
# Enhanced haze scattering slope
RunDefinition.add_parameter('kappa_0',True,
transform_prior_cube_coordinate = \
lambda x : -4.+14.*x)
RunDefinition.add_parameter('gamma_scat',True,
transform_prior_cube_coordinate = \
lambda x : -20.+22.*x)
# Cloud fraction
RunDefinition.add_parameter('patchiness',True,
transform_prior_cube_coordinate = \
lambda x : x)
# Data offsets
RunDefinition.add_parameter('JWST/NIRSPEC/PRISM_offset',True,
transform_prior_cube_coordinate = \
lambda x : gaussian_prior(x, 0, 1e-4))
RunDefinition.add_parameter('JWST/NIRISSSOSS/O2_offset',True,
transform_prior_cube_coordinate = \
lambda x : gaussian_prior(x, 0, 1e-4))
Opacities
[6]:
RunDefinition.set_rayleigh_species(['H2', 'He'])
RunDefinition.set_continuum_opacities(['H2-H2', 'H2-He'])
# Here we setup the line species for a free retrieval, setting the prior bounds with the abund_lim parameter
# The retrieved value is the log mass fraction.
# Let's use the most up-to-date line lists
RunDefinition.set_line_species(["H2O_Exomol", "CO_all_iso_HITEMP", "CO2", "CH4"], eq=False, abund_lim = (-8.0,0.0))
Plotting
Let’s set up some plotting details so that we can generate nice outputs.
Each parameter can be added to the corner plot, its label changed, and the values transformed to more digestible units (e.g., the planet radius in jupiter radii, rather than cm). We can also set the bounds and scaling on the best-fit spectrum plot and the limits for the P-T profile plot. With this complete, our retrieval is ready to go.
Most parameters include a default setting, so the plots will be created even if you don’t set any plotting parameters, but the outputs might not be very informative. In general, the possible arguments to plot_kwargs follow the naming conventions of matplotlib arguments and functions, with some additions. Full details of the plotting can be found in the Retrieval class, see the API documentation.
[7]:
##################################################################
# Define what to put into corner plot if run_mode == 'evaluate'
##################################################################
for key, value in RunDefinition.parameters.items():
value.plot_in_corner = True
value.corner_label = key.replace("_"," ")
##################################################################
# Define axis properties of spectral plot if run_mode == 'evaluate'
##################################################################
RunDefinition.plot_kwargs["spec_xlabel"] = 'Wavelength [micron]'
RunDefinition.plot_kwargs["spec_ylabel"] = r'$(R_{\rm P}/R_*)^2$ [ppm]'
RunDefinition.plot_kwargs["y_axis_scaling"] = 1e6 # so we have units of ppm
RunDefinition.plot_kwargs["xscale"] = 'linear'
RunDefinition.plot_kwargs["yscale"] = 'linear'
# Use at least ~100 samples to plot 3 sigma curves
RunDefinition.plot_kwargs["nsample"] = 10
##################################################################
# Define from which observation object to take P-T
# in evaluation mode (if run_mode == 'evaluate'),
# add PT-envelope plotting options
##################################################################
RunDefinition.plot_kwargs["take_PTs_from"] = 'JWST/NIRSPEC/PRISM'
RunDefinition.plot_kwargs["temp_limits"] = [150, 3000]
RunDefinition.plot_kwargs["press_limits"] = [1e1, 1e-7]
# If in doubt, define all of the plot_kwargs used here.
Running the retrieval
At this point, we are ready to run the retrieval! All we need to do is pass the RunDefinition we just created to the Retrieval class and call its run() method. There are a few additional parameters that can be adjusted, but the defaults should work well for almost all use cases. In general it may not be wise to run a retrieval on your laptop, as it can be quite computationally expensive. However, the included HST example should be able to be run fairly easily! (~an hour).
NOTE: This retrieval example only uses 40 live points so that it can be run quickly. However, this should be increased to at least 400 live points for a full retrieval, and will probably require a server or cluster.
Once the retrieval is complete, we can use the plot_all function to generate plots of the best fit spectrum, the pressure-temperature profile and the corner plots. Most of the various parameters used to control pyMultiNest or Ultranest can be set in the retrieval.run() function, see the API documentation.
[ ]:
# If you want to run the retrieval, you need to choose a different output directory name,
# due to pRT requirements.
output_dir = f"./"
retrieval = Retrieval(RunDefinition,
output_dir = output_dir,
sample_spec = False, # Output the spectrum from nsample random samples.
pRT_plot_style=True, # We think that our plots look nice.
ultranest=False) # Let's use pyMultiNest rather than Ultranest
# Default pymultinest retrieval setup, but with only 40 live points
retrieval.run(sampling_efficiency=0.8,
const_efficiency_mode=False,
n_live_points=40,
log_z_convergence=0.5,
step_sampler=False,
warmstart_max_tau=0.5,
n_iter_before_update=50,
resume=False,
max_iters=0,
frac_remain=0.1,
importance_nested_sampling = True,
Lepsilon=0.3)
[ ]:
# Automatically generate all of the standard output plots. The contribution
# argument means that the PT profile and abundance profile plots will display
# the contribution function. The mode argument means we'll be plotting a model
# based on the median retrieved parameter values, rather than the minimum likelihood
# model.
retrieval.plot_all(contribution = True, mode = 'median')
Contact
If you need any additional help, don’t hesitate to contact Evert Nasedkin.
[ ]:
[ ]: