API Reference

Main Package

Redback-JAX: A lightweight JAX-only version of the redback electromagnetic transient analysis package.

This package provides JAX-based implementations for electromagnetic transient modeling and Bayesian inference, focusing on performance and automatic differentiation capabilities.

redback_jax.__getattr__(name)[source]: Lazy imports to avoid enabling JAX x64 at package load time.

Models

JAX-based transient models for electromagnetic counterparts.

redback_jax.models.arnett_bolometric(time, f_nickel, mej, *, vej=None, kappa=None, kappa_gamma=None)[source]

Bolometric Arnett (1982) light curve with Ni/Co decay engine + diffusion.

Parameters:

time – time in days
f_nickel – fraction of nickel mass
mej – total ejecta mass in solar masses
kappa – opacity in cm^2/g (required)
kappa_gamma – gamma-ray opacity in cm^2/g (required)
vej – ejecta velocity in km/s (required)

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.arnett_with_features_cosmology(f_nickel, mej, *, redshift=0.0, cosmo_H0=PLANCK18_H0, cosmo_Om0=PLANCK18_OM0, vej=None, kappa=None, kappa_gamma=None, temperature_floor=None, features=NO_SED_FEATURES)[source]

Arnett model with cosmological luminosity distance calculation.

Parameters:

redshift – source redshift
f_nickel – fraction of nickel mass
mej – total ejecta mass in solar masses
cosmo_H0 – Hubble constant (km/s/Mpc)
cosmo_Om0 – matter density parameter
kappa – opacity in cm^2/g (required)
kappa_gamma – gamma-ray opacity in cm^2/g (required)
vej – ejecta velocity in km/s (required)
temperature_floor – floor temperature in K
features – SEDFeatures object

Returns:

namedtuple(time, lambdas, spectra)

redback_jax.models.blackbody_to_flux_density(temperature, r_photosphere, dl, frequency)[source]

A general blackbody_to_flux_density formula

Parameters:

temperature – effective temperature in kelvin
r_photosphere – photosphere radius in cm
dl – luminosity_distance in cm
frequency – frequency to calculate in Hz

Returns:

flux_density in erg/s/Hz/cm^2

redback_jax.models.magnetar_powered_bolometric(time, p0, bp, mass_ns, theta_pb, mej, kappa, kappa_gamma, vej)[source]

Bolometric light curve of a magnetar-powered supernova (Arnett diffusion).

Parameters:

time – source-frame time in days
p0 – initial spin period in milliseconds
bp – polar B-field in units of 10^14 G
mass_ns – NS mass in solar masses
theta_pb – spin–B-field angle in radians
mej – ejecta mass in solar masses
kappa – optical opacity in cm^2/g
kappa_gamma – gamma-ray opacity in cm^2/g
vej – ejecta velocity in km/s

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.magnetar_nickel_bolometric(time, f_nickel, mej, p0, bp, mass_ns, theta_pb, kappa, kappa_gamma, vej)[source]

Bolometric light curve powered by both a magnetar and Ni/Co radioactive decay (Arnett diffusion). The two luminosity sources are added before diffusion.

