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,
    ssp_resolution=None,
    resolution_source=None,
    isoc_type=None,
    spec_library=None,
    imf_type=None,
    fsps_version=None,
    fsps_kwargs=dict(),
    wave_min=None,
    wave_max=None,
    schema_version=None,
    _extra_datasets=(),
)

Immutable container for the SSP interpolation grids plus static provenance.

Parameters:

Name	Type	Description	Default
ssp_lgmet	array (n_met,) -- log10 absolute metallicity Z, NOT log10 Z/Zsun		required
ssp_lg_age_gyr	array(n_ages) - -log10(age / Gyr)		required
ssp_wave	(array(n_wave), Angstrom)		required
ssp_flux	array (n_met, n_ages, n_wave), L_sun/Hz per M_sun formed		required
ssp_resolution	ndarray (n_wave,), km/s -- library sigma_v(lambda) on ssp_wave, NaN where unknown; optional in memory, required by save()/load()		None

with_resolution

with_resolution(
    *, sigma_v=None, segments=None, source=None
)

Return a copy carrying a library resolution curve from exactly one of sigma_v (km/s on ssp_wave, NaN where unknown) or segments.

Source code in ceridwen/ssps/ssp_data.py

def with_resolution(self, *, sigma_v=None, segments=None, source=None):
    """Return a copy carrying a library resolution curve from exactly one of ``sigma_v`` (km/s on ssp_wave, NaN where unknown) or ``segments``."""
    import dataclasses as _dc
    from .library_resolution import sigma_v_from_segments
    if (sigma_v is None) == (segments is None):
        raise ValueError(
            "with_resolution: pass exactly one of sigma_v= or segments=")
    if segments is not None:
        sigma_v = sigma_v_from_segments(
            np.asarray(self.ssp_wave), segments)
    return _dc.replace(
        self,
        ssp_resolution=np.asarray(sigma_v, dtype=np.float64),
        resolution_source=(str(source) if source is not None else None),
    )

display

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

Print (or return as str when return_str) a summary of the grid and its provenance.

Source code in ceridwen/ssps/ssp_data.py

def display(self, *, return_str: bool = False, file=None):
    """Print (or return as str when ``return_str``) a summary of the grid and its provenance."""
    import sys as _sys

    wave  = np.asarray(self.ssp_wave)

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

    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 = [
        self._display_title,
        "-" * 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",
    ] + self._display_grid_lines(size_str)
    if self.ssp_resolution is None:
        lines += ["  library resolution       : MISSING "
                  "(cannot be saved; attach with with_resolution)"]
    else:
        res = np.asarray(self.ssp_resolution, dtype=np.float64)
        fin = np.isfinite(res)
        if fin.any():
            cov_lo = wave[fin].min(); cov_hi = wave[fin].max()
            lines += [
                f"  library resolution       : sigma_v "
                f"[{res[fin].min():.1f}, {res[fin].max():.1f}] km/s over "
                f"[{cov_lo:.0f}, {cov_hi:.0f}] AA "
                f"({100.0 * fin.mean():.0f}% of pixels; NaN elsewhere)",
            ]
        else:
            lines += ["  library resolution       : all-NaN "
                      "(unknown everywhere; no subtraction will occur)"]
        if self.resolution_source:
            lines += [f"  resolution source        : {self.resolution_source}"]
    lines += self._display_note_lines()

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

save

save(filename)

Write grids, resolution curve and provenance attrs to HDF5 (overwrites); raises ValueError if ssp_resolution is None.

Source code in ceridwen/ssps/ssp_data.py

def save(self, filename):
    """Write grids, resolution curve and provenance attrs to HDF5 (overwrites); raises ValueError if ``ssp_resolution`` is None."""
    if self.ssp_resolution is None:
        raise ValueError(
            f"{type(self).__name__}.save(): this grid carries no library "
            f"resolution curve (ssp_resolution is None).  "
            f"{self._schema_label} files require one — attach it with "
            "with_resolution(segments=...) or with_resolution(sigma_v=...) "
            "before saving."
        )
    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.create_dataset('ssp_resolution',
                         data=np.asarray(self.ssp_resolution,
                                         dtype=np.float64))

        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'
        f.attrs['units_resolution']   = 'sigma_v [km/s]; NaN = unknown'
        self._save_extra(f)
        if self.resolution_source is not None:
            f.attrs['resolution_source'] = str(self.resolution_source)

        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)
        f.attrs['fsps_kwargs_json'] = json.dumps(self.fsps_kwargs or {})

load classmethod

load(filename, flux_dtype=None)

Load a grid from HDF5; raises ValueError if the file lacks ssp_resolution. flux_dtype casts the flux cube on read.

Source code in ceridwen/ssps/ssp_data.py

@classmethod
def load(cls, filename, flux_dtype=None):
    """Load a grid from HDF5; raises ValueError if the file lacks ``ssp_resolution``. ``flux_dtype`` casts the flux cube on read."""
    arrays, _extra, meta = cls._read_h5(filename, flux_dtype=flux_dtype)
    return cls(**arrays, **meta)

from_fsps classmethod

from_fsps(
    save_to=None,
    resolution_segments=None,
    resolution_source=None,
    **fsps_kwargs
)

Build a grid from the SPS backend, attach the library resolution curve and record provenance.

Only library/IMF-defining kwargs (LIBRARY_IMF_KWARGS) are accepted. The resolution curve is the element-wise maximum of the 2-pixel sampling floor of ssp_wave and resolution_segments (if given, which then require resolution_source). Raises ValueError otherwise.

Source code in ceridwen/ssps/ssp_data.py

@classmethod
def from_fsps(cls, save_to: Optional[str] = None,
              resolution_segments=None,
              resolution_source: Optional[str] = None,
              **fsps_kwargs) -> "SSPData":
    """Build a grid from the SPS backend, attach the library resolution curve and record provenance.

    Only library/IMF-defining kwargs (``LIBRARY_IMF_KWARGS``) are accepted.
    The resolution curve is the element-wise maximum of the 2-pixel sampling
    floor of ``ssp_wave`` and ``resolution_segments`` (if given, which then
    require ``resolution_source``). Raises ValueError otherwise.
    """
    if resolution_source is not None and resolution_segments is None:
        raise ValueError(
            "from_fsps(): resolution_source was given without "
            "resolution_segments.  The sampling-floor provenance is "
            "recorded automatically; a source note only accompanies "
            "explicit library-LSF segments."
        )
    if resolution_segments is not None and resolution_source is None:
        raise ValueError(
            "from_fsps(): resolution_segments were given without "
            "resolution_source.  Cite where the LSF numbers come from "
            "(e.g. resolution_source='MILES FWHM 2.54A "
            "(Falcon-Barroso et al. 2011)') — uncited resolution "
            "numbers must not ship in a released grid."
        )
    from .library_resolution import combined_sigma_v, combined_source
    data = collect_ssp_data_wrapper(**fsps_kwargs)
    sigma_v = combined_sigma_v(np.asarray(data.ssp_wave),
                               segments=resolution_segments)
    data = data.with_resolution(
        sigma_v=sigma_v, source=combined_source(resolution_source))
    if save_to is not None:
        data.save(save_to)
    return data

ceridwen.ssps.SSPDataAfe dataclass

SSPDataAfe(
    ssp_lgmet,
    ssp_lg_age_gyr,
    ssp_wave,
    ssp_flux,
    ssp_resolution=None,
    resolution_source=None,
    isoc_type=None,
    spec_library=None,
    imf_type=None,
    fsps_version=None,
    fsps_kwargs=dict(),
    wave_min=None,
    wave_max=None,
    schema_version=None,
    _extra_datasets=(),
    *,
    ssp_afe
)

Bases: SSPData

Immutable alpha-enhanced SSP grid: an SSPData whose flux cube has a leading [alpha/Fe] axis.

Parameters:

Name	Type	Description	Default
ssp_lgmet	jnp.ndarray (n_met,) -- log10 absolute TOTAL metallicity Z, NOT [Fe/H]		required
ssp_afe	jnp.ndarray (n_afe,), keyword-only -- [alpha/Fe] grid, strictly increasing		required
ssp_lg_age_gyr	ndarray(n_ages) - -log10(age / Gyr)		required
ssp_wave	ndarray(n_wave) - -Angstrom		required
ssp_flux	jnp.ndarray (n_afe, n_met, n_ages, n_wave) -- Lsun / Hz per Msun formed		required
ssp_resolution	np.ndarray (n_wave,) or None -- sigma_v(lambda) [km/s] on ``ssp_wave``; required by save/load		None

load classmethod

load(filename, flux_dtype=None)

Load from HDF5; a 3-D grid without ssp_afe is promoted to n_afe = 1 at [alpha/Fe] = 0.

Source code in ceridwen/ssps/ssp_data_afe.py

@classmethod
def load(cls, filename, flux_dtype=None):
    """Load from HDF5; a 3-D grid without ``ssp_afe`` is promoted to n_afe = 1 at [alpha/Fe] = 0."""
    arrays, extra, meta = cls._read_h5(filename, flux_dtype=flux_dtype)
    ssp_afe = extra['ssp_afe']
    if ssp_afe is None:
        ssp_afe = jnp.zeros(1)
        arrays['ssp_flux'] = arrays['ssp_flux'][None, ...]
    return cls(**arrays, ssp_afe=ssp_afe, **meta)

from_fsps classmethod

from_fsps(
    save_to=None,
    afe_values=None,
    resolution_segments=None,
    resolution_source=None,
    **fsps_kwargs
)

Build the grid by looping get_spectrum(tage=0, zmet=i) over the [alpha/Fe] planes.

Parameters:

Name	Type	Description	Default
afe_values	array-like, optional -- [alpha/Fe] of each plane in ``afeindx`` order; defaults exist for n_afe 5 and 1 only		None
resolution_segments	list, optional -- documented library LSF segments, max-combined with the 2-pixel sampling floor		None
resolution_source	str, optional -- citation for ``resolution_segments``; each raises if given without the other		None

Source code in ceridwen/ssps/ssp_data_afe.py

@classmethod
def from_fsps(cls, save_to: Optional[str] = None,
              afe_values=None,
              resolution_segments=None,
              resolution_source: Optional[str] = None,
              **fsps_kwargs) -> "SSPDataAfe":
    """Build the grid by looping ``get_spectrum(tage=0, zmet=i)`` over the [alpha/Fe] planes.

    Parameters
    ----------
    afe_values : array-like, optional -- [alpha/Fe] of each plane in ``afeindx`` order; defaults exist for n_afe 5 and 1 only
    resolution_segments : list, optional -- documented library LSF segments, max-combined with the 2-pixel sampling floor
    resolution_source : str, optional -- citation for ``resolution_segments``; each raises if given without the other
    """
    if resolution_source is not None and resolution_segments is None:
        raise ValueError(
            "from_fsps(): resolution_source was given without "
            "resolution_segments.  The sampling-floor provenance is "
            "recorded automatically; a source note only accompanies "
            "explicit library-LSF segments."
        )
    if resolution_segments is not None and resolution_source is None:
        raise ValueError(
            "from_fsps(): resolution_segments were given without "
            "resolution_source.  Cite where the LSF numbers come from "
            "— uncited resolution numbers must not ship in a released "
            "grid."
        )
    for bad in ('afe', 'afeindx'):
        if bad in fsps_kwargs:
            raise ValueError(
                f"from_fsps() rejects the FSPS kwarg {bad!r}: the "
                f"[alpha/Fe] axis is spanned by the grid itself "
                f"(ssp_afe); CSPBasis_afe samples it through "
                f"theta['afe'], so a fixed alpha must not be set at "
                f"build time."
            )
    kwargs = _validate_fsps_kwargs(fsps_kwargs)

    fsps = _import_fsps()

    ssp = fsps.StellarPopulation(zcontinuous=0, sfh=0, **kwargs)

    n_afe = int(getattr(ssp, "n_afe", 1))
    if afe_values is not None:
        ssp_afe = np.atleast_1d(np.asarray(afe_values, dtype=float))
        if ssp_afe.size != n_afe:
            raise ValueError(
                f"afe_values has {ssp_afe.size} entries but the "
                f"compiled FSPS grid has n_afe = {n_afe}."
            )
    elif n_afe == 5:
        ssp_afe = FSPS_AFE_VALUES_NAFE5.copy()
    elif n_afe == 1:
        ssp_afe = np.zeros(1)
    else:
        raise ValueError(
            f"Compiled FSPS grid has n_afe = {n_afe}, which does not "
            f"match the documented aMIST/C3K layout (5) or a "
            f"solar-scaled build (1); pass afe_values= explicitly "
            f"(the [alpha/Fe] of each afeindx plane, in order)."
        )

    ssp_lgmet      = jnp.log10(ssp.zlegend)
    nzmet          = int(ssp_lgmet.size)
    ssp_lg_age_gyr = jnp.array(ssp.log_age - 9.0)

    planes = []
    _wave = None
    for afe_indx in range(1, n_afe + 1):               # afeindx is 1-based
        if n_afe > 1:
            ssp.params["afeindx"] = afe_indx
            print(f"[alpha/Fe] plane {afe_indx}/{n_afe} "
                  f"(afe = {ssp_afe[afe_indx - 1]:+.1f})")
        spectrum_collector = []
        for zmet_indx in range(1, nzmet + 1):
            print(f"...retrieving metallicity {zmet_indx}/{nzmet} "
                  f"[Z = {ssp.zlegend[zmet_indx - 1]:.4f}]")
            _wave, _fluxes = ssp.get_spectrum(
                tage=0.0, zmet=zmet_indx, peraa=False)
            spectrum_collector.append(_fluxes)
        planes.append(np.array(spectrum_collector))

        _z_now = np.log10(np.asarray(ssp.zlegend))
        if not np.allclose(_z_now, np.asarray(ssp_lgmet)):
            raise RuntimeError(
                f"zlegend changed between alpha planes (afeindx="
                f"{afe_indx}); the Z grids are ragged across "
                f"[alpha/Fe] and cannot form a rectangular "
                f"(afe, Z, age, wave) grid."
            )

    ssp_wave = jnp.array(_wave)
    ssp_flux = jnp.array(np.stack(planes, axis=0))

    meta = _read_fsps_provenance(ssp, kwargs, ssp_wave)
    meta['schema_version'] = SSP_AFE_SCHEMA_VERSION

    from .library_resolution import combined_sigma_v, combined_source
    sigma_v = combined_sigma_v(np.asarray(ssp_wave),
                               segments=resolution_segments)

    data = cls(ssp_lgmet, ssp_lg_age_gyr, ssp_wave, ssp_flux,
               ssp_afe=jnp.array(ssp_afe),
               ssp_resolution=sigma_v,
               resolution_source=combined_source(resolution_source),
               **meta)
    if save_to is not None:
        data.save(save_to)
    return data

ceridwen.ssps.fetch_grid

fetch_grid(name, *, force=False, quiet=False)

Return a local, checksum-verified path to the registry grid name, downloading into :func:grid_cache_dir on first use (force re-downloads).

Source code in ceridwen/ssps/grid_fetch.py

def fetch_grid(name: str, *, force: bool = False, quiet: bool = False) -> Path:
    """Return a local, checksum-verified path to the registry grid ``name``,
    downloading into :func:`grid_cache_dir` on first use (``force`` re-downloads)."""
    if name not in REGISTRY:
        raise KeyError(
            f"Unknown grid {name!r}. Available: {sorted(REGISTRY)}."
        )
    entry = REGISTRY[name]
    if entry["url"] is None:
        raise RuntimeError(
            f"Grid {name!r} is defined but not yet published (no URL in the "
            f"registry). Build it locally with scripts_afe/build_afe_grid.py, "
            f"or publish it with scripts_afe/publish_grid_zenodo.py and "
            f"paste the printed REGISTRY entry. Notes: {entry['notes']}"
        )

    dest = grid_cache_dir() / f"{name}.h5"

    if dest.exists() and not force:
        got = _sha256(dest)
        if entry["sha256"] and got != entry["sha256"]:
            raise RuntimeError(
                f"Cached grid {dest} fails its checksum "
                f"(got {got[:12]}..., expected {entry['sha256'][:12]}...). "
                f"Delete it or call fetch_grid({name!r}, force=True)."
            )
        return dest

    if not quiet:
        size = f" (~{entry['size_mb']} MB)" if entry.get("size_mb") else ""
        print(f"[ceridwen] fetching grid {name!r}{size} -> {dest}",
              file=sys.stderr)

    fd, tmp = tempfile.mkstemp(dir=dest.parent, suffix=".part")
    os.close(fd)
    tmp = Path(tmp)
    try:
        with urllib.request.urlopen(entry["url"]) as r, open(tmp, "wb") as f:
            shutil.copyfileobj(r, f, length=1 << 20)
        got = _sha256(tmp)
        if entry["sha256"] and got != entry["sha256"]:
            raise RuntimeError(
                f"Downloaded grid {name!r} fails its checksum "
                f"(got {got[:12]}..., expected {entry['sha256'][:12]}...). "
                f"The remote file changed or the download was corrupted; "
                f"not installing it."
            )
        tmp.replace(dest)
    finally:
        if tmp.exists():
            tmp.unlink()

    return dest

ceridwen.ssps.available_grids

available_grids(published_only=False)

Map of grid name -> one-line description.

Source code in ceridwen/ssps/grid_fetch.py

def available_grids(published_only: bool = False) -> dict[str, str]:
    """Map of grid name -> one-line description."""
    return {
        k: v["notes"] for k, v in REGISTRY.items()
        if v["url"] is not None or not published_only
    }

Composite stellar population (forward model)

ceridwen.csp.CSPBasis

CSPBasis(
    SSPData,
    theta=None,
    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",
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    fesc_geometry="runaway_bc",
    cosmo=None,
    **kwargs
)

Composite stellar population basis. predict(theta, observations) projects the model onto observations.

Parameters:

Name	Type	Description	Default
SSPData	SSPData -- SSP grids (wave, flux, ages, metallicities, optional resolution curve).		required
theta	dict -- initial values; must contain "lookback_time" (Gyr, increasing, index 0 = today,	= 2 nodes) and "sfh", plus "Z" (zh_const=True) or "zh" (zh_const=False). Mutually exclusive with lookback_time=.	None
lookback_time	array -- shortcut for ``theta``: the node grid only, neutral initial values	(sfh = 1, metallicity = grid median).	None
sfh_per_bin	bool -- with ``lookback_time=``, one SFR per bin (n_time-1,) instead of per node.		False
zh_const	bool -- constant metallicity ("Z") or a metallicity history ("zh", one per node).		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
add_igm	bool -- physics switches.		True
sps_home	str -- data directory for the nebular and dust-emission grids; defaults to $SPS_HOME.		None
init_neb_params	dict -- forwarded to NebularModel / Dust. ``isoc_type`` is	taken from the SSP grid's provenance when recorded.	None
init_dust_params	dict -- forwarded to NebularModel / Dust. ``isoc_type`` is	taken from the SSP grid's provenance when recorded.	None
sfh_interp	(step, linear)	piecewise-linear (analytic log-age integral, small negative weights clipped) SFH.	'step'
track_zred_age	bool -- with a sampled ``zred``, rescale the lookback grid so its oldest node	is the age of the Universe at that redshift.	False
fesc_geometry	(runaway_bc, picket)		'runaway_bc'
cosmo	Cosmology -- required; used for the flux factor and for the age of the Universe.		None
nebemlineinspec	bool -- default of ``include_lines`` in ``get_spectrum`` only.		False

Source code in ceridwen/csp/csp.py

def __init__(
    self,
    SSPData,
    theta=None,
    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',
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    fesc_geometry="runaway_bc",
    cosmo=None,
    **kwargs,
):
    self.verbose = bool(verbose)
    if kwargs:
        hint = ""
        if "sigma_losvd_kms" in kwargs:
            hint = (" ('sigma_losvd_kms' was removed: the galaxy velocity "
                    "dispersion is set once on the model, "
                    "SedModel(kinematics=Kinematics(sigma_gal=...)))")
        elif "tuniv" in kwargs:
            hint = (" ('tuniv' was removed: the age of the Universe comes from "
                    "the cosmology, csp.age_at(z) / cosmo.age(z))")
        raise TypeError(
            f"CSPBasis got unexpected keyword argument(s) {sorted(kwargs)}{hint}")
    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:
        init_neb_params = {"cloudy_dust": True}
    if init_dust_params is None:
        init_dust_params = {'bin_edges': [(-jnp.inf, -1.97)], 'laws': ['powerlaw']}

    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
    self.lib_resolution = (
        (np.asarray(SSPData.ssp_wave, dtype=np.float64),
         np.asarray(SSPData.ssp_resolution, dtype=np.float64))
        if getattr(SSPData, "ssp_resolution", None) is not None else None)
    self.zlegend   = 10 ** self.zmet                   # linear metallicity
    self.ssp_ages_lgyr = self.ages + 9                 # log10(yr)

    self._ssp_isoc_type    = getattr(SSPData, "isoc_type", None)
    self._ssp_spec_library = getattr(SSPData, "spec_library", None)

    if fesc_geometry not in ("runaway_bc", "picket"):
        raise ValueError(
            f"fesc_geometry must be 'runaway_bc' or 'picket', got {fesc_geometry!r}")
    self.fesc_geometry = str(fesc_geometry)

    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)

    self._ssp_lo_yr = 10.0 ** self._logage_hi   # (n_age-1,)
    self._ssp_hi_yr = 10.0 ** self._logage_lo   # (n_age-1,)

    _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]
    )
    self._ssp_voronoi_hi = jnp.concatenate(
        [_voro_mid, jnp.array([_voro_hi_ext])]
    )

    self.tiny_logt  = tiny_logt
    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
    from ..cosmology import Cosmology as _Cosmology
    if cosmo is None:
        raise TypeError(
            f"{type(self).__name__} needs an explicit cosmology: pass "
            "cosmo=Cosmology.planck18(), Cosmology.wmap9(), "
            "Cosmology.flat(H0, Om0), Cosmology.from_name(...) or "
            "Cosmology.from_astropy(...)  (from ceridwen import Cosmology)")
    if not isinstance(cosmo, _Cosmology):
        raise TypeError(
            f"cosmo must be a ceridwen Cosmology, got {type(cosmo).__name__}; "
            "for an astropy cosmology use Cosmology.from_astropy(...)")
    self._cosmo = cosmo
    self.track_zred_age = bool(track_zred_age)
    self.nebemlineinspec = bool(nebemlineinspec)

    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)

    if add_diffuse_dust or add_dust:
        self.set_attenuation_function(add_diffuse_dust, add_dust)

    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)

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

    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()})

    self.check_param_ranges(self.theta_init)

all_params property

all_params

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

cosmo property writable

cosmo

The cosmology; fixed at construction (assignment raises).

initialize_model_structure

initialize_model_structure(theta)

Validate theta (grid, sfh shape, metallicity key) and build theta_init / param_names.

Source code in ceridwen/csp/csp.py

def initialize_model_structure(self, theta):
    """Validate ``theta`` (grid, sfh shape, metallicity key) and build ``theta_init`` / ``param_names``."""
    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'])."
        )

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

    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)."
        )

    _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))
    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)."
        )

    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,
        )

    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


    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

    self.theta_init['sfh'] = sfh

    self.param_names = list(self.theta_init.keys())

    self._known_theta_keys = set(self.param_names) | {
        'lookback_time', 'Z', 'zh',
        'logmass', 'zred', 'lumdist_mpc', 'igm_factor', 'eline_scaling',
        'frac_obrun', 'spectrum_scaling', 'spectrum_calib',
    }

register_known_theta_keys

register_known_theta_keys(keys)

Add keys that _warn_unknown_theta_keys must accept (model-level parameters consumed by transforms).

Source code in ceridwen/csp/csp.py

def register_known_theta_keys(self, keys):
    """Add keys that ``_warn_unknown_theta_keys`` must accept (model-level parameters consumed by transforms)."""
    self._known_theta_keys |= set(keys)

check_param_ranges

check_param_ranges(theta=None, warn=True)

Messages (and warnings) for metallicity / nebular parameters outside the interpolation grids, where the model clamps silently.

Source code in ceridwen/csp/csp.py

def check_param_ranges(self, theta=None, warn=True):
    """Messages (and warnings) for metallicity / nebular parameters outside the interpolation grids, where the model clamps silently."""
    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."
                )

    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)

Assign self.attenuate_dust(wave, theta) -> (attn_birthcloud (n_bins, n_wave), attn_diffuse (n_wave,)) as optical depths.

Source code in ceridwen/csp/csp.py

def set_attenuation_function(self, add_diffuse_dust, add_dust):
    """Assign ``self.attenuate_dust(wave, theta) -> (attn_birthcloud (n_bins, n_wave), attn_diffuse (n_wave,))`` as optical depths."""
    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
        if self.verbose:
            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((wave.shape[0],))
            return attn, attn_diffuse
        if self.verbose:
            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
        if self.verbose:
            print("Using only diffuse dust attenuation.")
        self.attenuate_dust = attenuate_diffuse_only

get_spectrum_components

get_spectrum_components(theta)

(continuum, lines) on the rest-frame grid, unscaled (no mass, distance or IGM): continuum is get_spectrum(include_lines=False) and lines the difference to the full spectrum (the painted lines; with dust emission also their re-emitted energy).

Source code in ceridwen/csp/csp.py

def get_spectrum_components(self, theta: dict) -> tuple:
    """``(continuum, lines)`` on the rest-frame grid, unscaled (no mass, distance or IGM):
    ``continuum`` is ``get_spectrum(include_lines=False)`` and ``lines`` the difference to the
    full spectrum (the painted lines; with dust emission also their re-emitted energy).
    """
    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,
    *,
    kinematics=None,
    broaden_photometry=False,
    eline_system=None
)

