jaxspec 0.1.3__py3-none-any.whl → 0.2.0__py3-none-any.whl

jaxspec/model/background.py

@@ -1,22 +1,27 @@
 from abc import ABC, abstractmethod
 
+import jax
 import jax.numpy as jnp
 import numpyro
 import numpyro.distributions as dist
 
 from jax.scipy.integrate import trapezoid
+from numpyro.distributions import Poisson
 from tinygp import GaussianProcess, kernels
 
+from .._fit._build_model import build_prior, forward_model
+from .abc import SpectralModel
+
 
 class BackgroundModel(ABC):
     """
-    This class handles the modelization of backgrounds in our spectra. This is handled in a separate class for now
+    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, energy, observed_counts, name: str = "bkg", observed=True):
+    def numpyro_model(self, observation, name: str = "", observed=True):
         """
         Build the model for the background.
         """
@@ -25,7 +30,7 @@ class BackgroundModel(ABC):
 
 class SubtractedBackground(BackgroundModel):
     """
-    This class is to use when implying that the observed background should be simply subtracted from the observed.
+    Define a model where the observed background is simply subtracted from the observed.
 
     !!! danger
 
@@ -35,93 +40,40 @@ class SubtractedBackground(BackgroundModel):
 
     """
 
-    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)
+    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 jnp.zeros_like(observed_counts)
+        return observed_counts
 
 
 class BackgroundWithError(BackgroundModel):
     """
-    This class is to use when implying that the observed background should be simply subtracted from the observed. It
+    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.
-
-    !!! warning
-        This is the same as [`ConjugateBackground`][jaxspec.model.background.ConjugateBackground]
-        but slower since it performs the fit using MCMC instead of analytical solution.
     """
 
-    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
+    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"{name}_params", dist.Gamma(alpha, rate=beta))
+        countrate = numpyro.sample(f"_bkg_{name}_countrate", dist.Gamma(alpha, rate=beta))
 
-        with numpyro.plate(f"{name}_plate", len(observed_counts)):
+        with numpyro.plate(f"bkg_{name}_plate", len(observed_counts)):
             numpyro.sample(
-                f"{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
+                f"bkg_{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
-    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}
-    \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
-        counts are 0, and add a small scatter even if the measured background is effectively null.
-
-    ??? abstract "References"
-
-        - https://en.wikipedia.org/wiki/Conjugate_prior
-        - https://www.acsu.buffalo.edu/~adamcunn/probability/gamma.html
-        - https://bayesiancomputationbook.com/markdown/chp_01.html?highlight=conjugate#conjugate-priors
-        - https://vioshyvo.github.io/Bayesian_inference/conjugate-distributions.html
-
-    """
-
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
-        # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)
-        # alpha = observed_counts + 1
-        # beta = 1
-
-        with numpyro.plate(f"{name}_plate", len(observed_counts)):
-            countrate = numpyro.sample(f"{name}", dist.Gamma(2 * observed_counts + 1, 2), obs=None)
-
-        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):
     """
-    This class use a Gaussian Process to model the background. The GP is built using the
+    Define a Gaussian Process to model the background. The GP is built using the
     [`tinygp`](https://tinygp.readthedocs.io/en/stable/guide.html) library.
     """
 
@@ -146,16 +98,7 @@ class GaussianProcessBackground(BackgroundModel):
         self.n_nodes = n_nodes
         self.kernel = kernel
 
-    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
-        """
-        Build the model for the background.
-
-        Parameters:
-            energy: The energy bins lower and upper values (e_low, e_high).
-            observed_counts: The observed counts in each energy bin.
-            name: The name of the background model for parameters disambiguation.
-            observed: Whether the model is observed or not. Useful for `numpyro.infer.Predictive` calls.
-        """
+    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)):
@@ -163,28 +106,92 @@ class GaussianProcessBackground(BackgroundModel):
                 "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"{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0))