Reference: Gomez et al. 2018 (https://ui.adsabs.harvard.edu/abs/2018ApJS..236….6G/abstract)

Parameters:

time – source-frame time in days
f_nickel – nickel mass fraction (M_Ni = f_nickel * mej)
mej – total ejecta mass in solar masses
p0 – initial spin period in milliseconds
bp – polar B-field in units of 10^14 G
mass_ns – NS mass in solar masses
theta_pb – spin–B-field angle in radians
kappa – optical opacity in cm^2/g
kappa_gamma – gamma-ray opacity in cm^2/g
vej – ejecta velocity in km/s

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.csm_interaction_bolometric(time, mej, csm_mass, vej, eta, rho, kappa, r0, nn=12, delta=1, efficiency=0.5)[source]

Bolometric CSM-interaction light curve (Chevalier 1982 shocks + diffusion).

Parameters:

time – source-frame time in days
mej – ejecta mass in solar masses
csm_mass – CSM mass in solar masses
vej – ejecta velocity in km/s
eta – CSM density profile exponent
rho – CSM density amplitude in g/cm^3
kappa – opacity in cm^2/g
r0 – inner CSM radius in AU
nn – ejecta density power-law slope (default 12)
delta – inner ejecta density slope (default 1)
efficiency – kinetic-to-luminosity efficiency (default 0.5)

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.shock_cooling_bolometric(time, log10_mass, log10_radius, log10_energy, nn, delta, kappa)[source]

Bolometric shock-cooling light curve following Piro (2021).

All large quantities (mass, radius, energy) are passed as log10 to stay float32-safe throughout; all intermediate computations stay in log10 space.

Parameters:

time – source-frame time in days
log10_mass – log10 envelope mass in solar masses
log10_radius – log10 envelope radius in cm
log10_energy – log10 explosion energy in erg
nn – outer density power-law slope
delta – inner density power-law slope
kappa – opacity in cm^2/g

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.shocked_cocoon_bolometric(time, mej, vej, eta, tshock, shocked_fraction, cos_theta_cocoon, kappa)[source]

Bolometric light curve of a shocked jet cocoon (Piro & Kollmeier 2018). All large intermediate quantities computed in log10 space for float32 safety.

Parameters:

time – source-frame time in days
mej – ejecta mass in solar masses
vej – ejecta velocity in units of c (speed of light)
eta – ejecta density power-law slope
tshock – shock time in seconds
shocked_fraction – fraction of ejecta mass that is shocked
cos_theta_cocoon – cosine of cocoon opening half-angle
kappa – gray opacity in cm^2/g

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.shock_cooling_and_arnett_bolometric(time, log10_mass, log10_radius, log10_energy, nn, delta, f_nickel, mej, vej, kappa, kappa_gamma)[source]

Combined bolometric light curve: shock cooling (Piro 2021) + Ni/Co decay (Arnett 1982).

The two components are summed in luminosity space (log-sum-exp), which is float32-safe.

Parameters:

time – source-frame time in days
log10_mass – log10 envelope mass in solar masses
log10_radius – log10 envelope radius in cm
log10_energy – log10 explosion energy in erg
nn – outer density power-law slope
delta – inner density power-law slope
f_nickel – nickel mass fraction
mej – total ejecta mass in solar masses
vej – ejecta velocity in km/s
kappa – optical opacity in cm^2/g
kappa_gamma – gamma-ray opacity in cm^2/g

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.tde_analytical_bolometric(time, log10_l0, t_0_turn, mej, vej, kappa, kappa_gamma)[source]

Bolometric TDE light curve: t^{-5/3} fallback engine + Arnett diffusion.

Parameters:

time – source-frame time in days
log10_l0 – log10 of bolometric luminosity at 1 second in erg/s
t_0_turn – turn-on time in days
mej – ejecta mass in solar masses
vej – ejecta velocity in km/s
kappa – optical opacity in cm^2/g
kappa_gamma – gamma-ray opacity in cm^2/g

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.metzger_kilonova_bolometric(time, mej, vej, beta, kappa, vmax=0.7, neutron_precursor=True)[source]

Bolometric kilonova light curve (Metzger 2017) with 200 shells and Barnes+16 thermalisation, solved via a sequential Euler ODE with jax.lax.scan.

Parameters:

time – source-frame time in days (must be strictly increasing, ≥2 points)
mej – ejecta mass in solar masses
vej – minimum ejecta velocity in units of c
beta – velocity power-law slope (M ∝ v^{-beta})
kappa – gray opacity in cm^2/g
vmax – maximum ejecta velocity in units of c (default 0.7)
neutron_precursor – include neutron precursor emission (default True)

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.magnetar_boosted_kilonova_bolometric(time, mej, vej, beta, kappa, p0, bp, mass_ns, theta_pb, thermalisation_efficiency=1.0, vmax=0.7, neutron_precursor=True)[source]

Bolometric kilonova light curve with magnetar spin-down energy injection.

Parameters:

time – source-frame time in days (strictly increasing, ≥2 points)
mej – ejecta mass in solar masses
vej – minimum ejecta velocity in units of c
beta – velocity power-law slope
kappa – gray opacity in cm^2/g
p0 – initial spin period in milliseconds
bp – polar B-field in units of 10^14 G
mass_ns – neutron star mass in solar masses
theta_pb – angle between spin and B-field axes in radians
thermalisation_efficiency – fraction of magnetar luminosity thermalised (default 1.0)
vmax – maximum ejecta velocity in units of c (default 0.7)
neutron_precursor – include neutron precursor emission (default True)

Returns:

log10 of bolometric luminosity in erg/s

redback_jax.models.general_magnetar_driven_supernova_bolometric(time, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, f_nickel=0.0, solver='diffrax', rtol=1e-05, atol=1e-08, n_grid=2000)[source]

Bolometric light curve of a general magnetar-driven supernova.

Translated from redback’s general_magnetar_driven_supernova_bolometric (Sarin et al. 2022) into JAX, enabling JIT compilation and gradients.

Parameters:

time (array_like, days) – Source-frame times at which to evaluate the model.
mej (float, M_sun) – Ejecta mass.
E_sn (float, erg) – Explosion kinetic energy.
kappa (float, cm²/g) – Optical opacity.
l0 (float, erg/s) – Initial magnetar spin-down luminosity.
tau_sd (float, s) – Magnetar spin-down timescale.
nn (float) – Magnetar braking index (3 = dipole).
kappa_gamma (float, cm²/g) – Gamma-ray opacity for thermalisation efficiency.
f_nickel (float, optional) – Ni-56 mass fraction of ejecta. Default 0.
solver (str, optional) – ODE backend. 'diffrax' (default) uses the adaptive Tsit5 integrator (~380× faster than redback). 'euler' uses the fixed-step Euler scan (n_grid points).
rtol (float, optional) – Tolerances for the diffrax solver (ignored when solver=’euler’).
atol (float, optional) – Tolerances for the diffrax solver (ignored when solver=’euler’).
n_grid (int, optional) – Grid points for the Euler scan (ignored when solver=’diffrax’). Static — each unique value triggers a separate compilation.

Returns:

log10(L_bol) in erg/s evaluated at each element of time.

Return type:

jnp.ndarray

redback_jax.models.general_magnetar_driven_supernova_bolometric_and_vej(time, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, f_nickel=0.0, solver='diffrax', rtol=1e-05, atol=1e-08, n_grid=2000)[source]

Bolometric light curve + time-varying ejecta velocity.

Returns the same ODE outputs as general_magnetar_driven_supernova_bolometric but also exposes the ejecta velocity v_ej(t) = β(t)·c derived from the ODE Lorentz factor. Pass vej_kms to compute_temperature_floor_log10 for a time-varying photosphere (consistent with the redback reference model).

Parameters:

general_magnetar_driven_supernova_bolometric) ((same as)

Returns:

log10_lbol (jnp.ndarray — log10(L_bol) in erg/s, shape (T,))
vej_kms (jnp.ndarray — ejecta velocity in km/s, shape (T,)) – Both arrays have the same dtype as time.

redback_jax.models.general_magnetar_driven_supernova_bolometric_batched(time, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, f_nickel=None, n_grid=2000)[source]

Evaluate B parameter samples simultaneously via jax.vmap.

Runs general_magnetar_driven_supernova_bolometric on B independent parameter vectors in a single JIT-compiled kernel. On GPU this maps each ODE trajectory to a separate set of CUDA cores; on CPU it benefits from better cache reuse compared to sequential calls.

Parameters:

time (array_like, shape (T,)) – Shared source-frame times in days (broadcast over all B samples).
mej (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
E_sn (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
kappa (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
l0 (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
tau_sd (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
nn (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
kappa_gamma (array_like, shape (B,)) – Physical parameters — one scalar value per sample.
f_nickel (array_like, shape (B,), optional) – Ni-56 mass fractions. Defaults to zeros for all samples if None.
n_grid (int, optional) – ODE grid points (static — triggers recompile on change). Default 2000.

Returns:

log10(L_bol) in erg/s — row i corresponds to sample i.

Return type:

jnp.ndarray, shape (B, T)

redback_jax.models.general_magnetar_driven_supernova_bolometric_diffrax(time, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, f_nickel=0.0, rtol=1e-05, atol=1e-08)[source]

Bolometric light curve using the diffrax Tsit5 adaptive ODE solver.

Functionally equivalent to general_magnetar_driven_supernova_bolometric (with solver='diffrax') but exposed as a standalone function for backward compatibility. Typically 3–10× faster than the Euler scan at the same accuracy once compiled.

Parameters:

time (array_like, source-frame days)
general_magnetar_driven_supernova_bolometric) ((other params same as)
rtol (float, default 1e-5)
atol (float, default 1e-8)

Returns:

log10_lbol

Return type:

jnp.ndarray — log10(L_bol) in erg/s

redback_jax.models.general_magnetar_driven_supernova_bolometric_and_vej_diffrax(time, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, f_nickel=0.0, rtol=1e-05, atol=1e-08)[source]

Return (log10_lbol, vej_kms) from the diffrax Tsit5 adaptive ODE.

Counterpart to general_magnetar_driven_supernova_bolometric_and_vej using the diffrax backend. Required internally by general_magnetar_supernova_spectra_diffrax (the spectra factory needs vej derived from the ODE, not as a free parameter).

Parameters:

time (array_like, source-frame days)
general_magnetar_driven_supernova_bolometric_diffrax) ((other params same as)

Returns:

log10_lbol (jnp.ndarray — log10(L_bol) in erg/s)
vej_kms (jnp.ndarray — ejecta velocity in km/s)

redback_jax.models.general_magnetar_driven_supernova(time, frequency, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, temperature_floor, luminosity_distance, redshift, cutoff_wavelength=3000.0, f_nickel=0.0, solver='diffrax', rtol=1e-05, atol=1e-08, n_grid=2000, alpha_uv=1.0)[source]

Full multiband general magnetar-driven supernova model.

1:1 JAX translation of redback’s general_magnetar_driven_supernova with output_format='flux_density', using the TemperatureFloor photosphere and CutoffBlackbody SED.

Pipeline

K-correction: freq_src = freq_obs·(1+z), t_src = t_obs/(1+z)
ODE: (log10_lbol, vej_kms) from chosen backend
Photosphere: (T_ph, log10_r_ph) from TemperatureFloor
SED: F_mjy from CutoffBlackbody
Return: F_mjy · (1+z)

param time:
type time:: (N,) observer-frame days
param frequency:
type frequency:: (N,) observer-frame Hz
param mej:
type mej:: float, M_sun
param E_sn:
type E_sn:: float, erg
param kappa:
type kappa:: float, cm²/g
param l0:
type l0:: float, erg/s
param tau_sd:
type tau_sd:: float, s
param nn:
type nn:: float
param kappa_gamma:
type kappa_gamma:: float, cm²/g
param temperature_floor:
type temperature_floor:: float, K
param luminosity_distance:
type luminosity_distance:: float, cm
param redshift:
type redshift:: float
param cutoff_wavelength:
type cutoff_wavelength:: float, Å (default 3000)
param f_nickel:
type f_nickel:: float (default 0)
param solver:: 'diffrax' (default, adaptive Tsit5) or 'euler' (fixed-step).
type solver:: str, optional
param rtol:: diffrax tolerances (ignored when solver=’euler’).
type rtol:: float, optional
param atol:: diffrax tolerances (ignored when solver=’euler’).
type atol:: float, optional
param n_grid:: Euler grid points (ignored when solver=’diffrax’).
type n_grid:: int, optional
param alpha_uv:: UV power-law suppression index for CutoffBlackbody SED. λ < λ_c contributes Planck × (λ/λ_c)^alpha_uv. Valid range [0, 4).
type alpha_uv:: float, default 1.0
returns:: F_mjy
rtype:: (N,) mJy — observer-frame flux density

redback_jax.models.general_magnetar_driven_supernova_diffrax(time, frequency, mej, E_sn, kappa, l0, tau_sd, nn, kappa_gamma, temperature_floor, luminosity_distance, redshift, cutoff_wavelength=3000.0, f_nickel=0.0, rtol=1e-05, atol=1e-08, alpha_uv=1.0)[source]

Full multiband model using the diffrax Tsit5 adaptive ODE solver.

Standalone backward-compatible variant of general_magnetar_driven_supernova that always uses the diffrax backend. Equivalent to calling general_magnetar_driven_supernova(..., solver='diffrax').

Parameters:

n_grid) ((all params same as general_magnetar_driven_supernova except)
rtol (float, default 1e-5)
atol (float, default 1e-8)
alpha_uv (float, default 1.0)

Returns:

F_mjy

Return type:

(N,) mJy

redback_jax.models.general_magnetar_supernova_spectra_diffrax(redshift, lum_dist, temperature_floor, cutoff_wavelength=default_cutoff_wavelength, alpha_uv=default_alpha_uv, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs)

CutoffBlackbody spectra model wrapping general_magnetar_driven_supernova_bolometric_and_vej_diffrax.

Parameters:

redshift – source redshift
lum_dist – luminosity distance in cm
temperature_floor – floor temperature in K
cutoff_wavelength – UV cutoff wavelength in Å (default 3000, can be inferred)
alpha_uv – UV power-law suppression index (default 1.0, can be inferred)
features – SEDFeatures (default NO_SED_FEATURES)
**bolometric_kwargs – forwarded to the bolometric+vej function

Returns:

namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)

redback_jax.models.make_spectra_model(bolometric_fn)[source]

Wrap a bolometric model function to produce a full spectra model.

The returned function has signature:

spectra_model(redshift, lum_dist, vej, temperature_floor,
              features=NO_SED_FEATURES, **bolometric_kwargs)
-> namedtuple(time, lambdas, spectra)

Parameters:: bolometric_fn (callable) – Any function f(time_days, **kwargs) -> log10_lbol (log10 erg/s). time_days must be its first positional argument.
Returns:: A spectra model with the same photosphere/SED pipeline.
Return type:: callable

redback_jax.models.make_cutoff_spectra_model(bolometric_and_vej_fn, default_cutoff_wavelength=3000.0, default_alpha_uv=1.0)[source]

Wrap a (log10_lbol, vej_kms) bolometric function into a CutoffBlackbody spectra model.

Unlike make_spectra_model, this factory: - Expects the bolometric function to return (log10_lbol, vej_kms) (vej

is derived from the ODE, not a free parameter).

Uses a CutoffBlackbody SED with user-controllable cutoff_wavelength and alpha_uv that can be placed in the inference prior.

Parameters:

bolometric_and_vej_fn (callable) – f(time_days, **kwargs) -> (log10_lbol, vej_kms).
default_cutoff_wavelength (float, Å (default 3000))
default_alpha_uv (float (default 1.0))

Returns:

spectra_model –

``spectra_model(redshift, lum_dist, temperature_floor,: cutoff_wavelength=…, alpha_uv=…, features=NO_SED_FEATURES, **bolometric_kwargs)``

returning namedtuple(time, lambdas, spectra).

Return type:

callable

redback_jax.models.arnett_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping arnett_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.magnetar_powered_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping magnetar_powered_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.magnetar_nickel_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping magnetar_nickel_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.csm_interaction_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping csm_interaction_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.shock_cooling_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping shock_cooling_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.shocked_cocoon_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping shocked_cocoon_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.shock_cooling_and_arnett_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping shock_cooling_and_arnett_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.tde_analytical_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping tde_analytical_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.metzger_kilonova_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping metzger_kilonova_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

redback_jax.models.magnetar_boosted_kilonova_spectra(redshift, lum_dist, vej, temperature_floor, features=NO_SED_FEATURES, _time_observer_frame_grid=None, _lambda_observer_frame_grid=None, **bolometric_kwargs): Spectra model wrapping magnetar_boosted_kilonova_bolometric.nnArgs:n redshift: source redshiftn lum_dist: luminosity distance in cmn vej: ejecta velocity in km/s (photosphere)n temperature_floor: floor temperature in Kn features: SEDFeatures (default NO_SED_FEATURES)n **bolometric_kwargs: forwarded to the bolometric functionnnReturns:n namedtuple with fields time (days), lambdas (Angstrom), spectra (erg/s/cm^2/Angstrom)n

class redback_jax.models.SEDFeatures(rest_wavelengths, sigmas, amplitudes, t_starts, t_ends, rise_times=2.0 * 24.0 * 3600.0, fall_times=5.0 * 24.0 * 3600.0)[source]

Bases: object

A class representing a spectral feature(s) in the SED. This implements the PyTree interface for JAX, so it can be used in JIT-compiled functions.

Parameters:

rest_wavelengths – Central wavelengths in Angstroms
sigmas – Gaussian widths in Angstroms
amplitudes – Amplitudes (negative=absorption, positive=emission), percentage of continuum (e.g., -0.4 = 40% absorption)
t_starts – Start time in seconds
t_ends – End time in seconds
rise_times – Rise times in seconds (default: 2 days)
fall_times – Fall times in seconds (default: 5 days)

classmethod from_feature_list(feature_list)[source]

Create SEDFeatures from a list of feature dictionaries.

Parameters:: feature_list – List of dictionaries, each with keys: ‘rest_wavelength’, ‘sigma’, ‘amplitude’, ‘t_start’, ‘t_end’, ‘rise_time’, ‘fall_time’
Returns:: SEDFeatures object

tree_flatten()[source]

classmethod tree_unflatten(aux_data, children)[source]

calculate_smooth_evolution(time)[source]

Calculate smooth transitions for a set of features.

Parameters:: time – Time array in seconds (source frame)
Returns:: time_factors array in [0, 1] representing the evolution of the feature over time

redback_jax.models.apply_sed_feature(features, base_flux, frequency, time)[source]

Apply spectral features completely vectorized.

Parameters:: features – SEDFeatures object

:param base_flux as a 2-d array (time, wavelength) in erg/s/Hz/cm^2 :param frequency: frequency to calculate in Hz - Must be same length as time array or a single number.

In source frame.

Parameters:: time – time array in seconds (source frame).
Returns:: modified flux_density as a 2-d array (time, wavelength) in erg/s/Hz/cm^2

redback_jax.models.register_model(name, fn)[source]

Register a model function under name.

Parameters:: name (str)

redback_jax.models.get_model(name)[source]

Retrieve a registered model by name, or raise KeyError.

Parameters:: name (str)

redback_jax.models.load_plugins()[source]: Discover and load all redback_jax.models entry-point plugins.

Utilities

Utility functions for redback-jax.

redback_jax.utils.luminosity_distance_cm(redshift, H0=PLANCK18_H0, Om0=PLANCK18_OM0)[source]

Luminosity distance in cm.

Parameters:

redshift (float) – Source redshift.
H0 (float, optional) – Hubble constant in km/s/Mpc (default: Planck18).
Om0 (float, optional) – Matter density parameter (default: Planck18).

Returns:

Luminosity distance in cm.

Return type:

float

Examples

>>> dl_cm = luminosity_distance_cm(0.01)

Inference

JAX-based Bayesian inference tools using BlackJAX.

Clean API

from redback_jax.inference import Prior, Uniform, LogUniform, Gaussian from redback_jax.inference import SpectralLikelihood, NestedSampler, NSResult

Legacy API (kept for backward compatibility)

from redback_jax.inference import (
SamplerResult, create_uniform_prior, create_gaussian_likelihood, run_nested_sampling, run_mcmc, fit_transient, HAS_BLACKJAX,

)

class redback_jax.inference.Uniform(minimum, maximum, *, name)[source]

Bases: _Distribution

Uniform prior between minimum and maximum.

Parameters:

minimum (float) – Support of the distribution.
maximum (float) – Support of the distribution.
name (str) – Parameter name (used as dict key in likelihood calls).

Examples

>>> p = Uniform(0.05, 0.2, name='f_nickel')
>>> p.log_prob(jnp.array(0.1))   # log(1 / (0.2 - 0.05))

property low: float

property high: float

sample(key)[source]

Parameters:: key (Array)
Return type:: Array

log_prob(value)[source]

Parameters:: value (Array)
Return type:: Array

class redback_jax.inference.LogUniform(minimum, maximum, *, name)[source]

Bases: _Distribution

Log-uniform (Jeffreys) prior between minimum and maximum.

The density is proportional to 1/x, normalised over [minimum, maximum].

Parameters:

minimum (float) – Support of the distribution (must be > 0).
maximum (float) – Support of the distribution (must be > 0).
name (str) – Parameter name.

Examples

>>> p = LogUniform(1e-2, 1e2, name='kappa')

property low: float

property high: float

sample(key)[source]

Parameters:: key (Array)
Return type:: Array

log_prob(value)[source]

Parameters:: value (Array)
Return type:: Array

class redback_jax.inference.Gaussian(mu, sigma, *, name, minimum=-jnp.inf, maximum=jnp.inf)[source]

Bases: _Distribution

Gaussian prior (truncated to finite support by the sampler’s hard bounds).

Parameters:

mu (float) – Mean and standard deviation.
sigma (float) – Mean and standard deviation.
name (str) – Parameter name.
minimum (float)
maximum (float)

property low: float

property high: float

sample(key)[source]

Parameters:: key (Array)
Return type:: Array

log_prob(value)[source]

Parameters:: value (Array)
Return type:: Array

class redback_jax.inference.Prior(distributions)[source]

Bases: object

A composite prior built from a list of 1-D distributions.

Parameters:: distributions (list of _Distribution) – One distribution per free parameter. The list order determines the column order of the parameter vector passed to the likelihood.

Examples

>>> prior = Prior([
...     Uniform(58580, 58620, name='t0'),
...     Uniform(0.05, 0.20,   name='f_nickel'),
...     Uniform(0.8,  2.0,    name='mej'),
...     Uniform(3000, 8000,   name='vej'),
... ])
>>> particles = prior.sample_n(jax.random.PRNGKey(0), 100)  # (100, 4)
>>> log_p = prior.log_prob(particles[0])                    # scalar

sample(key)[source]

Draw one sample; returns a dict {name: scalar}.

Parameters:: key (Array)
Return type:: dict

sample_n(key, n)[source]

Draw n samples; returns an array of shape (n, n_params).

Parameters:

key (Array)
n (int)

Return type:

Array

log_prob(params)[source]

Evaluate the joint log-prior for a parameter vector of shape (n_params,).

This is JAX-traceable and can be used inside @jax.jit.

Parameters:: params (Array)
Return type:: Array

log_prob_fn()[source]: Return a pure JAX-traceable function params -> log_prior.

params_to_dict(params)[source]

Convert a parameter vector (n_params,) to a name-keyed dict.

Parameters:: params (Array)
Return type:: dict

dict_to_params(d)[source]

Convert a name-keyed dict to a parameter vector.

Parameters:: d (dict)
Return type:: Array

class redback_jax.inference.Likelihood(model, transient, fixed_params, t0_key='t0', evaluation_mode='full', compact_time_grid_size=256, compact_grid_pad_days=5.0)[source]

Bases: object

Gaussian photometric likelihood using a spectra model pipeline.

Parameters:

model (str or callable) – Model name (e.g. 'arnett_spectra') or a callable with signature f(redshift, lum_dist, vej, temperature_floor, **kwargs) returning a namedtuple (time, lambdas, spectra).
transient (Transient) – Data container with .time, .y, .y_err, .bands.
fixed_params (dict) – Parameters held fixed during inference — everything the model needs that is not a free parameter in the prior.
t0_key (str or None, optional) – Name of the MJD explosion-time free parameter (default 't0'). When present in the prior the likelihood converts transient.time from observer-frame MJD to source-frame days automatically. Set to None if times are already in source-frame days.
evaluation_mode ({"full", "compact_source", "direct_photometry"}, optional) – "full" preserves the existing model-default source grid. "compact_source" uses a dataset-specific source phase grid while still going through jax_supernovae.timeseries_multiband_flux. "direct_photometry" bypasses full source-cube materialization for factory-built blackbody spectra models and integrates directly through the bandpasses.
compact_time_grid_size (int)
compact_grid_pad_days (float)

class redback_jax.inference.FluxDensityLikelihood(model, time, frequency, flux_obs, flux_err, fixed_params)[source]

Bases: object

Gaussian likelihood for models returning observed-frame flux density (mJy).

Designed for use with general_magnetar_driven_supernova_diffrax and any other model with signature f(time, frequency, **params) -> jnp.ndarray. No jax_supernovae dependency — operates directly on flux residuals.

This class follows the same interface as Likelihood and is fully compatible with NestedSampler.

Parameters:

model (callable) – f(time, frequency, **params) -> jnp.ndarray (mJy), where time is observer-frame days and frequency is observer-frame Hz.
time (array-like) – Observer-frame times (days), shape (N,).
frequency (array-like) – Observer-frame frequencies (Hz), shape (N,).
flux_obs (array-like) – Observed flux density (mJy), shape (N,).
flux_err (array-like) – Flux density uncertainties (mJy), shape (N,).
fixed_params (dict) – Parameters held fixed during inference (e.g. luminosity_distance, redshift, kappa).

class redback_jax.inference.NestedSampler(likelihood, prior, outdir='results/', n_live=125, n_delete=20, num_mcmc_steps_multiplier=5, termination_dlogz=-3.0, verbose=True)[source]

Bases: object

BlackJAX nested sampler with a clean redback-style interface.

Parameters:

likelihood (Likelihood) – A Likelihood instance.
prior (Prior) – Composite prior object.
outdir (str, optional) – Directory for output files. Created if it does not exist. Set to None to disable file output.
n_live (int, optional) – Number of live points (default 125).
n_delete (int, optional) – Number of points to remove per iteration (default 20).
num_mcmc_steps_multiplier (int, optional) – MCMC steps per iteration = n_params × multiplier (default 5).
termination_dlogz (float, optional) – Stop when logZ_live - logZ < termination_dlogz (default -3).
verbose (bool, optional) – Show a tqdm progress bar (default True).

Examples

>>> sampler = NestedSampler(likelihood, prior, outdir='results/')
>>> result  = sampler.run(jax.random.PRNGKey(42))
>>> result.summary()

run(key)[source]

Run nested sampling.

Parameters:: key (jax.Array) – JAX random key.
Returns:: Posterior samples and evidence estimate.
Return type:: NSResult

plot_corner(result, truth=None, filename=None, **kwargs)[source]

Make a corner plot using anesthetic.

Parameters:

result (NSResult) – Output of run().
truth (dict, optional) – True parameter values to mark on the plot.
filename (str, optional) – Path to save the figure. Defaults to {outdir}/corner.png.

class redback_jax.inference.NSResult(logZ, samples, dead, log_weights, param_names)[source]

Bases: object

Container for nested sampling results.

logZ

Log evidence estimate.

Type:: float

samples

Posterior samples as {name: jnp.ndarray}.

Type:: dict

dead

Raw dead-point pytree from BlackJAX (for expert use).

Type:: object

log_weights

Log importance weights (shape (n_dead,)).

Type:: jnp.ndarray

param_names

Ordered parameter names.

Type:: list of str

summary()[source]: Print a parameter summary table.

class redback_jax.inference.MCMCSampler(likelihood, prior, n_warmup=500, n_samples=2000, n_chains=4, step_size=0.05, verbose=True)[source]

Bases: object

BlackJAX NUTS sampler with a clean redback-style interface.

The log-posterior is log_likelihood + log_prior, evaluated in the original parameter space. A reflected boundary is used to keep samples inside the prior support.

Parameters:

likelihood (Likelihood) – A Likelihood instance.
prior (Prior) – Composite prior object.
n_warmup (int, optional) – Number of warmup (adaptation) steps per chain (default 500).
n_samples (int, optional) – Number of post-warmup samples per chain (default 2000).
n_chains (int, optional) – Number of independent chains (default 4).
step_size (float, optional) – Initial NUTS step size (default 0.05).
verbose (bool, optional) – Show a progress bar (default True).

Examples

>>> sampler = MCMCSampler(likelihood, prior, n_warmup=500, n_samples=2000)
>>> result  = sampler.run(jax.random.PRNGKey(0))
>>> result.summary()

run(key)[source]

Run MCMC.

Parameters:: key (jax.Array) – JAX random key.
Return type:: MCMCResult

class redback_jax.inference.MCMCResult(samples_per_chain, param_names)[source]

Bases: object

Container for MCMC results.

samples

Posterior samples as {name: jnp.ndarray} — shape (n_chains * n_samples,).

Type:: dict

samples_per_chain

Samples per chain as {name: jnp.ndarray} — shape (n_chains, n_samples).

Type:: dict

param_names

Ordered parameter names.

Type:: list of str

n_chains

Number of chains.

Type:: int

n_samples

Number of post-warmup samples per chain.

Type:: int

summary()[source]: Print a parameter summary table.

class redback_jax.inference.SamplerResult(samples, log_likelihoods, log_weights, log_evidence, log_evidence_error, metadata)[source]

Bases: NamedTuple

Results from nested sampling run.

Parameters:

samples (Dict[str, Array])
log_likelihoods (Array)
log_weights (Array)
log_evidence (float)
log_evidence_error (float)
metadata (Dict[str, Any])

samples

Dictionary mapping parameter names to sample arrays

Type:: dict

log_likelihoods

Log likelihood values for each sample

Type:: jnp.ndarray

log_weights

Log weights for each sample (for nested sampling)

Type:: jnp.ndarray

log_evidence

Log evidence estimate

Type:: float

log_evidence_error

Error on log evidence estimate

Type:: float

metadata

Additional metadata from the sampling run

Type:: dict

samples: Dict[str, Array]: Alias for field number 0

log_likelihoods: Array: Alias for field number 1

log_weights: Array: Alias for field number 2

log_evidence: float: Alias for field number 3

log_evidence_error: float: Alias for field number 4

metadata: Dict[str, Any]: Alias for field number 5

redback_jax.inference.create_uniform_prior(prior_bounds)[source]

Create a uniform prior function from parameter bounds.

Parameters:: prior_bounds (dict) – Dictionary mapping parameter names to (low, high) bounds
Returns:: Function that transforms unit hypercube to parameter space
Return type:: callable

redback_jax.inference.create_gaussian_likelihood(model_fn, observed_data, errors, reduce_fn=None)[source]

Create a Gaussian likelihood function.

Parameters:

model_fn (callable) – Function that takes parameter dict and returns model predictions
observed_data (jnp.ndarray) – Observed data array
errors (jnp.ndarray) – Error array (standard deviations)
reduce_fn (callable, optional) – Function to reduce data (e.g., for rescaling errors)

Returns:

Log-likelihood function

Return type:

callable

redback_jax.inference.run_nested_sampling(loglikelihood_fn, prior_bounds, n_particles=500, num_mcmc_steps=20, max_iterations=100, rng_key=None, verbose=True)[source]

Run Sequential Monte Carlo (SMC) sampling using BlackJAX.

This uses adaptive tempered SMC which provides evidence estimates similar to nested sampling. The algorithm gradually increases the temperature from the prior to the posterior while tracking the normalizing constant.

Parameters:

loglikelihood_fn (callable) – Log-likelihood function that takes parameter dict
prior_bounds (dict) – Dictionary mapping parameter names to (low, high) bounds
n_particles (int, optional) – Number of particles (default: 500)
num_mcmc_steps (int, optional) – Number of MCMC steps per iteration (default: 20)
max_iterations (int, optional) – Maximum number of temperature steps (default: 100)
rng_key (jax.Array, optional) – JAX random key (default: None, will create one)
verbose (bool, optional) – Print progress information (default: True)

Returns:

Results from the SMC sampling run with evidence estimate

Return type:

SamplerResult

Raises:

ImportError – If blackjax is not installed

redback_jax.inference.run_mcmc(loglikelihood_fn, prior_bounds, n_samples=10000, n_warmup=1000, n_chains=4, step_size=0.01, rng_key=None, verbose=True)[source]

Run MCMC sampling using BlackJAX’s NUTS sampler.

Parameters:

loglikelihood_fn (callable) – Log-likelihood function that takes parameter dict
prior_bounds (dict) – Dictionary mapping parameter names to (low, high) bounds
n_samples (int, optional) – Number of samples to draw (default: 10000)
n_warmup (int, optional) – Number of warmup/burnin steps (default: 1000)
n_chains (int, optional) – Number of parallel chains (default: 4)
step_size (float, optional) – Initial step size for NUTS (default: 0.01)
rng_key (jax.Array, optional) – JAX random key (default: None, will create one)
verbose (bool, optional) – Print progress information (default: True)

Returns:

Results from the MCMC run

Return type:

SamplerResult

Raises:

ImportError – If blackjax is not installed

redback_jax.inference.fit_transient(transient, model_fn, prior_bounds, sampler='nested', sampler_kwargs=None, rng_key=None, verbose=True)[source]

Fit a transient model to observational data.

This is the main high-level interface for parameter inference, following the redback API style.

Parameters:

transient (Transient) – Transient object with observational data
model_fn (callable) – Model function that takes parameter dict and returns model predictions. Should be compatible with the transient’s data structure.
prior_bounds (dict) – Dictionary mapping parameter names to (low, high) bounds
sampler (str, optional) – Sampler to use: “nested” for nested sampling, “mcmc” for NUTS (default: “nested”)
sampler_kwargs (dict, optional) – Additional keyword arguments for the sampler
rng_key (jax.Array, optional) – JAX random key
verbose (bool, optional) – Print progress information

Returns:

Results from the sampling run

Return type:

SamplerResult

Examples

>>> from redback_jax import Transient
>>> from redback_jax.sources import PrecomputedSpectraSource
>>> import jax.numpy as jnp
>>>
>>> # Create transient data
>>> transient = Transient(
...     name='test_sn',
...     times=jnp.array([0, 5, 10, 15, 20]),
...     magnitudes=jnp.array([18.0, 17.5, 17.0, 17.5, 18.0]),
...     magnitude_errors=jnp.array([0.1, 0.1, 0.1, 0.1, 0.1]),
...     bands=['g'] * 5
... )
>>>
>>> # Create model function
>>> source = PrecomputedSpectraSource.from_arnett_model(...)
>>> bridges, band_to_idx = source.prepare_bridges(['g'])
>>> band_indices = jnp.array([0, 0, 0, 0, 0])
>>>
>>> def model_fn(params):
...     return source.bandmag(params, None, transient.times,
...                           band_indices=band_indices,
...                           bridges=bridges,
...                           unique_bands=['g'])
>>>
>>> # Define priors
>>> prior_bounds = {
...     'amplitude': (0.1, 10.0),
... }
>>>
>>> # Run inference
>>> result = fit_transient(transient, model_fn, prior_bounds)
>>> print(f"Log evidence: {result.log_evidence:.2f}")

redback_jax.inference.to_anesthetic_samples(result)[source]

Convert SamplerResult to anesthetic NestedSamples.

Parameters:: result (SamplerResult) – Results from nested sampling
Returns:: Anesthetic samples object with posterior samples
Return type:: anesthetic.NestedSamples or anesthetic.Samples
Raises:: ImportError – If anesthetic is not installed

redback_jax.inference.summarize_result(result)[source]

Summarize sampling results with posterior statistics.

Parameters:: result (SamplerResult) – Results from sampling
Returns:: Dictionary with parameter statistics (mean, std, median, etc.)
Return type:: dict