{obs.name: prediction}: Photometry -> maggies (n_filters,), Spectrum -> F_nu [erg/s/cm^2/Hz] on the observed pixels, Lines -> integrated fluxes (n_lines,). Observed-frame only when theta carries zred.

kinematics (Kinematics or None) supplies sigma_gal / sigma_gas; broaden_photometry applies them to the spectrum entering each Photometry's _broadener. Lines are painted on the model grid only when a consumer needs them there (free-z Photometry, Photometry with a sampled sigma_gas and broaden_photometry, or _force_paint_lines); otherwise fixed-z Photometry adds them through the static basis G = obs._T @ (IGM * profiles), Spectrum paints them on the observed pixels and Lines reads them from the grid.

eline_system (ElineSystem of SedModel, line marginalisation): returns (predictions, aux) with the fitted lines removed from every prediction and aux = {"prior_mean": CLOUDY fluxes of the fitted lines, "cols": {obs.name: design}}.

Source code in ceridwen/csp/csp.py

def predict(self, theta: dict, observations: list, *, kinematics=None,
            broaden_photometry=False, eline_system=None) -> dict:
    """``{obs.name: prediction}``: Photometry -> maggies (n_filters,), Spectrum -> F_nu [erg/s/cm^2/Hz]
    on the observed pixels, Lines -> integrated fluxes (n_lines,).  Observed-frame only when
    ``theta`` carries ``zred``.

    ``kinematics`` (``Kinematics`` or None) supplies sigma_gal / sigma_gas; ``broaden_photometry``
    applies them to the spectrum entering each Photometry's ``_broadener``.  Lines are painted on
    the model grid only when a consumer needs them there (free-z Photometry, Photometry with a
    sampled sigma_gas and ``broaden_photometry``, or ``_force_paint_lines``); otherwise fixed-z
    Photometry adds them through the static basis ``G = obs._T @ (IGM * profiles)``, Spectrum
    paints them on the observed pixels and Lines reads them from the grid.

    ``eline_system`` (``ElineSystem`` of ``SedModel``, line marginalisation): returns
    ``(predictions, aux)`` with the fitted lines removed from every prediction and
    ``aux = {"prior_mean": CLOUDY fluxes of the fitted lines, "cols": {obs.name: design}}``.
    """
    from ..observation.observation import (
        Spectrum as _Spectrum, Photometry as _Photometry,
    )
    if kinematics is not None:
        self._known_theta_keys |= set(kinematics.free_keys)
    has_neb = getattr(self, "neb", None) is not None
    gas_free = (kinematics is not None
                and isinstance(kinematics.effective_sigma_gas, str))
    paint_lines = has_neb and (
        getattr(self, "_force_paint_lines", False)
        or any(
            isinstance(o, _Photometry)
            and ((getattr(o, "free_z", False) and "zred" in theta)
                 or (gas_free and broaden_photometry))
            for o in observations
        )
    )
    cont, lines = self._assemble_components(theta, paint_lines)
    cont, lines_s = self._apply_mass_redshift_igm(
        cont, cont if lines is None else lines, theta)
    return self._project_observations(
        cont if lines is None else cont + lines_s, cont,
        observations, theta, paint_lines=paint_lines,
        line_component=None if lines is None else lines_s,
        kinematics=kinematics, broaden_photometry=broaden_photometry,
        eline_system=eline_system,
    )

predict_line_fluxes

predict_line_fluxes(
    theta, *, for_photometry=False, for_spectrum=False
)

Observed-frame fluxes of every nebular grid line (n_lines,), through the same weights, dust, escape, mass and distance factors as the spectrum.

Default: integrated fluxes [erg/s/cm^2] with the 1/(1+z) Jacobian, IGM at the line wavelength and eline_scaling (the Lines observation). for_spectrum: the same without eline_scaling (painted by the Spectrum projector). for_photometry: per-Hz amplitudes with the full f_nu flux factor and no IGM (applied inside the photometric line basis) and no eline_scaling.

Source code in ceridwen/csp/csp.py

def predict_line_fluxes(self, theta, *, for_photometry=False, for_spectrum=False):
    """Observed-frame fluxes of every nebular grid line (n_lines,), through the same weights,
    dust, escape, mass and distance factors as the spectrum.

    Default: integrated fluxes [erg/s/cm^2] with the 1/(1+z) Jacobian, IGM at the line
    wavelength and ``eline_scaling`` (the Lines observation).  ``for_spectrum``: the same
    without ``eline_scaling`` (painted by the Spectrum projector).  ``for_photometry``: per-Hz
    amplitudes with the full f_nu flux factor and no IGM (applied inside the photometric line
    basis) and no ``eline_scaling``.
    """
    W = self.calculate_ssp_weights(theta=theta)          # (n_z, n_age)
    logZ_gas = theta["gas_logz"]
    logU     = theta["gas_logu"]
    line_lum = self.neb.evaluate_batch_line_lum(
        logZ_gas, logU, self._neb_ages_young, self._neb_logqq_young,
    )                                                    # (n_z, n_young, n_lines)
    if "frac_obrun" in theta:
        f_esc = jnp.ravel(theta["frac_obrun"])[0]
        line_lum = line_lum * (1.0 - f_esc)

    lam = self.neb.nebem_line_pos                        # (n_lines,)
    li = jnp.clip(jnp.searchsorted(self.wave, lam) - 1,
                  0, self.wave.shape[0] - 2)
    lf = jnp.clip((lam - self.wave[li])
                  / (self.wave[li + 1] - self.wave[li]), 0.0, 1.0)

    n_young = self._neb_young_idx.shape[0]
    attn_age_lines = jnp.ones((n_young, lam.shape[0]))
    diff_lines = jnp.ones(lam.shape[0])
    if hasattr(self, "attenuate_dust"):
        attn, attn_diffuse = self.attenuate_dust(self.wave, theta)
        if hasattr(self, "_age_bin_mix"):
            M = self._age_bin_mix
            tau_age = jnp.einsum("ab,bw->aw", M, attn)   # (n_age, n_wave)
            tau_lines = ((1.0 - lf)[None, :] * tau_age[:, li]
                         +        lf[None, :] * tau_age[:, li + 1])
            aal = jnp.exp(-tau_lines)                    # (n_age, n_lines)
            if "frac_obrun" in theta and self.fesc_geometry != "picket":
                fo = jnp.ravel(theta["frac_obrun"])[0]
                aal = (1.0 - fo) * aal + fo
            attn_age_lines = aal[self._neb_young_idx, :]
        diff_lines = jnp.exp(-((1.0 - lf) * attn_diffuse[li]
                               + lf * attn_diffuse[li + 1]))

    F = jnp.einsum("zy,zyl,yl->l",
                   W[:, self._neb_young_idx], line_lum, attn_age_lines)
    F = F * diff_lines

    if "logmass" in theta:
        F = F * 10.0 ** jnp.ravel(theta["logmass"])[0]
    if "zred" in theta:
        z_scalar = jnp.ravel(theta["zred"])[0]
        ff = self._flux_factor(theta)
        F = F * (ff if for_photometry else ff / (1.0 + z_scalar))
        if self.igm is not None and not for_photometry:
            if "igm_factor" in theta:
                ig_factor = jnp.ravel(theta["igm_factor"])[0]
            else:
                ig_factor = jnp.float32(self.igm_factor)
            trans = self.igm.attenuation(self.wave, z_scalar,
                                         factor=ig_factor)
            F = F * ((1.0 - lf) * trans[li] + lf * trans[li + 1])
    if not for_photometry and not for_spectrum and "eline_scaling" in theta:
        F = F * jnp.ravel(theta["eline_scaling"])[0]
    return F

get_line_spec

get_line_spec(theta)

Painted line component alone on the rest-frame grid, scaled by mass, flux factor and IGM as in predict, times eline_scaling (the Lines-observation aperture factor).

Source code in ceridwen/csp/csp.py

def get_line_spec(self, theta):
    """Painted line component alone on the rest-frame grid, scaled by mass, flux factor and IGM as in
    ``predict``, times ``eline_scaling`` (the Lines-observation aperture factor)."""
    if not hasattr(self, "neb") or self.neb is None:
        return jnp.zeros_like(self.wave)

    _continuum, line_only = self.get_spectrum_components(theta)

    if "logmass" in theta:
        mass_scale = jnp.float32(10.0 ** theta["logmass"][0])
        line_only = line_only * mass_scale
    if "zred" in theta:
        z_scalar = jnp.ravel(theta["zred"])[0]
        line_only = line_only * jnp.float32(self._flux_factor(theta))
        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)
    if "eline_scaling" in theta:
        line_only = line_only * jnp.ravel(theta["eline_scaling"])[0]
    return line_only

age_at

age_at(z)

Age of the Universe [Gyr] at redshift z under self.cosmo (a float for a scalar, a JAX array otherwise).

Source code in ceridwen/csp/csp.py

def age_at(self, z):
    """Age of the Universe [Gyr] at redshift ``z`` under ``self.cosmo``
    (a float for a scalar, a JAX array otherwise)."""
    from ceridwen.cosmology import age_gyr
    out = age_gyr(jnp.asarray(z, dtype=float), self.cosmo)
    return float(out) if jnp.ndim(out) == 0 else out

display_sfh

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

Plot the SFH exactly as _ssp_weights interprets it (step: constant per bin; linear: chords between nodes); theta defaults to theta_init, units in {"Gyr", "Myr", "yr"}. Returns the axes.

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 exactly as ``_ssp_weights`` interprets it (step: constant per bin; linear:
    chords between nodes); ``theta`` defaults to ``theta_init``, ``units`` in {"Gyr", "Myr", "yr"}.
    Returns the axes.
    """
    import matplotlib.pyplot as plt

    theta = self.theta_init if theta is None else theta

    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)

    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}"
        )

    if per_bin:
        bar_psi = psi
    else:
        bar_psi = 0.5 * (psi[:-1] + psi[1:])

    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":
        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}}]$")


    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"
    )

    return ax

calculate_ssp_weights_const_zh

calculate_ssp_weights_const_zh(theta)

Weights for constant metallicity, piecewise-linear SFH.

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_const_zh(self, theta):
    """Weights for constant metallicity, piecewise-linear SFH."""
    return self._ssp_weights(theta, zh_mode="const", sfh_mode="linear")

calculate_ssp_weights_const_zh_step

calculate_ssp_weights_const_zh_step(theta)

Weights for constant metallicity, piecewise-constant SFH.

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_const_zh_step(self, theta):
    """Weights for constant metallicity, piecewise-constant SFH."""
    return self._ssp_weights(theta, zh_mode="const", sfh_mode="step")

calculate_ssp_weights_var_zh

calculate_ssp_weights_var_zh(theta)

Weights for a metallicity history, piecewise-linear SFH.

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_var_zh(self, theta):
    """Weights for a metallicity history, piecewise-linear SFH."""
    return self._ssp_weights(theta, zh_mode="var", sfh_mode="linear")

calculate_ssp_weights_var_zh_step

calculate_ssp_weights_var_zh_step(theta)

Weights for a metallicity history, piecewise-constant SFH.

Source code in ceridwen/csp/csp.py

def calculate_ssp_weights_var_zh_step(self, theta):
    """Weights for a metallicity history, piecewise-constant SFH."""
    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. include_lines None -> self.nebemlineinspec.

Source code in ceridwen/csp/csp.py

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

    if self.fesc_geometry == "picket" and "frac_obrun" in theta:
        return self._spectrum_picket_nodem(theta, include_lines)

    ion_mult, neb_amp = self._ion_multiplier(theta)

    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)

    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   # runaway fraction skips the birth cloud
        attn_age = jnp.where(self.kill_ion, jnp.float32(1.0), attn_age)   # escaped LyC: no birth-cloud dust at all

    W_f32 = W.astype(jnp.float32)
    spectrum = jnp.einsum("za,zaw,aw->w", W_f32, self.flux,
                          ion_mult * attn_age)
    spectrum = spectrum + self._neb_spectrum_term(
        W_f32, theta, include_lines=include_lines,
        attn_age=attn_age, amplitude=neb_amp)
    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)

    if self.fesc_geometry == "picket" and "frac_obrun" in theta:
        return self._spectrum_picket_dem(theta, include_lines)

    ion_mult, neb_amp = self._ion_multiplier(theta)

    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))

    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   # runaway fraction skips the birth cloud
        attn_age = jnp.where(self.kill_ion, jnp.float32(1.0), attn_age)   # escaped LyC: no birth-cloud dust at all

    W_f32 = W.astype(jnp.float32)
    neb_v, neb_base = self._neb_weights_and_base(
        W_f32, theta, include_lines=include_lines, amplitude=neb_amp)
    yi = self._neb_young_idx
    spectrum_dust_free = (
        jnp.einsum("za,zaw,aw->w", W_f32, self.flux, ion_mult)
        + jnp.einsum("y,yw->w", neb_v, neb_base))
    attenuated = (
        jnp.einsum("za,zaw,aw->w", W_f32, self.flux, ion_mult * attn_age)
        + jnp.einsum("y,yw,yw->w", neb_v, neb_base, attn_age[yi, :]))
    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 only (include_lines ignored).

Source code in ceridwen/csp/csp.py

def get_spectrum_dattn_nodem_noneb(self, theta, *, include_lines=None):
    """Dust attenuation only (``include_lines`` 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)

    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 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`` 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))

    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.

Source code in ceridwen/csp/csp.py

def get_spectrum_nodattn_nodem_noneb(self, theta, *, include_lines=None):
    """Stellar continuum only."""
    _ = 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, no dust.

Source code in ceridwen/csp/csp.py

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

    ion_mult, neb_amp = self._ion_multiplier(theta)

    W_f32 = W.astype(jnp.float32)
    return (jnp.einsum("za,zaw,aw->w", W_f32, self.flux, ion_mult)
            + self._neb_spectrum_term(W_f32, theta,
                                      include_lines=include_lines,
                                      amplitude=neb_amp))

ceridwen.csp.CSPBasis_afe

CSPBasis_afe(
    SSPData,
    theta=None,
    tiny_logt=-70,
    zh_const=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",
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    cosmo=None,
    **kwargs
)

Bases: CSPBasis

Composite stellar population basis with [alpha/Fe] interpolation and no nebular model.

Same constructor and interface as CSPBasis minus the nebular arguments; requires an SSPDataAfe grid (4-D ssp_flux with an ssp_afe axis). theta["afe"] is optional: absent, the plane closest to [alpha/Fe] = 0 is used. The SFH weights, dust, projection, flux factor and cosmology are those of CSPBasis.

Source code in ceridwen/csp/csp_afe.py

def __init__(
    self,
    SSPData,
    theta=None,
    tiny_logt=-70,
    zh_const=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',
    track_zred_age=False,
    lookback_time=None,
    sfh_per_bin=False,
    cosmo=None,
    **kwargs,
):
    self.verbose = bool(verbose)
    if kwargs:
        hint = ""
        if "sigma_losvd_kms" in kwargs:
            hint = (" ('sigma_losvd_kms' was removed: the galaxy velocity "
                    "dispersion is set once on the model, "
                    "SedModel(kinematics=Kinematics(sigma_gal=...)))")
        elif "tuniv" in kwargs:
            hint = (" ('tuniv' was removed: the age of the Universe comes from "
                    "the cosmology, csp.age_at(z) / cosmo.age(z))")
        raise TypeError(
            f"CSPBasis_afe got unexpected keyword argument(s) {sorted(kwargs)}{hint}")
    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_dust_params is None:
        init_dust_params = {'bin_edges': [(-jnp.inf, -1.97)], 'laws': ['powerlaw']}

    _flux_in = jnp.asarray(SSPData.ssp_flux)
    _afe_in  = getattr(SSPData, "ssp_afe", None)
    if _flux_in.ndim != 4 or _afe_in is None:
        raise TypeError(
            "CSPBasis_afe requires an alpha-enhanced SSP grid: a 4-D "
            "ssp_flux (n_afe, n_z, n_age, n_wave) WITH an ssp_afe axis "
            "(ceridwen.ssps.ssp_data_afe.SSPDataAfe). Got "
            f"ssp_flux.ndim={_flux_in.ndim} and ssp_afe="
            f"{'absent' if _afe_in is None else 'present'} "
            f"({type(SSPData).__name__}). For solar-scaled 3-D grids "
            "switch to the matching basis:\n"
            "    from ceridwen.csp import CSPBasis   # solar-scaled, "
            "with nebular model\n"
            "or rebuild the grid with SSPDataAfe.from_fsps (python-fsps "
            ">= 4.0, AFE_FLAG=1) to fit [alpha/Fe]."
        )
    _afe_in = jnp.atleast_1d(jnp.asarray(_afe_in, dtype=float))
    if _afe_in.size != _flux_in.shape[0]:
        raise ValueError(
            f"ssp_afe has {_afe_in.size} points but ssp_flux leads "
            f"with {_flux_in.shape[0]}; grid is inconsistent."
        )
    if _afe_in.size > 1 and not bool(np.all(np.diff(np.asarray(_afe_in)) > 0)):
        raise ValueError(
            "ssp_afe must be strictly increasing (required by the "
            "searchsorted-based interpolation in _flux_at_afe); got "
            f"{np.asarray(_afe_in).tolist()}."
        )
    self.flux      = jnp.array(_flux_in, dtype=jnp.float32)  # (n_afe, n_z, n_age, n_wave)
    self.afe_grid  = _afe_in                           # (n_afe,) [alpha/Fe]
    self._n_afe    = int(_afe_in.size)
    self._afe_solar_idx = int(np.argmin(np.abs(np.asarray(_afe_in))))
    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
    self.zlegend   = 10 ** self.zmet                   # linear metallicity
    self.lib_resolution = (
        (np.asarray(SSPData.ssp_wave, dtype=np.float64),
         np.asarray(SSPData.ssp_resolution, dtype=np.float64))
        if getattr(SSPData, "ssp_resolution", None) is not None else None)
    self.ssp_ages_lgyr = self.ages + 9                 # log10(yr)

    self._ssp_isoc_type    = getattr(SSPData, "isoc_type", None)
    self._ssp_spec_library = getattr(SSPData, "spec_library", None)

    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)

    self._ssp_lo_yr = 10.0 ** self._logage_hi   # (n_age-1,)
    self._ssp_hi_yr = 10.0 ** self._logage_lo   # (n_age-1,)

    _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]
    )
    self._ssp_voronoi_hi = jnp.concatenate(
        [_voro_mid, jnp.array([_voro_hi_ext])]
    )

    self.tiny_logt  = tiny_logt
    if sps_home is None:
        sps_home = os.environ.get("SPS_HOME")
    if add_dust_emission and not sps_home:
        raise ValueError(
            "sps_home is required for 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
    from ..cosmology import Cosmology as _Cosmology
    if cosmo is None:
        raise TypeError(
            f"{type(self).__name__} needs an explicit cosmology: pass "
            "cosmo=Cosmology.planck18(), Cosmology.wmap9(), "
            "Cosmology.flat(H0, Om0), Cosmology.from_name(...) or "
            "Cosmology.from_astropy(...)  (from ceridwen import Cosmology)")
    if not isinstance(cosmo, _Cosmology):
        raise TypeError(
            f"cosmo must be a ceridwen Cosmology, got {type(cosmo).__name__}; "
            "for an astropy cosmology use Cosmology.from_astropy(...)")
    self._cosmo = cosmo
    self.track_zred_age = bool(track_zred_age)
    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)

    if add_diffuse_dust or add_dust:
        self.set_attenuation_function(add_diffuse_dust, add_dust)

    theta = self.initialize_dust_components(
        add_dust, add_diffuse_dust, add_dust_emission,
        theta, init_dust_params, diffuse_law, sps_home,
    )

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

    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()})

    self.check_param_ranges(self.theta_init)

initialize_model_structure

initialize_model_structure(theta)

CSPBasis.initialize_model_structure plus the optional scalar theta['afe'].

Source code in ceridwen/csp/csp_afe.py

def initialize_model_structure(self, theta):
    """``CSPBasis.initialize_model_structure`` plus the optional scalar theta['afe']."""
    super().initialize_model_structure(theta)
    if 'afe' in theta:
        afe = jnp.atleast_1d(jnp.asarray(theta['afe'], dtype=float))
        assert afe.shape == (1,), \
            "'afe' must be a scalar (wrapped in shape-(1,) array)"
        if self._n_afe == 1:
            warnings.warn(
                "theta contains 'afe' but the SSP grid has a single "
                "[alpha/Fe] plane (n_afe=1, legacy or AFE_FLAG=0 grid); "
                "'afe' is IGNORED. Build an alpha-enhanced grid with "
                "SSPDataAfe.from_fsps to fit alpha enhancement.",
                stacklevel=3,
            )
    self._known_theta_keys = (self._known_theta_keys | {'afe'}) - {'eline_scaling'}

check_param_ranges

check_param_ranges(theta=None, warn=True)

CSPBasis.check_param_ranges plus theta['afe'] against the [alpha/Fe] grid.

Source code in ceridwen/csp/csp_afe.py

def check_param_ranges(self, theta=None, warn=True):
    """``CSPBasis.check_param_ranges`` plus theta['afe'] against the [alpha/Fe] grid."""
    if theta is None:
        theta = self.theta_init
    msgs = super().check_param_ranges(theta, warn=False)
    if 'afe' in theta and self._n_afe > 1:
        alo = float(self.afe_grid.min())
        ahi = float(self.afe_grid.max())
        v = np.asarray(theta['afe'], float)
        if v.size and (np.nanmin(v) < alo or np.nanmax(v) > ahi):
            msgs.append(
                f"theta['afe'] has values outside the SSP [alpha/Fe] grid "
                f"[{alo:+.2f}, {ahi:+.2f}]; these are silently clamped to "
                f"the nearest grid edge (aMIST/C3K support is "
                f"-0.2 .. +0.6)."
            )
    if warn:
        for m in msgs:
            warnings.warn(m, stacklevel=2)
    return msgs

get_spectrum_components

get_spectrum_components(theta)

(continuum, zeros): no nebular model, so the line component is identically zero.

Source code in ceridwen/csp/csp_afe.py

def get_spectrum_components(self, theta: dict) -> tuple:
    """``(continuum, zeros)``: no nebular model, so the line component is identically zero."""
    self._warn_unknown_theta_keys(theta)
    continuum = self.get_spectrum(theta=theta, include_lines=False)
    return continuum, jnp.zeros_like(continuum)

get_line_spec

get_line_spec(theta)

Zeros (no nebular model).

Source code in ceridwen/csp/csp_afe.py

def get_line_spec(self, theta):
    """Zeros (no nebular model)."""
    _ = theta
    return jnp.zeros_like(self.wave)

get_spectrum_dattn_nodem_noneb

get_spectrum_dattn_nodem_noneb(
    theta, *, include_lines=None
)

Dust attenuation only (include_lines ignored).

Source code in ceridwen/csp/csp_afe.py

def get_spectrum_dattn_nodem_noneb(self, theta, *, include_lines=None):
    """Dust attenuation only (``include_lines`` ignored)."""
    _ = include_lines
    flux = self._flux_at_afe(theta)               # (n_z, n_age, n_wave)
    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)

    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, 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 (include_lines ignored).

Source code in ceridwen/csp/csp_afe.py

