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

jaxspec/model/additive.py

@@ -1,8 +1,10 @@
 from __future__ import annotations
 
+from functools import partial
+
 import astropy.constants
 import astropy.units as u
-import haiku as hk
+import flax.nnx as nnx
 import interpax
 import jax
 import jax.numpy as jnp
@@ -10,7 +12,6 @@ import jax.scipy as jsp
 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
@@ -23,17 +24,18 @@ class Powerlaw(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = K \left( \frac{E}{E_0} \right)^{-\alpha}$$
 
-    ??? abstract "Parameters"
-        * $\alpha$ : Photon index of the power law $\left[\text{dimensionless}\right]$
-        * $E_0$ : Reference energy fixed at 1 keV $\left[ \mathrm{keV}\right]$
-        * $K$ : Normalization at the reference energy (1 keV) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $\alpha$ (`alpha`) $\left[\text{dimensionless}\right]$ : Photon index of the power law
+        * $E_0$ $\left[ \mathrm{keV}\right]$ : Reference energy fixed at 1 keV
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization at the reference energy (1 keV)
     """
 
-    def continuum(self, energy):
-        alpha = hk.get_parameter("alpha", [], float, init=HaikuConstant(1.3))
-        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1e-4))
+    def __init__(self):
+        self.alpha = nnx.Param(1.7)
+        self.norm = nnx.Param(1e-4)
 
-        return norm * energy ** (-alpha)
+    def continuum(self, energy):
+        return self.norm * energy ** (-self.alpha)
 
 
 class Additiveconstant(AdditiveComponent):
@@ -42,19 +44,15 @@ class Additiveconstant(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = K$$
 
-    ??? abstract "Parameters"
-        * $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
-    def continuum(self, energy):
-        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
-
-        return norm * jnp.ones_like(energy)
-
-    def primitive(self, energy):
-        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+    def __init__(self):
+        self.norm = nnx.Param(1.0)
 
-        return norm * energy
+    def integrated_continuum(self, e_low, e_high):
+        return (e_high - e_low) * self.norm
 
 
 class Lorentz(AdditiveComponent):
@@ -63,18 +61,24 @@ class Lorentz(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = K\frac{\frac{\sigma}{2\pi}}{(E-E_L)^2 + \left(\frac{\sigma}{2}\right)^2}$$
 
-    ??? abstract "Parameters"
-        - $E_L$ : Energy of the line $\left[\text{keV}\right]$
-        - $\sigma$ : FWHM of the line $\left[\text{keV}\right]$
-        - $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        - $E_L$ (`E_l`) $\left[\text{keV}\right]$ : Energy of the line
+        - $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : FWHM of the line
+        - $K$ (`norm`) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
-    def continuum(self, energy):
-        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))
+    def __init__(self):
+        self.E_l = jnp.asarray(nnx.Param(1.0), dtype=jnp.float64)
+        self.sigma = jnp.asarray(nnx.Param(1e-3), dtype=jnp.float64)
+        self.norm = jnp.asarray(nnx.Param(1.0), dtype=jnp.float64)
 
-        return norm * sigma / (2 * jnp.pi) / ((energy - line_energy) ** 2 + (sigma / 2) ** 2)
+    def continuum(self, energy):
+        return (
+            self.norm
+            * self.sigma
+            / (2 * jnp.pi)
+            / ((energy - self.E_l) ** 2 + (self.sigma / 2) ** 2)
+        )
 
 
 class Logparabola(AdditiveComponent):
@@ -83,23 +87,23 @@ class Logparabola(AdditiveComponent):
 
     $$
     \mathcal{M}\left( E \right) = K  \left( \frac{E}{E_{\text{Pivot}}} \right)
-    ^{-(\alpha + \beta ~ \log (E/E_{\text{Pivot}})) }
+    ^{-(\alpha - \beta ~ \log (E/E_{\text{Pivot}})) }
     $$
 
-    ??? abstract "Parameters"
-        * $a$ : Slope of the LogParabola at the pivot energy $\left[\text{dimensionless}\right]$
-        * $b$ : Curve parameter of the LogParabola $\left[\text{dimensionless}\right]$
-        * $E_{\text{Pivot}}$ : Pivot energy fixed at 1 keV $\left[ \mathrm{keV}\right]$
-        * $K$ : Normalization at the pivot energy (1keV) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $a$ (`a`) $\left[\text{dimensionless}\right]$ : Slope of the LogParabola at the pivot energy
+        * $b$ (`b`) $\left[\text{dimensionless}\right]$ : Curve parameter of the LogParabola
+        * $E_{\text{Pivot}}$ $\left[ \mathrm{keV}\right]$ : Pivot energy fixed at 1 keV
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
-    # TODO : conform with xspec definition
-    def continuum(self, energy):
-        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))
+    def __init__(self):
+        self.a = nnx.Param(1.0)
+        self.b = nnx.Param(1.0)
+        self.norm = nnx.Param(1.0)
 
-        return norm * energy ** (-(a + b * jnp.log(energy)))
+    def continuum(self, energy):
+        return self.norm * energy ** (-(self.a - self.b * jnp.log(energy)))
 
 
 class Blackbody(AdditiveComponent):
@@ -108,17 +112,20 @@ class Blackbody(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = \frac{K \times 8.0525 E^{2}}{(k_B T)^{4}\left(\exp(E/k_BT)-1\right)}$$
 
-    ??? abstract "Parameters"
-        * $k_B T$ : Temperature $\left[\text{keV}\right]$
-        * $K$ : $L_{39}/D_{10}^{2}$, where $L_{39}$ is the source luminosity in units of $10^{39}$ erg/s
-        and $D_{10}$ is the distance to the source in units of 10 kpc
+    !!! abstract "Parameters"
+        * $k_B T$ (`kT`) $\left[\text{keV}\right]$ : Temperature
+        * $K$ (`norm`) $\left[\text{dimensionless}\right]$ : $L_{39}/D_{10}^{2}$, where $L_{39}$ is the source luminosity [$10^{39} \frac{\text{erg}}{\text{s}}$]
+        and $D_{10}$ is the distance to the source [$10 \text{kpc}$]
     """
 
-    def continuum(self, energy):
-        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(11 / 3))
-        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+    # TODO : rewrite constant as a astropy unit stuff
 
