jaxspec 0.3.1__tar.gz → 0.3.3__tar.gz

{jaxspec-0.3.1 → jaxspec-0.3.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: jaxspec
-Version: 0.3.1
+Version: 0.3.3
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 Project-URL: Homepage, https://github.com/renecotyfanboy/jaxspec
 Project-URL: Documentation, https://jaxspec.readthedocs.io/en/latest/
@@ -16,19 +16,15 @@ Requires-Dist: cmasher<2,>=1.6.3
 Requires-Dist: flax>0.10.5
 Requires-Dist: interpax<0.4,>=0.3.5
 Requires-Dist: jax<0.7,>=0.5.0
-Requires-Dist: jaxns<3,>=2.6.7
-Requires-Dist: jaxopt<0.9,>=0.8.3
 Requires-Dist: matplotlib<4,>=3.8.0
 Requires-Dist: mendeleev<1.2,>=0.15
 Requires-Dist: networkx~=3.1
 Requires-Dist: numpy<3.0.0
 Requires-Dist: numpyro<0.20,>=0.17.0
-Requires-Dist: optimistix<0.0.11,>=0.0.10
 Requires-Dist: pandas<3,>=2.2.0
 Requires-Dist: pooch<2,>=1.8.2
 Requires-Dist: scipy<1.16
-Requires-Dist: seaborn<0.14,>=0.13.1
-Requires-Dist: simpleeval<1.1.0,>=0.9.13
+Requires-Dist: seaborn>=0.13.2
 Requires-Dist: sparse>0.15
 Requires-Dist: tinygp<0.4,>=0.3.0
 Requires-Dist: watermark<3,>=2.4.3

{jaxspec-0.3.1 → jaxspec-0.3.3}/docs/examples/index.md

@@ -54,5 +54,11 @@
 
     Use `jaxspec` likelihood in other packages
 
+-   [__Cross-calibration between instruments__](calibration.ipynb)
+
+    ---
+
+    Add extra gains and shifts to account for the miscalibrations
+
 
 </div>

{jaxspec-0.3.1 → jaxspec-0.3.3}/docs/references/fitting.md

@@ -13,3 +13,11 @@
       show_root_full_path: false
       show_root_toc_entry: true
       heading_level: 3
+
+::: jaxspec.fit.VIFitter
+    options:
+      show_root_heading: true
+      show_root_full_path: false
+      show_root_toc_entry: true
+      heading_level: 3
+

jaxspec-0.3.3/docs/references/instrument.md

@@ -0,0 +1,4 @@
+::: jaxspec.model.instrument
+    options:
+      show_root_heading: false
+      show_root_toc_entry: false

{jaxspec-0.3.1 → jaxspec-0.3.3}/mkdocs.yml

@@ -18,14 +18,13 @@ nav:
     - Add a background : examples/background.md
     - Good practices for MCMC: examples/work_with_arviz.ipynb
     - Interface with other frameworks : examples/external_samplers.ipynb
-  #- Models:
-  #    - models/index.md
-  #    - APEC: models/apec.md
+    - Cross calibration : examples/calibration.ipynb
   - API Reference:
     - Spectral model base: references/model.md
     - Additive models: references/additive.md
     - Multiplicative models: references/multiplicative.md
     - Background models: references/background.md
+    - Instrument models: references/instrument.md
     - Data containers: references/data.md
     - Fitting: references/fitting.md
     - Result containers: references/results.md
@@ -87,7 +86,7 @@ theme:
 plugins:
   - search
   - autorefs
-#  - typeset
+  - typeset
   - mkdocs-jupyter:
       include_source: True
       ignore_h1_titles: True

{jaxspec-0.3.1 → jaxspec-0.3.3}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "jaxspec"
-version = "0.3.1"
+version = "0.3.3"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = [{ name = "sdupourque", email = "sdupourque@irap.omp.eu" }]
 requires-python = ">=3.10,<3.13"
@@ -16,21 +16,17 @@ dependencies = [
     "matplotlib>=3.8.0,<4",
     "arviz>=0.17.1,<0.23.0",
     "chainconsumer>=1.1.2,<2",
-    "simpleeval>=0.9.13,<1.1.0",
     "cmasher>=1.6.3,<2",
-    "jaxopt>=0.8.3,<0.9",
     "tinygp>=0.3.0,<0.4",
-    "seaborn>=0.13.1,<0.14",
     "sparse>0.15",
-    "optimistix>=0.0.10,<0.0.11",
     "scipy<1.16",
     "mendeleev>=0.15,<1.2",
-    "jaxns>=2.6.7,<3",
     "pooch>=1.8.2,<2",
     "interpax>=0.3.5,<0.4",
     "watermark>=2.4.3,<3",
     "catppuccin>=2.3.4,<3",
     "flax>0.10.5",
+    "seaborn>=0.13.2",
 ]
 
 [project.urls]
@@ -44,14 +40,14 @@ jaxspec-debug-info = "jaxspec.scripts.debug:debug_info"
 docs = [
     "mkdocs>=1.6.1,<2",
     "mkdocs-material>=9.4.6,<10",
-    "mkdocstrings[python]>=0.24,<0.28",
+    "mkdocstrings[python]>=0.24,<1.1",
     "mkdocs-jupyter>=0.25.0,<0.26",
 ]
 test = [
     "chex>=0.1.83,<0.2",
     "mktestdocs>=0.2.1,<0.3",
     "coverage>=7.3.2,<8",
-    "pytest-cov>=4.1,<7.0",
+    "pytest-cov>=4.1,<8.0",
     "flake8>=7.0.0,<8",
     "pytest>=8.0.0,<9",
     "testbook>=0.4.2,<0.5",

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/analysis/results.py

@@ -77,8 +77,9 @@ class FitResult:
         r"""
         Convergence of the chain as computed by the $\hat{R}$ statistic.
         """
+        rhat = az.rhat(self.inference_data)
 
-        return all(az.rhat(self.inference_data) < 1.01)
+        return bool((rhat.to_array() < 1.01).all())
 
     def _ppc_folded_branches(self, obs_id):
         obs = self.obsconfs[obs_id]
@@ -167,6 +168,7 @@ class FitResult:
         e_max: float,
         unit: Unit = u.photon / u.cm**2 / u.s,
         register: bool = False,
+        n_points: int = 100,
     ) -> ArrayLike:
         """
         Compute the unfolded photon flux in a given energy band. The flux is then added to
@@ -177,6 +179,7 @@ class FitResult:
             e_max: The upper bound of the energy band in observer frame.
             unit: The unit of the photon flux.
             register: Whether to register the flux with the other posterior parameters.
+            n_points: The number of points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
@@ -188,18 +191,25 @@ class FitResult:
         def vectorized_flux(*pars):
             parameters_pytree = jax.tree.unflatten(pytree_def, pars)
             return self.model.photon_flux(
-                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=100
+                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=n_points
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree)
-        conversion_factor = (u.photon / u.cm**2 / u.s).to(unit)
-        value = flux * conversion_factor
+        conversion_factor = float((u.photon / u.cm**2 / u.s).to(unit))
+        value = np.asarray(flux * conversion_factor)
 
         if register:
-            self.inference_data.posterior[f"photon_flux_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/photon_flux_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/photon_flux_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -210,6 +220,7 @@ class FitResult:
         e_max: float,
         unit: Unit = u.erg / u.cm**2 / u.s,
         register: bool = False,
+        n_points: int = 100,
     ) -> ArrayLike:
         """
         Compute the unfolded energy flux in a given energy band. The flux is then added to
@@ -220,6 +231,7 @@ class FitResult:
             e_max: The upper bound of the energy band in observer frame.
             unit: The unit of the energy flux.
             register: Whether to register the flux with the other posterior parameters.
+            n_points: The number of points to use for computing the unfolded spectrum.
 
         !!! warning
             Computation of the folded flux is not implemented yet. Feel free to open an
@@ -231,18 +243,25 @@ class FitResult:
         def vectorized_flux(*pars):
             parameters_pytree = jax.tree.unflatten(pytree_def, pars)
             return self.model.energy_flux(
-                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=100
+                parameters_pytree, jnp.asarray([e_min]), jnp.asarray([e_max]), n_points=n_points
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree)
-        conversion_factor = (u.keV / u.cm**2 / u.s).to(unit)
-        value = flux * conversion_factor
+        conversion_factor = float((u.keV / u.cm**2 / u.s).to(unit))
+        value = np.asarray(flux * conversion_factor)
 
         if register:
-            self.inference_data.posterior[f"energy_flux_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -257,6 +276,7 @@ class FitResult:
         cosmology: Cosmology = Planck18,
         unit: Unit = u.erg / u.s,
         register: bool = False,
+        n_points=100,
     ) -> ArrayLike:
         """
         Compute the luminosity of the source specifying its redshift. The luminosity is then added to
@@ -294,17 +314,26 @@ class FitResult:
                 parameters_pytree,
                 jnp.asarray([e_min]) * (1 + redshift),
                 jnp.asarray([e_max]) * (1 + redshift),
-                n_points=100,
+                n_points=n_points,
             )[0]
 
         flat_tree, pytree_def = jax.tree.flatten(self.input_parameters)
         flux = vectorized_flux(*flat_tree) * (u.keV / u.cm**2 / u.s)
-        value = (flux * (4 * np.pi * cosmology.luminosity_distance(redshift) ** 2)).to(unit)
+        value = np.asarray(
+            (flux * (4 * np.pi * cosmology.luminosity_distance(redshift) ** 2)).to(unit)
+        )
 
         if register:
-            self.inference_data.posterior[f"luminosity_{e_min:.1f}_{e_max:.1f}"] = (
-                list(self.inference_data.posterior.coords),
-                value,
+            self.inference_data.posterior[f"mod/~/luminosity_{e_min:.1f}_{e_max:.1f}"] = (
+                xr.DataArray(
+                    value,
+                    dims=self.inference_data.posterior.dims,
+                    coords={
+                        "chain": self.inference_data.posterior.chain,
+                        "draw": self.inference_data.posterior.draw,
+                    },
+                    name=f"mod/~/luminosity_{e_min:.1f}_{e_max:.1f}",
+                )
             )
 
         return value
@@ -315,10 +344,13 @@ class FitResult:
 
         Parameters:
             name: The name of the chain.
+            parameter_kind: The kind of parameters to keep.
         """
 
         keys_to_drop = [
-            key for key in self.inference_data.posterior.keys() if not key.startswith("mod")
+            key
+            for key in self.inference_data.posterior.keys()
+            if not key.startswith(parameter_kind)
         ]
 
         reduced_id = az.extract(
@@ -403,6 +435,7 @@ class FitResult:
         title: str | None = None,
         figsize: tuple[float, float] = (6, 6),
         x_lims: tuple[float, float] | None = None,
+        rescale_background: bool = False,
     ) -> list[plt.Figure]:
         r"""
         Plot the posterior predictive distribution of the model. It also features a residual plot, defined using the
@@ -423,6 +456,7 @@ class FitResult:
             title: The title of the plot.
             figsize: The size of the figure.
             x_lims: The limits of the x-axis.
+            rescale_background: Whether to rescale the background model to the data with backscal ratio.
 
         Returns:
             A list of matplotlib figures for each observation in the model.
@@ -573,10 +607,14 @@ class FitResult:
                         )
                     )
 
+                    rescale_background_factor = (
+                        obsconf.folded_backratio.data if rescale_background else 1.0
+                    )
+
                     model_bkg_plot = _plot_binned_samples_with_error(
                         ax[0],
                         xbins.value,
-                        y_samples_bkg.value,
+                        y_samples_bkg.value * rescale_background_factor,
                         color=BACKGROUND_COLOR,
                         alpha_envelope=alpha_envelope,
                         n_sigmas=n_sigmas,
@@ -585,9 +623,9 @@ class FitResult:
                     true_bkg_plot = _plot_poisson_data_with_error(
                         ax[0],
                         xbins.value,
-                        y_observed_bkg.value,
-                        y_observed_bkg_low.value,
-                        y_observed_bkg_high.value,
+                        y_observed_bkg.value * rescale_background_factor,
+                        y_observed_bkg_low.value * rescale_background_factor,
+                        y_observed_bkg_high.value * rescale_background_factor,
                         color=BACKGROUND_DATA_COLOR,
                         alpha=0.7,
                     )

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/data/instrument.py

@@ -1,4 +1,5 @@
 import os
+import sparse
 
 import numpy as np
 import xarray as xr
@@ -92,7 +93,7 @@ class Instrument(xr.Dataset):
 
         else:
             specresp = rmf.matrix.sum(axis=0)
-            rmf.sparse_matrix /= specresp
+            rmf.sparse_matrix = sparse.COO( rmf.matrix / specresp )
 
         return cls.from_matrix(
             rmf.sparse_matrix, specresp, rmf.energ_lo, rmf.energ_hi, rmf.e_min, rmf.e_max

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/data/obsconf.py

@@ -163,10 +163,8 @@ class ObsConfiguration(xr.Dataset):
 
         if observation.folded_background is not None:
             folded_background = observation.folded_background.data[row_idx]
-            folded_background_unscaled = observation.folded_background_unscaled.data[row_idx]
         else:
             folded_background = np.zeros_like(folded_counts)
-            folded_background_unscaled = np.zeros_like(folded_counts)
 
         data_dict = {
             "transfer_matrix": (
@@ -208,14 +206,6 @@ class ObsConfiguration(xr.Dataset):
                     "unit": "photons",
                 },
             ),
-            "folded_background_unscaled": (
-                ["folded_channel"],
-                folded_background_unscaled,
-                {
-                    "description": "To be done",
-                    "unit": "photons",
-                },
-            ),
         }
 
         return cls(

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/data/observation.py

@@ -46,16 +46,14 @@ class Observation(xr.Dataset):
         quality,
         exposure,
         background=None,
-        background_unscaled=None,
         backratio=1.0,
         attributes: dict | None = None,
     ):
         if attributes is None:
             attributes = {}
 
-        if background is None or background_unscaled is None:
+        if background is None:
             background = np.zeros_like(counts, dtype=np.int64)
-            background_unscaled = np.zeros_like(counts, dtype=np.int64)
 
         data_dict = {
             "counts": (
@@ -86,7 +84,9 @@ class Observation(xr.Dataset):
             ),
             "folded_backratio": (
                 ["folded_channel"],
-                np.asarray(np.ma.filled(grouping @ backratio), dtype=float),
+                np.asarray(
+                    np.ma.filled(grouping @ backratio) / grouping.sum(axis=1).todense(), dtype=float
+                ),
                 {"description": "Background scaling after grouping"},
             ),
             "background": (
@@ -94,11 +94,6 @@ class Observation(xr.Dataset):
                 np.asarray(background, dtype=np.int64),
                 {"description": "Background counts", "unit": "photons"},
             ),
-            "folded_background_unscaled": (
-                ["folded_channel"],
-                np.asarray(np.ma.filled(grouping @ background_unscaled), dtype=np.int64),
-                {"description": "Background counts", "unit": "photons"},
-            ),
             "folded_background": (
                 ["folded_channel"],
                 np.asarray(np.ma.filled(grouping @ background), dtype=np.float64),
@@ -147,8 +142,7 @@ class Observation(xr.Dataset):
             pha.quality,
             pha.exposure,
             backratio=backratio,
-            background=bkg.counts * backratio if bkg is not None else None,
-            background_unscaled=bkg.counts if bkg is not None else None,
+            background=bkg.counts if bkg is not None else None,
             attributes=metadata,
         )
 

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/data/util.py

@@ -152,11 +152,11 @@ def forward_model_with_multiple_inputs(
         transfer_matrix = BCOO.from_scipy_sparse(
             obs_configuration.transfer_matrix.data.to_scipy_sparse().tocsr()
         )
+        expected_counts = transfer_matrix @ flux_func(parameters).T
 
     else:
         transfer_matrix = np.asarray(obs_configuration.transfer_matrix.data.todense())
-
-    expected_counts = jnp.matvec(transfer_matrix, flux_func(parameters))
+        expected_counts = jnp.matvec(transfer_matrix, flux_func(parameters))
 
     # The result is clipped at 1e-6 to avoid 0 round-off and diverging likelihoods
     return jnp.clip(expected_counts, a_min=1e-6)

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/fit/_bayesian_model.py

@@ -135,7 +135,7 @@ class BayesianModel(nnx.Module):
                 with numpyro.plate("obs_plate/~/" + name, len(observation.folded_counts)):
                     numpyro.sample(
                         "obs/~/" + name,
-                        Poisson(obs_countrate + bkg_countrate / observation.folded_backratio.data),
+                        Poisson(obs_countrate + bkg_countrate * observation.folded_backratio.data),
                         obs=observation.folded_counts.data if observed else None,
                     )
 

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/fit/_fitter.py

@@ -9,6 +9,7 @@ import matplotlib.pyplot as plt
 import numpyro
 
 from jax import random
+from jax.numpy import concatenate
 from jax.random import PRNGKey
 from numpyro.infer import AIES, ESS, MCMC, NUTS, SVI, Predictive, Trace_ELBO
 from numpyro.infer.autoguide import AutoMultivariateNormal
@@ -52,9 +53,18 @@ class BayesianModelFitter(BayesianModel, ABC):
         )
 
         log_likelihood = numpyro.infer.log_likelihood(numpyro_model, posterior_samples)
+        if len(log_likelihood.keys()) > 1:
+            log_likelihood["full"] = concatenate([ll for _, ll in log_likelihood.items()], axis=1)
+            log_likelihood["obs/~/all"] = concatenate(
+                [ll for k, ll in log_likelihood.items() if "obs" in k], axis=1
+            )
+            if self.background_model is not None:
+                log_likelihood["bkg/~/all"] = concatenate(
+                    [ll for k, ll in log_likelihood.items() if "bkg" in k], axis=1
+                )
 
         seeded_model = numpyro.handlers.substitute(
-            numpyro.handlers.seed(numpyro_model, keys[3]),
+            numpyro.handlers.seed(numpyro_model, keys[2]),
             substitute_fn=numpyro.infer.init_to_sample,
         )
 
@@ -108,12 +118,10 @@ class BayesianModelFitter(BayesianModel, ABC):
         predictive_parameters = []
 
         for key, value in self._observation_container.items():
+            predictive_parameters.append(f"obs/~/{key}")
             if self.background_model is not None:
-                predictive_parameters.append(f"obs/~/{key}")
                 predictive_parameters.append(f"bkg/~/{key}")
             #                predictive_parameters.append(f"ins/~/{key}")
-            else:
-                predictive_parameters.append(f"obs/~/{key}")
         #                predictive_parameters.append(f"ins/~/{key}")
 
         inference_data.posterior_predictive = inference_data.posterior_predictive[
@@ -215,13 +223,29 @@ class VIFitter(BayesianModelFitter):
         self,
         rng_key: int = 0,
         num_steps: int = 10_000,
-        optimizer=numpyro.optim.Adam(step_size=0.0005),
-        loss=Trace_ELBO(),
+        optimizer: numpyro.optim._NumPyroOptim = numpyro.optim.Adam(step_size=0.0005),
+        loss: numpyro.infer.elbo.ELBO = Trace_ELBO(),
         num_samples: int = 1000,
-        guide=None,
+        guide: numpyro.infer.autoguide.AutoGuide | None = None,
         use_transformed_model: bool = True,
         plot_diagnostics: bool = False,
     ) -> FitResult:
+        """
+        Fit the model to the data using a variational inference approach from numpyro.
+
+        Parameters:
+            rng_key: the random key used to initialize the sampler.
+            num_steps: the number of steps for VI.
+            optimizer: the optimizer to use.
+            num_samples: the number of samples to draw.
+            loss: the loss function to use.
+            guide: the guide to use.
+            use_transformed_model: whether to use the transformed model to build the InferenceData.
+            plot_diagnostics: plot the loss during VI.
+
+        Returns:
+            A [`FitResult`][jaxspec.analysis.results.FitResult] instance containing the results of the fit.
+        """
         bayesian_model = (
             self.transformed_numpyro_model if use_transformed_model else self.numpyro_model
         )
@@ -231,7 +255,7 @@ class VIFitter(BayesianModelFitter):
 
         svi = SVI(bayesian_model, guide, optimizer, loss=loss)
 
-        keys = random.split(random.PRNGKey(rng_key), 3)
+        keys = random.split(random.PRNGKey(rng_key), 2)
         svi_result = svi.run(keys[0], num_steps)
         params = svi_result.params
 

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/model/abc.py

@@ -372,9 +372,10 @@ class AdditiveComponent(ModelComponent):
         continuum = self.continuum(energy)
         integrated_continuum = self.integrated_continuum(e_low, e_high)
 
-        return jsp.integrate.trapezoid(
-            continuum * energy**2, jnp.log(energy), axis=-1
-        ) + integrated_continuum * (e_high - e_low)
+        return (
+            jsp.integrate.trapezoid(continuum * energy**2, jnp.log(energy), axis=-1)
+            + integrated_continuum * (e_high + e_low) / 2.0
+        )
 
     @partial(jax.jit, static_argnums=0, static_argnames="n_points")
     def photon_flux(self, params, e_low, e_high, n_points=2):

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/model/instrument.py

@@ -8,13 +8,26 @@ from numpyro.distributions import Distribution
 
 
 class GainModel(ABC, nnx.Module):
+    """
+    Generic class for a gain model
+    """
+
     @abstractmethod
     def numpyro_model(self, observation_name: str):
         pass
 
 
 class ConstantGain(GainModel):
+    """
+    A constant gain model
+    """
+
     def __init__(self, prior_distribution: Distribution):
+        """
+        Parameters:
+            prior_distribution: the prior distribution for the gain value.
+        """
+
         self.prior_distribution = prior_distribution
 
     def numpyro_model(self, observation_name: str):
@@ -27,13 +40,25 @@ class ConstantGain(GainModel):
 
 
 class ShiftModel(ABC, nnx.Module):
+    """
+    Generic class for a shift model
+    """
+
     @abstractmethod
     def numpyro_model(self, observation_name: str):
         pass
 
 
 class ConstantShift(ShiftModel):
+    """
+    A constant shift model
+    """
+
     def __init__(self, prior_distribution: Distribution):
+        """
+        Parameters:
+            prior_distribution: the prior distribution for the shift value.
+        """
         self.prior_distribution = prior_distribution
 
     def numpyro_model(self, observation_name: str):
@@ -52,6 +77,15 @@ class InstrumentModel(nnx.Module):
         gain_model: GainModel | None = None,
         shift_model: ShiftModel | None = None,
     ):
+        """
+        Encapsulate an instrument model, build as a combination of a shift and gain model.
+
+        Parameters:
+            reference_observation_name : The observation to use as a reference
+            gain_model : The gain model
+            shift_model : The shift model
+        """
+
         self.reference = reference_observation_name
         self.gain_model = gain_model
         self.shift_model = shift_model

{jaxspec-0.3.1 → jaxspec-0.3.3}/src/jaxspec/util/online_storage.py

@@ -25,4 +25,5 @@ table_manager = pooch.create(
         "example_data/NGC7793_ULX4/MOS2.arf": "sha256:a126ff5a95a5f4bb93ed846944cf411d6e1c448626cb73d347e33324663d8b3f",
         "example_data/NGC7793_ULX4/PNbackground_spectrum.fits": "sha256:55e017e0c19b324245fef049dff2a7a2e49b9a391667ca9c4f667c4f683b1f49",
     },
+    retry_if_failed=10,
 )

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/conftest.py

@@ -70,6 +70,7 @@ pooch_dataset = pooch.create(
     base_url="https://github.com/HEACIT/curated-test-data/raw/main/",
     path=str(data_directory),
     registry=data_hash,
+    retry_if_failed=10,
 )
 
 for file in data_hash.keys():
@@ -112,6 +113,15 @@ def get_individual_mcmc_results(obs_model_prior):
     return [MCMCFitter(model, prior, obsconf).fit(num_samples=5000) for obsconf in obsconfs]
 
 
+@pytest.fixture(scope="session")
+def get_failed_mcmc_results(obs_model_prior):
+    obsconfs, model, prior = obs_model_prior
+
+    return [
+        MCMCFitter(model, prior, obsconf).fit(num_warmup=10, num_samples=10) for obsconf in obsconfs
+    ]
+
+
 @pytest.fixture(scope="session")
 def get_joint_mcmc_result(obs_model_prior):
     obsconfs, model, prior = obs_model_prior

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/test_background.py

@@ -14,9 +14,9 @@ spectral_model_background = Powerlaw() + Blackbodyrad()
 
 prior_background = {
     "powerlaw_1_alpha": dist.Uniform(0, 5),
-    "powerlaw_1_norm": dist.LogUniform(1e-8, 1e-3),
+    "powerlaw_1_norm": dist.LogUniform(1e-7, 1e-3),
     "blackbodyrad_1_kT": dist.Uniform(0, 5),
-    "blackbodyrad_1_norm": dist.LogUniform(1e-6, 1e-1),
+    "blackbodyrad_1_norm": dist.LogUniform(1e-5, 1e-1),
 }
 
 
@@ -36,7 +36,7 @@ def test_background_model(obs_model_prior, bkg_model):
     obs_list, model, prior = obs_model_prior
     forward = MCMCFitter(model, prior, obs_list[0], background_model=bkg_model)
     result = forward.fit(
-        num_chains=4, num_warmup=100, num_samples=100, mcmc_kwargs={"progress_bar": False}
+        num_chains=4, num_warmup=1000, num_samples=1000, mcmc_kwargs={"progress_bar": False}
     )
     result.plot_ppc(title=f"Test {bkg_model.__class__.__name__}")
 

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/test_fakeit.py

@@ -93,6 +93,19 @@ def test_fakeits_parallel(obsconfs, model, sharded_parameters):
     chex.assert_type(spectra, int)
 
 
+def test_fakeits_sparsify(obsconfs, model, sharded_parameters):
+    obsconf = obsconfs[0]
+    spectra = fakeit_for_multiple_parameters(
+        obsconf, model, sharded_parameters, apply_stat=False, sparsify_matrix=False
+    )
+    chex.assert_type(spectra, float)
+
+    spectra = fakeit_for_multiple_parameters(
+        obsconf, model, sharded_parameters, apply_stat=False, sparsify_matrix=True
+    )
+    chex.assert_type(spectra, float)
+
+
 def test_fakeits_multiple_observation(obsconfs, model, multidimensional_parameters):
     obsconf = obsconfs[0]
     spectra = fakeit_for_multiple_parameters(

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/test_mcmc.py

@@ -3,10 +3,13 @@ import pytest
 from jaxspec.fit import BayesianModel
 
 
-def test_convergence(get_individual_mcmc_results, get_joint_mcmc_result):
+def test_convergence(get_individual_mcmc_results, get_joint_mcmc_result, get_failed_mcmc_results):
     for result in get_individual_mcmc_results + get_joint_mcmc_result:
         assert result.converged
 
+    for result in get_failed_mcmc_results:
+        assert not result.converged
+
 
 def test_ns(obs_model_prior):
     NSFitter = pytest.importorskip("jaxspec.fit.NSFitter")

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/test_results.py

@@ -54,14 +54,18 @@ def test_posterior_photon_flux(get_joint_mcmc_result):
     result = get_joint_mcmc_result[0]
     e_min, e_max = 0.7, 1.2
     result.photon_flux(e_min, e_max, register=True)
-    assert f"photon_flux_{e_min:.1f}_{e_max:.1f}" in list(result.inference_data.posterior.keys())
+    assert f"mod/~/photon_flux_{e_min:.1f}_{e_max:.1f}" in list(
+        result.inference_data.posterior.keys()
+    )
 
 
 def test_posterior_energy_flux(get_joint_mcmc_result):
     result = get_joint_mcmc_result[0]
     e_min, e_max = 0.7, 1.2
     result.energy_flux(e_min, e_max, register=True)
-    assert f"energy_flux_{e_min:.1f}_{e_max:.1f}" in list(result.inference_data.posterior.keys())
+    assert f"mod/~/energy_flux_{e_min:.1f}_{e_max:.1f}" in list(
+        result.inference_data.posterior.keys()
+    )
 
 
 def test_posterior_luminosity(get_joint_mcmc_result):
@@ -76,4 +80,6 @@ def test_posterior_luminosity(get_joint_mcmc_result):
 
     result.luminosity(e_min, e_max, redshift=0.1, register=True)
 
-    assert f"luminosity_{e_min:.1f}_{e_max:.1f}" in list(result.inference_data.posterior.keys())
+    assert f"mod/~/luminosity_{e_min:.1f}_{e_max:.1f}" in list(
+        result.inference_data.posterior.keys()
+    )

{jaxspec-0.3.1 → jaxspec-0.3.3}/tests/test_xspec.py

@@ -1,8 +1,10 @@
 import os
 
+import astropy.units as u
 import numpy as np
 import pytest
 
+from jaxspec.model.additive import Powerlaw
 from jaxspec.util.online_storage import table_manager
 
 xspec = pytest.importorskip("xspec")
@@ -78,3 +80,40 @@ def test_bins(load_xspec_data, load_jaxspec_data):
     assert np.isclose(
         folding.in_energies, xspec_in_energies
     ).all(), f"The unfolded channel energy bins are not the same as XSPEC {file_pha}"
+
+
+def test_flux_computation():
+    xspec.AllData.clear()
+    xspec.AllModels.clear()
+
+    xspec.AllData.dummyrsp(lowE=0.2, highE=1.7, nBins=10_000)
+    m = xspec.Model("powerlaw")
+    m.powerlaw.PhoIndex = 2.0
+    m.powerlaw.norm = 1.0
+
+    xspec.AllModels.calcFlux("0.5 1.5")
+
+    phflux_xspec = m.flux[3]  # ph/cm^2/s
+    eflux_xspec = m.flux[0]  # erg/cm^2/s
+
+    factor = (1 * u.keV).to(u.erg).value
+    phflux_jaxspec = Powerlaw().photon_flux(
+        {"powerlaw_1_norm": 1.0, "powerlaw_1_alpha": 2.0}, e_low=0.5, e_high=1.5, n_points=10_000
+    )
+
+    eflux_jaxspec = (
+        Powerlaw().energy_flux(
+            {"powerlaw_1_norm": 1.0, "powerlaw_1_alpha": 2.0},
+            e_low=0.5,
+            e_high=1.5,
+            n_points=10_000,
+        )
+        * factor
+    )
+
+    assert np.isclose(
+        phflux_xspec, phflux_jaxspec
+    ), f"Mismatch between XSPEC and jaxspec on photon flux, got {phflux_xspec} and {phflux_jaxspec}"
+    assert np.isclose(
+        eflux_xspec, eflux_jaxspec
+    ), f"Mismatch between XSPEC and jaxspec on energy flux, got {eflux_xspec} and {eflux_jaxspec}"