Skip to content

Background models

BackgroundModel

Bases: ABC

Handles the background modelling in our spectra. This is handled in a separate class for now since backgrounds can be phenomenological models fitted directly on the folded spectrum. This is not the case for the source model, which is fitted on the unfolded spectrum. This might be changed later.

Source code in src/jaxspec/model/background.py 
16
17
18
19
20
21
22
23
24
25
26
27
28
class BackgroundModel(ABC):
    """
    Handles the background modelling in our spectra. This is handled in a separate class for now
    since backgrounds can be phenomenological models fitted directly on the folded spectrum. This is not the case for
    the source model, which is fitted on the unfolded spectrum. This might be changed later.
    """

    @abstractmethod
    def numpyro_model(self, observation, name: str = "", observed=True):
        """
        Build the model for the background.
        """
        pass

numpyro_model(observation, name='', observed=True) abstractmethod

Build the model for the background.

Source code in src/jaxspec/model/background.py 
23
24
25
26
27
28
@abstractmethod
def numpyro_model(self, observation, name: str = "", observed=True):
    """
    Build the model for the background.
    """
    pass

BackgroundWithError

Bases: BackgroundModel

Define a model where the observed background is subtracted from the observed accounting for its intrinsic spread. It fits a countrate for each background bin assuming a Poisson distribution.

Source code in src/jaxspec/model/background.py 
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
class BackgroundWithError(BackgroundModel):
    """
    Define a model where the observed background is subtracted from the observed accounting for its intrinsic spread. It
    fits a countrate for each background bin assuming a Poisson distribution.
    """

    def numpyro_model(self, obs, name: str = "", observed=True):
        # We can't use the build_prior_function method here because the parameter size varies
        # with the current observation. It must be instantiated in place.
        # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)

        _, observed_counts = obs.out_energies, obs.folded_background.data
        alpha = observed_counts + 1
        beta = 1
        countrate = numpyro.sample(f"_bkg_{name}_countrate", dist.Gamma(alpha, rate=beta))

        with numpyro.plate(f"bkg_{name}_plate", len(observed_counts)):
            numpyro.sample(
                f"bkg_{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
            )

        return countrate

GaussianProcessBackground

Bases: BackgroundModel

Define a Gaussian Process to model the background. The GP is built using the tinygp library.

Source code in src/jaxspec/model/background.py 
 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
class GaussianProcessBackground(BackgroundModel):
    """
    Define a Gaussian Process to model the background. The GP is built using the
    [`tinygp`](https://tinygp.readthedocs.io/en/stable/guide.html) library.
    """

    def __init__(
        self,
        e_min: float,
        e_max: float,
        n_nodes: int = 30,
        kernel: kernels.Kernel = kernels.Matern52,
    ):
        """
        Build the Gaussian Process background model.

        Parameters:
            e_min: The lower bound of the energy range.
            e_max: The upper bound of the energy range.
            n_nodes: The number of nodes used by the GP, must be lower than the number of channel.
            kernel: The kernel used by the GP.
        """
        self.e_min = e_min
        self.e_max = e_max
        self.n_nodes = n_nodes
        self.kernel = kernel

    def numpyro_model(self, obs, name: str = "", observed=True):
        energy, observed_counts = obs.out_energies, obs.folded_background.data

        if (observed_counts is not None) and (self.n_nodes >= len(observed_counts)):
            raise RuntimeError(
                "More nodes than channels in the observation associated with GaussianProcessBackground."
            )

        else:
            observed_counts = jnp.asarray(observed_counts)

        # The parameters of the GP model
        mean = numpyro.sample(
            f"_bkg_{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0)
        )
        sigma = numpyro.sample(f"_bkg_{name}_sigma", dist.HalfNormal(3.0))
        rho = numpyro.sample(f"_bkg_{name}_rho", dist.HalfNormal(10.0))

        # Set up the kernel and GP objects
        kernel = sigma**2 * self.kernel(rho)
        nodes = jnp.linspace(0, 1, self.n_nodes)
        gp = GaussianProcess(kernel, nodes, diag=1e-5 * jnp.ones_like(nodes), mean=mean)

        log_rate = numpyro.sample(f"_bkg_{name}_log_rate_nodes", gp.numpyro_dist())

        interp_count_rate = jnp.exp(
            jnp.interp(energy, nodes * (self.e_max - self.e_min) + self.e_min, log_rate)
        )
        count_rate = trapezoid(interp_count_rate, energy, axis=0)

        # Finally, our observation model is Poisson
        with numpyro.plate("bkg_plate_" + name, len(observed_counts)):
            numpyro.sample(
                f"bkg_{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
            )

        return count_rate

__init__(e_min, e_max, n_nodes=30, kernel=kernels.Matern52)

Build the Gaussian Process background model.

Parameters:

Name Type Description Default
e_min float

The lower bound of the energy range.

 required
e_max float

The upper bound of the energy range.

 required
n_nodes int

The number of nodes used by the GP, must be lower than the number of channel.

 30
kernel Kernel

The kernel used by the GP.

 Matern52
Source code in src/jaxspec/model/background.py 
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def __init__(
    self,
    e_min: float,
    e_max: float,
    n_nodes: int = 30,
    kernel: kernels.Kernel = kernels.Matern52,
):
    """
    Build the Gaussian Process background model.

    Parameters:
        e_min: The lower bound of the energy range.
        e_max: The upper bound of the energy range.
        n_nodes: The number of nodes used by the GP, must be lower than the number of channel.
        kernel: The kernel used by the GP.
    """
    self.e_min = e_min
    self.e_max = e_max
    self.n_nodes = n_nodes
    self.kernel = kernel

SubtractedBackground

Bases: BackgroundModel

Define a model where the observed background is simply subtracted from the observed.

Danger

This is not a good way to model the background, as it does not account for the fact that the measured background is a Poisson realisation of the true background's countrate. This is why we prefer a [ConjugateBackground][jaxspec.model.background.ConjugateBackground].

Source code in src/jaxspec/model/background.py 
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
class SubtractedBackground(BackgroundModel):
    """
    Define a model where the observed background is simply subtracted from the observed.

    !!! danger

        This is not a good way to model the background, as it does not account for the fact that the measured
        background is a Poisson realisation of the true background's countrate. This is why we prefer a
        [`ConjugateBackground`][jaxspec.model.background.ConjugateBackground].

    """

    def numpyro_model(self, observation, name: str = "", observed=True):
        _, observed_counts = observation.out_energies, observation.folded_background.data
        numpyro.deterministic(f"bkg_{name}", observed_counts)

        return observed_counts