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

jaxspec/analysis/_plot.py

@@ -0,0 +1,35 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
+from jax.typing import ArrayLike
+from scipy.stats import nbinom
+
+
+def _plot_poisson_data_with_error(
+    ax: plt.Axes,
+    x_bins: ArrayLike,
+    y: ArrayLike,
+    percentiles: tuple = (16, 84),
+):
+    """
+    Plot Poisson data with error bars. We extrapolate the intrinsic error of the observation assuming a prior rate
+    distributed according to a Gamma RV.
+    """
+    y_low = nbinom.ppf(percentiles[0] / 100, y, 0.5)
+    y_high = nbinom.ppf(percentiles[1] / 100, y, 0.5)
+
+    ax_to_plot = ax.errorbar(
+        np.sqrt(x_bins[0] * x_bins[1]),
+        y,
+        xerr=np.abs(x_bins - np.sqrt(x_bins[0] * x_bins[1])),
+        yerr=[
+            y - y_low,
+            y_high - y,
+        ],
+        color="black",
+        linestyle="none",
+        alpha=0.3,
+        capsize=2,
+    )
+
+    return ax_to_plot

jaxspec/data/obsconf.py

@@ -229,5 +229,39 @@ class ObsConfiguration(xr.Dataset):
             attrs=observation.attrs | instrument.attrs,
         )
 
+    @classmethod
+    def mock_from_instrument(
+        cls,
+        instrument: Instrument,
+        exposure: float,
+        low_energy: float = 1e-20,
+        high_energy: float = 1e20,
+    ):
+        """
+        Create a mock observation configuration from an instrument object. The fake observation will have zero counts.
+
+        Parameters:
+            instrument: The instrument object.
+            exposure: The total exposure of the mock observation.
+            low_energy: The lower bound of the energy range to consider.
+            high_energy: The upper bound of the energy range to consider.
+        """
+
+        n_channels = len(instrument.coords["instrument_channel"])
+
+        observation = Observation.from_matrix(
+            np.zeros(n_channels),
+            sparse.eye(n_channels),
+            np.arange(n_channels),
+            np.zeros(n_channels, dtype=bool),
+            exposure,
+            backratio=np.ones(n_channels),
+            attributes={"description": "Mock observation"} | instrument.attrs,
+        )
+
+        return cls.from_instrument(
+            instrument, observation, low_energy=low_energy, high_energy=high_energy
+        )
+
     def plot_counts(self, **kwargs):
         return self.folded_counts.plot.step(x="e_min_folded", where="post", **kwargs)

jaxspec/fit.py

@@ -10,29 +10,27 @@ import arviz as az
 import haiku as hk
 import jax
 import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
 import numpyro
-import optimistix as optx
 
 from jax import random
 from jax.experimental.sparse import BCOO
-from jax.flatten_util import ravel_pytree
 from jax.random import PRNGKey
 from jax.tree_util import tree_map
 from jax.typing import ArrayLike
 from numpyro.contrib.nested_sampling import NestedSampler
 from numpyro.distributions import Distribution, Poisson, TransformedDistribution
 from numpyro.infer import AIES, ESS, MCMC, NUTS, Predictive
-from numpyro.infer.initialization import init_to_value
 from numpyro.infer.inspect import get_model_relations
 from numpyro.infer.reparam import TransformReparam
-from numpyro.infer.util import constrain_fn, log_density
-from scipy.stats import Covariance, multivariate_normal
+from numpyro.infer.util import log_density
 
+from .analysis._plot import _plot_poisson_data_with_error
 from .analysis.results import FitResult
 from .data import ObsConfiguration
 from .model.abc import SpectralModel
 from .model.background import BackgroundModel
-from .util import catchtime
 from .util.typing import PriorDictModel, PriorDictType
 
 