def get_spectrum_dattn_dem_noneb(self, theta, *, include_lines=None):
    """Dust attenuation + dust emission (``include_lines`` ignored)."""
    _ = include_lines
    flux = self._flux_at_afe(theta)               # (n_z, n_age, n_wave)
    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))

    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, flux)
    attenuated        = jnp.einsum("za,zaw,aw->w", weights, 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.

Source code in ceridwen/csp/csp_afe.py

def get_spectrum_nodattn_nodem_noneb(self, theta, *, include_lines=None):
    """Stellar continuum only."""
    _ = include_lines
    flux     = self._flux_at_afe(theta)           # (n_z, n_age, n_wave)
    weights  = self.calculate_ssp_weights(theta=theta).astype(jnp.float32)
    return jnp.einsum("za,zaw->w", weights, flux)

Broadening

ceridwen.broadening.Kinematics dataclass

Kinematics(sigma_gal, sigma_gas=TIED, sigma_max=2000.0)

Galaxy velocity dispersions [km/s]; set once on the model, shared by all observations.

sigma_gal : float (fixed) or str (theta key of a free parameter). No default here; SedModel defaults to DEFAULT_KINEMATICS (300 km/s, stars and gas). sigma_gas : float, str, or TIED (= sigma_gal). sigma_max : largest value the compiled kernels support; free values are clipped to it.

validate_theta

validate_theta(theta, priors=None)

Free keys are in theta (a dict or a set of names); a bounded prior on a free key stays within sigma_max; fixed widths are not also in theta.

Source code in ceridwen/broadening.py

def validate_theta(self, theta, priors: Optional[dict] = None) -> None:
    """Free keys are in ``theta`` (a dict or a set of names); a bounded prior on a
    free key stays within sigma_max; fixed widths are not also in theta."""
    for k in self.free_keys:
        if k not in theta:
            raise KeyError(
                f"Kinematics expects theta['{k}'] (free width) but it is not "
                f"in theta; add it (with a prior) or give a float to fix it")
        if priors is not None and k in priors:
            b = getattr(priors[k], "bounds", None)
            if callable(b):
                b = b()
            if b is not None:
                hi = float(np.max(np.asarray(b[1], dtype=float)))
                if np.isfinite(hi) and hi > self.sigma_max:
                    raise ValueError(
                        f"prior on '{k}' allows up to {hi} km/s > "
                        f"sigma_max={self.sigma_max}; raise sigma_max")
    for name, v in (("sigma_gal", self.sigma_gal),
                    ("sigma_gas", self.effective_sigma_gas)):
        if not isinstance(v, str) and name in theta:
            raise ValueError(
                f"{name} is fixed to {v} km/s in Kinematics but theta also "
                f"contains '{name}'; remove one of them")

resolve

resolve(theta)

(sigma_gal, sigma_gas) as JAX scalars; branching on Python types only.

Source code in ceridwen/broadening.py

def resolve(self, theta: dict):
    """(sigma_gal, sigma_gas) as JAX scalars; branching on Python types only."""
    def _get(v):
        if isinstance(v, str):
            return jnp.clip(jnp.ravel(jnp.asarray(theta[v]))[0], 0.0, self.sigma_max)
        return jnp.asarray(float(v))
    return _get(self.sigma_gal), _get(self.effective_sigma_gas)

ceridwen.broadening.Instrument dataclass

Instrument(kind, value, wave=None)

Line-spread function of a spectrograph; unit and convention are in the constructor name.

Instrument.R_fwhm(2700) R = lambda/FWHM (datasheet) Instrument.R_sigma(6358) R = lambda/sigma (same LSF as R_fwhm(2700)) Instrument.fwhm_aa(2.5) FWHM [A], observed frame Instrument.sigma_aa(1.06) sigma [A], observed frame Instrument.sigma_kms(47.0) sigma [km/s] Instrument.fwhm_kms(111.0) FWHM [km/s] Instrument.(array, wave=obs_wave) per-pixel value on observed wavelengths

sigma_kms_at

sigma_kms_at(wave_obs)

sigma [km/s] at each observed wavelength (NumPy, setup only).

Source code in ceridwen/broadening.py

def sigma_kms_at(self, wave_obs) -> np.ndarray:
    """sigma [km/s] at each observed wavelength (NumPy, setup only)."""
    wave_obs = np.asarray(wave_obs, dtype=np.float64)
    if self.value.ndim == 1:
        if wave_obs.min() < self.wave[0] or wave_obs.max() > self.wave[-1]:
            raise ValueError(
                "Instrument width array does not cover the observed "
                f"wavelengths [{wave_obs.min():.1f}, {wave_obs.max():.1f}] "
                f"(covers [{self.wave[0]:.1f}, {self.wave[-1]:.1f}])")
        v = np.interp(wave_obs, self.wave, self.value)
    else:
        v = np.full_like(wave_obs, float(self.value))
    if self.kind == "R_fwhm":
        return CKMS * FWHM_TO_SIGMA / v
    if self.kind == "R_sigma":
        return CKMS / v
    if self.kind == "fwhm_aa":
        return CKMS * FWHM_TO_SIGMA * v / wave_obs
    if self.kind == "sigma_aa":
        return CKMS * v / wave_obs
    if self.kind == "fwhm_kms":
        return FWHM_TO_SIGMA * v
    return v

ceridwen.broadening.DEFAULT_KINEMATICS module-attribute

DEFAULT_KINEMATICS = Kinematics(sigma_gal=300.0)

ceridwen.broadening.SpectralProjector dataclass

SpectralProjector(
    kinematics,
    instrument,
    subtract_library,
    window,
    J,
    W,
    paint,
    line_idx,
    sigma_inst_kms,
    sigma_lib_kms,
    sigma_fix_kms,
    zred_range=None,
    opz_ref=1.0,
    wave_obs=None,
    line_wave_rest=None,
    line_sigma_table_kms=None,
)

Built once per Spectrum by setup_for_model; predict() is the per-likelihood call.

With zred_range the projector serves a sampled redshift: the log grid is the one of the reference redshift extended at the same dv to cover every z in the range (so at opz_ref it reproduces the fixed-z projector), the response is built at opz_ref and the model is read at wave_log * opz_ref / opz per call (the sigma_gal kernel commutes with the shift); the lines are painted at line_wave_rest * opz. The library width in the fixed kernel is the one at opz_ref.

build classmethod

build(
    kinematics,
    instrument,
    wave_model,
    wave_obs,
    zred,
    lib_sigma_kms=None,
    line_wave_rest=None,
    subtract_library=True,
    zred_range=None,
)

wave_model rest-frame model grid [A] wave_obs observed pixel centres [A] zred fixed redshift, or the reference redshift inside zred_range lib_sigma_kms SSPData.ssp_resolution on wave_model, NaN = unknown (or None) line_wave_rest rest wavelengths of all model lines, in predict_line_fluxes order (or None) zred_range (z_min, z_max) of a SAMPLED redshift (None: fixed at zred)

Source code in ceridwen/broadening.py

@classmethod
def build(cls, kinematics: Kinematics, instrument: Optional[Instrument],
          wave_model, wave_obs, zred, lib_sigma_kms=None,
          line_wave_rest=None, subtract_library=True, zred_range=None):
    """
    wave_model      rest-frame model grid [A]
    wave_obs        observed pixel centres [A]
    zred            fixed redshift, or the reference redshift inside ``zred_range``
    lib_sigma_kms   SSPData.ssp_resolution on wave_model, NaN = unknown (or None)
    line_wave_rest  rest wavelengths of all model lines, in predict_line_fluxes order (or None)
    zred_range      (z_min, z_max) of a SAMPLED redshift (None: fixed at ``zred``)
    """
    b = kinematics
    if instrument is not None and not isinstance(instrument, Instrument):
        raise TypeError("instrument must be an Instrument (Instrument.R_fwhm(...), "
                        "Instrument.fwhm_aa(...), ...) or None")
    wave_model = np.asarray(wave_model, dtype=np.float64)
    wave_obs = np.asarray(wave_obs, dtype=np.float64)
    if np.any(np.diff(wave_obs) <= 0.0):
        raise ValueError("observed wavelengths must be strictly increasing")
    opz = 1.0 + float(zred)
    if zred_range is not None:
        z_lo, z_hi = float(zred_range[0]), float(zred_range[1])
        if not (np.isfinite(z_lo) and np.isfinite(z_hi) and -1.0 < z_lo < z_hi):
            raise ValueError(f"zred_range must be finite with -1 < z_min < z_max, got {zred_range}")
        if not (z_lo <= float(zred) <= z_hi):
            raise ValueError(f"the reference redshift {float(zred):g} lies outside "
                             f"zred_range = ({z_lo:g}, {z_hi:g})")
        zred_range = (z_lo, z_hi)
        opz_lo, opz_hi = 1.0 + z_lo, 1.0 + z_hi
    else:
        opz_lo = opz_hi = opz

    if instrument is None:
        s_inst = np.zeros_like(wave_obs)
        subtract_library = False
    else:
        s_inst = instrument.sigma_kms_at(wave_obs)
    if subtract_library and lib_sigma_kms is not None:
        lib = np.asarray(lib_sigma_kms, dtype=np.float64)
        if lib.shape != wave_model.shape:
            raise ValueError("lib_sigma_kms must be on the model grid")
        lib = np.nan_to_num(lib, nan=0.0)
        s_lib = np.interp(wave_obs / opz, wave_model, lib)
    else:
        lib = None
        s_lib = np.zeros_like(wave_obs)
    s_fix2 = s_inst ** 2 - s_lib ** 2
    n_bad = int(np.sum(s_fix2 < 0.0))
    if n_bad:
        warnings.warn(
            f"instrument narrower than the SSP library at {n_bad} of "
            f"{wave_obs.size} pixels ({100 * n_bad / wave_obs.size:.1f} %, "
            f"{wave_obs[s_fix2 < 0].min():.0f}-{wave_obs[s_fix2 < 0].max():.0f} A): "
            "the continuum is delivered at library resolution there, which is "
            "correct for a coarse grid; if the instrument is really finer, use a "
            "finer grid; if the number is wrong, check the Instrument unit")
    s_fix = np.sqrt(np.clip(s_fix2, 0.0, None))
    if zred_range is not None and lib is not None:
        worst = 0.0
        for o in (opz_lo, opz_hi):
            s_end = np.sqrt(np.clip(s_inst ** 2 - np.interp(wave_obs / o, wave_model, lib) ** 2,
                                    0.0, None))
            worst = max(worst, float(np.max(np.abs(s_end - s_fix) / np.maximum(s_fix, 1e-3))))
        if worst > 0.1:
            warnings.warn(
                f"the library width in the fixed kernel is taken at z = {opz - 1:g}; over "
                f"zred_range it changes the kernel by up to {100 * worst:.0f} % at some "
                "pixel. Narrow the redshift prior or accept the approximation")

    marg = BAND_NSIGMA * (b.sigma_max + float(s_fix.max())) / CKMS
    wmin_ref = wave_obs[0] / opz * np.exp(-marg)
    wmax_ref = wave_obs[-1] / opz * np.exp(+marg)
    wmin = wave_obs[0] / opz_hi * np.exp(-marg)
    wmax = wave_obs[-1] / opz_lo * np.exp(+marg)
    if wmin < wave_model[0] or wmax > wave_model[-1]:
        raise ValueError(
            f"model grid [{wave_model[0]:.0f}, {wave_model[-1]:.0f}] A does not "
            f"cover the observed window plus kernel margin "
            f"[{wmin:.0f}, {wmax:.0f}] A (rest"
            + (f", zred in [{z_lo:g}, {z_hi:g}]" if zred_range else "")
            + "); trim the observation, narrow the redshift range or reduce sigma_max")
    window = WindowSmoother(LogGrid.build(
        wave_model, wmin_ref, wmax_ref, b.sigma_max,
        extend_to=None if zred_range is None else (wmin, wmax)))
    J, W = build_response(window.grid, np.log(wave_obs) - np.log(opz), s_fix)

    paint, line_idx = None, np.zeros(0, dtype=np.int64)
    lw_kept = s_table = None
    if line_wave_rest is not None:
        lwr = np.asarray(line_wave_rest, dtype=np.float64)
        marg_l = BAND_NSIGMA * (b.sigma_max + float(s_inst.max())) / CKMS
        keep = ((lwr * opz_hi > wave_obs[0] * np.exp(-marg_l))
                & (lwr * opz_lo < wave_obs[-1] * np.exp(marg_l)))
        line_idx = np.flatnonzero(keep)
        if line_idx.size:
            if instrument is None:
                # no LSF: a line cannot be narrower than the pixel it lands on
                s_table = 0.5 * CKMS * np.gradient(np.log(wave_obs))
            else:
                s_table = s_inst
            lw_kept = lwr[keep]
            if zred_range is None:
                lw = lwr[keep] * opz
                paint = make_line_painter(wave_obs, lw, np.interp(lw, wave_obs, s_table))
            else:
                paint = make_line_painter_free_z(wave_obs, lwr[keep], s_table)

    return cls(kinematics=b, instrument=instrument,
               subtract_library=bool(subtract_library), window=window,
               J=jnp.asarray(J), W=jnp.asarray(W), paint=paint,
               line_idx=line_idx, sigma_inst_kms=s_inst,
               sigma_lib_kms=s_lib, sigma_fix_kms=s_fix,
               zred_range=zred_range, opz_ref=opz, wave_obs=wave_obs,
               line_wave_rest=lw_kept, line_sigma_table_kms=s_table)

continuum

continuum(spec_rest, sigma_gal_kms, opz=None)

Continuum on the observed pixels; opz = 1 + z (traced) when free_z.

Source code in ceridwen/broadening.py

def continuum(self, spec_rest, sigma_gal_kms, opz=None):
    """Continuum on the observed pixels; ``opz`` = 1 + z (traced) when ``free_z``."""
    scale = None if opz is None else self.opz_ref / opz
    spec_log = self.window.smooth(self.window.to_log(spec_rest, scale), sigma_gal_kms)
    return apply_response(self.J, self.W, spec_log)

line_basis

line_basis(sigma_gas_kms, opz_line)

(n_pix, n_kept) f_nu of UNIT-flux lines on the observed pixels, centred at line_wave_rest * opz_line (traced), width sqrt(sigma_gas^2 + sigma_inst^2) with the instrument width interpolated at the line; the painter's profile, one column per line.

Source code in ceridwen/broadening.py

def line_basis(self, sigma_gas_kms, opz_line):
    """(n_pix, n_kept) f_nu of UNIT-flux lines on the observed pixels, centred at
    ``line_wave_rest * opz_line`` (traced), width sqrt(sigma_gas^2 + sigma_inst^2) with the
    instrument width interpolated at the line; the painter's profile, one column per line."""
    wo = np.asarray(self.wave_obs, dtype=np.float64)
    lam = jnp.asarray(np.asarray(self.line_wave_rest, dtype=np.float64))
    lnl = jnp.log(lam) + jnp.log(opz_line)
    s_inst = jnp.interp(lam * opz_line, jnp.asarray(wo),
                        jnp.asarray(np.asarray(self.line_sigma_table_kms, dtype=np.float64)))
    s = jnp.sqrt(sigma_gas_kms ** 2 + s_inst ** 2) / CKMS
    x = (jnp.asarray(np.log(wo))[:, None] - lnl[None, :]) / s[None, :]
    phi = jnp.exp(-0.5 * x * x) / (jnp.sqrt(2.0 * jnp.pi) * s[None, :])
    return phi * jnp.asarray(wo / C_AA_S)[:, None]

predict_with_line_basis

predict_with_line_basis(
    spec_rest,
    line_flux_obs_all,
    theta,
    fit_pos=None,
    basis=None,
)

(prediction, A): continuum plus the lines of line_flux_obs_all painted with :meth:line_basis at (1 + zred + eline_delta_zred) (theta key, default 0), and the unit-flux columns A (n_pix, len(fit_pos)) of the kept lines at positions fit_pos (None when fit_pos is None).

Source code in ceridwen/broadening.py

def predict_with_line_basis(self, spec_rest, line_flux_obs_all, theta, fit_pos=None,
                            basis=None):
    """``(prediction, A)``: continuum plus the lines of ``line_flux_obs_all`` painted with
    :meth:`line_basis` at ``(1 + zred + eline_delta_zred)`` (theta key, default 0), and the
    unit-flux columns ``A`` (n_pix, len(fit_pos)) of the kept lines at positions ``fit_pos``
    (None when ``fit_pos`` is None)."""
    s_gal, s_gas = self.kinematics.resolve(theta)
    opz = None
    if self.free_z:
        if "zred" not in theta:
            raise KeyError("projector built with zred_range needs theta['zred']")
        opz = 1.0 + jnp.ravel(jnp.asarray(theta["zred"]))[0]
    out = self.continuum(spec_rest, s_gal, opz)
    if self.line_wave_rest is None or self.line_idx.size == 0:
        return out, None
    if basis is None:           # a precomputed basis is only passed for fixed z, width, dz = 0
        opz_line = self.opz_ref if opz is None else opz
        if "eline_delta_zred" in theta:
            opz_line = opz_line + jnp.ravel(jnp.asarray(theta["eline_delta_zred"]))[0]
        basis = self.line_basis(s_gas, opz_line)
    else:
        basis = jnp.asarray(basis)
    if line_flux_obs_all is not None:
        out = out + basis @ line_flux_obs_all[self._line_idx_j]
    A = None if fit_pos is None else basis[:, np.asarray(fit_pos, dtype=np.int64)]
    return out, A

ceridwen.broadening.PhotometricBroadener dataclass

PhotometricBroadener(window)

sigma_gal broadening of spec_rest over the rest range the filters cover; pixels outside pass through.

Observations

ceridwen.observation.Photometry

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

Bases: Observation

Broadband photometric observation; flux and uncertainty in AB maggies (1 maggie = 3631 Jy).

Parameters:

Name	Type	Description	Default
filters	list of str or Filter -- names are resolved in the filter library		[]
flux	(array(n_filters), maggies)		required
uncertainty	(array(n_filters), maggies - -1 - sigma)		required
mask	bool array (n_filters,) -- True = used in the fit		required
upper_limit	bool array (n_filters,) -- True = one-sided (model > data only) chi-squared penalty		None

Source code in ceridwen/observation/photometry.py

def __init__(self, filters=[], name=None, upper_limit=None, **kwargs):
    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 from filter-name strings or Filter objects.

Source code in ceridwen/observation/photometry.py

def set_filters(self, filters):
    """Set the filter list from filter-name strings or ``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]
    self._wavelength = jnp.asarray([f.wave_effective for f in self.filters])

get_maggies

get_maggies(model_wave, model_fnu)

Synthetic maggies, shape (n_filters,), of an F_nu spectrum on model_wave [Å] (reference path; normalisation follows the input flux units).

Source code in ceridwen/observation/photometry.py

def get_maggies(self, model_wave, model_fnu):
    """Synthetic maggies, shape (n_filters,), of an F_nu spectrum on ``model_wave`` [Å]
    (reference path; normalisation follows the input flux units)."""
    if self.filterset is None:
        raise ValueError("No FilterSet configured; call set_filters() first.")
    _c         = jnp.array(2.998e18)
    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 the (n_filters, n_wave) float32 projection matrix _T (F_nu -> maggies) for the rest-frame grid wave_model [Å] observed at fixed zred; required before predict.

Source code in ceridwen/observation/photometry.py

def setup_for_model(self, wave_model, zred: float = 0.0):
    """Precompute the ``(n_filters, n_wave)`` float32 projection matrix ``_T`` (F_nu -> maggies)
    for the rest-frame grid ``wave_model`` [Å] observed at fixed ``zred``; required before ``predict``."""
    wm_rest = np.asarray(wave_model, dtype=np.float64)
    opz = 1.0 + float(zred)
    wm = opz * wm_rest
    n_wave = len(wm)
    _c = 2.998e18

    fnu_to_flam = _c / wm**2

    lam_filt = np.asarray(self.filterset.lam, dtype=np.float64)
    n_lam = len(lam_filt)

    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)

    outside = (lam_filt < wm[0]) | (lam_filt > wm[-1])
    frac[outside] = 0.0

    trans = np.asarray(self.filterset.trans, dtype=np.float64)  # (n_filt, n_lam)
    inside = ~outside
    TH = np.zeros((trans.shape[0], n_wave), dtype=np.float64)
    np.add.at(TH.T, idx[inside],     (trans[:, inside] * (1.0 - frac[inside])).T)
    np.add.at(TH.T, idx[inside] + 1, (trans[:, inside] * frac[inside]).T)
    T  = TH * fnu_to_flam[None, :]

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

setup_broadening

setup_broadening(
    wave_model, zred, kinematics, *, free_z=False, neb=None
)

Build the sigma_gal broadener over the rest range the filters cover (the whole 912-25000 A window when free_z) and, when neb is given and sigma_gas is fixed, the static line basis at that width. Both are None for Kinematics.none(); call after setup_for_model.

Source code in ceridwen/observation/photometry.py

def setup_broadening(self, wave_model, zred, kinematics, *, free_z=False, neb=None):
    """Build the sigma_gal broadener over the rest range the filters cover (the
    whole 912-25000 A window when ``free_z``) and, when ``neb`` is given and
    sigma_gas is fixed, the static line basis at that width.  Both are None
    for ``Kinematics.none()``; call after ``setup_for_model``."""
    from ..broadening import PhotometricBroadener
    self._broadener = None
    self._line_basis = None
    if kinematics is None:
        return
    wm = np.asarray(wave_model, dtype=np.float64)
    static_zero = (kinematics.is_static
                   and float(kinematics.sigma_gal) == 0.0
                   and float(kinematics.effective_sigma_gas) == 0.0)
    if not static_zero:
        if free_z or self.filterset is None:
            wmin, wmax = max(912.0, wm[0]), min(25000.0, wm[-1])
        else:
            opz = 1.0 + float(zred)
            lam = np.asarray(self.filterset.lam, dtype=np.float64)
            wmin = max(lam.min() / opz, wm[0])
            wmax = min(lam.max() / opz, wm[-1])
        if wmin < wmax:
            self._broadener = PhotometricBroadener.build(kinematics, wm, wmin, wmax)
    gas = kinematics.effective_sigma_gas
    if neb is not None and not isinstance(gas, str) and float(gas) > 0.0:
        self._line_basis = np.asarray(neb.line_profiles(float(gas)), dtype=np.float32)

predict

predict(spectrum, wave_model)

Synthetic AB maggies, shape (n_filters,), as _T @ spectrum (F_nu on wave_model).

Source code in ceridwen/observation/photometry.py

def predict(self, spectrum, wave_model):
    """Synthetic AB maggies, shape (n_filters,), as ``_T @ spectrum`` (F_nu on ``wave_model``)."""
    if not getattr(self, "_has_precomputed_T", False):
        raise RuntimeError(
            "Photometry.predict() called before setup_for_model(): call "
            "phot.setup_for_model(wave_model, zred=...) once before the first "
            "predict / JIT trace (get_maggies(wave, spectrum) is the "
            "rest-frame reference path).")
    return self._T @ spectrum

predict_at_redshift

predict_at_redshift(spectrum_fnu_observed, wave_rest, zred)

Observer-frame AB maggies, shape (n_filters,), for a traced zred: spectrum_fnu_observed must already be observer-frame F_nu (flux factor and IGM applied) on the rest-frame grid wave_rest [Å].

Source code in ceridwen/observation/photometry.py

def predict_at_redshift(self, spectrum_fnu_observed, wave_rest, zred):
    """Observer-frame AB maggies, shape (n_filters,), for a traced ``zred``: ``spectrum_fnu_observed``
    must already be observer-frame F_nu (flux factor and IGM applied) on the rest-frame grid ``wave_rest`` [Å]."""
    if self.filterset is None:
        raise ValueError("No FilterSet configured; call set_filters() first.")
    wave_obs = (1.0 + jnp.asarray(zred)) * jnp.asarray(wave_rest)
    flux_flam = spectrum_fnu_observed * (2.998e18 / (wave_obs * wave_obs))
    return self.filterset.get_sed_maggies(flux_flam, sourcewave=wave_obs)

chi_sq

chi_sq(model_maggies)

Chi-squared (float) over unmasked bands; upper-limit bands are penalised only when model > data.

Source code in ceridwen/observation/photometry.py

def chi_sq(self, model_maggies):
    """Chi-squared (float) over unmasked bands; upper-limit bands are penalised only when model > data."""
    mf    = jnp.asarray(model_maggies, dtype=float)
    resid = (self.flux - mf) / self.uncertainty

    if self.upper_limit is not None:
        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, shape (n_filters,); masked bands NaN, upper-limit bands 0 when model < data.

Source code in ceridwen/observation/photometry.py

def residuals(self, model_maggies):
    """Per-filter (data - model) / sigma, shape (n_filters,); masked bands NaN, upper-limit bands 0 when model < data."""
    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,
    instrument=None,
    subtract_library=True,
    calibration=None,
    logify_spectrum=False,
    sky=None,
    noise_floor=0.0,
    zred_range=None,
    marginalize_elines=False,
    eline_prior_width=0.0,
    elines_to_fit=None,
    elines_to_fix=None,
    elines_to_ignore=None,
    **kwargs
)

Bases: Observation

Spectroscopic observation on an observed-frame pixel grid.

Parameters:

Name	Type	Description	Default
wavelength	array-like (n_pix,), Å, vacuum, OBSERVED frame.		None
flux	array-like (n_pix,) -- same units as the model spectra	passed to chi_sq / residuals.	None
uncertainty	array-like (n_pix,) -- same units as the model spectra	passed to chi_sq / residuals.	None
mask	array-like of bool or slice -- True for pixels that are USED.		slice(None)
instrument	Instrument or None -- the spectrograph LSF; unit and convention	are the constructor name (Instrument.R_fwhm(2700), Instrument.fwhm_aa(2.5), Instrument.sigma_kms(arr, wave=w), ...). None = no instrumental broadening (and no library subtraction).	None
subtract_library	bool -- remove the SSP library resolution in quadrature	from the instrument width (default True; needs the grid's resolution curve).	True
calibration	array-like (n_pix,) -- multiplicative model correction	(model * calibration ~ data), applied in chi_sq/residuals/log_likelihood.	None
logify_spectrum	bool -- residuals in log-flux space.		False
sky	array-like (n_pix,) -- subtracted from data in chi_sq/residuals/	log_likelihood only (not in predict).	None
noise_floor	float -- fractional floor: sigma_eff^2 = sigma^2 + (floor*|model|)^2.		0.0
noise	GaussianProcess -- adds a correlated-residual term in ``log_likelihood``	only.	None
zred_range	(z_min, z_max) -- support of a SAMPLED redshift; otherwise taken from	the finite bounds of the model's zred prior.	None
marginalize_elines	bool -- marginalise analytically over the fluxes of the nebular	lines in this spectrum (jointly with the model's Photometry and Lines), instead of fixing them at the CLOUDY prediction; needs an instrument; with a nebular model its parameters must not be sampled, without one (add_neb=False, CSPBasis_afe) the lines come from $SPS_HOME/data/emlines_info.dat and the prior is flat. See docs/eline_marginalisation.md.	False
eline_prior_width	float -- 0 (default): flat prior on each fitted line flux;	0: Gaussian prior centred on the CLOUDY flux with this FRACTIONAL width. Ly-alpha is always flat.	0.0
elines_to_fit	sequences of FSPS line names	($SPS_HOME/data/emlines_info.dat, e.g. "[O III] 5007"). Default: fit every grid line covered by the spectrum; fixed lines keep their CLOUDY flux; ignored lines are removed from every observation.	None
elines_to_fix	sequences of FSPS line names	($SPS_HOME/data/emlines_info.dat, e.g. "[O III] 5007"). Default: fit every grid line covered by the spectrum; fixed lines keep their CLOUDY flux; ignored lines are removed from every observation.	None
elines_to_ignore	sequences of FSPS line names	($SPS_HOME/data/emlines_info.dat, e.g. "[O III] 5007"). Default: fit every grid line covered by the spectrum; fixed lines keep their CLOUDY flux; ignored lines are removed from every observation.	None
The			required
they			required

Source code in ceridwen/observation/spectrum.py

def __init__(
    self,
    wavelength   = None,
    flux         = None,
    uncertainty  = None,
    mask         = slice(None),
    noise        = None,
    name         = None,
    instrument   = None,
    subtract_library = True,
    calibration  = None,
    logify_spectrum = False,
    sky          = None,
    noise_floor  = 0.0,
    zred_range   = None,
    marginalize_elines = False,
    eline_prior_width  = 0.0,
    elines_to_fit      = None,
    elines_to_fix      = None,
    elines_to_ignore   = None,
    **kwargs,
):
    for k in list(kwargs):
        if k in self._removed_kwargs:
            raise TypeError(
                f"Spectrum(): '{k}' was removed; use {self._removed_kwargs[k]}")
    if "eline_sigma" in kwargs:
        raise TypeError(
            "Spectrum(): there is no eline_sigma. The marginalised lines have the same "
            "width as every other line, sqrt(sigma_gas^2 + sigma_inst^2): set sigma_gas "
            "(a velocity DISPERSION in km/s, not a FWHM) once on the model, "
            "SedModel(kinematics=Kinematics(sigma_gal=..., sigma_gas=...))")
    if instrument is not None and not isinstance(instrument, Instrument):
        raise TypeError(
            "Spectrum(instrument=...) takes a ceridwen.broadening.Instrument "
            "(Instrument.R_fwhm(2700), Instrument.fwhm_aa(2.5), "
            "Instrument.sigma_kms(120.0), ...), got "
            f"{type(instrument).__name__}")
    self._wavelength    = (
        None if wavelength is None
        else jnp.asarray(wavelength, dtype=float)
    )
    self.instrument       = instrument
    self.subtract_library = bool(subtract_library)
    self.calibration    = (
        None if calibration is None
        else jnp.asarray(calibration, dtype=float)
    )
    self.logify_spectrum = logify_spectrum
    self.sky             = (None if sky is None
                            else jnp.asarray(sky, dtype=float))
    self.noise_floor     = float(noise_floor)
    self.zred_range      = (None if zred_range is None
                            else (float(zred_range[0]), float(zred_range[1])))
    self._proj = None
    self._set_eline_options(marginalize_elines, eline_prior_width,
                            elines_to_fit, elines_to_fix, elines_to_ignore, noise)

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

setup_for_model

setup_for_model(
    wave_model,
    zred=0.0,
    kinematics=None,
    lib_resolution=None,
    line_wave_rest=None,
    zred_range=None,
)

Build the projector from the rest-frame model grid wave_model [Å] redshifted by (1 + zred) onto self.wavelength; call once, outside JIT.

kinematics : Kinematics -- galaxy widths (None = Kinematics.none()) lib_resolution : (wave_rest [Å], sigma_v [km/s]) -- SSP library resolution curve on wave_model (NaN = unknown); used when subtract_library and an instrument are set line_wave_rest : (n_lines,) rest wavelengths of the nebular grid lines in predict_line_fluxes order, or None (no lines painted) zred_range : (z_min, z_max) -- for a SAMPLED redshift (default self.zred_range): the projection is then read at theta['zred'] on every call and zred is the reference redshift inside the range

Source code in ceridwen/observation/spectrum.py

def setup_for_model(self, wave_model, zred: float = 0.0,
                    kinematics=None, lib_resolution=None,
                    line_wave_rest=None, zred_range=None):
    """Build the projector from the rest-frame model grid ``wave_model`` [Å]
    redshifted by (1 + zred) onto ``self.wavelength``; call once, outside JIT.

    kinematics : Kinematics -- galaxy widths (None = ``Kinematics.none()``)
    lib_resolution : (wave_rest [Å], sigma_v [km/s]) -- SSP library
        resolution curve on ``wave_model`` (NaN = unknown); used when
        ``subtract_library`` and an instrument are set
    line_wave_rest : (n_lines,) rest wavelengths of the nebular grid lines
        in ``predict_line_fluxes`` order, or None (no lines painted)
    zred_range : (z_min, z_max) -- for a SAMPLED redshift (default ``self.zred_range``):
        the projection is then read at ``theta['zred']`` on every call and ``zred``
        is the reference redshift inside the range
    """
    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)."
        )
    if kinematics is None:
        kinematics = Kinematics.none()
    if zred_range is None:
        zred_range = self.zred_range
    lib = None
    if lib_resolution is not None and self.instrument is not None and self.subtract_library:
        lw = np.asarray(lib_resolution[0], dtype=np.float64)
        ls = np.asarray(lib_resolution[1], dtype=np.float64)
        wm = np.asarray(wave_model, dtype=np.float64)
        if lw.shape == wm.shape and np.allclose(lw, wm):
            lib = ls
        else:
            fin = np.isfinite(ls)
            lib = (np.interp(wm, lw[fin], ls[fin], left=np.nan, right=np.nan)
                   if fin.any() else None)
    self._proj = SpectralProjector.build(
        kinematics, self.instrument, wave_model, self._wavelength, zred,
        lib_sigma_kms=lib, line_wave_rest=line_wave_rest,
        subtract_library=self.subtract_library, zred_range=zred_range)

predict

predict(
    spectrum, wave_model=None, line_flux=None, theta=None
)

Model F_nu on the observed pixels (n_pix,): the rest-frame CONTINUUM spectrum (n_wave,) broadened by sigma_gal, the instrument LSF (library width removed) and resampled, plus the emission lines painted from their observed-frame integrated fluxes line_flux (all grid lines, or None) with sigma_gas + instrument. theta supplies the free widths and, for a projector built with zred_range, the redshift.

Source code in ceridwen/observation/spectrum.py

def predict(self, spectrum, wave_model=None, line_flux=None, theta=None):
    """Model F_nu on the observed pixels (n_pix,): the rest-frame CONTINUUM
    ``spectrum`` (n_wave,) broadened by sigma_gal, the instrument LSF (library
    width removed) and resampled, plus the emission lines painted from their
    observed-frame integrated fluxes ``line_flux`` (all grid lines, or None)
    with sigma_gas + instrument.  ``theta`` supplies the free widths and, for a
    projector built with ``zred_range``, the redshift."""
    if self._proj is None:
        raise RuntimeError(
            "Spectrum.predict() called before setup_for_model(): the "
            "projector has not been built. Call "
            "spec.setup_for_model(wave_model, zred=..., kinematics=...) once "
            "(before the first predict / JIT trace)."
        )
    return self._proj.predict(spectrum, line_flux, {} if theta is None else theta)

synthetic_photometry

synthetic_photometry(filterset)

Synthetic maggies (n_filters,) of this F_nu spectrum through filterset; None if there is no data.

Source code in ceridwen/observation/spectrum.py

def synthetic_photometry(self, filterset):
    """Synthetic maggies (n_filters,) of this F_nu spectrum through
    ``filterset``; None if there is 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 wavelength in [wave_min, wave_max] Å (inclusive).

Source code in ceridwen/observation/spectrum.py

def mask_wavelength_range(self, wave_min, wave_max):
    """Mask pixels with wavelength in [wave_min, wave_max] Å (inclusive)."""
    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 +/- dv km/s around each rest-frame line wavelength [Å], redshifted by (1 + zred) onto the observed grid.

Source code in ceridwen/observation/spectrum.py

def mask_lines(self, line_waves, dv=1000.0, zred=0.0):
    """Mask +/- ``dv`` km/s around each rest-frame line wavelength [Å],
    redshifted by (1 + zred) onto the observed grid."""
    if self._wavelength is None:
        return
    c_kms = 2.998e5
    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)