-        sigma = numpyro.sample(f"{name}_sigma", dist.HalfNormal(3.0))
-        rho = numpyro.sample(f"{name}_rho", dist.HalfNormal(10.0))
+        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"_{name}_log_rate_nodes", gp.numpyro_dist())
+        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(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
+        with numpyro.plate("bkg_plate_" + name, len(observed_counts)):
             numpyro.sample(
-                f"{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
+                f"bkg_{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
             )
 
         return count_rate
+
+
+class SpectralModelBackground(BackgroundModel):
+    def __init__(self, spectral_model: "SpectralModel", prior_distributions, sparse=False):
+        self.spectral_model = spectral_model
+        self.prior = prior_distributions
+        self.sparse = sparse
+
+    def numpyro_model(self, observation, name: str = "", observed=True):
+        params = build_prior(self.prior, prefix=f"_bkg_{name}_")
+        bkg_model = jax.jit(
+            lambda par: forward_model(self.spectral_model, par, observation, sparse=self.sparse)
+        )
+        bkg_countrate = bkg_model(params)
+
+        with numpyro.plate("bkg_plate_" + name, len(observation.folded_background)):
+            numpyro.sample(
+                "bkg_" + name,
+                Poisson(bkg_countrate),
+                obs=observation.folded_background.data if observed else None,
+            )
+
+        return bkg_countrate
+
+
+'''
+class ConjugateBackground(BackgroundModel):
+    r"""
+    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}
+    \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
+        counts are 0, and add a small scatter even if the measured background is effectively null.
+
+    ??? abstract "References"
+
+        - https://en.wikipedia.org/wiki/Conjugate_prior
+        - https://www.acsu.buffalo.edu/~adamcunn/probability/gamma.html
+        - https://bayesiancomputationbook.com/markdown/chp_01.html?highlight=conjugate#conjugate-priors
+        - https://vioshyvo.github.io/Bayesian_inference/conjugate-distributions.html
+
+    """
+
+    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+        # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)
+        # alpha = observed_counts + 1
+        # beta = 1
+
+        with numpyro.plate(f"{name}_plate", len(observed_counts)):
+            countrate = numpyro.sample(f"{name}", dist.Gamma(2 * observed_counts + 1, 2), obs=None)
+
+        return countrate
+'''

jaxspec/model/multiplicative.py

@@ -1,16 +1,31 @@
 from __future__ import annotations
 
-import haiku as hk
+import flax.nnx as nnx
 import jax.numpy as jnp
 import numpy as np
 
 from astropy.table import Table
-from haiku.initializers import Constant as HaikuConstant
 
 from ..util.online_storage import table_manager
 from .abc import MultiplicativeComponent
 
 
+class MultiplicativeConstant(MultiplicativeComponent):
+    r"""
+    A multiplicative constant
+
+    !!! abstract "Parameters"
+        * $K$ (`norm`) $\left[\text{dimensionless}\right]$: The multiplicative constant.
+
+    """
+
+    def __init__(self):
+        self.norm = nnx.Param(1.0)
+
+    def factor(self, energy):
+        return self.norm
+
+
 class Expfac(MultiplicativeComponent):
     r"""
     An exponential modification of a spectrum.
@@ -20,19 +35,20 @@ class Expfac(MultiplicativeComponent):
     \text{if $E>E_c$}\\1 & \text{if $E<E_c$}\end{cases}
     $$
 
-    ??? abstract "Parameters"
-        * $A$ : Amplitude of the modification $\left[\text{dimensionless}\right]$
-        * $f$ : Exponential factor $\left[\text{keV}^{-1}\right]$
-        * $E_c$ : Start energy of modification $\left[\text{keV}\right]$
+    !!! abstract "Parameters"
+        * $A$ (`A`) $\left[\text{dimensionless}\right]$ : Amplitude of the modification
+        * $f$ (`f`) $\left[\text{keV}^{-1}\right]$ : Exponential factor
+        * $E_c$ (`E_c`) $\left[\text{keV}\right]$: Start energy of modification
 
     """
 
-    def continuum(self, energy):
-        amplitude = hk.get_parameter("A", [], float, init=HaikuConstant(1))
-        factor = hk.get_parameter("f", [], float, init=HaikuConstant(1))
-        pivot = hk.get_parameter("E_c", [], float, init=HaikuConstant(1))
+    def __init__(self):
+        self.A = nnx.Param(1.0)
+        self.f = nnx.Param(1.0)
+        self.E_c = nnx.Param(1.0)
 
-        return jnp.where(energy >= pivot, 1.0 + amplitude * jnp.exp(-factor * energy), 1.0)
+    def factor(self, energy):
+        return jnp.where(energy >= self.E_c, 1.0 + self.A * jnp.exp(-self.f * energy), 1.0)
 
 
 class Tbabs(MultiplicativeComponent):
@@ -45,49 +61,47 @@ class Tbabs(MultiplicativeComponent):
         \mathcal{M}(E) = \exp^{-N_{\text{H}}\sigma(E)}
     $$
 
-    ??? abstract "Parameters"
-        * $N_{\text{H}}$ : Equivalent hydrogen column density
-            $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$
+    !!! abstract "Parameters"
+        * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
+
 
     !!! note
         Abundances and cross-sections $\sigma$ can be found in Wilms et al. (2000).
 
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self):
         table = Table.read(table_manager.fetch("xsect_tbabs_wilm.fits"))
