jaxspec 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

jaxspec/model/additive.py

@@ -1,16 +1,22 @@
 from __future__ import annotations
 
-
+import astropy.constants
+import astropy.units as u
 import haiku as hk
+import interpax
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
-import astropy.units as u
-import astropy.constants
+import numpy as np
+
+from astropy.table import Table
 from haiku.initializers import Constant as HaikuConstant
+
 from ..util.integrate import integrate_interval
+from ..util.online_storage import table_manager
+
+# from ._additive.apec import APEC
 from .abc import AdditiveComponent
-# from ._additive.apec import APEC  # noqa: F401
 
 
 class Powerlaw(AdditiveComponent):
@@ -26,8 +32,8 @@ class Powerlaw(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        alpha = hk.get_parameter("alpha", [], init=HaikuConstant(1.3))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1e-4))
+        alpha = hk.get_parameter("alpha", [], float, init=HaikuConstant(1.3))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1e-4))
 
         return norm * energy ** (-alpha)
 
@@ -43,12 +49,12 @@ class AdditiveConstant(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * jnp.ones_like(energy)
 
     def primitive(self, energy):
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * energy
 
@@ -66,9 +72,9 @@ class Lorentz(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        line_energy = hk.get_parameter("E_l", [], init=HaikuConstant(1))
-        sigma = hk.get_parameter("sigma", [], init=HaikuConstant(1))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        line_energy = hk.get_parameter("E_l", [], float, init=HaikuConstant(1))
+        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * sigma / (2 * jnp.pi) / ((energy - line_energy) ** 2 + (sigma / 2) ** 2)
 
@@ -91,9 +97,9 @@ class Logparabola(AdditiveComponent):
 
     # TODO : conform with xspec definition
     def continuum(self, energy):
-        a = hk.get_parameter("a", [], init=HaikuConstant(11 / 3))
-        b = hk.get_parameter("b", [], init=HaikuConstant(0.2))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        a = hk.get_parameter("a", [], float, init=HaikuConstant(11 / 3))
+        b = hk.get_parameter("b", [], float, init=HaikuConstant(0.2))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * energy ** (-(a + b * jnp.log(energy)))
 
@@ -111,8 +117,8 @@ class Blackbody(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        kT = hk.get_parameter("kT", [], init=HaikuConstant(11 / 3))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(11 / 3))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * 8.0525 * energy**2 / ((kT**4) * (jnp.exp(energy / kT) - 1))
 
@@ -130,8 +136,8 @@ class Blackbodyrad(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        kT = hk.get_parameter("kT", [], init=HaikuConstant(11 / 3))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(11 / 3))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * 1.0344e-3 * energy**2 / (jnp.exp(energy / kT) - 1)
 
@@ -150,9 +156,9 @@ class Gauss(AdditiveComponent):
     """
 
     def continuum(self, energy) -> (jax.Array, jax.Array):
-        line_energy = hk.get_parameter("E_l", [], init=HaikuConstant(1))
-        sigma = hk.get_parameter("sigma", [], init=HaikuConstant(1))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        line_energy = hk.get_parameter("E_l", [], float, init=HaikuConstant(1))
+        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * jsp.stats.norm.pdf(energy, loc=line_energy, scale=sigma)
 
@@ -198,9 +204,9 @@ class APEC(AdditiveComponent):
         )
 
     def mono_fine_structure(self, e_low, e_high) -> (jax.Array, jax.Array):
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
-        kT = hk.get_parameter("kT", [], init=HaikuConstant(1))
-        Z = hk.get_parameter("Z", [], init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(1))
+        Z = hk.get_parameter("Z", [], float, init=HaikuConstant(1))
 
         idx = jnp.searchsorted(self.kT_ref, kT, side="left") - 1
 
@@ -229,9 +235,9 @@ class APEC(AdditiveComponent):
 
     @partial(jnp.vectorize, excluded=(0,))
     def continuum(self, energy):
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
-        kT = hk.get_parameter("kT", [], init=HaikuConstant(1))
-        Z = hk.get_parameter("Z", [], init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(1))
+        Z = hk.get_parameter("Z", [], float, init=HaikuConstant(1))
 
         idx = jnp.searchsorted(self.kT_ref, kT, side="left") - 1  # index of left value
 
@@ -269,9 +275,9 @@ class Cutoffpl(AdditiveComponent):
     """
 
     def continuum(self, energy):
-        alpha = hk.get_parameter("alpha", [], init=HaikuConstant(1.3))
-        beta = hk.get_parameter("beta", [], init=HaikuConstant(15))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1e-4))
+        alpha = hk.get_parameter("alpha", [], float, init=HaikuConstant(1.3))
+        beta = hk.get_parameter("beta", [], float, init=HaikuConstant(15))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1e-4))
 
         return norm * energy ** (-alpha) * jnp.exp(-energy / beta)
 