@@ -101,27 +99,6 @@ def build_numpyro_model_for_single_obs(
     return numpyro_model
 
 
-def filter_inference_data(
-    inference_data, observation_container, background_model=None
-) -> az.InferenceData:
-    predictive_parameters = []
-
-    for key, value in observation_container.items():
-        if background_model is not None:
-            predictive_parameters.append(f"obs_{key}")
-            predictive_parameters.append(f"bkg_{key}")
-        else:
-            predictive_parameters.append(f"obs_{key}")
-
-    inference_data.posterior_predictive = inference_data.posterior_predictive[predictive_parameters]
-
-    parameters = [x for x in inference_data.posterior.keys() if not x.endswith("_base")]
-    inference_data.posterior = inference_data.posterior[parameters]
-    inference_data.prior = inference_data.prior[parameters]
-
-    return inference_data
-
-
 class CountForwardModel(hk.Module):
     """
     A haiku module which allows to build the function that simulates the measured counts
@@ -154,7 +131,8 @@ class CountForwardModel(hk.Module):
 
 class BayesianModel:
     """
-    Class to fit a model to a given set of observation.
+    Base class for a Bayesian model. This class contains the necessary methods to build a model, sample from the prior
+    and compute the log-likelihood and posterior probability.
     """
 
     def __init__(
@@ -166,6 +144,8 @@ class BayesianModel:
         sparsify_matrix: bool = False,
     ):
         """
+        Build a Bayesian model for a given spectral model and observations.
+
         Parameters:
             model: the spectral model to fit.
             prior_distributions: a nested dictionary containing the prior distributions for the model parameters, or a
@@ -191,7 +171,7 @@ class BayesianModel:
             prior_distributions_func = prior_distributions
 
         self.prior_distributions_func = prior_distributions_func
-        self.init_params = self.get_initial_params()
+        self.init_params = self.prior_samples()
 
     @cached_property
     def observation_container(self) -> dict[str, ObsConfiguration]:
@@ -215,9 +195,6 @@ class BayesianModel:
     def numpyro_model(self) -> Callable:
         """
         Build the numpyro model using the observed data, the prior distributions and the spectral model.
-
-        Returns:
-            A model function that can be used with numpyro.
         """
 
         def model(observed=True):
@@ -257,9 +234,6 @@ class BayesianModel:
     def log_likelihood_per_obs(self) -> Callable:
         """
         Build the log likelihood function for each bins in each observation.
-
-        Returns:
-            Callable log-likelihood function.
         """
 
         @jax.jit
@@ -316,6 +290,16 @@ class BayesianModel:
         observed_sites = relations["observed"]
         return [site for site in all_sites if site not in observed_sites]
 
+    @cached_property
+    def observation_names(self) -> list[str]:
+        """
+        List of the observations.
+        """
+        relations = get_model_relations(self.numpyro_model)
+        all_sites = relations["sample_sample"].keys()
+        observed_sites = relations["observed"]
+        return [site for site in all_sites if site in observed_sites]
+
     def array_to_dict(self, theta):
         """
         Convert an array of parameters to a dictionary of parameters.
@@ -339,7 +323,7 @@ class BayesianModel:
 
         return theta
 
-    def get_initial_params(self, key: PRNGKey = PRNGKey(0), num_samples: int = 1):
+    def prior_samples(self, key: PRNGKey = PRNGKey(0), num_samples: int = 100):
         """
         Get initial parameters for the model by sampling from the prior distribution
 
@@ -348,9 +332,84 @@ class BayesianModel:
             num_samples: the number of samples to draw from the prior.
         """
 
-        return Predictive(
-            self.numpyro_model, return_sites=self.parameter_names, num_samples=num_samples
-        )(key, observed=False)
+        @jax.jit
+        def prior_sample(key):
+            return Predictive(
+                self.numpyro_model, return_sites=self.parameter_names, num_samples=num_samples
+            )(key, observed=False)
+
+        return prior_sample(key)
+
+    def mock_observations(self, parameters, key: PRNGKey = PRNGKey(0)):
+        @jax.jit
+        def fakeit(key, parameters):
+            return Predictive(
+                self.numpyro_model,
+                return_sites=self.observation_names,
+                posterior_samples=parameters,
+            )(key, observed=False)
+
+        return fakeit(key, parameters)
+
+    def prior_predictive_coverage(
+        self, key: PRNGKey = PRNGKey(0), num_samples: int = 1000, percentiles: tuple = (16, 84)
+    ):
+        """
+        Check if the prior distribution include the observed data.
+        """
+        key_prior, key_posterior = jax.random.split(key, 2)
+        prior_params = self.prior_samples(key=key_prior, num_samples=num_samples)
+        posterior_observations = self.mock_observations(prior_params, key=key_posterior)
+
+        for key, value in self.observation_container.items():
+            fig, axs = plt.subplots(
+                nrows=2, ncols=1, sharex=True, figsize=(8, 8), height_ratios=[3, 1]
+            )
+
+            _plot_poisson_data_with_error(
+                axs[0],
+                value.out_energies,
+                value.folded_counts.values,
+                percentiles=percentiles,
+            )
+
+            axs[0].stairs(
+                np.max(posterior_observations["obs_" + key], axis=0),
+                edges=[*list(value.out_energies[0]), value.out_energies[1][-1]],
+                baseline=np.min(posterior_observations["obs_" + key], axis=0),
+                alpha=0.3,
+                fill=True,
+                color=(0.15, 0.25, 0.45),
+            )
+
+            # rank = np.vstack((posterior_observations["obs_" + key], value.folded_counts.values)).argsort(axis=0)[-1] / (num_samples) * 100
+            counts = posterior_observations["obs_" + key]
+            observed = value.folded_counts.values
+
+            num_samples = counts.shape[0]
+
+            less_than_obs = (counts < observed).sum(axis=0)
+            equal_to_obs = (counts == observed).sum(axis=0)
+
+            rank = (less_than_obs + 0.5 * equal_to_obs) / num_samples * 100
+
+            axs[1].stairs(rank, edges=[*list(value.out_energies[0]), value.out_energies[1][-1]])
+
+            axs[1].plot(
+                (value.out_energies.min(), value.out_energies.max()),
+                (50, 50),
+                color="black",
+                linestyle="--",
+            )
+
+            axs[1].set_xlabel("Energy (keV)")
+            axs[0].set_ylabel("Counts")
+            axs[1].set_ylabel("Rank (%)")
+            axs[1].set_ylim(0, 100)
+            axs[0].set_xlim(value.out_energies.min(), value.out_energies.max())
+            axs[0].loglog()
+            plt.suptitle(f"Prior Predictive coverage for {key}")
+            plt.show()
 
 
 class BayesianModelFitter(BayesianModel, ABC):
@@ -359,11 +418,20 @@ class BayesianModelFitter(BayesianModel, ABC):
         posterior_samples,
         num_chains: int = 1,
         num_predictive_samples: int = 1000,
-        key: PRNGKey = PRNGKey(0),
+        key: PRNGKey = PRNGKey(42),
         use_transformed_model: bool = False,
+        filter_inference_data: bool = True,
     ) -> az.InferenceData:
         """
-        Build an InferenceData object from the posterior samples.
+        Build an [InferenceData][arviz.InferenceData] object from posterior samples.
+
+        Parameters:
+            posterior_samples: the samples from the posterior distribution.
+            num_chains: the number of chains used to sample the posterior.
+            num_predictive_samples: the number of samples to draw from the prior.
+            key: the random key used to initialize the sampler.
+            use_transformed_model: whether to use the transformed model to build the InferenceData.
+            filter_inference_data: whether to filter the InferenceData to keep only the relevant parameters.
         """
 
         numpyro_model = (
@@ -409,7 +477,7 @@ class BayesianModelFitter(BayesianModel, ABC):
             key: reshape_first_dimension(value) for key, value in log_likelihood.items()
         }
 
-        return az.from_dict(
+        inference_data = az.from_dict(
             posterior_samples,
             prior=prior,
             posterior_predictive=posterior_predictive,
@@ -417,81 +485,42 @@ class BayesianModelFitter(BayesianModel, ABC):
             observed_data=observations,
         )
 
-    @abstractmethod
-    def fit(self, **kwargs) -> FitResult: ...
-
-
-class NUTSFitter(BayesianModelFitter):
-    """
-    A class to fit a model to a given set of observation using a Bayesian approach. This class uses the NUTS sampler
-    from numpyro to perform the inference on the model parameters.
-    """
+        return (
+            self.filter_inference_data(inference_data) if filter_inference_data else inference_data
+        )
 
-    def fit(
+    def filter_inference_data(
         self,
-        rng_key: int = 0,
-        num_chains: int = len(jax.devices()),
-        num_warmup: int = 1000,
-        num_samples: int = 1000,
-        max_tree_depth: int = 10,
-        target_accept_prob: float = 0.8,
-        dense_mass: bool = False,
-        kernel_kwargs: dict = {},
-        mcmc_kwargs: dict = {},
-    ) -> FitResult:
+        inference_data: az.InferenceData,
+    ) -> az.InferenceData:
         """
-        Fit the model to the data using NUTS sampler from numpyro.
+        Filter the inference data to keep only the relevant parameters for the observations.
 
-        Parameters:
-            rng_key: the random key used to initialize the sampler.
-            num_chains: the number of chains to run.
-            num_warmup: the number of warmup steps.
-            num_samples: the number of samples to draw.
-            max_tree_depth: the recursion depth of NUTS sampler.
-            target_accept_prob: the target acceptance probability for the NUTS sampler.
-            dense_mass: whether to use a dense mass for the NUTS sampler.
-            mcmc_kwargs: additional arguments to pass to the MCMC sampler. See [`MCMC`][numpyro.infer.mcmc.MCMC] for more details.
-
-        Returns:
-            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
+        - Removes predictive parameters from deterministic random variables (e.g. kernel of background GP)
+        - Removes parameters build from reparametrised variables (e.g. ending with `"_base"`)
         """
 
-        bayesian_model = self.transformed_numpyro_model
-        # bayesian_model = self.numpyro_model(prior_distributions)
-
-        chain_kwargs = {
-            "num_warmup": num_warmup,
-            "num_samples": num_samples,
-            "num_chains": num_chains,
-        }
-
-        kernel = NUTS(
-            bayesian_model,
-            max_tree_depth=max_tree_depth,
-            target_accept_prob=target_accept_prob,
-            dense_mass=dense_mass,
-            **kernel_kwargs,
-        )
+        predictive_parameters = []
 
-        mcmc = MCMC(kernel, **(chain_kwargs | mcmc_kwargs))
-        keys = random.split(random.PRNGKey(rng_key), 3)
+        for key, value in self.observation_container.items():
+            if self.background_model is not None:
+                predictive_parameters.append(f"obs_{key}")
+                predictive_parameters.append(f"bkg_{key}")
+            else:
+                predictive_parameters.append(f"obs_{key}")
 
-        mcmc.run(keys[0])
+        inference_data.posterior_predictive = inference_data.posterior_predictive[
+            predictive_parameters
+        ]
 
-        posterior = mcmc.get_samples()
+        parameters = [x for x in inference_data.posterior.keys() if not x.endswith("_base")]
+        inference_data.posterior = inference_data.posterior[parameters]
+        inference_data.prior = inference_data.prior[parameters]
 
-        inference_data = filter_inference_data(
-            self.build_inference_data(posterior, num_chains=num_chains),
-            self.observation_container,
-            self.background_model,
-        )
+        return inference_data
 
-        return FitResult(
-            self,
-            inference_data,
-            self.model.params,
-            background_model=self.background_model,
-        )
+    @abstractmethod
+    def fit(self, **kwargs) -> FitResult: ...
 
 
 class MCMCFitter(BayesianModelFitter):
@@ -513,6 +542,7 @@ class MCMCFitter(BayesianModelFitter):
         num_warmup: int = 1000,
         num_samples: int = 1000,
         sampler: Literal["nuts", "aies", "ess"] = "nuts",
+        use_transformed_model: bool = True,
         kernel_kwargs: dict = {},
         mcmc_kwargs: dict = {},
     ) -> FitResult:
@@ -524,17 +554,18 @@ class MCMCFitter(BayesianModelFitter):
             num_chains: the number of chains to run.
             num_warmup: the number of warmup steps.
             num_samples: the number of samples to draw.
-            max_tree_depth: the recursion depth of NUTS sampler.
-            target_accept_prob: the target acceptance probability for the NUTS sampler.
-            dense_mass: whether to use a dense mass for the NUTS sampler.
+            sampler: the sampler to use. Can be one of "nuts", "aies" or "ess".
+            use_transformed_model: whether to use the transformed model to build the InferenceData.
+            kernel_kwargs: additional arguments to pass to the kernel. See [`NUTS`][numpyro.infer.mcmc.MCMCKernel] for more details.
             mcmc_kwargs: additional arguments to pass to the MCMC sampler. See [`MCMC`][numpyro.infer.mcmc.MCMC] for more details.
 
         Returns:
             A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
         """
 
-        bayesian_model = self.transformed_numpyro_model
-        # bayesian_model = self.numpyro_model(prior_distributions)
+        bayesian_model = (
+            self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
+        )
 
         chain_kwargs = {
             "num_warmup": num_warmup,
@@ -557,10 +588,8 @@ class MCMCFitter(BayesianModelFitter):
 
         posterior = mcmc.get_samples()
 
-        inference_data = filter_inference_data(
-            self.build_inference_data(posterior, num_chains=num_chains),
-            self.observation_container,
-            self.background_model,
+        inference_data = self.build_inference_data(
+            posterior, num_chains=num_chains, use_transformed_model=True
         )
 
         return FitResult(
@@ -571,175 +600,22 @@ class MCMCFitter(BayesianModelFitter):
         )
 
 
-class MinimizationFitter(BayesianModelFitter):
-    """
-    A class to fit a model to a given set of observation using a minimization algorithm. This class uses the L-BFGS
-    algorithm from jaxopt to perform the minimization on the model parameters. The uncertainties are computed using the
-    Hessian of the log-log_likelihood, assuming that it is a multivariate Gaussian in the unbounded space defined by
-    numpyro.
-    """
-
-    def fit(
-        self,
-        rng_key: int = 0,
-        num_iter_max: int = 100_000,
-        num_samples: int = 1_000,
-        solver: Literal["bfgs", "levenberg_marquardt"] = "bfgs",
-        init_params=None,
-        refine_first_guess=True,
-    ) -> FitResult:
-        """
-        Fit the model to the data using L-BFGS algorithm.
-
-        Parameters:
-            rng_key: the random key used to initialize the sampler.
-            num_iter_max: the maximum number of iteration in the minimization algorithm.
-            num_samples: the number of sample to draw from the best-fit covariance.
-
-        Returns:
-            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
-        """
-
-        bayesian_model = self.numpyro_model
-        keys = jax.random.split(PRNGKey(rng_key), 4)
-
-        if init_params is not None:
-            # We initialize the parameters by randomly sampling from the prior
-            local_keys = jax.random.split(keys[0], 2)
-
-            with numpyro.handlers.seed(rng_seed=local_keys[0]):
-                starting_value = self.prior_distributions_func()
-
-            # We update the starting value with the provided init_params
-            for m, n, val in hk.data_structures.traverse(init_params):
-                if f"{m}_{n}" in starting_value.keys():
-                    starting_value[f"{m}_{n}"] = val
-
-            init_params, _ = numpyro.infer.util.find_valid_initial_params(
-                local_keys[1], bayesian_model, init_strategy=init_to_value(values=starting_value)
-            )
-
-        else:
-            init_params, _ = numpyro.infer.util.find_valid_initial_params(keys[0], bayesian_model)
-
-        init_params = init_params[0]
-
-        @jax.jit
-        def nll(unconstrained_params, _):
-            constrained_params = constrain_fn(
-                bayesian_model, tuple(), dict(observed=True), unconstrained_params
-            )
-
-            log_likelihood = numpyro.infer.util.log_likelihood(
-                model=bayesian_model, posterior_samples=constrained_params
-            )
-
-            # We solve a least square problem, this function ensure that the total residual is indeed the nll
-            return jax.tree.map(lambda x: jnp.sqrt(-x), log_likelihood)
-
-        """
-        if refine_first_guess:
-            with catchtime("Refine_first"):
-                solution = optx.least_squares(
-                    nll,
-                    optx.BestSoFarMinimiser(optx.OptaxMinimiser(optax.adam(1e-4), 1e-6, 1e-6)),
-                    init_params,
-                    max_steps=1000,
-                    throw=False
-                )
-            init_params = solution.value
-        """
-
-        if solver == "bfgs":
-            solver = optx.BestSoFarMinimiser(optx.BFGS(1e-6, 1e-6))
-        elif solver == "levenberg_marquardt":
-            solver = optx.BestSoFarLeastSquares(optx.LevenbergMarquardt(1e-6, 1e-6))
-        else:
-            raise NotImplementedError(f"{solver} is not implemented")
-
-        with catchtime("Minimization"):
-            solution = optx.least_squares(
-                nll,
-                solver,
-                init_params,
-                max_steps=num_iter_max,
-            )
-
-        params = solution.value
-        value_flat, unflatten_fun = ravel_pytree(params)
-
-        with catchtime("Compute error"):
-            precision = jax.hessian(
-                lambda p: jnp.sum(ravel_pytree(nll(unflatten_fun(p), None))[0] ** 2)
-            )(value_flat)
-
-            cov = Covariance.from_precision(precision)
-
-            samples_flat = multivariate_normal.rvs(mean=value_flat, cov=cov, size=num_samples)
-
-        samples = jax.vmap(unflatten_fun)(samples_flat)
-        posterior_samples = jax.jit(
-            jax.vmap(lambda p: constrain_fn(bayesian_model, tuple(), dict(observed=True), p))
-        )(samples)
-
-        with catchtime("Posterior"):
-            posterior_predictive = Predictive(bayesian_model, posterior_samples)(
-                keys[2], observed=False
-            )
-            prior = Predictive(bayesian_model, num_samples=num_samples)(keys[3], observed=False)
-            log_likelihood = numpyro.infer.log_likelihood(bayesian_model, posterior_samples)
-
-        def sanitize_chain(chain):
-            """
-            reshape the samples so that it is arviz compliant with an extra starting dimension
-            """
-            return tree_map(lambda x: x[None, ...], chain)
-
-        # We export the observed values to the inference_data
-        seeded_model = numpyro.handlers.substitute(
-            numpyro.handlers.seed(bayesian_model, jax.random.PRNGKey(0)),
-            substitute_fn=numpyro.infer.init_to_sample,
-        )
-        trace = numpyro.handlers.trace(seeded_model).get_trace()
-        observations = {
-            name: site["value"]
-            for name, site in trace.items()
-            if site["type"] == "sample" and site["is_observed"]
-        }
-
-        with catchtime("InferenceData wrapping"):
-            inference_data = az.from_dict(
-                sanitize_chain(posterior_samples),
-                prior=sanitize_chain(prior),
-                posterior_predictive=sanitize_chain(posterior_predictive),
-                log_likelihood=sanitize_chain(log_likelihood),
-                observed_data=observations,
-            )
-
-        inference_data = filter_inference_data(
-            inference_data, self.observation_container, self.background_model
-        )
-
-        return FitResult(
-            self,
-            inference_data,
-            self.model.params,
-            background_model=self.background_model,
-        )
-
-
-class NestedSamplingFitter(BayesianModelFitter):
+class NSFitter(BayesianModelFitter):
     r"""
     A class to fit a model to a given set of observation using the Nested Sampling algorithm. This class uses the
     [`DefaultNestedSampler`][jaxns.DefaultNestedSampler] from [`jaxns`](https://jaxns.readthedocs.io/en/latest/) which
     implements the [Phantom-Powered Nested Sampling](https://arxiv.org/abs/2312.11330) algorithm.
-    Add Citation to jaxns
+
+    !!! info
+        Ensure large prior volume is covered by the prior distributions to ensure the algorithm yield proper results.
+
     """
 
     def fit(
         self,
         rng_key: int = 0,
         num_samples: int = 1000,
+        num_live_points: int = 1000,
         plot_diagnostics=False,
         termination_kwargs: dict | None = None,
         verbose=True,
@@ -750,6 +626,10 @@ class NestedSamplingFitter(BayesianModelFitter):
         Parameters:
             rng_key: the random key used to initialize the sampler.
             num_samples: the number of samples to draw.
+            num_live_points: the number of live points to use at the start of the NS algorithm.
+            plot_diagnostics: whether to plot the diagnostics of the NS algorithm.
+            termination_kwargs: additional arguments to pass to the termination criterion of the NS algorithm.
+            verbose: whether to print the progress of the NS algorithm.
 
         Returns:
             A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
@@ -766,7 +646,7 @@ class NestedSamplingFitter(BayesianModelFitter):
                 difficult_model=True,
                 max_samples=1e6,
                 parameter_estimation=True,
-                num_live_points=1_000,
+                num_live_points=num_live_points,
             ),
             termination_kwargs=termination_kwargs if termination_kwargs else dict(),
         )
@@ -776,41 +656,9 @@ class NestedSamplingFitter(BayesianModelFitter):
         if plot_diagnostics:
             ns.diagnostics()
 
-        posterior_samples = ns.get_samples(keys[1], num_samples=num_samples)
-        log_likelihood = numpyro.infer.log_likelihood(bayesian_model, posterior_samples)
-        posterior_predictive = Predictive(bayesian_model, posterior_samples)(
-            keys[2], observed=False
-        )
-
-        prior = Predictive(bayesian_model, num_samples=num_samples)(keys[3], observed=False)
-
-        seeded_model = numpyro.handlers.substitute(
-            numpyro.handlers.seed(bayesian_model, jax.random.PRNGKey(0)),
-            substitute_fn=numpyro.infer.init_to_sample,
-        )
-        trace = numpyro.handlers.trace(seeded_model).get_trace()
-        observations = {
-            name: site["value"]
-            for name, site in trace.items()
-            if site["type"] == "sample" and site["is_observed"]
-        }
-
-        def sanitize_chain(chain):
-            """
-            reshape the samples so that it is arviz compliant with an extra starting dimension
-            """
-            return tree_map(lambda x: x[None, ...], chain)
-
-        inference_data = az.from_dict(
-            sanitize_chain(posterior_samples),
-            prior=sanitize_chain(prior),
-            posterior_predictive=sanitize_chain(posterior_predictive),
-            log_likelihood=sanitize_chain(log_likelihood),
-            observed_data=observations,
-        )
-
-        inference_data = filter_inference_data(
-            inference_data, self.observation_container, self.background_model
+        posterior = ns.get_samples(keys[1], num_samples=num_samples)
+        inference_data = self.build_inference_data(
+            posterior, num_chains=1, use_transformed_model=True
         )
 
         return FitResult(

jaxspec/model/abc.py

@@ -7,9 +7,11 @@ import haiku as hk
 import jax
 import jax.numpy as jnp
 import networkx as nx
+import rich
 
 from haiku._src import base
 from jax.scipy.integrate import trapezoid
+from rich.table import Table
 from simpleeval import simple_eval
 
 
@@ -110,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.

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
 
 

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxspec
-Version: 0.1.1
+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.27)
 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.1.dist-info → jaxspec-0.1.2.dist-info}/RECORD

@@ -1,21 +1,19 @@
 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=lfeqn1HFlZSFSvVbBL0obV9cCwMlbBWpJz_nw7JI0WY,29552
+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=nQZUmtUzXjW94gv3BJg1lHXHZtgrHoOlAR4a6G2a9VQ,20234
-jaxspec/model/additive.py,sha256=xD5E30nd5pqa-swQireA52ch1czxnqRosnh-dsp5xL0,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=GCQ6JRz92QqbzDBFwWxGZ9SUqTJZQpD7B6ji9VEFXWo,8135
@@ -26,8 +24,8 @@ jaxspec/util/abundance.py,sha256=fsC313taIlGzQsZNwbYsJupDWm7ZbqzGhY66Ku394Mw,854
 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=8qK1aJlsqTcVKjYN-BxsDx20BTwtnS-wMw6Bdurpm-o,2459
-jaxspec-0.1.1.dist-info/LICENSE.md,sha256=2q5XoWzddts5IqzIcgYYMOL21puU3MfO8gvT3Ype1eQ,1073
-jaxspec-0.1.1.dist-info/METADATA,sha256=vQMbkUPdTyiuLuHxEGLEitAg-2GAk65lvIPQJDfutY8,3572
-jaxspec-0.1.1.dist-info/WHEEL,sha256=sP946D7jFCHeNz5Iq4fL4Lu-PrWrFsgfLXbbkciIZwg,88
-jaxspec-0.1.1.dist-info/entry_points.txt,sha256=kzLG2mGlCWITRn4Q6zKG_idx-_RKAncvA0DMNYTgHAg,71
-jaxspec-0.1.1.dist-info/RECORD,,
+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

jaxspec/model/_additive/apec_loaders.py

@@ -1,73 +0,0 @@
-"""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
-
-from ...util.online_storage import table_manager
-
-
-@jax.jit
-def temperature_table_getter():
-    with h5netcdf.File(table_manager.fetch("apec.nc"), "r") as f:
-        temperature = jnp.asarray(f["/temperature"])
-
-    return temperature
-
-
-@jax.jit
-def get_temperature(kT):
-    temperature = temperature_table_getter()
-    idx = jnp.searchsorted(temperature, kT) - 1
-
-    return idx, temperature[idx], temperature[idx + 1]
-
-
-@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"])
-
-    return continuum_energy, continuum_emissivity, continuum_end_index
-
-
-@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"])
-
-    return pseudo_energy, pseudo_emissivity, pseudo_end_index
-
-
-@jax.jit
-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):
-    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):
-    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):
-    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]