Sum of squared normalised residuals over unmasked pixels for model_flux (n_pix,) on the observed grid.

Source code in ceridwen/observation/spectrum.py

def chi_sq(self, model_flux):
    """Sum of squared normalised residuals over unmasked pixels for
    ``model_flux`` (n_pix,) on the observed grid."""
    resid = self._compute_residuals(model_flux)
    return float(jnp.sum(jnp.where(self.mask, resid ** 2, 0.0)))

residuals

residuals(model_flux)

Per-pixel normalised residuals (n_pix,); masked pixels are NaN.

Source code in ceridwen/observation/spectrum.py

def residuals(self, model_flux):
    """Per-pixel normalised residuals (n_pix,); masked pixels are NaN."""
    resid = self._compute_residuals(model_flux)
    return jnp.where(self.mask, resid, jnp.nan)

log_likelihood

log_likelihood(model_flux)

Gaussian log-likelihood of model_flux (n_pix,), including the sigma_eff normalisation and the GP term when self.noise is set.

Source code in ceridwen/observation/spectrum.py

def log_likelihood(self, model_flux):
    """Gaussian log-likelihood of ``model_flux`` (n_pix,), including the
    sigma_eff normalisation and the GP term when ``self.noise`` is set."""
    resid, sigma_r = self._compute_residuals(model_flux, return_sigma=True)
    mask = self.mask

    safe_sigma = jnp.where(mask, sigma_r, 1.0)
    lognorm    = float(
        -0.5 * jnp.sum(jnp.where(mask, jnp.log(safe_sigma ** 2), 0.0))
    )

    if self.noise is not None and self._wavelength is not None:
        # GP term already includes the white-noise identity and -n/2 ln(2 pi)
        lnL_struct = float(self.noise.log_likelihood(
            np.array(resid),
            np.array(self._wavelength),
            np.array(mask),
        ))
    else:
        r = jnp.where(mask, resid, 0.0)
        n = float(jnp.sum(mask))
        lnL_struct = float(
            -0.5 * jnp.sum(r ** 2) - 0.5 * n * jnp.log(2.0 * jnp.pi)
        )

    return float(lnL_struct + lognorm)

fit_polynomial_calibration

fit_polynomial_calibration(model_flux, order=3)

Weighted least-squares Chebyshev calibration P(lambda) with data ~ P * model_flux over unmasked pixels; returns (coeffs (order+1,), P * model_flux (n_pix,)).

Source code in ceridwen/observation/spectrum.py

def fit_polynomial_calibration(self, model_flux, order: int = 3):
    """Weighted least-squares Chebyshev calibration P(lambda) with
    data ~ P * model_flux over unmasked pixels; returns
    (coeffs (order+1,), P * model_flux (n_pix,))."""
    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)

    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)

    A = np.polynomial.chebyshev.chebvander(x, order)  # (n_pix, order+1)

    A_w = (A * mf[:, None]) / sigma[:, None]
    y_w = data / sigma

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

    coeffs, _, _, _ = np.linalg.lstsq(A_wm, y_wm, rcond=None)

    poly_vals       = A @ coeffs
    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,
    components=None,
    sigma_v=200.0,
    **kwargs
)

Bases: Observation

Observed emission-line fluxes with line indices, rest wavelengths and uncertainties.

Parameters:

Name	Type	Description	Default
line_ind	array-like of int -- indices of the observed lines in the SPS emission-line array		required
line_names	list of str, optional -- one name per line; needed by ``mask_by_name`` / ``select_by_name``		None
wavelength	array-like of float, vacuum rest-frame Angstrom		None
flux	array-like of float, same units as the model prediction (typically erg/s/cm^2)		required
uncertainty	array-like of float, same units as the model prediction (typically erg/s/cm^2)		required
components	list of sequences of float, optional -- per observed line, the vacuum rest wavelengths	[Angstrom] of all grid lines summed into that measurement (unresolved doublets); default one per line	None
upper_limit	array-like of bool, optional -- True treats the line as a non-detection: chi^2 penalises	only model > data	None
sigma_v	float, km/s -- Gaussian aperture width of ``predict`` (integration of a painted	spectrum); CSPBasis.predict does not use it, it reads the line fluxes from the grid	200.0

Source code in ceridwen/observation/lines.py

def __init__(
    self,
    line_ind,
    line_names  = None,
    wavelength  = None,
    name        = None,
    upper_limit = None,
    components  = None,
    sigma_v     = 200.0,
    **kwargs,
):
    self.sigma_v = float(sigma_v)
    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)
    )
    n_lines = int(np.atleast_1d(line_ind).size)
    if components is None:
        self.line_components = [
            (float(w),) for w in np.atleast_1d(np.asarray(wavelength, dtype=float))
        ]
    else:
        comps = [tuple(float(x) for x in np.atleast_1d(c)) for c in components]
        if len(comps) != n_lines:
            raise ValueError(
                f"components has {len(comps)} entries but there are "
                f"{n_lines} lines")
        if any(len(c) == 0 for c in comps):
            raise ValueError("every components entry needs >= 1 wavelength")
        self.line_components = comps
    super().__init__(name=name, **kwargs)

has_blends property

has_blends

True when at least one observed line sums several grid lines.

to_json

to_json()

Base JSON plus line_names and line_components.

Source code in ceridwen/observation/lines.py

def to_json(self):
    """Base JSON plus ``line_names`` and ``line_components``."""
    d = json.loads(super().to_json())
    d["line_names"] = self.line_names
    d["line_components"] = [list(c) for c in self.line_components]
    return json.dumps(d)

setup_for_model

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

Record the model wavelength grid [Angstrom, rest, increasing], aperture sigma_v [km/s] and redshift used to build the _W aperture matrix; must precede predict.

Source code in ceridwen/observation/lines.py

def setup_for_model(self, wave_model, zred: float = 0.0, sigma_v=None):
    """Record the model wavelength grid [Angstrom, rest, increasing], aperture ``sigma_v`` [km/s]
    and redshift used to build the ``_W`` aperture matrix; must precede ``predict``."""
    self._W_args = (np.asarray(wave_model, dtype=np.float64),
                    float(self.sigma_v if sigma_v is None else sigma_v), float(zred))
    self.__dict__.pop("_W", None)

predict

predict(spectrum, wave_model)

Return _W @ spectrum: Gaussian-aperture line fluxes, shape (n_lines,), from an F_nu model spectrum; wave_model is unused, setup_for_model must have been called.

Source code in ceridwen/observation/lines.py

def predict(self, spectrum, wave_model):
    """Return ``_W @ spectrum``: Gaussian-aperture line fluxes, shape (n_lines,), from an F_nu
    model spectrum; ``wave_model`` is unused, ``setup_for_model`` must have been called."""
    if getattr(self, "_W_args", None) is None:
        raise RuntimeError(
            "Lines.predict() called before setup_for_model(): call "
            "lines.setup_for_model(wave_model) once before the first "
            "predict / JIT trace.")
    return jnp.asarray(self._W) @ spectrum

chi_sq

chi_sq(model_fluxes)

Return chi^2 over unmasked lines; upper-limit lines are penalised only when model > data.

Source code in ceridwen/observation/lines.py

def chi_sq(self, model_fluxes):
    """Return chi^2 over unmasked lines; upper-limit lines are penalised only when model > data."""
    mf    = jnp.asarray(model_fluxes, dtype=float)
    resid = (self.flux - mf) / self.uncertainty

    if self.upper_limit is not None:
        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)

Return per-line (data - model) / sigma, shape (n_lines,); masked lines NaN, upper-limit lines with model <= data set to 0.

Source code in ceridwen/observation/lines.py

def residuals(self, model_fluxes):
    """Return per-line ``(data - model) / sigma``, shape (n_lines,); masked lines NaN,
    upper-limit lines with model <= data set to 0."""
    mf    = jnp.asarray(model_fluxes, dtype=float)
    resid = (self.flux - mf) / self.uncertainty

    if self.upper_limit is not None:
        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)

Set mask = False for lines whose name is in names; no-op without line_names.

Source code in ceridwen/observation/lines.py

def mask_by_name(self, names):
    """Set ``mask = False`` for lines whose name is in ``names``; no-op without ``line_names``."""
    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 containing only the named lines (KeyError on an unknown name).

Source code in ceridwen/observation/lines.py

def select_by_name(self, names):
    """Return a new ``Lines`` containing only the named lines (KeyError on an unknown name)."""
    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,
        components  = [self.line_components[i] for i in idx],
        name        = self.name + "_sel",
    )

Model

ceridwen.model.SedModel

SedModel(
    csp,
    observations,
    priors=None,
    transforms=None,
    free_param_init=None,
    zred=0.0,
    cosmo=None,
    lumdist_mpc=None,
    kinematics=None,
    broaden_photometry=True,
)

Parameter manager and prediction layer: predict(theta) returns a dict keyed by observation name, log_prob(theta) the summed log-prior on free parameters.

Parameters:

Name	Type	Description	Default
priors	dict[str, Prior] -- free-parameter name -> prior; absent names get a flat improper prior		None
transforms	dict[str, callable] -- derived CSP parameter name -> ``fn(free_theta)``; derived names leave the free set		None
free_param_init	dict[str, Array] -- initial values of the free parameters replacing derived ones		None
zred	float -- fixed redshift; zred = 0 without ``lumdist_mpc`` applies NO flux factor, so predictions	are L_sun/Hz x 10^logmass, not maggies	0.0
cosmo			None
lumdist_mpc	float, Mpc -- explicit luminosity distance replacing D_L(zred) in the flux factor		None
kinematics	Kinematics -- the galaxy's stellar / gas velocity dispersions [km/s], fixed	(float) or sampled (theta key); default DEFAULT_KINEMATICS = 300 km/s, stars and gas	None
broaden_photometry	bool -- apply the kinematic broadening to the spectrum entering the	filters (default True; below 5e-4 mag for broad bands, per cent for a narrow band on a line)	True

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,
    cosmo=None,
    lumdist_mpc: float | None = None,
    kinematics: Kinematics | None = None,
    broaden_photometry: bool = True,
):
    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)
    if kinematics is None:
        kinematics = DEFAULT_KINEMATICS
    if not isinstance(kinematics, Kinematics):
        raise TypeError("kinematics must be a ceridwen.broadening.Kinematics "
                        f"(e.g. Kinematics(sigma_gal=250.0)), got {type(kinematics).__name__}")
    self.kinematics   = kinematics
    self.broaden_photometry = bool(broaden_photometry)

    if not hasattr(csp, "cosmo"):
        raise TypeError(
            f"{type(csp).__name__} carries no .cosmo; build the CSP with "
            "cosmo=Cosmology.planck18() (or another Cosmology)")
    if cosmo is not None and cosmo != csp.cosmo:
        raise ValueError(
            "SedModel(cosmo=...) differs from the CSP's cosmology, and the "
            "CSP is what evaluates distances and ages:\n"
            f"    CSP     : {csp.cosmo.describe()}\n"
            f"    SedModel: {cosmo.describe()}\n"
            "Set it once, at CSP construction: CSPBasis(ssp, ..., cosmo=...)"
        )

    self.lumdist_mpc = None
    if lumdist_mpc is not None:
        self.lumdist_mpc = float(lumdist_mpc)
        if not (self.lumdist_mpc > 0.0) or self.lumdist_mpc != self.lumdist_mpc:
            raise ValueError(f"lumdist_mpc must be a finite positive distance "
                             f"in Mpc, got {lumdist_mpc}")
        if "zred" in self.transforms:
            raise ValueError("lumdist_mpc cannot be combined with a 'zred' transform")
        if self.zred > 0.0:
            warnings.warn(
                f"lumdist_mpc = {self.lumdist_mpc:g} Mpc replaces D_L(zred = "
                f"{self.zred:g}) = {float(csp.cosmo.luminosity_distance(self.zred)):.1f} "
                "Mpc in the flux factor; (1+zred) still comes from zred",
                stacklevel=2)

    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}"
        )

    self.theta_init  = dict(csp.theta_init)
    self.param_names = list(csp.param_names)
    self.wave        = csp.wave

    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)

    unknown = sorted(set(self.priors) - set(self.param_names))
    if unknown:
        raise ValueError(
            f"priors given for {unknown}, which are not sampled parameters "
            f"(a prior on a derived or misspelled name would be silently ignored); "
            f"the sampled parameters are {self.param_names}")
    unpriored = [p for p in self.param_names if p not in self.priors]
    if unpriored:
        warnings.warn(
            f"no prior for sampled parameter(s) {unpriored}: NUTS treats them as "
            "improper flat, nested sampling refuses them", stacklevel=2)

    self.kinematics.validate_theta(set(self.theta_init) | set(self.transforms), self.priors)
    if hasattr(self.csp, "register_known_theta_keys"):
        self.csp.register_known_theta_keys(
            set(self.param_names) | set(self.priors) | set(self.transforms)
            | set(self.kinematics.free_keys)
        )

    self._zred_fixed = None
    self._lumdist_fixed = None
    if (self.zred != 0.0 or self.lumdist_mpc is not None) and "zred" not in self.transforms:
        self._zred_fixed = jnp.array([self.zred])
        if self.lumdist_mpc is not None:
            self._lumdist_fixed = jnp.array([self.lumdist_mpc])

    self.zred_is_free = ("zred" in self.param_names) or ("zred" in self.transforms)
    grid_rescaled = (bool(getattr(csp, "track_zred_age", False))
                     and self._zred_fixed is not None) or ("lookback_time" in self.transforms)
    if (not self.zred_is_free and not grid_rescaled
            and hasattr(csp, "sfh_times") and hasattr(csp, "age_at")):
        oldest = float(csp.sfh_times[-1]) / 1e9
        age = float(csp.age_at(self.zred))
        if oldest > age * (1.0 + 5e-3):
            raise ValueError(
                f"the oldest SFH node, {oldest:.3f} Gyr of lookback time, "
                f"predates the Universe at zred = {self.zred:g}: age = "
                f"{age:.3f} Gyr under {csp.cosmo.describe()}.  Build the "
                f"grid from csp.age_at(zred) (or cosmo.age(zred)), e.g. "
                f"lookback_time=jnp.linspace(0.0, {age:.3f}, n)"
            )

    self.setup_observations()

    if (self.zred == 0.0 and self.lumdist_mpc is None and not self.zred_is_free
            and "zred" not in self.transforms and self.observations):
        warnings.warn(
            "SedModel(zred=0) applies NO flux factor: predictions are in "
            "L_sun/Hz x 10^logmass, not maggies.  Give zred= (Hubble flow) or "
            "lumdist_mpc= (nearby object) for physical units, or ignore this if "
            "the data are in the same unitless convention",
            stacklevel=2)

cosmo property

cosmo

The CSP's cosmology (read-only).

obs_dict property

obs_dict

Observations keyed by obs.name.

n_obs property

n_obs

Number of registered observation objects.

setup_observations

setup_observations()

Build every observation's projection for this model's grid, redshift and kinematics (called by __init__ and by fitSED when it replaces the observations); drops the cached jitted predictors.

Source code in ceridwen/model/model.py

def setup_observations(self):
    """Build every observation's projection for this model's grid, redshift and
    kinematics (called by ``__init__`` and by ``fitSED`` when it replaces the
    observations); drops the cached jitted predictors."""
    names = [o.name for o in self.observations]
    if len(set(names)) != len(names):
        raise ValueError(f"observation names must be unique, got {names}")
    neb = getattr(self.csp, "neb", None)
    lib = getattr(self.csp, "lib_resolution", None)
    for obs in self.observations:
        kind = getattr(obs, "_kind", None)
        if kind == "spectrum":
            zr = self._spectrum_zred_range(obs) if self.zred_is_free else None
            lines_rest = None if neb is None else neb.nebem_line_pos
            if neb is None and getattr(obs, "marginalize_elines", False):
                # no nebular grid: the marginalised lines come from FSPS's line list
                from ..likelihood.eline_marginal import line_table_for, refuse_without_grid
                refuse_without_grid(self.csp, obs, self.observations)
                lines_rest = line_table_for(self.csp)["wave"]
            obs.setup_for_model(
                self.wave, zred=(self._spectrum_zred_ref(zr) if zr else self.zred),
                kinematics=self.kinematics, lib_resolution=lib,
                line_wave_rest=lines_rest,
                zred_range=zr)
        elif kind == "photometry":
            obs.setup_for_model(self.wave, zred=self.zred)
            obs.free_z = bool(self.zred_is_free)
            obs.setup_broadening(
                self.wave, self.zred,
                self.kinematics if self.broaden_photometry else None,
                free_z=bool(getattr(obs, "free_z", False)), neb=neb)
        else:
            obs.setup_for_model(self.wave, zred=self.zred)
    for cached in ("_predict_jit_fn", "_predict_vmap_fn"):
        self.__dict__.pop(cached, None)
    from ..likelihood.eline_marginal import build_eline_system
    self._eline_system = build_eline_system(self)

apply_transforms

apply_transforms(free_theta)

Return free_theta plus every derived parameter fn(free_theta).

Source code in ceridwen/model/model.py

def apply_transforms(self, free_theta: dict[str, Array]) -> dict[str, Array]:
    """Return ``free_theta`` plus every derived parameter ``fn(free_theta)``."""
    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)

Predictions keyed by observation name from the free-parameter dict. A fixed zred/lumdist_mpc is injected into the CSP theta here; without a 'zred' the CSP applies no flux factor and the outputs are not maggies.

Source code in ceridwen/model/model.py

def predict(self, theta: dict[str, Array]) -> dict[str, Array]:
    """Predictions keyed by observation name from the free-parameter dict.
    A fixed zred/lumdist_mpc is injected into the CSP theta here; without a
    'zred' the CSP applies no flux factor and the outputs are not maggies."""
    model_theta = self.apply_transforms(theta)
    if self._zred_fixed is not None and "zred" not in model_theta:
        if model_theta is theta:
            model_theta = dict(model_theta)
        model_theta["zred"] = self._zred_fixed
    if self._lumdist_fixed is not None and "lumdist_mpc" not in model_theta:
        if model_theta is theta:
            model_theta = dict(model_theta)
        model_theta["lumdist_mpc"] = self._lumdist_fixed
    return self.csp.predict(model_theta, self.observations,
                            kinematics=self.kinematics,
                            broaden_photometry=self.broaden_photometry)

predict_with_elines

predict_with_elines(theta)

(predictions, aux) for the emission-line marginalisation: the predictions with the fitted lines removed, and aux = {"prior_mean", "cols"} (their CLOUDY fluxes and the per-observation design columns). Needs a Spectrum with marginalize_elines=True.

Source code in ceridwen/model/model.py

