jaxspec 0.1.0__py3-none-any.whl → 0.1.2__py3-none-any.whl

jaxspec/model/abc.py

@@ -1,12 +1,17 @@
 from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from uuid import uuid4
+
 import haiku as hk
 import jax
 import jax.numpy as jnp
 import networkx as nx
+import rich
+
 from haiku._src import base
-from uuid import uuid4
 from jax.scipy.integrate import trapezoid
-from abc import ABC
+from rich.table import Table
 from simpleeval import simple_eval
 
 
@@ -107,6 +112,30 @@ class SpectralModel:
     def params(self):
         return self.transformed_func_photon.init(None, jnp.ones(10), jnp.ones(10))
 
+    def __rich_repr__(self):
+        table = Table(title=str(self))
+
+        table.add_column("Component", justify="right", style="bold", no_wrap=True)
+        table.add_column("Parameter")
+
+        params = self.params
+
+        for component in params.keys():
+            once = True
+
+            for parameters in params[component].keys():
+                table.add_row(component if once else "", parameters)
+                once = False
+
+        return table
+
+    def __repr_html_(self):
+        return self.__rich_repr__()
+
+    def __repr__(self):
+        rich.print(self.__rich_repr__())
+        return ""
+
     def photon_flux(self, params, e_low, e_high, n_points=2):
         r"""
         Compute the expected counts between $E_\min$ and $E_\max$ by integrating the model.
@@ -215,16 +244,18 @@ class SpectralModel:
                 continuum[node_id] = runtime_modules[node_id].continuum(energies)
 
             elif node and node["type"] == "operation":
-                component_1 = list(self.graph.in_edges(node_id))[0][0]
+                component_1 = list(self.graph.in_edges(node_id))[0][0]  # noqa: RUF015
                 component_2 = list(self.graph.in_edges(node_id))[1][0]
-                continuum[node_id] = node["function"](continuum[component_1], continuum[component_2])
+                continuum[node_id] = node["function"](
+                    continuum[component_1], continuum[component_2]
+                )
 
         if n_points == 2:
-            flux_1D = continuum[list(self.graph.in_edges("out"))[0][0]]
+            flux_1D = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
             flux = jnp.stack((flux_1D[:-1], flux_1D[1:]))
 
         else:
-            flux = continuum[list(self.graph.in_edges("out"))[0][0]]
+            flux = continuum[list(self.graph.in_edges("out"))[0][0]]  # noqa: RUF015
 
         if energy_flux:
             continuum_flux = trapezoid(
@@ -234,7 +265,9 @@ class SpectralModel:
             )
 
         else:
-            continuum_flux = trapezoid(flux * energies_to_integrate, x=jnp.log(energies_to_integrate), axis=0)
+            continuum_flux = trapezoid(
+                flux * energies_to_integrate, x=jnp.log(energies_to_integrate), axis=0
+            )
 
         # Iterate from the root nodes to the output node and
         # compute the fine structure contribution for each component
@@ -249,14 +282,18 @@ class SpectralModel:
             path = nx.shortest_path(self.graph, source=root_node_id, target="out")
             nodes_id_in_path = [node_id for node_id in path]
 
-            flux_from_component, mean_energy = runtime_modules[root_node_id].emission_lines(e_low, e_high)
+            flux_from_component, mean_energy = runtime_modules[root_node_id].emission_lines(
+                e_low, e_high
+            )
 
             multiplicative_nodes = []
 
             # Search all multiplicative components connected to this node
             # and apply them at mean energy
             for node_id in nodes_id_in_path[::-1]:
-                multiplicative_nodes.extend([node_id for node_id in self.find_multiplicative_components(node_id)])
+                multiplicative_nodes.extend(
+                    [node_id for node_id in self.find_multiplicative_components(node_id)]
+                )
 
             for mul_node in multiplicative_nodes:
                 flux_from_component *= runtime_modules[mul_node].continuum(mean_energy)
@@ -309,7 +346,10 @@ class SpectralModel:
         if component.type == "additive":
 
             def lam_func(e):
-                return component(**kwargs).continuum(e) + component(**kwargs).emission_lines(e, e + 1)[0]
+                return (
+                    component(**kwargs).continuum(e)
+                    + component(**kwargs).emission_lines(e, e + 1)[0]
+                )
 
         elif component.type == "multiplicative":
 
@@ -342,7 +382,9 @@ class SpectralModel:
 
         return cls(graph, labels)
 
-    def compose(self, other: SpectralModel, operation=None, function=None, name=None) -> SpectralModel:
+    def compose(
+        self, other: SpectralModel, operation=None, function=None, name=None
+    ) -> SpectralModel:
         """
         This function operate a composition between the operation graph of two models
         1) It fuses the two graphs using which joins at the 'out' nodes
