Retrievals: Warm-Started, Multi-Chain Emission Retrievals with NumPyro NUTS#

Written by Evert Nasedkin. Please cite pRT’s retrieval package (Nasedkin et al. 2024) in addition to pRT (Mollière et al. 2019) if you make use of the retrieval package for your work.

This advanced tutorial builds on the NumPyro NUTS emission retrieval tutorial, which you should work through first. We reuse the same HR8799e forward model (the differentiable molliere_2020_emission implementation of Mollière et al. 2020) and the same data, and demonstrate two advanced features of pRT’s gradient-based retrieval path:

Optimal estimation warm-starts. Before sampling, an optax gradient optimisation locates the maximum-a-posteriori (MAP) point, and the Hessian of the log-posterior at that optimum — the Laplace, or Rodgers optimal estimation, approximation of the posterior covariance — is used to precondition the NUTS mass matrix. This starts the chains inside the bulk of the posterior and shortens warmup adaptation.
Multi-chain sampling. Running several NUTS chains, either vectorized inside a single process or as independent chains orchestrated by Retrieval.run_multichain, and combining them into one posterior with convergence diagnostics (split-R-hat and effective sample size).

As in the base tutorial, emission retrievals with scattering and multiple datasets are computationally expensive: this notebook outlines the setup, but we advise running production chains on a workstation or cluster.

Getting started#

Prerequisites: the “Basic Retrieval Tutorial” and the NumPyro NUTS emission retrieval tutorial. The forward model, data, priors, and opacity setup below are identical to the base tutorial and are not re-explained here.

One structural difference: instead of building the RetrievalConfig inline, we wrap the whole configuration in a top-level build_config() function. This is the pattern required by Retrieval.run_multichain, which constructs a fresh configuration in each chain worker — and for the "processes" execution mode the builder must be picklable by reference (a top-level function in an importable module, not a closure or lambda). Defining it this way from the start lets us use the same configuration for the single-process warm-started run and for the multi-chain orchestration.

[ ]:

import os
# To not have numpy start parallelizing on its own
os.environ["OMP_NUM_THREADS"] = "1"

from jax import config
config.update("jax_enable_x64", True)

import tensorflow_probability.substrates.jax as tfp
from petitRADTRANS.retrieval import Retrieval, RetrievalConfig
from petitRADTRANS.retrieval.models import molliere_2020_emission
from petitRADTRANS import physical_constants as cst

tfpd = tfp.distributions

Retrieval configuration#

The builder below reproduces the configuration of the NUTS emission tutorial: the GRAVITY spectrum plus photometry, the Mollière et al. (2020) temperature/chemistry/cloud parameterisation with TensorFlow Probability priors, and the same opacity sources. See the base tutorial for a cell-by-cell discussion of each ingredient.

[ ]:

def build_config():
    """Build the HR8799e retrieval configuration.

    Defined as a top-level function so that ``Retrieval.run_multichain`` can
    construct a fresh configuration in every chain worker. For the
    ``chain_execution="processes"`` mode this function must live in an
    importable module (a script run under ``if __name__ == "__main__":``,
    not a notebook); the in-notebook modes used below work as-is.
    """
    retrieval_config = RetrievalConfig(
        retrieval_name="HR8799e_advanced_example",
        run_mode="retrieval",  # 'retrieval' to run, or 'evaluate' to make plots
        sampler_type="numpyro_nuts",
        adaptive_mesh_refinement=True,  # adaptive mesh refinement, slower if True
        scattering_in_emission=True,  # add scattering for emission spectra clouds
    )

    # --- Data -------------------------------------------------------------
    path_to_data = "./"

    retrieval_config.add_data(
        name="GRAVITY",
        path_to_observations=f"{path_to_data}/retrievals/emission/observations/HR8799e_Spectra.fits",
        data_resolution=500,
        model_resolution=1000,
        model_generating_function=molliere_2020_emission,
    )

    retrieval_config.add_photometry(
        path_to_data + "retrievals/emission/observations/hr8799e_photometry.txt",
        molliere_2020_emission,
        model_resolution=40,
    )

    # --- Parameters and priors ---------------------------------------------
    # Distance to the planet in cm
    retrieval_config.add_parameter(
        name="system_distance", is_free_parameter=False, value=41.2925 * cst.pc
    )

    # Log of the surface gravity in cgs units.
    retrieval_config.add_parameter(
        "log_g",
        True,
        distribution=tfpd.Uniform(low=2.0, high=5.5),
    )

    # Planet radius in cm
    retrieval_config.add_parameter(
        "planet_radius",
        True,
        distribution=tfpd.Uniform(low=0.7 * cst.r_jup_mean, high=2.0 * cst.r_jup_mean),
    )

    # Interior temperature in Kelvin
    retrieval_config.add_parameter(
        "T_int",
        True,
        distribution=tfpd.Uniform(low=300.0, high=2300.0),
        value=0.0,
    )

    # Spline temperature structure parameters. T1 < T2 < T3
    # As these priors depend on each other, they are implemented in the model function
    retrieval_config.add_parameter("T3", True, distribution=tfpd.Uniform(low=0.0, high=1.0), value=0.0)
    retrieval_config.add_parameter("T2", True, distribution=tfpd.Uniform(low=0.0, high=1.0), value=0.0)
    retrieval_config.add_parameter("T1", True, distribution=tfpd.Uniform(low=0.0, high=1.0))
    # Optical depth model: tau = delta * press_cgs**alpha
    retrieval_config.add_parameter("alpha", True, distribution=tfpd.Uniform(low=1.0, high=2.0))
    retrieval_config.add_parameter("log_delta", True, distribution=tfpd.Uniform(low=0.0, high=1.0))
    # Chemistry: (dis)equilibrium model bulk parameters
    retrieval_config.add_parameter("log_pquench", True, distribution=tfpd.Uniform(low=-6.0, high=3.0))
    retrieval_config.add_parameter("Fe/H", True, distribution=tfpd.Uniform(low=-1.5, high=1.5))
    retrieval_config.add_parameter("C/O", True, distribution=tfpd.Uniform(low=0.1, high=1.6))
    # Ackermann-Marley (2001) cloud model
    retrieval_config.add_parameter("sigma_lnorm", True, distribution=tfpd.Uniform(low=1.05, high=3.0))
    retrieval_config.add_parameter("log_kzz", True, distribution=tfpd.Uniform(low=5.0, high=13.0))

    # Let's set up different sedimentation fractions for each cloud species.
    retrieval_config.add_parameter("f_sed_MgSiO3(s)", True, distribution=tfpd.Uniform(low=1.0, high=11.0))
    retrieval_config.add_parameter("f_sed_Fe(s)", True, distribution=tfpd.Uniform(low=1.0, high=11.0))

    # --- Opacity sources ----------------------------------------------------
    retrieval_config.set_rayleigh_species(("H2", "He"))
    retrieval_config.set_continuum_opacities(("H2--H2", "H2--He"))
    retrieval_config.set_line_species(
        (
            "H2O__POKAZATEL",
            "CO-NatAbund",
            "CH4",
            "CO2",
            "HCN",
            "FeH",
            "H2S",
            "NH3",
            "PH3",
            "Na__NewAllard",
            "K__Allard",
            "TiO",
            "VO",
            "SiO",
        ),
        use_equilibrium_chemistry=True,
    )

    retrieval_config.add_cloud_species(
        "Fe(s)_crystalline_000__DHS",
        use_equilibrium_chemistry=True,
        equilbrium_mass_fraction_scaling_factor=(-3.5, 1.0),
    )

    retrieval_config.add_cloud_species(
        "MgSiO3(s)_crystalline_000__DHS",
        use_equilibrium_chemistry=True,
        equilbrium_mass_fraction_scaling_factor=(-3.5, 1.0),
    )

    return retrieval_config

[ ]:

# Build the configuration and set up plotting, as in the base tutorial.
retrieval_config = build_config()

