API reference

Auto-generated from the source docstrings. The public API is the symbols below; treat everything else as internal.

Stellar populations

ceridwen.ssps.SSPData dataclass

SSPData(
    ssp_lgmet,
    ssp_lg_age_gyr,
    ssp_wave,
    ssp_flux,
    isoc_type=None,
    spec_library=None,
    imf_type=None,
    fsps_version=None,
    fsps_kwargs=dict(),
    wave_min=None,
    wave_max=None,
    schema_version=None,
)

Immutable container for the SSP interpolation grids (+ provenance).

Attributes:

Name	Type	Description
ssp_lgmet	(ndarray, shape(n_met))	log10 of the absolute metallicity grid. Z is the mass fraction of elements heavier than helium. Typical range ~-2.3 to +0.2 dex.
ssp_lg_age_gyr	(ndarray, shape(n_ages))	log10(age / Gyr).
ssp_wave	(ndarray, shape(n_wave))	Wavelength grid in Angstroms.
ssp_flux	(ndarray, shape(n_met, n_ages, n_wave))	SSP flux density in Lsun / Hz per Msun of initial stellar mass.
isoc_type	str or None	Isochrone library the grid was built with (e.g. 'mist'), read from FSPS's compiled-in library set. None for legacy grids. :class:ceridwen.csp.CSPBasis uses this to pick the matching nebular CLOUDY grid automatically.
spec_library	str or None	Spectral library (e.g. 'miles'). None for legacy grids.
imf_type	int or None	FSPS IMF selector the grid was built with. None for legacy.
fsps_version	str or None	python-fsps version string used to build the grid.
fsps_kwargs	dict	The (whitelisted) FSPS build kwargs actually used. {} for legacy grids.
wave_min, wave_max	float or None	Wavelength range (Å) of ssp_wave at build time.
schema_version	str or None	On-disk metadata schema tag. None for legacy grids.

Notes

The provenance fields are ordinary Python objects (str / int / dict / None); they are never JAX arrays and never enter a @jit kernel. They are excluded from equality/hashing (compare=False).

No log_qq table is stored; the nebular model computes the ionising photon rate internally. HDF5 files that contain a log_qq dataset are loaded transparently — the field is simply ignored.

__post_init__

__post_init__()

Validate grid consistency.

Source code in ceridwen/ssps/ssp_data.py

def __post_init__(self):
    """Validate grid consistency."""
    if self.ssp_flux.shape != (self.ssp_lgmet.size,
                               self.ssp_lg_age_gyr.size,
                               self.ssp_wave.size):
        raise ValueError(
            f"SSP flux grid shape mismatch: expected "
            f"({self.ssp_lgmet.size}, {self.ssp_lg_age_gyr.size}, "
            f"{self.ssp_wave.size}) but got {self.ssp_flux.shape}.  "
            f"Grid dimensions must be consistent (n_met, n_ages, n_wave)."
        )

display

display(*, return_str=False, file=None)

Print a summary of the grid and its provenance.

Intended as a sanity check: call it right after from_fsps or load to confirm the isochrone set, spectral library, IMF, and grid coverage are what you expect before you build a CSPBasis. Purely diagnostic — none of this is touched by the forward model.

Parameters:

Name	Type	Description	Default
return_str	bool	Return the formatted string instead of printing it. Default False.	False
file	file - like	Destination for the print (default sys.stdout).	None

Returns:

Type	Description
str or None	The formatted string if return_str=True, else None.

Source code in ceridwen/ssps/ssp_data.py

def display(self, *, return_str: bool = False, file=None):
    """Print a summary of the grid and its provenance.

    Intended as a sanity check: call it right after ``from_fsps`` or
    ``load`` to confirm the isochrone set, spectral library, IMF, and
    grid coverage are what you expect before you build a ``CSPBasis``.
    Purely diagnostic — none of this is touched by the forward model.

    Parameters
    ----------
    return_str : bool, optional
        Return the formatted string instead of printing it.  Default False.
    file : file-like, optional
        Destination for the print (default ``sys.stdout``).

    Returns
    -------
    str or None
        The formatted string if ``return_str=True``, else ``None``.
    """
    import sys as _sys

    lgmet = np.asarray(self.ssp_lgmet)
    lgage = np.asarray(self.ssp_lg_age_gyr)
    wave  = np.asarray(self.ssp_wave)
    n_met, n_age, n_wave = self.ssp_flux.shape
    age_gyr = 10.0 ** lgage

    def _fmt(v, na="—"):
        return na if v is None else str(v)

    # Human-readable in-memory size of the (dominant) flux array.
    size = float(np.asarray(self.ssp_flux).nbytes)
    for unit in ("B", "KB", "MB", "GB", "TB"):
        if size < 1024.0 or unit == "TB":
            size_str = f"{size:.1f} {unit}"
            break
        size /= 1024.0

    lines = [
        "SSPData",
        "-" * 66,
        "provenance",
        f"  isochrones (isoc_type)   : {_fmt(self.isoc_type)}",
        f"  spectral library         : {_fmt(self.spec_library)}",
        f"  IMF (imf_type)           : {_fmt(self.imf_type)}",
        f"  FSPS version             : {_fmt(self.fsps_version)}",
        f"  schema version           : {_fmt(self.schema_version)}",
        f"  recorded wave_min/max    : {_fmt(self.wave_min)} / {_fmt(self.wave_max)}",
        f"  build kwargs             : {self.fsps_kwargs or '{}'}",
        "grids",
        f"  metallicity  log10 Z     : {n_met:>4d} pts   "
        f"[{lgmet.min():+.3f}, {lgmet.max():+.3f}]  (absolute Z, NOT Z/Zsun)",
        f"  age          log10(Gyr)  : {n_age:>4d} pts   "
        f"[{lgage.min():+.3f}, {lgage.max():+.3f}]  "
        f"= [{age_gyr.min():.3g}, {age_gyr.max():.3g}] Gyr",
        f"  wavelength   Angstrom    : {n_wave:>4d} pts   "
        f"[{wave.min():.1f}, {wave.max():.1f}]",
        f"  flux (n_met,n_age,n_wave): {tuple(int(s) for s in self.ssp_flux.shape)}  "
        f"[L_sun Hz^-1 M_sun^-1]  {np.asarray(self.ssp_flux).dtype}  {size_str}",
    ]
    if self.isoc_type is None:
        lines += [
            "note",
            "  isoc_type is None (legacy grid, built before provenance "
            "tracking):",
            "  CSPBasis will warn and fall back to 'mist' for the nebular grid.",
        ]

    txt = "\n".join(lines)
    if return_str:
        return txt
    print(txt, file=file or _sys.stdout)
    return None

save

save(filename)

Serialise the SSP grids (and provenance metadata) to HDF5.

Provenance fields are written to the file attrs; any that are None are simply omitted, so re-saving a legacy grid does not invent metadata. fsps_kwargs is stored as a JSON string.

Parameters:

Name	Type	Description	Default
filename	str or Path	Output file path. Will be overwritten if it exists.	required

Source code in ceridwen/ssps/ssp_data.py

def save(self, filename):
    """
    Serialise the SSP grids (and provenance metadata) to HDF5.

    Provenance fields are written to the file ``attrs``; any that are
    ``None`` are simply omitted, so re-saving a legacy grid does not
    invent metadata.  ``fsps_kwargs`` is stored as a JSON string.

    Parameters
    ----------
    filename : str or Path
        Output file path.  Will be overwritten if it exists.
    """
    with h5py.File(filename, 'w') as f:
        f.create_dataset('ssp_lgmet',      data=np.array(self.ssp_lgmet))
        f.create_dataset('ssp_lg_age_gyr', data=np.array(self.ssp_lg_age_gyr))
        f.create_dataset('ssp_wave',       data=np.array(self.ssp_wave))
        f.create_dataset('ssp_flux',       data=np.array(self.ssp_flux))

        f.attrs['description']        = 'FSPS SSP interpolation grids'
        f.attrs['units_lgmet']        = 'log10(absolute_metallicity)'
        f.attrs['units_lg_age_gyr']   = 'log10(age/Gyr)'
        f.attrs['units_wave']         = 'Angstrom'
        f.attrs['units_flux']         = 'L_sun Hz^-1 M_sun^-1'

        # --- provenance -------------------------------------------------
        for key in ('schema_version', 'isoc_type', 'spec_library',
                    'fsps_version'):
            val = getattr(self, key)
            if val is not None:
                f.attrs[key] = str(val)
        if self.imf_type is not None:
            f.attrs['imf_type'] = int(self.imf_type)
        if self.wave_min is not None:
            f.attrs['wave_min'] = float(self.wave_min)
        if self.wave_max is not None:
            f.attrs['wave_max'] = float(self.wave_max)
        # Always record the build-kwargs dict (possibly empty) as JSON.
        f.attrs['fsps_kwargs_json'] = json.dumps(self.fsps_kwargs or {})

load classmethod

load(filename)

Load an :class:SSPData from an HDF5 file.

Backward compatible: files written before provenance tracking carry none of the metadata attrs, so every provenance field is populated as None / {} rather than raising. Any legacy log_qq dataset is silently ignored — the nebular model computes its own ionising-photon rate from ssp_flux.

Source code in ceridwen/ssps/ssp_data.py

@classmethod
def load(cls, filename):
    """
    Load an :class:`SSPData` from an HDF5 file.

    Backward compatible: files written before provenance tracking carry
    none of the metadata attrs, so every provenance field is populated
    as ``None`` / ``{}`` rather than raising.  Any legacy ``log_qq``
    dataset is silently ignored — the nebular model computes its own
    ionising-photon rate from ``ssp_flux``.
    """
    def _decode(v):
        if isinstance(v, (bytes, bytearray)):
            return v.decode()
        return v

    with h5py.File(filename, 'r') as f:
        ssp_lgmet      = jnp.array(f['ssp_lgmet'][:])
        ssp_lg_age_gyr = jnp.array(f['ssp_lg_age_gyr'][:])
        ssp_wave       = jnp.array(f['ssp_wave'][:])
        ssp_flux       = jnp.array(f['ssp_flux'][:])

        a = f.attrs
        meta = {
            'schema_version': _decode(a['schema_version'])
                              if 'schema_version' in a else None,
            'isoc_type':      _decode(a['isoc_type'])
                              if 'isoc_type' in a else None,
            'spec_library':   _decode(a['spec_library'])
                              if 'spec_library' in a else None,
            'fsps_version':   _decode(a['fsps_version'])
                              if 'fsps_version' in a else None,
            'imf_type':       int(a['imf_type']) if 'imf_type' in a else None,
            'wave_min':       float(a['wave_min']) if 'wave_min' in a else None,
            'wave_max':       float(a['wave_max']) if 'wave_max' in a else None,
        }
        if 'fsps_kwargs_json' in a:
            meta['fsps_kwargs'] = json.loads(_decode(a['fsps_kwargs_json']))
        else:
            meta['fsps_kwargs'] = {}

    return cls(ssp_lgmet, ssp_lg_age_gyr, ssp_wave, ssp_flux, **meta)

from_fsps classmethod

from_fsps(save_to=None, **fsps_kwargs)

Build an :class:SSPData directly from FSPS, recording provenance.

Only kwargs that define the stellar library / IMF are accepted — the things that legitimately belong at SSP-build time. Anything the CSP forward model applies itself (star-formation history, dust, nebular emission, IGM, redshift, LOSVD smoothing, or a fixed metallicity) raises :class:ValueError, so the grid can never be silently double-processed or made inconsistent with :class:ceridwen.csp.CSPBasis. zcontinuous and sfh are fixed internally (the grid is built on FSPS's discrete zlegend metallicity points).

Parameters:

Name	Type	Description	Default
save_to	str or Path	If given, the result is also persisted to this path via :meth:save so subsequent runs can use :meth:load.	None
**fsps_kwargs		Forwarded to :class:fsps.StellarPopulation. Allowed keys are the IMF parameters (imf_type, imf1, imf2, imf3, imf_lower_limit, imf_upper_limit, vdmc, mdave) and stellar-evolution / isochrone-phase knobs (tpagb_norm_type, agb, pagb, redgb, fbhb, sbss, delt, dell, evtype, masscut, use_wr_spectra, logt_wmb_hot, add_stellar_remnants, fcstar). Common choice: imf_type=1 (Chabrier).	{}

Raises:

Type	Description
ValueError	If any kwarg is not a library/IMF-defining parameter.

Source code in ceridwen/ssps/ssp_data.py

@classmethod
def from_fsps(cls, save_to: Optional[str] = None,
              **fsps_kwargs) -> "SSPData":
    """
    Build an :class:`SSPData` directly from FSPS, recording provenance.

    Only kwargs that define the **stellar library / IMF** are accepted —
    the things that legitimately belong at SSP-build time.  Anything the
    CSP forward model applies itself (star-formation history, dust,
    nebular emission, IGM, redshift, LOSVD smoothing, or a fixed
    metallicity) raises :class:`ValueError`, so the grid can never be
    silently double-processed or made inconsistent with
    :class:`ceridwen.csp.CSPBasis`.  ``zcontinuous`` and ``sfh`` are
    fixed internally (the grid is built on FSPS's discrete ``zlegend``
    metallicity points).

    Parameters
    ----------
    save_to : str or Path, optional
        If given, the result is also persisted to this path via
        :meth:`save` so subsequent runs can use :meth:`load`.
    **fsps_kwargs
        Forwarded to :class:`fsps.StellarPopulation`.  Allowed keys are
        the IMF parameters (``imf_type``, ``imf1``, ``imf2``, ``imf3``,
        ``imf_lower_limit``, ``imf_upper_limit``, ``vdmc``, ``mdave``)
        and stellar-evolution / isochrone-phase knobs (``tpagb_norm_type``,
        ``agb``, ``pagb``, ``redgb``, ``fbhb``, ``sbss``, ``delt``,
        ``dell``, ``evtype``, ``masscut``, ``use_wr_spectra``,
        ``logt_wmb_hot``, ``add_stellar_remnants``, ``fcstar``).
        Common choice: ``imf_type=1`` (Chabrier).

    Raises
    ------
    ValueError
        If any kwarg is not a library/IMF-defining parameter.
    """
    data = collect_ssp_data_wrapper(**fsps_kwargs)
    if save_to is not None:
        data.save(save_to)
    return data

Composite stellar population (forward model)

ceridwen.csp.CSPBasis

CSPBasis(
    SSPData,
    theta=None,
    tuniv=13.8,
    tiny_logt=-70,
    zh_const=False,
    add_neb=True,
    init_neb_params=None,
    nebemlineinspec=False,
    add_dust=True,
    add_diffuse_dust=True,
    add_dust_emission=False,
    add_igm=False,
    igm_model="madau1995",
    igm_factor=1.0,
    sps_home=None,
    init_dust_params=None,
    diffuse_law="kriek_conroy",
    verbose=True,
    sfh_interp="step",
    sigma_losvd_kms=300.0,
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    **kwargs
)

Composite Stellar Population basis using a dict-valued theta.

The public interface is identical to csp.CSPBasis except that predict(theta) now expects (and theta_init now is) a dict[str, Array] rather than a flat 1-D jnp.ndarray.

Parameters:

Name	Type	Description	Default
SSPData	SSPData	Frozen dataclass with SSP grids (wave, flux, ages, zmet). The nebular model computes its own ionising-photon rate from SSPData.ssp_flux at construction time, so log_qq is no longer carried in the SSP grid container.	required
theta	dict	Initial parameter values. Must contain "sfh" and "lookback_time" (>= 2 nodes; n_time nodes define n_time-1 SFH bins). All other keys are optional. Instead of theta you can pass the lookback_time= shortcut (below), which fills the initial values with neutral defaults — use theta only when you want control over the initial values themselves (e.g. a specific initial SFH for display_sfh or a chosen starting point). "lookback_time" is required even when the redshift is sampled: with track_zred_age=True the forward pass rescales this grid to track age_gyr(zred), preserving its length and relative node spacing — the construction-time grid is the template that fixes n_time and the bin structure. The construction-time grid is a default, not a straitjacket: an explicit theta["lookback_time"] passed to predict / get_spectrum takes precedence and is used verbatim in the weight kernel (see _ssp_weights), e.g. for transform-derived grids computed from a sampled zred. Per-call grids are traced values and therefore CANNOT be validated (monotonicity, range) inside the compiled path — that fail-fast validation happens only on the concrete construction-time grid, which is why one is required here. Lookback-time convention theta["lookback_time"] is monotonically increasing in Gyr, with index 0 the present-day node (≈ 0 Gyr) and the last index the oldest sampled node (≤ tuniv). theta["sfh"] is indexed to match: sfh[0] is the SFR at today (per-node input) or the SFR of the youngest bin (per-bin / FastStepBasis input). Likewise theta["zh"] (per-node) has zh[0] = today's metallicity. Example:: T_univ = 13.8 lookback = jnp.linspace(0.0, T_univ, 10) # today → oldest sfh = jnp.exp(-lookback / 1.0) # late-assembly burst theta = {"lookback_time": lookback, "sfh": sfh, "Z": jnp.array([-0.5])} A decreasing grid (``lookback = T_univ - t_grid``) raises a ``ValueError`` at construction — see the assertion in ``initialize_model_structure``.	None
tuniv	float	Age of the Universe in Gyr. Default 13.8.	13.8
zh_const	bool	If True, use constant metallicity (requires key "Z"). If False, use time-varying metallicity (requires key "zh").	False
add_neb	bool	Physics switches.	True
add_dust	bool	Physics switches.	True
add_diffuse_dust	bool	Physics switches.	True
add_dust_emission	bool	Physics switches.	True
sps_home	str	Path to the FSPS data directory (needed for nebular and dust-emission grid loading). Defaults to the $SPS_HOME environment variable that FSPS users set on install; pass explicitly to override. Required only when add_neb or add_dust_emission is True.	None
init_neb_params	dict	Keyword arguments forwarded to NebularModel / Dust. isoc_type no longer needs to be set here: it is taken automatically from the SSP grid's recorded provenance (SSPData.isoc_type), so the nebular CLOUDY grid always matches the SSP isochrone set. Passing an isoc_type that conflicts with the grid raises ValueError; passing one for a legacy grid with no recorded library is honoured. A legacy grid with no recorded library and no explicit isoc_type warns and falls back to 'mist'.	None
init_dust_params	dict	Keyword arguments forwarded to NebularModel / Dust. isoc_type no longer needs to be set here: it is taken automatically from the SSP grid's recorded provenance (SSPData.isoc_type), so the nebular CLOUDY grid always matches the SSP isochrone set. Passing an isoc_type that conflicts with the grid raises ValueError; passing one for a legacy grid with no recorded library is honoured. A legacy grid with no recorded library and no explicit isoc_type warns and falls back to 'mist'.	None
diffuse_law	str	Attenuation law name for the diffuse dust component.	'kriek_conroy'
verbose	bool	Print parameter summary after initialization.	True
lookback_time	array - like	Shortcut alternative to theta: the static SFH node grid (Gyr, monotonically increasing, index 0 = today, >= 2 nodes). Initial values are filled with neutral defaults (sfh = 1 in every node/bin; metallicity = the median of the SSP grid, guaranteed in-grid). Mutually exclusive with theta. Example:: csp = CSPBasis(ssp, lookback_time=jnp.linspace(0.0, 12.0, 6), zh_const=True, add_neb=False)	None
sfh_per_bin	bool	Only used with the lookback_time= shortcut: if True, the SFH is one SFR per bin (shape (n_time-1,), FastStepBasis / prospector convention) instead of one per node (shape (n_time,)). Default False. (With theta= the convention is inferred from the shape of theta['sfh'].)	False

sfh_interp : {'step', 'linear'} Controls the SFH integration scheme used when computing SSP weights.

``'step'`` (default) — piecewise-constant (FastStepBasis-style).
    The SFR is held at the mean of the two endpoint values within
    each SFH time bin.  The weight of each SSP age bin is the
    product of that constant SFR and the linear-time overlap between
    the SFH bin and the SSP age bin.  Weights are non-negative by
    construction — no clipping is ever needed.

``'linear'`` — piecewise-linear.
    Analytically integrates a linearly-interpolated SFH against the
    SSP age bins in log-age space (``intsfwght``).  Higher-order
    accurate, but can produce small negative weights for steep SFH
    gradients, which are then clipped.

To switch at runtime::

csp.calculate_ssp_weights = csp.calculate_ssp_weights_const_zh_step
# or
csp.calculate_ssp_weights = csp.calculate_ssp_weights_const_zh

Source code in ceridwen/csp/csp.py

def __init__(
    self,
    SSPData,
    theta=None,
    tuniv=13.8,
    tiny_logt=-70,
    zh_const=False,
    add_neb=True,
    init_neb_params=None,
    nebemlineinspec=False,
    add_dust=True,
    add_diffuse_dust=True,
    add_dust_emission=False,
    add_igm=False,
    igm_model="madau1995",
    igm_factor=1.0,
    sps_home=None,
    init_dust_params=None,
    diffuse_law='kriek_conroy',
    verbose=True,
    sfh_interp='step',
    sigma_losvd_kms=300.0,
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    **kwargs,
):
    """
    sfh_interp : {'step', 'linear'}
        Controls the SFH integration scheme used when computing SSP weights.

        ``'step'`` (default) — piecewise-constant (FastStepBasis-style).
            The SFR is held at the mean of the two endpoint values within
            each SFH time bin.  The weight of each SSP age bin is the
            product of that constant SFR and the linear-time overlap between
            the SFH bin and the SSP age bin.  Weights are non-negative by
            construction — no clipping is ever needed.

        ``'linear'`` — piecewise-linear.
            Analytically integrates a linearly-interpolated SFH against the
            SSP age bins in log-age space (``intsfwght``).  Higher-order
            accurate, but can produce small negative weights for steep SFH
            gradients, which are then clipped.

    To switch at runtime::

        csp.calculate_ssp_weights = csp.calculate_ssp_weights_const_zh_step
        # or
        csp.calculate_ssp_weights = csp.calculate_ssp_weights_const_zh
    """
    # --- Shortcut construction: lookback_time= instead of theta= --------
    # The init theta exists to fix STATIC structure (n_time + node spacing,
    # the per-node/per-bin sfh convention, the Z-vs-zh metallicity mode) —
    # things the JIT-compiled kernels bake in at trace time.  The initial
    # VALUES only seed theta_init and the early range check, so the
    # shortcut fills them with neutral defaults: sfh = 1 everywhere and
    # the median metallicity of the SSP grid (always in-grid).  Every
    # predict/get_spectrum call still takes its own theta as usual.
    if lookback_time is not None:
        if theta is not None:
            raise ValueError(
                "Pass either theta= (full control over the initial "
                "parameter values) or the lookback_time= shortcut, not "
                "both."
            )
        _lb = jnp.atleast_1d(jnp.asarray(lookback_time, dtype=float))
        _n = int(_lb.size)
        theta = {
            'lookback_time': _lb,
            'sfh': jnp.ones(max(_n - 1, 1) if sfh_per_bin else _n),
        }
        _z_mid = float(jnp.median(jnp.asarray(SSPData.ssp_lgmet)))
        if zh_const:
            theta['Z'] = jnp.array([_z_mid])
        else:
            theta['zh'] = jnp.full((_n,), _z_mid)
    if theta is None:
        raise ValueError(
            "CSPBasis needs the static SFH grid structure. Pass either\n"
            "  lookback_time=jnp.linspace(0.0, T_oldest, n_nodes)   "
            "(shortcut; neutral initial values), or\n"
            "  theta={'lookback_time': ..., 'sfh': ..., 'Z' or 'zh': ...} "
            "(full control).\n"
            "lookback_time is in Gyr, monotonically increasing, index 0 = "
            "today, >= 2 nodes."
        )
    if init_neb_params is None:
        # isoc_type is intentionally NOT set here: it is taken from the
        # SSP grid's recorded provenance (see initialize_neb). Pass
        # init_neb_params={'isoc_type': ...} only to override.
        init_neb_params = {"cloudy_dust": True}
    if init_dust_params is None:
        init_dust_params = {'bin_edges': [(-jnp.inf, -1.97)], 'laws': ['powerlaw']}

    # --- SSP grids (static, never part of theta) -----------------------
    self.flux      = jnp.array(SSPData.ssp_flux, dtype=jnp.float32)  # (n_z, n_age, n_wave)
    self.wave      = jnp.array(SSPData.ssp_wave)       # (n_wave,)
    self.ages      = jnp.array(SSPData.ssp_lg_age_gyr) # (n_age,)  log10(Gyr)
    self.zmet      = jnp.array(SSPData.ssp_lgmet)      # (n_z,) log10 absolute Z
    # The nebular model computes the ionising-photon rate from ``self.flux``
    # internally (see ``initialize_neb`` / ``NebularModel.compute_log_qq``).
    self.zlegend   = 10 ** self.zmet                   # linear metallicity
    self.ssp_ages_lgyr = self.ages + 9                 # log10(yr)

    # Static provenance carried by the SSP grid (Python-level only, never a
    # JAX leaf). ``isoc_type`` is auto-propagated to the nebular model so
    # users never have to set it by hand; ``None`` for legacy grids.
    self._ssp_isoc_type    = getattr(SSPData, "isoc_type", None)
    self._ssp_spec_library = getattr(SSPData, "spec_library", None)

    # Precomputed constants for calculate_ssp_weights (all static)
    self._logage_lo  = self.ssp_ages_lgyr[1:]
    self._logage_hi  = self.ssp_ages_lgyr[:-1]
    self._dlogage    = jnp.diff(self.ssp_ages_lgyr)
    self._j_range    = jnp.arange(self.ssp_ages_lgyr.size)
    self._age_clip_lo = 10.0 ** (-70)                  # floor for log-time clipping
    self._age_clip_hi = 10.0 ** self.ssp_ages_lgyr[-1] # ceiling
    self._n_z   = len(self.zmet)
    self._n_age = len(self.ages)

    # SSP bin edges in linear years. _ssp_lo_yr is the younger (smaller)
    # edge; _ssp_hi_yr is the older (larger) edge of each SSP age bin.
    self._ssp_lo_yr = 10.0 ** self._logage_hi   # (n_age-1,)
    self._ssp_hi_yr = 10.0 ** self._logage_lo   # (n_age-1,)

    # Voronoi cell boundaries for the step-function weight scheme.
    # Each SSP age POINT j owns the linear-time interval
    #   [_ssp_voronoi_lo[j], _ssp_voronoi_hi[j]]
    # where the boundaries are the midpoints to the neighbouring age points.
    # For a piecewise-constant SFH the weight at SSP j equals the
    # SFR * (width of its Voronoi cell in yr); this matches FSPS
    # FastStepBasis's internal ±ε offset scheme.
    #
    # Boundary handling:
    #   - youngest SSP (j=0): lower bound set to 0.
    #   - oldest  SSP (j=-1): upper bound set to 2× the last inter-point
    #     spacing, which safely exceeds any realistic SFH extent.
    _ssp_age_yr  = 10.0 ** self.ssp_ages_lgyr           # (n_age,) linear yr
    _voro_mid    = 0.5 * (_ssp_age_yr[:-1] + _ssp_age_yr[1:])  # (n_age-1,)
    _voro_hi_ext = _ssp_age_yr[-1] + (_ssp_age_yr[-1] - _ssp_age_yr[-2])
    self._ssp_voronoi_lo = jnp.concatenate(
        [jnp.zeros(1), _voro_mid]
    )   # (n_age,)  — lower boundary of each Voronoi cell
    self._ssp_voronoi_hi = jnp.concatenate(
        [_voro_mid, jnp.array([_voro_hi_ext])]
    )   # (n_age,)  — upper boundary of each Voronoi cell

    self.tuniv      = tuniv
    self.tiny_logt  = tiny_logt
    # Resolve the FSPS data directory: explicit arg wins, else $SPS_HOME
    # (the variable FSPS users already set on install).
    if sps_home is None:
        sps_home = os.environ.get("SPS_HOME")
    if (add_neb or add_dust_emission) and not sps_home:
        raise ValueError(
            "sps_home is required for nebular / dust emission but was not "
            "given and $SPS_HOME is unset. Set `export SPS_HOME=/path/to/fsps` "
            "(your FSPS data directory) or pass sps_home=... explicitly."
        )
    self.sps_home   = sps_home
    # Free-redshift age-grid tracking.  When True AND theta carries a
    # sampled ``zred`` (and NO explicit ``lookback_time``), the SFH
    # lookback grid is rescaled inside the forward pass so its oldest node
    # tracks the age of the universe at the sampled redshift, via the
    # differentiable :func:`ceridwen.cosmology.age_gyr` (see
    # :meth:`_lookback_from_zred` and :meth:`_ssp_weights`).  Default
    # False keeps the fixed-z path bit-for-bit unchanged.
    self.track_zred_age = bool(track_zred_age)
    # Prospector-style ``nebemlineinspec`` switch.  It governs ONE
    # thing only: the default of the single-array public
    # ``get_spectrum(theta)``.  When False (default), that call
    # returns the line-free continuum (stellar + nebular continuum);
    # when True it returns continuum + emission lines.  It does NOT
    # affect ``predict`` or ``get_spectrum_components``, which always
    # compute the full ``(continuum, lines)`` decomposition so the
    # observations always see the lines (Photometry at true strength,
    # Spectrum / Lines scaled by ``eline_scaling``).  To force lines
    # on/off explicitly, pass ``include_lines=`` to ``get_spectrum`` or
    # use ``get_spectrum_components``.
    self.nebemlineinspec = bool(nebemlineinspec)

    # --- IGM attenuation model (optional) ------------------------------
    # ``add_igm=False`` leaves ``self.igm`` as None; ``CSPBasis.predict``
    # then skips the multiplicative step entirely (zero Python
    # branches in the traced hot path — the ``is None`` is a
    # compile-time decision).  When ``add_igm=True`` the model (by
    # default Madau 1995, identical to FSPS's ``igm_absorb.f90``)
    # is applied whenever ``theta['zred']`` is present, with
    # optional runtime strength override via ``theta['igm_factor']``.
    if add_igm:
        from ..igm import make_igm_model
        self.igm = make_igm_model(igm_model)
    else:
        self.igm = None
    self.igm_factor = float(igm_factor)

    # --- Dust attenuation function (set before dust init) --------------
    if add_diffuse_dust or add_dust:
        self.set_attenuation_function(add_diffuse_dust, add_dust)

    # --- Sub-model init (populates defaults into theta) -----------
    theta = self.initialize_dust_components(
        add_dust, add_diffuse_dust, add_dust_emission,
        theta, init_dust_params, diffuse_law, sps_home,
    )
    theta = self.initialize_neb(add_neb, theta, init_neb_params, sps_home)

    # Pre-compute the FSPS-default LOSVD smoothing kernel (sigma_smooth
    # default 300 km/s, velocity-space Gaussian on rest-frame
    # 912 < lambda < 25000 AA; matches prospect/models/sedmodel.py
    # losvd_smoothing).  Must run BEFORE configure_spectrum_model so the
    # wrap can see whether to install the smoother.
    self.sigma_losvd_kms = float(sigma_losvd_kms)
    self._setup_losvd_kernel()

    self.configure_spectrum_model(
        add_dust, add_diffuse_dust, add_dust_emission, add_neb, sps_home
    )

    # --- SFH integration scheme selection ---------------------------------
    # 'step'   → piecewise-constant, guaranteed non-negative (default)
    # 'linear' → piecewise-linear log-age integration (original scheme)
    if sfh_interp not in ('step', 'linear'):
        raise ValueError(
            f"sfh_interp must be 'step' or 'linear', got {sfh_interp!r}"
        )
    self.sfh_interp = sfh_interp
    self.zh_const = bool(zh_const)
    if zh_const:
        if sfh_interp == 'step':
            self.calculate_ssp_weights = self.calculate_ssp_weights_const_zh_step
        else:
            self.calculate_ssp_weights = self.calculate_ssp_weights_const_zh
    else:
        if sfh_interp == 'step':
            self.calculate_ssp_weights = self.calculate_ssp_weights_var_zh_step
        else:
            self.calculate_ssp_weights = self.calculate_ssp_weights_var_zh
    if verbose:
        print(f"SFH integration scheme : {sfh_interp}")

    self.initialize_model_structure(theta)

    if verbose:
        print("\nCSPBasis (dict theta) — registered parameters:")
        pprint.pprint({k: v.shape for k, v in self.theta_init.items()})

    # Early (construction-time) warning if the initial parameters already
    # sit outside the interpolation grids (silent edge-clamping).  Cheap,
    # non-jitted; users can re-run check_param_ranges() on sampled theta.
    self.check_param_ranges(self.theta_init)

all_params property

all_params

Return the initial theta dict (one entry per free parameter).

initialize_model_structure

initialize_model_structure(theta)

Validate the incoming theta and store self.theta_init. The dict is validated, converted to JAX arrays, and stored directly.

Required keys

"sfh" : shape (n_time,) "lookback_time" : shape (n_time,)

Either "Z" (scalar, constant metallicity) or "zh" (shape (n_time,), time-varying metallicity) must be present, depending on the zh_const flag set during __init__.

Source code in ceridwen/csp/csp.py

def initialize_model_structure(self, theta):
    """
    Validate the incoming theta and store ``self.theta_init``.
    The dict is validated, converted to JAX arrays, and stored directly.

    Required keys
    -------------
    ``"sfh"``          : shape ``(n_time,)``
    ``"lookback_time"`` : shape ``(n_time,)``

    Either ``"Z"`` (scalar, constant metallicity) or ``"zh"`` (shape
    ``(n_time,)``, time-varying metallicity) must be present, depending on
    the ``zh_const`` flag set during ``__init__``.
    """
    # --- Required keys: nice errors instead of raw KeyErrors -----------
    if 'lookback_time' not in theta:
        raise ValueError(
            "theta must contain 'lookback_time' — the static SFH node grid "
            "(Gyr, monotonically increasing, index 0 = today). It is "
            "required even when the redshift is a free parameter: with "
            "track_zred_age=True the grid is rescaled inside the forward "
            "pass to track age(zred), but its LENGTH and RELATIVE spacing "
            "come from this construction-time grid, so it defines n_time "
            "and the bin structure rather than a fixed absolute age range."
        )
    if 'sfh' not in theta:
        raise ValueError(
            "theta must contain 'sfh' — star-formation-rate values, either "
            "one per lookback node (shape (n_time,)) or one per bin "
            "(shape (n_time-1,), FastStepBasis convention), where n_time = "
            "len(theta['lookback_time'])."
        )

    # --- sfh_times: static (not part of theta) -------------------------
    self.sfh_times = jnp.atleast_1d(
        jnp.asarray(theta['lookback_time'], dtype=float)
    ) * 1e9   # Gyr → yr
    self.n_time = self.sfh_times.size

    # --- Minimum grid size ---------------------------------------------
    # n_time nodes define n_time-1 SFH bins; with fewer than 2 nodes there
    # is no bin to integrate over. (This must precede the monotonicity
    # check: np.diff of a single node is empty and np.all([]) is True, so
    # a 1-node grid would otherwise slip through and fail later inside a
    # jitted weight kernel with a cryptic shape error.)
    if self.n_time < 2:
        raise ValueError(
            f"theta['lookback_time'] has {self.n_time} node(s); at least "
            "2 are required (n_time nodes define n_time-1 SFH bins). "
            "Typical fits use 5-10 nodes, e.g. "
            "jnp.linspace(0.0, T_UNIV, 6)."
        )

    # Convention check: lookback_time must be monotonically *increasing*,
    # starting at 0 (today).  A decreasing grid trips here loudly rather
    # than silently producing wrong-physics weights.
    _lb = np.asarray(self.sfh_times, dtype=np.float64)
    _diffs = np.diff(_lb)
    if not (np.all(_diffs > 0.0) and _lb[0] >= 0.0 and _lb[0] < 1e8):
        raise ValueError(
            "theta['lookback_time'] must be monotonically *increasing* "
            "(NEW convention, post-2026-06-03 refactor):\n"
            f"  - index 0 = today (≈ 0 Gyr): got {_lb[0]/1e9:.3f} Gyr\n"
            f"  - index -1 = oldest (≈ T_univ): got {_lb[-1]/1e9:.3f} Gyr\n"
            f"  - first three values [Gyr]: {(_lb[:3]/1e9).tolist()}\n"
            "If you see this from a pre-refactor script, replace e.g.\n"
            "    lookback = T_UNIV - jnp.linspace(eps, T_UNIV, N)\n"
            "with\n"
            "    lookback = jnp.linspace(0.0, T_UNIV, N)\n"
            "and reverse theta['sfh'] (and theta['zh'] if present) to match."
        )

    sfh = jnp.atleast_1d(jnp.asarray(theta['sfh'], dtype=float))
    # ``sfh`` may carry either of two conventions:
    #
    # 1. FastStepBasis (prospector-compatible) — one SFR value per
    #    bin, length ``n_time - 1``.  ``calculate_ssp_weights_*_step``
    #    uses each entry directly, with no inter-edge averaging.
    # 2. Node-based legacy — one SFR value per lookback grid point,
    #    length ``n_time``.  ``calculate_ssp_weights_*_step``
    #    averages consecutive entries to recover per-bin SFR.
    #
    # Both shapes are accepted; the convention is stored as a flag the
    # weight calculators consult.  With the FastStepBasis convention,
    # the same parameter numbers mean the same physical SFH in ceridwen
    # and prospector.
    if sfh.shape == (self.n_time,):
        self.sfh_per_bin = False
    elif sfh.shape == (self.n_time - 1,):
        self.sfh_per_bin = True
    else:
        raise AssertionError(
            f"'sfh' shape {sfh.shape} must be either "
            f"({self.n_time},)  (node-based, legacy)  or "
            f"({self.n_time - 1},)  (per-bin, FastStepBasis)."
        )

    # --- Static SFH sanity (construction-time; non-jitted) -------------
    # A NaN/Inf SFH silently propagates to a NaN spectrum, and an all- or
    # partly-negative SFH is silently clipped to >=0 (≈zero flux), so flag
    # both here rather than letting them pass into the hot path unnoticed.
    sfh_np = np.asarray(sfh)
    if not np.all(np.isfinite(sfh_np)):
        raise ValueError(
            "theta['sfh'] contains non-finite (NaN/Inf) values; this would "
            "silently produce a NaN spectrum."
        )
    if np.any(sfh_np < 0):
        warnings.warn(
            "theta['sfh'] contains negative values. SFR is clipped to >=0 "
            "internally, so negative bins contribute ~zero flux (no error is "
            "raised at evaluation time).",
            stacklevel=3,
        )

    # --- Metallicity mode detection + validation -----------------------
    # The metallicity key MUST match the zh_const mode chosen at __init__.
    # A mismatch otherwise constructs silently and only fails later with a
    # cryptic KeyError deep inside a jitted get_spectrum/predict trace.
    if self.zh_const:
        if 'Z' not in theta:
            raise ValueError(
                "zh_const=True requires a constant metallicity theta['Z'] "
                "(shape-(1,) array, log10 absolute metallicity in ssp_lgmet "
                "grid units); none was provided. Either add theta['Z'], or "
                "construct with zh_const=False and provide a time-varying "
                "theta['zh'] of shape (n_time,)."
            )
        if 'zh' in theta:
            warnings.warn(
                "zh_const=True but theta also contains 'zh'; 'zh' is ignored "
                "in constant-metallicity mode (only 'Z' is used).",
                stacklevel=3,
            )
    else:
        if 'zh' not in theta:
            raise ValueError(
                "zh_const=False requires a time-varying metallicity history "
                "theta['zh'] of shape (n_time,) (log10 absolute metallicity "
                "in ssp_lgmet grid units, same as theta['Z']); none was "
                "provided. Either add theta['zh'], or construct with "
                "zh_const=True and provide a scalar theta['Z']."
            )
        if 'Z' in theta:
            warnings.warn(
                "zh_const=False but theta also contains 'Z'; 'Z' is ignored "
                "in time-varying-metallicity mode (only 'zh' is used).",
                stacklevel=3,
            )

    self.zh_is_scalar = None
    if 'zh' in theta:
        zh = jnp.atleast_1d(jnp.asarray(theta['zh'], dtype=float))
        assert zh.shape == (self.n_time,), "'zh' must match 'lookback_time' length"
        self.zh_is_scalar = False
    elif 'Z' in theta:
        Z = jnp.atleast_1d(jnp.asarray(theta['Z'], dtype=float))
        assert Z.shape == (1,), "'Z' must be a scalar (wrapped in shape-(1,) array)"
        self.zh_is_scalar = True


    # --- Build theta_init: all params except lookback_time -------------
    self.theta_init = {}
    for k, v in theta.items():
        if k == 'lookback_time':
            continue   # static grid — not a free parameter
        arr = jnp.atleast_1d(jnp.asarray(v, dtype=float))
        self.theta_init[k] = arr

    # Ensure sfh has the correct shape stored in theta_init
    self.theta_init['sfh'] = sfh

    # Ordered list of parameter names (for printing / sampling setup)
    self.param_names = list(self.theta_init.keys())

    # Recognized theta keys (for the trace-time typo guard).  Everything
    # the physics consumes is already in param_names (dust/neb defaults were
    # merged into theta before this point); the rest are optional
    # runtime-only scalars read by predict / get_line_spec.
    self._known_theta_keys = set(self.param_names) | {
        'lookback_time', 'Z', 'zh',
        'logmass', 'zred', 'igm_factor', 'eline_scaling',
        'sigma_smooth', 'frac_obrun',
    }

register_known_theta_keys

register_known_theta_keys(keys)

Register additional recognized theta keys so they are not mis-flagged as typos by :meth:_warn_unknown_theta_keys.

SedModel calls this with its model-level free parameters (e.g. logsfr_ratios, which the sfh transform consumes): those keys are forwarded through to predict in the full theta dict but are not CSP parameters, so without this they would trigger a spurious typo warning.

Source code in ceridwen/csp/csp.py

def register_known_theta_keys(self, keys):
    """Register additional recognized theta keys so they are not mis-flagged
    as typos by :meth:`_warn_unknown_theta_keys`.

    ``SedModel`` calls this with its model-level free parameters (e.g.
    ``logsfr_ratios``, which the ``sfh`` transform consumes): those keys are
    forwarded through to ``predict`` in the full theta dict but are not CSP
    parameters, so without this they would trigger a spurious typo warning.
    """
    self._known_theta_keys |= set(keys)

check_param_ranges

check_param_ranges(theta=None, warn=True)

Diagnostic (NON-jitted): list parameters that fall outside the interpolation grids, where the model silently clamps to the nearest grid edge and thus hides extrapolation.

Intended to be called once on your theta (or theta bounds) before a fit; it is never invoked from the hot path. Returns the list of human-readable messages (and emits them as warnings when warn).

Source code in ceridwen/csp/csp.py

def check_param_ranges(self, theta=None, warn=True):
    """Diagnostic (NON-jitted): list parameters that fall outside the
    interpolation grids, where the model silently clamps to the nearest
    grid edge and thus hides extrapolation.

    Intended to be called once on your theta (or theta bounds) before a
    fit; it is never invoked from the hot path.  Returns the list of
    human-readable messages (and emits them as warnings when ``warn``).
    """
    if theta is None:
        theta = self.theta_init
    msgs = []

    zlo, zhi = float(self.zmet.min()), float(self.zmet.max())
    for key in ('Z', 'zh'):
        if key in theta:
            v = np.asarray(theta[key], float)
            if v.size and (np.nanmin(v) < zlo or np.nanmax(v) > zhi):
                msgs.append(
                    f"theta['{key}'] has values outside the SSP metallicity "
                    f"grid [{zlo:.3f}, {zhi:.3f}]; these are silently clamped "
                    f"to the nearest grid edge. NOTE: this grid is in the "
                    f"same units as SSPData.ssp_lgmet (log10 of absolute "
                    f"metallicity), NOT log10(Z/Zsun) -- so Z=0.0 is out of "
                    f"range; use a value within the printed bounds."
                )

    # Nebular gas parameters vs the CLOUDY grid, if the neb model exposes
    # its axis arrays (defensive: skipped if the attribute names differ).
    neb = getattr(self, 'neb', None)
    if neb is not None:
        for key, attrs in (
            ('gas_logz', ('logZ_grid', 'logz_grid', '_logZ', 'nebem_logz')),
            ('gas_logu', ('logU_grid', 'logu_grid', '_logU', 'nebem_logu')),
        ):
            if key in theta:
                grid = next((getattr(neb, a) for a in attrs if hasattr(neb, a)),
                            None)
                if grid is not None:
                    g = np.asarray(grid, float)
                    glo, ghi = float(g.min()), float(g.max())
                    v = np.asarray(theta[key], float)
                    if v.size and (np.nanmin(v) < glo or np.nanmax(v) > ghi):
                        msgs.append(
                            f"theta['{key}'] outside the nebular grid "
                            f"[{glo:.3f}, {ghi:.3f}]; silently clamped."
                        )

    if warn:
        for m in msgs:
            warnings.warn(m, stacklevel=2)
    return msgs

set_attenuation_function

set_attenuation_function(add_diffuse_dust, add_dust)

Build and assign self.attenuate_dust(wave, theta) → (attn, attn_diffuse).

With the dict theta, each dust model simply reads the keys it knows about from the shared theta dict. No NamedTuple construction needed.

Source code in ceridwen/csp/csp.py

def set_attenuation_function(self, add_diffuse_dust, add_dust):
    """
    Build and assign ``self.attenuate_dust(wave, theta) → (attn, attn_diffuse)``.

    With the dict theta, each dust model simply reads the keys it knows
    about from the shared theta dict.  No NamedTuple construction needed.
    """
    if add_diffuse_dust and add_dust:
        def attenuate(wave, theta):
            attn         = self.dust_attn.compute_attenuation(wave, theta)
            attn_diffuse = self.diff_dust.compute_attenuation(wave, theta)
            return attn, attn_diffuse
        print("Using combined (binwise + diffuse) dust attenuation.")
        self.attenuate_dust = attenuate

    elif add_dust and not add_diffuse_dust:
        def attenuate_without_diffuse(wave, theta):
            attn         = self.dust_attn.compute_attenuation(wave, theta)
            attn_diffuse = jnp.zeros((1, wave.shape[0]))
            return attn, attn_diffuse
        print("Using only binwise dust attenuation.")
        self.attenuate_dust = attenuate_without_diffuse

    elif add_diffuse_dust and not add_dust:
        self.bin_low  = jnp.array([-jnp.inf])
        self.bin_high = jnp.array([jnp.inf])
        def attenuate_diffuse_only(wave, theta):
            attn_diffuse = self.diff_dust.compute_attenuation(wave, theta)
            attn         = jnp.zeros((1, wave.shape[0]))
            return attn, attn_diffuse
        print("Using only diffuse dust attenuation.")
        self.attenuate_dust = attenuate_diffuse_only

get_spectrum_components

get_spectrum_components(theta)

Return the canonical (continuum, lines) line decomposition.

Both arrays are on the rest-frame model grid self.wave and are unscaled -- mass / redshift / IGM factors are applied downstream by predict and get_line_spec, exactly as for get_spectrum.

• continuum -- the line-free spectrum (stellar continuum + nebular continuum), dust-attenuated. Identical to get_spectrum(theta, include_lines=False).
• lines -- the broadened nebular emission-line component alone, carried through the same dust attenuation, so the full SED is recovered as continuum + lines.

This is the single source of truth for "spectrum with vs. without emission lines": predict builds the photometry and slit spectra from it and get_line_spec returns its lines term. nebemlineinspec does not affect this method -- it only sets the default of the single-array public get_spectrum. With add_neb=False there is no nebular module and lines is identically zero.

Source code in ceridwen/csp/csp.py

def get_spectrum_components(self, theta: dict) -> tuple:
    """Return the canonical ``(continuum, lines)`` line decomposition.

    Both arrays are on the rest-frame model grid ``self.wave`` and are
    *unscaled* -- mass / redshift / IGM factors are applied downstream
    by ``predict`` and ``get_line_spec``, exactly as for ``get_spectrum``.

    - ``continuum`` -- the line-free spectrum (stellar continuum +
      nebular *continuum*), dust-attenuated.  Identical to
      ``get_spectrum(theta, include_lines=False)``.
    - ``lines`` -- the broadened nebular emission-line component alone,
      carried through the same dust attenuation, so the full SED is
      recovered as ``continuum + lines``.

    This is the single source of truth for "spectrum with vs. without
    emission lines": ``predict`` builds the photometry and slit spectra
    from it and ``get_line_spec`` returns its ``lines`` term.
    ``nebemlineinspec`` does not affect this method -- it only sets the
    default of the single-array public ``get_spectrum``.  With
    ``add_neb=False`` there is no nebular module and ``lines`` is
    identically zero.
    """
    # Trace-time-only typo guard (operates on static dict keys; costs
    # nothing in the compiled hot path).  Also covers predict().
    self._warn_unknown_theta_keys(theta)
    continuum = self.get_spectrum(theta=theta, include_lines=False)
    full      = self.get_spectrum(theta=theta, include_lines=True)
    return continuum, full - continuum

predict

predict(theta, observations)

Compute the CSP spectrum and project it onto every observation.

This method is the primary hot-path entry point for the sampler. It is designed to be fully JAX JIT-compatible with zero Python if / isinstance branches in the traced code path:

• get_spectrum(theta) is pure JAX.
• The Python for loop over observations is unrolled at trace time because observations is a static Python list (part of the closure, not a traced argument).
• obs.predict(spectrum, self.wave) dispatches through Python's method resolution order (static at trace time) to the appropriate subclass implementation — either a dense matrix–vector multiply (Spectrum, Lines) or a filter-set convolution (Photometry). The XLA kernel contains no conditional branches.

Pre-condition: every Observation in observations must have had obs.setup_for_model(self.wave) called before the first JIT trace. SedModel.__init__ does this automatically.

For a raw model spectrum without projection, use get_spectrum(theta) directly; for the separate line-free continuum and emission-line component, use get_spectrum_components(theta).

Parameters:

Name	Type	Description	Default
theta	dict[str, Array]	Free-parameter dict. Must contain at minimum "sfh" and the metallicity key ("Z" or "zh"), plus any dust / nebular parameters required by the active physics model.	required
observations	list of Observation	Observations to project onto. Must be the same Python objects (same list structure, same types) on every call — changing the list forces a retrace.	required

Returns:

Name	Type	Description
predictions	dict[str, Array]	Keyed by obs.name for each observation. Values: Photometry → shape (n_filters,), synthetic AB maggies Spectrum → shape (n_pix,), model F_nu interpolated onto the observed pixel grid Lines → shape (n_lines,), Gaussian-aperture fluxes
	.. warning::	The outputs are observed-frame AB maggies only if theta contains "zred": that key gates the cosmological flux factor (1+z) (10pc/D_L)^2 (and the L_sun/Hz → cgs conversion) inside :func:ceridwen.cosmology.flux_factor_maggies. Without it the values are raw 10 pc-frame numbers, ~6e21 too bright at z = 0.1. SedModel.predict injects its fixed zred automatically; only direct callers of this method (and of get_line_spec) need to supply it themselves.

Source code in ceridwen/csp/csp.py

def predict(self, theta: dict, observations: list) -> dict:
    """
    Compute the CSP spectrum and project it onto every observation.

    This method is the primary hot-path entry point for the sampler.
    It is designed to be fully JAX JIT-compatible with zero Python
    ``if`` / ``isinstance`` branches in the traced code path:

    - ``get_spectrum(theta)`` is pure JAX.
    - The Python ``for`` loop over ``observations`` is unrolled at
      trace time because ``observations`` is a static Python list
      (part of the closure, not a traced argument).
    - ``obs.predict(spectrum, self.wave)`` dispatches through Python's
      method resolution order (static at trace time) to the appropriate
      subclass implementation — either a dense matrix–vector multiply
      (``Spectrum``, ``Lines``) or a filter-set convolution
      (``Photometry``).  The XLA kernel contains no conditional branches.

    **Pre-condition:** every ``Observation`` in ``observations`` must have
    had ``obs.setup_for_model(self.wave)`` called before the first JIT
    trace.  ``SedModel.__init__`` does this automatically.

    For a raw model spectrum without projection, use ``get_spectrum(theta)``
    directly; for the separate line-free continuum and emission-line
    component, use ``get_spectrum_components(theta)``.

    Parameters
    ----------
    theta : dict[str, Array]
        Free-parameter dict.  Must contain at minimum ``"sfh"`` and the
        metallicity key (``"Z"`` or ``"zh"``), plus any dust / nebular
        parameters required by the active physics model.
    observations : list of Observation
        Observations to project onto.  Must be the same Python objects
        (same list structure, same types) on every call — changing the
        list forces a retrace.

    Returns
    -------
    predictions : dict[str, Array]
        Keyed by ``obs.name`` for each observation.  Values:

        - ``Photometry`` → shape (n_filters,), synthetic AB maggies
        - ``Spectrum``   → shape (n_pix,), model F_nu interpolated onto
          the observed pixel grid
        - ``Lines``      → shape (n_lines,), Gaussian-aperture fluxes

    .. warning::
        The outputs are observed-frame AB maggies **only if** ``theta``
        contains ``"zred"``: that key gates the cosmological flux factor
        ``(1+z) (10pc/D_L)^2`` (and the L_sun/Hz → cgs conversion)
        inside :func:`ceridwen.cosmology.flux_factor_maggies`.  Without
        it the values are raw 10 pc-frame numbers, ~6e21 too bright at
        z = 0.1.  ``SedModel.predict`` injects its fixed ``zred``
        automatically; only direct callers of this method (and of
        ``get_line_spec``) need to supply it themselves.
    """
    spectrum_phot, spectrum_slit = self._assemble_observer_spectra(theta)
    spectrum_phot, spectrum_slit = self._apply_mass_redshift_igm(
        spectrum_phot, spectrum_slit, theta
    )
    return self._project_observations(
        spectrum_phot, spectrum_slit, observations, theta
    )

get_line_spec

get_line_spec(theta)

Return the broadened-emission-line component of the model spectrum.

Computed as get_spectrum(include_lines=True) - get_spectrum(include_lines=False), which gives the contribution of the nebular lines alone -- the prospector-style "line-only" spectrum. Mass + redshift + IGM scaling are applied identically to CSPBasis.predict so the output is at the same physical scale as the observation arrays.

add_neb=False makes this return zero (no nebular module).

Source code in ceridwen/csp/csp.py

def get_line_spec(self, theta):
    """Return the broadened-emission-line component of the model spectrum.

    Computed as ``get_spectrum(include_lines=True) - get_spectrum(include_lines=False)``,
    which gives the contribution of the nebular lines alone -- the
    prospector-style "line-only" spectrum.  Mass + redshift + IGM
    scaling are applied identically to ``CSPBasis.predict`` so the
    output is at the same physical scale as the observation arrays.

    ``add_neb=False`` makes this return zero (no nebular module).
    """
    if not hasattr(self, "neb") or self.neb is None:
        return jnp.zeros_like(self.wave)

    _continuum, line_only = self.get_spectrum_components(theta)

    # Mirror the mass + zred + IGM scaling block of CSPBasis.predict.
    if "logmass" in theta:
        mass_scale = jnp.float32(10.0 ** theta["logmass"][0])
        line_only = line_only * mass_scale
    if "zred" in theta:
        from ..cosmology import flux_factor_maggies
        z_scalar = jnp.ravel(theta["zred"])[0]
        line_only = line_only * jnp.float32(flux_factor_maggies(z_scalar))
        if self.igm is not None:
            if "igm_factor" in theta:
                ig_factor = jnp.ravel(theta["igm_factor"])[0]
            else:
                ig_factor = jnp.float32(self.igm_factor)
            transmission = self.igm.attenuation(
                self.wave, z_scalar, factor=ig_factor,
            )
            line_only = line_only * transmission.astype(line_only.dtype)
    # Apply the eline_scaling fraction (1.0 = no loss) so the line-only
    # spectrum is consistent with the Lines.predict fluxes from predict().
    if "eline_scaling" in theta:
        line_only = line_only * jnp.ravel(theta["eline_scaling"])[0]
    return line_only

display_sfh

display_sfh(
    theta=None,
    ax=None,
    *,
    overlay_nodes=True,
    show_bin_edges=False,
    units="Gyr",
    **plot_kwargs
)

Plot the SFH against lookback time, rendered identically to the interpretation used by :meth:_ssp_weights.

For sfh_interp == "step" this draws a piecewise-constant function with one horizontal segment per bin [T_{i+1}, T_i] at height :math:\bar\psi_i (the per-bin SFR consumed by calculate_ssp_weights_*_step). For sfh_interp == "linear" it draws the piecewise-linear interpolant between per-node SFR values -- the same function whose analytic integral against the SSP age grid is computed by intsfwght.

Lookback time is read from theta["lookback_time"] if supplied (units: Gyr) and otherwise falls back to self.sfh_times (which is stored in years and converted back to Gyr here). theta_init intentionally does NOT carry lookback_time -- it is a static grid, not a free parameter -- so the default fallback path is the common case.

The x-axis runs left-to-right in increasing lookback time: present day (T = 0) sits at the origin on the left, and the oldest sampled node sits on the right. This matches the natural index order of theta["lookback_time"].

Parameters:

Name	Type	Description	Default
theta	dict	Parameter dict to display. Defaults to self.theta_init. If it carries a "lookback_time" entry, that takes precedence over self.sfh_times for the x-axis grid.	None
ax	Axes	Axes to draw into. If None, a new figure is created.	None
overlay_nodes	bool	If True, mark per-bin SFR values at bin midpoints (step mode) or per-node SFR values at lookback nodes (linear mode).	True
show_bin_edges	bool	If True, draw vertical dotted lines at every node T_i.	False
units	(Gyr, yr, Myr)	X-axis units for the lookback-time axis. The SFR axis is always [M_sun / yr].	"Gyr"
**plot_kwargs		Forwarded to the per-segment ax.plot calls (e.g. color, lw, linestyle, label).	{}

Returns:

Name	Type	Description
ax	Axes	The axes containing the plot.

Raises:

Type	Description
AssertionError	If the per-bin integral of the displayed SFR disagrees with the per-bin mass m_target used by :meth:_ssp_weights by more than 1e-6 relative. This pins the visual to the weight code so future refactors of either side cannot silently diverge.

Source code in ceridwen/csp/csp.py

def display_sfh(self, theta=None, ax=None, *,
                overlay_nodes=True, show_bin_edges=False,
                units="Gyr", **plot_kwargs):
    """Plot the SFH against lookback time, rendered identically to the
    interpretation used by :meth:`_ssp_weights`.

    For ``sfh_interp == "step"`` this draws a piecewise-constant function
    with one horizontal segment per bin ``[T_{i+1}, T_i]`` at height
    :math:`\\bar\\psi_i` (the per-bin SFR consumed by
    ``calculate_ssp_weights_*_step``).  For ``sfh_interp == "linear"`` it
    draws the piecewise-linear interpolant between per-node SFR values --
    the same function whose analytic integral against the SSP age grid
    is computed by ``intsfwght``.

    Lookback time is read from ``theta["lookback_time"]`` if supplied
    (units: Gyr) and otherwise falls back to ``self.sfh_times`` (which
    is stored in years and converted back to Gyr here).  ``theta_init``
    intentionally does NOT carry ``lookback_time`` -- it is a static
    grid, not a free parameter -- so the default fallback path is the
    common case.

    The x-axis runs left-to-right in increasing lookback time: present
    day (T = 0) sits at the origin on the left, and the oldest sampled
    node sits on the right.  This matches the natural index order of
    ``theta["lookback_time"]``.

    Parameters
    ----------
    theta : dict, optional
        Parameter dict to display.  Defaults to ``self.theta_init``.
        If it carries a ``"lookback_time"`` entry, that takes
        precedence over ``self.sfh_times`` for the x-axis grid.
    ax : matplotlib.axes.Axes, optional
        Axes to draw into.  If None, a new figure is created.
    overlay_nodes : bool
        If True, mark per-bin SFR values at bin midpoints (step mode)
        or per-node SFR values at lookback nodes (linear mode).
    show_bin_edges : bool
        If True, draw vertical dotted lines at every node ``T_i``.
    units : {"Gyr", "yr", "Myr"}
        X-axis units for the lookback-time axis.  The SFR axis is
        always [M_sun / yr].
    **plot_kwargs
        Forwarded to the per-segment ``ax.plot`` calls (e.g.
        ``color``, ``lw``, ``linestyle``, ``label``).

    Returns
    -------
    ax : matplotlib.axes.Axes
        The axes containing the plot.

    Raises
    ------
    AssertionError
        If the per-bin integral of the displayed SFR disagrees with the
        per-bin mass ``m_target`` used by :meth:`_ssp_weights` by more
        than 1e-6 relative.  This pins the visual to the weight code so
        future refactors of either side cannot silently diverge.
    """
    import matplotlib.pyplot as plt

    theta = self.theta_init if theta is None else theta

    # Lookback-time grid (Gyr).  theta_init does not carry it (stripped
    # in initialize_model_structure), so the default path falls back to
    # self.sfh_times (yr) converted to Gyr.
    if "lookback_time" in theta:
        T_gyr = np.asarray(theta["lookback_time"], dtype=float)
    else:
        T_gyr = np.asarray(self.sfh_times, dtype=float) / 1e9
    T_gyr = np.atleast_1d(T_gyr).ravel()
    n_time = T_gyr.size

    psi = np.atleast_1d(np.asarray(theta["sfh"], dtype=float)).ravel()
    if psi.size not in (n_time, n_time - 1):
        raise AssertionError(
            f"theta['sfh'] has length {psi.size}; expected {n_time} "
            f"(per-node) or {n_time - 1} (per-bin, FastStepBasis)."
        )
    per_bin = (psi.size == n_time - 1)

    # Bin widths in years (physical units for the mass-conservation check).
    # Lookback strictly INCREASING, so dt > 0 via T_yr[1:]-T_yr[:-1].
    T_yr = T_gyr * 1e9
    dt_yr = T_yr[1:] - T_yr[:-1]
    if not np.all(dt_yr > 0):
        raise AssertionError(
            "lookback-time grid must be strictly increasing (today at "
            f"index 0, oldest last); got dt_yr = {dt_yr}"
        )

    # Per-bin SFR -- same branch as _ssp_weights in step mode.  sfh[:-1] is
    # the younger-side node, sfh[1:] the older-side node.
    if per_bin:
        bar_psi = psi
    else:
        bar_psi = 0.5 * (psi[:-1] + psi[1:])

    # Per-node SFR for the linear interpolant.  For per-bin input
    # (non-canonical in linear mode -- _ssp_weights expects per-node),
    # invert the step-mode collapse: interior nodes are the mean of
    # the two adjacent per-bin values; endpoint nodes take the
    # neighbouring bin's value.
    if per_bin:
        psi_nodes = np.empty(n_time, dtype=float)
        psi_nodes[0]    = psi[0]
        psi_nodes[-1]   = psi[-1]
        psi_nodes[1:-1] = 0.5 * (psi[:-1] + psi[1:])
    else:
        psi_nodes = psi

    if units == "Gyr":
        scale, xlabel = 1.0,   "Lookback time [Gyr]"
    elif units == "Myr":
        scale, xlabel = 1e3,   "Lookback time [Myr]"
    elif units == "yr":
        scale, xlabel = 1e9,   "Lookback time [yr]"
    else:
        raise ValueError(
            f"units must be 'Gyr', 'Myr', or 'yr'; got {units!r}"
        )
    T_plot = T_gyr * scale

    if ax is None:
        _, ax = plt.subplots(figsize=(6.0, 4.0))

    style = {"color": "C0", "lw": 1.5}
    style.update(plot_kwargs)
    marker_color = style.get("color", "C0")

    n_bin = n_time - 1

    if self.sfh_interp == "step":
        # One horizontal segment per bin -- exactly the piecewise-constant
        # function the step-mode weight calculator integrates against the
        # SSP Voronoi cells.
        for i in range(n_bin):
            ax.plot([T_plot[i + 1], T_plot[i]],
                    [bar_psi[i],   bar_psi[i]],
                    **style)
        if overlay_nodes:
            T_mid = 0.5 * (T_plot[:-1] + T_plot[1:])
            ax.scatter(T_mid, bar_psi, marker="o",
                       color=marker_color, s=20, zorder=3)
    else:  # "linear"
        for i in range(n_bin):
            ax.plot([T_plot[i + 1], T_plot[i]],
                    [psi_nodes[i + 1], psi_nodes[i]],
                    **style)
        if overlay_nodes:
            ax.scatter(T_plot, psi_nodes, marker="o",
                       color=marker_color, s=20, zorder=3)

    if show_bin_edges:
        for t in T_plot:
            ax.axvline(t, color="grey", lw=0.5, linestyle=":")

    ax.set_xlabel(xlabel)
    ax.set_ylabel(r"$\dot{M}_\star\;[\mathrm{M_\odot\,yr^{-1}}]$")

    # T = 0 (today) sits at the origin on the left; lookback time
    # increases to the right.  No axis inversion.

    total_mass = float(np.sum(bar_psi * dt_yr))
    ax.set_title(
        f"sfh_interp={self.sfh_interp!r}, n_time={n_time}, "
        f"M_total = {total_mass:.3e} M_sun"
    )

    # ------------------------------------------------------------------
    # Mass-conservation contract.
    #
    # m_target  : per-bin mass m2 that _ssp_weights distributes onto
    #             the SSP grid (linear m2 reduces analytically to the
    #             trapezoid between psi nodes; step m2 is bar_psi * dt).
    # m_display : trapezoidal integral of the polyline this method just
    #             drew, segment by segment.  For step we drew a constant
    #             over each bin; for linear we drew the chord between
    #             adjacent nodes.  The integrals must agree by the same
    #             formulas -- a future refactor that changes either side
    #             alone will trip this assertion.
    # ------------------------------------------------------------------
    if self.sfh_interp == "step":
        m_target  = bar_psi * dt_yr
        m_display = bar_psi * dt_yr
    else:
        m_target  = 0.5 * (psi_nodes[:-1] + psi_nodes[1:]) * dt_yr
        m_display = 0.5 * (psi_nodes[:-1] + psi_nodes[1:]) * dt_yr

    rel = np.abs(m_display - m_target) / np.maximum(np.abs(m_target), 1e-30)
    if not np.all(rel < 1e-6):
        raise AssertionError(
            "display_sfh: integrated displayed SFR disagrees with the "
            f"per-bin mass used by _ssp_weights ({self.sfh_interp!r} "
            f"mode); max rel diff = {float(rel.max()):.3e}.  This means "
            "the plot and the weight calculation have drifted out of "
            "sync -- one of them was refactored without the other."
        )

    return ax

calculate_ssp_weights_const_zh

calculate_ssp_weights_const_zh(theta)

Constant-metallicity, piecewise-linear SFH weights.

Thin wrapper over :meth:_ssp_weights; reads theta["sfh"] (shape (n_time,), linear SFR) and theta["Z"] (shape (1,), log10 absolute metallicity on the self.zmet / ssp_lgmet grid — NOT log10 Z/Zsun). Same units as the var-zh variants' theta["zh"].

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_const_zh(self, theta):
    """Constant-metallicity, piecewise-linear SFH weights.

    Thin wrapper over :meth:`_ssp_weights`; reads ``theta["sfh"]``
    (shape ``(n_time,)``, linear SFR) and ``theta["Z"]`` (shape ``(1,)``,
    log10 absolute metallicity on the ``self.zmet`` / ``ssp_lgmet`` grid —
    NOT log10 Z/Zsun).  Same units as the var-zh variants' ``theta["zh"]``.
    """
    return self._ssp_weights(theta, zh_mode="const", sfh_mode="linear")

calculate_ssp_weights_const_zh_step

calculate_ssp_weights_const_zh_step(theta)

Constant-metallicity, piecewise-constant (FastStepBasis-style) SFH weights. Thin wrapper over :meth:_ssp_weights. Reads theta["sfh"] and theta["Z"].

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_const_zh_step(self, theta):
    """Constant-metallicity, piecewise-constant (FastStepBasis-style) SFH
    weights.  Thin wrapper over :meth:`_ssp_weights`.  Reads
    ``theta["sfh"]`` and ``theta["Z"]``.
    """
    return self._ssp_weights(theta, zh_mode="const", sfh_mode="step")

calculate_ssp_weights_var_zh

calculate_ssp_weights_var_zh(theta)

Time-varying-metallicity, piecewise-linear SFH weights.

Thin wrapper over :meth:_ssp_weights; reads theta["sfh"] (shape (n_time,), linear SFR) and theta["zh"] (shape (n_time,), log10 absolute metallicity at each lookback time, on the self.zmet / ssp_lgmet grid — NOT log10 Z/Zsun). Identical units to the const-zh variants' theta["Z"].

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_var_zh(self, theta):
    """Time-varying-metallicity, piecewise-linear SFH weights.

    Thin wrapper over :meth:`_ssp_weights`; reads ``theta["sfh"]``
    (shape ``(n_time,)``, linear SFR) and ``theta["zh"]`` (shape
    ``(n_time,)``, log10 absolute metallicity at each lookback time, on the
    ``self.zmet`` / ``ssp_lgmet`` grid — NOT log10 Z/Zsun).  Identical units
    to the const-zh variants' ``theta["Z"]``.
    """
    return self._ssp_weights(theta, zh_mode="var", sfh_mode="linear")

calculate_ssp_weights_var_zh_step

calculate_ssp_weights_var_zh_step(theta)

Time-varying-metallicity, piecewise-constant (FastStepBasis-style) SFH weights. Thin wrapper over :meth:_ssp_weights. Reads theta["sfh"] and theta["zh"].

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_var_zh_step(self, theta):
    """Time-varying-metallicity, piecewise-constant (FastStepBasis-style)
    SFH weights.  Thin wrapper over :meth:`_ssp_weights`.  Reads
    ``theta["sfh"]`` and ``theta["zh"]``.
    """
    return self._ssp_weights(theta, zh_mode="var", sfh_mode="step")

get_spectrum_dattn_nodem_neb

get_spectrum_dattn_nodem_neb(theta, *, include_lines=None)

Dust attenuation + nebular emission, no dust emission.

include_lines: None (default) -> use self.nebemlineinspec. True / False -> override (csp.predict always passes True).

Source code in ceridwen/csp/csp.py

def get_spectrum_dattn_nodem_neb(self, theta, *, include_lines=None):
    """Dust attenuation + nebular emission, no dust emission.

    ``include_lines``:
      None (default) -> use ``self.nebemlineinspec``.
      True / False  -> override (csp.predict always passes True).
    """
    if include_lines is None:
        include_lines = self.nebemlineinspec
    W = self.calculate_ssp_weights(theta=theta)   # (n_z, n_age)

    neb_all = self._build_neb_array(theta, include_lines=include_lines)

    # Lyman / ionising-photon escape fraction (prospector-equivalent
    # ``frac_obrun``).  When present in theta:
    #   * fraction (1 - f_esc) of ionising photons is absorbed in the
    #     HII region and reprocessed → nebular continuum + lines,
    #     so the nebular grid amplitude scales by (1 - f_esc).
    #   * fraction f_esc of the stellar ionising flux escapes and is
    #     restored to the spectrum (the kill_ion mask zeroed it
    #     out by default, assuming f_esc = 0).
    # Absent → defaults to f_esc = 0 (no kill-ion restoration, no nebular
    # scaling).  The check is a Python-static dict-key lookup, free at
    # trace time.
    if "frac_obrun" in theta:
        f_esc = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        stellar_fluxes  = jnp.where(
            self.kill_ion[None, :, :],
            f_esc * self.flux,                        # restore frac_obrun
            self.flux,
        )
        neb_all = neb_all * (jnp.float32(1.0) - f_esc)
    else:
        stellar_fluxes = jnp.where(self.kill_ion[None, :, :], 0.0, self.flux)
    combined_fluxes = stellar_fluxes + neb_all                      # (n_z, n_age, n_wave) float32

    attn, attn_diffuse = self.attenuate_dust(self.wave, theta)
    # Cast dust curves to float32 — keeps the forward model in single
    # precision until the likelihood.
    M        = self._age_bin_mix
    tau_age  = jnp.einsum("ab,bw->aw", M, attn.astype(jnp.float32))
    attn_age = jnp.exp(-tau_age)

    # FSPS-style OB-runaway dust escape (``add_dust.f90`` L93-94).
    # A fraction ``frac_obrun`` of the young-star flux bypasses the
    # birth-cloud (``attn_age``) attenuation while still passing through
    # the diffuse component below (the second physical effect of FSPS's
    # ``frac_obrun`` knob; the first -- LyC escape + ``Q`` scaling -- is
    # applied above).  Old SSP ages already have ``attn_age = 1``, so the
    # mix is a no-op for them.  When ``frac_obrun`` is absent or 0, this is
    # identically ``attn_age``.
    if "frac_obrun" in theta:
        fo = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        attn_age = (jnp.float32(1.0) - fo) * attn_age + fo

    W_f32 = W.astype(jnp.float32)
    spectrum = jnp.einsum("za,zaw,aw->w", W_f32, combined_fluxes, attn_age)
    spectrum = spectrum * jnp.exp(-attn_diffuse.astype(jnp.float32))

    return spectrum.reshape((-1,))

get_spectrum_dattn_dem_neb

get_spectrum_dattn_dem_neb(theta, *, include_lines=None)

Dust attenuation + nebular emission + dust emission.

Source code in ceridwen/csp/csp.py

def get_spectrum_dattn_dem_neb(self, theta, *, include_lines=None):
    """Dust attenuation + nebular emission + dust emission."""
    if include_lines is None:
        include_lines = self.nebemlineinspec
    W = self.calculate_ssp_weights(theta=theta)   # (n_z, n_age)

    neb_all = self._build_neb_array(theta, include_lines=include_lines)

    # Lyman / ionising-photon escape fraction.  See the long
    # comment on this block in get_spectrum_dattn_nodem_neb above.
    if "frac_obrun" in theta:
        f_esc = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        stellar_fluxes  = jnp.where(
            self.kill_ion[None, :, :],
            f_esc * self.flux,
            self.flux,
        )
        neb_all = neb_all * (jnp.float32(1.0) - f_esc)
    else:
        stellar_fluxes = jnp.where(self.kill_ion[None, :, :], 0.0, self.flux)
    combined_fluxes = stellar_fluxes + neb_all                      # float32

    attn, attn_diffuse = self.attenuate_dust(self.wave, theta)
    M             = self._age_bin_mix
    tau_age       = jnp.einsum("ab,bw->aw", M, attn.astype(jnp.float32))
    attn_age      = jnp.exp(-tau_age)
    diffuse_curve = jnp.exp(-attn_diffuse.astype(jnp.float32))

    # FSPS-style OB-runaway dust escape (``add_dust.f90`` L93-94).
    # See the long comment in ``get_spectrum_dattn_nodem_neb``.
    if "frac_obrun" in theta:
        fo = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        attn_age = (jnp.float32(1.0) - fo) * attn_age + fo

    W_f32 = W.astype(jnp.float32)
    spectrum_dust_free = jnp.einsum("za,zaw->w",       W_f32, combined_fluxes)
    attenuated         = jnp.einsum("za,zaw,aw->w",    W_f32, combined_fluxes, attn_age)
    attenuated         = attenuated * diffuse_curve

    dust_emi_spectrum, _mdust, _tduste = self.dust_emi.compute_dust_emission(
        spec_attn     = attenuated,
        spec_dustfree = spectrum_dust_free,
        spec_lambda   = self.wave,
        diffuse_curve = diffuse_curve,
        duste_qpah    = theta["duste_qpah"],
        duste_umin    = theta["duste_umin"],
        duste_gamma   = theta["duste_gamma"],
    )

    return dust_emi_spectrum

get_spectrum_dattn_nodem_noneb

get_spectrum_dattn_nodem_noneb(
    theta, *, include_lines=None
)

Dust attenuation, no nebular, no dust emission.

include_lines is accepted but ignored: there is no nebular emission to include / exclude when add_neb=False.

Source code in ceridwen/csp/csp.py

def get_spectrum_dattn_nodem_noneb(self, theta, *, include_lines=None):
    """Dust attenuation, no nebular, no dust emission.

    ``include_lines`` is accepted but ignored: there is no nebular
    emission to include / exclude when ``add_neb=False``.
    """
    _ = include_lines
    attn, attn_diffuse = self.attenuate_dust(self.wave, theta)

    M       = self._age_bin_mix
    tau_age = jnp.einsum("ab,bw->aw", M, attn.astype(jnp.float32))
    attn_age= jnp.exp(-tau_age)

    # FSPS-style OB-runaway dust escape (``add_dust.f90`` L93-94).
    # See the long comment in ``get_spectrum_dattn_nodem_neb``.
    if "frac_obrun" in theta:
        fo = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        attn_age = (jnp.float32(1.0) - fo) * attn_age + fo

    weights  = self.calculate_ssp_weights(theta).astype(jnp.float32)
    spectrum = jnp.einsum("za,zaw,aw->w", weights, self.flux, attn_age)
    spectrum *= jnp.exp(-attn_diffuse.astype(jnp.float32))

    return spectrum.reshape((-1,))

get_spectrum_dattn_dem_noneb

get_spectrum_dattn_dem_noneb(theta, *, include_lines=None)

Dust attenuation + dust emission, no nebular. include_lines accepted but ignored.

Source code in ceridwen/csp/csp.py

def get_spectrum_dattn_dem_noneb(self, theta, *, include_lines=None):
    """Dust attenuation + dust emission, no nebular.  ``include_lines`` accepted but ignored."""
    _ = include_lines
    attn, attn_diffuse = self.attenuate_dust(self.wave, theta)

    M             = self._age_bin_mix
    tau_age       = jnp.einsum("ab,bw->aw", M, attn.astype(jnp.float32))
    attn_age      = jnp.exp(-tau_age)
    diffuse_curve = jnp.exp(-attn_diffuse.astype(jnp.float32))

    # FSPS-style OB-runaway dust escape (``add_dust.f90`` L93-94).
    # See the long comment in ``get_spectrum_dattn_nodem_neb``.
    if "frac_obrun" in theta:
        fo = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        attn_age = (jnp.float32(1.0) - fo) * attn_age + fo

    weights           = self.calculate_ssp_weights(theta).astype(jnp.float32)
    spectrum_dust_free= jnp.einsum("za,zaw->w", weights, self.flux)
    attenuated        = jnp.einsum("za,zaw,aw->w", weights, self.flux, attn_age)
    attenuated       *= diffuse_curve

    dust_emi_spectrum, _mdust, _tduste = self.dust_emi.compute_dust_emission(
        spec_attn      = attenuated,
        spec_dustfree  = spectrum_dust_free,
        spec_lambda    = self.wave,
        diffuse_curve  = diffuse_curve,
        duste_qpah     = theta["duste_qpah"],
        duste_umin     = theta["duste_umin"],
        duste_gamma    = theta["duste_gamma"],
    )
    return dust_emi_spectrum

get_spectrum_nodattn_nodem_noneb

get_spectrum_nodattn_nodem_noneb(
    theta, *, include_lines=None
)

Stellar continuum only — no dust, no nebular. include_lines ignored.

Source code in ceridwen/csp/csp.py

def get_spectrum_nodattn_nodem_noneb(self, theta, *, include_lines=None):
    """Stellar continuum only — no dust, no nebular.  ``include_lines`` ignored."""
    _ = include_lines
    weights  = self.calculate_ssp_weights(theta=theta).astype(jnp.float32)
    return jnp.einsum("za,zaw->w", weights, self.flux)

get_spectrum_nodattn_nodem_neb

get_spectrum_nodattn_nodem_neb(
    theta, *, include_lines=None
)

Nebular emission only, no dust.

Source code in ceridwen/csp/csp.py

def get_spectrum_nodattn_nodem_neb(self, theta, *, include_lines=None):
    """Nebular emission only, no dust."""
    if include_lines is None:
        include_lines = self.nebemlineinspec
    W = self.calculate_ssp_weights(theta=theta)   # (n_z, n_age)

    neb_all = self._build_neb_array(theta, include_lines=include_lines)

    # Suppress stellar ionizing part — with the optional
    # frac_obrun escape fraction (see get_spectrum_dattn_nodem_neb).
    if "frac_obrun" in theta:
        f_esc = jnp.ravel(theta["frac_obrun"])[0].astype(jnp.float32)
        stellar_fluxes = jnp.where(
            self.kill_ion[None, :, :],
            f_esc * self.flux,
            self.flux,
        )
        neb_all = neb_all * (jnp.float32(1.0) - f_esc)
    else:
        stellar_fluxes = jnp.where(self.kill_ion[None, :, :], 0.0, self.flux)

    combined_fluxes = stellar_fluxes + neb_all  # float32
    W_f32 = W.astype(jnp.float32)
    return jnp.einsum("za,zaw->w", W_f32, combined_fluxes)

Observations

ceridwen.observation.Photometry

Photometry(
    filters=[], name=None, upper_limit=None, **kwargs
)

Bases: Observation

Broadband photometric observation in AB maggies.

Flux and uncertainty are stored in maggies (linear AB flux units; 1 maggie = 3631 Jy). Filter information is held in a sedpy_jax.observate.FilterSet.

Parameters:

Name	Type	Description	Default
filters	list of str or list of Filter objects	Filters to include. Strings are resolved to .par files in the sedpy_jax filter library.	[]
flux	(array - like, shape(n_filters))	Observed maggies.	required
uncertainty	(array - like, shape(n_filters))	1-sigma uncertainties in maggies.	required
mask	array-like of bool, shape (n_filters,)	True for filters that should be included in the fit.	required

Examples:

>>> phot = Photometry(
...     filters=["sdss_u0", "sdss_g0", "sdss_r0", "sdss_i0", "sdss_z0"],
...     flux=obs_maggies,
...     uncertainty=obs_maggies_unc,
... )
>>> model_maggies = phot.get_maggies(model_wave, model_fnu)
>>> chi2 = phot.chi_sq(model_maggies)

Parameters:

Name	Type	Description	Default
upper_limit	array-like of bool, shape (n_filters,)	Per-band non-detection flags. If True for band i the photometric likelihood treats that band as an upper limit: a chi-squared penalty is applied only when the model flux exceeds the observed value, i.e. .. math:: \chi^2_{\rm UL} = \left[\max(m - d, 0) \,/\, \sigma\right]^2. This mirrors the convention already used in :class:ceridwen.observation.Lines and matches Prospector's recommended treatment of non-detections (the simple flux=0, sigma=1-sigma-limit approximation). None (default) treats every band as a positive detection.	None

Source code in ceridwen/observation/photometry.py

def __init__(self, filters=[], name=None, upper_limit=None, **kwargs):
    """
    Parameters
    ----------
    upper_limit : array-like of bool, shape (n_filters,), optional
        Per-band non-detection flags.  If True for band ``i`` the
        photometric likelihood treats that band as an upper limit:
        a chi-squared penalty is applied *only* when the model flux
        exceeds the observed value, i.e.

        .. math::

            \\chi^2_{\\rm UL}
                = \\left[\\max(m - d, 0) \\,/\\, \\sigma\\right]^2.

        This mirrors the convention already used in
        :class:`ceridwen.observation.Lines` and matches Prospector's
        recommended treatment of non-detections (the simple flux=0,
        sigma=1-sigma-limit approximation).  ``None`` (default) treats
        every band as a positive detection.
    """
    self.set_filters(filters)
    self.upper_limit = (
        None if upper_limit is None
        else jnp.asarray(np.atleast_1d(upper_limit), dtype=bool)
    )
    super().__init__(name=name, **kwargs)

wavelength property

wavelength

Effective wavelengths of the filters [Å], shape (n_filters,).

set_filters

set_filters(filters)

Set the filter list. filters may be a list of filter-name strings or of sedpy_jax Filter objects.

Source code in ceridwen/observation/photometry.py

def set_filters(self, filters):
    """
    Set the filter list.  ``filters`` may be a list of filter-name
    strings or of sedpy_jax ``Filter`` objects.
    """
    if not filters:
        self.filters     = []
        self.filternames = []
        self.filterset   = None
        return

    try:
        self.filternames = [f.name for f in filters]
    except (AttributeError, TypeError):
        self.filternames = list(filters)

    self.filterset = FilterSet(self.filternames)
    self.filters   = list(self.filterset.filters)
    self.wave_eff = [f.wave_effective for f in self.filters]

get_maggies

get_maggies(model_wave, model_fnu)

Project a model spectrum onto the filters and return synthetic maggies.

The model spectrum is expected in F_nu units (e.g. L_sun Hz^{-1} M_sun^{-1} as returned by CSPBasis.get_spectrum). Internally the spectrum is converted to F_lambda before being projected through the AB-normalised FilterSet transmission matrix, so the output has the same relative normalisation as a standard AB photometric integral.

Parameters:

Name	Type	Description	Default
model_wave	(array - like, shape(n_wave))	Wavelength grid [Å].	required
model_fnu	(array - like, shape(n_wave))	Model spectrum in F_nu units (L_sun/Hz/M_sun or erg/s/Hz/cm^2).	required

Returns:

Name	Type	Description
maggies	(ndarray, shape(n_filters))	Synthetic photometry with the same relative normalisation as the input flux.

Notes

The AB normalisation constant in sedpy_jax cancels dimensionally when both the Ceridwen and FSPS spectra are expressed in the same units, making model/data comparisons unit-independent.

Source code in ceridwen/observation/photometry.py

def get_maggies(self, model_wave, model_fnu):
    """
    Project a model spectrum onto the filters and return synthetic maggies.

    The model spectrum is expected in **F_nu units** (e.g. L_sun Hz^{-1}
    M_sun^{-1} as returned by ``CSPBasis.get_spectrum``).  Internally the
    spectrum is converted to F_lambda before being projected through the
    AB-normalised FilterSet transmission matrix, so the output has the
    same relative normalisation as a standard AB photometric integral.

    Parameters
    ----------
    model_wave : array-like, shape (n_wave,)
        Wavelength grid [Å].
    model_fnu : array-like, shape (n_wave,)
        Model spectrum in F_nu units (L_sun/Hz/M_sun or erg/s/Hz/cm^2).

    Returns
    -------
    maggies : jnp.ndarray, shape (n_filters,)
        Synthetic photometry with the same relative normalisation as the
        input flux.

    Notes
    -----
    The AB normalisation constant in sedpy_jax cancels dimensionally when
    both the Ceridwen and FSPS spectra are expressed in the same units,
    making model/data comparisons unit-independent.
    """
    if self.filterset is None:
        raise ValueError("No FilterSet configured; call set_filters() first.")
    _c         = jnp.array(2.998e18)        # Å/s
    wave       = jnp.asarray(model_wave,  dtype=float)
    flux_flam  = jnp.asarray(model_fnu,   dtype=float) * _c / wave**2
    return self.filterset.get_sed_maggies(flux_flam, sourcewave=wave)

setup_for_model

setup_for_model(wave_model, zred=0.0)

Precompute a (n_filters, n_wave) projection matrix _T so that predict reduces to a single GEMV: maggies = _T @ F_nu.

The matrix folds together three operations that FilterSet.get_sed_maggies does per call:

1. F_nu -> F_lambda conversion: F_lam = F_nu * c / lam^2
2. Interpolation from the model wavelength grid onto the FilterSet's internal grid (interp_source)
3. Dot product with the precomputed FilterSet.trans matrix

By composing these into a single static matrix _T of shape (n_filters, n_wave_model), all three steps collapse into one GEMV at predict time.

Must be called once before predict (and before JIT compilation). SedModel.__init__ calls this automatically.

Parameters:

Name	Type	Description	Default
wave_model	array - like	Rest-frame wavelength grid of the model spectrum [Å].	required
zred	float	Fixed redshift at which to precompute the filter projection. Defaults to 0 (rest-frame). For a non-zero zred, the grid is effectively taken in the observed frame (wave_effective = (1 + zred) * wave_model) before filter integration — this keeps the predict-time GEMV path unchanged and preserves sampling-hot-path speed. Combine with :func:ceridwen.cosmology.flux_factor_maggies (applied inside CSPBasis.predict when theta['zred'] is present) to get correctly calibrated observed-frame maggies.	0.0

Source code in ceridwen/observation/photometry.py

def setup_for_model(self, wave_model, zred: float = 0.0):
    """
    Precompute a (n_filters, n_wave) projection matrix ``_T`` so that
    ``predict`` reduces to a single GEMV: ``maggies = _T @ F_nu``.

    The matrix folds together three operations that
    ``FilterSet.get_sed_maggies`` does per call:

    1. F_nu -> F_lambda conversion: ``F_lam = F_nu * c / lam^2``
    2. Interpolation from the model wavelength grid onto the
       FilterSet's internal grid (``interp_source``)
    3. Dot product with the precomputed ``FilterSet.trans`` matrix

    By composing these into a single static matrix ``_T`` of shape
    ``(n_filters, n_wave_model)``, all three steps collapse into one
    GEMV at predict time.

    Must be called once before ``predict`` (and before JIT compilation).
    ``SedModel.__init__`` calls this automatically.

    Parameters
    ----------
    wave_model : array-like
        Rest-frame wavelength grid of the model spectrum [Å].
    zred : float, optional
        Fixed redshift at which to precompute the filter projection.
        Defaults to 0 (rest-frame).  For a non-zero ``zred``, the grid
        is effectively taken in the observed frame
        (``wave_effective = (1 + zred) * wave_model``) before filter
        integration — this keeps the predict-time GEMV path unchanged
        and preserves sampling-hot-path speed.  Combine with
        :func:`ceridwen.cosmology.flux_factor_maggies` (applied inside
        ``CSPBasis.predict`` when ``theta['zred']`` is present) to get
        correctly calibrated observed-frame maggies.
    """
    wm_rest = np.asarray(wave_model, dtype=np.float64)   # (n_wave,)
    # Effective ("observed-frame") grid used for the maggies integral.
    # At zred = 0 this equals wm_rest.
    opz = 1.0 + float(zred)
    wm = opz * wm_rest
    n_wave = len(wm)
    _c = 2.998e18  # speed of light [A/s]

    # F_nu -> F_lambda factor per model wavelength bin
    fnu_to_flam = _c / wm**2                         # (n_wave,)

    # Build interpolation matrix H: (n_lam_filter, n_wave_model)
    # such that  F_lam_filtergrid = H @ F_lam_modelgrid
    # This is the linear interpolation that interp_source does per call.
    lam_filt = np.asarray(self.filterset.lam, dtype=np.float64)  # (n_lam,)
    n_lam = len(lam_filt)

    # Construct sparse interpolation weights
    # For each point in lam_filt, find the bracketing indices in wm
    # and the interpolation fraction.
    idx = np.searchsorted(wm, lam_filt, side="right") - 1
    idx = np.clip(idx, 0, n_wave - 2)
    frac = (lam_filt - wm[idx]) / (wm[idx + 1] - wm[idx])
    frac = np.clip(frac, 0.0, 1.0)

    # Zero out entries outside the model wavelength range
    outside = (lam_filt < wm[0]) | (lam_filt > wm[-1])
    frac[outside] = 0.0

    # Build H as a dense matrix (n_lam, n_wave)
    H = np.zeros((n_lam, n_wave), dtype=np.float64)
    for j in range(n_lam):
        if outside[j]:
            continue
        H[j, idx[j]]     = (1.0 - frac[j])
        H[j, idx[j] + 1] = frac[j]

    # FilterSet.trans is (n_filters, n_lam): already includes
    # R * lam * dlam / ab_zero_counts normalisation.
    trans = np.asarray(self.filterset.trans, dtype=np.float64)  # (n_filt, n_lam)

    # Compose: _T = trans @ H @ diag(fnu_to_flam)
    #   maggies = trans @ (H @ (F_nu * fnu_to_flam))
    #           = (trans @ H @ diag(fnu_to_flam)) @ F_nu
    #           = _T @ F_nu
    TH = trans @ H                                   # (n_filt, n_wave)
    T  = TH * fnu_to_flam[None, :]                   # (n_filt, n_wave)

    self._T = jnp.array(T.astype(np.float32))
    self._has_precomputed_T = True

predict

predict(spectrum, wave_model)

Project a model F_nu spectrum onto the filters.

If setup_for_model has been called, this is a single GEMV (_T @ spectrum). Otherwise falls back to get_maggies.

Parameters:

Name	Type	Description	Default
spectrum	(Array, shape(n_wave))	Model spectrum in F_nu units.	required
wave_model	(Array, shape(n_wave))	Model wavelength grid [Å].	required

Returns:

Type	Description
(Array, shape(n_filters))	Synthetic AB maggies.

Source code in ceridwen/observation/photometry.py

def predict(self, spectrum, wave_model):
    """
    Project a model F_nu spectrum onto the filters.

    If ``setup_for_model`` has been called, this is a single GEMV
    (``_T @ spectrum``).  Otherwise falls back to ``get_maggies``.

    Parameters
    ----------
    spectrum : jax.Array, shape (n_wave,)
        Model spectrum in F_nu units.
    wave_model : jax.Array, shape (n_wave,)
        Model wavelength grid [Å].

    Returns
    -------
    jax.Array, shape (n_filters,)
        Synthetic AB maggies.
    """
    if getattr(self, "_has_precomputed_T", False):
        return self._T @ spectrum
    return self.get_maggies(wave_model, spectrum)

predict_at_redshift

predict_at_redshift(spectrum_fnu_observed, wave_rest, zred)

Project an observer-frame F_nu spectrum through the filters when the redshift is a traced (sampled) JAX scalar.

This is the free-redshift counterpart of :meth:predict. The GEMV fast path baked by :meth:setup_for_model assumes a single Python-scalar zred was known at trace time and bakes the observed-frame wavelength grid into the projection matrix _T; that path cannot be used for sampling. Here the observed-frame wavelength grid is reconstructed per-sample as wave_obs = (1 + zred) * wave_rest and the spectrum is projected via :meth:FilterSet.get_sed_maggies with the traced sourcewave.

Pre-condition: spectrum_fnu_observed is the observer-frame F_nu, i.e. CSPBasis.predict has already multiplied by flux_factor_maggies(zred) and (optionally) by the IGM transmission. This method only handles the wavelength-grid bookkeeping and the filter integral.

Parameters:

Name	Type	Description	Default
spectrum_fnu_observed	(Array, shape(n_wave))	Observer-frame F_nu on the rest-frame model wavelength grid (the standard ceridwen output of get_spectrum + mass + flux-factor + IGM).	required
wave_rest	(Array, shape(n_wave))	Rest-frame model wavelength grid [Å] (typically csp.wave).	required
zred	(Array, scalar)	Sampled redshift. May be a traced array; the entire path below is JIT-compatible and differentiable in zred (sedpy_jax's interp_source uses jnp.interp which has a defined gradient w.r.t. its xp argument).	required

Returns:

Type	Description
(Array, shape(n_filters))	Synthetic AB maggies in observer frame.

Notes

• Cost is one filter interpolation + one trans-matrix dot per sample, vs the single GEMV of the fixed-z path. For a 14-d NUTS / 4000-particle NS / 20 000-step SVI run on a 40 GB A100 this is ~10-20x slower than the GEMV but still saturates the GPU.
• Numerically equivalent to ``setup_for_model(wave_rest, zred=z)
• predict(spectrum, wave_rest)`` evaluated at the same z, to float32 precision.
• Works regardless of whether setup_for_model has been called. When both paths are wired (e.g. for compare-mode plots), prefer the GEMV path for any fixed-z observation and this method for any free-z observation.

Source code in ceridwen/observation/photometry.py

def predict_at_redshift(self, spectrum_fnu_observed, wave_rest, zred):
    """
    Project an observer-frame F_nu spectrum through the filters when
    the redshift is a *traced* (sampled) JAX scalar.

    This is the free-redshift counterpart of :meth:`predict`.  The
    GEMV fast path baked by :meth:`setup_for_model` assumes a single
    Python-scalar ``zred`` was known at trace time and bakes the
    observed-frame wavelength grid into the projection matrix
    ``_T``; that path cannot be used for sampling.  Here the
    observed-frame wavelength grid is reconstructed per-sample as
    ``wave_obs = (1 + zred) * wave_rest`` and the spectrum is
    projected via :meth:`FilterSet.get_sed_maggies` with the
    traced ``sourcewave``.

    Pre-condition: ``spectrum_fnu_observed`` is the observer-frame
    F_nu, i.e. ``CSPBasis.predict`` has already multiplied by
    ``flux_factor_maggies(zred)`` and (optionally) by the IGM
    transmission.  This method only handles the wavelength-grid
    bookkeeping and the filter integral.

    Parameters
    ----------
    spectrum_fnu_observed : jax.Array, shape (n_wave,)
        Observer-frame F_nu on the rest-frame model wavelength grid
        (the standard ceridwen output of get_spectrum + mass +
        flux-factor + IGM).
    wave_rest : jax.Array, shape (n_wave,)
        Rest-frame model wavelength grid [Å] (typically ``csp.wave``).
    zred : jax.Array, scalar
        Sampled redshift.  May be a traced array; the entire path
        below is JIT-compatible and differentiable in ``zred``
        (sedpy_jax's ``interp_source`` uses ``jnp.interp`` which
        has a defined gradient w.r.t. its xp argument).

    Returns
    -------
    jax.Array, shape (n_filters,)
        Synthetic AB maggies in observer frame.

    Notes
    -----
    - Cost is one filter interpolation + one trans-matrix dot per
      sample, vs the single GEMV of the fixed-z path.  For a 14-d
      NUTS / 4000-particle NS / 20 000-step SVI run on a 40 GB A100
      this is ~10-20x slower than the GEMV but still saturates
      the GPU.
    - Numerically equivalent to ``setup_for_model(wave_rest, zred=z)
      + predict(spectrum, wave_rest)`` evaluated at the same z, to
      float32 precision.
    - Works regardless of whether ``setup_for_model`` has been
      called.  When both paths are wired (e.g. for compare-mode
      plots), prefer the GEMV path for any fixed-z observation
      and this method for any free-z observation.
    """
    if self.filterset is None:
        raise ValueError("No FilterSet configured; call set_filters() first.")
    _c = jnp.array(2.998e18, dtype=spectrum_fnu_observed.dtype)
    wave_obs = (jnp.float32(1.0) + zred.astype(spectrum_fnu_observed.dtype)) \
        * jnp.asarray(wave_rest, dtype=spectrum_fnu_observed.dtype)
    flux_flam = jnp.asarray(spectrum_fnu_observed,
                             dtype=spectrum_fnu_observed.dtype) \
        * _c / (wave_obs * wave_obs)
    return self.filterset.get_sed_maggies(flux_flam, sourcewave=wave_obs)

chi_sq

chi_sq(model_maggies)

Chi-squared contribution from this photometric observation.

For bands flagged as upper limits (self.upper_limit[i] = True), the contribution is one-sided: a penalty is applied only when the model flux exceeds the observed upper-limit value. Matches the convention used in :class:Lines and Prospector's recommended treatment of non-detections.

Parameters:

Name	Type	Description	Default
model_maggies	(array - like, shape(n_filters))	Synthetic photometry on the same filter set.	required

Returns:

Name	Type	Description
chi2	float

Source code in ceridwen/observation/photometry.py

def chi_sq(self, model_maggies):
    """
    Chi-squared contribution from this photometric observation.

    For bands flagged as upper limits (``self.upper_limit[i] = True``),
    the contribution is one-sided: a penalty is applied only when the
    model flux exceeds the observed upper-limit value.  Matches the
    convention used in :class:`Lines` and Prospector's recommended
    treatment of non-detections.

    Parameters
    ----------
    model_maggies : array-like, shape (n_filters,)
        Synthetic photometry on the same filter set.

    Returns
    -------
    chi2 : float
    """
    mf    = jnp.asarray(model_maggies, dtype=float)
    resid = (self.flux - mf) / self.uncertainty       # (data - model) / sigma

    if self.upper_limit is not None:
        # For upper-limit bands: only penalise when model > data,
        # i.e. when resid < 0.  Identical to the Lines convention.
        resid_sq = jnp.where(
            self.upper_limit,
            jnp.where(resid < 0.0, resid ** 2, 0.0),
            resid ** 2,
        )
    else:
        resid_sq = resid ** 2

    return float(jnp.sum(jnp.where(self.mask, resid_sq, 0.0)))

residuals

residuals(model_maggies)

Per-filter (data − model) / sigma. Masked filters are set to NaN.

For bands flagged as upper limits, residuals are clipped to 0 when the model is safely below the limit (positive residual), so the returned vector matches what enters chi_sq band-by-band.

Returns:

Name	Type	Description
res	(ndarray, shape(n_filters))

Source code in ceridwen/observation/photometry.py

def residuals(self, model_maggies):
    """
    Per-filter (data − model) / sigma.  Masked filters are set to NaN.

    For bands flagged as upper limits, residuals are clipped to 0 when
    the model is safely below the limit (positive residual), so the
    returned vector matches what enters ``chi_sq`` band-by-band.

    Returns
    -------
    res : jnp.ndarray, shape (n_filters,)
    """
    res = (self.flux - jnp.asarray(model_maggies, dtype=float)) / self.uncertainty

    if self.upper_limit is not None:
        res = jnp.where(
            self.upper_limit & (res >= 0.0),
            0.0,
            res,
        )
    return jnp.where(self.mask, res, jnp.nan)

ceridwen.observation.Spectrum

Spectrum(
    wavelength=None,
    flux=None,
    uncertainty=None,
    mask=slice(None),
    noise=None,
    name=None,
    resolution=None,
    calibration=None,
    logify_spectrum=False,
    smoothtype=None,
    inres=0.0,
    sky=None,
    noise_floor=0.0,
    sigma_losvd=None,
    fit_sigma_smooth=False,
    **kwargs
)

Bases: Observation

Spectroscopic observation.

Stores a densely-sampled spectrum together with optional resolution and multiplicative flux-calibration arrays. Provides helpers for masking spectral regions, computing chi-squared residuals, and projecting the spectrum onto broadband filters.

Parameters:

Name	Type	Description	Default
wavelength	(array - like, shape(n_pix))	Wavelength grid [Å], vacuum, observed frame (as delivered by the instrument). setup_for_model(wave_model, zred=...) redshifts the rest-frame model grid by (1 + zred) and interpolates onto these pixels; at zred = 0 observed and rest frame coincide.	None
flux	(array - like, shape(n_pix))	Observed flux. Units must be consistent with uncertainty and any model spectra passed to chi_sq / residuals. Ceridwen model spectra are in L_sun Hz^{-1} M_sun^{-1}.	None
uncertainty	(array - like, shape(n_pix))	1-sigma uncertainty, same units as flux.	None
mask	array-like of bool or slice, shape (n_pix,)	True for pixels that are used (not masked).	slice(None)
resolution	float or array - like	Instrumental smoothing width. Interpretation depends on smoothtype: smoothtype=None — stored but never applied. "vel" — scalar σ_v [km/s]. "R" — scalar resolving power R = λ/σ_λ. "lambda" — scalar σ_λ [Å]. "lsf" — 1-D array of σ(λ) [Å] at each observed pixel.	None
smoothtype	(vel, R, 'lambda', lsf)	Type of instrumental broadening to apply in predict. See the __init__ docstring and setup_for_model for full details. Default None (no smoothing, backward-compatible).	"vel"
inres	float	Intrinsic resolution of the model library, subtracted in quadrature before applying the target smoothing. Units match smoothtype: km/s for vel/R, Å for lambda/lsf. Default 0.0.	0.0
calibration	(array - like, shape(n_pix))	Multiplicative flux-calibration vector (model × calibration ≈ data). When set, it is applied as a per-pixel multiplicative correction to the model inside chi_sq, residuals, and log_likelihood (and inside the noise-floor inflation in _effective_sigma). For the alternative analytic-marginalisation path, see fit_polynomial_calibration.	None
logify_spectrum	bool	If True, chi_sq and residuals operate in log-flux space (residuals are Δln f / σ_ln f).	False
sky	(array - like, shape(n_pix))	Observed sky background spectrum, same units and pixel grid as flux. When provided, the sky is subtracted from the data before computing chi-squared residuals: residual = (flux - sky - model) / sigma. The sky vector is not propagated through predict; it enters only in chi_sq, residuals, and log_likelihood.	None
noise_floor	float	Fractional uncertainty floor applied to the model flux. The effective per-pixel sigma used in chi-squared becomes: .. math:: \sigma_{\rm eff}^2 = \sigma^2 + (f_{\rm floor}\,|m|)^2 where :math:f_{\rm floor} is noise_floor and :math:m is the model flux. Prevents chi-squared from being dominated by pixels where the photon-noise uncertainty is smaller than calibration systematics. Default 0.0 (disabled).	0.0
sigma_losvd	float or None	Galaxy line-of-sight velocity dispersion [km/s]. When set, an additional velocity-broadening step is applied to the model spectrum before any instrumental smoothing specified by smoothtype. Useful when sigma_losvd is a free parameter of the SED fit. Requires a call to setup_for_model whenever this value changes. Default None (disabled).	None
noise	GaussianProcess or None	If a GaussianProcess instance is provided, its log-likelihood contribution is added to the standard Gaussian log-likelihood in log_likelihood(), accounting for correlated residual structure. chi_sq is not affected by this attribute (it remains a simple diagonal chi-squared).	None

Examples:

>>> spec = Spectrum(
...     wavelength=wave_aa,
...     flux=obs_fnu,
...     uncertainty=obs_fnu_unc,
...     resolution=100.0,
...     smoothtype="vel",   # 100 km/s instrumental broadening
...     noise_floor=0.01,   # 1% calibration floor
...     sigma_losvd=150.0,  # 150 km/s galaxy velocity dispersion
... )
>>> spec.setup_for_model(model_wave)
>>> predicted = spec.predict(model_spectrum, model_wave)
>>> spec.mask_lines([6563., 4861.], dv=500.)     # mask Hα, Hβ
>>> chi2 = spec.chi_sq(model_fnu)
>>> coeffs, cal_model = spec.fit_polynomial_calibration(predicted, order=4)
>>> phot = spec.synthetic_photometry(filterset)

Parameters:

Name	Type	Description	Default
smoothtype	(vel, R, 'lambda', lsf)	Which instrumental smoothing kernel to apply in predict: "vel" Constant velocity dispersion. resolution is σ_v [km/s]. Uses a log-λ FFT so the kernel is shift-invariant in velocity. "R" Constant spectral resolving power R = λ/Δλ = c/σ_v. resolution is the scalar R value; converted internally to σ_v = c/R [km/s]. "lambda" Constant wavelength dispersion. resolution is σ_λ [Å]. Uses a linear-λ FFT. "lsf" Wavelength-dependent line-spread function. resolution must be a 1-D array of σ(λ) [Å] evaluated at the observed pixel wavelengths (self.wavelength). The kernel is interpolated to the model wavelength grid inside setup_for_model. None (default) No smoothing applied; predict performs pure linear interpolation (_H @ spectrum). resolution is stored but unused in this mode.	"vel"
inres	float	Intrinsic (library) resolution of the input model spectrum, subtracted in quadrature before applying the target smoothing. Units must match smoothtype: km/s for "vel"/"R", Å for "lambda"/"lsf". Default 0.0.	0.0

Source code in ceridwen/observation/spectrum.py

def __init__(
    self,
    wavelength   = None,
    flux         = None,
    uncertainty  = None,
    mask         = slice(None),
    noise        = None,
    name         = None,
    resolution   = None,
    calibration  = None,
    logify_spectrum = False,
    smoothtype   = None,
    inres        = 0.0,
    sky          = None,
    noise_floor  = 0.0,
    sigma_losvd  = None,
    fit_sigma_smooth = False,
    **kwargs,
):
    """
    Parameters
    ----------
    smoothtype : {"vel", "R", "lambda", "lsf"} or None
        Which instrumental smoothing kernel to apply in ``predict``:

        ``"vel"``
            Constant velocity dispersion.  ``resolution`` is σ_v [km/s].
            Uses a log-λ FFT so the kernel is shift-invariant in velocity.
        ``"R"``
            Constant spectral resolving power R = λ/Δλ = c/σ_v.
            ``resolution`` is the scalar R value; converted internally to
            σ_v = c/R [km/s].
        ``"lambda"``
            Constant wavelength dispersion.  ``resolution`` is σ_λ [Å].
            Uses a linear-λ FFT.
        ``"lsf"``
            Wavelength-dependent line-spread function.  ``resolution``
            must be a 1-D array of σ(λ) [Å] evaluated at the *observed*
            pixel wavelengths (``self.wavelength``).  The kernel is
            interpolated to the model wavelength grid inside
            ``setup_for_model``.
        ``None`` (default)
            No smoothing applied; ``predict`` performs pure linear
            interpolation (``_H @ spectrum``).  ``resolution`` is stored
            but unused in this mode.

    inres : float, optional
        Intrinsic (library) resolution of the input model spectrum,
        subtracted in quadrature before applying the target smoothing.
        Units must match ``smoothtype``: km/s for "vel"/"R", Å for
        "lambda"/"lsf".  Default 0.0.
    """
    # Store wavelength via the property setter so subclasses can override.
    self._wavelength    = (
        None if wavelength is None
        else jnp.asarray(wavelength, dtype=float)
    )
    self.resolution     = resolution
    self.calibration    = (
        None if calibration is None
        else jnp.asarray(calibration, dtype=float)
    )
    self.logify_spectrum = logify_spectrum
    self.smoothtype      = smoothtype
    self.inres           = float(inres)
    self.sky             = (None if sky is None
                            else jnp.asarray(sky, dtype=float))
    self.noise_floor     = float(noise_floor)
    # ── Galaxy LOSVD (sigma_smooth in Prospector convention) ──────────
    # When ``fit_sigma_smooth=False`` (default), ``sigma_losvd`` is
    # baked into ``_predict_fn`` at ``setup_for_model`` time as a
    # Python float.  When ``fit_sigma_smooth=True``, the closure
    # instead accepts a runtime ``sigma_smooth`` jnp scalar (km/s) and
    # the caller (CSPBasis.predict) passes ``theta["sigma_smooth"]``
    # through.  In that fittable mode ``sigma_losvd`` is only used as
    # the warmup / smoother-init value, so we default to 200 km/s
    # (Prospector ``TemplateLibrary["spectral_smoothing"]`` init) when
    # the user does not supply one.
    self.fit_sigma_smooth = bool(fit_sigma_smooth)
    if self.fit_sigma_smooth and sigma_losvd is None:
        sigma_losvd = 200.0
    self.sigma_losvd     = (None if sigma_losvd is None
                            else float(sigma_losvd))

    super().__init__(
        flux        = flux,
        uncertainty = uncertainty,
        mask        = mask,
        noise       = noise,
        name        = name,
        **kwargs,
    )

setup_for_model

setup_for_model(wave_model, zred=0.0)

Precompute projection matrices and/or smoothing kernels, then build _predict_fn — the single callable used by predict.

Must be called once (Python-level, outside JIT) after the model wavelength grid is known. SedModel.__init__ calls this automatically.

Two behaviours depending on self.smoothtype:

No smoothing (smoothtype=None) Builds the dense (n_pix, n_wave) linear-interpolation matrix _H and sets _predict_fn(spec) = _H @ spec.

With instrumental smoothing (smoothtype in {"vel", "R", "lambda", "lsf"}) Uses a factory function from sedpy_jax.smoothing to precompute all FFT grid transforms. The returned closure is fully JAX-JIT-compilable with respect to the spectrum. _predict_fn(spec) applies smoothing and interpolation to the observed pixel grid in one call. _H is also built (used only by the no-smoothing fast path).

Parameters:

Name	Type	Description	Default
wave_model	(array - like, shape(n_wave))	Rest-frame model wavelength grid [Å], strictly increasing.	required
zred	float	Fixed redshift at which to precompute the spectral projection. Default 0 (rest-frame). For zred > 0 the interpolation maps from the observed-frame grid (1 + zred) * wave_model onto self.wavelength (which is interpreted as observed-frame pixel wavelengths), preserving the predict-time GEMV fast path.	0.0

Source code in ceridwen/observation/spectrum.py

def setup_for_model(self, wave_model, zred: float = 0.0):
    """
    Precompute projection matrices and/or smoothing kernels, then build
    ``_predict_fn`` — the single callable used by ``predict``.

    Must be called once (Python-level, outside JIT) after the model
    wavelength grid is known.  ``SedModel.__init__`` calls this
    automatically.

    Two behaviours depending on ``self.smoothtype``:

    **No smoothing** (``smoothtype=None``)
        Builds the dense (n_pix, n_wave) linear-interpolation matrix
        ``_H`` and sets ``_predict_fn(spec) = _H @ spec``.

    **With instrumental smoothing** (``smoothtype`` in
    ``{"vel", "R", "lambda", "lsf"}``)
        Uses a factory function from ``sedpy_jax.smoothing`` to
        precompute all FFT grid transforms.  The returned closure is
        fully JAX-JIT-compilable with respect to the spectrum.
        ``_predict_fn(spec)`` applies smoothing *and* interpolation to
        the observed pixel grid in one call.  ``_H`` is also built (used
        only by the no-smoothing fast path).

    Parameters
    ----------
    wave_model : array-like, shape (n_wave,)
        Rest-frame model wavelength grid [Å], strictly increasing.
    zred : float, optional
        Fixed redshift at which to precompute the spectral projection.
        Default 0 (rest-frame).  For ``zred > 0`` the interpolation
        maps from the observed-frame grid
        ``(1 + zred) * wave_model`` onto ``self.wavelength`` (which is
        interpreted as observed-frame pixel wavelengths), preserving
        the predict-time GEMV fast path.
    """
    if self._wavelength is None:
        raise ValueError(
            "Spectrum.setup_for_model() needs the observed pixel "
            "wavelength grid, but this Spectrum has wavelength=None. "
            "Pass wavelength= (Å, vacuum, observed frame) at "
            "construction — flux/uncertainty may be added later (e.g. a "
            "predictive container for mock generation)."
        )
    wm_rest = np.asarray(wave_model, dtype=np.float64)
    opz = 1.0 + float(zred)
    wm = opz * wm_rest
    wo = np.asarray(self._wavelength, dtype=np.float64)
    n_wave = len(wm)
    n_pix  = len(wo)

    # ── Always build the dense interpolation matrix _H ────────────────
    # (used by the no-smoothing fast path)
    j_hi = np.searchsorted(wm, wo, side='right')
    j_hi = np.clip(j_hi, 1, n_wave - 1)
    j_lo = j_hi - 1

    dw    = wm[j_hi] - wm[j_lo]
    alpha = np.where(dw > 0, (wo - wm[j_lo]) / dw, 0.0)
    alpha = np.clip(alpha, 0.0, 1.0)

    H = np.zeros((n_pix, n_wave), dtype=np.float32)
    rows = np.arange(n_pix)
    H[rows, j_lo] += (1.0 - alpha).astype(np.float32)
    H[rows, j_hi] += alpha.astype(np.float32)
    self._H = jnp.array(H)

    # ── Build _predict_fn ──────────────────────────────────────────────
    st         = self.smoothtype
    has_instr  = (st is not None) and (self.resolution is not None)
    has_losvd  = self.sigma_losvd is not None
    # When the galaxy LOSVD is a free parameter, the closure has a
    # runtime ``sigma_smooth`` argument; ``predict`` switches its
    # signature accordingly.  The constructor injects a 200 km/s
    # default when ``fit_sigma_smooth=True`` and ``sigma_losvd is
    # None``, so a valid smoother can always be built here.
    fit_lo     = self.fit_sigma_smooth and has_losvd

    if not has_instr and not has_losvd:
        # No smoothing: pure interpolation via _H.
        _H = self._H
        self._predict_fn = lambda spec: _H @ spec

    else:
        # ── Trim model grid to observed wavelength range ────────────
        # The factory functions build a uniform log/linear FFT grid
        # over the full span of ``wave_model``.  If that grid covers
        # 100–25000 Å but the observation covers only 3700–7200 Å,
        # the FFT pixel is ~800 km/s wide.  Nebular emission lines
        # have σ ≈ 2 Å ≈ 0.09 log-pixels and are completely aliased
        # — their flux is redistributed away and the line disappears.
        #
        # Fix: restrict the model grid to the observed spectral window
        # (plus a generous buffer for smoothing-kernel wings) so the
        # FFT grid pixel is small enough to resolve the line profiles.
        _buf     = max(500.0, 0.15 * float(wo.max() - wo.min()))
        _trim    = ((wm >= float(wo.min()) - _buf) &
                    (wm <= float(wo.max()) + _buf))
        _wm_trim = wm[_trim]

        # Integer index array for JIT-safe gather inside the closure:
        # spec[_idx] selects only the trimmed wavelength range.  The
        # shape of _idx is statically known at Python level, so XLA
        # can lower this to a static gather with no shape ambiguity.
        _idx = jnp.array(np.where(_trim)[0])

        # ── Optional LOSVD pre-smoothing stage ─────────────────────
        # If sigma_losvd is set, a velocity-broadening step (galaxy
        # line-of-sight velocity dispersion) is applied to the model
        # spectrum BEFORE instrumental smoothing.  The LOSVD smoother
        # maps _wm_trim → _wm_trim when chained with an instrumental
        # smoother, or _wm_trim → wo when it is the only smoothing
        # step.
        if has_losvd:
            _sv_losvd = float(self.sigma_losvd)
            if has_instr:
                # Output stays on trimmed model grid for chaining with
                # the instrumental smoother below.
                _losvd_sm = make_vel_smoother(_wm_trim, _wm_trim, inres=0.0)
                _sv       = _sv_losvd   # capture scalar before lambda
                def _apply_losvd(spec_trim, _s=_losvd_sm, _v=_sv):
                    return _s(spec_trim, _v)
                # Runtime-sigma variant: takes (spec_trim, sigma_v)
                # and is differentiable w.r.t. sigma_v (sedpy_jax's
                # make_vel_smoother already supports this).
                def _apply_losvd_rt(spec_trim, sigma_v, _s=_losvd_sm):
                    return _s(spec_trim, sigma_v)
            else:
                # LOSVD only: output goes directly to observed grid.
                _losvd_sm = make_vel_smoother(_wm_trim, wo, inres=0.0)
                _sv       = _sv_losvd
                def _apply_losvd(spec_trim, _s=_losvd_sm, _v=_sv):
                    return _s(spec_trim, _v)
                def _apply_losvd_rt(spec_trim, sigma_v, _s=_losvd_sm):
                    return _s(spec_trim, sigma_v)

        if has_instr:
            if st == "vel":
                # Constant velocity dispersion σ_v [km/s].
                sigma_v   = float(self.resolution)
                _instr_sm = make_vel_smoother(_wm_trim, wo, inres=self.inres)
                if fit_lo:
                    _Lrt = _apply_losvd_rt
                    self._predict_fn = (
                        lambda spec, sigma_lo, _sm=_instr_sm, _sv=sigma_v,
                               _L=_Lrt:
                            _sm(_L(spec[_idx], sigma_lo), _sv)
                    )
                elif has_losvd:
                    _L = _apply_losvd
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sv=sigma_v, _L=_L:
                            _sm(_L(spec[_idx]), _sv)
                    )
                else:
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sv=sigma_v:
                            _sm(spec[_idx], _sv)
                    )

            elif st == "R":
                # Constant resolving power R = λ/σ_λ = c/σ_v → σ_v = c/R.
                sigma_v   = float(_CKMS / self.resolution)
                _instr_sm = make_vel_smoother(_wm_trim, wo, inres=self.inres)
                if fit_lo:
                    _Lrt = _apply_losvd_rt
                    self._predict_fn = (
                        lambda spec, sigma_lo, _sm=_instr_sm, _sv=sigma_v,
                               _L=_Lrt:
                            _sm(_L(spec[_idx], sigma_lo), _sv)
                    )
                elif has_losvd:
                    _L = _apply_losvd
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sv=sigma_v, _L=_L:
                            _sm(_L(spec[_idx]), _sv)
                    )
                else:
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sv=sigma_v:
                            _sm(spec[_idx], _sv)
                    )

            elif st == "lambda":
                # Constant wavelength dispersion σ_λ [Å].
                sigma_l   = float(self.resolution)
                _instr_sm = make_wave_smoother(_wm_trim, wo, inres=self.inres)
                if fit_lo:
                    _Lrt = _apply_losvd_rt
                    self._predict_fn = (
                        lambda spec, sigma_lo, _sm=_instr_sm, _sl=sigma_l,
                               _L=_Lrt:
                            _sm(_L(spec[_idx], sigma_lo), _sl)
                    )
                elif has_losvd:
                    _L = _apply_losvd
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sl=sigma_l, _L=_L:
                            _sm(_L(spec[_idx]), _sl)
                    )
                else:
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _sl=sigma_l:
                            _sm(spec[_idx], _sl)
                    )

            elif st == "lsf":
                # Wavelength-dependent LSF: resolution is σ(λ) [Å] at
                # the *observed* pixel grid.  Interpolate to trimmed grid.
                res_obs        = np.asarray(self.resolution, dtype=np.float64)
                sigma_lsf_trim = np.interp(_wm_trim, wo, res_obs)
                _instr_sm      = make_lsf_smoother(_wm_trim, sigma_lsf_trim, wo,
                                                   inres=self.inres)
                if fit_lo:
                    _Lrt = _apply_losvd_rt
                    self._predict_fn = (
                        lambda spec, sigma_lo, _sm=_instr_sm, _L=_Lrt:
                            _sm(_L(spec[_idx], sigma_lo))
                    )
                elif has_losvd:
                    _L = _apply_losvd
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm, _L=_L:
                            _sm(_L(spec[_idx]))
                    )
                else:
                    self._predict_fn = (
                        lambda spec, _sm=_instr_sm:
                            _sm(spec[_idx])
                    )

            else:
                raise ValueError(
                    f"Spectrum.smoothtype={st!r} is not recognised.  "
                    "Valid choices: None, 'vel', 'R', 'lambda', 'lsf'."
                )

        else:
            # LOSVD only (no instrumental smoothing); _apply_losvd
            # already maps _wm_trim → wo (observed grid).
            if fit_lo:
                _Lrt = _apply_losvd_rt
                self._predict_fn = (
                    lambda spec, sigma_lo, _L=_Lrt:
                        _L(spec[_idx], sigma_lo)
                )
            else:
                _L = _apply_losvd
                self._predict_fn = lambda spec, _L=_L: _L(spec[_idx])

predict

predict(spectrum, wave_model, sigma_smooth=None)

Project the model spectrum onto the observed pixel grid, applying instrumental smoothing if configured.

Calls _predict_fn(spectrum[, sigma_smooth]) which was constructed by setup_for_model. Depending on self.smoothtype:

• None — pure linear interpolation (_H @ spectrum).
• "vel" / "R" — constant-velocity FFT broadening then interpolation to observed pixels.
• "lambda" — constant-wavelength FFT broadening then interpolation.
• "lsf" — wavelength-dependent LSF broadening (CDF-transform FFT) then interpolation.

In all smoothing cases, the full smooth→interpolate pipeline is a single closure that is JAX-JIT-compilable with respect to spectrum.

Must call setup_for_model(wave_model) before this method.

Parameters:

Name	Type	Description	Default
spectrum	(Array, shape(n_wave))	Model spectrum in F_nu units on the model wavelength grid.	required
wave_model	(Array, shape(n_wave))	Model wavelength grid [Å] (accepted for interface consistency; the grid mapping was precomputed by setup_for_model).	required
sigma_smooth	jax.Array scalar	Runtime galaxy LOSVD [km/s] -- the Prospector sigma_smooth parameter. Only consulted when this Spectrum was constructed with fit_sigma_smooth=True; ignored otherwise (the static sigma_losvd baked at setup_for_model time is used instead). Passing the value from theta makes the LOSVD differentiable and fittable inside JIT/SVI/NUTS.	None

Returns:

Type	Description
(Array, shape(n_pix))	Model F_nu (smoothed and) interpolated onto self.wavelength.

Source code in ceridwen/observation/spectrum.py

def predict(self, spectrum, wave_model, sigma_smooth=None):
    """
    Project the model spectrum onto the observed pixel grid, applying
    instrumental smoothing if configured.

    Calls ``_predict_fn(spectrum[, sigma_smooth])`` which was constructed
    by ``setup_for_model``.  Depending on ``self.smoothtype``:

    * ``None``  — pure linear interpolation (``_H @ spectrum``).
    * ``"vel"`` / ``"R"`` — constant-velocity FFT broadening then
      interpolation to observed pixels.
    * ``"lambda"`` — constant-wavelength FFT broadening then interpolation.
    * ``"lsf"`` — wavelength-dependent LSF broadening (CDF-transform FFT)
      then interpolation.

    In all smoothing cases, the full smooth→interpolate pipeline is a
    single closure that is JAX-JIT-compilable with respect to ``spectrum``.

    Must call ``setup_for_model(wave_model)`` before this method.

    Parameters
    ----------
    spectrum : jax.Array, shape (n_wave,)
        Model spectrum in F_nu units on the model wavelength grid.
    wave_model : jax.Array, shape (n_wave,)
        Model wavelength grid [Å] (accepted for interface consistency;
        the grid mapping was precomputed by ``setup_for_model``).
    sigma_smooth : jax.Array scalar, optional
        Runtime galaxy LOSVD [km/s] -- the Prospector ``sigma_smooth``
        parameter.  Only consulted when this Spectrum was constructed
        with ``fit_sigma_smooth=True``; ignored otherwise (the static
        ``sigma_losvd`` baked at ``setup_for_model`` time is used
        instead).  Passing the value from ``theta`` makes the LOSVD
        differentiable and fittable inside JIT/SVI/NUTS.

    Returns
    -------
    jax.Array, shape (n_pix,)
        Model F_nu (smoothed and) interpolated onto ``self.wavelength``.
    """
    # Clear error instead of a cryptic AttributeError on the projection
    # closure if setup was skipped.  ``hasattr`` is a static Python check,
    # so inside a jit trace it resolves at compile time (no hot-path cost).
    if not hasattr(self, "_predict_fn"):
        raise RuntimeError(
            "Spectrum.predict() called before setup_for_model(): the "
            "projection/smoothing closure has not been built. Call "
            "spec.setup_for_model(wave_model) once (before the first "
            "predict / JIT trace)."
        )
    if self.fit_sigma_smooth:
        if sigma_smooth is None:
            # Fall back to the constructor default; lets a caller
            # invoke ``predict(spec, wave)`` for warmup / debugging
            # even when the closure is the runtime-sigma variant.
            sigma_smooth = self.sigma_losvd
        # Pass sigma through without forcing a dtype -- JAX's
        # dtype-promotion rules with jax_enable_x64=True will
        # promote to float64 to match the cached smoother grids,
        # giving the same precision the static fast path achieved.
        # Forcing float32 here would round-trip-degrade the smoother
        # output even when the user is fitting in double precision.
        sv = jnp.asarray(sigma_smooth).reshape(())
        return self._predict_fn(spectrum, sv)
    return self._predict_fn(spectrum)

synthetic_photometry

synthetic_photometry(filterset)

Project this spectrum onto a FilterSet and return synthetic maggies.

The spectrum is assumed to be in F_nu units (e.g. L_sun/Hz/M_sun) and is converted to F_lambda before the AB-normalised projection.

Parameters:

Name	Type	Description	Default
filterset	FilterSet		required

Returns:

Name	Type	Description
maggies	(ndarray, shape(n_filters))	Returns None if the spectrum has no data.

Source code in ceridwen/observation/spectrum.py

def synthetic_photometry(self, filterset):
    """
    Project this spectrum onto a FilterSet and return synthetic maggies.

    The spectrum is assumed to be in **F_nu units** (e.g. L_sun/Hz/M_sun)
    and is converted to F_lambda before the AB-normalised projection.

    Parameters
    ----------
    filterset : sedpy_jax.observate.FilterSet

    Returns
    -------
    maggies : jnp.ndarray, shape (n_filters,)
        Returns ``None`` if the spectrum has no data.
    """
    if self.flux is None or self._wavelength is None:
        return None
    _c        = jnp.array(2.998e18)   # Å/s
    flux_flam = self.flux * _c / self._wavelength**2
    return filterset.get_sed_maggies(flux_flam, sourcewave=self._wavelength)

mask_wavelength_range

mask_wavelength_range(wave_min, wave_max)

Mask pixels with wavelengths in [wave_min, wave_max] Å (inclusive).

Sets self.mask[i] = False for all pixels whose wavelength falls inside the specified range.

Parameters:

Name	Type	Description	Default
wave_min	float	Wavelength bounds [Å].	required
wave_max	float	Wavelength bounds [Å].	required

Source code in ceridwen/observation/spectrum.py

def mask_wavelength_range(self, wave_min, wave_max):
    """
    Mask pixels with wavelengths in [wave_min, wave_max] Å (inclusive).

    Sets ``self.mask[i] = False`` for all pixels whose wavelength falls
    inside the specified range.

    Parameters
    ----------
    wave_min, wave_max : float
        Wavelength bounds [Å].
    """
    if self._wavelength is None:
        return
    in_range  = (self._wavelength >= wave_min) & (self._wavelength <= wave_max)
    self.mask = self.mask & ~in_range

mask_lines

mask_lines(line_waves, dv=1000.0, zred=0.0)

Mask spectral lines by zeroing the mask within ±dv km/s of each line.

Parameters:

Name	Type	Description	Default
line_waves	array - like	Rest-frame central wavelengths [Å]; redshifted internally by (1 + zred) before masking against the observed-frame pixel grid.	required
dv	float	Half-width to mask on each side [km/s]. Default 1000 km/s.	1000.0
zred	float	Redshift to apply to line_waves. Default 0.0 (preserves pre-fix behaviour for callers already passing observed-frame wavelengths).	0.0

Source code in ceridwen/observation/spectrum.py

def mask_lines(self, line_waves, dv=1000.0, zred=0.0):
    """
    Mask spectral lines by zeroing the mask within ±dv km/s of each line.

    Parameters
    ----------
    line_waves : array-like
        Rest-frame central wavelengths [Å]; redshifted internally by
        (1 + ``zred``) before masking against the observed-frame pixel
        grid.
    dv : float
        Half-width to mask on each side [km/s].  Default 1000 km/s.
    zred : float
        Redshift to apply to ``line_waves``.  Default 0.0 (preserves
        pre-fix behaviour for callers already passing observed-frame
        wavelengths).
    """
    if self._wavelength is None:
        return
    c_kms = 2.998e5   # km/s
    opz   = 1.0 + float(zred)
    for lam0_rest in np.asarray(line_waves).ravel():
        lam0_obs = opz * float(lam0_rest)
        dlam     = lam0_obs * dv / c_kms
        self.mask_wavelength_range(lam0_obs - dlam, lam0_obs + dlam)

chi_sq

chi_sq(model_flux)

Chi-squared contribution from this spectrum.

Accounts for sky subtraction (self.sky) and a fractional noise floor (self.noise_floor). If self.logify_spectrum is True, residuals are computed in log-flux space: Δln f / (σ_eff / f).

Parameters:

Name	Type	Description	Default
model_flux	(array - like, shape(n_pix))	Model flux on the observed pixel grid (output of predict).	required

Returns:

Name	Type	Description
chi2	float	Sum of squared normalised residuals over unmasked pixels.

Source code in ceridwen/observation/spectrum.py

def chi_sq(self, model_flux):
    """
    Chi-squared contribution from this spectrum.

    Accounts for sky subtraction (``self.sky``) and a fractional noise
    floor (``self.noise_floor``).  If ``self.logify_spectrum`` is True,
    residuals are computed in log-flux space: Δln f / (σ_eff / f).

    Parameters
    ----------
    model_flux : array-like, shape (n_pix,)
        Model flux on the observed pixel grid (output of ``predict``).

    Returns
    -------
    chi2 : float
        Sum of squared normalised residuals over unmasked pixels.
    """
    resid = self._compute_residuals(model_flux)
    return float(jnp.sum(jnp.where(self.mask, resid ** 2, 0.0)))

residuals

residuals(model_flux)

Per-pixel (sky-corrected data − model) / sigma_eff. Masked pixels are set to NaN.

Returns:

Name	Type	Description
res	(ndarray, shape(n_pix))

Source code in ceridwen/observation/spectrum.py

def residuals(self, model_flux):
    """
    Per-pixel (sky-corrected data − model) / sigma_eff.
    Masked pixels are set to NaN.

    Returns
    -------
    res : jnp.ndarray, shape (n_pix,)
    """
    resid = self._compute_residuals(model_flux)
    return jnp.where(self.mask, resid, jnp.nan)

log_likelihood

log_likelihood(model_flux)

Full log-likelihood for this spectrum.

Combines the standard pixel-independent Gaussian log-likelihood with an optional Gaussian Process (GP) correction for correlated residuals if self.noise is a GaussianProcess instance.

.. math::

\log\mathcal{L} =
    -\tfrac{1}{2}\sum_{i\,\in\,\rm mask} r_i^2
    + \log p_{\rm GP}(\mathbf{r} \mid \mathrm{GP})

where :math:r_i = ((d_i - s_i) - c_i m_i)/\sigma_{\rm eff,i} (sky :math:s_i and calibration :math:c_i are optional; the noise floor in :math:\sigma_{\rm eff,i} is scaled against the calibrated model :math:c_i m_i when calibration is set), and the GP term is zero if no noise model is set.

Parameters:

Name	Type	Description	Default
model_flux	(array - like, shape(n_pix))	Model flux on the observed pixel grid.	required

Returns:

Type	Description
float	Log-likelihood (larger is better).

Source code in ceridwen/observation/spectrum.py

def log_likelihood(self, model_flux):
    """
    Full log-likelihood for this spectrum.

    Combines the standard pixel-independent Gaussian log-likelihood with
    an optional Gaussian Process (GP) correction for correlated residuals
    if ``self.noise`` is a ``GaussianProcess`` instance.

    .. math::

        \\log\\mathcal{L} =
            -\\tfrac{1}{2}\\sum_{i\\,\\in\\,\\rm mask} r_i^2
            + \\log p_{\\rm GP}(\\mathbf{r} \\mid \\mathrm{GP})

    where :math:`r_i = ((d_i - s_i) - c_i m_i)/\\sigma_{\\rm eff,i}`
    (sky :math:`s_i` and calibration :math:`c_i` are optional; the
    noise floor in :math:`\\sigma_{\\rm eff,i}` is scaled against the
    calibrated model :math:`c_i m_i` when ``calibration`` is set), and
    the GP term is zero if no noise model is set.

    Parameters
    ----------
    model_flux : array-like, shape (n_pix,)
        Model flux on the observed pixel grid.

    Returns
    -------
    float
        Log-likelihood (larger is better).
    """
    resid  = self._compute_residuals(model_flux)
    chi2   = float(jnp.sum(jnp.where(self.mask, resid ** 2, 0.0)))
    log_ll = -0.5 * chi2

    if self.noise is not None and self._wavelength is not None:
        log_ll += self.noise.log_likelihood(
            np.array(resid),
            np.array(self._wavelength),
            np.array(self.mask),
        )
    return float(log_ll)

fit_polynomial_calibration

fit_polynomial_calibration(model_flux, order=3)

Analytically fit a Chebyshev multiplicative calibration polynomial P(λ) such that data ≈ P(λ) × model_flux.

The polynomial coefficients are solved at each call via weighted linear least squares, making this suitable for marginalising out the calibration at every likelihood evaluation without a parameter-space penalty.

The polynomial is evaluated in a normalised wavelength coordinate :math:x \in [-1, 1] using Chebyshev basis functions :math:T_n(x), which are numerically stable for high orders.

Parameters:

Name	Type	Description	Default
model_flux	(array - like, shape(n_pix))	Model flux on the observed pixel grid (output of predict).	required
order	int	Polynomial order. 0 = constant, 1 = linear, etc. Default 3.	3

Returns:

Name	Type	Description
coeffs	(ndarray, shape(order + 1))	Chebyshev polynomial coefficients.
calibrated_flux	(ndarray, shape(n_pix))	P(λ) × model_flux — the calibration-corrected model prediction to be compared with self.flux.

Notes

Only unmasked pixels enter the least-squares fit. The returned calibrated_flux is evaluated over the full pixel grid.

Source code in ceridwen/observation/spectrum.py

def fit_polynomial_calibration(self, model_flux, order: int = 3):
    """
    Analytically fit a Chebyshev multiplicative calibration polynomial
    P(λ) such that ``data ≈ P(λ) × model_flux``.

    The polynomial coefficients are solved at each call via weighted
    linear least squares, making this suitable for marginalising out the
    calibration at every likelihood evaluation without a parameter-space
    penalty.

    The polynomial is evaluated in a normalised wavelength coordinate
    :math:`x \\in [-1, 1]` using Chebyshev basis functions
    :math:`T_n(x)`, which are numerically stable for high orders.

    Parameters
    ----------
    model_flux : array-like, shape (n_pix,)
        Model flux on the observed pixel grid (output of ``predict``).
    order : int, optional
        Polynomial order.  0 = constant, 1 = linear, etc.  Default 3.

    Returns
    -------
    coeffs : np.ndarray, shape (order + 1,)
        Chebyshev polynomial coefficients.
    calibrated_flux : jnp.ndarray, shape (n_pix,)
        ``P(λ) × model_flux`` — the calibration-corrected model
        prediction to be compared with ``self.flux``.

    Notes
    -----
    Only unmasked pixels enter the least-squares fit.  The returned
    ``calibrated_flux`` is evaluated over the full pixel grid.
    """
    mf     = np.asarray(model_flux, dtype=np.float64)
    data   = np.asarray(self._sky_corrected_data(), dtype=np.float64)
    sigma  = np.asarray(self.uncertainty, dtype=np.float64)
    wav    = np.asarray(self._wavelength,  dtype=np.float64)
    mask   = np.asarray(self.mask,         dtype=bool)

    # Normalise wavelength axis to [-1, 1] for numerical stability
    wav_mid  = 0.5 * (wav.max() + wav.min())
    wav_half = 0.5 * (wav.max() - wav.min())
    x        = (wav - wav_mid) / (wav_half if wav_half > 0 else 1.0)

    # Chebyshev design matrix: A[i, n] = T_n(x_i)
    A = np.polynomial.chebyshev.chebvander(x, order)  # (n_pix, order+1)

    # Weight by model flux and 1/sigma so we minimise
    # sum_mask ((data_i - P(x_i) * model_i) / sigma_i)^2
    A_w = (A * mf[:, None]) / sigma[:, None]   # (n_pix, order+1)
    y_w = data / sigma                          # (n_pix,)

    # Apply mask
    A_wm = A_w[mask]
    y_wm = y_w[mask]

    # Solve linear least squares
    coeffs, _, _, _ = np.linalg.lstsq(A_wm, y_wm, rcond=None)

    # Evaluate calibration polynomial on the full pixel grid
    poly_vals       = A @ coeffs                      # (n_pix,)
    calibrated_flux = jnp.asarray(poly_vals * mf)

    return coeffs, calibrated_flux

ceridwen.observation.Lines

Lines(
    line_ind,
    line_names=None,
    wavelength=None,
    name=None,
    upper_limit=None,
    **kwargs
)

Bases: Observation

Observed nebular emission-line fluxes.

Stores a set of emission-line fluxes together with their FSPS line-array indices, vacuum rest-frame wavelengths, and per-line 1-sigma uncertainties. The interface is deliberately compatible with prospect.observation.Lines: the line_ind attribute holds integer indices into the FSPS emline_luminosity array, and the alias mapping exposes "line_inds" as an alias for line_ind so that existing Prospector model code can address this object without modification.

Beyond Prospector, this class adds JAX-native chi_sq / residuals (JIT-compilable through a fitter), and mask_by_name / select_by_name helpers that operate on human-readable line names.

Parameters:

Name	Type	Description	Default
line_ind	array-like of int	Indices of the observed lines in the FSPS emission-line array ($SPS_HOME/data/emlines_info.dat). Required.	required
line_names	list of str	Human-readable names, one per line (e.g. "Halpha", "[OIII]5007"). Required for mask_by_name and select_by_name.	None
wavelength	array-like of float	Vacuum rest-frame wavelengths [Å], length = len(line_ind).	None
flux	array-like of float	Observed line fluxes. Units should be consistent with any model prediction passed to chi_sq / residuals (typically erg s⁻¹ cm⁻²).	required
uncertainty	array-like of float	1-sigma line-flux uncertainties, same units as flux.	required
mask	array-like of bool	True for lines to include in chi-squared. Defaults to all-True.	required
upper_limit	array-like of bool, shape (n_lines,)	If True for a given line, that line is treated as a non-detection upper limit rather than a positive detection. The chi-squared contribution for such lines is one-sided: a penalty is applied only when the model flux exceeds the observed value (i.e., the model predicts more emission than the upper limit allows): .. math:: \chi^2_{\rm UL} = \begin{cases} \left(\frac{d - m}{\sigma}\right)^2 & m > d \\ 0 & m \leq d \end{cases} where :math:d is the observed upper-limit flux and :math:m is the model prediction. Physically this corresponds to integrating the likelihood over all undetected flux values below the upper limit. Default None (all lines treated as detections).	None

Examples:

>>> lines = Lines(
...     line_ind   = [59, 63, 71],
...     line_names = ["Hbeta", "[OIII]5007", "Halpha"],
...     wavelength = [4861., 5007., 6563.],
...     flux       = obs_fluxes,
...     uncertainty= obs_unc,
...     upper_limit= [False, True, False],  # [OIII]5007 is a non-detection
... )
>>> lines.mask_by_name(["[OIII]5007"])   # exclude one line
>>> chi2 = lines.chi_sq(model_fluxes)
>>> subset = lines.select_by_name(["Hbeta", "Halpha"])

Source code in ceridwen/observation/lines.py

def __init__(
    self,
    line_ind,
    line_names  = None,
    wavelength  = None,
    name        = None,
    upper_limit = None,
    **kwargs,
):
    if line_ind is None:
        raise ValueError(
            "line_ind is required: pass the indices of the observed lines "
            "in the FSPS emline_luminosity array."
        )
    if wavelength is None:
        raise ValueError(
            "wavelength is required: pass the wavelengths of the observed lines."
        )
    self.line_ind   = jnp.asarray(np.atleast_1d(line_ind), dtype=int)
    self.line_names = list(line_names) if line_names is not None else None
    self._wavelength = (
        None if wavelength is None
        else jnp.asarray(np.atleast_1d(wavelength), dtype=float)
    )
    self.upper_limit = (
        None if upper_limit is None
        else jnp.asarray(np.atleast_1d(upper_limit), dtype=bool)
    )
    super().__init__(name=name, **kwargs)

setup_for_model

setup_for_model(wave_model, sigma_v=200.0, zred=0.0)

Precompute the (n_lines, n_wave) Gaussian-aperture weight matrix _W that extracts line fluxes from a model spectrum via a single matrix–vector multiply.

Must be called once before predict (and before JIT-compiling any function containing predict). SedModel.__init__ calls this automatically.

Physical description

For each emission line centred at wavelength :math:\lambda_k, the integrated line flux is estimated as a Gaussian-weighted integral over the model spectrum:

.. math::

F_k = \int w_k(\lambda)\, f_\nu(\lambda)\, \mathrm{d}\lambda

where

.. math::

w_k(\lambda) = \exp\!\left[-\frac{1}{2}
\left(\frac{\lambda - \lambda_k}{\sigma_k}\right)^2\right],
\quad \sigma_k = \lambda_k\, \frac{\sigma_v}{c}

Discretised with the trapezoidal rule on the model wavelength grid, this becomes _W @ spectrum where _W[k, j] = w_k(wave_j) * dlambda_j and dlambda_j are the trapezoidal quadrature weights.

Parameters:

Name	Type	Description	Default
wave_model	(array - like, shape(n_wave))	Model wavelength grid [Å], strictly increasing.	required
sigma_v	float	1-sigma Gaussian aperture width [km/s]. Default 200 km/s. Sufficient to capture narrow nebular lines as generated by NebularGridModel while excluding continuum and neighbouring lines spaced by more than ~600 km/s. The same aperture is applied to both model and data so the absolute calibration cancels in the likelihood. sigma_v is a construction-time hyperparameter; it is not part of theta and is not differentiable through predict. Catalogued line fluxes are single scalars per line and carry no shape information, so HMC cannot constrain it.	200.0

Source code in ceridwen/observation/lines.py

def setup_for_model(self, wave_model, sigma_v=200.0, zred: float = 0.0):
    """
    Precompute the (n_lines, n_wave) Gaussian-aperture weight matrix
    ``_W`` that extracts line fluxes from a model spectrum via a single
    matrix–vector multiply.

    Must be called once before ``predict`` (and before JIT-compiling any
    function containing ``predict``).  ``SedModel.__init__`` calls this
    automatically.

    Physical description
    --------------------
    For each emission line centred at wavelength :math:`\\lambda_k`, the
    integrated line flux is estimated as a Gaussian-weighted integral over
    the model spectrum:

    .. math::

        F_k = \\int w_k(\\lambda)\\, f_\\nu(\\lambda)\\, \\mathrm{d}\\lambda

    where

    .. math::

        w_k(\\lambda) = \\exp\\!\\left[-\\frac{1}{2}
        \\left(\\frac{\\lambda - \\lambda_k}{\\sigma_k}\\right)^2\\right],
        \\quad \\sigma_k = \\lambda_k\\, \\frac{\\sigma_v}{c}

    Discretised with the trapezoidal rule on the model wavelength grid,
    this becomes ``_W @ spectrum`` where
    ``_W[k, j] = w_k(wave_j) * dlambda_j`` and ``dlambda_j`` are the
    trapezoidal quadrature weights.

    Parameters
    ----------
    wave_model : array-like, shape (n_wave,)
        Model wavelength grid [Å], strictly increasing.
    sigma_v : float, optional
        1-sigma Gaussian aperture width [km/s].  Default 200 km/s.
        Sufficient to capture narrow nebular lines as generated by
        ``NebularGridModel`` while excluding continuum and neighbouring
        lines spaced by more than ~600 km/s.  The same aperture is applied
        to both model and data so the absolute calibration cancels in the
        likelihood.

        ``sigma_v`` is a construction-time hyperparameter; it is not part
        of ``theta`` and is not differentiable through ``predict``.
        Catalogued line fluxes are single scalars per line and carry no
        shape information, so HMC cannot constrain it.
    """
    wm_rest = np.asarray(wave_model,        dtype=np.float64)   # (n_wave,)
    lam0_rest = np.asarray(self._wavelength, dtype=np.float64)   # (n_lines,)
    c_kms  = 2.998e5  # km/s
    opz = 1.0 + float(zred)

    # Both the model grid and the line centres move together into the
    # observed frame by the (1 + zred) factor.  The Gaussian shape is
    # preserved because the velocity aperture sigma_v is defined in
    # velocity units — at higher redshift the wavelength sigma grows
    # proportionally with the line wavelength, so (lambda - lambda_0) /
    # sigma is invariant.
    wm   = opz * wm_rest
    lam0 = opz * lam0_rest

    # Trapezoidal quadrature weights along the (observed-frame) model
    # wavelength axis.  At zred > 0 these pick up one factor of
    # (1 + zred) naturally — this is the dlambda_obs = (1+z) dlambda_rest
    # Jacobian — so integrated line fluxes scale with (1+z) as expected
    # for a redshift-preserving Gaussian aperture.
    dlam        = np.empty(len(wm), dtype=np.float64)
    dlam[1:-1]  = 0.5 * (wm[2:] - wm[:-2])
    dlam[0]     = 0.5 * (wm[1]  - wm[0])
    dlam[-1]    = 0.5 * (wm[-1] - wm[-2])

    # Absolute-flux normalisation.  Without this per-line factor,
    # ``_W @ F_nu`` returns ``F_line * lambda_obs**2 / c`` (units:
    # erg s^-1 cm^-2 Hz^-1 * A — a mixed-unit "aperture proxy"), NOT
    # the integrated line flux in erg s^-1 cm^-2.  The raw proxy is
    # fine if you feed *both* data and model through the same
    # aperture (the docstring's "calibration cancels" regime), but
    # catalogue emission-line tables almost always quote already-
    # reduced integrated line fluxes in erg s^-1 cm^-2 — so we
    # normalise once here and have ``_W @ spectrum`` return flux
    # in the catalogue's own unit system.
    #
    # Derivation: FSPS Cloudy lines are added to the spectrum with
    # a Gaussian of width sigma_v_model = nebular_smooth_init km/s
    # (floor of ~2 pixel widths).  This is typically narrower than
    # the sigma_v = 200 km/s aperture used here.  In the narrow-
    # model-line limit the aperture integral reduces to
    #   _W @ F_nu ≈ F_line * lambda_obs^2 / c .
    # Multiplying each row by c / lambda_obs^2 restores
    #   _W @ F_nu ≈ F_line [erg s^-1 cm^-2] ,
    # letting observed data in the same units be passed in directly
    # as ``Lines.flux``.
    c_aa_s = 2.998e18                          # speed of light [Å/s]
    norm = c_aa_s / (lam0 ** 2)                # (n_lines,)

    # Bake W into a static (n_lines, n_wave) JAX constant.  XLA
    # constant-folds at trace time.
    diff     = wm[None, :] - lam0[:, None]         # (n_lines, n_wave)
    sigma_aa = lam0 * (sigma_v / c_kms)            # (n_lines,)
    W = np.exp(-0.5 * (diff / sigma_aa[:, None]) ** 2)
    W = (W * dlam[None, :]).astype(np.float32)     # (n_lines, n_wave)
    W = (W * norm[:, None].astype(np.float32))
    self._W = jnp.array(W)

predict

predict(spectrum, wave_model)

Extract emission-line fluxes from the model spectrum via Gaussian- aperture integration: computes _W @ spectrum where _W was precomputed once in setup_for_model. On GPU this is a single GEMV; XLA constant-folds _W into the compiled graph.

Must call setup_for_model(wave_model, sigma_v=...) first.

Parameters:

Name	Type	Description	Default
spectrum	(Array, shape(n_wave))	Model spectrum in F_nu units.	required
wave_model	(Array, shape(n_wave))	Accepted for interface consistency; not used inside this method.	required

Returns:

Type	Description
(Array, shape(n_lines))	Gaussian-aperture integrated flux for each line.

Source code in ceridwen/observation/lines.py

def predict(self, spectrum, wave_model):
    """
    Extract emission-line fluxes from the model spectrum via Gaussian-
    aperture integration: computes ``_W @ spectrum`` where ``_W`` was
    precomputed once in ``setup_for_model``.  On GPU this is a single
    GEMV; XLA constant-folds ``_W`` into the compiled graph.

    Must call ``setup_for_model(wave_model, sigma_v=...)`` first.

    Parameters
    ----------
    spectrum : jax.Array, shape (n_wave,)
        Model spectrum in F_nu units.
    wave_model : jax.Array, shape (n_wave,)
        Accepted for interface consistency; not used inside this method.

    Returns
    -------
    jax.Array, shape (n_lines,)
        Gaussian-aperture integrated flux for each line.
    """
    if not hasattr(self, "_W"):
        raise RuntimeError(
            "Lines.predict() called before setup_for_model(): the "
            "Gaussian aperture weight matrix has not been built. "
            "Call lines.setup_for_model(wave_model, sigma_v=...) "
            "once before the first predict / JIT trace."
        )
    return self._W @ spectrum

chi_sq

chi_sq(model_fluxes)

Chi-squared contribution from the observed line fluxes.

For lines flagged as upper limits (self.upper_limit[k] = True), the contribution is one-sided: a penalty is applied only when the model flux exceeds the observed upper-limit value.

Parameters:

Name	Type	Description	Default
model_fluxes	(array - like, shape(n_lines))	Predicted line fluxes, same units as self.flux.	required

Returns:

Name	Type	Description
chi2	float

Source code in ceridwen/observation/lines.py

def chi_sq(self, model_fluxes):
    """
    Chi-squared contribution from the observed line fluxes.

    For lines flagged as upper limits (``self.upper_limit[k] = True``),
    the contribution is one-sided: a penalty is applied only when the
    model flux exceeds the observed upper-limit value.

    Parameters
    ----------
    model_fluxes : array-like, shape (n_lines,)
        Predicted line fluxes, same units as ``self.flux``.

    Returns
    -------
    chi2 : float
    """
    mf    = jnp.asarray(model_fluxes, dtype=float)
    resid = (self.flux - mf) / self.uncertainty       # (data - model)/sigma

    if self.upper_limit is not None:
        # For upper-limit lines: only penalise when model > data,
        # i.e., when resid < 0  (model exceeded the observed limit).
        resid_sq = jnp.where(
            self.upper_limit,
            jnp.where(resid < 0.0, resid ** 2, 0.0),
            resid ** 2,
        )
    else:
        resid_sq = resid ** 2

    return float(jnp.sum(jnp.where(self.mask, resid_sq, 0.0)))

residuals

residuals(model_fluxes)

Per-line (data − model) / sigma. Masked lines are set to NaN. Upper-limit lines where the model does not exceed the limit are set to zero (no tension) rather than showing a negative residual.

Returns:

Name	Type	Description
res	(ndarray, shape(n_lines))

Source code in ceridwen/observation/lines.py

def residuals(self, model_fluxes):
    """
    Per-line ``(data − model) / sigma``.  Masked lines are set to NaN.
    Upper-limit lines where the model does not exceed the limit are set
    to zero (no tension) rather than showing a negative residual.

    Returns
    -------
    res : jnp.ndarray, shape (n_lines,)
    """
    mf    = jnp.asarray(model_fluxes, dtype=float)
    resid = (self.flux - mf) / self.uncertainty

    if self.upper_limit is not None:
        # Show zero residual when model is safely below the upper limit
        resid = jnp.where(
            self.upper_limit & (resid >= 0.0),
            0.0,
            resid,
        )

    return jnp.where(self.mask, resid, jnp.nan)

mask_by_name

mask_by_name(names)

Exclude lines whose name appears in names from chi-squared.

Sets self.mask[i] = False for all lines whose entry in self.line_names matches any element of names. A no-op if self.line_names is not set.

Parameters:

Name	Type	Description	Default
names	list of str		required

Source code in ceridwen/observation/lines.py

def mask_by_name(self, names):
    """
    Exclude lines whose name appears in ``names`` from chi-squared.

    Sets ``self.mask[i] = False`` for all lines whose entry in
    ``self.line_names`` matches any element of ``names``.  A no-op if
    ``self.line_names`` is not set.

    Parameters
    ----------
    names : list of str
    """
    if self.line_names is None:
        return
    names_set = set(names)
    exclude = jnp.array(
        [n in names_set for n in self.line_names], dtype=bool
    )
    self.mask = self.mask & ~exclude

select_by_name

select_by_name(names)

Return a new Lines instance containing only the named lines.

Parameters:

Name	Type	Description	Default
names	list of str	Must all be present in self.line_names.	required

Returns:

Type	Description
Lines

Raises:

Type	Description
ValueError	If self.line_names is not set.
KeyError	If any element of names is absent from self.line_names.

Source code in ceridwen/observation/lines.py

def select_by_name(self, names):
    """
    Return a new ``Lines`` instance containing only the named lines.

    Parameters
    ----------
    names : list of str
        Must all be present in ``self.line_names``.

    Returns
    -------
    Lines

    Raises
    ------
    ValueError
        If ``self.line_names`` is not set.
    KeyError
        If any element of ``names`` is absent from ``self.line_names``.
    """
    if self.line_names is None:
        raise ValueError(
            "line_names not set on this Lines object; "
            "cannot select by name."
        )
    idx = []
    for n in names:
        if n not in self.line_names:
            raise KeyError(
                f"Line '{n}' not found.  Available: {self.line_names}"
            )
        idx.append(self.line_names.index(n))
    idx = np.array(idx)

    def _pick(arr):
        return None if arr is None else np.array(arr)[idx]

    return Lines(
        line_ind    = np.array(self.line_ind)[idx],
        line_names  = [self.line_names[i] for i in idx],
        wavelength  = _pick(self._wavelength),
        flux        = _pick(self.flux),
        uncertainty = _pick(self.uncertainty),
        mask        = np.array(self.mask)[idx],
        upper_limit = _pick(self.upper_limit) if self.upper_limit is not None else None,
        name        = self.name + "_sel",
    )

Model

ceridwen.model.SedModel

SedModel(
    csp,
    observations,
    priors=None,
    transforms=None,
    free_param_init=None,
    zred=0.0,
)

Parameter manager + prediction layer for Ceridwen SED fitting.

Parameters:

Name	Type	Description	Default
csp	CSPBasis	Initialised composite stellar population model. Must expose csp.wave, csp.theta_init, and csp.predict(theta, observations).	required
observations	list of Observation	Data containers (Photometry, Spectrum, Lines). Each must have a unique .name attribute.	required
priors	dict[str, Prior]	Mapping from free-parameter name to a prior object implementing logpdf(x) -> Array. Parameters absent from this dict receive no prior contribution (flat improper prior).	None
transforms	dict[str, callable]	Mapping from derived (CSP) parameter name to a callable that computes its value from the free-parameter dict:: model_theta[derived] = fn(free_theta) Derived parameters listed here are removed from the free-parameter list and replaced by the new free parameters supplied via free_param_init. This mirrors Prospector's depends_on mechanism. Example — fitting log-ratios of SFR bins instead of raw SFH:: from ceridwen.model.transforms import logsfr_ratios_to_sfh transforms = { "sfh": lambda t: logsfr_ratios_to_sfh( t["logsfr_ratios"], sfh_times_yr=csp.sfh_times, ) }	None
free_param_init	dict[str, Array]	Initial values for free parameters that replace derived ones. Keys in this dict are added to param_names and theta_init; the corresponding derived params (transforms keys) are removed. Required when transforms is not empty.	None

Attributes:

Name	Type	Description
theta_init	dict[str, Array]	Initial values for the free parameters only (derived params are absent; their replacements from free_param_init are present).
param_names	list[str]	Ordered list of free-parameter names.
transforms	dict[str, callable]	Registered transforms (empty dict if none).
obs_dict	dict[str, Observation]	Observations keyed by obs.name.
wave	(Array, shape(n_wave))	Model wavelength grid [Å].

Source code in ceridwen/model/model.py

def __init__(
    self,
    csp,
    observations: Sequence[Observation],
    priors: dict[str, Any] | None = None,
    transforms: dict[str, Callable] | None = None,
    free_param_init: dict[str, Any] | None = None,
    zred: float = 0.0,
):
    self.csp          = csp
    self.observations = list(observations)
    self.priors       = dict(priors) if priors is not None else {}
    self.transforms   = dict(transforms) if transforms is not None else {}
    self.zred         = float(zred)

    # Validate that all observation names are unique
    names = [obs.name for obs in self.observations]
    if len(names) != len(set(names)):
        dups = [n for n in names if names.count(n) > 1]
        raise ValueError(
            f"Observation names must be unique.  Duplicates found: {dups}"
        )

    # Start from the CSP's full parameter set
    self.theta_init  = dict(csp.theta_init)
    self.param_names = list(csp.param_names)
    self.wave        = csp.wave

    # Apply transforms bookkeeping:
    #   1. Remove derived parameters (they are outputs of transforms, not
    #      free parameters that the sampler proposes).
    #   2. Add the new free parameters supplied via free_param_init.
    if self.transforms:
        _derived = set(self.transforms.keys())
        for p in _derived:
            if p in self.theta_init:
                del self.theta_init[p]
            if p in self.param_names:
                self.param_names.remove(p)

        if free_param_init is not None:
            for p, v in free_param_init.items():
                arr = jnp.atleast_1d(jnp.asarray(v, dtype=float))
                self.theta_init[p] = arr
                if p not in self.param_names:
                    self.param_names.append(p)

    # Tell the CSP unknown-key guard about the model-level free parameters
    # (e.g. logsfr_ratios consumed by the sfh transform). apply_transforms
    # forwards the full free-parameter dict to csp.predict, so without this
    # those keys would be mis-flagged as typos on every fit.
    if hasattr(self.csp, "register_known_theta_keys"):
        self.csp.register_known_theta_keys(
            set(self.param_names) | set(self.priors) | set(self.transforms)
        )

    # Precompute static projection matrices for Spectrum and Lines.
    # This must happen once, at Python level, BEFORE any JIT trace of
    # predict().  Each observation's setup_for_model() stores a constant
    # JAX array (_H for Spectrum, _W for Lines) that XLA constant-folds
    # at trace time — meaning the GPU kernel contains no matrix construction,
    # only a single GEMV.  Photometry.setup_for_model() is a no-op.
    #
    # ``zred`` bakes the (1+z) wavelength stretch into the projection
    # matrices so the GEMV fast path stays the same shape for non-zero
    # fixed redshift.  The matching cosmological flux factor is applied
    # inside CSPBasis.predict via the ``zred`` entry that
    # :meth:`predict` injects into the CSP theta (see below) — both
    # things are needed for observed-frame calibration.
    for obs in self.observations:
        obs.setup_for_model(self.wave, zred=self.zred)

    # If the user supplied a non-trivial fixed redshift, store it so
    # :meth:`predict` can inject it into the CSP theta at trace time.
    #
    # A fixed zred must NOT be seeded into ``theta_init``: theta_init
    # is the sampled free-parameter pytree handed to the sampler
    # (``run_sampler`` -> ``adapter.run(..., model.theta_init, ...)``),
    # so a seeded entry silently becomes an UNPRIORED sampled
    # dimension.  Nothing bounds it; once the sampler drifts it to
    # z <= 0 the 10 pc fallback inside ``flux_factor_maggies`` erases
    # the entire (10pc/D_L)^2 dimming — a factor ~2e15 in flux at
    # z = 0.1 — without raising any error.  Injecting at predict time
    # instead guarantees EVERY prediction path (mock generation,
    # loglike_fn, predict_jit, predict_vmap) applies the same
    # cosmological normalisation, while keeping zred out of the
    # sampled parameter vector.
    #
    # When astropy is installed we prefer its Planck18 luminosity
    # distance for this one-off scalar computation (it includes
    # neutrinos + radiation and matches published tables to <0.1%),
    # and bake the resulting flux factor into a static JAX scalar.
    # The sampled path (when zred is free) continues to use the
    # native differentiable backend, so NUTS gradients still work.
    #
    # GOTCHA: only inject when there is no user-supplied ``zred``
    # transform.  If the user registered transforms={"zred": ...}
    # they are explicitly injecting zred at predict time from a
    # fixed external value; adding it here on top would double-route
    # the parameter.
    self._zred_fixed = None
    self.flux_factor_astropy = None
    if self.zred != 0.0 and "zred" not in self.transforms:
        self._zred_fixed = jnp.array([self.zred])
        try:
            from ..cosmology import (
                flux_factor_maggies, have_astropy,
            )
            if have_astropy():
                ff = float(flux_factor_maggies(
                    self.zred, backend="astropy"))
                # Stored for diagnostics; the free-z fit path ignores
                # this and recomputes via the native JAX backend.
                self.flux_factor_astropy = ff
        except Exception:
            # Non-fatal: fall through to the native backend.
            self.flux_factor_astropy = None

obs_dict property

obs_dict

Observations as a dict keyed by obs.name.

Pass this to MultiObservationLikelihood.make_lnprobfn as the observations argument::

lnprobfn = multi_lhood.make_lnprobfn(model.obs_dict, model, model)

n_obs property

n_obs

Number of registered observation objects.

apply_transforms

apply_transforms(free_theta)

Apply all registered transforms to produce a CSP-compatible model_theta.

Starts from a shallow copy of free_theta and computes each derived parameter by calling the corresponding transform callable::

model_theta[derived] = transform_fn(free_theta)

The free-parameter keys (e.g. "logsfr_ratios") are kept in model_theta alongside the derived ones; the CSP simply ignores any keys it does not recognise.

Parameters:

Name	Type	Description	Default
free_theta	dict[str, Array]	Free-parameter dict as used by the sampler.	required

Returns:

Name	Type	Description
model_theta	dict[str, Array]	Extended dict suitable for csp.predict. Contains all entries of free_theta plus the derived parameter values.

Source code in ceridwen/model/model.py

def apply_transforms(self, free_theta: dict[str, Array]) -> dict[str, Array]:
    """
    Apply all registered transforms to produce a CSP-compatible model_theta.

    Starts from a shallow copy of ``free_theta`` and computes each
    derived parameter by calling the corresponding transform callable::

        model_theta[derived] = transform_fn(free_theta)

    The free-parameter keys (e.g. ``"logsfr_ratios"``) are kept in
    ``model_theta`` alongside the derived ones; the CSP simply ignores
    any keys it does not recognise.

    Parameters
    ----------
    free_theta : dict[str, Array]
        Free-parameter dict as used by the sampler.

    Returns
    -------
    model_theta : dict[str, Array]
        Extended dict suitable for ``csp.predict``.  Contains all
        entries of ``free_theta`` plus the derived parameter values.
    """
    if not self.transforms:
        return free_theta
    model_theta = dict(free_theta)
    for derived_param, fn in self.transforms.items():
        model_theta[derived_param] = fn(free_theta)
    return model_theta

predict

predict(theta)

Project the CSP model spectrum onto all observations.

If transforms are registered, theta is treated as the free-parameter dict; apply_transforms is called first to obtain the CSP-compatible model_theta before forwarding to csp.predict.

Internally calls csp.predict(model_theta, observations) which computes the spectrum once and projects it onto each observation:

• Photometry → synthetic AB maggies via filter convolution
• Spectrum → model F_ν interpolated onto observed wavelength grid
• Lines → Gaussian-aperture integrated line fluxes

Mass scaling — if "logmass" is present in theta, the spectrum is multiplied by 10 ** logmass inside csp.predict() before projection. The logsfr_ratios_to_sfh transform normalises the SFH so that the trapezoidal integral of SFR over the lookback grid equals 1 M⊙ (Prospector / FSPS convention), so this factor sets the physical amplitude for a galaxy with stellar mass M = 10^logmass M⊙. Scaling once before projection is more efficient than scaling each observation separately.

Parameters:

Name	Type	Description	Default
theta	dict[str, Array]	Free-parameter dict (before any transforms). May optionally include "logmass" (shape (1,)).	required

Returns:

Type	Description
dict[str, Array]	Keyed by obs.name for each observation in self.observations.

Source code in ceridwen/model/model.py

def predict(self, theta: dict[str, Array]) -> dict[str, Array]:
    """
    Project the CSP model spectrum onto all observations.

    If transforms are registered, ``theta`` is treated as the
    *free*-parameter dict; ``apply_transforms`` is called first to
    obtain the CSP-compatible model_theta before forwarding to
    ``csp.predict``.

    Internally calls ``csp.predict(model_theta, observations)`` which
    computes the spectrum once and projects it onto each observation:

    - ``Photometry`` → synthetic AB maggies via filter convolution
    - ``Spectrum``   → model F_ν interpolated onto observed wavelength grid
    - ``Lines``      → Gaussian-aperture integrated line fluxes

    **Mass scaling** — if ``"logmass"`` is present in ``theta``, the
    spectrum is multiplied by ``10 ** logmass`` inside ``csp.predict()``
    *before* projection.  The ``logsfr_ratios_to_sfh`` transform
    normalises the SFH so that the trapezoidal integral of SFR over
    the lookback grid equals 1 M⊙ (Prospector / FSPS convention), so
    this factor sets the physical amplitude for a galaxy with stellar
    mass ``M = 10^logmass`` M⊙.  Scaling once before projection is
    more efficient than scaling each observation separately.

    Parameters
    ----------
    theta : dict[str, Array]
        Free-parameter dict (before any transforms).  May optionally
        include ``"logmass"`` (shape ``(1,)``).

    Returns
    -------
    dict[str, Array]
        Keyed by ``obs.name`` for each observation in ``self.observations``.
    """
    model_theta = self.apply_transforms(theta)
    # Fixed-redshift injection: guarantee the cosmological flux factor
    # (1+z) * (10pc/D_L)^2 is applied inside csp.predict for EVERY
    # caller (mock generation, loglike_fn, predict_jit, predict_vmap).
    # Without this, a theta lacking "zred" silently skips the entire
    # distance normalisation and the returned "maggies" are the raw
    # 10 pc-frame numbers (~6e21 too bright at z = 0.1).  The dict-key
    # check is Python-static, and ``self._zred_fixed`` is a concrete
    # closure constant, so this folds out at JIT trace time — zero
    # cost in the compiled hot path.  A caller-supplied (e.g. sampled)
    # ``zred`` always wins over the fixed value.
    if self._zred_fixed is not None and "zred" not in model_theta:
        if model_theta is theta:          # apply_transforms may not copy
            model_theta = dict(model_theta)
        model_theta["zred"] = self._zred_fixed
    # Mass scaling is handled inside csp.predict() — the spectrum is
    # scaled once before projection, rather than per-observation.
    return self.csp.predict(model_theta, self.observations)

predict_jit

predict_jit(theta)

JIT-compiled version of :meth:predict.

Identical semantics, but the first call triggers XLA compilation and subsequent calls with the same dict structure hit the compiled cache. Use this for interactive evaluation (sanity checks, posterior predictive checks) outside the sampler hot path, where run_sampler already wraps the full log-posterior in @jax.jit.

For vectorised evaluation over many parameter draws, prefer :meth:predict_vmap.

Source code in ceridwen/model/model.py

def predict_jit(self, theta: dict[str, Array]) -> dict[str, Array]:
    """
    JIT-compiled version of :meth:`predict`.

    Identical semantics, but the first call triggers XLA compilation
    and subsequent calls with the same dict structure hit the compiled
    cache.  Use this for interactive evaluation (sanity checks,
    posterior predictive checks) outside the sampler hot path, where
    ``run_sampler`` already wraps the full log-posterior in ``@jax.jit``.

    For vectorised evaluation over many parameter draws, prefer
    :meth:`predict_vmap`.
    """
    # Built once on first access via cached_property (avoids tracing at
    # __init__ time, before observations are set up) and cached on the
    # instance thereafter.
    return self._predict_jit_fn(theta)

predict_vmap

predict_vmap(theta_batch)

Vectorised prediction over a batch of parameter dicts.

Parameters:

Name	Type	Description	Default
theta_batch	dict[str, Array]	Each value has a leading batch dimension, e.g. {"logsfr_ratios": (N, 4), "Z": (N, 1), "logmass": (N, 1)}.	required

Returns:

Type	Description
dict[str, Array]	Each value has a leading batch dimension, e.g. {"optical_spec": (N, n_pix)}.

Source code in ceridwen/model/model.py

def predict_vmap(
    self,
    theta_batch: dict[str, Array],
) -> dict[str, Array]:
    """
    Vectorised prediction over a batch of parameter dicts.

    Parameters
    ----------
    theta_batch : dict[str, Array]
        Each value has a leading batch dimension, e.g.
        ``{"logsfr_ratios": (N, 4), "Z": (N, 1), "logmass": (N, 1)}``.

    Returns
    -------
    dict[str, Array]
        Each value has a leading batch dimension, e.g.
        ``{"optical_spec": (N, n_pix)}``.
    """
    return self._predict_vmap_fn(theta_batch)

ln_prior

ln_prior(theta)

Evaluate the log-prior for all registered free parameters.

For each parameter p in self.priors, computes sum(prior.logpdf(theta[p])) (the sum handles vector-valued parameters such as a non-parametric SFH) and accumulates the total. Parameters absent from self.priors contribute 0 (flat prior).

Parameters:

Name	Type	Description	Default
theta	dict[str, Array]		required

Returns:

Name	Type	Description
lnp	(Array, scalar)

Source code in ceridwen/model/model.py

def ln_prior(self, theta: dict[str, Array]) -> Array:
    """
    Evaluate the log-prior for all registered free parameters.

    For each parameter ``p`` in ``self.priors``, computes
    ``sum(prior.logpdf(theta[p]))`` (the ``sum`` handles vector-valued
    parameters such as a non-parametric SFH) and accumulates the total.
    Parameters absent from ``self.priors`` contribute 0 (flat prior).

    Parameters
    ----------
    theta : dict[str, Array]

    Returns
    -------
    lnp : Array, scalar
    """
    lnp = jnp.zeros(())
    for param_name, prior in self.priors.items():
        if param_name in theta:
            lnp = lnp + jnp.sum(prior.logpdf(theta[param_name]))
    return lnp

log_prob

log_prob(theta)

Alias for ln_prior.

Required by DiagonalGaussianLikelihood.make_lnprobfn and MultiObservationLikelihood.make_lnprobfn, which call prior.log_prob(theta).

Parameters:

Name	Type	Description	Default
theta	dict[str, Array]		required

Returns:

Type	Description
(Array, scalar)

Source code in ceridwen/model/model.py

def log_prob(self, theta: dict[str, Array]) -> Array:
    """
    Alias for ``ln_prior``.

    Required by ``DiagonalGaussianLikelihood.make_lnprobfn`` and
    ``MultiObservationLikelihood.make_lnprobfn``, which call
    ``prior.log_prob(theta)``.

    Parameters
    ----------
    theta : dict[str, Array]

    Returns
    -------
    Array, scalar
    """
    return self.ln_prior(theta)

summary

summary()

Return a multi-line human-readable summary of the model configuration.

Covers: registered free parameters (with shapes and prior types), active transforms (free → derived param mappings), all observation objects, and the CSP physics switch configuration.

Source code in ceridwen/model/model.py

def summary(self) -> str:
    """
    Return a multi-line human-readable summary of the model configuration.

    Covers: registered free parameters (with shapes and prior types),
    active transforms (free → derived param mappings), all observation
    objects, and the CSP physics switch configuration.
    """
    lines = [
        "SedModel",
        "=" * 50,
        f"CSP spectrum model : {self.csp.get_spectrum.__name__}",
        f"Wavelength range   : {float(self.wave.min()):.0f} – "
                               f"{float(self.wave.max()):.0f} Å",
        "",
        "Free Parameters",
        "-" * 40,
    ]
    for name in self.param_names:
        val   = self.theta_init.get(name)
        shape = getattr(val, "shape", "(scalar)")
        prior = self.priors.get(name)
        prior_str = repr(prior) if prior is not None else "flat (no prior)"
        lines.append(f"  {name:<28s}: shape {shape}  |  {prior_str}")

    if self.transforms:
        lines += ["", "Transforms  (free → derived)", "-" * 40]
        for derived, fn in self.transforms.items():
            fn_name = getattr(fn, "__name__", repr(fn))
            lines.append(f"  {derived:<20s} ← {fn_name}")

    if "logmass" in self.param_names:
        lines += [
            "",
            "Mass scaling",
            "-" * 40,
            "  logmass ∈ free params → predicted flux × 10^logmass",
            "  (SFH transform `logsfr_ratios_to_sfh` enforces "
            "∫SFR dt = 1 M⊙;",
            "   logmass therefore equals log10 of the total formed "
            "stellar mass.)",
            "  Convention matches Prospector / FSPS.",
        ]

    lines += ["", "Observations", "-" * 40]
    for obs in self.observations:
        lines.append(f"  {obs!r}")

    return "\n".join(lines)

display

display(ax=None, figsize=None, return_fig=False)

Draw a publication-quality probabilistic graphical model (PGM) diagram.

The diagram follows standard PGM conventions:

• Open circles — stochastic latent variables (free parameters θᵢ)
• Stacked circles — vector-valued parameters (e.g. SFH weight vector)
• Double-bordered rectangle — deterministic SED computation f_ν(λ)
• Coloured rectangles — observation projection operators
• Filled circles — observed data (shaded = conditioned upon)
• Dashed arrows — prior ↦ parameter dependency (ε ≡ stochastic edge)
• Solid arrows — deterministic dependency (θ → f_SED → ŷ → y)

The figure adapts dynamically to however many parameters and observations are registered, making it immediately suitable for inclusion in a paper.

Parameters:

Name	Type	Description	Default
ax	Axes	Axes to draw into. If None a new figure is created.	None
figsize	(float, float)	Figure size in inches (width, height). Defaults scale automatically with the number of free parameters.	None
return_fig	bool	If True return (fig, ax); otherwise call plt.show() and return None.	False

Returns:

Type	Description
(fig, ax) or None	Only returned when return_fig=True.

Examples:

>>> model.display(return_fig=True)[0].savefig("pgm.pdf", dpi=300)

Source code in ceridwen/model/model.py

def display(
    self,
    ax=None,
    figsize: tuple[float, float] | None = None,
    return_fig: bool = False,
):
    """
    Draw a publication-quality probabilistic graphical model (PGM) diagram.

    The diagram follows standard PGM conventions:

    - Open circles       — stochastic latent variables (free parameters θᵢ)
    - Stacked circles    — vector-valued parameters (e.g. SFH weight vector)
    - Double-bordered rectangle — deterministic SED computation f_ν(λ)
    - Coloured rectangles — observation projection operators
    - Filled circles     — observed data (shaded = conditioned upon)
    - Dashed arrows      — prior ↦ parameter dependency (ε ≡ stochastic edge)
    - Solid arrows       — deterministic dependency (θ → f_SED → ŷ → y)

    The figure adapts dynamically to however many parameters and
    observations are registered, making it immediately suitable for
    inclusion in a paper.

    Parameters
    ----------
    ax : matplotlib.axes.Axes, optional
        Axes to draw into.  If *None* a new figure is created.
    figsize : (float, float), optional
        Figure size in inches ``(width, height)``.  Defaults scale
        automatically with the number of free parameters.
    return_fig : bool, optional
        If *True* return ``(fig, ax)``; otherwise call ``plt.show()``
        and return *None*.

    Returns
    -------
    (fig, ax) or None
        Only returned when ``return_fig=True``.

    Examples
    --------
    >>> model.display(return_fig=True)[0].savefig("pgm.pdf", dpi=300)
    """
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from matplotlib.patches import FancyBboxPatch
    import numpy as np

    # ── colour palette ─────────────────────────────────────────────────
    C = dict(
        bg        = "#FFFFFF",
        # parameter nodes
        param_fc  = "#FFFFFF",
        param_ec  = "#222222",
        vec_fc    = "#EEF3FF",
        vec_ec    = "#556BBB",
        # SED deterministic node
        sed_fc    = "#E6F2FB",
        sed_ec    = "#1A6098",
        # observation type nodes
        phot_fc   = "#FFF4E6",  phot_ec = "#C95800",
        spec_fc   = "#EDFAED",  spec_ec = "#276929",
        line_fc   = "#F5EEFF",  line_ec = "#6A22A8",
        # observed data nodes (filled = conditioned on)
        data_fc   = "#37474F",
        data_ec   = "#1A252B",
        data_tc   = "#FFFFFF",
        # arrows
        arr_prior = "#BBBBBB",
        arr_fwd   = "#555555",
        arr_obs   = "#777777",
    )

    def _obs_colors(obs):
        # Prefer isinstance (handles subclasses); fall back to class name.
        try:
            if isinstance(obs, Photometry):
                return (C["phot_fc"], C["phot_ec"])
            if isinstance(obs, Spectrum):
                return (C["spec_fc"], C["spec_ec"])
            if isinstance(obs, Lines):
                return (C["line_fc"], C["line_ec"])
        except Exception:
            pass
        return {
            "Photometry": (C["phot_fc"], C["phot_ec"]),
            "Spectrum":   (C["spec_fc"], C["spec_ec"]),
            "Lines":      (C["line_fc"], C["line_ec"]),
        }.get(type(obs).__name__, (C["param_fc"], C["param_ec"]))

    # ── label helpers ──────────────────────────────────────────────────
    _LATEX = {
        "sfh":         r"$\mathbf{w}_\mathrm{SFH}$",
        "logzsol":     r"$\log Z_\star/Z_\odot$",
        "Z":           r"$Z$",
        "zred":        r"$z$",
        "tau_dust":    r"$\hat{\tau}$",
        "tau_1":       r"$\hat{\tau}_1$",
        "tau_2":       r"$\hat{\tau}_2$",
        "dust_index":  r"$\delta_\mathrm{dust}$",
        "dust_ratio":  r"$f_\mathrm{dust}$",
        "gas_logz":    r"$\log Z_\mathrm{neb}$",
        "gas_logu":    r"$\log U$",
        "sigma_v":     r"$\sigma_v$",
        "f_agn":       r"$f_\mathrm{AGN}$",
        "agn_tau":     r"$\tau_\mathrm{AGN}$",
        "duste_qpah":  r"$q_\mathrm{PAH}$",
        "duste_umin":  r"$U_\mathrm{min}$",
        "duste_gamma": r"$\gamma_e$",
        "mass":        r"$\log M_\star$",
        "logmass":     r"$\log_{10}\,M_\star$",
    }

    def _param_label(name: str) -> str:
        if name in _LATEX:
            return _LATEX[name]
        for k, v in _LATEX.items():
            if k in name:
                return v
        safe = name.replace("_", r"\_")
        return rf"$\theta_{{\mathrm{{{safe}}}}}$"

    def _prior_label(prior) -> str:
        if prior is None:
            return "flat"
        cls = type(prior).__name__
        p   = prior.params
        try:
            if cls in ("Uniform", "TopHat"):
                lo = float(p["low"]);  hi = float(p["high"])
                return rf"$\mathcal{{U}}({lo:.3g},\,{hi:.3g})$"
            if cls == "Normal":
                mu = float(p["mean"]); sg = float(p["sigma"])
                return rf"$\mathcal{{N}}({mu:.3g},\,{sg:.3g})$"
            if cls == "ClippedNormal":
                mu = float(p["mean"]); sg = float(p["sigma"])
                return rf"$\mathcal{{N}}_c({mu:.3g},\,{sg:.3g})$"
            if cls == "LogNormal":
                return r"$\mathrm{LogNorm}$"
            if cls == "StudentT":
                return r"$\mathrm{Student}\text{-}t$"
            if "Multivariate" in cls:
                d = int(p["mean"].shape[0])
                return rf"$\mathcal{{N}}_{{{d}d}}$"
        except Exception:
            pass
        return r"$p(\theta)$"

    def _is_vector(name: str) -> bool:
        val = self.theta_init.get(name)
        if val is None:
            return False
        shape = getattr(val, "shape", ())
        return bool(shape) and shape[0] > 1

    def _shape_str(name: str) -> str:
        val = self.theta_init.get(name)
        if val is None:
            return ""
        shape = getattr(val, "shape", ())
        if not shape or (len(shape) == 1 and shape[0] == 1):
            return ""
        if len(shape) == 1:
            return rf"$\times\,{shape[0]}$"
        return str(shape)

    def _obs_dim(obs) -> str:
        try:
            n = int(jnp.size(obs.flux))
            if isinstance(obs, Photometry):
                return rf"$n_{{\mathrm{{filt}}}}={n}$"
            if isinstance(obs, Spectrum):
                return rf"$n_{{\mathrm{{pix}}}}={n}$"
            if isinstance(obs, Lines):
                return rf"$n_{{\mathrm{{lines}}}}={n}$"
            # fallback: match by class name substring
            cls = type(obs).__name__
            if "Phot" in cls:
                return rf"$n_{{\mathrm{{filt}}}}={n}$"
            if "Spec" in cls:
                return rf"$n_{{\mathrm{{pix}}}}={n}$"
            if "Line" in cls:
                return rf"$n_{{\mathrm{{lines}}}}={n}$"
            return rf"$n={n}$"
        except Exception:
            pass
        return rf"$\hat{{y}}$"

    # ── transform colour ───────────────────────────────────────────────
    C["tr_fc"] = "#FFF8E1"   # warm amber fill
    C["tr_ec"] = "#E65100"   # deep orange border
    C["arr_tr"] = "#E65100"  # transform arrows

    # ── figure geometry ────────────────────────────────────────────────
    n_p = len(self.param_names)
    n_o = len(self.observations)
    has_transforms = bool(self.transforms)

    # 1 data-unit ≡ 1 inch when aspect='equal'
    fw = max(9.0, n_p * 1.10 + 2.0)
    fh = 8.2 if has_transforms else 7.4
    if figsize is not None:
        fw, fh = figsize

    if ax is None:
        fig = plt.figure(figsize=(fw, fh), facecolor=C["bg"])
        ax  = fig.add_axes([0.01, 0.01, 0.98, 0.98],
                           facecolor=C["bg"])
        _created = True
    else:
        fig = ax.get_figure()
        _created = False

    ax.set_xlim(0, fw)
    ax.set_ylim(0, fh)
    ax.set_aspect("equal", adjustable="box")
    ax.axis("off")

    # y levels — shift everything up by 0.8 when transforms are present
    _tr_shift = 0.8 if has_transforms else 0.0
    y_prior = fh - 0.80
    y_param = fh - 2.05
    y_trans = y_param - 1.15   # transform node row (only used when transforms exist)
    y_sed   = fh / 2.0 + 0.10 + (_tr_shift / 2)
    y_obs   = 1.90
    y_data  = 0.62

    r_p  = 0.34   # scalar-param radius
    r_v  = 0.36   # vector-param radius
    r_d  = 0.30   # data-node radius

    # x positions of parameters (spread across 85% of figure width)
    x0, x1 = fw * 0.07, fw * 0.93
    x_p = ([fw / 2] if n_p == 1
            else list(np.linspace(x0, x1, n_p)))

    # x positions of observation nodes
    ox0, ox1 = fw * 0.15, fw * 0.85
    x_o = ([fw / 2]           if n_o == 1
           else [fw*0.33, fw*0.67] if n_o == 2
           else list(np.linspace(ox0, ox1, n_o)))

    x_sed = fw / 2.0

    # ── drawing primitives ─────────────────────────────────────────────
    def circle(x, y, r, fc, ec, lw=1.3, zorder=3, ls="-", alpha=1.0):
        ax.add_patch(mpatches.Circle(
            (x, y), r, facecolor=fc, edgecolor=ec,
            linewidth=lw, zorder=zorder, linestyle=ls, alpha=alpha,
        ))

    def rect(x, y, w, h, fc, ec, lw=1.6, zorder=3, rr=0.12):
        ax.add_patch(FancyBboxPatch(
            (x - w / 2, y - h / 2), w, h,
            boxstyle=f"round,pad=0,rounding_size={rr}",
            facecolor=fc, edgecolor=ec,
            linewidth=lw, zorder=zorder,
        ))

    def arrow(x1, y1, x2, y2, color, lw=0.9,
              style="->", rad=0.0, ls="solid", zorder=2):
        ax.annotate(
            "", xy=(x2, y2), xytext=(x1, y1),
            arrowprops=dict(
                arrowstyle=style, color=color, lw=lw,
                connectionstyle=f"arc3,rad={rad}",
                linestyle=ls,
            ),
            zorder=zorder,
        )

    def txt(x, y, s, ha="center", va="center",
            fs=9, color="#222222", weight="normal",
            style="normal", zorder=6, **kw):
        ax.text(x, y, s, ha=ha, va=va, fontsize=fs, color=color,
                fontweight=weight, fontstyle=style,
                zorder=zorder, **kw)

    # ── 1 ·  prior labels ──────────────────────────────────────────────
    for i, pname in enumerate(self.param_names):
        xp    = x_p[i]
        prior = self.priors.get(pname)
        plbl  = _prior_label(prior)
        txt(xp, y_prior, plbl, fs=7.5, color="#444444",
            style="italic" if prior is None else "normal")
        # dashed stochastic edge: prior distribution → parameter
        arrow(xp, y_prior - 0.16,
              xp, y_param + (r_v if _is_vector(pname) else r_p) + 0.05,
              C["arr_prior"], lw=0.75, ls="dashed")

    # ── 2 ·  parameter nodes ───────────────────────────────────────────
    for i, pname in enumerate(self.param_names):
        xp  = x_p[i]
        vec = _is_vector(pname)

        if vec:
            # stacked-card visual: two offset circles
            circle(xp + 0.06, y_param - 0.06, r_v,
                   fc=C["vec_fc"], ec=C["vec_ec"],
                   lw=0.7, zorder=3, alpha=0.7)
            circle(xp, y_param, r_v,
                   fc=C["vec_fc"], ec=C["vec_ec"],
                   lw=1.3, zorder=4)
            # dimension annotation
            ds = _shape_str(pname)
            if ds:
                txt(xp + r_v + 0.07, y_param + r_v - 0.08,
                    ds, fs=6.5, color=C["vec_ec"], ha="left",
                    style="italic")
        else:
            circle(xp, y_param, r_p,
                   fc=C["param_fc"], ec=C["param_ec"], lw=1.3)

        # parameter symbol inside circle
        lbl = _param_label(pname)
        txt(xp, y_param, lbl, fs=8.5 if not vec else 8.0,
            weight="bold", color="#111111")

    # ── 2b ·  transform nodes (deterministic diamonds) ────────────────
    if has_transforms:
        n_tr    = len(self.transforms)
        tr_x0   = fw * 0.15
        tr_x1   = fw * 0.85
        x_tr    = ([fw / 2] if n_tr == 1
                   else list(np.linspace(tr_x0, tr_x1, n_tr)))
        tr_w, tr_h = 1.70, 0.52

        for j, (derived_name, fn) in enumerate(self.transforms.items()):
            xt  = x_tr[j]
            fn_name = getattr(fn, "__name__", "fn")

            # Diamond shape approximated via rotated rectangle (FancyBboxPatch)
            ax.add_patch(FancyBboxPatch(
                (xt - tr_w / 2, y_trans - tr_h / 2), tr_w, tr_h,
                boxstyle="round,pad=0,rounding_size=0.08",
                facecolor=C["tr_fc"], edgecolor=C["tr_ec"],
                linewidth=1.5, zorder=3, linestyle="--",
            ))
            txt(xt, y_trans + 0.10,
                rf"$\mathtt{{{fn_name}}}$",
                fs=7.5, color=C["tr_ec"], weight="bold")
            txt(xt, y_trans - 0.12,
                rf"$\rightarrow$ {derived_name}",
                fs=6.5, color="#555555", style="italic")

            # Arrows: all free params that feed into this transform → transform node
            # (draw from all free params since we don't know which ones each fn uses)
            for i, pname in enumerate(self.param_names):
                xp  = x_p[i]
                vec = _is_vector(pname)
                r   = r_v if vec else r_p
                arrow(xp, y_param - r - 0.03,
                      xt,  y_trans + tr_h / 2 + 0.05,
                      C["arr_tr"], lw=0.7, ls="dashed")

            # Arrow: transform node → SED node
            arrow(xt, y_trans - tr_h / 2 - 0.04,
                  x_sed, y_sed + (max(3.8, min(fw * 0.42, n_p * 0.88))) / 2 * 0.0 + 0.42,
                  C["tr_ec"], lw=1.0)

    # ── 3 ·  SED computation node ──────────────────────────────────────
    sed_w = max(3.8, min(fw * 0.42, n_p * 0.88))
    sed_h = 0.84

    # outer box
    rect(x_sed, y_sed, sed_w, sed_h,
         C["sed_fc"], C["sed_ec"], lw=2.2)
    # inner double-border (convention for deterministic node)
    rect(x_sed, y_sed, sed_w - 0.13, sed_h - 0.13,
         "none", C["sed_ec"], lw=0.6, zorder=4)

    txt(x_sed, y_sed + 0.18,
        r"$f_\nu(\lambda\,;\,\boldsymbol{\theta})$",
        fs=13, weight="bold", color=C["sed_ec"])

    variant_raw = getattr(
        getattr(self.csp, "get_spectrum", None), "__name__", "get_spectrum"
    )
    _VARIANT_MAP = {
        "dattn_dem_neb":       "dust  ·  dust-em.  ·  neb.",
        "dattn_nodem_neb":     "dust  ·  neb.",
        "dattn_dem_noneb":     "dust  ·  dust-em.",
        "dattn_nodem_noneb":   "dust  (no neb.)",
        "nodattn_nodem_neb":   "neb. only",
        "nodattn_nodem_noneb": "stellar continuum only",
    }
    variant_key  = variant_raw.replace("get_spectrum_", "")
    variant_disp = _VARIANT_MAP.get(
        variant_key, variant_key.replace("_", " · "))
    txt(x_sed, y_sed - 0.20,
        rf"$\mathtt{{get\_spectrum}}(\boldsymbol{{\theta}})$"
        rf"  ·  {variant_disp}",
        fs=7.2, color="#2B5F8A", style="italic")

    # ── 4 ·  arrows: parameters → SED ─────────────────────────────────
    # When transforms are present, free-param arrows go to transform
    # nodes (drawn in section 2b).  Non-transformed params still connect
    # directly to the SED node.
    for i, pname in enumerate(self.param_names):
        if has_transforms:
            # Skip — arrows already drawn in section 2b
            continue
        xp   = x_p[i]
        vec  = _is_vector(pname)
        r    = r_v if vec else r_p
        # fan tip into box proportionally
        xt   = x_sed + (xp - x_sed) * 0.30
        arrow(xp, y_param - r - 0.03,
              xt,  y_sed + sed_h / 2 + 0.04,
              C["arr_fwd"], lw=0.85)

    # ── 5 ·  observation nodes + data nodes ───────────────────────────
    for j, obs in enumerate(self.observations):
        xo       = x_o[j]
        fc, ec  = _obs_colors(obs)
        ow, oh  = 1.60, 0.64

        # canonical type name and projection symbol
        if isinstance(obs, Photometry):
            obs_type_lbl = "Photometry"
            proj_lbl     = r"$\mathbf{T}_\mathrm{filt}\!\cdot\!f_\nu$"
        elif isinstance(obs, Spectrum):
            obs_type_lbl = "Spectrum"
            proj_lbl     = r"$\mathbf{H}\!\cdot\!f_\nu$"
        elif isinstance(obs, Lines):
            obs_type_lbl = "Lines"
            proj_lbl     = r"$\mathbf{W}\!\cdot\!f_\nu$"
        else:
            obs_type_lbl = type(obs).__name__
            proj_lbl     = r"$\hat{y}$"

        txt(xo, y_obs + oh / 2 + 0.26, proj_lbl,
            fs=7.5, color=ec, style="italic")

        # SED → observation arrow
        xt = x_sed + (xo - x_sed) * 0.22
        arrow(xt, y_sed - sed_h / 2 - 0.04,
              xo, y_obs + oh / 2 + 0.05,
              ec, lw=1.1)

        # observation box
        rect(xo, y_obs, ow, oh, fc, ec, lw=1.7)
        txt(xo, y_obs + 0.13, obs_type_lbl,
            fs=8.5, weight="bold", color=ec)
        txt(xo, y_obs - 0.13, obs.name,
            fs=7.0, color="#555555",
            family="monospace")

        # observation → data arrow (with noise annotation)
        arrow(xo, y_obs - oh / 2 - 0.04,
              xo, y_data + r_d + 0.04,
              ec, lw=1.1)

        # noise annotation feeding into data node
        noise_x = xo + r_d + 0.44
        noise_y = y_data + 0.20
        txt(noise_x, noise_y, r"$\sigma_k$",
            fs=8.5, color="#888888")
        arrow(noise_x - 0.06, noise_y - 0.13,
              xo + r_d + 0.04, y_data + 0.05,
              "#BBBBBB", lw=0.75)

        # data node (filled = observed / conditioned upon)
        circle(xo, y_data, r_d,
               fc=C["data_fc"], ec=C["data_ec"], lw=1.5)
        txt(xo, y_data, _obs_dim(obs),
            fs=7.2, color=C["data_tc"], weight="bold")

        # label below data node
        txt(xo, y_data - r_d - 0.26,
            rf"$\mathbf{{y}}_\mathrm{{{obs.name}}}$",
            fs=8, color="#333333", style="italic")

    # ── 6 ·  legend ───────────────────────────────────────────────────
    lg_y  = 0.28
    lg_r  = 0.13
    items = [
        (C["param_fc"], C["param_ec"], r"latent $\theta_i$"),
        (C["vec_fc"],   C["vec_ec"],   r"vector param"),
        (C["data_fc"],  C["data_ec"],  r"observed $y_k$"),
        (C["sed_fc"],   C["sed_ec"],   r"deterministic node"),
    ]
    if has_transforms:
        items.append((C["tr_fc"], C["tr_ec"], r"transform node"))
    n_leg = len(items)
    xs_leg = np.linspace(fw * 0.08, fw * 0.70, n_leg)
    for lx, (lfc, lec, llbl) in zip(xs_leg, items):
        circle(lx, lg_y, lg_r, fc=lfc, ec=lec, lw=1.0, zorder=5)
        txt(lx + 0.22, lg_y, llbl,
            fs=7.5, ha="left", color="#444444")
    # dashed-arrow legend entry
    ax.annotate(
        "", xy=(fw * 0.82, lg_y), xytext=(fw * 0.80, lg_y),
        arrowprops=dict(
            arrowstyle="->", color=C["arr_prior"], lw=0.9,
            linestyle="dashed",
        ),
        zorder=5,
    )
    txt(fw * 0.83, lg_y, r"stochastic edge",
        fs=7.5, ha="left", color="#444444")

    # ── 7 ·  title ────────────────────────────────────────────────────
    n_tr   = len(self.transforms)
    tr_str = (rf"  ·  ${n_tr}$ transform{'s' if n_tr != 1 else ''}"
              if has_transforms else "")
    ax.set_title(
        rf"SedModel — ${n_p}$ free parameters  ·  "
        rf"${n_o}$ observation{'s' if n_o != 1 else ''}{tr_str}",
        fontsize=9.5, color="#333333", pad=3,
    )

    # no tight_layout — axes already fill the figure via add_axes

    if return_fig:
        return fig, ax
    plt.show()
    return None

ceridwen.model.transforms.logsfr_ratios_to_sfh

logsfr_ratios_to_sfh(logsfr_ratios, sfh_times_yr=None)

Convert log-ratios of consecutive SFR bins to a unit-mass SFH weight vector.

Parameters:

Name	Type	Description	Default
logsfr_ratios	(array_like, shape(n - 1))	Log10 ratios of consecutive SFR bins. logsfr_ratios[i] = log10( SFR[i] / SFR[i+1] ). Under the lookback-time convention (index 0 = today, last index = oldest), SFR[0] is the most-recent SFR and SFR[i+1] is at a slightly older lookback time. Positive logsfr_ratios[i] therefore mean SFR[i] > SFR[i+1], i.e. the SFR is higher today than in the past — a late-assembly history; negative values mean an earlier burst with the SFR declining toward the present.	required
sfh_times_yr	(array_like, shape(n))	Lookback-time grid in years (same as CSPBasis.sfh_times). When provided, trapezoidal integration weights are used for normalisation so that the integral of the SFH equals 1 Msun. If None, the discrete sum is used instead.	None

Returns:

Name	Type	Description
sfh	(ndarray, shape(n))	Unit-mass SFH weight vector suitable for theta["sfh"].

Notes

The normalisation enforces sum(sfh * w) = 1 where w are the standard trapezoidal quadrature weights (half-width at boundaries), so that the trapezoidal integral of the SFH over the lookback-time grid equals 1 Msun:

.. math:: \int_0^{t_{\rm univ}} \mathrm{SFR}(t)\,dt = 1\;\mathrm{M_\odot}.

The total stellar mass of the model is then set by theta["logmass"] (applied as a multiplicative 10**logmass inside CSPBasis.predict), matching the Prospector / FSPS convention.

This is consistent with the per-bin mass integral m2 = sfh_mid * dt computed inside CSPBasis.calculate_ssp_weights_const_zh_step and the piecewise-linear integral used by calculate_ssp_weights_const_zh: on a shared node grid the trapezoid sum here and the midpoint sum there are algebraically identical, so no compensating rescale is needed inside the CSP.

JAX compatibility

The function is fully JIT-compatible: every operation is a jnp primitive on traced arrays, the sfh_times_yr is not None branch resolves at trace time (it is a Python-level check on the closure argument, not a runtime decision on a traced value), and there are no data-dependent shapes.

Source code in ceridwen/model/transforms.py

def logsfr_ratios_to_sfh(
    logsfr_ratios,
    sfh_times_yr=None,
):
    """
    Convert log-ratios of consecutive SFR bins to a unit-mass SFH weight
    vector.

    Parameters
    ----------
    logsfr_ratios : array_like, shape (n - 1,)
        Log10 ratios of consecutive SFR bins.
        ``logsfr_ratios[i] = log10( SFR[i] / SFR[i+1] )``.

        Under the lookback-time convention (index 0 =
        today, last index = oldest), ``SFR[0]`` is the most-recent SFR
        and ``SFR[i+1]`` is at a slightly older lookback time.  Positive
        ``logsfr_ratios[i]`` therefore mean ``SFR[i] > SFR[i+1]``, i.e.
        the SFR is *higher today than in the past* — a late-assembly
        history; negative values mean an earlier burst with the SFR
        declining toward the present.
    sfh_times_yr : array_like, shape (n,), optional
        Lookback-time grid in years (same as ``CSPBasis.sfh_times``).
        When provided, trapezoidal integration weights are used for
        normalisation so that the integral of the SFH equals 1 Msun.
        If *None*, the discrete sum is used instead.

    Returns
    -------
    sfh : jnp.ndarray, shape (n,)
        Unit-mass SFH weight vector suitable for ``theta["sfh"]``.

    Notes
    -----
    The normalisation enforces ``sum(sfh * w) = 1`` where ``w`` are
    the standard trapezoidal quadrature weights (half-width at
    boundaries), so that the trapezoidal integral of the SFH over the
    lookback-time grid equals **1 Msun**:

    .. math::
        \\int_0^{t_{\\rm univ}} \\mathrm{SFR}(t)\\,dt = 1\\;\\mathrm{M_\\odot}.

    The total stellar mass of the model is then set by
    ``theta["logmass"]`` (applied as a multiplicative ``10**logmass``
    inside ``CSPBasis.predict``), matching the Prospector / FSPS
    convention.

    This is consistent with the per-bin mass integral
    ``m2 = sfh_mid * dt`` computed inside
    ``CSPBasis.calculate_ssp_weights_const_zh_step`` and the
    piecewise-linear integral used by
    ``calculate_ssp_weights_const_zh``: on a shared node grid the
    trapezoid sum here and the midpoint sum there are algebraically
    identical, so no compensating rescale is needed inside the CSP.

    JAX compatibility
    -----------------
    The function is fully JIT-compatible: every operation is a
    ``jnp`` primitive on traced arrays, the ``sfh_times_yr is not
    None`` branch resolves at trace time (it is a Python-level check
    on the closure argument, not a runtime decision on a traced
    value), and there are no data-dependent shapes.
    """
    ratios  = jnp.asarray(logsfr_ratios, dtype=float)            # (n-1,)
    # Anchor log10(SFR[0]) = 0, then cumulate the negative ratios
    log_sfr = jnp.concatenate([jnp.zeros(1),
                                -jnp.cumsum(ratios)])             # (n,)
    sfr     = 10.0 ** log_sfr                                     # (n,)

    if sfh_times_yr is not None:
        times = jnp.asarray(sfh_times_yr, dtype=float)            # (n,)
        dt    = jnp.abs(jnp.diff(times))                          # (n-1,)
        # Standard trapezoidal quadrature weights (yr):
        #   w[0]   = 0.5 * dt[0]
        #   w[i]   = 0.5 * (dt[i-1] + dt[i])   for 0 < i < n-1
        #   w[n-1] = 0.5 * dt[n-2]
        w_lo  = jnp.concatenate([jnp.zeros(1), dt])               # (n,)
        w_hi  = jnp.concatenate([dt, jnp.zeros(1)])               # (n,)
        w     = 0.5 * (w_lo + w_hi)                               # (n,)
        # Unit-mass normalisation: ∫SFR dt = sum(sfr * w) = 1 Msun.
        # GOTCHA: normalise to total mass, NOT mean SFR — dividing by
        # sum(w) instead would leave an implicit factor of t_universe[yr]
        # (~1.4e10) in the spectrum and bias every logmass estimate by
        # ~10 dex at z=0 (less at higher z).
        total_mass = jnp.sum(sfr * w)
        sfh   = sfr / total_mass
    else:
        # Discrete fallback: sum(sfh) = 1.
        sfh = sfr / jnp.sum(sfr)

    return sfh

Priors

ceridwen.priors

Prior distributions for CERIDWEN models.

A clean, discoverable import path for the priors (implemented in ceridwen.sampler.priors)::

from ceridwen.priors import Uniform, Normal, ClippedNormal, LogNormal, StudentT

Each prior exposes log_prob, sample, unit_transform (inverse-CDF, used by nested sampling) and inverse_unit_transform (CDF), and is a JAX PyTree so it can flow through jit/grad.

Prior dataclass

Prior(parnames=(), name='', **kwargs)

Bases: ABC

JAX-friendly prior base class that delegates all probability operations to a TFP-JAX distribution. Subclasses must implement tfp_dist().

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

tfp_dist abstractmethod

tfp_dist()

Return a TFP-JAX distribution object built from self.params.

Must be implemented by subclasses.

Source code in ceridwen/sampler/priors.py

@abc.abstractmethod
def tfp_dist(self) -> tfd.Distribution:
    """
    Return a TFP-JAX distribution object built from `self.params`.

    Must be implemented by subclasses.
    """
    raise NotImplementedError

Uniform dataclass

Uniform(parnames=(), name='', **kwargs)

Bases: Prior

Uniform distribution on [low, high].

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

TopHat dataclass

TopHat(parnames=(), name='', **kwargs)

Bases: Uniform

Uniform distribution between two bounds, renamed for backwards compatibility :param low: Minimum of the distribution

:param high: Maximum of the distribution

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

Normal dataclass

Normal(parnames=(), name='', **kwargs)

Bases: Prior

A simple gaussian prior.

:param mean: Mean of the distribution

:param sigma: Standard deviation of the distribution

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

ClippedNormal dataclass

ClippedNormal(parnames=(), name='', **kwargs)

Bases: Prior

A Gaussian prior clipped to some range.

:param mean: Mean of the normal distribution

:param sigma: Standard deviation of the normal distribution

:param low: Minimum of the distribution

:param high: Maximum of the distribution

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

LogNormal dataclass

LogNormal(parnames=(), name='', **kwargs)

Bases: Prior

A log-normal prior, where the natural log of the variable is distributed normally. Useful for parameters that cannot be less than zero.

Note that LogNormal(np.exp(mode) / f) == LogNormal(np.exp(mode) * f) and f = np.exp(sigma) corresponds to "one sigma" from the peak.

:param mode: Natural log of the variable value at which the probability density is highest.

:param sigma: Standard deviation of the distribution of the natural log of the variable.

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

StudentT dataclass

StudentT(parnames=(), name='', **kwargs)

Bases: Prior

A Student's T distribution

:param mean: Mean of the distribution

:param scale: Size of the distribution, analogous to the standard deviation

:param df: Number of degrees of freedom

Source code in ceridwen/sampler/priors.py

def __init__(self,
             parnames: Sequence[str] = (),

             name: str = "",
             **kwargs: Any):

    prior_params: Tuple[str, ...] = getattr(type(self), "prior_params", None)
    if prior_params is None:
        raise ValueError(
            f"{type(self).__name__} must define class attribute "
            "`prior_params = (\"param1\", \"param2\", ...)`."
        )

    if not parnames:
        parnames = prior_params

    if len(parnames) != len(prior_params):
        raise ValueError(
            f"parnames has length {len(parnames)} but prior_params has "
            f"length {len(prior_params)}"
        )

    alias = dict(zip(prior_params, parnames))

    params: Dict[str, Array] = {}
    for intrinsic, external in alias.items():
        if external in kwargs:
            params[intrinsic] = jnp.asarray(kwargs[external])

    object.__setattr__(self, "alias", alias)
    object.__setattr__(self, "params", params)
    object.__setattr__(self, "name", name)

Likelihood

ceridwen.likelihood.DiagonalGaussianLikelihood dataclass

DiagonalGaussianLikelihood(
    noise_model=DiagonalNoiseModel(),
)

Bases: LikelihoodBase

Gaussian log-likelihood with an independent (diagonal) noise model.

This class sequences three operations:

1. Call noise_model.compute(sigma_obs, mu, mask, params) to obtain per-datum inv_var and log_det.
2. Call lnlike_diag_gaussian(y, mu, inv_var, log_det, mask) to get the scalar log-likelihood and diagnostics.

Because both steps are pure JAX functions, the entire __call__ is JIT-compilable and differentiable end-to-end.

Parameters:

Name	Type	Description	Default
noise_model	DiagonalNoiseModel	Noise model instance. Defaults to a plain observational-uncertainty- only model (no jitter, no fractional error).	DiagonalNoiseModel()

Examples:

Basic usage with no nuisance parameters::

lhood = DiagonalGaussianLikelihood()
lnl, aux = lhood(y, mu, sigma_obs, mask)
# aux.chi      -- normalised residuals
# aux.lnl_pointwise  -- per-datum contributions

With jitter as a sampled parameter::

lhood = DiagonalGaussianLikelihood(
    noise_model=DiagonalNoiseModel(use_jitter=True)
)
lnl, aux = lhood(y, mu, sigma_obs, mask,
                 params={"log_jitter": theta["log_jitter"]})

Gradient of log-likelihood w.r.t. theta (for HMC/gradient-based samplers)::

def lnl_fn(theta):
    mu = model.predict(theta)
    return lhood(y, mu, sigma_obs, mask, params=theta)[0]

grad = jax.grad(lnl_fn)(theta)

Using has_aux=True to get diagnostics and gradient in one pass::

(lnl, aux), grad = jax.value_and_grad(
    lambda t: lhood(y, model.predict(t), sigma_obs, mask, t),
    has_aux=True,
)(theta)

__call__

__call__(y, mu, sigma_obs, mask, params=None)

Evaluate log-likelihood and return diagnostics.

This method is the intended hot path for sampling. It is safe inside jax.jit, jax.grad, and jax.vmap.

Parameters:

Name	Type	Description	Default
y	(Array, shape(n_data))		required
mu	(Array, shape(n_data))		required
sigma_obs	(Array, shape(n_data))		required
mask	Array of bool, shape (n_data,)		required
params	dict	Nuisance parameters expected by the noise model.	None

Returns:

Name	Type	Description
lnl_total	(Array, scalar)
aux	LikelihoodOutput

Source code in ceridwen/likelihood/likelihood.py

def __call__(
    self,
    y         : Array,
    mu        : Array,
    sigma_obs : Array,
    mask      : Array,
    params    : Optional[dict[str, Array]] = None,
) -> tuple[Array, LikelihoodOutput]:
    """
    Evaluate log-likelihood and return diagnostics.

    This method is the intended hot path for sampling.  It is safe inside
    ``jax.jit``, ``jax.grad``, and ``jax.vmap``.

    Parameters
    ----------
    y : Array, shape (n_data,)
    mu : Array, shape (n_data,)
    sigma_obs : Array, shape (n_data,)
    mask : Array of bool, shape (n_data,)
    params : dict, optional
        Nuisance parameters expected by the noise model.

    Returns
    -------
    lnl_total : Array, scalar
    aux : LikelihoodOutput
    """
    noise_out: NoiseModelOutput = self.noise_model.compute(
        sigma_obs, mu, mask, params
    )
    return lnlike_diag_gaussian(
        y, mu, noise_out.inv_var, noise_out.log_det, mask
    )

make_lnprobfn

make_lnprobfn(observations, model, prior)

Build a JIT-compiled log-posterior for a single observation type.

observations must have attributes .flux, .uncertainty, and .mask (the interface defined by ceridwen.observation.Observation).

The returned function has the exact signature expected by blackjax::

kernel = blackjax.nuts(lnprobfn, step_size)
state  = kernel.init(initial_theta)

Parameters:

Name	Type	Description	Default
observations	Observation	A single observation object (photometry or spectroscopy or lines). For multiple modalities use MultiObservationLikelihood.	required
model	object	Must implement predict(theta) -> Array.	required
prior	object	Must implement log_prob(theta) -> Array.	required

Returns:

Type	Description
Callable[[theta], Array]	JIT-compiled log-posterior.

Notes

observations, model, prior, and self are all closed over at factory time and become compile-time constants. Only theta is traced. This means the first call incurs a one-time XLA compilation cost; subsequent calls are fully compiled and have minimal Python overhead.

Source code in ceridwen/likelihood/likelihood.py

def make_lnprobfn(
    self,
    observations : Any,
    model        : Any,
    prior        : Any,
) -> Callable[[dict[str, Array]], Array]:
    """
    Build a JIT-compiled log-posterior for a single observation type.

    ``observations`` must have attributes ``.flux``, ``.uncertainty``, and
    ``.mask`` (the interface defined by ``ceridwen.observation.Observation``).

    The returned function has the exact signature expected by blackjax::

        kernel = blackjax.nuts(lnprobfn, step_size)
        state  = kernel.init(initial_theta)

    Parameters
    ----------
    observations : Observation
        A single observation object (photometry *or* spectroscopy *or*
        lines).  For multiple modalities use ``MultiObservationLikelihood``.
    model : object
        Must implement ``predict(theta) -> Array``.
    prior : object
        Must implement ``log_prob(theta) -> Array``.

    Returns
    -------
    Callable[[theta], Array]
        JIT-compiled log-posterior.

    Notes
    -----
    ``observations``, ``model``, ``prior``, and ``self`` are all closed
    over at factory time and become compile-time constants.  Only ``theta``
    is traced.  This means the first call incurs a one-time XLA compilation
    cost; subsequent calls are fully compiled and have minimal Python
    overhead.
    """
    # Hoist static data out of the closure so JAX sees them as constants.
    y         : Array = observations.flux
    sigma_obs : Array = observations.uncertainty
    mask      : Array = observations.mask
    noise_model       = self.noise_model

    @jax.jit
    def lnprobfn(theta: dict[str, Array]) -> Array:
        mu = model.predict(theta)
        # Pass the full theta as params.  The noise model extracts only
        # the keys it needs (log_jitter, log_f_calib); all other keys are
        # ignored.  This avoids conditional dict construction inside the
        # JIT-compiled hot path and is safe for any noise model subclass.
        noise_out = noise_model.compute(sigma_obs, mu, mask, theta)
        lnl, _    = lnlike_diag_gaussian(
            y, mu, noise_out.inv_var, noise_out.log_det, mask
        )
        lnp = prior.log_prob(theta)
        return lnl + lnp

    return lnprobfn

ceridwen.likelihood.MultiObservationLikelihood dataclass

MultiObservationLikelihood(
    keys=tuple(), likelihoods=tuple()
)

Bases: LikelihoodBase

Combine log-likelihoods across multiple independent data modalities.

In SED fitting, the posterior typically conditions on several observation types simultaneously -- broadband photometry (in maggies), a spectrum (in F_lambda), and emission-line fluxes. Each has different units, different noise properties, and potentially different noise model nuisance parameters.

This class owns a static mapping from string keys to likelihood objects. At trace time the Python loop over keys is unrolled by XLA; there is no runtime dispatch overhead.

Parameters:

Name	Type	Description	Default
keys	tuple of str	Ordered observation keys, e.g. ("phot", "spec", "lines").	tuple()
likelihoods	tuple of LikelihoodBase	One likelihood per key, in the same order. Different elements may have different noise models (e.g. photometry with fractional error, spectroscopy with jitter).	tuple()

Examples:

Combine photometry and spectroscopy::

multi = MultiObservationLikelihood(
    keys=("phot", "spec"),
    likelihoods=(
        DiagonalGaussianLikelihood(
            DiagonalNoiseModel(use_fractional=True)
        ),
        DiagonalGaussianLikelihood(
            DiagonalNoiseModel(use_jitter=True)
        ),
    ),
)
lnprobfn = multi.make_lnprobfn(observations, model, prior)

where observations is a dict {"phot": phot_obs, "spec": spec_obs} and model.predict(theta) returns a dict with the same keys.

__call__

__call__(y, mu, sigma_obs, mask, params=None)

Evaluate total log-likelihood across all modalities.

Parameters:

Name	Type	Description	Default
y	dict[str, Array]	Per-modality data arrays, keyed by the same strings as self.keys.	required
mu	dict[str, Array]	Per-modality data arrays, keyed by the same strings as self.keys.	required
sigma_obs	dict[str, Array]	Per-modality data arrays, keyed by the same strings as self.keys.	required
mask	dict[str, Array]	Per-modality data arrays, keyed by the same strings as self.keys.	required
params	dict[str, Array]	All nuisance parameters (shared across modalities). Each likelihood/noise model extracts only the keys it needs.	None

Returns:

Name	Type	Description
lnl_total	(Array, scalar)	Sum of log-likelihoods across all modalities.
aux	dict[str, LikelihoodOutput]	Per-modality diagnostics.

Source code in ceridwen/likelihood/likelihood.py

def __call__(
    self,
    y         : dict[str, Array],
    mu        : dict[str, Array],
    sigma_obs : dict[str, Array],
    mask      : dict[str, Array],
    params    : Optional[dict[str, Array]] = None,
) -> tuple[Array, dict[str, LikelihoodOutput]]:
    """
    Evaluate total log-likelihood across all modalities.

    Parameters
    ----------
    y, mu, sigma_obs, mask : dict[str, Array]
        Per-modality data arrays, keyed by the same strings as ``self.keys``.
    params : dict[str, Array], optional
        All nuisance parameters (shared across modalities).  Each
        likelihood/noise model extracts only the keys it needs.

    Returns
    -------
    lnl_total : Array, scalar
        Sum of log-likelihoods across all modalities.
    aux : dict[str, LikelihoodOutput]
        Per-modality diagnostics.
    """
    lnl_total = jnp.zeros(())
    aux: dict[str, LikelihoodOutput] = {}
    for key, lhood in zip(self.keys, self.likelihoods):
        lnl_i, aux_i = lhood(y[key], mu[key], sigma_obs[key], mask[key], params)
        lnl_total    = lnl_total + lnl_i
        aux[key]     = aux_i
    return lnl_total, aux

make_lnprobfn

make_lnprobfn(observations, model, prior)

Build a JIT-compiled log-posterior for multiple observation types.

Parameters:

Name	Type	Description	Default
observations	dict[str, Observation]	Keyed by the same strings as self.keys. Each value must have .flux, .uncertainty, and .mask attributes.	required
model	object	Must implement predict(theta) -> dict[str, Array].	required
prior	object	Must implement log_prob(theta) -> Array.	required

Returns:

Type	Description
Callable[[theta], Array]	JIT-compiled log-posterior, suitable for blackjax.

Source code in ceridwen/likelihood/likelihood.py

def make_lnprobfn(
    self,
    observations : dict[str, Any],
    model        : Any,
    prior        : Any,
) -> Callable[[dict[str, Array]], Array]:
    """
    Build a JIT-compiled log-posterior for multiple observation types.

    Parameters
    ----------
    observations : dict[str, Observation]
        Keyed by the same strings as ``self.keys``.  Each value must have
        ``.flux``, ``.uncertainty``, and ``.mask`` attributes.
    model : object
        Must implement ``predict(theta) -> dict[str, Array]``.
    prior : object
        Must implement ``log_prob(theta) -> Array``.

    Returns
    -------
    Callable[[theta], Array]
        JIT-compiled log-posterior, suitable for blackjax.
    """
    # Pre-extract static data arrays from observations to avoid attribute
    # lookups inside the JIT-compiled hot path.
    static_data: dict[str, tuple[Array, Array, Array]] = {
        key: (
            observations[key].flux,
            observations[key].uncertainty,
            observations[key].mask,
        )
        for key in self.keys
    }
    keys        = self.keys
    likelihoods = self.likelihoods

    @jax.jit
    def lnprobfn(theta: dict[str, Array]) -> Array:
        predictions: dict[str, Array] = model.predict(theta)
        lnl = jnp.zeros(())

        for key, lhood in zip(keys, likelihoods):
            y_k, sig_k, mask_k = static_data[key]
            mu_k = predictions[key]
            # Delegate entirely to each likelihood's __call__, which
            # handles noise model dispatch internally.  Pass the full
            # theta as params; each noise model silently ignores keys
            # it does not recognise.  This respects the LikelihoodBase
            # interface and works for any future subclass.
            lnl_k, _ = lhood(y_k, mu_k, sig_k, mask_k, params=theta)
            lnl = lnl + lnl_k

        lnp = prior.log_prob(theta)
        return lnl + lnp

    return lnprobfn

Fitting

ceridwen.fit.fitSED

fitSED(
    model,
    observations=None,
    output_dir=".",
    *,
    sampler="nested",
    rng_key=None,
    sampler_kwargs=None,
    vi=None,
    vi_kwargs=None,
    filename="ceridwen_result.h5",
    overwrite=True,
    verbose=True
)

Fit an SED model to observations and write results to HDF5.

This is the main user-facing function. It constructs the likelihood, configures and runs the requested sampler, and writes the posterior samples, model metadata, and observation data to disk.

Parameters:

Name	Type	Description	Default
model	SedModel	A fully configured SedModel instance. If observations is not None, model.observations is replaced before fitting.	required
observations	list[Observation]	Observations to fit. If None, uses model.observations. Passing observations here triggers obs.setup_for_model() automatically.	None
output_dir	str or Path	Directory for the output HDF5 file. Created if it does not exist.	'.'
sampler	str	"nested" (default) for BlackJAX nested sampling, or "nuts" for BlackJAX NUTS HMC.	'nested'
rng_key	Array	JAX PRNG key. Defaults to jax.random.PRNGKey(0).	None
sampler_kwargs	dict	Extra keyword arguments forwarded to the sampler adapter constructor (e.g. num_warmup, num_samples, num_chains, dense_mass, etc.). Any argument accepted by BlackJAXNestedSamplerAdapter or BlackJAXNUTSAdapter can be passed here.	None
vi	None, str, VariationalMap, or TrainedMap	Variational-inference preconditioning (only used when sampler='nuts'). See :class:BlackJAXNUTSAdapter for accepted values. In short: None (default): plain window-adaptation NUTS. 'tril': train a full-rank Gaussian transport map (paper's TriL baseline). 'iaf': train a stacked inverse-autoregressive-flow map (paper's NeuTra neural transport). When a map is supplied, NUTS samples in the whitened z-space, giving dramatically shorter warmup. See Hoffman et al. 2019, arXiv:1903.03704.	None
vi_kwargs	dict	Forwarded to :func:train_vi (e.g. num_steps=1500, batch_size=16, lr0=1e-2) and/or the VI map constructor (e.g. init_scale=0.1 for TriL, n_flows=3 for IAF).	None
filename	str	Name of the output HDF5 file. Default ceridwen_result.h5. A plain-text log with the same stem (ceridwen_result.log) is written alongside it: timestamped device/backend info, sampler configuration, timings, and the result summary. It is written regardless of verbose (which only controls console echo).	'ceridwen_result.h5'
overwrite	bool	If True (default), overwrite an existing file.	True
verbose	bool	Print progress to the console. Default True.	True

Returns:

Type	Description
SamplingResult	The sampling result object (same as returned by run_sampler). The HDF5 file and the .log file are written as side effects.

Source code in ceridwen/fit.py

def fitSED(
    model,
    observations: Sequence | None = None,
    output_dir: str | Path = ".",
    *,
    sampler: str = "nested",
    rng_key: Array | None = None,
    sampler_kwargs: dict[str, Any] | None = None,
    vi: "str | Any | None" = None,
    vi_kwargs: dict[str, Any] | None = None,
    filename: str = "ceridwen_result.h5",
    overwrite: bool = True,
    verbose: bool = True,
):
    """
    Fit an SED model to observations and write results to HDF5.

    This is the main user-facing function.  It constructs the likelihood,
    configures and runs the requested sampler, and writes the posterior
    samples, model metadata, and observation data to disk.

    Parameters
    ----------
    model : SedModel
        A fully configured ``SedModel`` instance.  If ``observations``
        is not None, ``model.observations`` is replaced before fitting.
    observations : list[Observation], optional
        Observations to fit.  If None, uses ``model.observations``.
        Passing observations here triggers ``obs.setup_for_model()``
        automatically.
    output_dir : str or Path
        Directory for the output HDF5 file.  Created if it does not
        exist.
    sampler : str
        ``"nested"`` (default) for BlackJAX nested sampling, or
        ``"nuts"`` for BlackJAX NUTS HMC.
    rng_key : Array, optional
        JAX PRNG key.  Defaults to ``jax.random.PRNGKey(0)``.
    sampler_kwargs : dict, optional
        Extra keyword arguments forwarded to the sampler adapter
        constructor (e.g. ``num_warmup``, ``num_samples``,
        ``num_chains``, ``dense_mass``, etc.).  Any argument accepted
        by ``BlackJAXNestedSamplerAdapter`` or ``BlackJAXNUTSAdapter``
        can be passed here.
    vi : None, str, VariationalMap, or TrainedMap, optional
        Variational-inference preconditioning (only used when
        ``sampler='nuts'``).  See :class:`BlackJAXNUTSAdapter` for
        accepted values.  In short:

        - ``None`` (default): plain window-adaptation NUTS.
        - ``'tril'``: train a full-rank Gaussian transport map
          (paper's TriL baseline).
        - ``'iaf'``: train a stacked inverse-autoregressive-flow map
          (paper's NeuTra neural transport).

        When a map is supplied, NUTS samples in the whitened
        z-space, giving dramatically shorter warmup.  See
        Hoffman et al. 2019, arXiv:1903.03704.

    vi_kwargs : dict, optional
        Forwarded to :func:`train_vi` (e.g. ``num_steps=1500``,
        ``batch_size=16``, ``lr0=1e-2``) and/or the VI map
        constructor (e.g. ``init_scale=0.1`` for TriL, ``n_flows=3``
        for IAF).
    filename : str
        Name of the output HDF5 file.  Default ``ceridwen_result.h5``.
        A plain-text log with the same stem (``ceridwen_result.log``) is
        written alongside it: timestamped device/backend info, sampler
        configuration, timings, and the result summary.  It is written
        regardless of ``verbose`` (which only controls console echo).
    overwrite : bool
        If True (default), overwrite an existing file.
    verbose : bool
        Print progress to the console.  Default True.

    Returns
    -------
    SamplingResult
        The sampling result object (same as returned by ``run_sampler``).
        The HDF5 file and the ``.log`` file are written as side effects.
    """
    from .likelihood.likelihood import (
        DiagonalGaussianLikelihood,
        MultiObservationLikelihood,
    )
    from .sampler.runner import run_sampler

    # ── Validate inputs ───────────────────────────────────────────────
    if rng_key is None:
        rng_key = jax.random.PRNGKey(0)

    if sampler_kwargs is None:
        sampler_kwargs = {}

    output_dir = Path(output_dir)
    output_dir.mkdir(parents=True, exist_ok=True)
    output_path = output_dir / filename

    if output_path.exists() and not overwrite:
        raise FileExistsError(
            f"{output_path} already exists.  Pass overwrite=True to replace."
        )

    # ── Attach observations if provided ───────────────────────────────
    if observations is not None:
        model.observations = list(observations)
        # Forward model.zred so the per-observation projection matrices
        # (Photometry._T, Spectrum._H, Lines._W) are baked at the
        # *observed-frame* wavelength grid (1+z) * wave_rest.  Calling
        # setup_for_model(...) with the default zred=0 here would clobber
        # the correctly-redshifted projection that SedModel.__init__
        # built (model.py L200), so any non-zero fixed redshift would
        # silently degrade to a rest-frame filter integral while
        # CSPBasis.predict still applied flux_factor_maggies(z).  At
        # z ~ 2.7 that produces 10-20 sigma photometric residuals
        # because the filters sample the wrong intrinsic wavelengths.
        for obs in model.observations:
            obs.setup_for_model(model.csp.wave, zred=model.zred)

    if not model.observations:
        raise ValueError("No observations attached to model.")

    # ── Build likelihood ──────────────────────────────────────────────
    _t0_likelihood = time.perf_counter()
    obs_dict = model.obs_dict
    keys = tuple(obs_dict.keys())
    likelihoods = tuple(DiagonalGaussianLikelihood() for _ in keys)
    multi_likelihood = MultiObservationLikelihood(
        keys=keys,
        likelihoods=likelihoods,
    )
    _t_likelihood = time.perf_counter() - _t0_likelihood

    # Route diagnostics through the package logger, with two sinks:
    #   * console (StreamHandler)  -- only when verbose=True (unchanged);
    #   * a per-fit log file next to the HDF5 result -- ALWAYS, so every
    #     fit leaves a text record (device, sampler config, timings,
    #     result summary) even when run non-verbose inside a batch job.
    # NOTE: sampler-internal progress (the "[vi] iter ..." lines, NUTS
    # warmup bars) is written with bare print() inside ceridwen.sampler
    # and is NOT captured here; use `python fit_script.py |& tee run.log`
    # for a full console transcript.
    logger.setLevel(logging.INFO)
    logger.propagate = False
    if verbose and not any(
        isinstance(h, logging.StreamHandler)
        and not isinstance(h, logging.FileHandler)
        for h in logger.handlers
    ):
        _handler = logging.StreamHandler()
        _handler.setFormatter(logging.Formatter("%(message)s"))
        logger.addHandler(_handler)

    log_path = output_path.with_suffix(".log")
    _file_handler = logging.FileHandler(log_path, mode="w")
    _file_handler.setFormatter(logging.Formatter(
        "%(asctime)s  %(message)s", datefmt="%Y-%m-%d %H:%M:%S"))
    logger.addHandler(_file_handler)

    try:
        # Report the XLA backend up front: a fit silently falling back to CPU
        # (missing CUDA jaxlib, JAX_PLATFORMS=cpu leaking from a mock script,
        # driver mismatch) looks identical except for a ~10-100x slowdown.
        _devices = jax.devices()
        _backend = jax.default_backend().upper()   # 'CPU', 'GPU', or 'TPU'
        _device_str = ", ".join(str(d) for d in _devices)

        logger.info(f"ceridwen.fitSED")
        logger.info(f"  Device      : {_backend}  ({_device_str})")
        if _backend == "CPU":
            logger.info(
                "                ^ running on CPU — for GPU check that a "
                "CUDA-enabled jaxlib is installed and JAX_PLATFORMS is unset"
            )
        logger.info(f"  Sampler     : {sampler}")
        logger.info(f"  Parameters  : {model.param_names}  ({sum(int(jnp.size(v)) for v in model.theta_init.values())} dims)")
        logger.info(f"  Observations: {list(keys)}")
        logger.info(f"  Output      : {output_path}")
        logger.info(f"  Log         : {log_path}")
        logger.info(f"  Likelihood build: {_t_likelihood:.3f} s")

        # ── Configure sampler ─────────────────────────────────────────
        _t0_adapter = time.perf_counter()
        adapter = _build_adapter(
            sampler, model, sampler_kwargs, verbose,
            vi=vi, vi_kwargs=vi_kwargs,
        )
        _t_adapter = time.perf_counter() - _t0_adapter
        logger.info(f"  Adapter build:    {_t_adapter:.3f} s")

        # ── Run ────────────────────────────────────────────────────────
        _t0_sampler = time.perf_counter()
        result = run_sampler(model, multi_likelihood, adapter, rng_key)
        _t_sampler = time.perf_counter() - _t0_sampler

        logger.info(f"\n{result.summary()}")

        # ── Write HDF5 ────────────────────────────────────────────────
        _t0_h5 = time.perf_counter()
        write_result_h5(output_path, model, result, verbose=verbose)
        _t_h5 = time.perf_counter() - _t0_h5

        _t_total = _t_likelihood + _t_adapter + _t_sampler + _t_h5
        logger.info(f"\n  fitSED timing breakdown:")
        logger.info(f"    Likelihood build : {_t_likelihood:>8.3f} s")
        logger.info(f"    Adapter build    : {_t_adapter:>8.3f} s")
        logger.info(f"    Sampler run      : {_t_sampler:>8.1f} s")
        logger.info(f"    HDF5 write       : {_t_h5:>8.3f} s")
        logger.info(f"    Total fitSED     : {_t_total:>8.1f} s")

        return result
    finally:
        # Detach the per-fit file handler so repeated fitSED calls in the
        # same session do not multiply handlers or write to stale files.
        logger.removeHandler(_file_handler)
        _file_handler.close()

ceridwen.sampler.run_sampler

run_sampler(model, likelihood, adapter, rng_key)

Build JAX callables from a Ceridwen model and dispatch to a sampler.

This is the primary user-facing entry point for posterior sampling in Ceridwen. It is intentionally thin:

1. Constructs a JIT-compiled loglike_fn that sums log-likelihoods over all registered observations (no prior contribution).
2. Wraps model.ln_prior as a JIT-compiled logprior_fn.
3. Delegates to adapter.run(loglike_fn, logprior_fn, theta_init, key).

Parameters:

Name	Type	Description	Default
model	SedModel	Initialised model with observations set up and priors registered. model.predict(theta) and model.ln_prior(theta) are called inside the JIT-compiled hot paths.	required
likelihood	MultiObservationLikelihood	Likelihood object for the registered observations. Its .keys and .likelihoods attributes are iterated statically at XLA trace time — they are Python constants, not traced values.	required
adapter	SamplerAdapter	Concrete backend (e.g. BlackJAXNestedSamplerAdapter).	required
rng_key	Array	JAX PRNGKey for reproducibility.	required

Returns:

Type	Description
SamplingResult	Call .to_anesthetic() for corner plots and evidence estimates.

Examples:

Nested sampling with BlackJAX NSS::

from ceridwen.sampler import run_sampler
from ceridwen.sampler.nested import BlackJAXNestedSamplerAdapter

adapter = BlackJAXNestedSamplerAdapter(
    priors          = model.priors,
    num_live        = 500,
    num_inner_steps = len(model.param_names) * 5,
)
result = run_sampler(model, multi_likelihood, adapter,
                     jax.random.PRNGKey(42))

ns = result.to_anesthetic(labels={"Z": r"$\log Z/Z_\odot$"})
print(result.summary())

Future HMC adapter (same call signature)::

adapter = BlackJAXHMCAdapter(step_size=0.01, n_warmup=500)
result  = run_sampler(model, multi_likelihood, adapter, key)

Source code in ceridwen/sampler/runner.py

def run_sampler(
    model      : Any,
    likelihood : Any,
    adapter    : SamplerAdapter,
    rng_key    : Array,
) -> SamplingResult:
    """
    Build JAX callables from a Ceridwen model and dispatch to a sampler.

    This is the **primary user-facing entry point** for posterior sampling
    in Ceridwen.  It is intentionally thin:

    1. Constructs a JIT-compiled ``loglike_fn`` that sums log-likelihoods
       over all registered observations (no prior contribution).
    2. Wraps ``model.ln_prior`` as a JIT-compiled ``logprior_fn``.
    3. Delegates to ``adapter.run(loglike_fn, logprior_fn, theta_init, key)``.

    Parameters
    ----------
    model : SedModel
        Initialised model with observations set up and priors registered.
        ``model.predict(theta)`` and ``model.ln_prior(theta)`` are called
        inside the JIT-compiled hot paths.
    likelihood : MultiObservationLikelihood
        Likelihood object for the registered observations.  Its ``.keys``
        and ``.likelihoods`` attributes are iterated statically at
        XLA trace time — they are Python constants, not traced values.
    adapter : SamplerAdapter
        Concrete backend (e.g. ``BlackJAXNestedSamplerAdapter``).
    rng_key : Array
        JAX PRNGKey for reproducibility.

    Returns
    -------
    SamplingResult
        Call ``.to_anesthetic()`` for corner plots and evidence estimates.

    Examples
    --------
    Nested sampling with BlackJAX NSS::

        from ceridwen.sampler import run_sampler
        from ceridwen.sampler.nested import BlackJAXNestedSamplerAdapter

        adapter = BlackJAXNestedSamplerAdapter(
            priors          = model.priors,
            num_live        = 500,
            num_inner_steps = len(model.param_names) * 5,
        )
        result = run_sampler(model, multi_likelihood, adapter,
                             jax.random.PRNGKey(42))

        ns = result.to_anesthetic(labels={"Z": r"$\\log Z/Z_\\odot$"})
        print(result.summary())

    Future HMC adapter (same call signature)::

        adapter = BlackJAXHMCAdapter(step_size=0.01, n_warmup=500)
        result  = run_sampler(model, multi_likelihood, adapter, key)
    """
    # ── Static data extracted once, before trace ──────────────────────────
    _obs_dict    = model.obs_dict
    _keys        = tuple(likelihood.keys)
    _likelihoods = tuple(likelihood.likelihoods)
    _static_data = {
        key: (
            _obs_dict[key].flux,
            _obs_dict[key].uncertainty,
            _obs_dict[key].mask,
        )
        for key in _keys
    }

    # ── Log-likelihood: sum over observations, no prior ───────────────────
    @jax.jit
    def loglike_fn(theta: dict[str, Array]) -> Array:
        predictions = model.predict(theta)
        lnl = jnp.zeros(())
        for key, lhood in zip(_keys, _likelihoods):
            y_k, sig_k, mask_k = _static_data[key]
            mu_k    = predictions[key]
            lnl_k, _ = lhood(y_k, mu_k, sig_k, mask_k, params=theta)
            lnl = lnl + lnl_k
        return lnl

    # ── Log-prior ─────────────────────────────────────────────────────────
    @jax.jit
    def logprior_fn(theta: dict[str, Array]) -> Array:
        return model.ln_prior(theta)

    # ── Delegate ──────────────────────────────────────────────────────────
    return adapter.run(loglike_fn, logprior_fn, model.theta_init, rng_key)

ceridwen.sampler.nested.BlackJAXNestedSamplerAdapter

BlackJAXNestedSamplerAdapter(
    priors,
    num_live=500,
    num_inner_steps=None,
    num_delete=None,
    logZ_tol=-3.0,
    verbose=True,
    checkpoint_interval_s=1200.0,
    checkpoint_dir=None,
)

Bases: SamplerAdapter

Adapter wrapping blackjax.nss for use with any Ceridwen SedModel.

Handles three concerns that are specific to nested sampling in Ceridwen:

1. Live-point initialisation — samples each free parameter independently from its registered prior using the Prior.sample method from ceridwen.sampler.priors.

2. Shape reconciliation — the BlackJAX NSS live-point dict has shape {name: (num_live, *param_shape)}. The step function vmaps over axis 0, delivering single-particle slices of shape (*param_shape,) to loglike_fn / logprior_fn. This matches the Ceridwen dict-theta convention exactly.

3. Evidence extraction — optionally uses anesthetic for a more accurate :math:\ln Z estimate with uncertainty.

Parameters:

Name	Type	Description	Default
priors	dict[str, Prior]	Mapping from free-parameter name to a Ceridwen Prior object. Every free parameter must have a prior — nested sampling requires a proper (normalisable) prior; an improper flat prior makes the evidence integral undefined.	required
num_live	int	Number of live points. Default 500.	500
num_inner_steps	int	Inner MCMC steps per NS iteration. Default n_dims * 5 where n_dims is the total scalar dimension count.	None
num_delete	int	Live points discarded per iteration. Default num_live // 2.	None
logZ_tol	float	Convergence threshold on :math:\ln(Z_\mathrm{live}/Z). Default -3.0.	-3.0
verbose	bool	Print a tqdm progress bar and convergence info. Default True.	True

Source code in ceridwen/sampler/nested.py

def __init__(
    self,
    priors          : dict,
    num_live        : int   = 500,
    num_inner_steps : Optional[int] = None,
    num_delete      : Optional[int] = None,
    logZ_tol        : float = -3.0,
    verbose         : bool  = True,
    checkpoint_interval_s : float = 1200.0,
    checkpoint_dir        : Optional[str] = None,
):
    self.priors          = dict(priors)
    self.num_live        = int(num_live)
    self._num_inner_steps = num_inner_steps   # None → auto
    self._num_delete      = num_delete         # None → num_live // 2
    self.logZ_tol        = float(logZ_tol)
    self.verbose         = bool(verbose)
    # Periodic checkpointing.  Every ``checkpoint_interval_s`` seconds
    # (default 1200 = 20 min; <= 0 disables) the accumulated dead points
    # are finalised against the current live ensemble and dumped to disk,
    # so a run killed by the scheduler wall-time, a node failure, or any
    # mid-run crash still yields a recoverable (partial) posterior --
    # BlackJAX provides no native checkpointing.  The destination is
    # resolved at run time from ``checkpoint_dir`` ->
    # $CERIDWEN_CHECKPOINT_DIR -> $CERIDWEN_RESCUE_DIR; when none is set
    # checkpointing is silently skipped (no surprise writes).  The same
    # snapshot format is written once more at convergence as the rescue
    # pickle, so :meth:`load_checkpoint` recovers either.
    self.checkpoint_interval_s = float(checkpoint_interval_s)
    self._checkpoint_dir       = checkpoint_dir

load_checkpoint staticmethod

load_checkpoint(path)

Load a checkpoint / rescue pickle written by this adapter.

Returns the dict {positions, loglikelihood, loglikelihood_birth, logZ, n_dead, partial}. A partial=True snapshot is a usable (under-converged) posterior from a run killed before convergence --- feed positions + loglikelihood + loglikelihood_birth to anesthetic.NestedSamples exactly as the end-of-run path does.

Source code in ceridwen/sampler/nested.py

@staticmethod
def load_checkpoint(path):
    """Load a checkpoint / rescue pickle written by this adapter.

    Returns the dict ``{positions, loglikelihood, loglikelihood_birth,
    logZ, n_dead, partial}``.  A ``partial=True`` snapshot is a usable
    (under-converged) posterior from a run killed before convergence ---
    feed ``positions`` + ``loglikelihood`` + ``loglikelihood_birth`` to
    ``anesthetic.NestedSamples`` exactly as the end-of-run path does.
    """
    import pickle as _pickle
    with open(path, "rb") as fh:
        return _pickle.load(fh)

run

run(loglike_fn, logprior_fn, theta_init, rng_key)

Run BlackJAX NSS and return a SamplingResult.

Parameters:

Name	Type	Description	Default
loglike_fn	callable	JIT-compiled log-likelihood (no prior).	required
logprior_fn	callable	JIT-compiled log-prior (must be proper).	required
theta_init	dict[str, Array]	Reference parameter dict (shapes / dtypes).	required
rng_key	Array		required

Returns:

Type	Description
SamplingResult

Source code in ceridwen/sampler/nested.py

def run(
    self,
    loglike_fn  : Callable[[dict[str, Array]], Array],
    logprior_fn : Callable[[dict[str, Array]], Array],
    theta_init  : dict[str, Array],
    rng_key     : Array,
) -> SamplingResult:
    """
    Run BlackJAX NSS and return a ``SamplingResult``.

    Parameters
    ----------
    loglike_fn : callable
        JIT-compiled log-likelihood (no prior).
    logprior_fn : callable
        JIT-compiled log-prior (must be proper).
    theta_init : dict[str, Array]
        Reference parameter dict (shapes / dtypes).
    rng_key : Array

    Returns
    -------
    SamplingResult
    """
    try:
        import blackjax
        import blackjax.ns.utils as ns_utils
    except ImportError as exc:
        raise ImportError(
            "BlackJAX with nested sampling (blackjax.ns) is required.\n"
            "Install: pip install 'git+https://github.com/blackjax-devs/blackjax@f73e12956'"
        ) from exc

    import tqdm

    n_dims          = self._n_dims(theta_init)
    num_inner_steps = (self._num_inner_steps
                       if self._num_inner_steps is not None
                       else n_dims * 5)
    num_delete      = (self._num_delete
                       if self._num_delete is not None
                       else self.num_live // 2)

    if self.verbose:
        print(
            f"BlackJAX NSS  |  n_dims={n_dims}  "
            f"num_live={self.num_live}  "
            f"num_inner_steps={num_inner_steps}  "
            f"num_delete={num_delete}"
        )

    # ── Initialise live points ────────────────────────────────────────
    rng_key, prior_key = jax.random.split(rng_key)
    particles = self._sample_prior(theta_init, prior_key)

    # ── Build NSS kernel ──────────────────────────────────────────────
    # loglike_fn / logprior_fn operate on a SINGLE particle (un-batched).
    # The NSS step_fn vmaps internally over the live-point ensemble.
    nested_sampler = blackjax.nss(
        logprior_fn      = logprior_fn,
        loglikelihood_fn = loglike_fn,
        num_delete       = num_delete,
        num_inner_steps  = num_inner_steps,
    )
    init_fn = jax.jit(nested_sampler.init)
    step_fn = jax.jit(nested_sampler.step)

    if self.verbose:
        print("  [timing] Calling init_fn (JIT compile + eval) ...",
              flush=True)
        _t0 = time.perf_counter()

    live = init_fn(particles)

    # ── BlackJAX version compatibility ───────────────────────────────
    # Three known layouts for logZ / logZ_live:
    #   v1  NSState (direct):           state.logZ, state.logZ_live
    #   v2  AdaptiveNSState (wrapper):  state.sampler_state.logZ, ...
    #   v3  AdaptiveNSState (integrator): state.integrator.logZ, ...
    def _build_logZ_accessors(state):
        """Return (get_logZ, get_logZ_live) callables for *state*."""
        # v1 – direct NSState
        if hasattr(state, "logZ") and hasattr(state, "logZ_live"):
            return (lambda s: float(s.logZ),
                    lambda s: float(s.logZ_live))
        # v3 – newest: integrator sub-object
        if hasattr(state, "integrator"):
            ig = state.integrator
            if hasattr(ig, "logZ") and hasattr(ig, "logZ_live"):
                return (lambda s: float(s.integrator.logZ),
                        lambda s: float(s.integrator.logZ_live))
        # v2 – older wrapper: sampler_state sub-object
        if hasattr(state, "sampler_state"):
            inner = state.sampler_state
            if hasattr(inner, "logZ") and hasattr(inner, "logZ_live"):
                return (lambda s: float(s.sampler_state.logZ),
                        lambda s: float(s.sampler_state.logZ_live))
        # Unknown layout – raise with diagnostics
        _fields = [f for f in dir(state) if not f.startswith("_")]
        raise AttributeError(
            f"Cannot locate logZ/logZ_live on {type(state).__name__}.\n"
            f"  Top-level fields : {_fields}\n"
            + "Please check your BlackJAX version."
        )

    _get_logZ, _get_logZ_live = _build_logZ_accessors(live)

    if self.verbose:
        _t1 = time.perf_counter()
        print(f"  [timing] init_fn done    ({_t1 - _t0:.1f} s)  "
              f"logZ={_get_logZ(live):.4f}  "
              f"logZ_live={_get_logZ_live(live):.4f}", flush=True)
        print(f"  (state type: {type(live).__name__})")

    # ── NS run loop ───────────────────────────────────────────────────
    dead_list    = []
    n_like_calls = 0
    t_start      = time.perf_counter()

    # Periodic-checkpoint bookkeeping (see __init__).
    _ckpt_dir   = self._resolve_ckpt_dir()
    _ckpt_on    = bool(_ckpt_dir) and self.checkpoint_interval_s > 0
    _last_ckpt  = t_start
    if _ckpt_on and self.verbose:
        print(f"  [checkpoint] every {self.checkpoint_interval_s:.0f} s "
              f"-> {_ckpt_dir}", flush=True)

    desc = "NS  (starting)"
    try:
        desc = f"NS  logZ={_get_logZ(live):.1f}"
    except (AttributeError, NameError):
        pass

    with tqdm.tqdm(
        desc=desc,
        unit=" dead",
        disable=not self.verbose,
    ) as pbar:
        _iter = 0
        while float(_get_logZ_live(live) - _get_logZ(live)) >= self.logZ_tol:
            rng_key, subkey = jax.random.split(rng_key)
            if _iter == 0 and self.verbose:
                print("  [step_fn] Compiling the step kernel (one-time JIT) "
                      "+ running the first iteration. This compile can be "
                      "slow on CPU (seconds to many minutes depending on "
                      "model size and hardware); subsequent steps are fast.",
                      flush=True)
            _t_iter = time.perf_counter()

            live, dead_info = step_fn(subkey, live)

            _dt_iter = time.perf_counter() - _t_iter
            _iter += 1
            if self.verbose:
                _logZ = _get_logZ(live)
                _dlogZ = _get_logZ_live(live) - _logZ
                print(
                    f"  [iter {_iter:>4d}]  {_dt_iter:6.1f} s  "
                    f"logZ={_logZ:+.3f}  ΔlogZ={_dlogZ:.3f}  "
                    f"dead={num_delete * _iter}",
                    flush=True,
                )
            dead_list.append(dead_info)
            n_like_calls += num_delete * num_inner_steps
            pbar.update(num_delete)
            try:
                pbar.set_description(
                    f"NS  logZ={_get_logZ(live):.2f}  "
                    f"ΔlogZ={_get_logZ_live(live) - _get_logZ(live):.2f}"
                )
            except AttributeError:
                pass

            # Periodic checkpoint: finalise + dump a partial snapshot so a
            # wall-time kill / node crash mid-run is recoverable.
            if _ckpt_on and (time.perf_counter() - _last_ckpt
                             >= self.checkpoint_interval_s):
                _p = self._dump_snapshot(
                    _ckpt_dir, live, dead_list, ns_utils,
                    _get_logZ(live), tag="checkpoint", partial=True)
                _last_ckpt = time.perf_counter()
                if _p and self.verbose:
                    print(f"  [checkpoint] iter {_iter}: {_p}", flush=True)

    wall_time = time.perf_counter() - t_start
    if self.verbose:
        print(
            f"  Converged  logZ = {_get_logZ(live):.3f}  "
            f"({wall_time:.1f} s,  {n_like_calls:,} likelihood calls)"
        )

    # ── Merge live points into dead set ───────────────────────────────
    # Version-aware finalise + dead-point unpack (see _finalise_dead).
    # The ``update_info`` kwarg only exists in newer BlackJAX; 0.1.0b0
    # ships ``finalise(live, dead)`` with no such param, so passing it
    # unconditionally raises ``TypeError`` only AFTER convergence,
    # losing a multi-hour run.  The guard lives in _finalise_dead so it
    # (and the periodic-checkpoint path) cannot drift.
    _dead_positions, _dead_logl, _dead_logl_birth = self._finalise_dead(
        live, dead_list, ns_utils)

    # ── Rescue pickle ─────────────────────────────────────────────────
    # Dump the finalised dead points so any failure further down the save
    # path (anesthetic, evidence, the caller's I/O) is recoverable rather
    # than discarding a multi-hour run.  Same format as the periodic
    # checkpoints; loadable via load_checkpoint().  partial=False marks a
    # fully-converged snapshot.
    _rescue_dir = self._resolve_ckpt_dir()
    if _rescue_dir:
        self._dump_snapshot(_rescue_dir, live, dead_list, ns_utils,
                            _get_logZ(live), tag="rescue", partial=False)

    # ── Evidence & importance weights (anesthetic preferred) ─────────
    log_Z        = float(_get_logZ(live))
    log_Z_err    = float("nan")
    log_weights  = None
    try:
        from anesthetic import NestedSamples
        import numpy as np

        _raw   = _dead_positions
        _names = [n for n in _raw if n in theta_init]
        _cols  = [
            np.asarray(_raw[n]).reshape(
                len(np.asarray(_dead_logl)), -1
            )
            for n in _names
        ]
        _data  = np.hstack(_cols)
        _ns    = NestedSamples(
            _data,
            logL       = np.asarray(_dead_logl),
            logL_birth = np.asarray(_dead_logl_birth),
            logzero    = float("nan"),
        )
        log_Z     = float(_ns.logZ())
        log_Z_err = float(_ns.logZ(12).std())
        # Extract proper NS importance weights from anesthetic.
        # These encode the prior-volume compression at each dead point.
        log_weights = jnp.asarray(np.asarray(_ns.logw()))
    except Exception:
        pass  # fall back to live.logZ and manual weights below

    # Fallback: compute log-weights from prior volume shrinkage if
    # anesthetic is unavailable or failed.
    if log_weights is None:
        n_dead  = len(jnp.asarray(_dead_logl))
        n_live  = self.num_live
        # Standard NS trapezoid rule: log(X_{i-1} - X_{i+1}) / 2
        # where X_i = exp(-i / n_live) is the prior volume fraction.
        log_vols = -jnp.arange(n_dead, dtype=float) / n_live
        log_dvol = jnp.log(
            jnp.exp(jnp.roll(log_vols, 1) - log_vols)
            - jnp.exp(jnp.roll(log_vols, -1) - log_vols)
        ) + log_vols
        # Fix boundary: first and last points
        log_dvol = log_dvol.at[0].set(
            jnp.log1p(-jnp.exp(-1.0 / n_live))
        )
        log_dvol = log_dvol.at[-1].set(
            log_vols[-1] - jnp.log(n_live)
        )
        log_weights = log_dvol

    # ── Pack samples ──────────────────────────────────────────────────
    # Squeeze trailing size-1 axes so scalar params have shape (n_dead,)
    # rather than (n_dead, 1), matching typical user expectation.
    samples = {}
    for name in theta_init:
        arr = jnp.asarray(_dead_positions[name])   # (n_dead, *shape)
        # Squeeze only if the parameter was a scalar (shape (1,))
        if arr.ndim > 1 and arr.shape[-1] == 1:
            arr = jnp.squeeze(arr, axis=-1)
        samples[name] = arr

    return SamplingResult(
        samples               = samples,
        log_evidence          = log_Z,
        log_evidence_err      = log_Z_err,
        log_weights           = log_weights,
        log_likelihoods       = jnp.asarray(_dead_logl),
        log_likelihoods_birth = jnp.asarray(_dead_logl_birth),
        param_names           = list(theta_init.keys()),
        n_likelihood_calls    = n_like_calls,
        wall_time_s           = wall_time,
        sampler_name          = "blackjax.nss",
        raw                   = {
            "positions": _dead_positions,
            "loglikelihood": _dead_logl,
            "loglikelihood_birth": _dead_logl_birth,
        },
    )

ceridwen.sampler.nuts.BlackJAXNUTSAdapter

BlackJAXNUTSAdapter(
    num_warmup=None,
    num_samples=2000,
    num_chains=4,
    initial_step_size=None,
    target_acceptance=0.95,
    max_num_doublings=10,
    dense_mass=None,
    bounds=None,
    vi=None,
    vi_kwargs=None,
    verbose=True,
)

Bases: SamplerAdapter

Adapter wrapping blackjax.nuts with window adaptation for Ceridwen.

Bounded (uniform-prior) parameters are automatically reparameterised onto an unconstrained space via sigmoid/logit, eliminating the hard boundary walls that cause divergent transitions.

Optionally, a variational transport map (see :mod:ceridwen.sampler.vi) may be supplied via vi. When set, the adapter trains the map against the unconstrained posterior and then runs NUTS on the whitened target :math:\log p(f(z)) + \log|\partial f/\partial z| (Hoffman et al. 2019, arXiv:1903.03704). In whitened space the target is approximately :math:\mathcal{N}(0, I) so identity mass matrix and step size :math:\mathcal{O}(1) are near-optimal, giving dramatically shorter warmup.

Parameters:

Name	Type	Description	Default
num_warmup	int	Number of warmup (adaptation) steps per chain. Default 1500 for native NUTS. When vi is set, warmup defaults to 200; very short warmup lets the dual-averaging adapter overshoot and produce many divergences on a ~14-D SED posterior.	None
num_samples	int	Number of post-warmup posterior draws per chain. Default 2000.	2000
num_chains	int	Number of independent chains. Default 4.	4
initial_step_size	float	Starting step size for the leapfrog integrator before adaptation. Default 0.01 for native NUTS. When vi is set, default is 0.5 — z-space is pre-whitened so the optimal :math:\varepsilon is :math:\mathcal{O}(1), and dual averaging is slow to shrink from 1.0.	None
target_acceptance	float	Target acceptance probability for dual averaging. Default 0.95; higher values reduce divergences in VI-preconditioned NUTS.	0.95
max_num_doublings	int	Maximum tree depth (2^max_num_doublings leapfrog steps). Default 10.	10
dense_mass	bool	Use a dense (full) inverse mass matrix. Default True for native NUTS. When vi is set, default is False — the VI map already orthogonalises the geometry, so an adapted diagonal mass matrix in z-space is sufficient and faster to fit.	None
bounds	dict	Maps parameter names to (low, high) tuples for bounded params. If None (default), bounds are auto-detected from the model priors passed through run_sampler. You can also pass them explicitly:: bounds={'Z': (-2.5, 0.2), 'logmass': (9.0, 12.0)}	None
vi	None, str, VariationalMap, or TrainedMap	Variational preconditioning mode. None (default): run native NUTS with window adaptation. 'tril': train a full-rank Gaussian map, then whiten. 'iaf': train a stacked-IAF neural-transport map, then whiten. :class:VariationalMap instance: train this map. :class:TrainedMap instance: skip training, use as-is.	None
vi_kwargs	dict	Forwarded either to the VI map constructor (when vi is a string) or to :func:ceridwen.sampler.vi.train_vi for training hyperparameters. Recognised keys: num_steps (default 1500), batch_size (16), lr0 (1e-2) plus map-specific kwargs (e.g. init_scale for 'tril', n_flows for 'iaf').	None
verbose	bool	Print progress information. Default True.	True

Source code in ceridwen/sampler/nuts.py

def __init__(
    self,
    num_warmup: int | None = None,
    num_samples: int = 2000,
    num_chains: int = 4,
    initial_step_size: float | None = None,
    target_acceptance: float = 0.95,
    max_num_doublings: int = 10,
    dense_mass: bool | None = None,
    bounds: dict[str, tuple[float, float]] | None = None,
    vi: "str | Any | None" = None,
    vi_kwargs: dict | None = None,
    verbose: bool = True,
):
    # VI-aware defaults: in whitened space, target geometry is
    # approximately isotropic Gaussian, so shorter warmup / larger
    # initial step / diagonal mass matrix are all appropriate.
    _has_vi = vi is not None
    if num_warmup is None:
        num_warmup = 200 if _has_vi else 1500
    if initial_step_size is None:
        initial_step_size = 0.5 if _has_vi else 0.01
    if dense_mass is None:
        dense_mass = not _has_vi

    self.num_warmup = int(num_warmup)
    self.num_samples = int(num_samples)
    self.num_chains = int(num_chains)
    self.initial_step_size = float(initial_step_size)
    self.target_acceptance = float(target_acceptance)
    self.max_num_doublings = int(max_num_doublings)
    self.dense_mass = bool(dense_mass)
    self.bounds = dict(bounds) if bounds is not None else None
    self.vi = vi
    self.vi_kwargs = dict(vi_kwargs) if vi_kwargs is not None else {}
    self.verbose = bool(verbose)
    # Filled during .run() so the caller can inspect the trained map
    # after sampling (e.g. for plotting the learnt covariance).
    self.trained_map = None

run

run(loglike_fn, logprior_fn, theta_init, rng_key)

Run BlackJAX NUTS with window adaptation.

Strategy (GPU-optimised): 1. Run ONE warmup to adapt step size + mass matrix. 2. Build a single NUTS kernel from the adapted parameters. 3. vmap the sampling across all chains in parallel. This compiles ONE XLA program and runs all chains simultaneously, fully utilising GPU parallelism.

Parameters:

Name	Type	Description	Default
loglike_fn	callable	JIT-compiled log-likelihood (no prior).	required
logprior_fn	callable	JIT-compiled log-prior.	required
theta_init	dict[str, Array]	Initial parameter values with correct shapes.	required
rng_key	Array		required

Returns:

Type	Description
SamplingResult

Source code in ceridwen/sampler/nuts.py

def run(
    self,
    loglike_fn: Callable[[dict[str, Array]], Array],
    logprior_fn: Callable[[dict[str, Array]], Array],
    theta_init: dict[str, Array],
    rng_key: Array,
) -> SamplingResult:
    """
    Run BlackJAX NUTS with window adaptation.

    Strategy (GPU-optimised):
      1. Run ONE warmup to adapt step size + mass matrix.
      2. Build a single NUTS kernel from the adapted parameters.
      3. vmap the sampling across all chains in parallel.
         This compiles ONE XLA program and runs all chains
         simultaneously, fully utilising GPU parallelism.

    Parameters
    ----------
    loglike_fn : callable
        JIT-compiled log-likelihood (no prior).
    logprior_fn : callable
        JIT-compiled log-prior.
    theta_init : dict[str, Array]
        Initial parameter values with correct shapes.
    rng_key : Array

    Returns
    -------
    SamplingResult
    """
    try:
        import blackjax
    except ImportError as exc:
        raise ImportError(
            "BlackJAX is required for NUTS sampling.\n"
            "Install: pip install git+https://github.com/blackjax-devs/blackjax"
        ) from exc

    # Dispatch: VI-preconditioned path is structurally different
    # (whitened target, chains-from-q init, no mass-matrix adapt in
    # z-space) so it lives in its own method.
    if self.vi is not None:
        return self._run_whitened(
            loglike_fn, logprior_fn, theta_init, rng_key,
        )

    n_dims = self._n_dims(theta_init)
    theta_template = theta_init
    bounds = self.bounds if self.bounds is not None else {}

    if self.verbose:
        print(
            f"BlackJAX NUTS  |  n_dims={n_dims}  "
            f"num_warmup={self.num_warmup}  "
            f"num_samples={self.num_samples}  "
            f"num_chains={self.num_chains}  "
            f"dense_mass={self.dense_mass}"
        )
        if bounds:
            for k, (a, b) in bounds.items():
                print(f"  Bounded: {k} -> sigmoid({a}, {b})")
        else:
            print("  No bounded parameters (consider passing bounds=...)")

    # ── Build reparameterisation layer ─────────────────────────────
    lo, hi, is_bounded = _build_transforms(bounds, theta_template)
    n_bounded = int(jnp.sum(is_bounded))

    if self.verbose and n_bounded > 0:
        print(f"  Reparameterised {n_bounded}/{n_dims} "
              f"bounded dimensions via sigmoid/logit")

    # ── Unconstrained log-posterior ────────────────────────────────
    @jax.jit
    def logposterior_flat(x):
        theta_flat = _to_constrained(x, lo, hi, is_bounded)
        theta = self._unflatten(theta_flat, theta_template)
        lnl = loglike_fn(theta)
        lnp = logprior_fn(theta)
        lnj = _log_jacobian(x, lo, hi, is_bounded)
        return lnl + lnp + lnj

    @jax.jit
    def loglike_constrained(x):
        """Log-likelihood in unconstrained coords (for diagnostics)."""
        theta_flat = _to_constrained(x, lo, hi, is_bounded)
        theta = self._unflatten(theta_flat, theta_template)
        return loglike_fn(theta)

    # ── Initial position in unconstrained space ────────────────────
    x_init_flat = self._flatten(theta_init)
    x_init = _to_unconstrained(x_init_flat, lo, hi, is_bounded)

    t_start = time.perf_counter()

    # ==============================================================
    #  Phase 1: Warmup (single chain, adapts step size + mass matrix)
    # ==============================================================
    if self.verbose:
        print(f"\n  Warmup ({self.num_warmup} steps, "
              f"adapting step size + {'dense' if self.dense_mass else 'diagonal'} mass matrix)...")

    # ``max_num_doublings`` forwards through window_adaptation to the
    # wrapped blackjax.nuts kernel used during warmup; it caps the
    # per-step leapfrog count at 2**max_num_doublings.  Must be passed
    # explicitly — BlackJAX otherwise falls back to its default of 10
    # (→ 1024 leapfrog steps), and at small adapted step sizes one
    # NUTS iteration can take seconds.
    warmup = blackjax.window_adaptation(
        blackjax.nuts,
        logposterior_flat,
        target_acceptance_rate=self.target_acceptance,
        initial_step_size=self.initial_step_size,
        progress_bar=self.verbose,
        is_mass_matrix_diagonal=not self.dense_mass,
        max_num_doublings=self.max_num_doublings,
    )

    warmup_key, sample_key = jax.random.split(rng_key)

    _t0_warmup = time.perf_counter()
    (warmup_state, parameters), _ = warmup.run(
        warmup_key,
        x_init,
        num_steps=self.num_warmup,
    )
    jax.block_until_ready(warmup_state.position)
    _t_warmup = time.perf_counter() - _t0_warmup

    step_size = parameters['step_size']
    if self.verbose:
        print(f"    Adapted step size: {float(step_size):.4f}")
        if self.dense_mass:
            im = np.asarray(parameters['inverse_mass_matrix'])
            cond = np.linalg.cond(im)
            print(f"    Mass matrix condition number: {cond:.1f}")
        print(f"    Warmup wall time: {_t_warmup:.1f} s  "
              f"(includes XLA compilation)")

    # ==============================================================
    #  Phase 2: Sampling
    #
    #  Strategy depends on the number of available devices:
    #
    #  Multi-GPU (n_devices >= n_chains):
    #    Use jax.pmap to run one chain per GPU in true parallel.
    #    Unlike vmap, pmap places each chain on a *separate device*,
    #    so each NUTS while_loop (tree building) runs independently
    #    with its own adaptive tree depth — no padding to max depth.
    #    This gives near-linear speedup with number of GPUs.
    #
    #  Single-GPU fallback:
    #    Run chains sequentially with a cached XLA kernel.
    #    The lax.scan compiles once on Chain 1 and is reused.
    # ==============================================================
    # ``parameters`` from window_adaptation already carries
    # ``max_num_doublings``, so it must not be passed a second time
    # here (``TypeError: got multiple values for keyword argument``).
    nuts_kernel = blackjax.nuts(
        logposterior_flat,
        **parameters,
    ).step

    def _nuts_step(state, key):
        state, info = nuts_kernel(key, state)
        return state, (state, info)

    @jax.jit
    def _run_one_chain(init_state, chain_key):
        """Run num_samples NUTS steps from init_state."""
        keys = jax.random.split(chain_key, self.num_samples)
        final_state, (states, infos) = jax.lax.scan(
            _nuts_step, init_state, keys
        )
        return states, infos

    chain_keys = jax.random.split(sample_key, self.num_chains)

    n_devices = len(jax.devices())
    use_pmap = (n_devices >= self.num_chains) and (self.num_chains > 1)

    if self.verbose:
        print(f"\n  Available devices: {n_devices}  |  "
              f"Chains: {self.num_chains}  |  "
              f"Strategy: {'pmap (one chain per GPU)' if use_pmap else 'sequential'}")

    all_chain_positions = []  # (num_samples, n_dims) per chain
    all_loglikelihoods = []
    all_divergences = []
    all_infos = []
    _t_sample_chains = []
    _t_postproc_chains = []

    if use_pmap:
        # ── Multi-GPU parallel path ──────────────────────────────
        # Replicate the warmup state across devices and pmap the scan.
        if self.verbose:
            print(f"\n  Running {self.num_chains} chains in parallel "
                  f"across {self.num_chains} GPUs...", flush=True)

        # pmap expects a leading device axis.  Replicate warmup_state
        # across chains (each chain starts from the same adapted state).
        def _replicate_state(state, n):
            """Replicate a NUTS state across n devices."""
            return jax.tree.map(
                lambda x: jnp.broadcast_to(x, (n,) + x.shape), state
            )

        pmap_init = _replicate_state(warmup_state, self.num_chains)

        # pmap the chain runner.  axis_name is used for potential
        # cross-device reductions (not needed here, but good practice).
        @jax.pmap
        def _run_chains_pmap(init_state, chain_key):
            keys = jax.random.split(chain_key, self.num_samples)
            final_state, (states, infos) = jax.lax.scan(
                _nuts_step, init_state, keys
            )
            return states, infos

        _t0_sample = time.perf_counter()
        pmap_states, pmap_infos = _run_chains_pmap(
            pmap_init, chain_keys
        )
        # Block until all GPUs finish
        jax.block_until_ready(pmap_states.position)
        _t_sample_total = time.perf_counter() - _t0_sample

        if self.verbose:
            print(f"    All chains wall time: {_t_sample_total:.1f} s  "
                  f"(parallel across {self.num_chains} GPUs)")

        # ── Post-processing: unpack pmap results ─────────────────
        _t0_pp = time.perf_counter()
        for ci in range(self.num_chains):
            x_chain = pmap_states.position[ci]  # (num_samples, n_dims)
            theta_chain = jax.vmap(
                lambda x: _to_constrained(x, lo, hi, is_bounded)
            )(x_chain)
            all_chain_positions.append(theta_chain)

            _chunk = min(50, self.num_samples)
            _lnl_parts = []
            for _i in range(0, self.num_samples, _chunk):
                _lnl_parts.append(
                    jax.vmap(loglike_constrained)(x_chain[_i:_i + _chunk])
                )
            chain_lnl = jnp.concatenate(_lnl_parts, axis=0)
            jax.block_until_ready(chain_lnl)
            all_loglikelihoods.append(chain_lnl)

            chain_infos = jax.tree.map(lambda x: x[ci], pmap_infos)
            all_infos.append(chain_infos)
            if hasattr(chain_infos, 'is_divergent'):
                n_div = int(jnp.sum(chain_infos.is_divergent))
                all_divergences.append(n_div)
                if self.verbose and n_div > 0:
                    print(f"    Chain {ci+1}: {n_div} divergent transitions!")
            else:
                all_divergences.append(0)

        _t_pp = time.perf_counter() - _t0_pp
        _t_sample_chains = [_t_sample_total / self.num_chains] * self.num_chains
        _t_postproc_chains = [_t_pp / self.num_chains] * self.num_chains

        if self.verbose:
            print(f"    Post-processing: {_t_pp:.1f} s")

    else:
        # ── Single-GPU sequential path ───────────────────────────
        for chain_idx in range(self.num_chains):
            if self.verbose:
                print(f"\n  Chain {chain_idx + 1}/{self.num_chains}  "
                      f"({self.num_samples} draws)...", flush=True)

            _t0_sample = time.perf_counter()
            states, infos = _run_one_chain(warmup_state, chain_keys[chain_idx])
            jax.block_until_ready(states.position)
            _t_sample = time.perf_counter() - _t0_sample
            _t_sample_chains.append(_t_sample)

            if self.verbose:
                print(f"    Sampling wall time: {_t_sample:.1f} s")

            _t0_pp = time.perf_counter()

            x_chain = states.position
            theta_chain = jax.vmap(
                lambda x: _to_constrained(x, lo, hi, is_bounded)
            )(x_chain)
            all_chain_positions.append(theta_chain)

            _chunk = min(50, self.num_samples)
            _lnl_parts = []
            for _i in range(0, self.num_samples, _chunk):
                _lnl_parts.append(
                    jax.vmap(loglike_constrained)(x_chain[_i:_i + _chunk])
                )
            chain_lnl = jnp.concatenate(_lnl_parts, axis=0)
            jax.block_until_ready(chain_lnl)
            all_loglikelihoods.append(chain_lnl)

            all_infos.append(infos)
            if hasattr(infos, 'is_divergent'):
                n_div = int(jnp.sum(infos.is_divergent))
                all_divergences.append(n_div)
                if self.verbose and n_div > 0:
                    print(f"    WARNING: {n_div} divergent transitions!")
            else:
                all_divergences.append(0)

            _t_pp = time.perf_counter() - _t0_pp
            _t_postproc_chains.append(_t_pp)

            if self.verbose:
                print(f"    Post-processing wall time: {_t_pp:.1f} s")

    n_like_calls = self.num_chains * (self.num_warmup + self.num_samples)
    wall_time = time.perf_counter() - t_start

    # ── Timing summary ────────────────────────────────────────────
    _t_merge_start = time.perf_counter()

    if self.verbose:
        _t_sample_total = sum(_t_sample_chains)
        _t_pp_total = sum(_t_postproc_chains)
        print("\n  " + "=" * 60)
        print("  NUTS timing breakdown")
        print("  " + "=" * 60)
        print(f"  {'Phase':<40s}  {'Time':>10s}")
        print("  " + "-" * 60)
        print(f"  {'Warmup (incl. XLA compilation)':<40s}  {_t_warmup:>9.1f}s")
        if use_pmap:
            print(f"  {'Sampling (pmap, all chains parallel)':<40s}  "
                  f"{_t_sample_total:>9.1f}s")
        else:
            for ci in range(self.num_chains):
                print(f"  {'  Chain ' + str(ci+1) + ' sampling':<40s}  "
                      f"{_t_sample_chains[ci]:>9.1f}s")
            print(f"  {'Sampling total (sequential)':<40s}  {_t_sample_total:>9.1f}s")
        print(f"  {'Post-processing total':<40s}  {_t_pp_total:>9.1f}s")
        print("  " + "-" * 60)
        print(f"  {'TOTAL':<40s}  {wall_time:>9.1f}s")
        if use_pmap:
            seq_est = _t_sample_total * self.num_chains
            print(f"  {'Estimated sequential time':<40s}  {seq_est:>9.1f}s")
            print(f"  {'pmap speedup':<40s}  {seq_est / _t_sample_total:>9.1f}x")
        print("  " + "=" * 60)

    # ── Merge chains (constrained space) ──────────────────────────
    merged_flat = jnp.concatenate(all_chain_positions, axis=0)

    merged_samples = {}
    idx = 0
    for name, template in theta_template.items():
        size = int(jnp.size(template))
        arr = merged_flat[:, idx:idx + size]
        if arr.shape[-1] == 1:
            arr = jnp.squeeze(arr, axis=-1)
        merged_samples[name] = arr
        idx += size

    merged_lnl = jnp.concatenate(all_loglikelihoods, axis=0)

    total_divergences = sum(all_divergences)
    _t_merge = time.perf_counter() - _t_merge_start

    if self.verbose:
        print(
            f"\n  Done  ({wall_time:.1f} s total,  "
            f"merge {_t_merge:.1f} s,  "
            f"{n_like_calls:,} likelihood calls"
            f"{f', {total_divergences} divergences' if total_divergences else ''})"
        )

    # ── Convergence diagnostics ───────────────────────────────────
    if self.verbose and self.num_chains >= 2:
        self._print_diagnostics(all_chain_positions, theta_template)

    return SamplingResult(
        samples=merged_samples,
        log_evidence=float("nan"),
        log_evidence_err=float("nan"),
        log_weights=jnp.zeros_like(merged_lnl),
        log_likelihoods=merged_lnl,
        param_names=list(theta_init.keys()),
        n_likelihood_calls=n_like_calls,
        wall_time_s=wall_time,
        sampler_name="blackjax.nuts",
        raw={
            "num_chains": self.num_chains,
            "num_warmup": self.num_warmup,
            "num_samples": self.num_samples,
            "total_divergences": total_divergences,
            "dense_mass": self.dense_mass,
            "per_chain_constrained": all_chain_positions,
            "per_chain_infos": all_infos,
            "used_pmap": use_pmap,
            "n_devices": n_devices,
        },
    )

Dust, nebular, IGM, cosmology

ceridwen.dust.DustModel

Dust

Dust(bin_edges=[(-jnp.inf, -1.97)], laws=['powerlaw'])

JAX-compatible modular dust model that supports multiple attenuation laws per bin.

Parameters are passed as plain dicts (dict[str, Array]).

Call Dust.describe_attenuation_laws() to list all available models.

Source code in ceridwen/dust/DustModel.py

def __init__(self, bin_edges=[(-jnp.inf, -1.97)], laws=['powerlaw']):
    assert len(bin_edges) == len(laws), "Must have one dust law per bin"
    self.bin_edges = jnp.array(bin_edges)
    self.num_bins = len(bin_edges)
    self.laws = laws

    self.law_names_resolved = []
    law_name_counter = defaultdict(int)
    law_occurrences = {name: laws.count(name) for name in set(laws)}

    self.law_funcs = []
    self.law_params = []

    for name in laws:
        count = law_name_counter[name]
        law_name_counter[name] += 1

        law_entry = ATTENUATION_LAWS[name]
        base_func = law_entry["func"]
        defaults = law_entry.get("defaults", {})
        params = law_entry.get("params", {})
        doc = law_entry.get("doc", "")

        if law_occurrences[name] == 1:
            resolved_name = name
            func = base_func
            param_dict = self._get_law_params(name)
        else:
            number = count + 1
            resolved_name = f"{name}{number}"
            func = modify_function(base_func, number, defaults)

            ATTENUATION_LAWS[resolved_name] = {
                "func": func,
                "defaults": {f"{k}{number}": v for k, v in defaults.items()},
                "params": {f"{k}{number}": d for k, d in params.items()},
                "doc": doc
            }

            renamed_keys = [f"{k}{number}" for k in defaults.keys()]
            ignored_keys = list(defaults.keys())
            print(f"[dust setup] Law '{name}' used multiple times → registered as '{resolved_name}'.")
            print(f"             Use parameters: {', '.join(renamed_keys)}")
            print(f"             Parameters like {', '.join(ignored_keys)} will be ignored.")

            param_dict = self._get_law_params(name)
            param_dict = {f"{k}{number}": v for k, v in param_dict.items()}

        sig = inspect.signature(func)
        ordered_param_names = [
            p.name for p in sig.parameters.values()
            if p.name != "wave"
            and p.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
        ]
        ordered_param_names = [n for n in ordered_param_names if n in param_dict]

        if not ordered_param_names:
            ordered_param_names = sorted(param_dict.keys())

        missing = [n for n in ordered_param_names if n not in param_dict]
        if missing:
            raise ValueError(
                f"Parameters {missing} expected for '{resolved_name}' but not found in param_dict"
            )

        law_defaults = ATTENUATION_LAWS.get(resolved_name, {}).get("defaults", {})
        wrapped_func = make_law_wrapper(func, ordered_param_names, law_defaults)

        self.law_names_resolved.append(resolved_name)
        self.law_funcs.append(wrapped_func)
        self.law_params.append(param_dict)

compute_attenuation

compute_attenuation(wave, fit_params)

Compute bin-wise attenuation curves.

Parameters:

Name	Type	Description	Default
wave	ndarray	Wavelength array in Angstroms.	required
fit_params	dict[str, Array]	Parameter dict. Each law wrapper extracts only the keys it needs.	required

Returns:

Type	Description
(ndarray, shape(num_bins, len(wave)))

Source code in ceridwen/dust/DustModel.py

def compute_attenuation(self, wave, fit_params):
    """
    Compute bin-wise attenuation curves.

    Parameters
    ----------
    wave : jnp.ndarray
        Wavelength array in Angstroms.
    fit_params : dict[str, Array]
        Parameter dict.  Each law wrapper extracts only the keys it needs.

    Returns
    -------
    jnp.ndarray, shape (num_bins, len(wave))
    """
    def curve_fn(i, wave):
        return lax.switch(i, self.law_funcs, wave, fit_params)

    return vmap(curve_fn, in_axes=(0, None))(jnp.arange(self.num_bins), wave)

get_default_fit_params

get_default_fit_params()

Return a plain dict of default fit parameters.

Keys are the parameter names used in the active dust laws; values are JAX scalars. Returned as a dict so it merges directly into the global theta dict.

Source code in ceridwen/dust/DustModel.py

def get_default_fit_params(self):
    """
    Return a plain dict of default fit parameters.

    Keys are the parameter names used in the active dust laws; values are
    JAX scalars.  Returned as a dict so it merges directly into the global
    theta dict.
    """
    defaults = {}
    for law in self.law_names_resolved:
        for k, v in ATTENUATION_LAWS[law].get("defaults", {}).items():
            defaults[k] = jnp.asarray(v)
    return defaults

DiffuseDust

DiffuseDust(law='kriek_conroy')

Bases: Dust

Single-bin dust model (one law covering all ages) with diffuse_ prefixed parameter names to avoid collisions with birth-cloud parameters in a shared theta dict.

Source code in ceridwen/dust/DustModel.py

def __init__(self, law="kriek_conroy"):
    super().__init__(bin_edges=[(-jnp.inf, jnp.inf)], laws=[law])

    old_name = self.law_names_resolved[0]
    param_dict = self.law_params[0]

    self.diffuse_param_map = {}
    renamed_param_dict = {}
    for k, v in param_dict.items():
        new_k = f"diffuse_{k}"
        renamed_param_dict[new_k] = v
        self.diffuse_param_map[new_k] = k

    self.law_params[0] = renamed_param_dict

    func = ATTENUATION_LAWS[old_name]["func"]
    sig = inspect.signature(func)
    param_names = [
        f"diffuse_{p.name}" for p in sig.parameters.values()
        if p.name != "wave"
        and p.kind not in (inspect.Parameter.VAR_POSITIONAL, inspect.Parameter.VAR_KEYWORD)
    ]
    wrapped_func = make_law_wrapper(func, param_names)
    self.law_funcs[0] = wrapped_func
    self.dust_param_names = param_names

get_default_params

get_default_params()

Return a plain dict of default diffuse-dust parameters (diffuse_* keys).

Source code in ceridwen/dust/DustModel.py

def get_default_params(self):
    """
    Return a plain dict of default diffuse-dust parameters (``diffuse_*`` keys).
    """
    defaults = {}
    law = self.law_names_resolved[0]
    for k, v in ATTENUATION_LAWS[law].get("defaults", {}).items():
        defaults[f"diffuse_{k}"] = jnp.asarray(v)
    return defaults

compute_attenuation

compute_attenuation(wave, fit_params)

Compute the diffuse attenuation curve (single bin).

Parameters:

Name	Type	Description	Default
wave	ndarray		required
fit_params	dict[str, Array]		required

Returns:

Type	Description
(ndarray, shape(len(wave)))

Source code in ceridwen/dust/DustModel.py

def compute_attenuation(self, wave, fit_params):
    """
    Compute the diffuse attenuation curve (single bin).

    Parameters
    ----------
    wave : jnp.ndarray
    fit_params : dict[str, Array]

    Returns
    -------
    jnp.ndarray, shape (len(wave),)
    """
    return self.law_funcs[0](wave, fit_params)

make_law_wrapper

make_law_wrapper(f, param_names, defaults=None)

Return a JAX-traceable wrapper that extracts named params from a dict.

Parameters missing from fit_params fall back to defaults (the registry defaults), so a law parameter can be pinned via theta without breaking older suites whose theta dicts never carried it. The membership test is a static Python check resolved at trace time.

Source code in ceridwen/dust/DustModel.py

def make_law_wrapper(f, param_names, defaults=None):
    """Return a JAX-traceable wrapper that extracts named params from a dict.

    Parameters missing from ``fit_params`` fall back to ``defaults`` (the
    registry defaults), so a law parameter can be pinned via theta without
    breaking older suites whose theta dicts never carried it.  The
    membership test is a static Python check resolved at trace time.
    """
    defaults = defaults or {}

    def wrapped(wave, fit_params):
        args = tuple(fit_params[name] if name in fit_params else defaults[name]
                     for name in param_names)
        return f(wave, *args)
    return wrapped

modify_function

modify_function(func, number, defaults_dict=None)

Return a parameter-renamed wrapper of func (no exec).

When the same attenuation law is used for multiple age bins, its parameters are suffixed with the bin number so they do not collide in the shared theta dict (e.g. tau_pow → tau_pow{number}). The returned wrapper is invoked positionally by :func:make_law_wrapper (func(wave, *args)); the renaming is purely so that inspect.signature reports the suffixed names (which :class:Dust reads to build the per-bin parameter-extraction list). We therefore forward the positional arguments unchanged and attach a renamed __signature__.

The single-law path never calls this function and is unaffected.

Source code in ceridwen/dust/DustModel.py

def modify_function(func, number, defaults_dict=None):
    """Return a parameter-renamed wrapper of ``func`` (no ``exec``).

    When the same attenuation law is used for multiple age bins, its parameters
    are suffixed with the bin ``number`` so they do not collide in the shared
    theta dict (e.g. ``tau_pow`` → ``tau_pow{number}``).  The returned wrapper
    is invoked *positionally* by :func:`make_law_wrapper`
    (``func(wave, *args)``); the renaming is purely so that
    ``inspect.signature`` reports the suffixed names (which :class:`Dust` reads
    to build the per-bin parameter-extraction list).  We therefore forward the
    positional arguments unchanged and attach a renamed ``__signature__``.

    The single-law path never calls this function and is unaffected.
    """
    sig = inspect.signature(func)

    new_params = []
    for name, param in sig.parameters.items():
        if param.kind == param.VAR_KEYWORD:
            continue
        new_name = name if name == "wave" else f"{name}{number}"
        if param.default is not inspect.Parameter.empty:
            default = param.default
        elif defaults_dict and name in defaults_dict:
            default = defaults_dict[name]
        else:
            default = inspect.Parameter.empty
        new_params.append(param.replace(name=new_name, default=default))
    new_params.append(inspect.Parameter("kwargs", inspect.Parameter.VAR_KEYWORD))
    new_sig = sig.replace(parameters=new_params)

    def wrapper(wave, *args, **kwargs):
        # make_law_wrapper calls positionally: wrapper(wave, *param_values).
        # Forward positionally to the original law; the parameter renaming is
        # only for the dict-key lookup, not for this call.
        return func(wave, *args)

    wrapper.__name__ = wrapper.__qualname__ = f"{func.__name__}{number}"
    wrapper.__signature__ = new_sig
    return wrapper

ceridwen.dust.DustEmission

DustEmission

DustEmission(
    duste_model="DL07",
    dust_file=None,
    spec_lambda=None,
    **kwargs
)

Initialize the DustEmission object with parameters for dust emission modeling.

Parameters:

Name	Type	Description	Default
duste_model	str	Dust emission model to use: 'DL07' or 'THEMIS'.	'DL07'
dust_file	str	Path to the dust emission file (required).	None
spec_lambda	ndarray	Wavelength grid over which dust emission will be evaluated (required).	None
kwargs	dict	Optional keyword arguments to override default dust parameters. Supported: duste_qpah, duste_umin, duste_gamma	{}

Source code in ceridwen/dust/DustEmission.py

def __init__(self, duste_model="DL07",
             dust_file=None, spec_lambda=None, **kwargs):
    """
    Initialize the DustEmission object with parameters for dust emission modeling.

    Parameters
    ----------
    duste_model : str
        Dust emission model to use: 'DL07' or 'THEMIS'.
    dust_file : str
        Path to the dust emission file (required).
    spec_lambda : ndarray
        Wavelength grid over which dust emission will be evaluated (required).
    kwargs : dict
        Optional keyword arguments to override default dust parameters.
        Supported: duste_qpah, duste_umin, duste_gamma
    """

    # Store model choice (e.g., 'DL07' or 'THEMIS')
    self.duste_model = duste_model

    # Dust parameter values
    self.duste_qpah = None
    self.duste_umin = None
    self.duste_gamma = None

    # Model grid arrays for allowed values of qPAH and Umin
    self.qpaharr = None
    self.uminarr = None

    # Placeholder for loaded dust emission spectra
    self.dustem2_dustem = None

    # File path and wavelength grid
    self.dust_file = None
    self.spec_lambda = None

    # Store any additional keyword arguments for later use
    self.dwargs = kwargs

    # Read in key dust parameters, allowing overrides via kwargs
    self.duste_qpah = kwargs.pop("duste_qpah", 3.5)
    self.duste_umin = kwargs.pop("duste_umin", 1.0)
    self.duste_gamma = kwargs.pop("duste_gamma", 0.01)

    # Set parameter grids based on selected dust model
    if self.duste_model == "DL07":
        self.qpaharr = jnp.array([0.47,1.12,1.77,2.50,3.19,3.90,4.58])
        self.uminarr = jnp.array([0.1,0.15,0.2,0.3,0.4,0.5,0.7,0.8,1.0,1.2,1.5,
                                  2.0,2.5,3.0,4.0,5.0,7.0,8.0,12.0,15.0,20.0,25.0])
        self.nqpah_dustem = self.qpaharr.size
        self.numin_dustem = self.uminarr.size
    elif self.duste_model == "THEMIS":
        # THEMIS model uses smaller qPAH values rescaled to percent
        self.qpaharr = jnp.array([0.02, 0.06, 0.10, 0.14, 0.17, 0.20, 0.24,
                                  0.28, 0.32, 0.36, 0.40]) / 2.2 * 100
        self.uminarr = jnp.array([
            0.1, 0.12, 0.15, 0.17, 0.2, 0.25, 0.3, 0.35, 0.4, 0.5, 0.6,
            0.7, 0.8, 1.0, 1.2, 1.5, 1.7, 2.0, 2.5, 3.0, 3.5, 4.0, 5.0,
            6.0, 7.0, 8.0, 10.0, 12.0, 15.0, 17.0, 20.0, 25.0, 30.0,
            35.0, 40.0, 50.0, 80.0
        ])
        self.nqpah_dustem = self.qpaharr.size
        self.numin_dustem = self.uminarr.size
    else:
        raise ValueError("Invalid duste_model. Choose 'DL07' or 'THEMIS'.")

    # Ensure that necessary data is provided
    if dust_file is None or spec_lambda is None:
        raise ValueError("If `duste=True`, both `dust_file` and `spec_lambda` must be provided.")

    self.dust_file = dust_file
    self.spec_lambda = spec_lambda

    # Load emission templates or model data from file
    self.load_dust_emission(dust_file, spec_lambda)

    # ── Precomputed integration machinery (static, depends only on wavelength grid) ──
    nu = CLIGHT_AA_S / jnp.asarray(spec_lambda)       # (n_wave,) Hz
    dnu = jnp.diff(nu)
    # Trapezoidal quadrature weights: trapezoid(f, nu) == dot(f, _trap_w)
    self._trap_w = jnp.concatenate([
        jnp.array([0.5 * dnu[0]]),
        0.5 * (dnu[:-1] + dnu[1:]),
        jnp.array([0.5 * dnu[-1]]),
    ])

__repr__

__repr__()

Custom string representation of the DustEmission object. Provides a readable summary of model settings and parameters.

Source code in ceridwen/dust/DustEmission.py

def __repr__(self):
    """
    Custom string representation of the DustEmission object.
    Provides a readable summary of model settings and parameters.
    """

    def format_array(arr):
        """Helper to format short arrays inline; longer arrays multiline."""
        if arr is None:
            return "None"
        if arr.ndim == 1 and len(arr) <= 5:
            return f"[{', '.join(map(str, arr))}]"
        return f"\n    " + "\n    ".join(map(str, arr))

    attributes = {
        "Duste model": self.duste_model,
        "DUST qPAH": self.duste_qpah,
        "DUST Umin": self.duste_umin,
        "DUST Gamma": self.duste_gamma,
        f"qpaharr ({self.duste_model})": format_array(self.qpaharr),
        f"uminarr ({self.duste_model})": format_array(self.uminarr),
        "dust_file": self.dust_file,
        "spec_lambda": self.spec_lambda.shape if self.spec_lambda is not None else None,
    }

    if self.dwargs:
        attributes["Extra parameters (dwargs)"] = self.dwargs

    attr_str = "\n".join(f"  {k:<30}: {v}" for k, v in attributes.items() if v is not None)
    return f"\nDustEmission Model:\n{'='*50}\n{attr_str}\n{'='*50}"

get_default_params

get_default_params()

Return a plain dict of default dust-emission fit parameters.

Returned as a dict so it merges directly into the global theta dict.

Source code in ceridwen/dust/DustEmission.py

def get_default_params(self):
    """
    Return a plain dict of default dust-emission fit parameters.

    Returned as a dict so it merges directly into the global theta dict.
    """
    return {
        "duste_qpah": jnp.asarray(self.duste_qpah),
        "duste_umin": jnp.asarray(self.duste_umin),
        "duste_gamma": jnp.asarray(self.duste_gamma),
    }

update_dust_params

update_dust_params(
    duste_qpah=3.5, duste_umin=1.0, duste_gamma=0.01
)

Update dust parameters for emission calculations.

Parameters: duste_qpah (float, optional): New PAH fraction. duste_umin (float, optional): New minimum U radiation field. duste_gamma (float, optional): New fraction of high U component.

Source code in ceridwen/dust/DustEmission.py

def update_dust_params(self, duste_qpah = 3.5, duste_umin = 1.0, duste_gamma = 0.01):
    """
    Update dust parameters for emission calculations.

    Parameters:
        duste_qpah (float, optional): New PAH fraction.
        duste_umin (float, optional): New minimum U radiation field.
        duste_gamma (float, optional): New fraction of high U component.
    """
    self.duste_qpah = duste_qpah
    self.duste_umin = duste_umin
    self.duste_gamma = duste_gamma

compute_dust_emission

compute_dust_emission(
    spec_attn,
    spec_dustfree,
    spec_lambda,
    diffuse_curve,
    duste_qpah,
    duste_umin,
    duste_gamma,
)

Compute dust emission using precomputed trapezoidal weights for fast integration (bolometric luminosities via jnp.dot against the static weight vector).

Parameters: spec_attn (jnp.ndarray): Attenuated spectrum after dust absorption. spec_dustfree (jnp.ndarray): Stellar spectrum before attenuation. spec_lambda (jnp.ndarray): Wavelength array in Angstroms. diffuse_curve (jnp.ndarray): exp(-tau_diffuse), shape (1, n_wave) or (n_wave,). duste_qpah (float): PAH fraction. duste_umin (float): Minimum U radiation field. duste_gamma (float): Fraction of high U component.

Returns: tuple: (spectrum with dust emission, dust mass, total dust emission)

Source code in ceridwen/dust/DustEmission.py

def compute_dust_emission(self, spec_attn, spec_dustfree, spec_lambda, diffuse_curve,
                                    duste_qpah, duste_umin, duste_gamma):
    """
    Compute dust emission using precomputed trapezoidal weights for fast
    integration (bolometric luminosities via jnp.dot against the static
    weight vector).

    Parameters:
        spec_attn (jnp.ndarray): Attenuated spectrum after dust absorption.
        spec_dustfree (jnp.ndarray): Stellar spectrum before attenuation.
        spec_lambda (jnp.ndarray): Wavelength array in Angstroms.
        diffuse_curve (jnp.ndarray): exp(-tau_diffuse), shape (1, n_wave) or (n_wave,).
        duste_qpah (float): PAH fraction.
        duste_umin (float): Minimum U radiation field.
        duste_gamma (float): Fraction of high U component.

    Returns:
        tuple: (spectrum with dust emission, dust mass, total dust emission)
    """
    tiny = 1e-70
    w = self._trap_w                                # (n_wave,) precomputed
    dc = diffuse_curve.ravel()                      # (n_wave,)
    f_attn = spec_attn.ravel()                      # (n_wave,)
    f_free = spec_dustfree.ravel()                  # (n_wave,)

    # --- Bolometric luminosities via dot product (replaces trapezoid) ---
    lbold = jnp.dot(f_attn, w)
    lboln = jnp.dot(f_free, w)

    # --- QPAH bilinear interpolation index ---
    qlo = jnp.clip(jnp.searchsorted(self.qpaharr, duste_qpah) - 1, 0, self.nqpah_dustem - 2)
    dq  = jnp.clip((duste_qpah - self.qpaharr[qlo]) / (self.qpaharr[qlo + 1] - self.qpaharr[qlo]), 0, 1)

    # --- Umin bilinear interpolation index ---
    ulo = jnp.clip(jnp.searchsorted(self.uminarr, duste_umin) - 1, 0, self.numin_dustem - 2)
    du  = jnp.clip((duste_umin - self.uminarr[ulo]) / (self.uminarr[ulo + 1] - self.uminarr[ulo]), 0, 1)

    gamma = jnp.clip(duste_gamma, 0.0, 1.0)

    # --- Bilinear weights (computed once, used for both dumin and dumax) ---
    w00 = (1 - dq) * (1 - du)
    w10 = dq * (1 - du)
    w01 = (1 - dq) * du
    w11 = dq * du

    i_lo = 2 * ulo
    i_hi = 2 * (ulo + 1)

    D = self.dustem2_dustem           # (n_wave, n_qpah, n_umin*2)

    dumin = (w00 * D[:, qlo, i_lo]     + w10 * D[:, qlo + 1, i_lo] +
             w11 * D[:, qlo + 1, i_hi] + w01 * D[:, qlo, i_hi])

    dumax = (w00 * D[:, qlo, i_lo + 1]     + w10 * D[:, qlo + 1, i_lo + 1] +
             w11 * D[:, qlo + 1, i_hi + 1] + w01 * D[:, qlo, i_hi + 1])

    # --- Dust emission template (combine Umin and Umax parts of P(U)dU) ---
    mduste = jnp.maximum((1 - gamma) * dumin + gamma * dumax, tiny).ravel()
    norm   = jnp.dot(mduste, w)
    mduste_norm = mduste / norm              # unit-luminosity template

    # --- Initial emission from absorbed stellar luminosity ---
    labs0  = lboln - lbold
    duste0 = jnp.maximum(mduste_norm * labs0, tiny)

    # --- Self-absorption iteration 1 ---
    duste0_atten = duste0 * dc
    # Absorbed luminosity = integral of (duste0 - duste0*dc) = integral of duste0*(1-dc)
    absorbed_1 = jnp.dot(duste0 * (1.0 - dc), w)
    duste1 = jnp.maximum(mduste_norm * absorbed_1, tiny)

    # --- Self-absorption iteration 2 (attenuate only; no further re-emission needed) ---
    duste1_atten = duste1 * dc

    # --- Final results ---
    tduste = duste0_atten + duste1_atten
    specdust = f_attn + tduste
    mdust = MDUST_PREFACTOR * labs0 / norm

    return specdust, mdust, tduste

ceridwen.neb.NebularModel

NebularModel(
    cloudy_dust,
    sps_home,
    csp_lambda,
    ssp_flux=None,
    ssp_ages_lgyr=None,
    isoc_type="mist",
    nebnz=11,
    nebnage=10,
    nebnip=7,
    smooth_velocity=True,
    sigma_smooth=100.0,
    res_floor_factor=2.0,
    nebular_smooth_init=None,
)

CLOUDY-grid nebular emission model — physically strict variant.

The continuum cube and the line cube are each interpolated against their own (logZ, age, logU) axes, ensuring that the returned nebular spectrum corresponds to the parameters CLOUDY was actually run at. This differs from FSPS's run-time convention; see the module docstring and :class:ceridwen.neb.NebularGridModel_fsps_match.NebularModelFSPSMatch for the FSPS-matching variant.

Parameters:

Name	Type	Description	Default
cloudy_dust	bool	False → load the dust-free grids ZAU_ND_<isoc>.{cont,lines}. True → load the dust-attenuated grids ZAU_WD_<isoc>.….	required
sps_home	str | Path	FSPS root directory; expects <sps_home>/nebular/ to contain the ZAU grids.	required
csp_lambda	(array, shape(nspec))	Wavelength grid (Å) the nebular spectrum will be projected onto.	required
ssp_flux	(array, shape(n_z, n_age, n_wave) or None)	Optional. If provided, log_qq is computed internally from the SSP fluxes via the FSPS-equivalent formula.	None
ssp_ages_lgyr	(array, shape(n_age) or None)	Optional. log10(age / yr) of every SSP in ssp_flux. Used to flag the SSPs whose age lies inside both CLOUDY grids (intersection of the two age ranges).	None
isoc_type	(mist, pdva, prsc, bpss)	Isochrone tag identifying the ZAU file suffix.	'mist'
nebnz	int	Dimensions of each CLOUDY grid (default 11, 10, 7 for MIST).	11
nebnage	int	Dimensions of each CLOUDY grid (default 11, 10, 7 for MIST).	11
nebnip	int	Dimensions of each CLOUDY grid (default 11, 10, 7 for MIST).	11
smooth_velocity	bool	True → sigma_smooth is in km/s; False → Å.	True
sigma_smooth	float	Line-broadening σ. Default 0.0 (matches FSPS).	100.0
nebular_smooth_init	float or None	Backwards-compatible alias for sigma_smooth.	None

Attributes:

Name	Type	Description
nebem_cont	(nspec, nebnz, nebnage, nebnip)	log10 of nebular continuum (L_ν / Q) on csp_lambda.
nebem_line	(nemline, nebnz, nebnage, nebnip)	log10 of per-line luminosities (Lsun / Q).
nebem_cont_logz, nebem_cont_age, nebem_cont_logu	1 - D	The .cont cube's own grid axes (ascending).
nebem_line_logz, nebem_line_age, nebem_line_logu	1 - D	The .lines cube's own grid axes (ascending).
nebem_logz, nebem_age, nebem_logu	1 - D	Legacy aliases — point at the line axes for compatibility with code (e.g. NebularGridModelSVD) that expects the FSPS canonical axis names.
nebem_line_pos	(nemline,)	Rest-frame emission-line centroids in Å.
gaussnebarr	(nspec, nemline)	Pre-computed Gaussian profiles including the FSPS λ²/c factor.
log_qq	(n_z, n_age) or None	Self-consistent log10(Q) derived from ssp_flux.
young_mask, young_idx		Boolean mask / index array over ssp_ages_lgyr selecting ages inside the intersection of both cubes' age ranges.

Source code in ceridwen/neb/NebularGridModel.py

def __init__(self,
             cloudy_dust,
             sps_home,
             csp_lambda,
             ssp_flux=None,
             ssp_ages_lgyr=None,
             isoc_type='mist',
             nebnz=11, nebnage=10, nebnip=7,
             smooth_velocity=True,
             sigma_smooth=100.0,
             res_floor_factor=2.0,
             nebular_smooth_init=None):
    # Defaults match FSPS verbatim:
    #   nebular_smooth_init = 100 km/s  (sps_vars.f90)
    #   smooth_velocity     = .true.    (sps_vars.f90)
    #   pixel floor         = neb_res_min * 2  (sps_setup.f90)
    # Prospector inherits this through python-fsps.
    if nebular_smooth_init is not None:
        sigma_smooth = float(nebular_smooth_init)

    self.csp_lambda = jnp.asarray(csp_lambda)
    self.nspec      = int(self.csp_lambda.size)

    self.smooth_velocity = bool(smooth_velocity)
    self.sigma_smooth    = float(sigma_smooth)
    self.res_floor_factor = float(res_floor_factor)

    self.nebnz   = int(nebnz)
    self.nebnage = int(nebnage)
    self.nebnip  = int(nebnip)

    suffix = 'WD' if cloudy_dust else 'ND'
    base = Path(sps_home) / 'nebular' / f'ZAU_{suffix}_{isoc_type}'
    self.cont_file = base.with_suffix('.cont')
    self.line_file = base.with_suffix('.lines')

    self._load_continuum()
    self._load_lines()
    self._compute_resolution_elements()
    self._build_gaussians()

    if ssp_flux is not None:
        self.log_qq = self.compute_log_qq(jnp.asarray(ssp_flux))
    else:
        self.log_qq = None

    # Restrict nebular emission to SSPs whose age lies inside BOTH
    # cubes' age ranges (intersection).  This is the safest choice
    # given that the two grids may differ.
    if ssp_ages_lgyr is not None:
        ages = jnp.asarray(ssp_ages_lgyr)
        max_age = min(float(self.nebem_cont_age[-1]),
                      float(self.nebem_line_age[-1]))
        young = ages <= max_age
        self.young_mask = young
        self.young_idx  = jnp.where(young)[0]
    else:
        self.young_mask = None
        self.young_idx  = None

compute_log_qq

compute_log_qq(ssp_flux)

log10(Q) for every (Z, age) SSP, matching FSPS's run-time formula::

qq = ∫ L_nu / lambda  dλ          (0 < λ < 912 Å)
Q  = (L_sun_erg / h_erg_s) × qq   (photons / s)

Forced to float64 because the deep-UV fluxes are tiny.

Source code in ceridwen/neb/NebularGridModel.py

def compute_log_qq(self, ssp_flux):
    """
    ``log10(Q)`` for every (Z, age) SSP, matching FSPS's run-time
    formula::

        qq = ∫ L_nu / lambda  dλ          (0 < λ < 912 Å)
        Q  = (L_sun_erg / h_erg_s) × qq   (photons / s)

    Forced to float64 because the deep-UV fluxes are tiny.
    """
    mask = self.csp_lambda < LYMAN_LIMIT_AA
    wave_ion = self.csp_lambda[mask].astype(jnp.float64)
    flux_ion = ssp_flux[..., mask].astype(jnp.float64)
    qq = jnp.trapezoid(flux_ion / wave_ion, x=wave_ion)
    scale = LSUN_ERG_S / HPLANK_ERG_S
    return jnp.log10(jnp.maximum(qq * scale, TINY))

evaluate

evaluate(logZ, logU, logage, logQ)

Single-point evaluation — returns (cont, lines) in Lsun/Hz.

The continuum cube is interpolated against the cont axes, the line cube against the line axes (each cube is therefore evaluated at the physical point CLOUDY was actually run at).

Source code in ceridwen/neb/NebularGridModel.py

def evaluate(self, logZ, logU, logage, logQ):
    """
    Single-point evaluation — returns ``(cont, lines)`` in Lsun/Hz.

    The continuum cube is interpolated against the **cont** axes,
    the line cube against the **line** axes (each cube is therefore
    evaluated at the physical point CLOUDY was actually run at).
    """
    # ── continuum cube on its own axes
    zc  = _locate(logZ,   self.nebem_cont_logz)
    dzc = _frac(logZ,     self.nebem_cont_logz, zc)
    uc  = _locate(logU,   self.nebem_cont_logu)
    duc = _frac(logU,     self.nebem_cont_logu, uc)
    ac  = _locate(logage, self.nebem_cont_age)
    dac = _frac(logage,   self.nebem_cont_age,  ac)
    log_cont = _trilinear(self.nebem_cont, zc, dzc, ac, dac, uc, duc)

    # ── line cube on its own axes
    zl  = _locate(logZ,   self.nebem_line_logz)
    dzl = _frac(logZ,     self.nebem_line_logz, zl)
    ul  = _locate(logU,   self.nebem_line_logu)
    dul = _frac(logU,     self.nebem_line_logu, ul)
    al  = _locate(logage, self.nebem_line_age)
    dal = _frac(logage,   self.nebem_line_age,  al)
    log_line = _trilinear(self.nebem_line, zl, dzl, al, dal, ul, dul)

    cont_flux = jnp.power(10.0, log_cont + logQ)
    line_lum  = jnp.power(10.0, log_line + logQ)
    line_spec = self.gaussnebarr @ line_lum
    return cont_flux, line_spec

evaluate_batch

evaluate_batch(
    logZ_gas,
    logU,
    ssp_ages_young,
    logqq_young,
    return_components=False,
)

Vectorised evaluation for all (Z_ssp, age_young) pairs at a single (logZ_gas, logU). Each cube is bilinearly collapsed in (Z, U) on its OWN axes, then linearly interpolated in age across the young-SSP set.

Parameters:

Name	Type	Description	Default
return_components	bool	When False (default), return cont_flux + line_spec as a single (n_z, n_wave, n_young) array -- the legacy behaviour consumed by the four CSPBasis.get_spectrum_*_neb variants. When True, return the tuple (cont_flux, line_spec) in the same layout, so the caller can decide whether to include the broadened emission lines in the continuum spectrum (the prospector nebemlineinspec switch).	False

Source code in ceridwen/neb/NebularGridModel.py

def evaluate_batch(self, logZ_gas, logU, ssp_ages_young, logqq_young,
                    return_components=False):
    """
    Vectorised evaluation for all (Z_ssp, age_young) pairs at a
    single ``(logZ_gas, logU)``.  Each cube is bilinearly collapsed
    in (Z, U) on its OWN axes, then linearly interpolated in age
    across the young-SSP set.

    Parameters
    ----------
    return_components : bool
        When ``False`` (default), return ``cont_flux + line_spec`` as a
        single ``(n_z, n_wave, n_young)`` array -- the legacy behaviour
        consumed by the four ``CSPBasis.get_spectrum_*_neb`` variants.
        When ``True``, return the tuple ``(cont_flux, line_spec)`` in the
        same layout, so the caller can decide whether to include the
        broadened emission lines in the continuum spectrum (the
        prospector ``nebemlineinspec`` switch).
    """
    logZ_gas = jnp.squeeze(logZ_gas)
    logU     = jnp.squeeze(logU)

    def _interp_cube(cube, logz_grid, age_grid, logu_grid):
        z1 = _locate(logZ_gas, logz_grid)
        dz = _frac(logZ_gas,   logz_grid, z1)
        u1 = _locate(logU,     logu_grid)
        du = _frac(logU,       logu_grid, u1)

        w00 = (1.0 - dz) * (1.0 - du)
        w01 = (1.0 - dz) *       du
        w10 =       dz   * (1.0 - du)
        w11 =       dz   *       du
        zu = (w00 * cube[..., z1,     :, u1    ]
            + w01 * cube[..., z1,     :, u1 + 1]
            + w10 * cube[..., z1 + 1, :, u1    ]
            + w11 * cube[..., z1 + 1, :, u1 + 1])

        a1 = jnp.clip(jnp.searchsorted(age_grid, ssp_ages_young) - 1,
                      0, age_grid.shape[0] - 2)
        da = jnp.clip(
            (ssp_ages_young - age_grid[a1])
            / (age_grid[a1 + 1] - age_grid[a1]),
            0.0, 1.0,
        )
        return (1.0 - da)[None, :] * zu[..., a1] + da[None, :] * zu[..., a1 + 1]

    log_cont = _interp_cube(self.nebem_cont,
                            self.nebem_cont_logz,
                            self.nebem_cont_age,
                            self.nebem_cont_logu)                     # (nspec , n_young)
    log_line = _interp_cube(self.nebem_line,
                            self.nebem_line_logz,
                            self.nebem_line_age,
                            self.nebem_line_logu)                     # (nlines, n_young)

    cont_flux = jnp.power(10.0, log_cont[None, :, :] + logqq_young[:, None, :])
    line_lum  = jnp.power(10.0, log_line[None, :, :] + logqq_young[:, None, :])
    line_spec = jnp.einsum('wl,zly->zwy', self.gaussnebarr, line_lum)
    if return_components:
        return (cont_flux.transpose(0, 2, 1),
                line_spec.transpose(0, 2, 1))
    return (cont_flux + line_spec).transpose(0, 2, 1)

get_default_params

get_default_params()

Return the FSPS default nebular parameters.

Source code in ceridwen/neb/NebularGridModel.py

def get_default_params(self):
    """Return the FSPS default nebular parameters."""
    return {'gas_logz': jnp.asarray(0.0),
            'gas_logu': jnp.asarray(-2.0)}

ceridwen.igm

ceridwen/igm.py

Intergalactic-medium (IGM) absorption models for redshift-aware SED fitting. Applied as a wavelength-dependent transmission curve to the rest-frame source spectrum as a function of the source redshift; the returned factor can simply be multiplied into the spectrum before the observation projection.

Interface

:class:IGMModel is the abstract base. Each concrete subclass implements tau(lam_rest, zred) returning the optical depth on the rest-frame wavelength grid at the given source redshift. The public :meth:IGMModel.attenuation method returns exp(-tau * factor) — the transmission curve — with an optional scalar factor matching FSPS's igm_factor fudge (so the IGM strength can be fit).

Models

:class:Madau1995 Line-by-line port of FSPS's igm_absorb.f90: 17 Lyman-series transitions (Ly-α through Ly-17) with blanketing index 3.46, metal blanketing on Ly-α, and the analytic Lyman-continuum approximation (Eq. 16 of Madau 1995). τ is capped at its peak short-wavelength value so the fitting-function breakdown below ~100 Å does not produce a non-monotonic curve.

:class:NoIGM Identity (τ ≡ 0). Useful as a null model so the forward pass does not need a Python-level branch on self.igm_model is None.

Extensibility

Adding e.g. Inoue et al. (2014) is a matter of subclassing :class:IGMModel, implementing tau, and registering in :data:_MODEL_REGISTRY. The public constructor :func:make_igm_model accepts either a string name or a pre-built instance:

from ceridwen.igm import make_igm_model, Madau1995 igm = make_igm_model("madau1995") igm = make_igm_model(Madau1995()) igm = make_igm_model(None) # NoIGM

Tied to CSPBasis

CSPBasis accepts add_igm and igm_model kwargs. When add_igm=True, CSPBasis.predict multiplies the post-mass- post-flux-factor spectrum by the transmission curve at theta["zred"] (and, optionally, theta["igm_factor"]) — so IGM strength can be fit.

References

Madau, P. (1995), ApJ, 441, 18. "Radiative Transfer in a Clumpy Universe: The Colors of High-Redshift Galaxies".

IGMModel

Bases: ABC

Abstract IGM attenuation model.

Concrete subclasses implement tau(lam_rest, zred). The :meth:attenuation method returns exp(-tau * factor) and is what callers (e.g. CSPBasis.predict) multiply into the spectrum.

tau abstractmethod

tau(lam_rest, zred)

Optical depth on the rest-frame wavelength grid at source redshift zred. Must be broadcastable to lam_rest's shape.

Source code in ceridwen/igm.py

@abc.abstractmethod
def tau(self, lam_rest: Array, zred: Array) -> Array:
    """Optical depth on the rest-frame wavelength grid at source
    redshift ``zred``.  Must be broadcastable to ``lam_rest``'s
    shape."""

attenuation

attenuation(lam_rest, zred, factor=1.0)

Transmission curve exp(-tau * factor).

Source code in ceridwen/igm.py

def attenuation(self, lam_rest: Array, zred: Array,
                factor: Union[Array, float] = 1.0) -> Array:
    """Transmission curve ``exp(-tau * factor)``."""
    return jnp.exp(-self.tau(lam_rest, zred) * factor)

NoIGM

Bases: IGMModel

Identity IGM: zero optical depth at every wavelength.

Madau1995

Bases: IGMModel

Madau (1995) IGM attenuation.

Direct port of FSPS's igm_absorb.f90. Three components:

1. Lyman-series line blanketing (17 transitions from Ly-α down to Ly-17):

.. math:: \tau_\mathrm{Ly\,series}(\lambda_\mathrm{rest}) = \sum_{i=1}^{17} A_i \left(\frac{\lambda_\mathrm{rest}(1+z)}{\lambda_i}\right)^{3.46},

applied only where :math:\lambda_\mathrm{rest} < \lambda_i.

1. Metal blanketing on Ly-α with a softer index 1.68:

.. math:: \tau_\mathrm{metal} = 0.0017 \bigl(\lambda_\mathrm{obs}/1215.67\bigr)^{1.68}.

1. Lyman-continuum photoelectric absorption (Madau 1995 Eq. 16 analytic approximation) for :math:\lambda_\mathrm{rest} < 911.75\,\text{Å}.

A small post-processing step caps τ at its peak short-wavelength value, suppressing the fitting-function breakdown at ~100 Å.

At zred = 0 τ is forced to zero — the Madau 1995 approximation is not physically valid for a line-of-sight of zero length and would otherwise give spurious residual attenuation to local- universe sources. This matches FSPS's own usage, which only invokes IGM_ABSORB when zred > 0.

make_igm_model

make_igm_model(name_or_model)

Build an :class:IGMModel from a string name, an instance, or None.

Parameters:

Name	Type	Description	Default
name_or_model	str, IGMModel, or None	String name (key of :data:_MODEL_REGISTRY), an already-built :class:IGMModel instance (returned as-is), or None (returns :class:NoIGM).	required

Examples:

>>> make_igm_model("madau1995")      # -> Madau1995()
>>> make_igm_model(None)             # -> NoIGM()
>>> make_igm_model(Madau1995())      # -> same instance

Source code in ceridwen/igm.py

def make_igm_model(name_or_model):
    """Build an :class:`IGMModel` from a string name, an instance, or
    ``None``.

    Parameters
    ----------
    name_or_model : str, IGMModel, or None
        String name (key of :data:`_MODEL_REGISTRY`), an already-built
        :class:`IGMModel` instance (returned as-is), or ``None``
        (returns :class:`NoIGM`).

    Examples
    --------
    >>> make_igm_model("madau1995")      # -> Madau1995()
    >>> make_igm_model(None)             # -> NoIGM()
    >>> make_igm_model(Madau1995())      # -> same instance
    """
    if name_or_model is None:
        return NoIGM()
    if isinstance(name_or_model, IGMModel):
        return name_or_model
    name = str(name_or_model).lower().strip()
    if name not in _MODEL_REGISTRY:
        raise ValueError(
            f"Unknown IGM model {name!r}.  "
            f"Available: {sorted(_MODEL_REGISTRY)}"
        )
    return _MODEL_REGISTRY[name]()

ceridwen.cosmology

ceridwen/cosmology.py

Redshift helpers for ceridwen, with two back-ends:

1. JAX native (default, always available) — flat LambdaCDM, matter only, Planck-2018-like parameters, Simpson integrator on a fixed 128-node grid. JIT-compilable, differentiable in z — required whenever zred is a sampled free parameter (NUTS needs gradients).

2. astropy (optional, auto-detected) — uses astropy.cosmology.Planck18.luminosity_distance and therefore includes radiation + massive neutrinos, closing the ~0.5-1% gap between the native integrator and published Planck18 tables. astropy's code is NumPy-based, not JAX — so it can only be used when z is a concrete Python scalar (not a traced value), typically for the one-off precompute step at SedModel(csp, obs, priors, zred=0.5). Set backend='astropy' in the top-level helpers to opt in explicitly.

References

Hogg, "Distance measures in cosmology" (astro-ph/9905116). Planck Collaboration 2020, A&A 641, A6 (Planck 2018 VI). astropy Cosmology docs: https://docs.astropy.org/en/stable/cosmology/

Conventions

• distances in megaparsec (Mpc); luminosity distance D_L(z) in Mpc.
• Native back-end integrates :math:\int_0^z dz'/E(z') with Simpson's rule on a fixed n_nodes grid — fully differentiable, constant XLA graph cost independent of z.

Use

from ceridwen.cosmology import luminosity_distance_mpc luminosity_distance_mpc(0.5) # JAX, any z luminosity_distance_mpc(0.5, backend='astropy') # Planck18 exact, scalar z only

Cosmology dataclass

Cosmology(
    H0=67.66,
    Om0=0.30966,
    Tcmb0=2.7255,
    Neff=3.046,
    m_nu_ev_sum=0.06,
)

Flat LambdaCDM cosmology with photons + effective neutrinos.

This is a deliberately-minimal JAX-friendly port of astropy's FlatLambdaCDM. The massive-neutrino contribution is folded into Om0 as cold matter (the late-time approximation; valid for m_nu <= 1 eV at z <= 10 to <0.1% on D_L), so we do NOT need the full relativistic transition that astropy evaluates point-by-point.

All properties are Python scalars and constant-fold into the XLA graph when this object is passed as a closure to a jitted function.

Ogamma0 property

Ogamma0

Photon density today :math:\Omega_\gamma.

Onu0_massive_as_matter property

Onu0_massive_as_matter

Massive-neutrino density today treated as cold matter (Mangano et al. 2005 relation :math:\Omega_\nu h^2 = \sum m_\nu/93.14\,\mathrm{eV}).

Onu0_relativistic property

Onu0_relativistic

Relativistic (massless-equivalent) neutrino density today.

For Planck18 with one 0.06 eV neutrino, Neff = 3.046 counts the total at CMB epoch but one species becomes non-relativistic at late times. Subtract one for the late-time relativistic count.

Om0_eff property

Om0_eff

Effective matter density today (CDM + baryons + massive nu).

Or0 property

Or0

Radiation + massless-neutrino density today.

Ode0 property

Ode0

Dark-energy density today (flat: 1 - Om - Or).

E_of_z

E_of_z(z, cosmo=DEFAULT_COSMO)

Dimensionless expansion rate

.. math:: E(z) = \sqrt{\Omega_m^\mathrm{eff}(1+z)^3 + \Omega_r (1+z)^4 + \Omega_\Lambda}.

Includes radiation (photons + massless neutrinos) exactly; the single massive neutrino is folded into :math:\Omega_m^\mathrm{eff} as cold matter. This matches :class:astropy.cosmology.FlatLambdaCDM with the same cosmological parameters to < 0.1% on :math:D_L(z) for :math:z \le 10.

The input is clamped to z >= 0 before use. VI ELBO training and NUTS leapfrog steps can transiently explore z < 0 while the proposal is being built, and at z < -1 the radicand :math:\Omega_m(1+z)^3 + \Omega_r(1+z)^4 + \Omega_\Lambda can go negative, giving sqrt(negative) = NaN. A single NaN in the ELBO gradient poisons the entire TriL map and causes 100% divergences downstream. Clamping at 0 keeps the trace finite; the proposal is still rejected by whatever prior / transform is in play.

The formula is identical to astropy's efunc implementation modulo the massive-neutrino approximation; structurally

.. math:: E(z)^2 = (1+z)^3\bigl[\Omega_r(1+z) + \Omega_m^\mathrm{eff}\bigr] + \Omega_\Lambda,

which is the form astropy uses to factor out one (1+z)^3 for numerical stability at high z.

Source code in ceridwen/cosmology.py

def E_of_z(z: Array, cosmo: Cosmology = DEFAULT_COSMO) -> Array:
    r"""Dimensionless expansion rate

    .. math::
        E(z) = \sqrt{\Omega_m^\mathrm{eff}(1+z)^3
                   + \Omega_r (1+z)^4
                   + \Omega_\Lambda}.

    Includes radiation (photons + massless neutrinos) exactly; the
    single massive neutrino is folded into :math:`\Omega_m^\mathrm{eff}`
    as cold matter.  This matches
    :class:`astropy.cosmology.FlatLambdaCDM` with the same cosmological
    parameters to < 0.1% on :math:`D_L(z)` for :math:`z \le 10`.

    The input is clamped to ``z >= 0`` before use.  VI ELBO training
    and NUTS leapfrog steps can transiently explore ``z < 0`` while
    the proposal is being built, and at ``z < -1`` the radicand
    :math:`\Omega_m(1+z)^3 + \Omega_r(1+z)^4 + \Omega_\Lambda` can go
    negative, giving ``sqrt(negative) = NaN``.  A single NaN in the
    ELBO gradient poisons the entire TriL map and causes 100%
    divergences downstream.  Clamping at 0 keeps the trace finite;
    the proposal is still rejected by whatever prior / transform is
    in play.

    The formula is identical to astropy's ``efunc`` implementation
    modulo the massive-neutrino approximation; structurally

    .. math::
        E(z)^2 = (1+z)^3\bigl[\Omega_r(1+z) + \Omega_m^\mathrm{eff}\bigr]
                 + \Omega_\Lambda,

    which is the form astropy uses to factor out one ``(1+z)^3`` for
    numerical stability at high z.
    """
    z = jnp.maximum(z, 0.0)
    opz = 1.0 + z
    return jnp.sqrt(
        opz * opz * opz * (cosmo.Or0 * opz + cosmo.Om0_eff)
        + cosmo.Ode0
    )

comoving_distance_mpc

comoving_distance_mpc(z, cosmo=DEFAULT_COSMO, n_nodes=128)

Comoving (line-of-sight) distance :math:D_C(z) = D_H \int_0^z dz'/E(z').

Source code in ceridwen/cosmology.py

def comoving_distance_mpc(z: Array, cosmo: Cosmology = DEFAULT_COSMO,
                          n_nodes: int = 128) -> Array:
    r"""Comoving (line-of-sight) distance
    :math:`D_C(z) = D_H \int_0^z dz'/E(z')`.
    """
    return cosmo.hubble_distance_mpc * _integrate_dz_over_E(z, cosmo, n_nodes)

age_gyr

age_gyr(z, cosmo=DEFAULT_COSMO, n_nodes=257)

Age of the Universe at redshift z in Gyr (JAX-native).

.. math:: t(z) = \frac{1}{H_0}\int_0^{a(z)}\frac{da'}{a'\,E(z'(a'))}, \qquad a = \frac{1}{1+z},\; z'(a') = 1/a' - 1.

Simpson's rule on a fixed scale-factor grid (traced shape static), so the result is differentiable in z and usable inside a JIT'd model (e.g. to recompute SFH age-bins from a sampled redshift). The integrand vanishes analytically as :math:a'\to 0 (radiation era), so the a'=0 node is set to zero to avoid 0/0. Matches :meth:astropy.cosmology.FlatLambdaCDM.age to <~0.3% over :math:0 \le z \le 20 (residual dominated by the massive-neutrino approximation in :func:E_of_z, not the integrator).

Source code in ceridwen/cosmology.py

def age_gyr(z: Array, cosmo: Cosmology = DEFAULT_COSMO,
            n_nodes: int = 257) -> Array:
    r"""Age of the Universe at redshift ``z`` in **Gyr** (JAX-native).

    .. math::
        t(z) = \frac{1}{H_0}\int_0^{a(z)}\frac{da'}{a'\,E(z'(a'))},
        \qquad a = \frac{1}{1+z},\; z'(a') = 1/a' - 1.

    Simpson's rule on a fixed scale-factor grid (traced shape static), so
    the result is differentiable in ``z`` and usable inside a JIT'd model
    (e.g. to recompute SFH age-bins from a sampled redshift).  The
    integrand vanishes analytically as :math:`a'\to 0` (radiation era),
    so the ``a'=0`` node is set to zero to avoid ``0/0``.  Matches
    :meth:`astropy.cosmology.FlatLambdaCDM.age` to <~0.3% over
    :math:`0 \le z \le 20` (residual dominated by the massive-neutrino
    approximation in :func:`E_of_z`, not the integrator).
    """
    if n_nodes % 2 == 0:
        n_nodes += 1
    z = jnp.asarray(z, dtype=float)
    a = 1.0 / (1.0 + jnp.maximum(z, 0.0))          # scale factor (...,)
    u = jnp.linspace(0.0, 1.0, n_nodes)            # (n_nodes,)
    x = a[..., None] * u                           # (..., n_nodes): a' from 0..a
    # Double-``where`` so neither the value nor the GRADIENT sees the 0/0 at
    # a'=0: evaluate with a safe x=1 there, then mask the contribution to 0.
    pos = x > 0.0
    x_safe = jnp.where(pos, x, 1.0)
    zp = 1.0 / x_safe - 1.0
    g = jnp.where(pos, 1.0 / (x_safe * E_of_z(zp, cosmo)), 0.0)
    w = jnp.ones(n_nodes).at[1::2].set(4.0).at[2:-1:2].set(2.0)
    h = a / (n_nodes - 1)
    integral = (h / 3.0) * jnp.sum(w * g, axis=-1)
    return (_INV_H0_GYR_NUM / cosmo.H0) * integral

luminosity_distance_mpc

luminosity_distance_mpc(
    z, cosmo=DEFAULT_COSMO, n_nodes=128, backend="native"
)

Luminosity distance :math:D_L(z) = (1+z) D_C(z) in Mpc.

Parameters:

Name	Type	Description	Default
z	Array or float	Redshift. JAX arrays are supported by the native backend; scalars only for the astropy backend.	required
cosmo	Cosmology	Only used by the native backend (astropy uses Planck18 internally).	DEFAULT_COSMO
n_nodes	int	Simpson-integrator node count (native only). 128 is accurate to ~1e-6 for the integral itself; the residual error vs astropy is dominated by missing radiation + neutrinos, not the integrator.	128
backend	('native', 'astropy')	'native' (default) uses the JAX Simpson integrator and is differentiable; 'astropy' calls :func:astropy.cosmology.Planck18.luminosity_distance and is limited to concrete scalar z but is "ground-truth" accurate.	'native'

Source code in ceridwen/cosmology.py

def luminosity_distance_mpc(z, cosmo: Cosmology = DEFAULT_COSMO,
                            n_nodes: int = 128,
                            backend: str = "native") -> Array:
    r"""Luminosity distance :math:`D_L(z) = (1+z) D_C(z)` in Mpc.

    Parameters
    ----------
    z : Array or float
        Redshift.  JAX arrays are supported by the native backend;
        scalars only for the astropy backend.
    cosmo : Cosmology
        Only used by the native backend (astropy uses Planck18 internally).
    n_nodes : int
        Simpson-integrator node count (native only).  128 is accurate to
        ~1e-6 for the integral itself; the residual error vs astropy is
        dominated by missing radiation + neutrinos, not the integrator.
    backend : {'native', 'astropy'}
        ``'native'`` (default) uses the JAX Simpson integrator and is
        differentiable; ``'astropy'`` calls
        :func:`astropy.cosmology.Planck18.luminosity_distance` and is
        limited to concrete scalar z but is "ground-truth" accurate.
    """
    if backend == "astropy":
        return _astropy_luminosity_distance_mpc(z)
    if backend != "native":
        raise ValueError(f"Unknown cosmology backend {backend!r}")
    return (1.0 + z) * comoving_distance_mpc(z, cosmo, n_nodes)

flux_factor

flux_factor(z, cosmo=DEFAULT_COSMO, n_nodes=128)

Cosmological flux factor to turn rest-frame per-Hz luminosity into observed-frame flux density, including the :math:(1+z) term for :math:F_\nu.

Specifically, for a rest-frame luminosity per unit frequency :math:L_\nu^\mathrm{rest} in erg s^{-1} Hz^{-1} one has

.. math:: F_\nu^\mathrm{obs}(\nu_\mathrm{obs}) = \frac{(1+z)\, L_\nu^\mathrm{rest}(\nu_\mathrm{rest})} {4 \pi\, D_L(z)^2}.

This helper returns :math:(1+z) / (4\pi D_L^2) in CGS-compatible units such that multiplying a spectrum in erg s^{-1} Hz^{-1} yields the observed-frame :math:F_\nu in erg s^{-1} cm^{-2} Hz^{-1}.

For the ceridwen "maggies"-style convention where the rest-frame spectrum is in L_sun Hz^{-1} M_sun^{-1} (as produced by FSPS / SSP tables), this factor should be composed with any unit conversion the forward model already does. See :func:flux_factor_maggies for the maggies-ready version.

Source code in ceridwen/cosmology.py

def flux_factor(z: Array, cosmo: Cosmology = DEFAULT_COSMO,
                n_nodes: int = 128) -> Array:
    r"""Cosmological flux factor to turn rest-frame per-Hz luminosity into
    observed-frame flux density, including the :math:`(1+z)` term for
    :math:`F_\nu`.

    Specifically, for a rest-frame luminosity per unit frequency
    :math:`L_\nu^\mathrm{rest}` in erg s^{-1} Hz^{-1} one has

    .. math::
        F_\nu^\mathrm{obs}(\nu_\mathrm{obs}) =
        \frac{(1+z)\, L_\nu^\mathrm{rest}(\nu_\mathrm{rest})}
             {4 \pi\, D_L(z)^2}.

    This helper returns :math:`(1+z) / (4\pi D_L^2)` in CGS-compatible units
    such that multiplying a spectrum in erg s^{-1} Hz^{-1} yields the
    observed-frame :math:`F_\nu` in erg s^{-1} cm^{-2} Hz^{-1}.

    For the ceridwen "maggies"-style convention where the rest-frame
    spectrum is in L_sun Hz^{-1} M_sun^{-1} (as produced by FSPS / SSP
    tables), this factor should be composed with any unit conversion
    the forward model already does.  See
    :func:`flux_factor_maggies` for the maggies-ready version.
    """
    dL_mpc = luminosity_distance_mpc(z, cosmo, n_nodes)
    dL_cm = dL_mpc * MPC_TO_CM
    return (1.0 + z) / (4.0 * jnp.pi * dL_cm * dL_cm)

flux_factor_maggies

flux_factor_maggies(
    z, cosmo=DEFAULT_COSMO, n_nodes=128, backend="native"
)

Redshift scaling relative to the CSP's fiducial d = 10 pc.

Multiplies a CSPBasis-computed (rest-frame, per-source) spectrum to produce the observed-frame spectrum at the luminosity distance for redshift z. At z = 0 the CSP normalisation already corresponds to 10 pc, so :math:\mathrm{ff}(0) = 1. For z > 0:

.. math:: \mathrm{ff}(z) = \frac{(1 + z)\, D_L(z=0)^2}{D_L(z)^2} = \frac{(1 + z)\,(10\,\mathrm{pc})^2}{D_L(z)^2}.

This assumes the observation flux unit is maggies relative to the same absolute magnitude convention.

Parameters:

Name	Type	Description	Default
z	Array or float	Redshift.	required
backend	('native', 'astropy')	See :func:luminosity_distance_mpc. Use 'astropy' for a one-off scalar at model construction time (e.g. inside SedModel.__init__ when zred is a fixed float); use 'native' inside the forward-evaluation loop so the call is JIT-compilable and differentiable. SedModel automatically applies the astropy backend when a float zred is provided and the package is installed.	'native'

Source code in ceridwen/cosmology.py

def flux_factor_maggies(z, cosmo: Cosmology = DEFAULT_COSMO,
                        n_nodes: int = 128,
                        backend: str = "native") -> Array:
    r"""Redshift scaling relative to the CSP's fiducial ``d = 10 pc``.

    Multiplies a CSPBasis-computed (rest-frame, per-source) spectrum to
    produce the observed-frame spectrum at the luminosity distance for
    redshift ``z``.  At ``z = 0`` the CSP normalisation already corresponds
    to 10 pc, so :math:`\mathrm{ff}(0) = 1`.  For ``z > 0``:

    .. math::
        \mathrm{ff}(z) = \frac{(1 + z)\, D_L(z=0)^2}{D_L(z)^2}
                      = \frac{(1 + z)\,(10\,\mathrm{pc})^2}{D_L(z)^2}.

    This assumes the observation flux unit is maggies relative to the
    same absolute magnitude convention.

    Parameters
    ----------
    z : Array or float
        Redshift.
    backend : {'native', 'astropy'}
        See :func:`luminosity_distance_mpc`.  Use ``'astropy'`` for a
        one-off scalar at model construction time (e.g. inside
        ``SedModel.__init__`` when ``zred`` is a fixed float); use
        ``'native'`` inside the forward-evaluation loop so the call is
        JIT-compilable and differentiable.  ``SedModel`` automatically
        applies the astropy backend when a ``float`` ``zred`` is
        provided and the package is installed.
    """
    dL_pc = 1e6 * luminosity_distance_mpc(z, cosmo, n_nodes, backend=backend)
    # z <= 0 has no cosmological dimming in this convention (the CSP
    # normalisation is already "source at 10 pc"), so the distance ratio is
    # pinned to exactly 1 there.  The gate is on ``z``, NOT on ``dL_pc``:
    # gating on ``dL_pc > 0`` (the previous behaviour) silently rerouted
    # every z > 0 evaluation whose distance came out non-positive or NaN
    # to the 10 pc fallback, erasing the entire (10pc/D_L)^2 dimming —
    # a factor ~2e15 in flux at z = 0.1 — with no error raised.  With the
    # gate on ``z``, a broken D_L at z > 0 now propagates as NaN/inf and
    # fails loudly.  Double-``where`` so neither the value nor the GRADIENT
    # touches the division at z <= 0 (where D_L = 0 would give 0/0).
    positive_z = z > 0
    safe_dL = jnp.where(positive_z, dL_pc, float(_MAGGIES_D_FID_PC))
    ff_distance = jnp.where(
        positive_z,
        (1.0 + z) * (_MAGGIES_D_FID_PC / safe_dL) ** 2,
        1.0,
    )
    # Close the L_sun/Hz -> erg/s/cm^2/Hz gap so downstream filter projection
    # via sedpy_jax lands in the AB zero-point frame used by the Photometry
    # data side.  Without this, ceridwen's predicted maggies are ~3e-7 x the
    # true values and any fit against real photometry pushes logmass straight
    # against its prior's upper bound.
    return ff_distance * _LSUN_HZ_TO_FNU_CGS_AT_10PC

have_astropy

have_astropy()

True if astropy is importable. Used by SedModel to decide whether to use the higher-accuracy Planck18 backend for fixed-z precompute.

Source code in ceridwen/cosmology.py

def have_astropy() -> bool:
    """True if astropy is importable.  Used by SedModel to decide whether
    to use the higher-accuracy Planck18 backend for fixed-z precompute."""
    try:
        import astropy.cosmology  # noqa: F401
        return True
    except Exception:
        return False