Skip to content

Instrument models

ConstantGain

Bases: GainModel

A scalar gain factor, independent of energy.

The factor lives as :attr:factor (an nnx.Param). Its prior is provided via the unified prior dict under the key "instrument.gain.factor" (shared across instrumented obs) or "instrument.gain.factor[*]" / "instrument.gain.factor[obs_name]" (per-obs).

Source code in src/jaxspec/model/instrument.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
class ConstantGain(GainModel):
    """A scalar gain factor, independent of energy.

    The factor lives as :attr:`factor` (an ``nnx.Param``). Its prior is provided
    via the unified prior dict under the key ``"instrument.gain.factor"``
    (shared across instrumented obs) or ``"instrument.gain.factor[*]"`` /
    ``"instrument.gain.factor[obs_name]"`` (per-obs).
    """

    def __init__(self):
        self.factor = nnx.Param(jnp.asarray(1.0))

    def __call__(self, energies: ArrayLike) -> ArrayLike:
        return self.factor[...]

ConstantShift

Bases: ShiftModel

An additive energy shift, constant across the spectrum.

The offset lives as :attr:offset (an nnx.Param). Its prior is provided via the unified prior dict under the key "instrument.shift.offset" (shared) or "instrument.shift.offset[*]" / "instrument.shift.offset[obs_name]" (per-obs).

Source code in src/jaxspec/model/instrument.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
class ConstantShift(ShiftModel):
    """An additive energy shift, constant across the spectrum.

    The offset lives as :attr:`offset` (an ``nnx.Param``). Its prior is provided
    via the unified prior dict under the key ``"instrument.shift.offset"``
    (shared) or ``"instrument.shift.offset[*]"`` / ``"instrument.shift.offset[obs_name]"``
    (per-obs).
    """

    def __init__(self):
        self.offset = nnx.Param(jnp.asarray(0.0))

    def __call__(self, energies: ArrayLike) -> ArrayLike:
        return energies + self.offset[...]

GainModel

Bases: Module

Generic gain model. __call__(energies) returns the per-energy gain factor.

Source code in src/jaxspec/model/instrument.py 
23
24
25
26
27
class GainModel(nnx.Module):
    """Generic gain model. ``__call__(energies)`` returns the per-energy gain factor."""

    @abstractmethod
    def __call__(self, energies: ArrayLike) -> ArrayLike: ...

InstrumentModel

Bases: Module

Per-observation instrument response.

Pass as a dict to :class:~jaxspec.fit.BayesianModel::

BayesianModel(
    spectral_model, prior, observations,
    instrument_model={
        "PN": None, # explicit reference
        "MOS1": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
        "MOS2": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
    },
)

None entries (or simply omitting an observation) apply the identity fold (transfer_matrix @ flux) — useful for the reference instrument.

Parameters:

Name Type Description Default
gain GainModel | None

Optional :class:GainModel (e.g. :class:ConstantGain). When None, no flux scaling is applied.

 None
shift ShiftModel | None

Optional :class:ShiftModel (e.g. :class:ConstantShift). When None, the input energies pass through unchanged.

 None
Source code in src/jaxspec/model/instrument.py 
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
class InstrumentModel(nnx.Module):
    """Per-observation instrument response.

    Pass as a dict to :class:`~jaxspec.fit.BayesianModel`::

        BayesianModel(
            spectral_model, prior, observations,
            instrument_model={
                "PN": None, # explicit reference
                "MOS1": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
                "MOS2": InstrumentModel(gain=ConstantGain(), shift=ConstantShift()),
            },
        )

    ``None`` entries (or simply omitting an observation) apply the identity
    fold (``transfer_matrix @ flux``) — useful for the reference instrument.

    Parameters:
        gain: Optional :class:`GainModel` (e.g. :class:`ConstantGain`). When
            ``None``, no flux scaling is applied.
        shift: Optional :class:`ShiftModel` (e.g. :class:`ConstantShift`). When
            ``None``, the input energies pass through unchanged.
    """

    #: When ``True``, :class:`~jaxspec.fit._forward_model.ForwardModel` builds
    #: the un-merged response components (``redistribution``, ``grouping``,
    #: ``area``, ``exposure``) into the per-observation cache passed to
    #: :meth:`fold`. Subclasses set this to ``True`` when their math needs the
    #: components separately (e.g. pileup, RMF calibration).
    requires_components = False

    def __init__(self, gain: GainModel | None = None, shift: ShiftModel | None = None):
        self.gain = gain
        self.shift = shift

    def default_prior(self, observation: ObsConfiguration, obs_name: str) -> dict:
        """Return data-dependent default priors scoped to this obs.

        Mirrors :meth:`~jaxspec.model.background.BackgroundModel.default_prior`:
        subclasses (e.g. pileup models with per-obs dead-time / grade-fraction
        parameters) override this to inject ``[obs_name]``-scoped defaults.
        User prior entries override these defaults.
        """
        return {}

    def apply_shift(self, energies: ArrayLike) -> ArrayLike:
        """Apply :attr:`shift` to ``energies`` and clip away non-positive values."""
        if self.shift is None:
            return energies
        return jnp.clip(self.shift(energies), min=1e-6)

    def apply_gain(self, flux, energies: ArrayLike):
        """Multiply ``flux`` (or each branch in a pytree) by :attr:`gain`'s factor."""
        if self.gain is None:
            return flux
        factor = jnp.clip(self.gain(energies), min=0.0)
        return jax.tree.map(lambda f: f * factor, flux)

    def fold(
        self,
        spectrum: ArrayLike,
        cache: dict,
        eval_energies: ArrayLike | None = None,
    ):
        """
        Fold the input spectrum (or branches of a pytree) into the instrument using the pre-computed transfer matrix.
        """

        # Current contract: shift is applied here if a grid is provided, else in
        # the forward model. The spectrum was integrated on the *unshifted*
        # ``eval_energies``, so the shift must go on the native grid: bin i
        # collects the flux from its shifted band, matching the no-grid path
        # (which evaluates flux_func on the shifted native energies directly).
        if eval_energies is not None:
            target_energies = self.apply_shift(cache["in_energies"])
            spectrum = jax.tree.map(
                lambda s: redistribute(s, *eval_energies, *target_energies), spectrum
            )

        spectrum = jax.tree.map(lambda s: self.apply_gain(s, cache["in_energies"]), spectrum)

        return jax.tree.map(lambda s: jnp.clip(cache["transfer_matrix"] @ s, min=1e-6), spectrum)