-        return norm * 8.0525 * energy**2 / ((kT**4) * (jnp.exp(energy / kT) - 1))
+    def __init__(self):
+        self.kT = nnx.Param(0.5)
+        self.norm = nnx.Param(1.0)
+
+    def continuum(self, energy):
+        return self.norm * 8.0525 * energy**2 / ((self.kT**4) * jnp.expm1(energy / self.kT))
 
 
 class Blackbodyrad(AdditiveComponent):
@@ -127,137 +134,52 @@ class Blackbodyrad(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = \frac{K \times 1.0344\times 10^{-3} E^{2}}{\left(\exp (E/k_BT)-1\right)}$$
 
-    ??? abstract "Parameters"
-        * $k_B T$ : Temperature $\left[\text{keV}\right]$
-        * $K$ : $R^2_{km}/D_{10}^{2}$, where $R_{km}$ is the source radius in km
-        and $D_{10}$ is the distance to the source in units of 10 kpc [dimensionless]
+    !!! abstract "Parameters"
+        * $k_B T$ (`kT`) $\left[\text{keV}\right]$ : Temperature
+        * $K$ (`norm`) [dimensionless] : $R^2_{km}/D_{10}^{2}$, where $R_{km}$ is the source radius [$\text{km}$]
+        and $D_{10}$ is the distance to the source [$10 \text{kpc}$]
     """
 
-    def continuum(self, energy):
-        kT = hk.get_parameter("kT", [], float, init=HaikuConstant(11 / 3))
-        norm = hk.get_parameter("norm", [], float, init=HaikuConstant(1))
+    def __init__(self):
+        self.kT = nnx.Param(0.5)
+        self.norm = nnx.Param(1.0)
 
-        return norm * 1.0344e-3 * energy**2 / (jnp.exp(energy / kT) - 1)
+    def continuum(self, energy):
+        return self.norm * 1.0344e-3 * energy**2 / jnp.expm1(energy / self.kT)
 
 
 class Gauss(AdditiveComponent):
     r"""
-    A Gaussian line profile. If the width is $$\leq 0$$ then it is treated as a delta function.
+    A Gaussian line profile. If the width is $\leq 0$ then it is treated as a delta function.
     The `Zgauss` variant computes a redshifted Gaussian.
 
     $$\mathcal{M}\left( E \right) = \frac{K}{\sigma \sqrt{2 \pi}}\exp\left(\frac{-(E-E_L)^2}{2\sigma^2}\right)$$
 
-    ??? abstract "Parameters"
-        * $E_L$ : Energy of the line $\left[\text{keV}\right]$
-        * $\sigma$ : Width of the line $\left[\text{keV}\right]$
-        * $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $E_L$ (`E_l`) $\left[\text{keV}\right]$ : Energy of the line
+        * $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : Width of the line
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
-    def continuum(self, energy) -> (jax.Array, jax.Array):
-        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)
-
-
-"""
-class APEC(AdditiveComponent):
-    def __init__(self, name="apec"):
-        super(APEC, self).__init__(name=name)
-
-        ref = importlib.resources.files("jaxspec") / "tables/apec_tab.npz"
-        with importlib.resources.as_file(ref) as path:
-            files = np.load(path)
-
-        self.kT_ref = files["kT_ref"]
-        self.e_ref = np.nan_to_num(files["continuum_energy"], nan=1e6)
-        self.c_ref = np.nan_to_num(files["continuum_emissivity"])
-        self.pe_ref = np.nan_to_num(files["pseudo_energy"], nan=1e6)
-        self.pc_ref = np.nan_to_num(files["pseudo_emissivity"])
-        self.energy_lines = np.nan_to_num(files["lines_energy"], nan=1e6)  # .astype(np.float32))
-        self.epsilon_lines = np.nan_to_num(files["lines_emissivity"])  # .astype(np.float32))
-        self.element_lines = np.nan_to_num(files["lines_element"])  # .astype(np.int32))
-
-        del files
-
-        self.trace_elements = jnp.array([3, 4, 5, 9, 11, 15, 17, 19, 21, 22, 23, 24, 25, 27, 29, 30]) - 1
-        self.metals = np.array([6, 7, 8, 10, 12, 13, 14, 16, 18, 20, 26, 28]) - 1  # Element number to python index
-        self.metals_one_hot = np.zeros((30,))
-        self.metals_one_hot[self.metals] = 1
-
-    def interp_on_cubes(self, energy, energy_cube, continuum_cube):
-        # Changer en loginterp
-        # Interpoler juste sur les points qui ne sont pas tabulés
-        # Ajouter l'info de la taille max des données (à resortir dans la routine qui trie les fichier apec)
-        return jnp.vectorize(
-            lambda ecube, ccube: jnp.interp(energy, ecube, ccube),
-            signature="(k),(k)->()",
-        )(energy_cube, continuum_cube)
-
-    def reduction_with_elements(self, Z, energy, energy_cube, continuum_cube):
-        return jnp.sum(
-            self.interp_on_cubes(energy, energy_cube, continuum_cube) * jnp.where(self.metals_one_hot, Z, 1.0)[None, :],
-            axis=-1,
-        )
-
-    def mono_fine_structure(self, e_low, e_high) -> (jax.Array, jax.Array):
-        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
-
-        energy = jax_slice(self.energy_lines, idx, 2)
-        epsilon = jax_slice(self.epsilon_lines, idx, 2)
-        element = jax_slice(self.element_lines, idx, 2) - 1
-
-        emissivity_in_bins = jnp.where((e_low < energy) & (energy < e_high), True, False) * epsilon
-        flux_at_edges = jnp.nansum(
-            jnp.where(jnp.isin(element, self.metals), Z, 1) * emissivity_in_bins,
-            axis=-1,
-        )  # Coeff for metallicity
-
-        return (
-            jnp.interp(kT, jax_slice(self.kT_ref, idx, 2), flux_at_edges) * 1e14 * norm,
-            (e_low + e_high) / 2,
-        )
-
-    def emission_lines(self, e_low, e_high) -> (jax.Array, jax.Array):
-        # Compute the fine structure lines with e_low and e_high as array, mapping the mono_fine_structure function
-        # over the various axes of e_low and e_high
-
-        return jnp.vectorize(self.mono_fine_structure)(e_low, e_high)
-
-        # return jnp.zeros_like(e_low), (e_low + e_high)/2
-
-    @partial(jnp.vectorize, excluded=(0,))
-    def continuum(self, energy):
-        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
-
-        continuum = jnp.interp(
-            kT,
-            jax_slice(self.kT_ref, idx, 2),
-            self.reduction_with_elements(Z, energy, jax_slice(self.e_ref, idx, 2), jax_slice(self.c_ref, idx, 2)),
-        )
-        pseudo = jnp.interp(
-            kT,
-            jax_slice(self.kT_ref, idx, 2),
-            self.reduction_with_elements(
-                Z,
-                energy,
-                jax_slice(self.pe_ref, idx, 2),
-                jax_slice(self.pc_ref, idx, 2),
-            ),
+    def __init__(self):
+        self.E_l = nnx.Param(2.0)
+        self.sigma = nnx.Param(1e-2)
+        self.norm = nnx.Param(1.0)
+
+    def integrated_continuum(self, e_low, e_high):
+        return self.norm * (
+            jsp.stats.norm.cdf(
+                e_high,
+                loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+                scale=jnp.asarray(self.sigma, dtype=jnp.float64),
+            )
+            - jsp.stats.norm.cdf(
+                e_low,
+                loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+                scale=jnp.asarray(self.sigma, dtype=jnp.float64),
+            )
         )
 
-        return (continuum + pseudo) * 1e14 * norm
-"""
-
 
 class Cutoffpl(AdditiveComponent):
     r"""
@@ -265,71 +187,40 @@ class Cutoffpl(AdditiveComponent):
 
     $$\mathcal{M}\left( E \right) = K \left( \frac{E}{E_0} \right)^{-\alpha} \exp(-E/\beta)$$
 
-    ??? abstract "Parameters"
-        * $\alpha$ : Photon index of the power law $\left[\text{dimensionless}\right]$
-        * $\beta$ : Folding energy of the exponential cutoff $\left[\text{keV}\right]$
-        * $E_0$ : Reference energy fixed at 1 keV $\left[ \mathrm{keV}\right]$
-        * $K$ : Normalization at the reference energy (1 keV) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $\alpha$ (`alpha`) $\left[\text{dimensionless}\right]$ : Photon index of the power law
+        * $\beta$ (`beta`) $\left[\text{keV}\right]$ : Folding energy of the exponential cutoff
+        * $E_0$ $\left[ \mathrm{keV}\right]$ : Reference energy fixed at 1 keV
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
-    def continuum(self, energy):
-        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)
-
-
-'''
-class Diskpbb(AdditiveComponent):
-    r"""
-    A multiple blackbody disk model where local disk temperature T(r) is proportional to $$r^{(-p)}$$,
-    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}}}$$
-
-    ??? 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]$,
-        $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]$
-    """
+    def __init__(self):
+        self.alpha = nnx.Param(1.7)
+        self.beta = nnx.Param(15.0)
+        self.norm = nnx.Param(1e-4)
 
     def continuum(self, energy):
-        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):
-            return 2.78e-3 * energy**2 * (kT / tin) ** (-2 / p - 1) / (jnp.exp(energy / kT) - 1)
-
-        func_vmapped = jax.vmap(lambda e: integrate_interval(lambda kT: integrand(kT, e), 0, tin, n=51))
-
-        return norm * (0.75 / p) * func_vmapped(energy)
-'''
+        return self.norm * energy ** (-self.alpha) * jnp.exp(-energy / self.beta)
 
 
 class Diskbb(AdditiveComponent):
     r"""
     `Diskpbb` with $p=0.75$
 
