Instrument models

ConstantGain

Bases: GainModel

A constant gain model

Source code in src/jaxspec/model/instrument.py

class ConstantGain(GainModel):
    """
    A constant gain model
    """

    def __init__(self, prior_distribution: Distribution):
        """
        Parameters:
            prior_distribution: the prior distribution for the gain value.
        """

        self.prior_distribution = prior_distribution

    def numpyro_model(self, observation_name: str):
        factor = numpyro.sample(f"ins/~/gain_{observation_name}", self.prior_distribution)

        def gain(energy):
            return factor

        return gain

__init__(prior_distribution)

Parameters:

Name	Type	Description	Default
prior_distribution	Distribution	the prior distribution for the gain value.	required

Source code in src/jaxspec/model/instrument.py

def __init__(self, prior_distribution: Distribution):
    """
    Parameters:
        prior_distribution: the prior distribution for the gain value.
    """

    self.prior_distribution = prior_distribution

ConstantShift

Bases: ShiftModel

A constant shift model

Source code in src/jaxspec/model/instrument.py

class ConstantShift(ShiftModel):
    """
    A constant shift model
    """

    def __init__(self, prior_distribution: Distribution):
        """
        Parameters:
            prior_distribution: the prior distribution for the shift value.
        """
        self.prior_distribution = prior_distribution

    def numpyro_model(self, observation_name: str):
        shift_offset = numpyro.sample(f"ins/~/shift_{observation_name}", self.prior_distribution)

        def shift(energy):
            return energy + shift_offset

        return shift

__init__(prior_distribution)

Parameters:

Name	Type	Description	Default
prior_distribution	Distribution	the prior distribution for the shift value.	required

Source code in src/jaxspec/model/instrument.py

def __init__(self, prior_distribution: Distribution):
    """
    Parameters:
        prior_distribution: the prior distribution for the shift value.
    """
    self.prior_distribution = prior_distribution

GainModel

Bases: ABC, Module

Generic class for a gain model

Source code in src/jaxspec/model/instrument.py

class GainModel(ABC, nnx.Module):
    """
    Generic class for a gain model
    """

    @abstractmethod
    def numpyro_model(self, observation_name: str):
        pass

InstrumentModel

Bases: Module

Source code in src/jaxspec/model/instrument.py

class InstrumentModel(nnx.Module):
    def __init__(
        self,
        reference_observation_name: str,
        gain_model: GainModel | None = None,
        shift_model: ShiftModel | None = None,
    ):
        """
        Encapsulate an instrument model, build as a combination of a shift and gain model.

        Parameters:
            reference_observation_name : The observation to use as a reference
            gain_model : The gain model
            shift_model : The shift model
        """

        self.reference = reference_observation_name
        self.gain_model = gain_model
        self.shift_model = shift_model

    def get_gain_and_shift_model(
        self, observation_name: str
    ) -> tuple[Callable | None, Callable | None]:
        """
        Return the gain and shift models for the given observation. It should be called within a numpyro model.
        """

        if observation_name == self.reference:
            return None, None

        else:
            gain = (
                self.gain_model.numpyro_model(observation_name)
                if self.gain_model is not None
                else None
            )
            shift = (
                self.shift_model.numpyro_model(observation_name)
                if self.shift_model is not None
                else None
            )

            return gain, shift

__init__(reference_observation_name, gain_model=None, shift_model=None)

Encapsulate an instrument model, build as a combination of a shift and gain model.

Parameters:

Name	Type	Description	Default
reference_observation_name		The observation to use as a reference	required
gain_model		The gain model	required
shift_model		The shift model	required

Source code in src/jaxspec/model/instrument.py

def __init__(
    self,
    reference_observation_name: str,
    gain_model: GainModel | None = None,
    shift_model: ShiftModel | None = None,
):
    """
    Encapsulate an instrument model, build as a combination of a shift and gain model.

    Parameters:
        reference_observation_name : The observation to use as a reference
        gain_model : The gain model
        shift_model : The shift model
    """

    self.reference = reference_observation_name
    self.gain_model = gain_model
    self.shift_model = shift_model

get_gain_and_shift_model(observation_name)

Return the gain and shift models for the given observation. It should be called within a numpyro model.

Source code in src/jaxspec/model/instrument.py

def get_gain_and_shift_model(
    self, observation_name: str
) -> tuple[Callable | None, Callable | None]:
    """
    Return the gain and shift models for the given observation. It should be called within a numpyro model.
    """

    if observation_name == self.reference:
        return None, None

    else:
        gain = (
            self.gain_model.numpyro_model(observation_name)
            if self.gain_model is not None
            else None
        )
        shift = (
            self.shift_model.numpyro_model(observation_name)
            if self.shift_model is not None
            else None
        )

        return gain, shift

ShiftModel

Bases: ABC, Module

Generic class for a shift model

Source code in src/jaxspec/model/instrument.py

class ShiftModel(ABC, nnx.Module):
    """
    Generic class for a shift model
    """

    @abstractmethod
    def numpyro_model(self, observation_name: str):
        pass