def predict_with_elines(self, theta: dict[str, Array]):
    """``(predictions, aux)`` for the emission-line marginalisation: the predictions with
    the fitted lines removed, and ``aux = {"prior_mean", "cols"}`` (their CLOUDY fluxes and
    the per-observation design columns).  Needs a Spectrum with marginalize_elines=True."""
    if getattr(self, "_eline_system", None) is None:
        raise ValueError("no Spectrum of this model has marginalize_elines=True")
    model_theta = self.apply_transforms(theta)
    if self._zred_fixed is not None and "zred" not in model_theta:
        model_theta = dict(model_theta)
        model_theta["zred"] = self._zred_fixed
    if self._lumdist_fixed is not None and "lumdist_mpc" not in model_theta:
        model_theta = dict(model_theta)
        model_theta["lumdist_mpc"] = self._lumdist_fixed
    return self.csp.predict(model_theta, self.observations,
                            kinematics=self.kinematics,
                            broaden_photometry=self.broaden_photometry,
                            eline_system=self._eline_system)

predict_jit

predict_jit(theta)

JIT-compiled :meth:predict (compiled on first call).

Source code in ceridwen/model/model.py

def predict_jit(self, theta: dict[str, Array]) -> dict[str, Array]:
    """JIT-compiled :meth:`predict` (compiled on first call)."""
    return self._predict_jit_fn(theta)

predict_vmap

predict_vmap(theta_batch)

Vectorised :meth:predict over a leading batch axis of every theta entry.

Source code in ceridwen/model/model.py

def predict_vmap(
    self,
    theta_batch: dict[str, Array],
) -> dict[str, Array]:
    """Vectorised :meth:`predict` over a leading batch axis of every theta entry."""
    return self._predict_vmap_fn(theta_batch)

ln_prior

ln_prior(theta)

Scalar sum of prior.logpdf(theta[p]) over registered priors.

Source code in ceridwen/model/model.py

def ln_prior(self, theta: dict[str, Array]) -> Array:
    """Scalar sum of ``prior.logpdf(theta[p])`` over registered priors."""
    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.

Source code in ceridwen/model/model.py

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

summary

summary()

Multi-line summary of parameters, transforms, observations and CSP setup.

Source code in ceridwen/model/model.py

def summary(self) -> str:
    """Multi-line summary of parameters, transforms, observations and CSP setup."""
    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} Å",
        f"Cosmology          : {self.cosmo.describe()}",
        f"Redshift           : {self._redshift_line()}",
        f"Kinematics         : {self._kinematics_line()}",
        "",
        "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.)",
        ]

    lines += ["", "Observations", "-" * 40]
    for obs in self.observations:
        lines.append(f"  {obs!r}")
    es = getattr(self, "_eline_system", None)
    if es is not None:
        lines += ["", "Emission lines", "-" * 40, f"  {es.describe()}"]

    return "\n".join(lines)

display

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

Draw the model as a PGM diagram; returns (fig, ax) when return_fig.

Source code in ceridwen/model/model.py

def display(
    self,
    ax=None,
    figsize: tuple[float, float] | None = None,
    return_fig: bool = False,
):
    """Draw the model as a PGM diagram; returns ``(fig, ax)`` when ``return_fig``."""
    import matplotlib.pyplot as plt
    import matplotlib.patches as mpatches
    from matplotlib.patches import FancyBboxPatch

    C = dict(
        bg        = "#FFFFFF",
        param_fc  = "#FFFFFF",
        param_ec  = "#222222",
        vec_fc    = "#EEF3FF",
        vec_ec    = "#556BBB",
        sed_fc    = "#E6F2FB",
        sed_ec    = "#1A6098",
        phot_fc   = "#FFF4E6",  phot_ec = "#C95800",
        spec_fc   = "#EDFAED",  spec_ec = "#276929",
        line_fc   = "#F5EEFF",  line_ec = "#6A22A8",
        data_fc   = "#37474F",
        data_ec   = "#1A252B",
        data_tc   = "#FFFFFF",
        arr_prior = "#BBBBBB",
        arr_fwd   = "#555555",
        arr_obs   = "#777777",
    )

    def _obs_colors(obs):
        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"]))

    _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}$"
            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}}$"

    C["tr_fc"] = "#FFF8E1"
    C["tr_ec"] = "#E65100"
    C["arr_tr"] = "#E65100"

    n_p = len(self.param_names)
    n_o = len(self.observations)
    has_transforms = bool(self.transforms)

    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")

    _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
    y_sed   = fh / 2.0 + 0.10 + (_tr_shift / 2)
    y_obs   = 1.90
    y_data  = 0.62

    r_p  = 0.34
    r_v  = 0.36
    r_d  = 0.30

    x0, x1 = fw * 0.07, fw * 0.93
    x_p = ([fw / 2] if n_p == 1
            else list(np.linspace(x0, x1, n_p)))

    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

    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)

    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")
        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")

    for i, pname in enumerate(self.param_names):
        xp  = x_p[i]
        vec = _is_vector(pname)

        if vec:
            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)
            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)

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

    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")

            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")

            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(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)

    sed_w = max(3.8, min(fw * 0.42, n_p * 0.88))
    sed_h = 0.84

    rect(x_sed, y_sed, sed_w, sed_h,
         C["sed_fc"], C["sed_ec"], lw=2.2)
    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")

    for i, pname in enumerate(self.param_names):
        if has_transforms:
            continue
        xp   = x_p[i]
        vec  = _is_vector(pname)
        r    = r_v if vec else r_p
        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)

    for j, obs in enumerate(self.observations):
        xo       = x_o[j]
        fc, ec  = _obs_colors(obs)
        ow, oh  = 1.60, 0.64

        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")

        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)

        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")

        arrow(xo, y_obs - oh / 2 - 0.04,
              xo, y_data + r_d + 0.04,
              ec, lw=1.1)

        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)

        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")

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

    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")
    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")

    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,
    )


    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)

Unit-mass SFH weight vector (n,) from logsfr_ratios[i] = log10(SFR[i]/SFR[i+1]) (n-1,), index 0 = today. With sfh_times_yr (lookback yr, (n,)) the trapezoidal integral is normalised to 1 Msun; otherwise sum(sfh) = 1.

Source code in ceridwen/model/transforms.py

def logsfr_ratios_to_sfh(
    logsfr_ratios,
    sfh_times_yr=None,
):
    """Unit-mass SFH weight vector (n,) from ``logsfr_ratios[i] = log10(SFR[i]/SFR[i+1])``
    (n-1,), index 0 = today. With ``sfh_times_yr`` (lookback yr, (n,)) the trapezoidal
    integral is normalised to 1 Msun; otherwise sum(sfh) = 1."""
    ratios  = jnp.asarray(logsfr_ratios, dtype=float)
    log_sfr = jnp.concatenate([jnp.zeros(1),
                                -jnp.cumsum(ratios)])
    sfr     = 10.0 ** log_sfr

    if sfh_times_yr is not None:
        times = jnp.asarray(sfh_times_yr, dtype=float)
        dt    = jnp.abs(jnp.diff(times))
        w_lo  = jnp.concatenate([jnp.zeros(1), dt])
        w_hi  = jnp.concatenate([dt, jnp.zeros(1)])
        w     = 0.5 * (w_lo + w_hi)
        # normalise to total mass (1 Msun), NOT to mean SFR
        total_mass = jnp.sum(sfr * w)
        sfh   = sfr / total_mass
    else:
        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 logpdf (also __call__), sample, unit_transform (inverse CDF) and inverse_unit_transform (CDF); bounds gives the support.

Prior dataclass

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

Bases: ABC

Prior base class delegating to a TFP-JAX distribution; subclasses define prior_params and 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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

serialize

serialize()

JSON-ready description: type, parameters (lists for arrays), name.

Source code in ceridwen/sampler/priors.py

def serialize(self) -> dict:
    """JSON-ready description: type, parameters (lists for arrays), name."""
    out = {"type": type(self).__name__, "name": self.name}
    for k, v in self.params.items():
        a = np.asarray(v)
        out[k] = a.tolist() if a.ndim else float(a)
    return out

tfp_dist abstractmethod

tfp_dist()

Return the TFP-JAX distribution built from self.params.

Source code in ceridwen/sampler/priors.py

@abc.abstractmethod
def tfp_dist(self) -> tfd.Distribution:
    """Return the TFP-JAX distribution built from ``self.params``."""
    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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

TopHat dataclass

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

Bases: Uniform

Alias of Uniform kept for backwards compatibility.

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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

Normal dataclass

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

Bases: Prior

Gaussian prior with parameters mean, sigma.

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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

MultivariateNormalPrior dataclass

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

Bases: Prior

Multivariate Gaussian prior.

Parameters:

Name	Type	Description	Default
mean	(d,)		required
Sigma	(d, d) -- covariance matrix		required

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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

range property

range

Mean -/+ 4 sigma per dimension (plotting range).

ClippedNormal dataclass

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

Bases: Prior

Gaussian prior truncated to [low, high]; parameters mean, sigma, 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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

LogNormal dataclass

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

Bases: Prior

Log-normal prior; mode and sigma are the mean and std of ln(x).

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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

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

StudentT dataclass

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

Bases: Prior

Student's t prior with parameters mean, scale, df (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.pop(external))
    if kwargs:
        raise TypeError(
            f"{type(self).__name__} got unknown argument(s) {sorted(kwargs)}; "
            f"it takes {list(alias.values())}")
    missing = [external for intrinsic, external in alias.items() if intrinsic not in params]
    if missing:
        raise TypeError(f"{type(self).__name__} is missing argument(s) {missing}")

    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.

__call__

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

Return (lnl_total, 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]:
    """Return ``(lnl_total, LikelihoodOutput)``."""
    noise_out: NoiseModelOutput = self.noise_model.compute(
        sigma_obs, mu, mask, params, data=y
    )
    return lnlike_diag_gaussian(
        y, mu, noise_out.inv_var, noise_out.log_det, mask
    )

make_lnprobfn

make_lnprobfn(observations, model, prior)

Return a jitted log-posterior for one observation (needs .flux, .uncertainty, .mask).

Source code in ceridwen/likelihood/likelihood.py

def make_lnprobfn(
    self,
    observations : Any,
    model        : Any,
    prior        : Any,
) -> Callable[[dict[str, Array]], Array]:
    """Return a jitted log-posterior for one observation (needs ``.flux``, ``.uncertainty``, ``.mask``)."""
    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)
        noise_out = noise_model.compute(sigma_obs, mu, mask, theta, data=y)
        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.DiagonalGaussianLikelihoodWithUpperLimits dataclass

DiagonalGaussianLikelihoodWithUpperLimits(
    noise_model=DiagonalNoiseModel(),
)

Bases: LikelihoodBase

Diagonal Gaussian log-likelihood honouring per-datum upper-limit flags (one-sided penalty; reduces to :class:DiagonalGaussianLikelihood when no flags are set).

__call__

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

Return (lnl_total, LikelihoodOutput); is_upper_limit=None means all detections.

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,
    is_upper_limit  : Optional[Array]            = None,
) -> tuple[Array, LikelihoodOutput]:
    """Return ``(lnl_total, LikelihoodOutput)``; ``is_upper_limit=None`` means all detections."""
    noise_out: NoiseModelOutput = self.noise_model.compute(
        sigma_obs, mu, mask, params, data=y
    )
    if is_upper_limit is None:
        is_upper_limit = jnp.zeros_like(mask, dtype=bool)
    return lnlike_diag_gaussian_with_upper_limits(
        y, mu, noise_out.inv_var, noise_out.log_det, mask, is_upper_limit,
    )

make_lnprobfn

make_lnprobfn(observations, model, prior)

Return a jitted log-posterior using observations.upper_limit (all-False if absent).

Source code in ceridwen/likelihood/likelihood.py

def make_lnprobfn(
    self,
    observations : Any,
    model        : Any,
    prior        : Any,
) -> Callable[[dict[str, Array]], Array]:
    """Return a jitted log-posterior using ``observations.upper_limit`` (all-False if absent)."""
    y         : Array = observations.flux
    sigma_obs : Array = observations.uncertainty
    mask      : Array = observations.mask
    is_ul = getattr(observations, "upper_limit", None)
    if is_ul is None:
        is_ul = jnp.zeros_like(mask, dtype=bool)
    else:
        is_ul = jnp.asarray(is_ul, dtype=bool)
    noise_model = self.noise_model

    @jax.jit
    def lnprobfn(theta: dict[str, Array]) -> Array:
        mu = model.predict(theta)
        noise_out = noise_model.compute(sigma_obs, mu, mask, theta, data=y)
        lnl, _    = lnlike_diag_gaussian_with_upper_limits(
            y, mu, noise_out.inv_var, noise_out.log_det, mask, is_ul,
        )
        lnp = prior.log_prob(theta)
        return lnl + lnp

    return lnprobfn

ceridwen.likelihood.MultiObservationLikelihood dataclass

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

Bases: LikelihoodBase

Sum of independent likelihoods over several observation keys.

Parameters:

Name	Type	Description	Default
keys	tuple of str -- observation keys, e.g. ``("phot", "spec", "lines")``		tuple()
likelihoods	tuple of LikelihoodBase -- one per key, same order		tuple()

__call__

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

Return (lnl_total, {key: LikelihoodOutput}); all inputs are dicts keyed like self.keys (is_upper_limit only needs the keys whose likelihood honours upper limits).

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,
    is_upper_limit : Optional[dict[str, Array]] = None,
) -> tuple[Array, dict[str, LikelihoodOutput]]:
    """Return ``(lnl_total, {key: LikelihoodOutput})``; all inputs are dicts keyed like ``self.keys``
    (``is_upper_limit`` only needs the keys whose likelihood honours upper limits)."""
    lnl_total = jnp.zeros(())
    aux: dict[str, LikelihoodOutput] = {}
    for key, lhood in zip(self.keys, self.likelihoods):
        ul = None if is_upper_limit is None else is_upper_limit.get(key)
        if ul is not None:
            lnl_i, aux_i = lhood(y[key], mu[key], sigma_obs[key], mask[key], params,
                                 is_upper_limit=ul)
        else:
            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)

Return a jitted log-posterior; observations and model.predict(theta) are dicts keyed like self.keys. Each observation's sky, calibration and upper_limit are honoured as in ceridwen.sampler.runner.run_sampler.

Source code in ceridwen/likelihood/likelihood.py

def make_lnprobfn(
    self,
    observations : dict[str, Any],
    model        : Any,
    prior        : Any,
) -> Callable[[dict[str, Array]], Array]:
    """Return a jitted log-posterior; ``observations`` and ``model.predict(theta)`` are dicts keyed
    like ``self.keys``.  Each observation's ``sky``, ``calibration`` and ``upper_limit`` are honoured
    as in ``ceridwen.sampler.runner.run_sampler``."""
    static_data = {}
    for key in self.keys:
        obs = observations[key]
        y = obs.flux
        sky = getattr(obs, "sky", None)
        if sky is not None:
            y = y - sky
        ul = getattr(obs, "upper_limit", None)
        ul = None if ul is None or not bool(jnp.any(ul)) else jnp.asarray(ul, dtype=bool)
        static_data[key] = (y, obs.uncertainty, obs.mask, getattr(obs, "calibration", None), ul)
    keys        = self.keys
    likelihoods = self.likelihoods

    if getattr(model, "_eline_system", None) is not None:
        from .eline_marginal import joint_loglike

        @jax.jit
        def lnprobfn_elines(theta: dict[str, Array]) -> Array:
            return (joint_loglike(model, keys, likelihoods, static_data, theta)
                    + prior.log_prob(theta))

        return lnprobfn_elines

    @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, calib_k, ul_k = static_data[key]
            mu_k = predictions[key]
            if calib_k is not None:
                mu_k = mu_k * calib_k
            if ul_k is not None:
                lnl_k, _ = lhood(y_k, mu_k, sig_k, mask_k, params=theta, is_upper_limit=ul_k)
            else:
                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 model to observations with sampler ("nested" or "nuts") and return the SamplingResult; writes output_dir/filename (HDF5) and a .log with the same stem.

Parameters:

Name	Type	Description	Default
observations	list[Observation] -- replaces ``model.observations`` and re-runs ``setup_for_model`` at ``model.zred``		None
sampler_kwargs	dict -- forwarded to the sampler adapter constructor		None
vi	None, 'tril', 'iaf', or a VI map -- NUTS-only variational preconditioning		None
vi_kwargs	dict -- forwarded to VI training / map constructor		None

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 ``model`` to ``observations`` with ``sampler`` ("nested" or "nuts") and return the
    ``SamplingResult``; writes ``output_dir/filename`` (HDF5) and a ``.log`` with the same stem.

    Parameters
    ----------
    observations : list[Observation] -- replaces ``model.observations`` and re-runs ``setup_for_model`` at ``model.zred``
    sampler_kwargs : dict -- forwarded to the sampler adapter constructor
    vi : None, 'tril', 'iaf', or a VI map -- NUTS-only variational preconditioning
    vi_kwargs : dict -- forwarded to VI training / map constructor
    """
    from .likelihood.likelihood import MultiObservationLikelihood
    from .sampler.runner import run_sampler

    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."
        )

    if observations is not None:
        model.observations = list(observations)
        model.setup_observations()

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

    _t0_likelihood = time.perf_counter()
    obs_dict = model.obs_dict
    keys = tuple(obs_dict.keys())
    likelihoods = tuple(_likelihood_for(obs_dict[k], model.param_names) for k in keys)
    multi_likelihood = MultiObservationLikelihood(
        keys=keys,
        likelihoods=likelihoods,
    )
    _t_likelihood = time.perf_counter() - _t0_likelihood

    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:
        _devices = jax.devices()
        _backend = jax.default_backend().upper()
        _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"  Cosmology   : {model.cosmo.describe()}")
        logger.info(f"  Redshift    : {model._redshift_line()}")
        logger.info(f"  Kinematics  : {model._kinematics_line()}")
        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)}")
        for k, lh in zip(keys, likelihoods):
            logger.info(f"    {k}: {_describe_likelihood(obs_dict[k], lh)}")
        if getattr(model, "_eline_system", None) is not None:
            logger.info(f"  Emission lines: {model._eline_system.describe()}")
        logger.info(f"  Output      : {output_path}")
        logger.info(f"  Log         : {log_path}")
        logger.info(f"  Likelihood build: {_t_likelihood:.3f} s")

        _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")
        logger.info(f"  Sampler settings: {_describe_adapter(adapter, model)}")

        _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()}")

        _t0_h5 = time.perf_counter()
        write_result_h5(output_path, model, result, verbose=verbose,
                        likelihood=multi_likelihood)
        _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:
        logger.removeHandler(_file_handler)
        _file_handler.close()

ceridwen.fit.load_result_h5

load_result_h5(path)

Rebuild a SamplingResult from a result HDF5 file (the forward model itself is not restored).

Source code in ceridwen/fit.py

def load_result_h5(path: str | Path):
    """Rebuild a ``SamplingResult`` from a result HDF5 file (the forward model itself is not restored)."""
    import h5py
    from .sampler.runner import SamplingResult

    path = Path(path)
    with h5py.File(path, "r") as f:
        samp = f["samples"]
        param_names = list(f["model"]["param_names"].asstr()[()])

        samples = {p: jnp.asarray(np.array(samp[p])) for p in param_names}
        log_likelihoods = jnp.asarray(np.array(samp["log_likelihoods"]))
        log_weights     = jnp.asarray(np.array(samp["log_weights"]))
        llb = (jnp.asarray(np.array(samp["log_likelihoods_birth"]))
               if "log_likelihoods_birth" in samp else None)

        raw = {}
        if "num_chains" in samp.attrs:
            raw["num_chains"] = int(samp.attrs["num_chains"])
        if "num_samples_per_chain" in samp.attrs:
            raw["num_samples"] = int(samp.attrs["num_samples_per_chain"])
        if "num_warmup" in samp.attrs:
            raw["num_warmup"] = int(samp.attrs["num_warmup"])

        return SamplingResult(
            samples               = samples,
            log_evidence          = float(samp.attrs.get("log_evidence", float("nan"))),
            log_evidence_err      = float(samp.attrs.get("log_evidence_err", float("nan"))),
            log_weights           = log_weights,
            log_likelihoods       = log_likelihoods,
            param_names           = param_names,
            n_likelihood_calls    = int(samp.attrs.get("n_likelihood_calls", -1)),
            wall_time_s           = float(samp.attrs.get("wall_time_s", float("nan"))),
            sampler_name          = str(samp.attrs.get("sampler_name", "unknown")),
            log_likelihoods_birth = llb,
            raw                   = raw or None,
        )

ceridwen.fit.read_result_h5

read_result_h5(path)

Read a result HDF5 file into a nested dict with keys 'obs', 'model', 'samples' (and 'elines' when the fit marginalised the emission lines).

Source code in ceridwen/fit.py

def read_result_h5(path: str | Path) -> dict:
    """Read a result HDF5 file into a nested dict with keys ``'obs'``, ``'model'``, ``'samples'``
    (and ``'elines'`` when the fit marginalised the emission lines)."""
    import h5py

    out = {"obs": {}, "model": {}, "samples": {}}

    with h5py.File(path, "r") as f:
        for obs_name in f["obs"]:
            og = f["obs"][obs_name]
            obs_data = {k: np.array(og[k]) for k in og}
            for attr_name in og.attrs:
                obs_data[attr_name] = og.attrs[attr_name]
            out["obs"][obs_name] = obs_data

        mod = f["model"]
        out["model"]["param_names"] = list(mod["param_names"].asstr()[()])
        out["model"]["wave"] = np.array(mod["wave"])

        out["model"]["theta_init"] = {}
        for name in mod["theta_init"]:
            out["model"]["theta_init"][name] = np.array(mod["theta_init"][name])

        out["model"]["priors"] = {}
        if "priors" in mod:
            for name in mod["priors"].attrs:
                raw = mod["priors"].attrs[name]
                try:
                    out["model"]["priors"][name] = json.loads(raw)
                except (json.JSONDecodeError, TypeError):
                    out["model"]["priors"][name] = raw

        for attr_name in mod.attrs:
            out["model"][attr_name] = mod.attrs[attr_name]
        if "cosmo_H0" in mod.attrs:
            from .cosmology import Cosmology
            out["model"]["cosmo"] = Cosmology.from_dict(dict(mod.attrs))

        samp = f["samples"]
        for dset_name in samp:
            out["samples"][dset_name] = np.array(samp[dset_name])

        for attr_name in samp.attrs:
            out["samples"][attr_name] = samp.attrs[attr_name]

        if "elines" in f:
            g = f["elines"]
            out["elines"] = {k: (list(g[k].asstr()[()]) if k == "names" else np.array(g[k]))
                             for k in g}
            for attr_name in g.attrs:
                out["elines"][attr_name] = g.attrs[attr_name]

    return out

ceridwen.fit.result_cosmology

result_cosmology(path)

The Cosmology a result file was fitted with (/model attrs cosmo_*); KeyError if absent.

Source code in ceridwen/fit.py

def result_cosmology(path: str | Path):
    """The ``Cosmology`` a result file was fitted with (``/model`` attrs ``cosmo_*``); KeyError if absent."""
    import h5py
    from .cosmology import Cosmology

    with h5py.File(Path(path), "r") as f:
        attrs = dict(f["model"].attrs)
    if "cosmo_H0" not in attrs:
        raise KeyError(f"{path} carries no cosmology attributes (written before "
                       "they were recorded); rebuild with the cosmology you used")
    return Cosmology.from_dict(attrs)

ceridwen.sampler.run_sampler

run_sampler(model, likelihood, adapter, rng_key)

Build JIT loglike_fn (summed over observations, no prior) and logprior_fn from model/likelihood and delegate to adapter.run.

Source code in ceridwen/sampler/runner.py

def run_sampler(
    model      : Any,
    likelihood : Any,
    adapter    : SamplerAdapter,
    rng_key    : Array,
) -> SamplingResult:
    """Build JIT ``loglike_fn`` (summed over observations, no prior) and
    ``logprior_fn`` from ``model``/``likelihood`` and delegate to ``adapter.run``."""
    _obs_dict    = model.obs_dict
    _keys        = tuple(likelihood.keys)
    _likelihoods = tuple(likelihood.likelihoods)
    _static_data = {}
    for key in _keys:
        obs = _obs_dict[key]
        y = obs.flux
        sky = getattr(obs, "sky", None)
        if sky is not None:
            y = y - sky
        ul = getattr(obs, "upper_limit", None)
        ul = None if ul is None or not bool(jnp.any(ul)) else jnp.asarray(ul, dtype=bool)
        _static_data[key] = (y, obs.uncertainty, obs.mask, getattr(obs, "calibration", None), ul)

    if getattr(model, "_eline_system", None) is not None:
        from ..likelihood.eline_marginal import joint_loglike

        @jax.jit
        def loglike_fn(theta: dict[str, Array]) -> Array:
            return joint_loglike(model, _keys, _likelihoods, _static_data, theta)

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

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

    @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, calib_k, ul_k = _static_data[key]
            mu_k = predictions[key]
            if calib_k is not None:
                mu_k = mu_k * calib_k
            if ul_k is not None:
                lnl_k, _ = lhood(y_k, mu_k, sig_k, mask_k, params=theta, is_upper_limit=ul_k)
            else:
                lnl_k, _ = lhood(y_k, mu_k, sig_k, mask_k, params=theta)
            lnl = lnl + lnl_k
        return lnl

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

    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=-5.0,
    verbose=True,
    checkpoint_interval_s=1200.0,
    checkpoint_dir=None,
)

Bases: SamplerAdapter

Nested-sampling adapter driving the NSS kernel with a Ceridwen SedModel.

Parameters:

Name	Type	Description	Default
priors	dict[str, Prior] -- every free parameter needs a proper prior		required
num_inner_steps	int -- inner MCMC steps per iteration; default n_dims * 5		None
num_delete	int -- live points removed per iteration; default max(1, num_live // 5), must be < num_live		None
logZ_tol	float -- stop when ln(Z_live / Z) < logZ_tol; default -5		-5.0
checkpoint_interval_s	float -- seconds between checkpoints; <= 0 disables		1200.0
checkpoint_dir	str -- falls back to $CERIDWEN_CHECKPOINT_DIR, then $CERIDWEN_RESCUE_DIR, else off		None

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 = -5.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
    self._num_delete      = num_delete
    self.logZ_tol        = float(logZ_tol)
    self.verbose         = bool(verbose)
    self.checkpoint_interval_s = float(checkpoint_interval_s)
    self._checkpoint_dir       = checkpoint_dir

load_checkpoint staticmethod

load_checkpoint(path)

Load a checkpoint/rescue pickle: {positions, loglikelihood, loglikelihood_birth, logZ, n_dead, partial}.

Source code in ceridwen/sampler/nested.py

@staticmethod
def load_checkpoint(path):
    """Load a checkpoint/rescue pickle: {positions, loglikelihood, loglikelihood_birth, logZ, n_dead, partial}."""
    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 nested sampling and return a SamplingResult; logprior_fn must be proper.

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 nested sampling and return a ``SamplingResult``; ``logprior_fn`` must be proper."""
    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 max(1, self.num_live // 5))

    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}"
        )

    rng_key, prior_key = jax.random.split(rng_key)
    particles = self._sample_prior(theta_init, prior_key)

    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)

    def _logz_fields(state):
        if hasattr(state, "logZ") and hasattr(state, "logZ_live"):
            return state.logZ, state.logZ_live
        ig = getattr(state, "integrator", None)
        if ig is not None and hasattr(ig, "logZ") and hasattr(ig, "logZ_live"):
            return ig.logZ, ig.logZ_live
        inner = getattr(state, "sampler_state", None)
        if inner is not None and hasattr(inner, "logZ") and hasattr(inner, "logZ_live"):
            return inner.logZ, inner.logZ_live
        raise AttributeError(
            f"Cannot locate logZ/logZ_live on {type(state).__name__} "
            f"(fields {[f for f in dir(state) if not f.startswith('_')]}); "
            "check your BlackJAX version.")

    def _logz(state):
        a, b = jax.device_get(_logz_fields(state))
        return float(a), float(b)

    _get_logZ = lambda s: _logz(s)[0]

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

    dead_list    = []
    n_like_calls = 0
    t_start      = time.perf_counter()

    _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)
    elif self.verbose:
        print("  [checkpoint] off (set checkpoint_dir= or $CERIDWEN_CHECKPOINT_DIR "
              "to write recoverable snapshots)", flush=True)

    logZ, logZ_live = _logz(live)
    with tqdm.tqdm(desc=f"NS  logZ={logZ:.1f}", unit=" dead",
                   disable=not self.verbose) as pbar:
        _iter = 0
        while logZ_live - logZ >= self.logZ_tol:
            rng_key, subkey = jax.random.split(rng_key)
            if _iter == 0 and self.verbose:
                pbar.write("  [step_fn] compiling the step kernel (one-time JIT)")
            _t_iter = time.perf_counter()
            live, dead_info = step_fn(subkey, live)
            logZ, logZ_live = _logz(live)
            _dt_iter = time.perf_counter() - _t_iter
            _iter += 1
            dead_list.append(dead_info)
            n_like_calls += num_delete * num_inner_steps
            pbar.update(num_delete)
            pbar.set_description(f"NS  logZ={logZ:.2f}  dlogZ={logZ_live - logZ:.2f}",
                                 refresh=False)
            pbar.set_postfix({"s/iter": f"{_dt_iter:.1f}"}, refresh=False)

            if _ckpt_on and (time.perf_counter() - _last_ckpt
                             >= self.checkpoint_interval_s):
                _p = self._dump_snapshot(
                    _ckpt_dir, live, dead_list, ns_utils,
                    logZ, tag="checkpoint", partial=True)
                _last_ckpt = time.perf_counter()
                if _p and self.verbose:
                    pbar.write(f"  [checkpoint] iter {_iter}: {_p}")

    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)"
        )

    _dead_positions, _dead_logl, _dead_logl_birth = self._finalise_dead(
        live, dead_list, ns_utils)

    _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,
                            finalised=(_dead_positions, _dead_logl, _dead_logl_birth))

    # dead set is in deletion order, NOT sorted by logL; weights follow sample order
    import numpy as np
    from .ns_weights import nested_log_weights, log_evidence_from_weights
    _logl_np    = np.asarray(_dead_logl)
    _birth_np   = np.asarray(_dead_logl_birth)
    _lw_np      = nested_log_weights(_logl_np, _birth_np)
    log_weights = jnp.asarray(_lw_np)
    log_Z       = None
    log_Z_err   = float("nan")
    try:
        from anesthetic import NestedSamples

        _raw   = _dead_positions
        _names = [n for n in _raw if n in theta_init]
        _cols  = [np.asarray(_raw[n]).reshape(_logl_np.shape[0], -1) for n in _names]
        _data  = np.hstack(_cols)
        _ns    = NestedSamples(
            _data,
            logL       = _logl_np,
            logL_birth = _birth_np,
            logzero    = float("nan"),
        )
        log_Z     = float(_ns.logZ())
        log_Z_err = float(_ns.logZ(12).std())
    except Exception:
        pass
    if log_Z is None:
        log_Z = log_evidence_from_weights(_lw_np)

    samples = {}
    for name in theta_init:
        arr = jnp.asarray(_dead_positions[name])
        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