@@ -283,22 +289,22 @@ class Diskpbb(AdditiveComponent):
     where $$p$$ is a free parameter. The standard disk model, diskbb, is recovered if $$p=0.75$$.
     If radial advection is important then $$p<0.75$$.
 
-    $$\mathcal{M}\left( E \right) = \frac{2\pi(\cos i)r^{2}_{\text{in}}}{pd^2} \int_{T_{\text{in}}}^{T_{\text{out}}}
-    \left( \frac{T}{T_{\text{in}}} \right)^{-2/p-1} \text{bbody}(E,T) \frac{dT}{T_{\text{in}}}$$
+    $$\\mathcal{M}\\left( E \right) = \frac{2\\pi(\\cos i)r^{2}_{\text{in}}}{pd^2} \\int_{T_{\text{in}}}^{T_{\text{out}}}
+    \\left( \frac{T}{T_{\text{in}}} \right)^{-2/p-1} \text{bbody}(E,T) \frac{dT}{T_{\text{in}}}$$
 
     ??? abstract "Parameters"
-        * $\text{norm}$ : $\cos i(r_{\text{in}}/d)^{2}$,
-        where $r_{\text{in}}$ is "an apparent" inner disk radius $\left[\text{km}\right]$,
+        * $\text{norm}$ : $\\cos i(r_{\text{in}}/d)^{2}$,
+        where $r_{\text{in}}$ is "an apparent" inner disk radius $\\left[\text{km}\right]$,
         $d$ the distance to the source in units of $10 \text{kpc}$,
         $i$ the angle of the disk ($i=0$ is face-on)
-        * $p$ : Exponent of the radial dependence of the disk temperature $\left[\text{dimensionless}\right]$
-        * $T_{\text{in}}$ : Temperature at inner disk radius $\left[ \mathrm{keV}\right]$
+        * $p$ : Exponent of the radial dependence of the disk temperature $\\left[\text{dimensionless}\right]$
+        * $T_{\text{in}}$ : Temperature at inner disk radius $\\left[ \\mathrm{keV}\right]$
     """
 
     def continuum(self, energy):
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
-        p = hk.get_parameter("p", [], init=HaikuConstant(0.75))
-        tin = hk.get_parameter("Tin", [], init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+        p = hk.get_parameter("p", [], float, init=HaikuConstant(0.75))
+        tin = hk.get_parameter("Tin", [], float, init=HaikuConstant(1))
 
         # Tout is set to 0 as it is evaluated at R=infinity
         def integrand(kT, energy):
@@ -324,15 +330,21 @@ class Diskbb(AdditiveComponent):
     def continuum(self, energy):
         p = 0.75
         tout = 0.0
-        tin = hk.get_parameter("Tin", [], init=HaikuConstant(1))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        tin = hk.get_parameter("Tin", [], float, init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         # Tout is set to 0 as it is evaluated at R=infinity
         def integrand(kT, e, tin, p):
             return e**2 * (kT / tin) ** (-2 / p - 1) / (jnp.exp(e / kT) - 1)
 
         integral = integrate_interval(integrand)
-        return norm * 2.78e-3 * (0.75 / p) / tin * jnp.vectorize(lambda e: integral(tout, tin, e, tin, p))(energy)
+        return (
+            norm
+            * 2.78e-3
+            * (0.75 / p)
+            / tin
+            * jnp.vectorize(lambda e: integral(tout, tin, e, tin, p))(energy)
+        )
 
 
 class Agauss(AdditiveComponent):
@@ -352,9 +364,9 @@ class Agauss(AdditiveComponent):
 
     def continuum(self, energy) -> (jax.Array, jax.Array):
         hc = (astropy.constants.h * astropy.constants.c).to(u.angstrom * u.keV).value
-        line_wavelength = hk.get_parameter("Lambda_l", [], init=HaikuConstant(hc))
-        sigma = hk.get_parameter("sigma", [], init=HaikuConstant(0.001))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
+        line_wavelength = hk.get_parameter("Lambda_l", [], float, init=HaikuConstant(hc))
+        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(0.001))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
 
         return norm * jsp.stats.norm.pdf(hc / energy, loc=line_wavelength, scale=sigma)
 
@@ -376,12 +388,16 @@ class Zagauss(AdditiveComponent):
 
     def continuum(self, energy) -> (jax.Array, jax.Array):
         hc = (astropy.constants.h * astropy.constants.c).to(u.angstrom * u.keV).value
-        line_wavelength = hk.get_parameter("Lambda_l", [], init=HaikuConstant(hc))
-        sigma = hk.get_parameter("sigma", [], init=HaikuConstant(0.001))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
-        redshift = hk.get_parameter("redshift", [], init=HaikuConstant(0))
+        line_wavelength = hk.get_parameter("Lambda_l", [], float, init=HaikuConstant(hc))
+        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(0.001))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+        redshift = hk.get_parameter("redshift", [], float, init=HaikuConstant(0))
 
-        return norm * (1 + redshift) * jsp.stats.norm.pdf((hc / energy) / (1 + redshift), loc=line_wavelength, scale=sigma)
+        return (
+            norm
+            * (1 + redshift)
+            * jsp.stats.norm.pdf((hc / energy) / (1 + redshift), loc=line_wavelength, scale=sigma)
+        )
 
 
 class Zgauss(AdditiveComponent):
@@ -399,9 +415,132 @@ class Zgauss(AdditiveComponent):
     """
 
     def continuum(self, energy) -> (jax.Array, jax.Array):