-    ??? abstract "Parameters"
-        * $T_{\text{in}}$ : Temperature at inner disk radius $\left[ \mathrm{keV}\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)
+    !!! abstract "Parameters"
+        * $T_{\text{in}}$ (`Tin`) $\left[ \mathrm{keV}\right]$: Temperature at inner disk radius
+        * $\text{norm}$ (`norm`) $\left[\text{dimensionless}\right]$ : $\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 [$10 \text{kpc}$], $i$ the angle of the disk ($i=0$ is face-on)
     """
 
+    def __init__(self):
+        self.Tin = nnx.Param(1.0)
+        self.norm = nnx.Param(1e-4)
+
     def continuum(self, energy):
         p = 0.75
         tout = 0.0
-        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):
@@ -337,11 +228,11 @@ class Diskbb(AdditiveComponent):
 
         integral = integrate_interval(integrand)
         return (
-            norm
+            self.norm
             * 2.78e-3
             * (0.75 / p)
-            / tin
-            * jnp.vectorize(lambda e: integral(tout, tin, e, tin, p))(energy)
+            / self.Tin
+            * jnp.vectorize(lambda e: integral(tout, self.Tin, e, self.Tin, p))(energy)
         )
 
 
@@ -354,19 +245,25 @@ class Agauss(AdditiveComponent):
     $$\mathcal{M}\left( \lambda \right) =
     \frac{K}{\sigma \sqrt{2 \pi}} \exp\left(\frac{-(\lambda - \lambda_L)^2}{2 \sigma^2}\right)$$
 
-    ??? abstract "Parameters"
-        * $\lambda_L$ : Wavelength of the line in Angström $\left[\unicode{x212B}\right]$
-        * $\sigma$ : Width of the line width in Angström  $\left[\unicode{x212B}\right]$
-        * $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
+    !!! abstract "Parameters"
+        * $\lambda_L$ (`lambda_l`) $\left[\unicode{x212B}\right]$ : Wavelength of the line in Angström
+        * $\sigma$ (`sigma`) $\left[\unicode{x212B}\right]$ : Width of the line width in Angström
+        * $K$ (`norm`) $\left[\frac{\unicode{x212B}~\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$: Normalization
     """
 