# Define what to put into corner plot if run_mode == 'evaluate'
retrieval_config.parameters["planet_radius"].plot_in_corner = True
retrieval_config.parameters["planet_radius"].corner_label = r"$R_{\rm P}$ ($\rm R_{Jup}$)"
retrieval_config.parameters["planet_radius"].corner_transform = lambda x: x / cst.r_jup_mean
retrieval_config.parameters["log_g"].plot_in_corner = True
retrieval_config.parameters["log_g"].corner_ranges = [2.0, 5.0]
retrieval_config.parameters["log_g"].corner_label = "log g"
retrieval_config.parameters["fsed"].plot_in_corner = True
retrieval_config.parameters["log_kzz"].plot_in_corner = True
retrieval_config.parameters["log_kzz"].corner_label = "log Kzz"
retrieval_config.parameters["C/O"].plot_in_corner = True
retrieval_config.parameters["Fe/H"].plot_in_corner = True
retrieval_config.parameters["log_pquench"].plot_in_corner = True
retrieval_config.parameters["log_pquench"].corner_label = "log pquench"

for spec in retrieval_config.cloud_species:
    cname = spec.split("_")[0]
    retrieval_config.parameters["eq_scaling_" + cname].plot_in_corner = True
    retrieval_config.parameters["eq_scaling_" + cname].corner_label = cname


def configure_retrieval_plotter(retrieval):
    retrieval.plotter.spec_xlabel = "Wavelength [micron]"
    retrieval.plotter.spec_ylabel = "Flux [W/m2/micron]"
    retrieval.plotter.y_axis_scaling = 1.0
    retrieval.plotter.xscale = "log"
    retrieval.plotter.yscale = "linear"
    retrieval.plotter.resolution = 100.0  # maximum resolution, will rebin the data
    retrieval.plotter.nsample = 100  # if we want a plot with many sampled spectra
    retrieval.plotter.reference_data_name = "GRAVITY"
    retrieval.plotter.temp_limits = [150, 3000]
    retrieval.plotter.press_limits = [1e2, 1e-5]

[ ]:

retrieval = Retrieval(retrieval_config, output_directory="./", evaluate_sample_spectra=False)
configure_retrieval_plotter(retrieval)

Warm-starting NUTS with optimal estimation#

By default, pRT initialises the NumPyro walkers at the prior median (the origin of the unconstrained sampling space). For a model like this one, that starting point can sit in a region where the log-posterior is extremely steep: the first leapfrog steps diverge, the acceptance rate collapses, and the warmup phase spends most of its budget just finding the typical set. The classic optimal estimation method of Rodgers addresses both problems, and pRT exposes it through two keyword arguments of Retrieval.run:

MAP optimisation (optimize_init=True): an optax optimiser (Adam by default, with global-norm gradient clipping so the enormous initial gradients produce bounded steps) maximises the log-posterior in the unconstrained space and replaces the prior-median initial position with the MAP estimate. The optimiser is configured through optimizer_kwargs, which is forwarded to `optimal_estimation.optimize_unconstrained_position <../../autoapi/petitRADTRANS/retrieval/optimal_estimation/index.html>`__: the relevant keys are optimizer, learning_rate, clip_norm, n_steps, patience, n_starts/start_jitter (random restarts to guard against local optima), and lbfgs_polish_steps (a second-order L-BFGS refinement of the best optimum, recommended for tightening flat or ill-conditioned directions).
Hessian preconditioning (init_inverse_mass_matrix_from_hessian=True): the inverse Hessian of the negative log-posterior at the MAP is the Laplace estimate of the posterior covariance — the optimal-estimation covariance matrix. Passing it to NUTS as the initial inverse mass matrix preconditions the Hamiltonian kernel to the posterior’s scales (and, with a dense matrix, its correlations) from the very first step, instead of starting from an identity metric. Set is_mass_matrix_diagonal=True (the default) to use only the diagonal of the inverse Hessian, or False for the full dense matrix. NumPyro still adapts the mass matrix during warmup; the Hessian just provides a far better starting metric, which is why we can afford a shorter num_warmup below.