-        line_energy = hk.get_parameter("E_l", [], init=HaikuConstant(1))
-        sigma = hk.get_parameter("sigma", [], init=HaikuConstant(1))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
-        redshift = hk.get_parameter("redshift", [], init=HaikuConstant(0))
+        line_energy = hk.get_parameter("E_l", [], float, init=HaikuConstant(1))
+        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(1))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+        redshift = hk.get_parameter("redshift", [], float, init=HaikuConstant(0))
+
+        return (norm / (1 + redshift)) * jsp.stats.norm.pdf(
+            energy * (1 + redshift), loc=line_energy, scale=sigma
+        )
+
+
+class NSatmos(AdditiveComponent):
+    r"""
+    A neutron star atmosphere model based on the `NSATMOS` model from `XSPEC`. See [this page](https://heasarc.gsfc.nasa.gov/docs/xanadu/xspec/manual/node205.html)
+
+    !!! warning
+        The boundary case of $R_{\text{NS}} < 1.125 R_{\text{S}}$ is handled with a null flux instead of a constant value as in `XSPEC`.
+
+    ??? abstract "Parameters"
+        * $T_{eff}$ : Effective temperature at the surface in K (No redshift applied)
+        * $M_{ns}$ : Mass of the NS in solar masses
+        * $R_∞$ : Radius at infinity (modulated by gravitational effects) in km
+        * $D$ : Distance to the neutron star in kpc
+        * norm : fraction of the neutron star surface emitting
+    """
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        entry_table = Table.read(table_manager.fetch("nsatmosdata.fits"), 1)
+
+        # Extract the table values. All this code could be summarize in two lines if we reformat the nsatmosdata.fits table
+        self.tab_temperature = np.asarray(entry_table["TEMP"][0], dtype=float)  # Logarithmic value
+        self.tab_gravity = np.asarray(entry_table["GRAVITY"][0], dtype=float)  # Logarithmic value
+        self.tab_mucrit = np.asarray(entry_table["MUCRIT"][0], dtype=float)
+        self.tab_energy = np.asarray(entry_table["ENERGY"][0], dtype=float)
+        self.tab_flux_flat = Table.read(table_manager.fetch("nsatmosdata.fits"), 2)["FLUX"]
+
+        tab_flux = np.empty(
+            (
+                self.tab_temperature.size,
+                self.tab_gravity.size,
+                self.tab_mucrit.size,
+                self.tab_energy.size,
+            )
+        )
+
+        for i in range(len(self.tab_temperature)):
+            for j in range(len(self.tab_gravity)):
+                for k in range(len(self.tab_mucrit)):
+                    tab_flux[i, j, k] = np.array(
+                        self.tab_flux_flat[
+                            i * len(self.tab_gravity) * len(self.tab_mucrit)
+                            + j * len(self.tab_mucrit)
+                            + k
+                        ]
+                    )
+
+        self.tab_flux = np.asarray(tab_flux, dtype=float)
+
+    def interp_flux_func(self, temperature_log, gravity_log, mu):
+        # Interpolate the tables to get the flux on the tabulated energy grid
+
+        return interpax.interp3d(
+            10.0**temperature_log,
+            10.0**gravity_log,
+            mu,
+            10.0**self.tab_temperature,
+            10.0**self.tab_gravity,
+            self.tab_mucrit,
+            self.tab_flux,
+            method="linear",
+        )
+
+    def continuum(self, energy):
+        temp_log = hk.get_parameter(
+            "Tinf", [], float, init=HaikuConstant(6.0)
+        )  # log10 of temperature in Kelvin
+
+        #         'Tinf': temp_log, # 5 to 6.5
+        #         'M': mass, # 0.5 to 3
+        #         'Rns': radius, # 5 to 30
+
+        mass = hk.get_parameter("M", [], float, init=HaikuConstant(1.4))
+        radius = hk.get_parameter("Rns", [], float, init=HaikuConstant(10.0))
+        distance = hk.get_parameter("dns", [], float, init=HaikuConstant(10.0))
+        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1.0))
+
+        # Derive parameters usable to retrive value in the flux table
+        Rcgs = 1e5 * radius  # Radius in cgs
+        r_schwarzschild = 2.95e5 * mass  # Schwarzschild radius in cgs
+        r_normalized = Rcgs / r_schwarzschild  # Ratio of the radius to the Schwarzschild radius
+        r_over_Dsql = 2 * jnp.log10(
+            Rcgs / (3.09e21 * distance)
+        )  # Log( (R/D)**2 ), 3.09e21 constant transforms radius in cgs to kpc
+        zred1 = 1 / jnp.sqrt(1 - (1 / r_normalized))  # Gravitational redshift
+        gravity = (6.67e-8 * 1.99e33 * mass / Rcgs**2) * zred1  # Gravity field g in cgs
+        gravity_log = jnp.log10(
+            gravity
+        )  # Log gravity because this is the format given in the table
+
+        # Not sure about mu yet, but it is linked to causality
+        cmu = jnp.where(
+            r_normalized < 1.5, jnp.sqrt(1.0 - 6.75 / r_normalized**2 + 6.75 / r_normalized**3), 0.0
+        )
+
+        # Interpolate the flux table to get the flux at the surface
+
+        flux = jax.jit(self.interp_flux_func)(temp_log, gravity_log, cmu)
+
+        # Rescale the photon energies and fluxes back to the correct local temperature
+        Tfactor = 10.0 ** (temp_log - 6.0)
+        fluxshift = 3.0 * (temp_log - 6.0)
+        E = self.tab_energy * Tfactor
+        flux += fluxshift
+
+        # Rescale applying redshift
+        fluxshift = -jnp.log10(zred1)
+        E = E / zred1
+        flux += fluxshift
+
+        # Convert to counts/keV (which corresponds to dividing by1.602e-9*EkeV)
+        # Multiply by the area of the star, and calculate the count rate at the observer
+        flux += r_over_Dsql
+        counts = 10.0 ** (flux - jnp.log10(1.602e-9) - jnp.log10(E))
+
+        true_flux = norm * jnp.exp(
+            interpax.interp1d(jnp.log(energy), jnp.log(E), jnp.log(counts), method="linear")
+        )
 