-        self.energy = jnp.asarray(np.array(table["ENERGY"]), dtype=np.float64)
-        self.sigma = jnp.asarray(np.array(table["SIGMA"]), dtype=np.float64)
+        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
+        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self.nh = nnx.Param(1.0)
 
-    def continuum(self, energy):
-        nh = hk.get_parameter("N_H", [], float, init=HaikuConstant(1))
+    def factor(self, energy):
         sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
 
-        return jnp.exp(-nh * sigma)
+        return jnp.exp(-self.nh * sigma)
 
 
 class Phabs(MultiplicativeComponent):
     r"""
     A photoelectric absorption model.
 
-    ??? abstract "Parameters"
-        * $N_{\text{H}}$ : Equivalent hydrogen column density
-            $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$
+    !!! abstract "Parameters"
+        * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
+
 
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self):
         table = Table.read(table_manager.fetch("xsect_phabs_aspl.fits"))
-        self.energy = jnp.asarray(np.array(table["ENERGY"]), dtype=np.float64)
-        self.sigma = jnp.asarray(np.array(table["SIGMA"]), dtype=np.float64)
+        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
+        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self.nh = nnx.Param(1.0)
 
-    def continuum(self, energy):
-        nh = hk.get_parameter("N_H", [], float, init=HaikuConstant(1))
-        sigma = jnp.interp(energy, self.energy, self.sigma, left=jnp.inf, right=0.0)
+    def factor(self, energy):
+        sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
 
-        return jnp.exp(-nh * sigma)
+        return jnp.exp(-self.nh * sigma)
 
 
 class Wabs(MultiplicativeComponent):
@@ -95,22 +109,19 @@ class Wabs(MultiplicativeComponent):
     A photo-electric absorption using Wisconsin (Morrison & McCammon 1983) cross-sections.
 
     ??? abstract "Parameters"
-        * $N_{\text{H}}$ : Equivalent hydrogen column density
-            $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$
-
+        * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self):
         table = Table.read(table_manager.fetch("xsect_wabs_angr.fits"))
-        self.energy = jnp.asarray(np.array(table["ENERGY"]), dtype=np.float64)
-        self.sigma = jnp.asarray(np.array(table["SIGMA"]), dtype=np.float64)
+        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
+        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self.nh = nnx.Param(1.0)
 
-    def continuum(self, energy):
-        nh = hk.get_parameter("N_H", [], float, init=HaikuConstant(1))
+    def factor(self, energy):
         sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
 
-        return jnp.exp(-nh * sigma)
+        return jnp.exp(-self.nh * sigma)
 
 
 class Gabs(MultiplicativeComponent):
@@ -122,23 +133,26 @@ class Gabs(MultiplicativeComponent):
         \left( -\frac{\left(E-E_0\right)^2}{2 \sigma^2} \right) \right)
     $$
 
-    ??? abstract "Parameters"
-        * $\tau$ : Absorption strength $\left[\text{dimensionless}\right]$
-        * $\sigma$ : Absorption width $\left[\text{keV}\right]$
-        * $E_0$ : Absorption center $\left[\text{keV}\right]$
+    !!! abstract "Parameters"
+        * $\tau$ (`tau`) $\left[\text{dimensionless}\right]$ : Absorption strength
+        * $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : Absorption width
+        * $E_0$ (`E0`) $\left[\text{keV}\right]$ : Absorption center
 
     !!! note
         The optical depth at line center is $\tau/(\sqrt{2 \pi} \sigma)$.
 
     """
 
-    def continuum(self, energy):
-        tau = hk.get_parameter("tau", [], float, init=HaikuConstant(1))
-        sigma = hk.get_parameter("sigma", [], float, init=HaikuConstant(1))
-        center = hk.get_parameter("E_0", [], float, init=HaikuConstant(1))
+    def __init__(self):
+        self.tau = nnx.Param(1.0)
+        self.sigma = nnx.Param(1.0)
+        self.E0 = nnx.Param(1.0)
 