Because we will run several chains, we also set "walker_jitter" inside optimizer_kwargs: chain 0 starts exactly at the MAP, while the remaining chains get Gaussian jitter of that width around it. Starting all chains at literally the same point would make multi-chain convergence diagnostics such as R-hat meaningless.

The MAP point found by the warm-up is printed, saved to out_NumPyro/<retrieval_name>_map_warmup.{npz,json}, and the achieved log-posterior is recorded in the run summary as warmup_best_logdensity, so you can verify that the optimisation actually moved the walkers into the bulk of the posterior.

Tip: the same machinery is available as a standalone backend with sampler_type="optimal_estimation" (aliases "optax", "map"), which returns the MAP plus draws from the Laplace (Gaussian) posterior. That is a fast way to get a first best fit and approximate error bars before committing to a full NUTS run; see the samplers documentation for details.

[ ]:

run_retrieval = False

[ ]:

if run_retrieval:
    retrieval.run(
        # NUTS sampling configuration
        num_warmup=500,  # warm-started, preconditioned chains need less adaptation
        num_samples=2000,
        num_chains=4,
        chain_method="vectorized",
        target_acceptance_rate=0.9,
        progress_bar=True,
        use_jit=True,
        seed=12345,  # remove or randomize for a production retrieval
        # Optimal estimation warm-up: find the MAP with optax before sampling
        optimize_init=True,
        optimizer_kwargs={
            "optimizer": "adam",
            "learning_rate": 1e-2,
            "clip_norm": 1.0,
            "n_steps": 400,
            "n_starts": 2,  # random restarts to guard against local optima
            "lbfgs_polish_steps": 50,  # second-order polish of the best optimum
            "walker_jitter": 0.05,  # spread chains 1-3 around the MAP (chain 0 stays at it)
            "seed": 12345,
        },
        # Precondition the NUTS metric with the inverse Hessian at the MAP
        init_inverse_mass_matrix_from_hessian=True,
        is_mass_matrix_diagonal=True,  # False for a dense (correlated) metric
    )

When the run starts you will see, in order: the optimiser progress (log-posterior and gradient norm per reported step, for each restart), the L-BFGS polish steps, a table of the MAP parameter values (also written to out_NumPyro/HR8799e_advanced_example_map_warmup.npz), a line reporting the diagonal range of the Hessian-derived inverse mass matrix, and finally the usual NumPyro warmup and sampling progress bars.

If the warm-up fails to improve the log-posterior substantially, increase n_steps or n_starts, or lower the learning_rate. Note that optimize_init is ignored if you pass an explicit initial_position, and the Hessian preconditioning is only applied together with optimize_init=True (and only if you have not supplied your own inverse_mass_matrix).

Multi-chain NUTS#

The run above already used several chains: with num_chains=4 and chain_method="vectorized", NumPyro vmaps all chains into a single JIT-compiled computation. That is the fastest option on a single device, but the per-step memory scales with the number of chains, which can be prohibitive for heavy forward models. The in-process alternatives are chain_method="sequential" (one chain after another, a single resident model copy) and chain_method="parallel" (one chain per JAX device, e.g. per GPU).

For CPU machines and clusters, the scalable approach is process-level parallelism: independent single-chain runs, each with its own seed and output name, whose posteriors are merged afterwards. Retrieval.run_multichain orchestrates all of this natively:

it derives a deterministic per-chain seed and retrieval name from chain_seed,
it runs the chains according to a chain_execution policy ("auto", "vectorised", "sharded", "sequential", or "processes"; "auto" picks sharded execution when GPUs are available and processes on CPU),
and it combines the per-chain posteriors into one result with convergence diagnostics: split-R-hat, effective sample size, and divergence fractions.

This is where the build_config function defined at the top of the notebook pays off: run_multichain calls it to create a fresh configuration in each worker.

Two practical notes for this notebook:

The "processes" mode spawns worker processes that re-import the builder, so it must be defined in an importable module and the call must be guarded by if __name__ == "__main__": — that mode is for scripts, not notebooks. Here we use "sequential", which runs the chains one after another in this process and still exercises the full combine-and-diagnose pipeline.
When optimize_init is combined with the per-chain execution modes, every chain repeats the MAP optimisation independently (each with its own derived seed, so restarts differ between chains). This costs some redundant optimisation time but means each chain starts from its own MAP estimate.