-        return (norm / (1 + redshift)) * jsp.stats.norm.pdf(energy * (1 + redshift), loc=line_energy, scale=sigma)
+        return jax.lax.select(r_normalized < 1.125, jnp.zeros_like(true_flux), true_flux)

jaxspec/model/background.py

@@ -1,9 +1,11 @@
 from abc import ABC, abstractmethod
-from tinygp import kernels, GaussianProcess
-from jax.scipy.integrate import trapezoid
-import numpyro.distributions as dist
+
 import jax.numpy as jnp
 import numpyro
+import numpyro.distributions as dist
+
+from jax.scipy.integrate import trapezoid
+from tinygp import GaussianProcess, kernels
 
 
 class BackgroundModel(ABC):
@@ -33,7 +35,8 @@ class SubtractedBackground(BackgroundModel):
 
     """
 
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
+        _, observed_counts = obs.out_energies, obs.folded_background.data
         numpyro.deterministic(f"{name}", observed_counts)
 
         return jnp.zeros_like(observed_counts)
@@ -49,32 +52,37 @@ class BackgroundWithError(BackgroundModel):
         but slower since it performs the fit using MCMC instead of analytical solution.
     """
 
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
         # 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"{name}_params", dist.Gamma(alpha, rate=beta))
 
         with numpyro.plate(f"{name}_plate", len(observed_counts)):
