build_forward_model

exo_skryer.build_model.build_forward_model(cfg, obs: Dict, stellar_flux: ndarray | None = None, return_highres: bool = False) → Callable[[Dict[str, Array]], Array | Dict[str, Array]]

Build a JIT-compiled forward model for atmospheric retrieval.

This function constructs a forward model by assembling physics kernels for vertical structure (temperature, chemistry, altitude), opacity sources (line, continuum, clouds), and radiative transfer. The returned function is JIT-compiled for efficient gradient-based inference.

Parameters:

cfgconfig object

Configuration object containing physics settings (cfg.physics), opacity configuration (cfg.opac), and retrieval parameters (cfg.params). Must specify schemes for vertical structure (vert_Tp, vert_alt, vert_chem, vert_mu), opacity sources (opac_line, opac_ray, opac_cia, opac_cloud, opac_special), and radiative transfer (rt_scheme).

obsdict

Observational data dictionary containing:

• ‘wl’ : Observed wavelengths in microns (for bandpass loading)

• ‘dwl’ : Wavelength bin widths in microns

stellar_fluxndarray, optional

Stellar flux array for emission spectroscopy calculations. Required when rt_scheme is ‘emission_1d’ and emission_mode is ‘planet’ (not brown dwarf). Should match the high-resolution wavelength grid.

return_highresbool, optional

If True, the forward model returns both high-resolution and binned spectra as a dictionary: {'hires': D_hires, 'binned': D_bin}. If False (default), returns only the binned spectrum as a 1D array.

Returns:

forward_modelcallable

A JIT-compiled function with signature: forward_model(params: Dict[str, jnp.ndarray]) -> Union[jnp.ndarray, Dict]

The function takes a parameter dictionary (free parameters from the retrieval) and returns:

• If return_highres=False: 1D array of binned transit depth or emission flux

• If return_highres=True: Dict with keys ‘hires’ (high-res spectrum) and ‘binned’ (convolved spectrum)