+    def factor(self, energy):
         return jnp.exp(
-            -tau / (jnp.sqrt(2 * jnp.pi) * sigma) * jnp.exp(-0.5 * ((energy - center) / sigma) ** 2)
+            -self.tau
+            / (jnp.sqrt(2 * jnp.pi) * self.sigma)
+            * jnp.exp(-0.5 * ((energy - self.E0) / self.sigma) ** 2)
         )
 
 
@@ -151,16 +165,17 @@ class Highecut(MultiplicativeComponent):
         \left( \frac{E_c - E}{E_f} \right)& \text{if $E > E_c$}\\ 1 & \text{if $E < E_c$}\end{cases}
     $$
 
-    ??? abstract "Parameters"
-        * $E_c$ : Cutoff energy $\left[\text{keV}\right]$
-        * $E_f$ : Folding energy $\left[\text{keV}\right]$
+    !!! abstract "Parameters"
+        * $E_c$ (`Ec`) $\left[\text{keV}\right]$ : Cutoff energy
+        * $E_f$ (`Ef`) $\left[\text{keV}\right]$ : Folding energy
     """
 
-    def continuum(self, energy):
-        cutoff = hk.get_parameter("E_c", [], float, init=HaikuConstant(1))
-        folding = hk.get_parameter("E_f", [], float, init=HaikuConstant(1))
+    def __init__(self):
+        self.Ec = nnx.Param(1.0)
+        self.Ef = nnx.Param(1.0)
 
-        return jnp.where(energy <= cutoff, 1.0, jnp.exp((cutoff - energy) / folding))
+    def factor(self, energy):
+        return jnp.where(energy <= self.Ec, 1.0, jnp.exp((self.Ec - energy) / self.Ef))
 
 
 class Zedge(MultiplicativeComponent):
@@ -172,18 +187,21 @@ class Zedge(MultiplicativeComponent):
         & \text{if $E > E_c$}\\ 1 & \text{if $E < E_c$}\end{cases}
     $$
 
-    ??? abstract "Parameters"
-        * $E_c$ : Threshold energy
-        * $E_f$ : Absorption depth at the threshold
-        * $z$ : Redshift [dimensionless]
+    !!! abstract "Parameters"
+        * $E_c$ (`Ec`) $\left[\text{keV}\right]$ : Threshold energy
+        * $D$ (`D`) $\left[\text{dimensionless}\right]$ : Absorption depth at the threshold
+        * $z$ (`z`) $\left[\text{dimensionless}\right]$ : Redshift
     """
 
-    def continuum(self, energy):
-        E_c = hk.get_parameter("E_c", [], float, init=HaikuConstant(1))
-        D = hk.get_parameter("D", [], float, init=HaikuConstant(1))
-        z = hk.get_parameter("z", [], float, init=HaikuConstant(0))
+    def __init__(self):
+        self.Ec = nnx.Param(1.0)
+        self.D = nnx.Param(1.0)
+        self.z = nnx.Param(0.0)
 
-        return jnp.where(energy <= E_c, 1.0, jnp.exp(-D * (energy * (1 + z) / E_c) ** 3))
+    def factor(self, energy):
+        return jnp.where(
+            energy <= self.Ec, 1.0, jnp.exp(-self.D * (energy * (1 + self.z) / self.Ec) ** 3)
+        )
 
 
 class Tbpcf(MultiplicativeComponent):
@@ -194,28 +212,25 @@ class Tbpcf(MultiplicativeComponent):
         \mathcal{M}(E) = f \exp^{-N_{\text{H}}\sigma(E)} + (1-f)
     $$
 
-    ??? abstract "Parameters"
-        * $N_{\text{H}}$ : Equivalent hydrogen column density
-            $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$
-        * $f$ : Partial covering fraction, ranges between 0 and 1 [dimensionless]
+    !!! abstract "Parameters"
+        * $N_{\text{H}}$ (`nh`) $\left[\frac{\text{atoms}~10^{22}}{\text{cm}^2}\right]$ : Equivalent hydrogen column density
+        * $f$ (`f`) $\left[\text{dimensionless}\right]$ : Partial covering fraction, ranges between 0 and 1
 
     !!! note
         Abundances and cross-sections $\sigma$ can be found in Wilms et al. (2000).
 
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self):
         table = Table.read(table_manager.fetch("xsect_tbabs_wilm.fits"))