NUTS adapter with window adaptation; bounded (uniform-prior) parameters are sampled in sigmoid/logit space, and an optional VI map (vi) whitens the target before sampling.

Parameters:

Name	Type	Description	Default
num_warmup	int -- adaptation steps per chain; default 1500 (200 with ``vi``)		None
num_samples	int -- post-warmup draws per chain		2000
initial_step_size	float -- leapfrog step before adaptation; default 0.01 (0.5 with ``vi``)		None
target_acceptance	float -- dual-averaging target		0.95
max_num_doublings	int -- max tree depth (2**n leapfrog steps)		10
dense_mass	bool -- full inverse mass matrix; default True (False with ``vi``)		None
bounds	dict -- name -> (low, high); None auto-detects from model priors		None
vi	None | 'tril' | 'iaf' | VariationalMap | TrainedMap -- variational preconditioning		None
vi_kwargs	dict -- map constructor kwargs plus ``num_steps``, ``batch_size``, ``lr0`` for training		None

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,
):
    _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)
    self.trained_map = None

run

run(loglike_fn, logprior_fn, theta_init, rng_key)

Run NUTS (one warmup, then all chains) and return a SamplingResult in constrained space.

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 NUTS (one warmup, then all chains) and return a ``SamplingResult`` in constrained space."""
    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

    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=...)")

    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")

    @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 _prior_and_jacobian(x):
        theta_flat = _to_constrained(x, lo, hi, is_bounded)
        theta = self._unflatten(theta_flat, theta_template)
        return logprior_fn(theta) + _log_jacobian(x, lo, hi, is_bounded)

    def _loglike_of(states_):
        return states_.logdensity - jax.vmap(_prior_and_jacobian)(states_.position)

    x_init_flat = self._flatten(theta_init)
    x_init = _to_unconstrained(x_init_flat, lo, hi, is_bounded)

    t_start = time.perf_counter()

    if self.verbose:
        print(f"\n  Warmup ({self.num_warmup} steps, "
              f"adapting step size + {'dense' if self.dense_mass else 'diagonal'} mass matrix)...")

    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)")

    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):
        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 = []
    all_loglikelihoods = []
    all_divergences = []
    all_infos = []
    _t_sample_chains = []
    _t_postproc_chains = []

    if use_pmap:
        if self.verbose:
            print(f"\n  Running {self.num_chains} chains in parallel "
                  f"across {self.num_chains} GPUs...", flush=True)

        def _replicate_state(state, n):
            return jax.tree.map(
                lambda x: jnp.broadcast_to(x, (n,) + x.shape), state
            )

        pmap_init = _replicate_state(warmup_state, self.num_chains)

        @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
        )
        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)")

        _t0_pp = time.perf_counter()
        for ci in range(self.num_chains):
            x_chain = pmap_states.position[ci]
            theta_chain = jax.vmap(
                lambda x: _to_constrained(x, lo, hi, is_bounded)
            )(x_chain)
            all_chain_positions.append(theta_chain)
            all_loglikelihoods.append(
                _loglike_of(jax.tree.map(lambda x: x[ci], pmap_states)))

            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:
        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)
            all_loglikelihoods.append(_loglike_of(states))

            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

    _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)

    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 ''})"
        )

    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,
        },
    )

Post-processing

ceridwen.postprocess.PostProcess

PostProcess(
    model,
    result,
    *,
    n_samples=None,
    seed=0,
    windows_myr=DEFAULT_WINDOWS_MYR,
    sfr=True,
    ssfr=True,
    uv=True,
    ionizing=True,
    predictions=True,
    derived=None,
    batch_size=256
)

Posterior post-processing of a fit.

Parameters:

Name	Type	Description	Default
model	SedModel -- the model the samples were drawn with; parameter names must match the result's		required
result	SamplingResult or path -- sampler output or the HDF5 file it was written to		required
n_samples	int, optional -- equal-weight draws (default 2000 for weighted results, all for uniform)		None
windows_myr	sequence of float, Myr -- averaging windows W for ``sfrW`` / ``ssfrW``		DEFAULT_WINDOWS_MYR
derived	dict[str, callable] -- ``name -> f(SpectrumSample) -> float or 1-D array``, evaluated per draw		None
batch_size	int -- draws per compiled batch through the forward model		256

Source code in ceridwen/postprocess.py

def __init__(self, model, result, *, n_samples: Optional[int] = None,
             seed: int = 0, windows_myr: Sequence[float] = DEFAULT_WINDOWS_MYR,
             sfr: bool = True, ssfr: bool = True, uv: bool = True,
             ionizing: bool = True, predictions: bool = True,
             derived: Optional[dict] = None, batch_size: int = 256):
    self.model = model
    self.csp = model.csp
    self.result = self._load_result(result)
    self.seed = int(seed)
    self.windows_myr = tuple(float(w) for w in windows_myr)
    if any(w <= 0 for w in self.windows_myr):
        raise ValueError("windows_myr must be positive")
    self.want = dict(sfr=bool(sfr), ssfr=bool(ssfr), uv=bool(uv),
                     ionizing=bool(ionizing), predictions=bool(predictions))
    if self.want["ssfr"] and not self.want["sfr"]:
        raise ValueError("ssfr=True needs sfr=True")
    self.derived = dict(derived or {})
    for k, f in self.derived.items():
        if not callable(f):
            raise TypeError(f"derived[{k!r}] is not callable")
        if k in ("sfh", "uv", "ionizing"):
            raise ValueError(f"derived name {k!r} collides with a built-in block")
    self.batch_size = int(batch_size)
    if self.batch_size < 1:
        raise ValueError("batch_size must be >= 1")
    self._samples_np = {k: np.asarray(v, dtype=float) for k, v in self.result.samples.items()}
    self._logl_np = np.asarray(self.result.log_likelihoods, dtype=float)
    self._check_names()
    self.log_weights = self._log_weights()
    self.n_samples = n_samples
    self.output: Optional[dict] = None

run

run()

Compute and return (and store in self.output) the nested dict of the module docstring.

Source code in ceridwen/postprocess.py

def run(self) -> dict:
    """Compute and return (and store in ``self.output``) the nested dict of the module docstring."""
    self._compile()
    idx = self._draw_indices()
    theta_batch = self._theta_batch(idx)
    raw = self._run_batches(theta_batch)
    out = self._derive(theta_batch, raw)
    out["theta"] = {k: np.asarray(v).reshape((v.shape[0], -1)).squeeze(axis=-1)
                    if np.asarray(v).shape[1:] == (1,) else np.asarray(v)
                    for k, v in theta_batch.items()}
    out["draw_index"] = idx
    out["log_likelihood"] = np.asarray(self.result.log_likelihoods, dtype=float)[idx]

    ll = np.asarray(self.result.log_likelihoods, dtype=float)
    ibest = int(np.nanargmax(ll))
    tb = self._theta_batch(np.array([ibest]))
    rb = self._run_batches(tb)
    best = self._derive(tb, rb)
    best["theta"] = {k: np.asarray(v)[0] for k, v in tb.items()}
    best["index"] = ibest
    best["log_likelihood"] = float(ll[ibest])
    best = _squeeze_leading(best)
    out["bestfit"] = best

    out["meta"] = {
        "n_samples": int(idx.size), "n_raw": int(self.n_raw), "seed": self.seed,
        "resampled": not self._uniform_weights(),
        "windows_myr": list(self.windows_myr), "param_names": list(self.model.param_names),
        "zred_fixed": float(self.model.zred), "cosmology": self.csp.cosmo.to_dict(),
        "sampler": getattr(self.result, "sampler_name", ""),
        "weights": self._weights_source,
        "log_evidence": float(getattr(self.result, "log_evidence", float("nan"))),
        "log_evidence_err": float(getattr(self.result, "log_evidence_err", float("nan"))),
        "observations": [(o.name, getattr(o, "_kind", "")) for o in self.model.observations],
        "sfh_interp": self.csp.sfh_interp, "sfh_per_bin": bool(getattr(self.csp, "sfh_per_bin", False)),
    }
    self.output = out
    return out

figures

figures(
    outdir, *, prefix="", title=None, truths=None, fmt="pdf"
)

Write the summary, corner and sampling-diagnostic figures (:mod:ceridwen.plotting) to outdir; returns their paths. truths ({name: value}) marks injected values in mock tests.

Source code in ceridwen/postprocess.py

def figures(self, outdir, *, prefix="", title=None, truths=None, fmt="pdf"):
    """Write the summary, corner and sampling-diagnostic figures
    (:mod:`ceridwen.plotting`) to ``outdir``; returns their paths.
    ``truths`` ({name: value}) marks injected values in mock tests."""
    from .plotting import make_figures
    if self.output is None:
        self.run()
    return make_figures(self.output, self.model, self.result, outdir,
                        prefix=prefix, title=title, truths=truths, fmt=fmt)

save

save(path)

Write self.output (running first if needed) as a flat .npz with '/'-joined keys.

Source code in ceridwen/postprocess.py

def save(self, path) -> Path:
    """Write ``self.output`` (running first if needed) as a flat ``.npz`` with '/'-joined keys."""
    if self.output is None:
        self.run()
    path = Path(path)
    if path.suffix != ".npz":
        path = path.with_suffix(".npz")
    flat = {}
    _flatten(self.output, "", flat)
    np.savez(path, **flat)
    return path

ceridwen.postprocess.SpectrumSample dataclass

SpectrumSample(
    wave_rest,
    full,
    intrinsic,
    dustfree,
    theta,
    zred,
    logmass,
    sfr,
    lookback_gyr,
    cosmo,
)

One posterior draw as handed to a derived function: rest-frame L_nu [L_sun/Hz] x 10**logmass on wave_rest [A]; theta after the model transforms.

index_of

index_of(wave_aa)

Index of the model pixel closest to wave_aa (rest frame).

Source code in ceridwen/postprocess.py

def index_of(self, wave_aa: float) -> int:
    """Index of the model pixel closest to ``wave_aa`` (rest frame)."""
    return int(np.argmin(np.abs(self.wave_rest - float(wave_aa))))

mean_lnu

mean_lnu(spectrum, lo_aa, hi_aa)

Wavelength-averaged L_nu of spectrum over [lo, hi] A (trapezoid).

Source code in ceridwen/postprocess.py

def mean_lnu(self, spectrum: np.ndarray, lo_aa: float, hi_aa: float) -> float:
    """Wavelength-averaged L_nu of ``spectrum`` over [lo, hi] A (trapezoid)."""
    m = (self.wave_rest >= lo_aa) & (self.wave_rest <= hi_aa)
    if m.sum() < 2:
        raise ValueError(f"fewer than two model pixels in [{lo_aa}, {hi_aa}] A")
    w = self.wave_rest[m]
    return float(_trapz(spectrum[m], w) / (w[-1] - w[0]))

luminosity

luminosity(spectrum, lo_aa, hi_aa)

int L_nu dnu over [lo, hi] A of spectrum [erg/s].

Source code in ceridwen/postprocess.py

def luminosity(self, spectrum: np.ndarray, lo_aa: float, hi_aa: float) -> float:
    """int L_nu dnu over [lo, hi] A of ``spectrum`` [erg/s]."""
    m = (self.wave_rest >= lo_aa) & (self.wave_rest <= hi_aa)
    if m.sum() < 2:
        raise ValueError(f"fewer than two model pixels in [{lo_aa}, {hi_aa}] A")
    nu = 2.99792458e18 / self.wave_rest[m]
    return float(-_trapz(spectrum[m], nu) * _LSUN_ERG_S)

ceridwen.postprocess.load_postprocess

load_postprocess(path)

Rebuild the nested dict written by :meth:PostProcess.save.

Source code in ceridwen/postprocess.py

def load_postprocess(path) -> dict:
    """Rebuild the nested dict written by :meth:`PostProcess.save`."""
    out: dict = {}
    with np.load(path, allow_pickle=False) as f:
        for key in f.files:
            parts = key.split("/")
            d = out
            for p in parts[:-1]:
                d = d.setdefault(p, {})
            v = f[key]
            if v.dtype.kind in ("U", "S"):
                try:
                    v = json.loads(str(v))
                except json.JSONDecodeError:
                    v = str(v)
            d[parts[-1]] = v
    return out

Figures

ceridwen.plotting.summary_figure

summary_figure(
    out,
    model,
    *,
    title=None,
    prior_draws=500,
    params=None,
    truths=None,
    savepath=None,
    figsize=(15, 11)
)

Summary page: SED with residuals, emission lines, SFH with posterior and prior bands, and 1-D marginals of the fitted parameters (median, 16-84 %, best fit). out is PostProcess.run()'s dict; model the fitted SedModel; truths ({name: value or array}) marks injected values in green.

Source code in ceridwen/plotting.py

def summary_figure(out, model, *, title=None, prior_draws=500, params=None, truths=None,
                   savepath=None, figsize=(15, 11)):
    """Summary page: SED with residuals, emission lines, SFH with posterior and prior
    bands, and 1-D marginals of the fitted parameters (median, 16-84 %, best fit).
    ``out`` is ``PostProcess.run()``'s dict; ``model`` the fitted ``SedModel``;
    ``truths`` ({name: value or array}) marks injected values in green."""
    import matplotlib.pyplot as plt
    from matplotlib.gridspec import GridSpec

    C = COLORS
    theta = out["theta"]
    cols = _flat_params(theta, params)
    best = _flat_point(out["bestfit"]["theta"], (params or list(theta.keys())))
    prior = None
    if prior_draws:
        pd = _prior_draws(model, int(prior_draws))
        if pd is not None:
            prior = {"theta": _flat_params(pd)}
            try:
                prior["sfr"], prior["T"] = _sfr_of(model, out, pd)
            except Exception:
                prior["sfr"] = None

    obs_by_kind = {"photometry": [], "spectrum": [], "lines": []}
    for o in model.observations:
        obs_by_kind.get(getattr(o, "_kind", ""), []).append(o)
    has_lines = bool(obs_by_kind["lines"])
    n_par = len(cols)
    n_marg_cols = 4
    n_marg_rows = int(np.ceil(n_par / n_marg_cols))

    fig = plt.figure(figsize=figsize)
    gs = GridSpec(2, 2, figure=fig, height_ratios=[1.0, 1.0], width_ratios=[1.35, 1.0],
                  left=0.06, right=0.98, top=0.90, bottom=0.07, hspace=0.32, wspace=0.22)
    gs_top = gs[0, 0].subgridspec(2, 1, height_ratios=[3, 1], hspace=0.05)
    ax_sed = fig.add_subplot(gs_top[0]); ax_chi = fig.add_subplot(gs_top[1], sharex=ax_sed)
    ax_lines = fig.add_subplot(gs[0, 1]) if has_lines else None
    ax_sfh = fig.add_subplot(gs[1, 0])
    gs_marg = gs[1, 1].subgridspec(n_marg_rows, n_marg_cols, hspace=0.9, wspace=0.35)

    chi2, ndata = 0.0, 0
    pred = out["prediction"]
    best_pred = out["bestfit"]["prediction"]
    truth_pt = _flat_point(truths, list(truths)) if truths else {}
    observed = "spectra_observed" in pred
    unit = _MAGGIE_TO_NJY if observed else 1.0

    wave_rest = np.asarray(pred["wave_rest"])
    if observed:
        z = np.asarray(pred["zred"])
        wobs = (1.0 + np.median(z)) * wave_rest
        spec = np.asarray(pred["spectra_observed"], dtype=float) * _CGS_FNU_TO_NJY
    else:
        wobs = wave_rest
        spec = np.asarray(pred["spectra_model"], dtype=float)
    lo, med, hi = _quantiles(spec)
    m = (med > 0) & np.isfinite(med)
    ax_sed.fill_between(wobs[m] / 1e4, lo[m], hi[m], color=C["band"], alpha=0.6, lw=0,
                        label="model (16$-$84%)")
    ax_sed.plot(wobs[m] / 1e4, med[m], color=C["posterior"], lw=0.9)
    unit_spec = _CGS_FNU_TO_NJY if observed else 1.0     # spectra are cgs f_nu, photometry maggies
    for o in obs_by_kind["spectrum"]:
        w = np.asarray(o.wavelength, dtype=float) / 1e4
        y = np.asarray(o.flux, dtype=float) * unit_spec
        s = np.asarray(o.uncertainty, dtype=float) * unit_spec
        mk = np.asarray(o.mask, dtype=bool)
        ax_sed.fill_between(w[mk], (y - s)[mk], (y + s)[mk], color=C["data"], alpha=0.15, lw=0)
        ax_sed.plot(w[mk], y[mk], color=C["data"], lw=0.6, label=f"observed spectrum ({o.name})")
        if o.name in pred["spectra"]:
            p = np.asarray(pred["spectra"][o.name], dtype=float) * unit_spec
            plo, pmed, phi = _quantiles(p)
            ax_sed.plot(w[mk], pmed[mk], color=C["bestfit"], lw=0.6, alpha=0.8, label="posterior spectrum")
            r = ((pmed - y) / s)[mk]
            ax_chi.plot(w[mk], r, color=C["posterior"], lw=0.5)
            pb = np.asarray(best_pred["spectra"][o.name], dtype=float) * unit_spec
            chi2 += float(np.sum(((pb - y) / s)[mk] ** 2)); ndata += int(mk.sum())
    for o in obs_by_kind["photometry"]:
        w = np.asarray(o.wavelength, dtype=float) / 1e4
        y = np.asarray(o.flux, dtype=float) * unit
        s = np.asarray(o.uncertainty, dtype=float) * unit
        mk = np.asarray(o.mask, dtype=bool)
        ul = np.asarray(o.upper_limit, dtype=bool) if getattr(o, "upper_limit", None) is not None else np.zeros_like(mk)
        det = mk & ~ul
        ax_sed.errorbar(w[det], y[det], yerr=s[det], fmt="o", color=C["data"], ms=5, capsize=2,
                        label="observed", zorder=5)
        if ul.any():
            ax_sed.errorbar(w[ul], y[ul], yerr=0.3 * y[ul], uplims=True, fmt="v", color=C["grid"],
                            ms=5, label="upper limit", zorder=5)
        if o.name in pred["photometry"]:
            p = np.asarray(pred["photometry"][o.name], dtype=float) * unit
            plo, pmed, phi = _quantiles(p)
            ax_sed.errorbar(w, pmed, yerr=[pmed - plo, phi - pmed], fmt="s", mfc="none",
                            color=C["bestfit"], ms=6, capsize=0, label="posterior photometry", zorder=6)
            r = (pmed - y) / s
            rlo, rhi = (plo - y) / s, (phi - y) / s
            ax_chi.errorbar(w[det], r[det], yerr=[(r - rlo)[det], (rhi - r)[det]], fmt="o",
                            color=C["bestfit"], ms=4, capsize=0)
            pb = np.asarray(best_pred["photometry"][o.name], dtype=float) * unit
            rb = (pb - y) / s
            chi2 += float(np.sum(np.where(ul, np.maximum(rb, 0.0), rb)[mk] ** 2)); ndata += int(mk.sum())
    ax_sed.set_xscale("log"); ax_sed.set_yscale("log")
    ax_sed.set_ylabel(r"$F_\nu$ [nJy]" if observed else r"$F_\nu$ [model units, no flux factor]")
    ax_sed.legend(fontsize=7, loc="lower right", frameon=False)
    ax_sed.tick_params(labelbottom=False)
    ax_chi.axhline(0, color=C["data"], lw=0.8)
    ax_chi.axhspan(-1, 1, color=C["prior"], alpha=0.6, lw=0)
    ax_chi.set_ylim(-4, 4); ax_chi.set_ylabel(r"$\chi$")
    ax_chi.set_xlabel(r"observed wavelength [$\mu$m]")
    for ax in (ax_sed, ax_chi):
        ax.grid(False)

    # emission lines: (model - observed) / sigma per line
    if has_lines:
        yy = 0
        for o in obs_by_kind["lines"]:
            names = list(o.line_names) if getattr(o, "line_names", None) else [f"{float(l):.0f}" for l in np.asarray(o.wavelength)]
            y = np.asarray(o.flux, dtype=float); s = np.asarray(o.uncertainty, dtype=float)
            mk = np.asarray(o.mask, dtype=bool)
            ul = np.asarray(o.upper_limit, dtype=bool) if getattr(o, "upper_limit", None) is not None else np.zeros_like(mk)
            p = np.asarray(pred["lines"][o.name], dtype=float) if o.name in pred["lines"] else None
            pb = np.asarray(best_pred["lines"][o.name], dtype=float) if o.name in best_pred["lines"] else None
            for i, nm in enumerate(names):
                if not mk[i]:
                    continue
                snr = y[i] / s[i]
                if p is not None:
                    lo, med, hi = _quantiles((p[:, i] - y[i]) / s[i])
                    ax_lines.errorbar(med, yy, xerr=[[med - lo], [hi - med]], fmt="o",
                                      color=C["bestfit"] if snr >= 3 else C["band"],
                                      mfc=C["bestfit"] if snr >= 3 else "white", ms=5, capsize=0)
                    if pb is not None:
                        rb = (pb[i] - y[i]) / s[i]
                        ax_lines.plot(rb, yy, marker="|", color=C["data"], ms=8, lw=0)
                        chi2 += float(max(rb, 0.0) ** 2 if ul[i] else rb ** 2); ndata += 1
                ax_lines.text(4.9, yy, f"{snr:.0f}", ha="right", va="center", fontsize=7, color=C["grid"])
                ax_lines.text(-5.2, yy, nm, ha="right", va="center", fontsize=8)
                yy += 1
        ax_lines.axvline(0, color=C["data"], lw=0.8)
        ax_lines.axvspan(-1, 1, color=C["prior"], alpha=0.6, lw=0)
        ax_lines.set_xlim(-5.4, 5.2); ax_lines.set_ylim(-0.7, yy - 0.3); ax_lines.invert_yaxis()
        ax_lines.set_yticks([]); ax_lines.set_xlabel(r"(model $-$ observed) / $\sigma$")
        ax_lines.text(5.05, -0.55, "S/N", ha="right", va="center", fontsize=7, color=C["grid"])
        ax_lines.plot([], [], "o", color=C["bestfit"], label=r"S/N $\geq$ 3")
        ax_lines.plot([], [], "o", color=C["band"], mfc="white", label="S/N < 3")
        ax_lines.plot([], [], "|", color=C["data"], label="best fit")
        ax_lines.legend(fontsize=7, loc="upper right", frameon=False, ncol=3, bbox_to_anchor=(1.0, 1.08))

    # SFH
    sfh = out["extras"]["sfh"]
    T = np.asarray(sfh["lookback_gyr"], dtype=float) * 1e3
    sfr = np.asarray(sfh["sfr"], dtype=float)
    per_bin = sfr.shape[1] == T.shape[1] - 1
    Tmed = np.median(T, axis=0)
    edges = Tmed if per_bin else np.concatenate([[Tmed[0]], 0.5 * (Tmed[1:] + Tmed[:-1]), [Tmed[-1]]])
    edges = np.maximum(edges, 1.0)
    lo, med, hi = _quantiles(sfr)
    lo2, _, hi2 = _quantiles(sfr, (0.025, 0.5, 0.975))
    if prior is not None and prior.get("sfr") is not None:
        plo, _, phi = _quantiles(prior["sfr"])
        ax_sfh.stairs(np.maximum(phi, 1e-30), edges, baseline=np.maximum(plo, 1e-30), fill=True,
                      color=C["prior"], alpha=0.8, lw=0, label=r"prior (16$-$84%) at posterior $M_\ast$")
    ax_sfh.stairs(np.maximum(hi2, 1e-30), edges, baseline=np.maximum(lo2, 1e-30), fill=True, color=C["band2"], lw=0)
    ax_sfh.stairs(np.maximum(hi, 1e-30), edges, baseline=np.maximum(lo, 1e-30), fill=True, color=C["band"],
                  alpha=0.9, lw=0, label=r"posterior (16$-$84%)")
    ax_sfh.stairs(np.maximum(med, 1e-30), edges, color=C["posterior"], lw=1.8, label="posterior median")
    bsfr = np.asarray(out["bestfit"]["extras"]["sfh"]["sfr"], dtype=float)
    ax_sfh.stairs(np.maximum(bsfr, 1e-30), edges, color=C["bestfit"], lw=1.0, label="best fit")
    for w in (10, 100):
        ax_sfh.axvline(w, color=C["grid"], lw=0.6, ls=":")
    ax_sfh.set_xscale("log"); ax_sfh.set_yscale("log")
    pos = med[med > 0]
    if pos.size:
        ax_sfh.set_ylim(max(pos.min() * 1e-2, 1e-6), max(hi.max() * 3, pos.max() * 3))
    ax_sfh.set_xlim(edges[0], edges[-1])
    ax_sfh.set_xlabel("lookback time [Myr]"); ax_sfh.set_ylabel(r"SFR [M$_\odot$ yr$^{-1}$]")
    if "sfr10" in sfh and "sfr100" in sfh:
        with np.errstate(divide="ignore", invalid="ignore"):
            r = np.log10(np.asarray(sfh["sfr10"]) / np.asarray(sfh["sfr100"]))
        r = r[np.isfinite(r)]
        if r.size:
            q = _quantiles(r)
            ax_sfh.text(0.03, 0.04, r"$\log\,\mathrm{SFR}_{10}/\mathrm{SFR}_{100}$ = " + _fmt_q(*q),
                        transform=ax_sfh.transAxes, fontsize=9)
    ax_sfh.legend(fontsize=7, loc="upper right", frameon=False)

    # marginals
    for i, (name, x) in enumerate(cols.items()):
        ax = fig.add_subplot(gs_marg[i // n_marg_cols, i % n_marg_cols])
        x = x[np.isfinite(x)]
        q = _quantiles(x)
        if prior is not None and name in prior["theta"]:
            plo, _, phi = _quantiles(prior["theta"][name])
            ax.axvspan(plo, phi, color=C["prior"], alpha=0.8, lw=0)
        ax.hist(x, bins=30, color=C["band"], histtype="stepfilled", alpha=0.9)
        ax.hist(x, bins=30, color=C["posterior"], histtype="step", lw=1.0)
        ax.axvline(q[1], color=C["posterior"], lw=1.2)
        if name in best:
            ax.axvline(best[name], color=C["bestfit"], lw=1.0)
        if name in truth_pt:
            ax.axvline(truth_pt[name], color=C["truth"], lw=1.2, ls="--")
        ax.set_yticks([]); ax.tick_params(axis="x", labelsize=7)
        ax.set_title(_label(name), fontsize=9, pad=3)
        ax.set_xlabel(_fmt_q(*q), fontsize=8, labelpad=1)
        for s_ in ("top", "right", "left"):
            ax.spines[s_].set_visible(False)

    # header
    meta = out["meta"]
    bits = []
    if "zred" in theta:
        q = _quantiles(theta["zred"]); bits.append(r"$z$ = " + _fmt_q(*q, nd=3))
    else:
        bits.append(rf"$z$ = {meta.get('zred_fixed', 0.0):.3f} (fixed)")
    if "logmass" in theta:
        bits.append(r"$\log M_\ast/\mathrm{M}_\odot$ = " + _fmt_q(*_quantiles(theta["logmass"])))
    if "sfr10" in sfh:
        with np.errstate(divide="ignore"):
            s10 = np.log10(np.asarray(sfh["sfr10"]))
        s10 = s10[np.isfinite(s10)]
        if s10.size:
            bits.append(r"$\log\,\mathrm{SFR}_{10}$ = " + _fmt_q(*_quantiles(s10)))
    nfree = sum(int(np.size(v)) for v in model.theta_init.values())
    if ndata:
        bits.append(rf"$\chi^2/\nu$ = {chi2 / max(ndata - nfree, 1):.2f} (best fit; $N_\mathrm{{data}}$ = {ndata})")
    if "log_evidence" in meta and np.isfinite(meta["log_evidence"]):
        bits.append(rf"$\ln Z$ = {meta['log_evidence']:.1f}")
    fig.text(0.06, 0.965, title or "CERIDWEN fit", fontsize=14, weight="bold", va="top")
    fig.text(0.06, 0.93, "    ".join(bits) + f"    sampler: {meta.get('sampler', '')}", fontsize=9.5, va="top")
    if savepath:
        fig.savefig(savepath, bbox_inches="tight")
    return fig

ceridwen.plotting.corner_figure

corner_figure(
    out,
    *,
    params=None,
    truths=None,
    savepath=None,
    bins=30,
    panel=1.6
)

Corner plot of all fitted parameters: 1-D marginals on the diagonal, 2-D histograms with 1 and 2 sigma contours below; posterior median (blue, dashed), maximum-likelihood sample (red) and, when given, truths (green) marked.

Source code in ceridwen/plotting.py

def corner_figure(out, *, params=None, truths=None, savepath=None, bins=30, panel=1.6):
    """Corner plot of all fitted parameters: 1-D marginals on the diagonal, 2-D
    histograms with 1 and 2 sigma contours below; posterior median (blue, dashed),
    maximum-likelihood sample (red) and, when given, ``truths`` (green) marked."""
    import matplotlib.pyplot as plt
    C = COLORS
    cols = _flat_params(out["theta"], params)
    best = _flat_point(out["bestfit"]["theta"], (params or list(out["theta"].keys())))
    truth_pt = _flat_point(truths, list(truths)) if truths else {}
    names = list(cols)
    K = len(names)
    X = np.column_stack([cols[n] for n in names])
    ok = np.all(np.isfinite(X), axis=1)
    X = X[ok]
    fig, axes = plt.subplots(K, K, figsize=(panel * K, panel * K))
    axes = np.atleast_2d(axes)
    lims = [(np.quantile(X[:, i], 0.001), np.quantile(X[:, i], 0.999)) for i in range(K)]
    lims = [(a - 0.05 * (b - a), b + 0.05 * (b - a)) if b > a else (a - 1, a + 1) for a, b in lims]
    for i in range(K):
        for j in range(K):
            ax = axes[i, j]
            if j > i:
                ax.set_visible(False); continue
            if i == j:
                x = X[:, i]
                q = _quantiles(x)
                ax.hist(x, bins=bins, range=lims[i], color=C["band"], histtype="stepfilled", alpha=0.9)
                ax.hist(x, bins=bins, range=lims[i], color=C["posterior"], histtype="step", lw=1.0)
                ax.axvline(q[1], color=C["posterior"], ls="--", lw=1.0)
                ax.axvline(q[0], color=C["posterior"], ls=":", lw=0.7); ax.axvline(q[2], color=C["posterior"], ls=":", lw=0.7)
                ax.axvline(best[names[i]], color=C["bestfit"], lw=1.2)
                if names[i] in truth_pt:
                    ax.axvline(truth_pt[names[i]], color=C["truth"], lw=1.2, ls="--")
                ax.set_title(_label(names[i]) + "\n" + _fmt_q(*q), fontsize=8)
                ax.set_yticks([]); ax.set_xlim(lims[i])
            else:
                x, y = X[:, j], X[:, i]
                H, xe, ye = np.histogram2d(x, y, bins=bins, range=[lims[j], lims[i]])
                Hs = _smooth(H.T)
                flat = np.sort(Hs.ravel())[::-1]
                csum = np.cumsum(flat) / flat.sum()
                levels = [flat[np.searchsorted(csum, f)] for f in (0.865, 0.393)]
                levels = sorted(set([max(l, 1e-12) for l in levels]))
                xc = 0.5 * (xe[1:] + xe[:-1]); yc = 0.5 * (ye[1:] + ye[:-1])
                ax.pcolormesh(xe, ye, np.ma.masked_where(Hs == 0, Hs), cmap="Blues", shading="flat", rasterized=True)
                if len(levels) >= 1:
                    ax.contour(xc, yc, Hs, levels=levels + [Hs.max() + 1], colors=C["posterior"], linewidths=0.8)
                ax.axvline(np.median(x), color=C["posterior"], ls="--", lw=0.7)
                ax.axhline(np.median(y), color=C["posterior"], ls="--", lw=0.7)
                ax.plot(best[names[j]], best[names[i]], "o", color=C["bestfit"], ms=4)
                if names[j] in truth_pt and names[i] in truth_pt:
                    ax.plot(truth_pt[names[j]], truth_pt[names[i]], "s", color=C["truth"], ms=4)
                ax.set_xlim(lims[j]); ax.set_ylim(lims[i])
            if i < K - 1:
                ax.tick_params(labelbottom=False)
            else:
                ax.set_xlabel(_label(names[j]), fontsize=8)
            if j > 0 or i == 0:
                ax.tick_params(labelleft=False)
            else:
                ax.set_ylabel(_label(names[i]), fontsize=8)
            ax.tick_params(labelsize=6)
    fig.text(0.62, 0.95, "dashed: posterior median (dotted: 16 / 84 %)\nred: maximum-likelihood sample"
             + ("\ngreen: injected truth" if truth_pt else ""), fontsize=9, color=C["data"], va="top")
    fig.subplots_adjust(left=0.07, right=0.98, bottom=0.07, top=0.95, hspace=0.08, wspace=0.08)
    if savepath:
        fig.savefig(savepath, bbox_inches="tight")
    return fig

ceridwen.plotting.diagnostic_figure

diagnostic_figure(
    result,
    out=None,
    *,
    params=None,
    savepath=None,
    max_points=4000
)

Sampling diagnostics. Nested sampling: every parameter's dead points in deletion order coloured by posterior weight, the log-likelihood run and the cumulative evidence. MCMC: per-chain traces with split-R-hat and ESS.

Source code in ceridwen/plotting.py

def diagnostic_figure(result, out=None, *, params=None, savepath=None, max_points=4000):
    """Sampling diagnostics.  Nested sampling: every parameter's dead points in
    deletion order coloured by posterior weight, the log-likelihood run and the
    cumulative evidence.  MCMC: per-chain traces with split-R-hat and ESS."""
    import matplotlib.pyplot as plt
    C = COLORS
    theta = {k: np.asarray(v) for k, v in result.samples.items()}
    if params:
        theta = {k: theta[k] for k in params}
    cols = _flat_params(theta)
    names = list(cols); K = len(names)
    raw = getattr(result, "raw", None) or {}
    n_chains = int(raw.get("num_chains", 0) or 0)
    is_mcmc = n_chains > 0 and getattr(result, "log_likelihoods_birth", None) is None
    ncol = 2
    nrow = int(np.ceil(K / ncol)) + 1
    fig, axes = plt.subplots(nrow, ncol, figsize=(12, 2.0 * nrow))
    axes = np.atleast_2d(axes)
    ll = np.asarray(result.log_likelihoods, dtype=float)
    n = ll.size
    if is_mcmc:
        S = n // n_chains
        for i, name in enumerate(names):
            ax = axes[i // ncol, i % ncol]
            ch = cols[name][: n_chains * S].reshape(n_chains, S)
            for c in range(n_chains):
                ax.plot(ch[c], color=C["chains"][c % len(C["chains"])], lw=0.4, alpha=0.8)
            rhat = _split_rhat(ch); ess = sum(_ess(ch[c]) for c in range(n_chains))
            ax.set_title(f"{_label(name)}    $\\hat R$ = {rhat:.3f}    ESS = {ess:.0f}", fontsize=8)
            ax.tick_params(labelsize=7)
        ax = axes[-1, 0]
        for c in range(n_chains):
            ax.plot(ll[c * S:(c + 1) * S], color=C["chains"][c % len(C["chains"])], lw=0.4, alpha=0.8)
        ax.set_title("log-likelihood per chain", fontsize=8)
        ax2 = axes[-1, 1]
        div = raw.get("total_divergences", None)
        ax2.axis("off")
        ax2.text(0.0, 0.8, f"{n_chains} chains x {S} draws" + (f", {div} divergences" if div is not None else ""),
                 fontsize=9, transform=ax2.transAxes)
        fig.suptitle("MCMC traces", fontsize=11)
    else:
        from .sampler.ns_weights import nested_log_weights
        births = getattr(result, "log_likelihoods_birth", None)
        lw = (nested_log_weights(ll, np.asarray(births)) if births is not None
              else np.asarray(result.log_weights, dtype=float))
        fin = np.isfinite(lw)
        w = np.zeros(n); w[fin] = np.exp(lw[fin] - lw[fin].max()); w /= w.sum()
        order = np.arange(n)
        step = max(1, n // max_points)
        sel = order[::step]
        for i, name in enumerate(names):
            ax = axes[i // ncol, i % ncol]
            ax.scatter(sel, cols[name][sel], c=np.clip(lw[sel] - lw[fin].max(), -12, 0), cmap="Blues",
                       s=3, vmin=-12, vmax=0, rasterized=True)
            qs = _quantiles(cols[name][fin], (0.16, 0.5, 0.84))
            ax.set_title(f"{_label(name)}    posterior " + _fmt_q(*qs), fontsize=8)
            ax.tick_params(labelsize=7)
        ax = axes[-1, 0]
        ax.plot(order, ll, color=C["posterior"], lw=0.6)
        ax.set_ylabel(r"$\ln L$", fontsize=8); ax.set_xlabel("dead point (deletion order)", fontsize=8)
        fl = ll[fin]
        if fl.size:
            ax.set_ylim(np.quantile(fl, 0.02), fl.max() + 0.05 * (fl.max() - np.quantile(fl, 0.02)))
        ax.set_title("log-likelihood run", fontsize=8)
        ax2 = axes[-1, 1]
        ax2.plot(order, np.cumsum(w), color=C["posterior"], lw=1.0)
        ax2.set_ylim(0, 1.02); ax2.set_xlabel("dead point (deletion order)", fontsize=8)
        ax2.set_ylabel("cumulative posterior weight", fontsize=8)
        ess = 1.0 / np.sum(w ** 2)
        lz = getattr(result, "log_evidence", np.nan); lze = getattr(result, "log_evidence_err", np.nan)
        ax2.set_title(rf"$\ln Z$ = {lz:.2f} $\pm$ {lze:.2f}    ESS = {ess:.0f} of {n}", fontsize=8)
        fig.suptitle("nested-sampling run: dead points coloured by posterior weight", fontsize=11)
    for k in range(K, (nrow - 1) * ncol):
        axes[k // ncol, k % ncol].axis("off")
    fig.tight_layout()
    if savepath:
        fig.savefig(savepath, bbox_inches="tight")
    return fig

ceridwen.plotting.make_figures

make_figures(
    out,
    model,
    result,
    outdir,
    *,
    prefix="",
    title=None,
    truths=None,
    fmt="pdf"
)

Write summary, corner and diagnostic figures to outdir; returns their paths.

Source code in ceridwen/plotting.py

def make_figures(out, model, result, outdir, *, prefix="", title=None, truths=None, fmt="pdf"):
    """Write summary, corner and diagnostic figures to ``outdir``; returns their paths."""
    import matplotlib.pyplot as plt
    outdir = Path(outdir); outdir.mkdir(parents=True, exist_ok=True)
    paths = {}
    for name, fn, args, kw in (("summary", summary_figure, (out, model), {"title": title, "truths": truths}),
                               ("corner", corner_figure, (out,), {"truths": truths}),
                               ("diagnostics", diagnostic_figure, (result, out), {})):
        p = outdir / f"{prefix}{name}.{fmt}"
        plt.close(fn(*args, savepath=p, **kw))
        paths[name] = p
    return paths

Dust, nebular, IGM, cosmology

ceridwen.dust.Dust

Dust(bin_edges=[(-jnp.inf, -1.97)], laws=['powerlaw'])

Modular dust attenuation model: one attenuation law per age bin, parameters read from a plain theta dict.

Parameters:

Name	Type	Description	Default
bin_edges	list of (lo, hi), log10(Gyr) -- age range of each bin		[(-inf, -1.97)]
laws	list of str -- attenuation-law name per bin (see ``describe_attenuation_laws``)		['powerlaw']

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)

Per-bin attenuation curves, shape (num_bins, len(wave)); wave in Angstroms.

Source code in ceridwen/dust/DustModel.py

def compute_attenuation(self, wave, fit_params):
    """Per-bin attenuation curves, shape ``(num_bins, len(wave))``; ``wave`` in Angstroms."""
    return jnp.stack([f(wave, fit_params) for f in self.law_funcs])

get_default_fit_params

get_default_fit_params()

Dict of default fit parameters (JAX scalars) for the active laws.

Source code in ceridwen/dust/DustModel.py

def get_default_fit_params(self):
    """Dict of default fit parameters (JAX scalars) for the active laws."""
    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

ceridwen.dust.DiffuseDust

DiffuseDust(law='kriek_conroy')

Bases: Dust

Single-bin dust model covering all ages, with diffuse_-prefixed parameter names.

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()

Dict of default diffuse-dust parameters (diffuse_* keys).

Source code in ceridwen/dust/DustModel.py

def get_default_params(self):
    """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)