+    def __init__(self):
+        self.lambda_l = nnx.Param(12.0)
+        self.sigma = nnx.Param(1e-2)
+        self.norm = nnx.Param(1.0)
+
     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", [], 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)
+        return self.norm * jsp.stats.norm.pdf(
+            hc / energy,
+            loc=jnp.asarray(self.lambda_l, dtype=jnp.float64),
+            scale=jnp.asarray(self.sigma, dtype=jnp.float64),
+        )
 
 
 class Zagauss(AdditiveComponent):
@@ -377,24 +274,32 @@ class Zagauss(AdditiveComponent):
     $$\mathcal{M}\left( \lambda \right) =
     \frac{K (1+z)}{\sigma \sqrt{2 \pi}} \exp\left(\frac{-(\lambda/(1+z) - \lambda_L)^2}{2 \sigma^2}\right)$$
 
-    ??? abstract "Parameters"
-        * $\lambda_L$ : Wavelength of the line in Angström $\left[\text{\AA}\right]$
-        * $\sigma$ : Width of the line width in Angström $\left[\text{\AA}\right]$
-        * $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
-        * $z$ : Redshift [dimensionless]
+    !!! abstract "Parameters"
+        * $\lambda_L$ (`lambda_l`) $\left[\unicode{x212B}\right]$  : Wavelength of the line in Angström
+        * $\sigma$ (`sigma`) $\left[\unicode{x212B}\right]$  : Width of the line width in Angström
+        * $z$ (`redshift`) $\left[\text{dimensionless}\right]$ : Redshift
+        * $K$ (`norm`) $\left[\frac{\unicode{x212B}~\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
+    def __init__(self):
+        self.lambda_l = nnx.Param(12.0)
+        self.sigma = nnx.Param(1e-2)
+        self.redshift = nnx.Param(0.0)
+        self.norm = nnx.Param(1.0)
+
     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", [], 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))
+
+        redshift = self.redshift
 
         return (
-            norm
+            self.norm
             * (1 + redshift)
-            * jsp.stats.norm.pdf((hc / energy) / (1 + redshift), loc=line_wavelength, scale=sigma)
+            * jsp.stats.norm.pdf(
+                (hc / energy) / (1 + redshift),
+                loc=jnp.asarray(self.lambda_l, dtype=jnp.float64),
+                scale=jnp.asarray(self.sigma, dtype=jnp.float64),
+            )
         )
 
 
@@ -405,21 +310,24 @@ class Zgauss(AdditiveComponent):
     $$\mathcal{M}\left( E \right) =
     \frac{K}{(1+z) \sigma \sqrt{2 \pi}}\exp\left(\frac{-(E(1+z)-E_L)^2}{2\sigma^2}\right)$$
 
-    ??? abstract "Parameters"
-        * $E_L$ : Energy of the line $\left[\text{keV}\right]$
-        * $\sigma$ : Width of the line $\left[\text{keV}\right]$
-        * $K$ : Normalization $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$
-        * $z$ : Redshift [dimensionless]
+    !!! abstract "Parameters"
+        * $E_L$ (`E_l`) $\left[\text{keV}\right]$ : Energy of the line
+        * $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : Width of the line
+        * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$ : Normalization
+        * $z$ (`redshift`) $\left[\text{dimensionless}\right]$ : Redshift
     """
 
-    def continuum(self, energy) -> (jax.Array, jax.Array):
-        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))
+    def __init__(self):
+        self.E_l = nnx.Param(2.0)
+        self.sigma = nnx.Param(1e-2)
+        self.redshift = nnx.Param(0.0)
+        self.norm = nnx.Param(1.0)
 
-        return (norm / (1 + redshift)) * jsp.stats.norm.pdf(
-            energy * (1 + redshift), loc=line_energy, scale=sigma
+    def continuum(self, energy) -> (jax.Array, jax.Array):
+        return (self.norm / (1 + self.redshift)) * jsp.stats.norm.pdf(
+            energy * (1 + self.redshift),
+            loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+            scale=jnp.asarray(self.sigma, dtype=jnp.float64),
         )
 
 
@@ -430,50 +338,57 @@ class NSatmos(AdditiveComponent):
     !!! 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
+    !!! abstract "Parameters"
+        * $T_{eff}$ (`Tinf`) $\left[\text{Kelvin}\right]$ : Effective temperature at the surface (No redshift applied)
+        * $M_{ns}$ (`mass`) $\left[M_{\odot}\right]$ : Mass of the NS
+        * $R_∞$  (`radius`) $\left[\text{km}\right]$ : Radius at infinity (modulated by gravitational effects)
+        * $D$ (`distance`) $\left[\text{kpc}\right]$ : Distance to the neutron star
+        * norm (`norm`) $\left[\text{dimensionless}\right]$ : fraction of the neutron star surface emitting
     """
 
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
+    def __init__(self):
+        self.Tinf = nnx.Param(6.0)
+        self.mass = nnx.Param(1.4)
+        self.radius = nnx.Param(10.0)
+        self.distance = nnx.Param(10.0)
+        self.norm = nnx.Param(1.0)
+
         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"]
+        # Extract the table values. All this code could be summarized in two lines if we reformat the nsatmosdata.fits table
+        tab_temperature = np.asarray(entry_table["TEMP"][0], dtype=float)  # Logarithmic value
+        tab_gravity = np.asarray(entry_table["GRAVITY"][0], dtype=float)  # Logarithmic value
+        tab_mucrit = np.asarray(entry_table["MUCRIT"][0], dtype=float)
+        tab_energy = np.asarray(entry_table["ENERGY"][0], dtype=float)
+        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,
+                tab_temperature.size,
+                tab_gravity.size,
+                tab_mucrit.size,
+                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)):
+        for i in range(len(tab_temperature)):
+            for j in range(len(tab_gravity)):
+                for k in range(len(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
+                        tab_flux_flat[
+                            i * len(tab_gravity) * len(tab_mucrit) + j * len(tab_mucrit) + k
                         ]
                     )
 
-        self.tab_flux = np.asarray(tab_flux, dtype=float)
+        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
+        self.tab_temperature = nnx.Variable(tab_temperature)
+        self.tab_gravity = nnx.Variable(tab_gravity)
+        self.tab_mucrit = nnx.Variable(tab_mucrit)
+        self.tab_energy = nnx.Variable(tab_energy)
+        self.tab_flux = nnx.Variable(tab_flux)
 
+    def interp_flux_func(self, temperature_log, gravity_log, mu):
         return interpax.interp3d(
             10.0**temperature_log,
             10.0**gravity_log,
@@ -485,19 +400,18 @@ class NSatmos(AdditiveComponent):
             method="linear",
         )
 
+    @partial(jnp.vectorize, excluded=(0,))
     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
+        # 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))
+        temp_log = self.Tinf
+        mass = self.mass
+        radius = self.radius
+        distance = self.distance
+        norm = self.norm
 
         # Derive parameters usable to retrive value in the flux table
         Rcgs = 1e5 * radius  # Radius in cgs
@@ -519,7 +433,7 @@ class NSatmos(AdditiveComponent):
 
         # Interpolate the flux table to get the flux at the surface
 
-        flux = jax.jit(self.interp_flux_func)(temp_log, gravity_log, cmu)
+        flux = 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)
@@ -542,3 +456,76 @@ class NSatmos(AdditiveComponent):
         )
 
         return jax.lax.select(r_normalized < 1.125, jnp.zeros_like(true_flux), true_flux)