-        self.energy = jnp.asarray(np.array(table["ENERGY"]), dtype=np.float64)
-        self.sigma = jnp.asarray(np.array(table["SIGMA"]), dtype=np.float64)
+        self.energy = nnx.Variable(np.asarray(table["ENERGY"], dtype=np.float64))
+        self.sigma = nnx.Variable(np.asarray(table["SIGMA"], dtype=np.float64))
+        self.nh = nnx.Param(1.0)
+        self.f = nnx.Param(0.2)
 
     def continuum(self, energy):
-        nh = hk.get_parameter("N_H", [], float, init=HaikuConstant(1))
-        f = hk.get_parameter("f", [], float, init=HaikuConstant(0.2))
         sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
-
-        return f * jnp.exp(-nh * sigma) + (1 - f)
+        return self.f * jnp.exp(-self.nh * sigma) + (1 - self.f)
 
 
 class FDcut(MultiplicativeComponent):
@@ -227,12 +242,13 @@ class FDcut(MultiplicativeComponent):
     $$
 
     ??? abstract "Parameters"
-        * $E_c$ : Cutoff energy $\left[\text{keV}\right]$
-        * $E_f$ : Folding energy $\left[\text{keV}\right]$
+        * $E_c$ (`Ec`) $\left[\text{keV}\right]$ : Cutoff energy
+        * $E_f$ (`Ef`) $\left[\text{keV}\right]$ : Folding energy
     """
 
-    def continuum(self, energy):
-        cutoff = hk.get_parameter("E_c", [], init=HaikuConstant(1))
-        folding = hk.get_parameter("E_f", [], init=HaikuConstant(1))
+    def __init__(self):
+        self.Ec = nnx.Param(1.0)
+        self.Ef = nnx.Param(3.0)
 
-        return (1 + jnp.exp((energy - cutoff) / folding)) ** -1
+    def continuum(self, energy):
+        return (1 + jnp.exp((energy - self.Ec) / self.Ef)) ** -1

jaxspec/scripts/debug.py

@@ -8,4 +8,4 @@ def debug_info():
 
     # Cimer CamilleTheBest pour l'idée
     print(watermark())
-    print(watermark(packages="jaxspec,jax,jaxlib,numpyro,haiku,numpy,scipy"))
+    print(watermark(packages="jaxspec,jax,jaxlib,numpyro,flax,numpy,scipy"))

jaxspec/util/__init__.py

@@ -1,45 +0,0 @@
-from collections.abc import Callable
-from contextlib import contextmanager
-from time import perf_counter
-
-import haiku as hk
-
-from jax.random import PRNGKey, split
-
-
-@contextmanager
-def catchtime(desc="Task", print_time=True) -> Callable[[], float]:
-    """
-    Context manager to measure time taken by a task.
-
-    Parameters
-    ----------
-        desc (str): Description of the task.
-        print_time (bool): Whether to print the time taken by the task.
-
-    Returns
-    -------
-        Callable[[], float]: Function to get the time taken by the task.
-    """
-
-    t1 = t2 = perf_counter()
-    yield lambda: t2 - t1
-    t2 = perf_counter()
-    if print_time:
-        print(f"{desc}: {t2 - t1:.3f} seconds")
-
-
-def sample_prior(dict_of_prior, key=PRNGKey(0), flat_parameters=False):
-    """
-    Sample the prior distribution from a dict of distributions
-    """
-
-    parameters = dict(hk.data_structures.to_haiku_dict(dict_of_prior))
-    parameters_flat = {}
-
-    for m, n, distribution in hk.data_structures.traverse(dict_of_prior):
-        key, subkey = split(key)
-        parameters[m][n] = distribution.sample(subkey)
-        parameters_flat[m + "_" + n] = distribution.sample(subkey)
-
-    return parameters if not flat_parameters else parameters_flat

jaxspec/util/misc.py

@@ -0,0 +1,25 @@
+from collections.abc import Callable
+from contextlib import contextmanager
+from time import perf_counter
+
+
+@contextmanager
+def catchtime(desc="Task", print_time=True) -> Callable[[], float]:
+    """
+    Context manager to measure time taken by a task.
+
+    Parameters
+    ----------
+        desc (str): Description of the task.
+        print_time (bool): Whether to print the time taken by the task.
+
+    Returns:
+    -------
+        Callable[[], float]: Function to get the time taken by the task.
+    """
+
+    t1 = t2 = perf_counter()
+    yield lambda: t2 - t1
+    t2 = perf_counter()
+    if print_time:
+        print(f"{desc}: {t2 - t1:.3f} seconds")