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

jaxspec/model/_additive/apec.py

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

jaxspec/model/_additive/apec_loaders.py

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