Diffuse attenuation curve, shape (len(wave),).

Source code in ceridwen/dust/DustModel.py

def compute_attenuation(self, wave, fit_params):
    """Diffuse attenuation curve, shape ``(len(wave),)``."""
    return self.law_funcs[0](wave, fit_params)

ceridwen.dust.DustEmission

DustEmission

DustEmission(
    duste_model="DL07",
    dust_file=None,
    spec_lambda=None,
    **kwargs
)

Dust emission templates ('DL07' or 'THEMIS') interpolated onto spec_lambda [Angstrom].

Parameters:

Name	Type	Description	Default
dust_file	str -- data root containing dust/dustem/		None
kwargs	duste_qpah, duste_umin, duste_gamma defaults		{}

Source code in ceridwen/dust/DustEmission.py

def __init__(self, duste_model="DL07",
             dust_file=None, spec_lambda=None, **kwargs):
    """Dust emission templates ('DL07' or 'THEMIS') interpolated onto spec_lambda [Angstrom].

    Parameters
    ----------
    dust_file : str -- data root containing dust/dustem/
    kwargs : duste_qpah, duste_umin, duste_gamma defaults
    """
    self.duste_model = duste_model

    self.duste_qpah = None
    self.duste_umin = None
    self.duste_gamma = None

    self.qpaharr = None
    self.uminarr = None

    self.dustem2_dustem = None

    self.dust_file = None
    self.spec_lambda = None

    self.dwargs = 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)

    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":
        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'.")

    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

    self.load_dust_emission(dust_file, spec_lambda)

    nu = CLIGHT_AA_S / jnp.asarray(spec_lambda)
    dnu = jnp.diff(nu)
    self._trap_w = jnp.concatenate([
        jnp.array([0.5 * dnu[0]]),
        0.5 * (dnu[:-1] + dnu[1:]),
        jnp.array([0.5 * dnu[-1]]),
    ])

get_default_params

get_default_params()

Return dict of default dust-emission fit parameters.

Source code in ceridwen/dust/DustEmission.py

def get_default_params(self):
    """Return dict of default dust-emission fit parameters."""
    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
)

Set the default dust-emission parameters.

Source code in ceridwen/dust/DustEmission.py

def update_dust_params(self, duste_qpah = 3.5, duste_umin = 1.0, duste_gamma = 0.01):
    """Set the default dust-emission parameters."""
    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,
)

Return (spec_attn + dust emission, dust mass, dust emission); diffuse_curve is exp(-tau_diffuse).

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):
    """Return (spec_attn + dust emission, dust mass, dust emission); diffuse_curve is exp(-tau_diffuse)."""
    tiny = 1e-70
    w = self._trap_w
    dc = diffuse_curve.ravel()
    f_attn = spec_attn.ravel()
    f_free = spec_dustfree.ravel()

    lbold = jnp.dot(f_attn, w)
    lboln = jnp.dot(f_free, w)

    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)

    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)

    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, 2*n_umin): Umin/Umax pairs interleaved

    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])

    mduste = jnp.maximum((1 - gamma) * dumin + gamma * dumax, tiny).ravel()
    norm   = jnp.dot(mduste, w)
    mduste_norm = mduste / norm

    labs0  = lboln - lbold
    duste0 = jnp.maximum(mduste_norm * labs0, tiny)

    duste0_atten = duste0 * dc
    absorbed_1 = jnp.dot(duste0 * (1.0 - dc), w)
    duste1 = jnp.maximum(mduste_norm * absorbed_1, tiny)

    duste1_atten = duste1 * dc

    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=0.0,
    res_floor_factor=2.0,
    nebular_smooth_init=None,
)

Nebular continuum and lines from the CLOUDY cubes under <sps_home>/nebular.

Parameters:

Name	Type	Description	Default
cloudy_dust	bool -- ``ZAU_WD`` (True) or ``ZAU_ND`` (False) grids.		required
csp_lambda	(nspec,) -- model wavelength grid [A].		required
ssp_flux	(n_z, n_age, n_wave) -- SSP L_nu [L_sun/Hz]; gives ``log_qq`` (n_z, n_age).		None
ssp_ages_lgyr	(n_age,) -- log10(age/yr) of the SSPs; ages inside both cubes are "young".		None
isoc_type	str -- ZAU file suffix.		'mist'
nebnz	int -- cube dimensions.		11
nebnage	int -- cube dimensions.		11
nebnip	int -- cube dimensions.		11
smooth_velocity	bool -- ``sigma_smooth`` in km/s (True) or A.		True
sigma_smooth	float -- intrinsic width of the painted lines (0: pixel floor only).		0.0
res_floor_factor	float -- minimum painted width in local pixels.		2.0
Attributes			required
are			required

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=0.0,
             res_floor_factor=2.0,
             nebular_smooth_init=None):
    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

    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]
        n_young = int(np.asarray(young).sum())
        oldest_young_yr = (10.0 ** float(np.asarray(ages)[
            np.asarray(young)].max()) if n_young else 0.0)
        if n_young == 0 or oldest_young_yr > 3.2e8:
            raise RuntimeError(
                "Nebular young-SSP mask is inconsistent with the CLOUDY "
                f"grid: {n_young} SSPs flagged young, oldest "
                f"{oldest_young_yr/1e6:.1f} Myr (grid max_age "
                f"{max_age:.3f} log10 yr). Check the cube age-axis "
                "units.")
    else:
        self.young_mask = None
        self.young_idx  = None

    self.emline_index_consistent = None
    self.nebem_line_names = None
    try:
        _info_path = str(Path(sps_home) / "data" / "emlines_info.dat")
        _info_wave = []
        _info_name = []
        with open(_info_path) as _f:
            for _row in _f:
                _parts = _row.split(",")
                if len(_parts) >= 2:
                    _info_wave.append(float(_parts[0]))
                    _info_name.append(_parts[1].strip())
        _info_wave = np.asarray(_info_wave)
        _pos = np.asarray(self.nebem_line_pos)
        # names per cube row, matched by wavelength (1 A) like CSPBasis matches Lines
        if _info_wave.size:
            _j = np.argmin(np.abs(_pos[:, None] - _info_wave[None, :]), axis=1)
            self.nebem_line_names = [
                _info_name[j] if abs(_pos[i] - _info_wave[j]) <= 1.0 else None
                for i, j in enumerate(_j)]
        if _info_wave.size != _pos.size:
            self.emline_index_consistent = False
            warnings.warn(
                f"emlines_info.dat lists {_info_wave.size} lines but the "
                f"nebular cube {getattr(self, 'line_file', '?')} has "
                f"{_pos.size} -- the two files are from different FSPS "
                "vintages. Raw line_ind indices into this cube are "
                "UNRELIABLE; ceridwen's own predictions match lines by "
                "wavelength and are unaffected.", stacklevel=2)
        else:
            _n = min(_info_wave.size, _pos.size)
            _bad = int(np.sum(np.abs(_info_wave[:_n] - _pos[:_n]) > 1.0))
            self.emline_index_consistent = (_bad == 0)
            if _bad:
                warnings.warn(
                    f"{_bad}/{_n} rows of emlines_info.dat disagree with "
                    "the nebular cube wavelengths by >1 A -- mixed FSPS "
                    "vintages in $SPS_HOME. Raw line_ind indices into "
                    "this cube are UNRELIABLE; ceridwen's own predictions "
                    "match lines by wavelength and are unaffected.",
                    stacklevel=2)
    except OSError:
        pass                              # no emlines_info.dat: nothing to check

