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

jaxspec/model/_additive/apec.py

@@ -0,0 +1,377 @@
+import jax.numpy as jnp
+import jax
+import haiku as hk
+import astropy.units as u
+from jax import lax
+from jax.lax import scan, fori_loop
+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
+
+
+@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
+
+
+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)
+
+
+def interp_flux(energy, energy_ref, continuum_ref, end_index, integrate=True):
+    """
+    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)
+
+        return carry, continuum
+
+    _, continuum = scan(scanned_func, 0.0, (energy[:-1], energy[1:]))
+
+    return continuum
+
+
+def interp_flux_elements(energy_ref, continuum_ref, end_index, energy, abundances, integrate=True):
+    """
+    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, integrate=integrate)
+
+        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 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 = 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.custom_jvp
+@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
+@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)
+
+    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.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)
+
+
+@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.
+
+    !!! 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(APEC, self).__init__(**kwargs)
+
+        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

jaxspec/model/_additive/apec_loaders.py

@@ -0,0 +1,90 @@
+""" 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 jax
+import jax.numpy as jnp
+import numpy as np
+import importlib.resources
+import xarray as xr
+
+
+apec_file = xr.open_dataset(importlib.resources.files("jaxspec") / "tables/apec.nc", engine="h5netcdf")
+
+
+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
+
+
+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])
+
+    return pseudo_energy_array, pseudo_emissivity_array, end_index_pseudo
+
+
+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])
+
+    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 line_energy_array, line_element_array, line_emissivity_array, end_index_lines
+
+
+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)))
+
+
+@jax.jit
+def get_temperature(kT):
+    return jax.pure_callback(temperature_table_getter, pure_callback_temperature_shape, kT)
+
+
+@jax.jit
+def get_continuum(idx):
+    return jax.pure_callback(continuum_table_getter, pure_callback_continuum_shape, idx)
+
+
+@jax.jit
+def get_pseudo(idx):
+    return jax.pure_callback(pseudo_table_getter, pure_callback_pseudo_shape, idx)
+
+
+@jax.jit
+def get_lines(idx):
+    return jax.pure_callback(lines_table_getter, pure_callback_line_shape, idx)

jaxspec/model/abc.py

@@ -123,9 +123,9 @@ class SpectralModel:
 
         !!! info
             This method is internally used in the inference process and should not be used directly. See
-            [`photon_flux`][jaxspec.analysis.results.ChainResult.photon_flux] to compute
+            [`photon_flux`][jaxspec.analysis.results.FitResult.photon_flux] to compute
             the photon flux associated with a set of fitted parameters in a
-            [`ChainResult`][jaxspec.analysis.results.ChainResult]
+            [`FitResult`][jaxspec.analysis.results.FitResult]
             instead.
         """
 
@@ -151,9 +151,9 @@ class SpectralModel:
 
         !!! info
             This method is internally used in the inference process and should not be used directly. See