-            numpyro.sample(f"{name}", dist.Poisson(countrate), obs=observed_counts if observed else None)
+            numpyro.sample(
+                f"{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
+            )
 
         return countrate
 
 
 '''
+# TODO: Implement this class and sample it with Gibbs Sampling
+
 class ConjugateBackground(BackgroundModel):
     r"""
-    This class fit an expected rate $\lambda$ in each bin of the background spectrum. Assuming a Gamma prior
+    This class fit an expected rate $\\lambda$ in each bin of the background spectrum. Assuming a Gamma prior
     distribution, we can analytically derive the posterior as a Negative binomial distribution.
 
-    $$ p(\lambda_{\text{Bkg}}) \sim \Gamma \left( \alpha, \beta \right) \implies
-    p\left(\lambda_{\text{Bkg}} | \text{Counts}_{\text{Bkg}}\right) \sim \text{NB}\left(\alpha, \frac{\beta}{\beta +1}
+    $$ p(\\lambda_{\text{Bkg}}) \\sim \\Gamma \\left( \alpha, \beta \right) \\implies
+    p\\left(\\lambda_{\text{Bkg}} | \text{Counts}_{\text{Bkg}}\right) \\sim \text{NB}\\left(\alpha, \frac{\beta}{\beta +1}
     \right) $$
 
     !!! info
         Here, $\alpha$ and $\beta$ are set to $\alpha = \text{Counts}_{\text{Bkg}} + 1$ and $\beta = 1$. Doing so,
-        the prior distribution is such that $\mathbb{E}[\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}} +1$ and
-        $\text{Var}[\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}}+1$. The +1 is to avoid numerical issues when the
+        the prior distribution is such that $\\mathbb{E}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}} +1$ and
+        $\text{Var}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}}+1$. The +1 is to avoid numerical issues when the
         counts are 0, and add a small scatter even if the measured background is effectively null.
 
     ??? abstract "References"
@@ -97,6 +105,19 @@ class ConjugateBackground(BackgroundModel):
         return countrate
 '''
 
+"""
+class SpectralBackgroundModel(BackgroundModel):
+    # I should pass the current spectral model as an argument to the background model
+    # In the numpyro model function
+    def __init__(self, model, prior):
+        self.model = model
+        self.prior = prior
+
+    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+        #TODO : keep the sparsification from top model
+        transformed_model = hk.without_apply_rng(hk.transform(lambda par: CountForwardModel(model, obs, sparse=False)(par)))
+"""
+
 
 class GaussianProcessBackground(BackgroundModel):
     """
@@ -104,7 +125,13 @@ class GaussianProcessBackground(BackgroundModel):
     [`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):
+    def __init__(
+        self,
+        e_min: float,
+        e_max: float,
+        n_nodes: int = 30,
+        kernel: kernels.Kernel = kernels.Matern52,
+    ):
         """
         Build the Gaussian Process background model.
 
@@ -119,7 +146,7 @@ class GaussianProcessBackground(BackgroundModel):
         self.n_nodes = n_nodes
         self.kernel = kernel
 
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
         """
         Build the model for the background.
 
@@ -129,9 +156,12 @@ class GaussianProcessBackground(BackgroundModel):
             name: The name of the background model for parameters disambiguation.
             observed: Whether the model is observed or not. Useful for `numpyro.infer.Predictive` calls.
         """
+        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.")
+            raise RuntimeError(
+                "More nodes than channels in the observation associated with GaussianProcessBackground."
+            )
 
         # The parameters of the GP model
         mean = numpyro.sample(f"{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0))
@@ -144,13 +174,17 @@ class GaussianProcessBackground(BackgroundModel):
         gp = GaussianProcess(kernel, nodes, diag=1e-5 * jnp.ones_like(nodes), mean=mean)
 
         log_rate = numpyro.sample(f"_{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))
+        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(f"{name}_plate", len(observed_counts)):
             # TODO : change to Poisson Likelihood when there is no background model
             # TODO : Otherwise clip the background model to 1e-6 to avoid numerical issues
-            numpyro.sample(f"{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None)
+            numpyro.sample(
+                f"{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
+            )
 
         return count_rate