[ ]:

run_multichain_demo = False

[ ]:

if run_multichain_demo:
    result = Retrieval.run_multichain(
        build_config,
        num_chains=4,
        output_directory="./",
        chain_execution="sequential",  # notebook-safe; use "processes" from a script
        chain_seed=12345,
        # Everything below is forwarded to each per-chain Retrieval.run call:
        num_warmup=500,
        num_samples=1000,
        target_acceptance_rate=0.9,
        progress_bar=True,
        optimize_init=True,
        optimizer_kwargs={
            "n_steps": 400,
            "lbfgs_polish_steps": 50,
        },
        init_inverse_mass_matrix_from_hessian=True,
    )

    print(f"execution mode   : {result['execution']}")
    print(f"succeeded chains : {result['succeeded_chain_ids']}")
    print(f"total samples    : {result['diagnostics']['n_samples_total']}")
    print(f"max R-hat        : {result['diagnostics']['max_rhat']:.4f}")
    print(f"converged        : {result['diagnostics']['converged']}")

run_multichain writes one artifact per chain, out_NumPyro/HR8799e_advanced_example_chain<c>_samples.npz, plus a combined posterior under the name HR8799e_advanced_example_combined (in the same format the usual Retrieval.get_samples and plotting code reads) and a HR8799e_advanced_example_combined_diagnostics.json with the convergence summary. A run is considered converged when the maximum split-R-hat across parameters is below rhat_threshold (1.01 by default); if it is not, increase num_samples or num_warmup and re-run. With resume=True, chains whose samples file already exists are skipped, so a re-submission only runs the missing chains and then re-combines. If you ran chains yourself, retrieval.combine_chains(num_chains=...) performs the merge in one call and retrieval.get_convergence_diagnostics() returns the diagnostics dictionary.

Scaling out: independent processes and SLURM#

On a multi-core workstation, move the builder into an importable module and switch to the "processes" mode, which keeps the parent lightweight (only the workers build the heavy radiative-transfer runtime) and pins each chain to a disjoint set of cores:

# run_hr8799e_chains.py
from petitRADTRANS.retrieval import Retrieval
from hr8799e_config import build_config  # the builder, in an importable module

if __name__ == "__main__":
    result = Retrieval.run_multichain(
        build_config,
        num_chains=8,
        output_directory="out",
        chain_execution="processes",
        chain_seed=12345,
        per_chain_memory_gb=18,  # informs the memory estimate and process concurrency
        num_warmup=1000,
        num_samples=2000,
        target_acceptance_rate=0.9,
        optimize_init=True,
        init_inverse_mass_matrix_from_hessian=True,
    )

For multi-node clusters, Retrieval.write_chain_launcher generates a SLURM job array (one task per chain plus a dependent combine job) from the same builder:

Retrieval.write_chain_launcher(
    build_config,  # imported as <module>.build_config
    num_chains=16,
    output_directory="out",
    script_dir="launch",
    partition="my_partition",
    cpus_per_task=8,
    time="24:00:00",
    env_setup_lines=["module load anaconda/3", "conda activate prt"],
    run_kwargs={
        "num_warmup": 1000,
        "num_samples": 2000,
        "optimize_init": True,
        "init_inverse_mass_matrix_from_hessian": True,
    },
    resume=True,
)

See the parallelisation documentation for the full discussion of execution policies, memory estimation, and failure handling.

Evaluating the results#

Once a retrieval is complete, the standard outputs are produced exactly as in the base tutorial. The chain diagnostics (acceptance rates, divergences, integration steps) are stored in out_NumPyro, alongside the MAP warm-up files, and a combined multi-chain run is read transparently by the same plotting code.

[ ]:

if run_retrieval:
    retrieval.plot_all(contribution=True)

Contact

If you need any additional help, don’t hesitate to contact Evert Nasedkin.