+
+
+class Band(AdditiveComponent):
+    r"""
+    A Band function model
+
+    $$
+    \mathcal{M}(E) = \begin{cases} K \left( \frac{E}{100 \, \text{keV}}\right)^{\alpha_1}\exp(-\frac{E}{E_c})  &
+    \text{if $E < E_c (\alpha_1 - \alpha_2)$} \\
+    K \left[ (\alpha_1 - \alpha_2) \frac{E_c}{100 \, \text{keV}} \right]^{\alpha_1-\alpha_2} \left( \frac{E}{100 \, \text{keV}}\right)^{\alpha_2} \exp(-(\alpha_1 - \alpha_2))  & \text{if $E > E_c (\alpha_1 - \alpha_2)$}
+    \end{cases}
+    $$
+
+    !!! abstract "Parameters"
+        * $\alpha_1$ (`alpha1`) $\left[\text{dimensionless}\right]$ : First powerlaw index
+        * $\alpha_2$  (`alpha2`) $\left[\text{dimensionless}\right]$ : Second powerlaw index
+        * $E_c$  (`Ec`) $\left[\text{keV}\right]$ : Radius at infinity (modulated by gravitational effects)
+        * norm (`norm`) $\left[\frac{\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization at the reference energy (100 keV)
+    """
+
+    def __init__(self):
+        self.alpha1 = nnx.Param(-1.0)
+        self.alpha2 = nnx.Param(-2.0)
+        self.Ec = nnx.Param(200.0)
+        self.norm = nnx.Param(1e-4)
+
+    def continuum(self, energy):
+        Epivot = 100.0
+        alpha_diff = jnp.asarray(self.alpha1) - jnp.asarray(self.alpha2)
+
+        spectrum = jnp.where(
+            energy < self.Ec * (alpha_diff),
+            (energy / Epivot) ** self.alpha1 * jnp.exp(-energy / self.Ec),
+            (alpha_diff * (self.Ec / Epivot)) ** (alpha_diff)
+            * (energy / 100) ** self.alpha2
+            * jnp.exp(-alpha_diff),
+        )
+
+        return self.norm * spectrum
+
+
+'''
+class Diskpbb(AdditiveComponent):
+    r"""
+    A multiple blackbody disk model where local disk temperature T(r) is proportional to $$r^{(-p)}$$,
+    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}}}$$
+
+    ??? 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]$,
+        $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]$
+    """
+
+    def continuum(self, energy):
+        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):
+            return 2.78e-3 * energy**2 * (kT / tin) ** (-2 / p - 1) / (jnp.exp(energy / kT) - 1)
+
+        func_vmapped = jax.vmap(lambda e: integrate_interval(lambda kT: integrand(kT, e), 0, tin, n=51))
+
+        return norm * (0.75 / p) * func_vmapped(energy)
+'''