jaxspec 0.0.5__py3-none-any.whl → 0.0.7__py3-none-any.whl

jaxspec/model/_additive/apec.py

@@ -1,16 +1,21 @@
-import jax.numpy as jnp
-import jax
-import haiku as hk
+import warnings
+
+from typing import Literal
+
 import astropy.units as u
+import haiku as hk
+import jax
+import jax.numpy as jnp
+
+from astropy.constants import c, m_p
+from haiku.initializers import Constant as HaikuConstant
 from jax import lax
-from jax.lax import scan, fori_loop
+from jax.lax import fori_loop, scan
 from jax.scipy.stats import norm as gaussian
-from typing import Literal
+
 from ...util.abundance import abundance_table, element_data
-from haiku.initializers import Constant as HaikuConstant
-from astropy.constants import c, m_p
 from ..abc import AdditiveComponent
-from .apec_loaders import get_temperature, get_continuum, get_pseudo, get_lines
+from .apec_loaders import get_continuum, get_lines, get_pseudo, get_temperature
 
 
 @jax.jit
@@ -51,7 +56,9 @@ def interp_and_integrate(energy_low, energy_high, energy_ref, continuum_ref, end
         # Within
 
         current_energy_is_between = (energy_low <= current_energy) * (current_energy < energy_high)
-        previous_energy_is_between = (energy_low <= previous_energy) * (previous_energy < energy_high)
+        previous_energy_is_between = (energy_low <= previous_energy) * (
+            previous_energy < energy_high
+        )
         energies_within_bins = (previous_energy <= energy_low) * (energy_high < current_energy)
 
         case = (
@@ -68,7 +75,9 @@ def interp_and_integrate(energy_low, energy_high, energy_ref, continuum_ref, end
                 lambda pe, pc, ce, cc, el, er: (cc + lerp(el, pe, ce, pc, cc)) * (ce - el) / 2,  # 2
                 lambda pe, pc, ce, cc, el, er: (pc + lerp(er, pe, ce, pc, cc)) * (er - pe) / 2,  # 3
                 lambda pe, pc, ce, cc, el, er: (pc + cc) * (ce - pe) / 2,  # 4
-                lambda pe, pc, ce, cc, el, er: (lerp(el, pe, ce, pc, cc) + lerp(er, pe, ce, pc, cc)) * (er - el) / 2,
+                lambda pe, pc, ce, cc, el, er: (lerp(el, pe, ce, pc, cc) + lerp(er, pe, ce, pc, cc))
+                * (er - el)
+                / 2,
                 # 5
             ],
             previous_energy,
@@ -86,23 +95,24 @@ def interp_and_integrate(energy_low, energy_high, energy_ref, continuum_ref, end
     return integrated_flux
 
 
+@jax.jit
 def interp(e_low, e_high, energy_ref, continuum_ref, end_index):
     energy_ref = jnp.where(jnp.arange(energy_ref.shape[0]) < end_index, energy_ref, jnp.nan)
 
-    return (jnp.interp(e_high, energy_ref, continuum_ref) - jnp.interp(e_low, energy_ref, continuum_ref)) / (e_high - e_low)
+    return (
+        jnp.interp(e_high, energy_ref, continuum_ref) - jnp.interp(e_low, energy_ref, continuum_ref)
+    ) / (e_high - e_low)
 
 
-def interp_flux(energy, energy_ref, continuum_ref, end_index, integrate=True):
+@jax.jit
+def interp_flux(energy, energy_ref, continuum_ref, end_index):
     """
     Iterate through an array of shape (energy_ref,) and compute the flux between the bins defined by energy
     """
 
     def scanned_func(carry, unpack):
         e_low, e_high = unpack
-        if integrate:
-            continuum = interp_and_integrate(e_low, e_high, energy_ref, continuum_ref, end_index)
-        else:
-            continuum = interp(e_low, e_high, energy_ref, continuum_ref, end_index)
+        continuum = interp_and_integrate(e_low, e_high, energy_ref, continuum_ref, end_index)
 
         return carry, continuum
 
@@ -111,7 +121,8 @@ def interp_flux(energy, energy_ref, continuum_ref, end_index, integrate=True):
     return continuum
 
 
-def interp_flux_elements(energy_ref, continuum_ref, end_index, energy, abundances, integrate=True):
+@jax.jit
+def interp_flux_elements(energy_ref, continuum_ref, end_index, energy, abundances):
     """
     Iterate through an array of shape (abundance, energy_ref) and compute the flux between the bins defined by energy
     and weight the flux depending on the abundance of each element
@@ -119,7 +130,7 @@ def interp_flux_elements(energy_ref, continuum_ref, end_index, energy, abundance
 
     def scanned_func(_, unpack):
         energy_ref, continuum_ref, end_idx = unpack
-        element_flux = interp_flux(energy, energy_ref, continuum_ref, end_idx, integrate=integrate)
+        element_flux = interp_flux(energy, energy_ref, continuum_ref, end_idx)
 
         return _, element_flux
 
@@ -136,23 +147,9 @@ def get_lines_contribution_broadening(
         # Notice the -1 in line element to match the 0-based indexing
         l_energy, l_emissivity, l_element = line_energy[i], line_emissivity[i], line_element[i] - 1
         broadening = l_energy * total_broadening[l_element]
-        l_flux = gaussian.cdf(energy[1:], l_energy, broadening) - gaussian.cdf(energy[:-1], l_energy, broadening)
-        l_flux = l_flux * l_emissivity * abundances[l_element]
-
-        return flux + l_flux
-
-    return fori_loop(0, end_index, body_func, jnp.zeros_like(energy[:-1]))
-
-
-@jax.jit
-def get_lines_contribution_broadening_derivative(
-    line_energy, line_element, line_emissivity, end_index, energy, abundances, total_broadening
-):
-    def body_func(i, flux):
-        # Notice the -1 in line element to match the 0-based indexing
-        l_energy, l_emissivity, l_element = line_energy[i], line_emissivity[i], line_element[i] - 1
-        broadening = l_energy * total_broadening[l_element]
-        l_flux = gaussian.cdf(energy[1:], l_energy, broadening) - gaussian.cdf(energy[:-1], l_energy, broadening)
+        l_flux = gaussian.cdf(energy[1:], l_energy, broadening) - gaussian.cdf(
+            energy[:-1], l_energy, broadening
+        )
         l_flux = l_flux * l_emissivity * abundances[l_element]
 
         return flux + l_flux
@@ -160,7 +157,6 @@ def get_lines_contribution_broadening_derivative(
     return fori_loop(0, end_index, body_func, jnp.zeros_like(energy[:-1]))
 
 
-@jax.custom_jvp
 @jax.jit
 def continuum_func(energy, kT, abundances):
     idx, kT_low, kT_high = get_temperature(kT)
@@ -170,96 +166,27 @@ def continuum_func(energy, kT, abundances):
     return lerp(kT, kT_low, kT_high, continuum_low, continuum_high)
 
 
-@jax.jit
-@continuum_func.defjvp
-def continuum_jvp(primals, tangents):
-    energy, kT, abundances = primals
-    energy_dot, kT_dot, abundances_dot = tangents
-
-    idx, kT_low, kT_high = get_temperature(kT)
-    continuum_low_pars = get_continuum(idx)
-    continuum_high_pars = get_continuum(idx + 1)
-
-    # Energy derivative
-    dcontinuum_low_denerg = interp_flux_elements(*continuum_low_pars, energy, abundances, integrate=False)
-    dcontinuum_high_denerg = interp_flux_elements(*continuum_high_pars, energy, abundances, integrate=False)
-    energy_derivative = lerp(kT, kT_low, kT_high, dcontinuum_low_denerg, dcontinuum_high_denerg) * jnp.diff(energy_dot)
-
-    # Temperature derivative
-    continuum_low = interp_flux_elements(*continuum_low_pars, energy, abundances)
-    continuum_high = interp_flux_elements(*continuum_high_pars, energy, abundances)
-    kT_derivative = (continuum_high - continuum_low) / (kT_high - kT_low) * kT_dot
-
-    # Abundances derivative
-    dcontinuum_low_dabund = interp_flux_elements(*continuum_low_pars, energy, abundances_dot)
-    dcontinuum_high_dabund = interp_flux_elements(*continuum_high_pars, energy, abundances_dot)
-    abundances_derivative = lerp(kT, kT_low, kT_high, dcontinuum_low_dabund, dcontinuum_high_dabund)
-
-    primals_out = continuum_func(*primals)
-
-    return primals_out, energy_derivative + kT_derivative + abundances_derivative
-
-
-@jax.custom_jvp
 @jax.jit
 def pseudo_func(energy, kT, abundances):
     idx, kT_low, kT_high = get_temperature(kT)
-    continuum_low = interp_flux_elements(*get_continuum(idx), energy, abundances)
-    continuum_high = interp_flux_elements(*get_continuum(idx + 1), energy, abundances)
+    continuum_low = interp_flux_elements(*get_pseudo(idx), energy, abundances)
+    continuum_high = interp_flux_elements(*get_pseudo(idx + 1), energy, abundances)
 
     return lerp(kT, kT_low, kT_high, continuum_low, continuum_high)
 
 
-@jax.jit
-@pseudo_func.defjvp
-def pseudo_jvp(primals, tangents):
-    energy, kT, abundances = primals
-    energy_dot, kT_dot, abundances_dot = tangents
-
-    idx, kT_low, kT_high = get_temperature(kT)
-    continuum_low_pars = get_pseudo(idx)
-    continuum_high_pars = get_pseudo(idx + 1)
-
-    # Energy derivative
-    dcontinuum_low_denerg = interp_flux_elements(*continuum_low_pars, energy, abundances, integrate=False)
-    dcontinuum_high_denerg = interp_flux_elements(*continuum_high_pars, energy, abundances, integrate=False)
-    energy_derivative = lerp(kT, kT_low, kT_high, dcontinuum_low_denerg, dcontinuum_high_denerg) * jnp.diff(energy_dot)
-
-    # Temperature derivative
-    continuum_low = interp_flux_elements(*continuum_low_pars, energy, abundances)
-    continuum_high = interp_flux_elements(*continuum_high_pars, energy, abundances)
-    kT_derivative = (continuum_high - continuum_low) / (kT_high - kT_low) * kT_dot
-
-    # Abundances derivative
-    dcontinuum_low_dabund = interp_flux_elements(*continuum_low_pars, energy, abundances_dot)
-    dcontinuum_high_dabund = interp_flux_elements(*continuum_high_pars, energy, abundances_dot)
-    abundances_derivative = lerp(kT, kT_low, kT_high, dcontinuum_low_dabund, dcontinuum_high_dabund)
-
-    primals_out = pseudo_func(*primals)
-
-    return primals_out, energy_derivative + kT_derivative + abundances_derivative
-
-
-@jax.custom_jvp
+# @jax.custom_jvp
 @jax.jit
 def lines_func(energy, kT, abundances, broadening):
     idx, kT_low, kT_high = get_temperature(kT)
     line_low = get_lines_contribution_broadening(*get_lines(idx), energy, abundances, broadening)
-    line_high = get_lines_contribution_broadening(*get_lines(idx + 1), energy, abundances, broadening)
+    line_high = get_lines_contribution_broadening(
+        *get_lines(idx + 1), energy, abundances, broadening
+    )
 
     return lerp(kT, kT_low, kT_high, line_low, line_high)
 
 
-@jax.jit
-@lines_func.defjvp
-def lines_jvp(primals, tangents):
-    energy, kT, abundances, broadening = primals
-    energy_dot, kT_dot, abundances_dot, broadening_dot = tangents
-
-    primals_out = lines_func(*primals)
-    return primals_out, jnp.zeros_like(primals_out)
-
-
 class APEC(AdditiveComponent):
     """
     APEC model implementation in pure JAX for X-ray spectral fitting.
@@ -276,11 +203,15 @@ class APEC(AdditiveComponent):
         thermal_broadening: bool = True,
         turbulent_broadening: bool = True,
         variant: Literal["none", "v", "vv"] = "none",
-        abundance_table: Literal["angr", "aspl", "feld", "aneb", "grsa", "wilm", "lodd", "lgpp", "lgps"] = "angr",
+        abundance_table: Literal[
+            "angr", "aspl", "feld", "aneb", "grsa", "wilm", "lodd", "lgpp", "lgps"
+        ] = "angr",
         trace_abundance: float = 1.0,
         **kwargs,
     ):
-        super(APEC, self).__init__(**kwargs)
+        super().__init__(**kwargs)
+
+        warnings.warn("Be aware that this APEC implementation is not meant to be used yet")
 
         self.atomic_weights = jnp.asarray(element_data["atomic_weight"].to_numpy())
 
@@ -325,14 +256,18 @@ class APEC(AdditiveComponent):
         """
         if self.turbulent_broadening:
             # This return value must be multiplied by the energy of the line to get actual broadening
-            return hk.get_parameter("Velocity", [], init=HaikuConstant(100.0)) / c.to(u.km / u.s).value
+            return (
+                hk.get_parameter("Velocity", [], init=HaikuConstant(100.0)) / c.to(u.km / u.s).value
+            )
         else:
             return 0.0
 
     def get_parameters(self):
         none_elements = ["C", "N", "O", "Ne", "Mg", "Al", "Si", "S", "Ar", "Ca", "Fe", "Ni"]
         v_elements = ["He", "C", "N", "O", "Ne", "Mg", "Al", "Si", "S", "Ar", "Ca", "Fe", "Ni"]
-        trace_elements = jnp.asarray([3, 4, 5, 9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30], dtype=int) - 1
+        trace_elements = (
+            jnp.asarray([3, 4, 5, 9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30], dtype=int) - 1
+        )
 
         # Set abundances of trace element (will be overwritten in the vv case)
         abund = jnp.ones((30,)).at[trace_elements].multiply(self.trace_abundance)
@@ -354,7 +289,9 @@ class APEC(AdditiveComponent):
                     abund = abund.at[i].set(Z)
 
         if abund != "angr":
-            abund = abund * jnp.asarray(abundance_table[self.abundance_table] / abundance_table["angr"])
+            abund = abund * jnp.asarray(
+                abundance_table[self.abundance_table] / abundance_table["angr"]
+            )
 
         # Set the temperature, redshift, normalisation
         kT = hk.get_parameter("kT", [], init=HaikuConstant(6.5))
@@ -372,6 +309,8 @@ class APEC(AdditiveComponent):
 
         continuum = continuum_func(energy, kT, abundances) if self.continuum_to_compute else 0.0
         pseudo_continuum = pseudo_func(energy, kT, abundances) if self.pseudo_to_compute else 0.0
-        lines = lines_func(energy, kT, abundances, total_broadening) if self.lines_to_compute else 0.0
+        lines = (
+            lines_func(energy, kT, abundances, total_broadening) if self.lines_to_compute else 0.0
+        )
 
         return (continuum + pseudo_continuum + lines) * norm * 1e14 / (1 + z), (e_low + e_high) / 2

jaxspec/model/_additive/apec_loaders.py

@@ -1,90 +1,73 @@
-""" This module contains the functions that load the APEC tables from the HDF5 file. They are implemented as JAX
-pure callback to enable reading data from the files without saturating the memory. """
+"""This module contains the functions that load the APEC tables from the HDF5 file. They are implemented as JAX
+pure callback to enable reading data from the files without saturating the memory."""
 
+import h5netcdf
 import jax
 import jax.numpy as jnp
-import numpy as np
-import importlib.resources
-import xarray as xr
 
+from ...util.online_storage import table_manager
 
-apec_file = xr.open_dataset(importlib.resources.files("jaxspec") / "tables/apec.nc", engine="h5netcdf")
 
+@jax.jit
+def temperature_table_getter():
+    with h5netcdf.File(table_manager.fetch("apec.nc"), "r") as f:
+        temperature = jnp.asarray(f["/temperature"])
 
-def temperature_table_getter(kT):
-    idx = np.searchsorted(apec_file.temperature, kT) - 1
-
-    if idx > len(apec_file.temperature) - 2:
-        return idx, 0.0, 0.0
-    else:
-        return idx, float(apec_file.temperature[idx]), float(apec_file.temperature[idx + 1])
-
-
-def continuum_table_getter(idx):
-    if idx > len(apec_file.temperature) - 2:
-        continuum_energy_array = jnp.zeros_like(apec_file.continuum_energy[0])
-        continuum_emissivity_array = jnp.zeros_like(apec_file.continuum_emissivity[0])
-        end_index_continuum = jnp.zeros_like(apec_file.continuum_end_index[0])
-
-    else:
-        continuum_energy_array = jnp.asarray(apec_file.continuum_energy[idx])
-        continuum_emissivity_array = jnp.asarray(apec_file.continuum_emissivity[idx])
-        end_index_continuum = jnp.asarray(apec_file.continuum_end_index[idx])
-
-    return continuum_energy_array, continuum_emissivity_array, end_index_continuum
-
+    return temperature
 
-def pseudo_table_getter(idx):
-    if idx > len(apec_file.temperature) - 2:
-        pseudo_energy_array = jnp.zeros_like(apec_file.pseudo_energy[0])
-        pseudo_emissivity_array = jnp.zeros_like(apec_file.pseudo_emissivity[0])
-        end_index_pseudo = jnp.zeros_like(apec_file.pseudo_end_index[0])
 
-    else:
-        pseudo_energy_array = jnp.asarray(apec_file.pseudo_energy[idx])
-        pseudo_emissivity_array = jnp.asarray(apec_file.pseudo_emissivity[idx])
-        end_index_pseudo = jnp.asarray(apec_file.pseudo_end_index[idx])
+@jax.jit
+def get_temperature(kT):
+    temperature = temperature_table_getter()
+    idx = jnp.searchsorted(temperature, kT) - 1
 
-    return pseudo_energy_array, pseudo_emissivity_array, end_index_pseudo
+    return idx, temperature[idx], temperature[idx + 1]
 
 
-def lines_table_getter(idx):
-    if idx > len(apec_file.temperature) - 2:
-        line_energy_array = jnp.zeros_like(apec_file.line_energy[0])
-        line_element_array = jnp.zeros_like(apec_file.line_element[0])
-        line_emissivity_array = jnp.zeros_like(apec_file.line_emissivity[0])
-        end_index_lines = jnp.zeros_like(apec_file.line_end_index[0])
+@jax.jit
+def continuum_table_getter():
+    with h5netcdf.File(table_manager.fetch("apec.nc"), "r") as f:
+        continuum_energy = jnp.asarray(f["/continuum_energy"])
+        continuum_emissivity = jnp.asarray(f["/continuum_emissivity"])
+        continuum_end_index = jnp.asarray(f["/continuum_end_index"])
 
-    else:
-        line_energy_array = jnp.asarray(apec_file.line_energy[idx])
-        line_element_array = jnp.asarray(apec_file.line_element[idx])
-        line_emissivity_array = jnp.asarray(apec_file.line_emissivity[idx])
-        end_index_lines = jnp.asarray(apec_file.line_end_index[idx])
+    return continuum_energy, continuum_emissivity, continuum_end_index
 
-    return line_energy_array, line_element_array, line_emissivity_array, end_index_lines
 
+@jax.jit
+def pseudo_table_getter():
+    with h5netcdf.File(table_manager.fetch("apec.nc"), "r") as f:
+        pseudo_energy = jnp.asarray(f["/pseudo_energy"])
+        pseudo_emissivity = jnp.asarray(f["/pseudo_emissivity"])
+        pseudo_end_index = jnp.asarray(f["/pseudo_end_index"])
 
-pure_callback_temperature_shape = jax.eval_shape(lambda: jax.tree.map(jnp.asarray, temperature_table_getter(10.0)))
-pure_callback_continuum_shape = jax.eval_shape(lambda: jax.tree.map(jnp.asarray, continuum_table_getter(0)))
-pure_callback_pseudo_shape = jax.eval_shape(lambda: jax.tree.map(jnp.asarray, pseudo_table_getter(0)))
-pure_callback_line_shape = jax.eval_shape(lambda: jax.tree.map(jnp.asarray, lines_table_getter(0)))
+    return pseudo_energy, pseudo_emissivity, pseudo_end_index
 
 
 @jax.jit
-def get_temperature(kT):
-    return jax.pure_callback(temperature_table_getter, pure_callback_temperature_shape, kT)
+def line_table_getter():
+    with h5netcdf.File(table_manager.fetch("apec.nc"), "r") as f:
+        line_energy = jnp.asarray(f["/line_energy"])
+        line_element = jnp.asarray(f["/line_element"])
+        line_emissivity = jnp.asarray(f["/line_emissivity"])
+        line_end_index = jnp.asarray(f["/line_end_index"])
+
+    return line_energy, line_element, line_emissivity, line_end_index
 
 
 @jax.jit
 def get_continuum(idx):
-    return jax.pure_callback(continuum_table_getter, pure_callback_continuum_shape, idx)
+    continuum_energy, continuum_emissivity, continuum_end_index = continuum_table_getter()
+    return continuum_energy[idx], continuum_emissivity[idx], continuum_end_index[idx]
 
 
 @jax.jit
 def get_pseudo(idx):
-    return jax.pure_callback(pseudo_table_getter, pure_callback_pseudo_shape, idx)
+    pseudo_energy, pseudo_emissivity, pseudo_end_index = pseudo_table_getter()
+    return pseudo_energy[idx], pseudo_emissivity[idx], pseudo_end_index[idx]
 
 
 @jax.jit
 def get_lines(idx):
-    return jax.pure_callback(lines_table_getter, pure_callback_line_shape, idx)
+    line_energy, line_element, line_emissivity, line_end_index = line_table_getter()
+    return line_energy[idx], line_element[idx], line_emissivity[idx], line_end_index[idx]

jaxspec/model/additive.py

@@ -1,16 +1,18 @@
 from __future__ import annotations
 
-
+import astropy.constants
+import astropy.units as u
 import haiku as hk
 import jax
 import jax.numpy as jnp
 import jax.scipy as jsp
-import astropy.units as u
-import astropy.constants
+
 from haiku.initializers import Constant as HaikuConstant
+
 from ..util.integrate import integrate_interval
+
+# from ._additive.apec import APEC
 from .abc import AdditiveComponent
-from ._additive.apec import APEC  # noqa: F401
 
 
 class Powerlaw(AdditiveComponent):
@@ -283,16 +285,16 @@ 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):
@@ -332,7 +334,13 @@ class Diskbb(AdditiveComponent):
             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):
@@ -381,7 +389,11 @@ class Zagauss(AdditiveComponent):
         norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
         redshift = hk.get_parameter("redshift", [], 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):
@@ -404,4 +416,6 @@ class Zgauss(AdditiveComponent):
         norm = hk.get_parameter("norm", [], init=HaikuConstant(1))
         redshift = hk.get_parameter("redshift", [], init=HaikuConstant(0))
 
-        return (norm / (1 + redshift)) * jsp.stats.norm.pdf(energy * (1 + redshift), loc=line_energy, scale=sigma)
+        return (norm / (1 + redshift)) * jsp.stats.norm.pdf(
+            energy * (1 + redshift), loc=line_energy, scale=sigma
+        )

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