apply_gain(flux, energies)

Multiply flux (or each branch in a pytree) by :attr:gain's factor.

Source code in src/jaxspec/model/instrument.py 
120
121
122
123
124
125
def apply_gain(self, flux, energies: ArrayLike):
    """Multiply ``flux`` (or each branch in a pytree) by :attr:`gain`'s factor."""
    if self.gain is None:
        return flux
    factor = jnp.clip(self.gain(energies), min=0.0)
    return jax.tree.map(lambda f: f * factor, flux)

apply_shift(energies)

Apply :attr:shift to energies and clip away non-positive values.

Source code in src/jaxspec/model/instrument.py 
114
115
116
117
118
def apply_shift(self, energies: ArrayLike) -> ArrayLike:
    """Apply :attr:`shift` to ``energies`` and clip away non-positive values."""
    if self.shift is None:
        return energies
    return jnp.clip(self.shift(energies), min=1e-6)

default_prior(observation, obs_name)

Return data-dependent default priors scoped to this obs.

Mirrors :meth:~jaxspec.model.background.BackgroundModel.default_prior: subclasses (e.g. pileup models with per-obs dead-time / grade-fraction parameters) override this to inject [obs_name]-scoped defaults. User prior entries override these defaults.

Source code in src/jaxspec/model/instrument.py 
104
105
106
107
108
109
110
111
112
def default_prior(self, observation: ObsConfiguration, obs_name: str) -> dict:
    """Return data-dependent default priors scoped to this obs.

    Mirrors :meth:`~jaxspec.model.background.BackgroundModel.default_prior`:
    subclasses (e.g. pileup models with per-obs dead-time / grade-fraction
    parameters) override this to inject ``[obs_name]``-scoped defaults.
    User prior entries override these defaults.
    """
    return {}

fold(spectrum, cache, eval_energies=None)

Fold the input spectrum (or branches of a pytree) into the instrument using the pre-computed transfer matrix.

Source code in src/jaxspec/model/instrument.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def fold(
    self,
    spectrum: ArrayLike,
    cache: dict,
    eval_energies: ArrayLike | None = None,
):
    """
    Fold the input spectrum (or branches of a pytree) into the instrument using the pre-computed transfer matrix.
    """

    # Current contract: shift is applied here if a grid is provided, else in
    # the forward model. The spectrum was integrated on the *unshifted*
    # ``eval_energies``, so the shift must go on the native grid: bin i
    # collects the flux from its shifted band, matching the no-grid path
    # (which evaluates flux_func on the shifted native energies directly).
    if eval_energies is not None:
        target_energies = self.apply_shift(cache["in_energies"])
        spectrum = jax.tree.map(
            lambda s: redistribute(s, *eval_energies, *target_energies), spectrum
        )

    spectrum = jax.tree.map(lambda s: self.apply_gain(s, cache["in_energies"]), spectrum)

    return jax.tree.map(lambda s: jnp.clip(cache["transfer_matrix"] @ s, min=1e-6), spectrum)

ShiftModel

Bases: Module

Generic shift model. __call__(energies) returns shifted energies.

Source code in src/jaxspec/model/instrument.py 
46
47
48
49
50
class ShiftModel(nnx.Module):
    """Generic shift model. ``__call__(energies)`` returns shifted energies."""

    @abstractmethod
    def __call__(self, energies: ArrayLike) -> ArrayLike: ...