compute_log_qq

compute_log_qq(ssp_flux)

log10 Q [photons/s] for every SSP: (L_sun/h) * int_{lambda<912} L_nu / lambda dlambda (float64).

Source code in ceridwen/neb/NebularGridModel.py

def compute_log_qq(self, ssp_flux):
    """log10 Q [photons/s] for every SSP: (L_sun/h) * int_{lambda<912} L_nu / lambda dlambda (float64)."""
    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))

line_profiles

line_profiles(sigma_kms=0.0)

(nspec, nemline) NumPy profiles like gaussnebarr with width sqrt(floor^2 + (lambda sigma_kms / c)^2): a line painted at the floor and then broadened.

Source code in ceridwen/neb/NebularGridModel.py

def line_profiles(self, sigma_kms=0.0):
    """(nspec, nemline) NumPy profiles like ``gaussnebarr`` with width
    sqrt(floor^2 + (lambda sigma_kms / c)^2): a line painted at the floor and then broadened.
    """
    pos = np.asarray(self.nebem_line_pos, dtype=np.float64)
    lam = np.asarray(self.csp_lambda, dtype=np.float64)
    floor = np.asarray(self.neb_res_min, dtype=np.float64) * self.res_floor_factor
    if self.smooth_velocity:
        base = pos * self.sigma_smooth / CLIGHT_AA_S * 1.0e13
    else:
        base = np.full_like(pos, self.sigma_smooth)
    dl0 = np.maximum(base, floor)
    dl = np.sqrt(dl0 ** 2 + (pos * float(sigma_kms) / CLIGHT_AA_S * 1.0e13) ** 2)
    prof = np.exp(-0.5 * ((lam[:, None] - pos[None, :]) / dl[None, :]) ** 2)
    return prof / (SQRT_2PI * dl[None, :]) * (pos[None, :] ** 2 / CLIGHT_AA_S)

evaluate

evaluate(logZ, logU, logage, logQ)

(cont (nspec,), lines (nspec,)) [L_sun/Hz] at one (logZ, logU, logage, logQ).

Source code in ceridwen/neb/NebularGridModel.py

def evaluate(self, logZ, logU, logage, logQ):
    """``(cont (nspec,), lines (nspec,))`` [L_sun/Hz] at one (logZ, logU, logage, logQ)."""
    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)

    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,
)

(n_z, n_young, nspec) nebular spectra at one (logZ_gas, logU) for the young SSP ages, or (cont, lines) in that layout with return_components. The metallicity dependence enters only through logqq_young; the per-age reference ref keeps 10**(...) in float32 range.

Source code in ceridwen/neb/NebularGridModel.py

def evaluate_batch(self, logZ_gas, logU, ssp_ages_young, logqq_young,
                    return_components=False):
    """(n_z, n_young, nspec) nebular spectra at one (logZ_gas, logU) for the young SSP ages,
    or ``(cont, lines)`` in that layout with ``return_components``.  The metallicity dependence
    enters only through ``logqq_young``; the per-age reference ``ref`` keeps 10**(...) in float32 range.
    """
    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)

    ref   = jnp.max(logqq_young, axis=0)                          # (n_young,)
    ref   = jnp.where(jnp.isfinite(ref), ref, 0.0)
    scale = jnp.power(10.0, logqq_young - ref[None, :])            # (n_z, n_young)

    cont_base = jnp.power(10.0, log_cont + ref[None, :])           # (nspec , n_young)
    line_base = jnp.einsum('wl,ly->wy', self.gaussnebarr,
                           jnp.power(10.0, log_line + ref[None, :]))
    cont_flux = cont_base[None, :, :] * scale[:, None, :]          # (n_z, nspec, n_young)
    line_spec = line_base[None, :, :] * scale[:, None, :]          # (n_z, nspec, n_young)
    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)

evaluate_batch_factored

evaluate_batch_factored(
    logZ_gas,
    logU,
    ssp_ages_young,
    logqq_young,
    include_lines=True,
)

(base (n_young, n_wave), scale (n_z, n_young)) with neb[z, y, w] == scale[z, y] * base[y, w] (same arithmetic as evaluate_batch, not expanded).

Source code in ceridwen/neb/NebularGridModel.py

def evaluate_batch_factored(self, logZ_gas, logU, ssp_ages_young,
                            logqq_young, include_lines=True):
    """``(base (n_young, n_wave), scale (n_z, n_young))`` with neb[z, y, w] == scale[z, y] * base[y, w]
    (same arithmetic as ``evaluate_batch``, not expanded).
    """
    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)
    ref   = jnp.max(logqq_young, axis=0)
    ref   = jnp.where(jnp.isfinite(ref), ref, 0.0)
    scale = jnp.power(10.0, logqq_young - ref[None, :])            # (n_z, n_young)

    base = jnp.power(10.0, log_cont + ref[None, :])                # (nspec, n_young)
    if include_lines:
        log_line = _interp_cube(self.nebem_line, self.nebem_line_logz,
                                self.nebem_line_age, self.nebem_line_logu)
        base = base + jnp.einsum(
            'wl,ly->wy', self.gaussnebarr,
            jnp.power(10.0, log_line + ref[None, :]))
    return base.T, scale                                           # (n_young, n_wave), (n_z, n_young)

evaluate_batch_line_lum

evaluate_batch_line_lum(
    logZ_gas, logU, ssp_ages_young, logqq_young
)

Line luminosities (n_z, n_young, nlines) [L_sun] at one (logZ_gas, logU), no painting.

Source code in ceridwen/neb/NebularGridModel.py

def evaluate_batch_line_lum(self, logZ_gas, logU, ssp_ages_young,
                            logqq_young):
    """Line luminosities (n_z, n_young, nlines) [L_sun] at one (logZ_gas, logU), no painting."""
    logZ_gas = jnp.squeeze(logZ_gas)
    logU     = jnp.squeeze(logU)
    cube      = self.nebem_line
    logz_grid = self.nebem_line_logz
    age_grid  = self.nebem_line_age
    logu_grid = self.nebem_line_logu

    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])          # (nlines, nage)

    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,
    )
    log_line = ((1.0 - da)[None, :] * zu[..., a1]
                +       da[None, :] * zu[..., a1 + 1])   # (nlines, n_young)

    line_lum = jnp.power(10.0, log_line[None, :, :]
                         + logqq_young[:, None, :])
    return line_lum.transpose(0, 2, 1)

get_default_params

get_default_params()

Default nebular parameters (gas_logz = 0, gas_logu = -2).

Source code in ceridwen/neb/NebularGridModel.py

def get_default_params(self):
    """Default nebular parameters (gas_logz = 0, gas_logu = -2)."""
    return {'gas_logz': jnp.asarray(0.0),
            'gas_logu': jnp.asarray(-2.0)}

ceridwen.igm

Intergalactic-medium absorption models: transmission curves exp(-tau * factor) on the rest-frame wavelength grid at a given source redshift.

IGMModel

Bases: ABC

Abstract IGM attenuation model; subclasses implement tau(lam_rest, zred).

tau abstractmethod

tau(lam_rest, zred)

Optical depth on the rest-frame wavelength grid [Å] at source redshift zred.

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``."""

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: 17 Lyman-series lines, Ly-α metal blanketing and Lyman-continuum absorption; tau is capped at its short-wavelength peak and is zero at zred = 0.

make_igm_model

make_igm_model(name_or_model)

Return an :class:IGMModel from a registry name, an instance (as-is), or None (NoIGM).

Source code in ceridwen/igm.py

def make_igm_model(name_or_model):
    """Return an :class:`IGMModel` from a registry name, an instance (as-is), or ``None`` (NoIGM)."""
    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

Flat LambdaCDM redshift helpers: :class:Cosmology presets, D_L(z), age(z) and the maggies flux factor, JAX-native (differentiable in z) with an optional scalar-only astropy backend.

Cosmology dataclass

Cosmology(
    H0=67.66,
    Om0=0.30966,
    Tcmb0=2.7255,
    Neff=3.046,
    m_nu_ev_sum=0.06,
    name=None,
)

Flat LambdaCDM cosmology with photons and effective neutrinos; the massive neutrino is folded into Om0 as cold matter. Defaults to Planck 2018.

Parameters:

Name	Type	Description	Default
H0	(float, km / s / Mpc)		67.66
Om0	float -- cold matter today, EXCLUDING massive neutrinos.		0.30966
Tcmb0	(float, K)		2.7255
Neff	float -- effective number of neutrino species.		3.046
m_nu_ev_sum	float, eV -- sum of neutrino masses.		0.06
name	str -- label only; ignored in equality.		None

Ogamma0 property

Ogamma0

Photon density today.

Onu0_massive_as_matter property

Onu0_massive_as_matter

Massive-neutrino density today, treated as cold matter.

Onu0_relativistic property

Onu0_relativistic

Relativistic neutrino density today (Neff - 1 species).

Om0_eff property

Om0_eff

Matter density today including massive neutrinos.

Or0 property

Or0

Radiation + relativistic-neutrino density today.

Ode0 property

Ode0

Dark-energy density today (flat: 1 - Om0_eff - Or0).

is_planck18 property

is_planck18

True when the parameters equal the Planck-2018 preset (name ignored).

wmap9 classmethod

wmap9()

WMAP9 preset.

Source code in ceridwen/cosmology.py

@classmethod
def wmap9(cls) -> "Cosmology":
    """WMAP9 preset."""
    return cls(**_PRESETS["WMAP9"], name="WMAP9")

flat classmethod

flat(
    H0,
    Om0,
    *,
    Tcmb0=2.7255,
    Neff=3.046,
    m_nu_ev_sum=0.06,
    name=None
)

Flat LCDM from H0 and Om0; radiation and neutrinos default to Planck 2018.

Source code in ceridwen/cosmology.py

@classmethod
def flat(cls, H0: float, Om0: float, *, Tcmb0: float = 2.7255,
         Neff: float = 3.046, m_nu_ev_sum: float = 0.06,
         name: Optional[str] = None) -> "Cosmology":
    """Flat LCDM from H0 and Om0; radiation and neutrinos default to Planck 2018."""
    return cls(H0=H0, Om0=Om0, Tcmb0=Tcmb0, Neff=Neff,
               m_nu_ev_sum=m_nu_ev_sum, name=name)

from_name classmethod

from_name(name)

Preset by name, case-insensitive (see available_cosmologies()).

Source code in ceridwen/cosmology.py

@classmethod
def from_name(cls, name: str) -> "Cosmology":
    """Preset by name, case-insensitive (see ``available_cosmologies()``)."""
    key = {k.lower(): k for k in _PRESETS}.get(str(name).strip().lower())
    if key is None:
        raise KeyError(
            f"unknown cosmology preset {name!r}; known: {', '.join(_PRESETS)}")
    return cls(**_PRESETS[key], name=key)

from_dict classmethod

from_dict(d)

Inverse of :meth:to_dict; also accepts cosmo_*-prefixed HDF5 attrs.

Source code in ceridwen/cosmology.py

@classmethod
def from_dict(cls, d) -> "Cosmology":
    """Inverse of :meth:`to_dict`; also accepts ``cosmo_*``-prefixed HDF5 attrs."""
    d = dict(d)
    if "H0" not in d and "cosmo_H0" in d:
        d = {k[len("cosmo_"):]: v for k, v in d.items() if k.startswith("cosmo_")}
    missing = [k for k in _FIELDS if k not in d]
    if missing:
        raise KeyError(f"Cosmology.from_dict: missing {missing} (have {sorted(d)})")
    name = d.get("name", None)
    if isinstance(name, bytes):
        name = name.decode()
    if name is not None:
        name = str(name).strip() or None
    return cls(**{k: float(d[k]) for k in _FIELDS}, name=name)

to_dict

to_dict()

Plain floats plus name ("" when unnamed).

Source code in ceridwen/cosmology.py

def to_dict(self) -> dict:
    """Plain floats plus ``name`` ("" when unnamed)."""
    d = {k: float(getattr(self, k)) for k in _FIELDS}
    d["name"] = self.name or ""
    return d

age

age(z)

Age of the Universe at z [Gyr].

Source code in ceridwen/cosmology.py

def age(self, z) -> Array:
    """Age of the Universe at ``z`` [Gyr]."""
    return age_gyr(z, self)

luminosity_distance

luminosity_distance(z)

D_L(z) [Mpc].

Source code in ceridwen/cosmology.py

def luminosity_distance(self, z) -> Array:
    """D_L(z) [Mpc]."""
    return luminosity_distance_mpc(z, self)

from_astropy classmethod

from_astropy(cosmo)

Build from a flat astropy cosmology; a non-flat input raises.

Source code in ceridwen/cosmology.py

@classmethod
def from_astropy(cls, cosmo) -> "Cosmology":
    """Build from a flat astropy cosmology; a non-flat input raises."""
    Ok0 = float(getattr(cosmo, "Ok0", 0.0))
    if abs(Ok0) > 1e-8:
        raise ValueError(
            f"ceridwen's cosmology integrator assumes flatness, but the "
            f"supplied cosmology has Ok0 = {Ok0:.3e}."
        )
    m_nu = getattr(cosmo, "m_nu", None)
    if m_nu is None:
        m_nu_sum = 0.0
    else:
        try:
            m_nu_sum = float(m_nu.to("eV").value.sum())
        except AttributeError:
            m_nu_sum = float(m_nu.to("eV").value)
    return cls(
        H0=float(cosmo.H0.to("km/(s Mpc)").value),
        Om0=float(cosmo.Om0),
        Tcmb0=float(cosmo.Tcmb0.to("K").value),
        Neff=float(cosmo.Neff),
        m_nu_ev_sum=m_nu_sum,
        name=(str(cosmo.name) if getattr(cosmo, "name", None) else None),
    )

to_astropy

to_astropy()

Return the matching astropy.cosmology.FlatLambdaCDM (requires astropy).

Source code in ceridwen/cosmology.py

def to_astropy(self):
    """Return the matching ``astropy.cosmology.FlatLambdaCDM`` (requires astropy)."""
    import math

    import astropy.units as u
    from astropy.cosmology import FlatLambdaCDM

    kwargs = dict(
        H0=self.H0 * u.km / u.s / u.Mpc,
        Om0=self.Om0,
        Tcmb0=self.Tcmb0 * u.K,
        Neff=self.Neff,
    )
    # astropy needs len(m_nu) == floor(Neff); total mass on the last species
    n_species = math.floor(self.Neff)
    if n_species > 0 and self.Tcmb0 > 0.0:
        masses = [0.0] * n_species
        masses[-1] = self.m_nu_ev_sum
        kwargs["m_nu"] = u.Quantity(masses, u.eV)
    return FlatLambdaCDM(**kwargs)

available_cosmologies

available_cosmologies()

Names accepted by :meth:Cosmology.from_name.

Source code in ceridwen/cosmology.py

def available_cosmologies() -> tuple:
    """Names accepted by :meth:`Cosmology.from_name`."""
    return tuple(_PRESETS)

resolve_cosmology

resolve_cosmology(cosmo=None)

Return cosmo, or Planck 2018 when None.

Source code in ceridwen/cosmology.py

def resolve_cosmology(cosmo: "Cosmology | None" = None) -> Cosmology:
    """Return ``cosmo``, or Planck 2018 when ``None``."""
    return DEFAULT_COSMO if cosmo is None else cosmo

E_of_z

E_of_z(z, cosmo=None)

Dimensionless expansion rate E(z) = H(z)/H0; z is clamped to >= 0 so transient negative proposals cannot produce NaN.

Source code in ceridwen/cosmology.py

def E_of_z(z: Array, cosmo: Cosmology | None = None) -> Array:
    r"""Dimensionless expansion rate E(z) = H(z)/H0; ``z`` is clamped to >= 0
    so transient negative proposals cannot produce NaN."""
    cosmo = resolve_cosmology(cosmo)
    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=None, n_nodes=128)

Line-of-sight comoving distance D_C(z) [Mpc].

Source code in ceridwen/cosmology.py

def comoving_distance_mpc(z: Array, cosmo: Cosmology | None = None,
                          n_nodes: int = 128) -> Array:
    r"""Line-of-sight comoving distance D_C(z) [Mpc]."""
    cosmo = resolve_cosmology(cosmo)
    return cosmo.hubble_distance_mpc * _integrate_dz_over_E(z, cosmo, n_nodes)

age_gyr

age_gyr(z, cosmo=None, n_nodes=257)

Age of the Universe at z [Gyr], JAX-native and differentiable in z.

Source code in ceridwen/cosmology.py

def age_gyr(z: Array, cosmo: Cosmology | None = None,
            n_nodes: int = 257) -> Array:
    r"""Age of the Universe at ``z`` [Gyr], JAX-native and differentiable in z."""
    cosmo = resolve_cosmology(cosmo)
    if n_nodes % 2 == 0:
        n_nodes += 1
    z = jnp.asarray(z, dtype=float)
    a = 1.0 / (1.0 + jnp.maximum(z, 0.0))
    u = jnp.linspace(0.0, 1.0, n_nodes)
    x = a[..., None] * u                           # (..., n_nodes): a' from 0..a
    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=None, n_nodes=128, backend="native"
)

Luminosity distance D_L(z) [Mpc].

Parameters:

Name	Type	Description	Default
backend	('native', 'astropy')	'astropy' is scalar-z only.	'native'

Source code in ceridwen/cosmology.py

def luminosity_distance_mpc(z, cosmo: Cosmology | None = None,
                            n_nodes: int = 128,
                            backend: str = "native") -> Array:
    r"""Luminosity distance D_L(z) [Mpc].

    Parameters
    ----------
    backend : {'native', 'astropy'} -- 'native' is JAX and differentiable;
        'astropy' is scalar-z only.
    """
    cosmo = resolve_cosmology(cosmo)
    if backend == "astropy":
        return _astropy_luminosity_distance_mpc(z, cosmo)
    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=None, n_nodes=128)

(1+z) / (4 pi D_L^2) with D_L in cm: converts rest-frame L_nu [erg/s/Hz] to observed F_nu [erg/s/cm^2/Hz].

Source code in ceridwen/cosmology.py

def flux_factor(z: Array, cosmo: Cosmology | None = None,
                n_nodes: int = 128) -> Array:
    r"""(1+z) / (4 pi D_L^2) with D_L in cm: converts rest-frame L_nu
    [erg/s/Hz] to observed F_nu [erg/s/cm^2/Hz]."""
    cosmo = resolve_cosmology(cosmo)
    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_cgs

flux_factor_cgs(
    z,
    cosmo=None,
    n_nodes=128,
    backend="native",
    lumdist_mpc=None,
)

Factor turning a CSP spectrum [L_sun/Hz/M_sun] into observed-frame F_nu [erg/s/cm^2/Hz] (cgs): (1+z) (10 pc / D_L)^2 times the 10 pc unit constant. z <= 0 is pinned to D_L = 10 pc. The AB-maggies photometry is obtained downstream by projecting this cgs spectrum through the filters (dividing by the 3631 Jy AB zero point); this factor itself is NOT maggies.

Parameters:

Name	Type	Description	Default
backend	('native', 'astropy')		'native'
lumdist_mpc	float or Array, Mpc -- explicit D_L replacing D_L(z); the	(1+z) term still comes from z and the z <= 0 pin does not apply.	None

Source code in ceridwen/cosmology.py

def flux_factor_cgs(z, cosmo: Cosmology | None = None,
                    n_nodes: int = 128,
                    backend: str = "native",
                    lumdist_mpc=None) -> Array:
    r"""Factor turning a CSP spectrum [L_sun/Hz/M_sun] into observed-frame
    F_nu [erg/s/cm^2/Hz] (cgs): (1+z) (10 pc / D_L)^2 times the 10 pc unit
    constant. z <= 0 is pinned to D_L = 10 pc.  The AB-maggies photometry is
    obtained downstream by projecting this cgs spectrum through the filters
    (dividing by the 3631 Jy AB zero point); this factor itself is NOT maggies.

    Parameters
    ----------
    backend : {'native', 'astropy'} -- see :func:`luminosity_distance_mpc`.
    lumdist_mpc : float or Array, Mpc -- explicit D_L replacing D_L(z); the
        (1+z) term still comes from z and the z <= 0 pin does not apply.
    """
    cosmo = resolve_cosmology(cosmo)
    if lumdist_mpc is not None:
        dL_pc = 1e6 * jnp.asarray(lumdist_mpc, dtype=float)
        return ((1.0 + z) * (_D_FID_10PC / dL_pc) ** 2
                * _LSUN_HZ_TO_FNU_CGS_AT_10PC)
    dL_pc = 1e6 * luminosity_distance_mpc(z, cosmo, n_nodes, backend=backend)
    positive_z = z > 0
    safe_dL = jnp.where(positive_z, dL_pc, float(_D_FID_10PC))
    ff_distance = jnp.where(
        positive_z,
        (1.0 + z) * (_D_FID_10PC / safe_dL) ** 2,
        1.0,
    )
    return ff_distance * _LSUN_HZ_TO_FNU_CGS_AT_10PC

have_astropy

have_astropy()

True if astropy is importable.

Source code in ceridwen/cosmology.py

def have_astropy() -> bool:
    """True if astropy is importable."""
    try:
        import astropy.cosmology  # noqa: F401
        return True
    except Exception:
        return False

Spectrophotometric calibration

ceridwen.csp.spectrum_calibration.spectrum_calibration_factor

spectrum_calibration_factor(obs, theta, dtype=None)

Factor spectrum_scaling * (1 + spectrum_calib . P(x)) multiplying the MODEL spectrum: a scalar (level only), an (n_pix,) vector (shape term), or None if neither key is in theta.

Source code in ceridwen/csp/spectrum_calibration.py

def spectrum_calibration_factor(obs, theta, dtype=None):
    """Factor ``spectrum_scaling * (1 + spectrum_calib . P(x))`` multiplying the MODEL spectrum:
    a scalar (level only), an ``(n_pix,)`` vector (shape term), or ``None`` if neither key is in theta.
    """
    has_level = "spectrum_scaling" in theta
    has_shape = "spectrum_calib" in theta
    if not has_level and not has_shape:
        return None
    factor = None
    if has_level:
        factor = jnp.ravel(theta["spectrum_scaling"])[0]
    if has_shape:
        coeff = jnp.ravel(theta["spectrum_calib"])
        order = int(coeff.shape[0])
        if order < 1:
            raise ValueError(
                "spectrum_calib must have at least one coefficient "
                "(shape (order,), order >= 1); the level term is "
                "spectrum_scaling."
            )
        design = _design(obs, order)
        shape_term = 1.0 + design @ coeff
        factor = shape_term if factor is None else factor * shape_term
    if dtype is not None:
        factor = factor.astype(dtype)
    return factor

ceridwen.csp.spectrum_calibration.legendre_design_matrix

legendre_design_matrix(wavelength, order)

(n_pix, order) float64 matrix of Legendre P_1..P_order on the pixel wavelengths mapped affinely onto [-1, 1] (range over ALL pixels, masked or not; no P_0 term).

Source code in ceridwen/csp/spectrum_calibration.py

def legendre_design_matrix(wavelength, order: int) -> np.ndarray:
    """``(n_pix, order)`` float64 matrix of Legendre P_1..P_order on the pixel wavelengths
    mapped affinely onto [-1, 1] (range over ALL pixels, masked or not; no P_0 term)."""
    lam = np.asarray(wavelength, dtype=np.float64)
    if lam.ndim != 1 or lam.size < 2:
        raise ValueError(
            "spectrum_calib needs a 1-D observed wavelength array with at "
            f"least 2 pixels; got shape {lam.shape}."
        )
    fin = np.isfinite(lam)
    lo, hi = float(lam[fin].min()), float(lam[fin].max())
    if not hi > lo:
        raise ValueError(
            "spectrum_calib: observed wavelength range is degenerate "
            f"(min = max = {lo})."
        )
    x = 2.0 * (lam - lo) / (hi - lo) - 1.0
    x = np.where(fin, x, 0.0)
    return np.polynomial.legendre.legvander(x, int(order))[:, 1:]