@@ -524,3 +566,10 @@ class AdditiveComponent(ModelComponent, ABC):
 
         return jnp.trapz(self(x) * dx, x=t)
     '''
+
+
+class MultiplicativeComponent(ModelComponent, ABC):
+    type = "multiplicative"
+
+    @abstractmethod
+    def continuum(self, energy): ...

jaxspec/model/additive.py

@@ -14,8 +14,6 @@ from haiku.initializers import Constant as HaikuConstant
 
 from ..util.integrate import integrate_interval
 from ..util.online_storage import table_manager
-
-# from ._additive.apec import APEC
 from .abc import AdditiveComponent
 
 
@@ -38,7 +36,7 @@ class Powerlaw(AdditiveComponent):
         return norm * energy ** (-alpha)
 
 
-class AdditiveConstant(AdditiveComponent):
+class Additiveconstant(AdditiveComponent):
     r"""
     A constant model
 

jaxspec/model/multiplicative.py

@@ -1,7 +1,5 @@
 from __future__ import annotations
 
-from abc import ABC, abstractmethod
-
 import haiku as hk
 import jax.numpy as jnp
 import numpy as np
@@ -10,14 +8,7 @@ from astropy.table import Table
 from haiku.initializers import Constant as HaikuConstant
 
 from ..util.online_storage import table_manager
-from .abc import ModelComponent
-
-
-class MultiplicativeComponent(ModelComponent, ABC):
-    type = "multiplicative"
-
-    @abstractmethod
-    def continuum(self, energy): ...
+from .abc import MultiplicativeComponent
 
 
 class Expfac(MultiplicativeComponent):
@@ -226,6 +217,7 @@ class Tbpcf(MultiplicativeComponent):
 
         return f * jnp.exp(-nh * sigma) + (1 - f)
 
+
 class FDcut(MultiplicativeComponent):
     r"""
     A Fermi-Dirac cutoff model.
@@ -243,4 +235,4 @@ class FDcut(MultiplicativeComponent):
         cutoff = hk.get_parameter("E_c", [], init=HaikuConstant(1))
         folding = hk.get_parameter("E_f", [], init=HaikuConstant(1))
 
-        return (1 + jnp.exp((energy - cutoff)/folding)) ** -1
+        return (1 + jnp.exp((energy - cutoff) / folding)) ** -1

jaxspec/util/typing.py

@@ -9,6 +9,13 @@ from pydantic import BaseModel, field_validator
 PriorDictType = dict[str, dict[str, dist.Distribution | ArrayLike]]
 
 
+def is_flat_dict(input_data: dict[str, Any]) -> bool:
+    """
+    Check if the input data is a flat dictionary with string keys and non-dictionary values.
+    """
+    return all(isinstance(k, str) and not isinstance(v, dict) for k, v in input_data.items())
+
+
 class PriorDictModel(BaseModel):
     """
     Pydantic model for a nested dictionary of NumPyro distributions or JAX arrays.
@@ -21,6 +28,23 @@ class PriorDictModel(BaseModel):
     class Config:  # noqa D106
         arbitrary_types_allowed = True
 
+    @classmethod
+    def from_dict(cls, input_prior: dict[str, Any]):
+        if is_flat_dict(input_prior):
+            nested_dict = {}
+
+            for key, obj in input_prior.items():
+                component, component_number, *parameter = key.split("_")
+
+                sub_dict = nested_dict.get(f"{component}_{component_number}", {})
+                sub_dict["_".join(parameter)] = obj
+
+                nested_dict[f"{component}_{component_number}"] = sub_dict
+
+            return cls(nested_dict=nested_dict)
+
+        return cls(nested_dict=input_prior)
+
     @field_validator("nested_dict", mode="before")
     def check_and_cast_nested_dict(cls, value: dict[str, Any]):
         if not isinstance(value, dict):
@@ -35,9 +59,10 @@ class PriorDictModel(BaseModel):
                     try:
                         # Attempt to cast to JAX array
                         value[key][inner_key] = jnp.array(obj, dtype=float)
+
                     except Exception as e:
                         raise ValueError(
-                            f'The value for key "{inner_key}" in inner dictionary must '
-                            f"be a NumPyro distribution or castable to JAX array. Error: {e}"
+                            f'The value for key "{inner_key}" in {key} be a NumPyro '
+                            f"distribution or castable to JAX array. Error: {e}"
                         )
         return value

{jaxspec-0.1.0.dist-info → jaxspec-0.1.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxspec
-Version: 0.1.0
+Version: 0.1.2
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 Home-page: https://github.com/renecotyfanboy/jaxspec
 License: MIT
@@ -15,19 +15,18 @@ Requires-Dist: arviz (>=0.17.1,<0.20.0)
 Requires-Dist: astropy (>=6.0.0,<7.0.0)
 Requires-Dist: chainconsumer (>=1.1.2,<2.0.0)
 Requires-Dist: cmasher (>=1.6.3,<2.0.0)
-Requires-Dist: dm-haiku (>=0.0.11,<0.0.13)
+Requires-Dist: dm-haiku (>=0.0.12,<0.0.13)
 Requires-Dist: gpjax (>=0.8.0,<0.9.0)
 Requires-Dist: interpax (>=0.3.3,<0.4.0)
-Requires-Dist: jax (>=0.4.30,<0.5.0)
+Requires-Dist: jax (>=0.4.33,<0.5.0)
 Requires-Dist: jaxlib (>=0.4.30,<0.5.0)
-Requires-Dist: jaxns (>=2.5.1,<3.0.0)
+Requires-Dist: jaxns (<2.6)
 Requires-Dist: jaxopt (>=0.8.1,<0.9.0)
 Requires-Dist: matplotlib (>=3.8.0,<4.0.0)
 Requires-Dist: mendeleev (>=0.15,<0.18)
-Requires-Dist: mkdocstrings (>=0.24,<0.26)
 Requires-Dist: networkx (>=3.1,<4.0)
 Requires-Dist: numpy (<2.0.0)
-Requires-Dist: numpyro (>=0.15.2,<0.16.0)
+Requires-Dist: numpyro (>=0.15.3,<0.16.0)
 Requires-Dist: optimistix (>=0.0.7,<0.0.8)
 Requires-Dist: pandas (>=2.2.0,<3.0.0)
 Requires-Dist: pooch (>=1.8.2,<2.0.0)
@@ -41,7 +40,14 @@ Requires-Dist: watermark (>=2.4.3,<3.0.0)
 Project-URL: Documentation, https://jaxspec.readthedocs.io/en/latest/
 Description-Content-Type: text/markdown
 
-# jaxspec
+<p align="center">
+  <img src="https://raw.githubusercontent.com/renecotyfanboy/jaxspec/main/docs/logo/logo_small.svg" alt="Logo" width="100" height="100">
+</p>
+
+<h1 align="center">
+  jaxspec
+</h1>
+
 
 [![PyPI - Version](https://img.shields.io/pypi/v/jaxspec?style=for-the-badge&logo=pypi&color=rgb(37%2C%20150%2C%20190))](https://pypi.org/project/jaxspec/)
 [![Python package](https://img.shields.io/pypi/pyversions/jaxspec?style=for-the-badge)](https://pypi.org/project/jaxspec/)

{jaxspec-0.1.0.dist-info → jaxspec-0.1.2.dist-info}/RECORD

@@ -1,33 +1,31 @@
 jaxspec/__init__.py,sha256=Sbn02lX6Y-zNXk17N8dec22c5jeypiS0LkHmGfz7lWA,126
 jaxspec/analysis/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+jaxspec/analysis/_plot.py,sha256=C4XljmuzQz8xQur_jQddgInrBDmKgTn0eugSreLoD5k,862
 jaxspec/analysis/compare.py,sha256=g2UFhmR9Zt-7cz5gQFOB6lXuklXB3yTyUvjTypOzoSY,725
 jaxspec/analysis/results.py,sha256=Kz3eryxS3N_hiajcFLTWS1dtgTQo5hlh-rDCnJ3A-3c,27811
 jaxspec/data/__init__.py,sha256=aantcYKC9kZFvaE-V2SIwSuLhIld17Kjrd9CIUu___Y,415
 jaxspec/data/grouping.py,sha256=hhgBt-voiH0DDSyePacaIGsaMnrYbJM_-ZeU66keC7I,622
 jaxspec/data/instrument.py,sha256=0pSf1p82g7syDMmKm13eVbYih-Veiq5DnwsyZe6_b4g,3890
-jaxspec/data/obsconf.py,sha256=0X9jR-pV-Pk4-EVuUdlVWgl_gBx8ZurVkRNrfKQWdC4,8663
+jaxspec/data/obsconf.py,sha256=gv14sL6azK2avRiMCWuTbyLBPulzm4PwvoLY6iWPEVE,9833
 jaxspec/data/observation.py,sha256=1UnFu5ihZp9z-vP_I7tsFY8jhhIJunv46JyuE-acrg0,6394
 jaxspec/data/ogip.py,sha256=sv9p00qHS5pzw61pzWyyF0nV-E-RXySdSFK2tUavokA,9545
 jaxspec/data/util.py,sha256=ycLPVE-cjn6VpUWYlBU1BGfw73ANXIBilyVAUOYOSj0,9540
-jaxspec/fit.py,sha256=C9hQxMJz1nLu47rHWkiKx7J7oS1bXow_kMKwswsJy8U,24791
+jaxspec/fit.py,sha256=hI0koMO4KsNpe9mLlaFm_tNLgm4BVAYVyiMb1E1eyZE,24553
 jaxspec/model/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxspec/model/_additive/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-jaxspec/model/_additive/apec.py,sha256=r7CQqscAgR0BXC_AJqF6B7CPq3Byoo65Z-h9XgACZeU,12460
-jaxspec/model/_additive/apec_loaders.py,sha256=jkUoH0ezeYdaNw3oV10V0L-jt848SKp2thanLWLWp9k,2412
-jaxspec/model/abc.py,sha256=SWjKOOsqU5UJsVy63Tt9dDq8H2eTIbvK2C9iqgiR0cY,19817
-jaxspec/model/additive.py,sha256=CT2K2DVVeHKN1tee9-J3MYdEPqEOolLB2E7HU-RJKZw,22485
+jaxspec/model/abc.py,sha256=MuxEyvn223QPwGoFIJiST8nRMgrZ08ZLkw33oep3tx4,20887
+jaxspec/model/additive.py,sha256=wjY2wL3Io3F45GJpz-UB8xYVnA-W1OFBnZMbj5pWPbQ,22449
 jaxspec/model/background.py,sha256=QSFFiuyUEvuzXBx3QfkvVneUR8KKEP-VaANEVXcavDE,7865
 jaxspec/model/list.py,sha256=0RPAoscVz_zM1CWdx_Gd5wfrQWV5Nv4Kd4bSXu2ayUA,860
-jaxspec/model/multiplicative.py,sha256=TG3PCgS7oCuHwJ4TM4whw6pz318oo9MVvjSs4sQZVPc,8300
+jaxspec/model/multiplicative.py,sha256=GCQ6JRz92QqbzDBFwWxGZ9SUqTJZQpD7B6ji9VEFXWo,8135
 jaxspec/scripts/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 jaxspec/scripts/debug.py,sha256=RIRykBW_kzt8__PhOohQn2xtDW6oAz49E8rmuR5ewAU,293
 jaxspec/util/__init__.py,sha256=vKurfp7p2hxHptJjXhXqFAXAikAGXAqISMJUqPeiGTw,1259
 jaxspec/util/abundance.py,sha256=fsC313taIlGzQsZNwbYsJupDWm7ZbqzGhY66Ku394Mw,8546
 jaxspec/util/integrate.py,sha256=_Ax_knpC7d4et2-QFkOUzVtNeQLX1-cwLvm-FRBxYcw,4505
 jaxspec/util/online_storage.py,sha256=vm56RfcbFKpkRVfr0bXO7J9aQxuBq-I_oEgA26YIhCo,2469
-jaxspec/util/typing.py,sha256=qwZMKHivZlozoo0ESsiaQNkG99Dh3PE2Z-5aOQD9zc0,1650
-jaxspec-0.1.0.dist-info/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
-jaxspec-0.1.0.dist-info/METADATA,sha256=ONyl92xPTT1LNXKNkonzdH4IlWFlkGdDo-09EEGur9c,3572
-jaxspec-0.1.0.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-jaxspec-0.1.0.dist-info/entry_points.txt,sha256=kzLG2mGlCWITRn4Q6zKG_idx-_RKAncvA0DMNYTgHAg,71
-jaxspec-0.1.0.dist-info/RECORD,,
+jaxspec/util/typing.py,sha256=8qK1aJlsqTcVKjYN-BxsDx20BTwtnS-wMw6Bdurpm-o,2459
+jaxspec-0.1.2.dist-info/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
+jaxspec-0.1.2.dist-info/METADATA,sha256=FE2bTAk-3Xryi6fplV4Y-F2eibUdLZgC9ET9_4HvdOA,3708
+jaxspec-0.1.2.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
+jaxspec-0.1.2.dist-info/entry_points.txt,sha256=kzLG2mGlCWITRn4Q6zKG_idx-_RKAncvA0DMNYTgHAg,71
+jaxspec-0.1.2.dist-info/RECORD,,

jaxspec/model/_additive/apec.py

@@ -1,316 +0,0 @@
-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 fori_loop, scan
-from jax.scipy.stats import norm as gaussian
-
-from ...util.abundance import abundance_table, element_data
-from ..abc import AdditiveComponent
-from .apec_loaders import get_continuum, get_lines, get_pseudo, get_temperature
-
-
-@jax.jit
-def lerp(x, x0, x1, y0, y1):
-    """
-    Linear interpolation routine
-    Return y(x) =  (y0 * (x1 - x) + y1 * (x - x0)) / (x1 - x0)
-    """
-    return (y0 * (x1 - x) + y1 * (x - x0)) / (x1 - x0)
-
-
-@jax.jit
-def interp_and_integrate(energy_low, energy_high, energy_ref, continuum_ref, end_index):
-    """
-    This function interpolate & integrate the values of a tabulated reference continuum between two energy limits
-    Sorry for the boilerplate here, but be sure that it works !
-
-    Parameters:
-        energy_low: lower limit of the integral
-        energy_high: upper limit of the integral
-        energy_ref: energy grid of the reference continuum
-        continuum_ref: continuum values evaluated at energy_ref
-
-    """
-    energy_ref = jnp.where(jnp.arange(energy_ref.shape[0]) < end_index, energy_ref, jnp.nan)
-    start_index = jnp.searchsorted(energy_ref, energy_low, side="left") - 1
-    end_index = jnp.searchsorted(energy_ref, energy_high, side="left") + 1
-
-    def body_func(index, value):
-        integrated_flux, previous_energy, previous_continuum = value
-        current_energy, current_continuum = energy_ref[index], continuum_ref[index]
-
-        # 5 cases
-        # Neither current and previous energies are within the integral limits > nothing is added to the integrated flux
-        # The left limit of the integral is between the current and previous energy > previous energy is set to the limit, previous continuum is interpolated, and then added to the integrated flux
-        # The right limit of the integral is between the current and previous energy > current energy is set to the limit, current continuum is interpolated, and then added to the integrated flux
-        # Both current and previous energies are within the integral limits -> add to the integrated flux
-        # 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
-        )
-        energies_within_bins = (previous_energy <= energy_low) * (energy_high < current_energy)
-
-        case = (
-            (1 - previous_energy_is_between) * current_energy_is_between * 1
-            + previous_energy_is_between * (1 - current_energy_is_between) * 2
-            + (previous_energy_is_between * current_energy_is_between) * 3
-            + energies_within_bins * 4
-        )
-
-        term_to_add = lax.switch(
-            case,
-            [
-                lambda pe, pc, ce, cc, el, er: 0.0,  # 1
-                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,
-                # 5
-            ],
-            previous_energy,
-            previous_continuum,
-            current_energy,
-            current_continuum,
-            energy_low,
-            energy_high,
-        )
-
-        return integrated_flux + term_to_add, current_energy, current_continuum
-
-    integrated_flux, _, _ = fori_loop(start_index, end_index, body_func, (0.0, 0.0, 0.0))
-
-    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)
-
-
-@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
-        continuum = interp_and_integrate(e_low, e_high, energy_ref, continuum_ref, end_index)
-
-        return carry, continuum
-
-    _, continuum = scan(scanned_func, 0.0, (energy[:-1], energy[1:]))
-
-    return continuum
-
-
-@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
-    """
-
-    def scanned_func(_, unpack):
-        energy_ref, continuum_ref, end_idx = unpack
-        element_flux = interp_flux(energy, energy_ref, continuum_ref, end_idx)
-
-        return _, element_flux
-
-    _, flux = scan(scanned_func, 0.0, (energy_ref, continuum_ref, end_index))
-
-    return abundances @ flux
-
-
-@jax.jit
-def get_lines_contribution_broadening(
-    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 = 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 continuum_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)
-
-    return lerp(kT, kT_low, kT_high, continuum_low, continuum_high)
-
-
-@jax.jit
-def pseudo_func(energy, kT, abundances):
-    idx, kT_low, kT_high = get_temperature(kT)
-    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.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
-    )
-
-    return lerp(kT, kT_low, kT_high, line_low, line_high)
-
-
-class APEC(AdditiveComponent):
-    """
-    APEC model implementation in pure JAX for X-ray spectral fitting.
-
-    !!! warning
-        This implementation is optimised for the CPU, it shows poor performance on the GPU.
-    """
-
-    def __init__(
-        self,
-        continuum: bool = True,
-        pseudo: bool = True,
-        lines: bool = True,
-        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",
-        trace_abundance: float = 1.0,
-        **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())
-
-        self.abundance_table = abundance_table
-        self.thermal_broadening = thermal_broadening
-        self.turbulent_broadening = turbulent_broadening
-        self.continuum_to_compute = continuum
-        self.pseudo_to_compute = pseudo
-        self.lines_to_compute = lines
-        self.trace_abundance = trace_abundance
-        self.variant = variant
-
-    def get_thermal_broadening(self):
-        r"""
-        Compute the thermal broadening $\sigma_T$ for each element using :
-
-        $$ \frac{\sigma_T}{E_{\text{line}}} = \frac{1}{c}\sqrt{\frac{k_{B} T}{A m_p}}$$
-
-        where $E_{\text{line}}$ is the energy of the line, $c$ is the speed of light, $k_{B}$ is the Boltzmann constant,
-        $T$ is the temperature, $A$ is the atomic weight of the element and $m_p$ is the proton mass.
-        """
-
-        if self.thermal_broadening:
-            kT = hk.get_parameter("kT", [], init=HaikuConstant(6.5))
-            factor = 1 / c * (1 / m_p) ** (1 / 2)
-            factor = factor.to(u.keV ** (-1 / 2)).value
-
-            # Multiply this factor by Line_Energy * sqrt(kT/A) to get the broadening for a line
-            # This return value must be multiplied by the energy of the line to get actual broadening
-            return factor * jnp.sqrt(kT / self.atomic_weights)
-
-        else:
-            return jnp.zeros((30,))
-
-    def get_turbulent_broadening(self):
-        r"""
-        Return the turbulent broadening  using :
-
-        $$\frac{\sigma_\text{turb}}{E_{\text{line}}} = \frac{\sigma_{v ~ ||}}{c}$$
-
-        where $\sigma_{v ~ ||}$ is the velocity dispersion along the line of sight in km/s.
-        """
-        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
-            )
-        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
-        )
-
-        # Set abundances of trace element (will be overwritten in the vv case)
-        abund = jnp.ones((30,)).at[trace_elements].multiply(self.trace_abundance)
-
-        if self.variant == "vv":
-            for i, element in enumerate(abundance_table["Element"]):
-                if element != "H":
-                    abund = abund.at[i].set(hk.get_parameter(element, [], init=HaikuConstant(1.0)))
-
-        elif self.variant == "v":
-            for i, element in enumerate(abundance_table["Element"]):
-                if element != "H" and element in v_elements:
-                    abund = abund.at[i].set(hk.get_parameter(element, [], init=HaikuConstant(1.0)))
-
-        else:
-            Z = hk.get_parameter("Abundance", [], init=HaikuConstant(1.0))
-            for i, element in enumerate(abundance_table["Element"]):
-                if element != "H" and element in none_elements:
-                    abund = abund.at[i].set(Z)
-
-        if abund != "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))
-        z = hk.get_parameter("Redshift", [], init=HaikuConstant(0.0))
-        norm = hk.get_parameter("norm", [], init=HaikuConstant(1.0))
-
-        return kT, z, norm, abund
-
-    def emission_lines(self, e_low, e_high):
-        # Get the parameters and extract the relevant data
-        energy = jnp.hstack([e_low, e_high[-1]])
-        kT, z, norm, abundances = self.get_parameters()
-        total_broadening = jnp.hypot(self.get_thermal_broadening(), self.get_turbulent_broadening())
-        energy = energy * (1 + z)
-
-        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
-        )
-
-        return (continuum + pseudo_continuum + lines) * norm * 1e14 / (1 + z), (e_low + e_high) / 2