-            [`energy_flux`](/references/results/#jaxspec.analysis.results.ChainResult.energy_flux) to compute
+            [`energy_flux`](/references/results/#jaxspec.analysis.results.FitResult.energy_flux) to compute
             the energy flux associated with a set of fitted parameters in a
-            [`ChainResult`](/references/results/#jaxspec.analysis.results.ChainResult)
+            [`FitResult`](/references/results/#jaxspec.analysis.results.FitResult)
             instead.
         """
 
@@ -309,7 +309,7 @@ class SpectralModel:
         if component.type == "additive":
 
             def lam_func(e):
-                return component().continuum(e) + component().emission_lines(e, e + 1)[0]
+                return component(**kwargs).continuum(e) + component(**kwargs).emission_lines(e, e + 1)[0]
 
         elif component.type == "multiplicative":
 
@@ -326,7 +326,7 @@ class SpectralModel:
             "component_type": component.type,
             "name": component.__name__.lower(),
             "component": component,
-            "params": hk.transform(lam_func).init(None, jnp.ones(1)),
+            # "params": hk.transform(lam_func).init(None, jnp.ones(1)),
             "fine_structure": False,
             "kwargs": kwargs,
             "depth": 0,
@@ -454,7 +454,7 @@ class ComponentMetaClass(type(hk.Module)):
     syntax while style enabling the components to be used as haiku modules.
     """
 
-    def __call__(self, **kwargs):
+    def __call__(self, **kwargs) -> SpectralModel:
         """
         This method enable to use model components as haiku modules when folded in a haiku transform
         function and also to instantiate them as SpectralModel when out of a haiku transform
@@ -476,3 +476,51 @@ class ModelComponent(hk.Module, ABC, metaclass=ComponentMetaClass):
 
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
+
+
+class AdditiveComponent(ModelComponent, ABC):
+    type = "additive"
+
+    def continuum(self, energy):
+        """
+        Method for computing the continuum associated to the model.
+        By default, this is set to 0, which means that the model has no continuum.
+        This should be overloaded by the user if the model has a continuum.
+        """
+
+        return jnp.zeros_like(energy)
+
+    def emission_lines(self, e_min, e_max) -> (jax.Array, jax.Array):
+        """
+        Method for computing the fine structure of an additive model between two energies.
+        By default, this is set to 0, which means that the model has no emission lines.
+        This should be overloaded by the user if the model has a fine structure.
+        """
+
+        return jnp.zeros_like(e_min), (e_min + e_max) / 2
+
+    '''
+    def integral(self, e_min, e_max):
+        r"""
+        Method for integrating an additive model between two energies. It relies on
+        double exponential quadrature for finite intervals to compute an approximation
+        of the integral of a model.
+
+        references
+        ----------
+        * $Takahasi and Mori (1974) <https://ems.press/journals/prims/articles/2686>$_
+        * $Mori and Sugihara (2001) <https://doi.org/10.1016/S0377-0427(00)00501-X>$_
+        * $Tanh-sinh quadrature <https://en.wikipedia.org/wiki/Tanh-sinh_quadrature>$_ from Wikipedia
+
+        """
+
+        t = jnp.linspace(-4, 4, 71) # The number of points used is hardcoded and this is not ideal
+        # Quadrature nodes as defined in reference
+        phi = jnp.tanh(jnp.pi / 2 * jnp.sinh(t))
+        dphi = jnp.pi / 2 * jnp.cosh(t) * (1 / jnp.cosh(jnp.pi / 2 * jnp.sinh(t)) ** 2)
+        # Change of variable to turn the integral from E_min to E_max into an integral from -1 to 1
+        x = (e_max - e_min) / 2 * phi + (e_max + e_min) / 2
+        dx = (e_max - e_min) / 2 * dphi
+
+        return jnp.trapz(self(x) * dx, x=t)
+    '''

jaxspec/model/additive.py

@@ -1,6 +1,5 @@
 from __future__ import annotations
 
-from abc import ABC
 
 import haiku as hk
 import jax
@@ -8,58 +7,10 @@ import jax.numpy as jnp
 import jax.scipy as jsp
 import astropy.units as u
 import astropy.constants
-
-from .abc import ModelComponent
 from haiku.initializers import Constant as HaikuConstant
 from ..util.integrate import integrate_interval
-
-
-class AdditiveComponent(ModelComponent, ABC):
-    type = "additive"
-
-    def continuum(self, energy):
-        """
-        Method for computing the continuum associated to the model.
-        By default, this is set to 0, which means that the model has no continuum.
-        This should be overloaded by the user if the model has a continuum.
-        """
-
-        return jnp.zeros_like(energy)
-
-    def emission_lines(self, e_min, e_max) -> (jax.Array, jax.Array):
-        """
-        Method for computing the fine structure of an additive model between two energies.
-        By default, this is set to 0, which means that the model has no emission lines.
-        This should be overloaded by the user if the model has a fine structure.
-        """
-
-        return jnp.zeros_like(e_min), (e_min + e_max) / 2
-
-    '''
-    def integral(self, e_min, e_max):
-        r"""
-        Method for integrating an additive model between two energies. It relies on
-        double exponential quadrature for finite intervals to compute an approximation
-        of the integral of a model.
-
-        references
-        ----------
-        * $Takahasi and Mori (1974) <https://ems.press/journals/prims/articles/2686>$_
-        * $Mori and Sugihara (2001) <https://doi.org/10.1016/S0377-0427(00)00501-X>$_
-        * $Tanh-sinh quadrature <https://en.wikipedia.org/wiki/Tanh-sinh_quadrature>$_ from Wikipedia
-
-        """
-
-        t = jnp.linspace(-4, 4, 71) # The number of points used is hardcoded and this is not ideal
-        # Quadrature nodes as defined in reference
-        phi = jnp.tanh(jnp.pi / 2 * jnp.sinh(t))
-        dphi = jnp.pi / 2 * jnp.cosh(t) * (1 / jnp.cosh(jnp.pi / 2 * jnp.sinh(t)) ** 2)
-        # Change of variable to turn the integral from E_min to E_max into an integral from -1 to 1
-        x = (e_max - e_min) / 2 * phi + (e_max + e_min) / 2
-        dx = (e_max - e_min) / 2 * dphi
-
-        return jnp.trapz(self(x) * dx, x=t)
-    '''
+from .abc import AdditiveComponent
+# from ._additive.apec import APEC  # noqa: F401
 
 
 class Powerlaw(AdditiveComponent):

jaxspec/tables/abundances.dat

@@ -0,0 +1,31 @@
+Element   angr        aspl      feld       aneb      grsa        wilm     lodd        lgpp        lgps
+H      1.00e+00   1.00e+00   1.00e+00   1.00e+00   1.00e+00   1.00e+00   1.00e+00   1.00E+00 	1.00E+00
+He     9.77e-02   8.51e-02   9.77e-02   8.01e-02   8.51e-02   9.77e-02   7.92e-02   8.41E-02 	9.69E-02
+Li     1.45e-11   1.12e-11   1.26e-11   2.19e-09   1.26e-11   0.00       1.90e-09   1.26E-11 	2.15E-09
+Be     1.41e-11   2.40e-11   2.51e-11   2.87e-11   2.51e-11   0.00       2.57e-11   2.40E-11 	2.36E-11
+B      3.98e-10   5.01e-10   3.55e-10   8.82e-10   3.55e-10   0.00       6.03e-10   5.01E-10 	7.26E-10
+C      3.63e-04   2.69e-04   3.98e-04   4.45e-04   3.31e-04   2.40e-04   2.45e-04   2.45E-04 	2.78E-04
+N      1.12e-04   6.76e-05   1.00e-04   9.12e-05   8.32e-05   7.59e-05   6.76e-05   7.24E-05 	8.19E-05
+O      8.51e-04   4.90e-04   8.51e-04   7.39e-04   6.76e-04   4.90e-04   4.90e-04   5.37E-04 	6.06E-04
+F      3.63e-08   3.63e-08   3.63e-08   3.10e-08   3.63e-08   0.00       2.88e-08   3.63E-08 	3.10E-08
+Ne     1.23e-04   8.51e-05   1.29e-04   1.38e-04   1.20e-04   8.71e-05   7.41e-05   1.12E-04 	1.27E-04
+Na     2.14e-06   1.74e-06   2.14e-06   2.10e-06   2.14e-06   1.45e-06   1.99e-06   2.00E-06 	2.23E-06
+Mg     3.80e-05   3.98e-05   3.80e-05   3.95e-05   3.80e-05   2.51e-05   3.55e-05   3.47E-05 	3.98E-05
+Al     2.95e-06   2.82e-06   2.95e-06   3.12e-06   2.95e-06   2.14e-06   2.88e-06   2.95E-06 	3.27E-06
+Si     3.55e-05   3.24e-05   3.55e-05   3.68e-05   3.55e-05   1.86e-05   3.47e-05   3.31E-05 	3.86E-05
+P      2.82e-07   2.57e-07   2.82e-07   3.82e-07   2.82e-07   2.63e-07   2.88e-07   2.88E-07 	3.20E-07
+S      1.62e-05   1.32e-05   1.62e-05   1.89e-05   2.14e-05   1.23e-05   1.55e-05   1.38E-05 	1.63E-05
+Cl     3.16e-07   3.16e-07   3.16e-07   1.93e-07   3.16e-07   1.32e-07   1.82e-07   3.16E-07 	2.00E-07
+Ar     3.63e-06   2.51e-06   4.47e-06   3.82e-06   2.51e-06   2.57e-06   3.55e-06   3.16E-06 	3.58E-06
+K      1.32e-07   1.07e-07   1.32e-07   1.39e-07   1.32e-07   0.00       1.29e-07   1.32E-07 	1.45E-07
+Ca     2.29e-06   2.19e-06   2.29e-06   2.25e-06   2.29e-06   1.58e-06   2.19e-06   2.14E-06 	2.33E-06
+Sc     1.26e-09   1.41e-09   1.48e-09   1.24e-09   1.48e-09   0.00       1.17e-09   1.26E-09 	1.33E-09
+Ti     9.77e-08   8.91e-08   1.05e-07   8.82e-08   1.05e-07   6.46e-08   8.32e-08   7.94E-08 	9.54E-08
+V      1.00e-08   8.51e-09   1.00e-08   1.08e-08   1.00e-08   0.00       1.00e-08   1.00E-08 	1.11E-08
+Cr     4.68e-07   4.37e-07   4.68e-07   4.93e-07   4.68e-07   3.24e-07   4.47e-07   4.37E-07 	5.06E-07
+Mn     2.45e-07   2.69e-07   2.45e-07   3.50e-07   2.45e-07   2.19e-07   3.16e-07   2.34E-07 	3.56E-07
+Fe     4.68e-05   3.16e-05   3.24e-05   3.31e-05   3.16e-05   2.69e-05   2.95e-05   2.82E-05 	3.27E-05
+Co     8.32e-08   9.77e-08   8.32e-08   8.27e-08   8.32e-08   8.32e-08   8.13e-08   8.32E-08 	9.07E-08
+Ni     1.78e-06   1.66e-06   1.78e-06   1.81e-06   1.78e-06   1.12e-06   1.66e-06   1.70E-06 	1.89E-06
+Cu     1.62e-08   1.55e-08   1.62e-08   1.89e-08   1.62e-08   0.00       1.82e-08   1.62E-08 	2.09E-08
+Zn     3.98e-08   3.63e-08   3.98e-08   4.63e-08   3.98e-08   0.00       4.27e-08   4.17E-08 	5.02E-08