jaxspec 0.2.0__tar.gz → 0.2.1.dev2__tar.gz

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: jaxspec
-Version: 0.2.0
+Version: 0.2.1.dev2
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 License: MIT
 Author: sdupourque
@@ -18,7 +18,7 @@ Requires-Dist: chainconsumer (>=1.1.2,<2.0.0)
 Requires-Dist: cmasher (>=1.6.3,<2.0.0)
 Requires-Dist: flax (>=0.10.1,<0.11.0)
 Requires-Dist: interpax (>=0.3.3,<0.4.0)
-Requires-Dist: jax (>=0.4.37,<0.5.0)
+Requires-Dist: jax (>=0.5.0,<0.6.0)
 Requires-Dist: jaxns (>=2.6.7,<3.0.0)
 Requires-Dist: jaxopt (>=0.8.1,<0.9.0)
 Requires-Dist: matplotlib (>=3.8.0,<4.0.0)
@@ -26,7 +26,7 @@ Requires-Dist: mendeleev (>=0.15,<0.20)
 Requires-Dist: networkx (>=3.1,<4.0)
 Requires-Dist: numpy (<2.0.0)
 Requires-Dist: numpyro (>=0.16.1,<0.17.0)
-Requires-Dist: optimistix (>=0.0.7,<0.0.10)
+Requires-Dist: optimistix (>=0.0.7,<0.0.11)
 Requires-Dist: pandas (>=2.2.0,<3.0.0)
 Requires-Dist: pooch (>=1.8.2,<2.0.0)
 Requires-Dist: scipy (<1.15)

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jaxspec"
-version = "0.2.0"
+version = "0.2.1dev-2"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = ["sdupourque <sdupourque@irap.omp.eu>"]
 license = "MIT"
@@ -11,7 +11,7 @@ documentation = "https://jaxspec.readthedocs.io/en/latest/"
 
 [tool.poetry.dependencies]
 python = ">=3.10,<3.13"
-jax = "^0.4.37"
+jax = "^0.5.0"
 numpy = "<2.0.0"
 pandas = "^2.2.0"
 astropy = "^6.0.0"
@@ -26,7 +26,7 @@ jaxopt = "^0.8.1"
 tinygp = "^0.3.0"
 seaborn = "^0.13.1"
 sparse = "^0.15.4"
-optimistix = ">=0.0.7,<0.0.10"
+optimistix = ">=0.0.7,<0.0.11"
 scipy = "<1.15"
 mendeleev = ">=0.15,<0.20"
 jaxns = "^2.6.7"
@@ -56,7 +56,7 @@ testbook = "^0.4.2"
 
 [tool.poetry.group.dev.dependencies]
 pre-commit = ">=3.5,<5.0"
-ruff = ">=0.2.1,<0.9.0"
+ruff = ">=0.2.1,<0.10.0"
 jupyterlab = "^4.0.7"
 notebook = "^7.0.6"
 ipywidgets = "^8.1.1"

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/src/jaxspec/analysis/results.py

@@ -391,6 +391,8 @@ class FitResult:
         alpha_envelope: (float, float) = (0.15, 0.25),
         style: str | Any = "default",
         title: str | None = None,
+        figsize: tuple[float, float] = (6, 6),
+        x_lims: tuple[float, float] | None = None,
     ) -> list[plt.Figure]:
         r"""
         Plot the posterior predictive distribution of the model. It also features a residual plot, defined using the
@@ -400,7 +402,7 @@ class FitResult:
         {(\text{Posterior counts})_{84\%}-(\text{Posterior counts})_{16\%}} $$
 
         Parameters:
-            percentile: The percentile of the posterior predictive distribution to plot.
+            n_sigmas: The number of sigmas to plot the envelops.
             x_unit: The units of the x-axis. It can be either a string (parsable by astropy.units) or an astropy unit. It must be homogeneous to either a length, a frequency or an energy.
             y_type: The type of the y-axis. It can be either "counts", "countrate", "photon_flux" or "photon_flux_density".
             plot_background: Whether to plot the background model if it is included in the fit.
@@ -408,6 +410,9 @@ class FitResult:
             scale: The axes scaling
             alpha_envelope: The transparency range for envelops
             style: The style of the plot. It can be either a string or a matplotlib style context.
+            title: The title of the plot.
+            figsize: The size of the figure.
+            x_lims: The limits of the x-axis.
 
         Returns:
             A list of matplotlib figures for each observation in the model.
@@ -436,7 +441,7 @@ class FitResult:
                 fig, ax = plt.subplots(
                     2,
                     1,
-                    figsize=(6, 6),
+                    figsize=figsize,
                     sharex="col",
                     height_ratios=[0.7, 0.3],
                 )
@@ -525,8 +530,10 @@ class FitResult:
                             alpha_envelope=alpha_envelope,
                         )
 
+                        name = component_name.split("*")[-1]
+
                         legend_plots += component_plot
-                        legend_labels.append(component_name)
+                        legend_labels.append(name)
 
                 if self.background_model is not None and plot_background:
                     # We plot the background only if it is included in the fit, i.e. by subtracting
@@ -617,6 +624,9 @@ class FitResult:
                         ax[0].set_xscale("log")
                         ax[0].set_yscale("log")
 
+                if x_lims is not None:
+                    ax[0].set_xlim(*x_lims)
+
                 fig.align_ylabels()
                 plt.subplots_adjust(hspace=0.0)
                 fig.tight_layout()
@@ -654,7 +664,7 @@ class FitResult:
         """
 
         consumer = ChainConsumer()
-        consumer.add_chain(self.to_chain(self.model.to_string()))
+        consumer.add_chain(self.to_chain("Results"))
         consumer.set_plot_config(config)
 
         # Context for default mpl style

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/src/jaxspec/data/util.py

@@ -1,14 +1,16 @@
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Literal, TypeVar
+from typing import TYPE_CHECKING, Literal, TypeVar
 
 import jax
+import jax.numpy as jnp
+import numpy as np
 import numpyro
 
 from astropy.io import fits
+from jax.experimental.sparse import BCOO
 from numpyro import handlers
 
-from .._fit._build_model import forward_model
 from ..model.abc import SpectralModel
 from ..util.online_storage import table_manager
 from . import Instrument, ObsConfiguration, Observation
@@ -16,6 +18,10 @@ from . import Instrument, ObsConfiguration, Observation
 K = TypeVar("K")
 V = TypeVar("V")
 
+if TYPE_CHECKING:
+    from ..data import ObsConfiguration
+    from ..model.abc import SpectralModel
+
 
 def load_example_pha(
     source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"],
@@ -124,8 +130,40 @@ def load_example_obsconf(source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"])
         raise ValueError(f"{source} not recognized.")
 
 
+def forward_model_with_multiple_inputs(
+    model: "SpectralModel",
+    parameters,
+    obs_configuration: "ObsConfiguration",
+    sparse=False,
+):
+    energies = np.asarray(obs_configuration.in_energies)
+    parameter_dims = next(iter(parameters.values())).shape
+
+    def flux_func(p):
+        return model.photon_flux(p, *energies)
+
+    for _ in parameter_dims:
+        flux_func = jax.vmap(flux_func)
+
+    flux_func = jax.jit(flux_func)
+
+    if sparse:
+        # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
+        transfer_matrix = BCOO.from_scipy_sparse(
+            obs_configuration.transfer_matrix.data.to_scipy_sparse().tocsr()
+        )
+
+    else:
+        transfer_matrix = np.asarray(obs_configuration.transfer_matrix.data.todense())
+
+    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)
+
+
 def fakeit_for_multiple_parameters(
-    instrument: ObsConfiguration | list[ObsConfiguration],
+    obsconfs: ObsConfiguration | list[ObsConfiguration],
     model: SpectralModel,
     parameters: Mapping[K, V],
     rng_key: int = 0,
@@ -134,10 +172,32 @@ def fakeit_for_multiple_parameters(
 ):
     """
     Convenience function to simulate multiple spectra from a given model and a set of parameters.
+    This is supposed to be somewhat optimized and can handle multiple parameters at once without blowing
+    up the memory. The parameters should be passed as a dictionary with the parameter name as the key and
+    the parameter values as the values, the value can be a scalar or a nd-array.
+
+    # Example:
+
+    ``` python
+    from jaxspec.data.util import fakeit_for_multiple_parameters
+    from numpy.random import default_rng
+
+    rng = default_rng(42)
+    size = (10, 30)
+
+    parameters = {
+        "tbabs_1_nh": rng.uniform(0.1, 0.4, size=size),
+        "powerlaw_1_alpha": rng.uniform(1, 3, size=size),
+        "powerlaw_1_norm": rng.exponential(10 ** (-0.5), size=size),
+        "blackbodyrad_1_kT": rng.uniform(0.1, 3.0, size=size),
+        "blackbodyrad_1_norm": rng.exponential(10 ** (-3), size=size)
+    }
 
+    spectra = fakeit_for_multiple_parameters(obsconf, model, parameters)
+    ```
 
     Parameters:
-        instrument: The instrumental setup.
+        obsconfs: The observational setup(s).
         model: The model to use.
         parameters: The parameters of the model.
         rng_key: The random number generator seed.
@@ -145,12 +205,12 @@ def fakeit_for_multiple_parameters(
         sparsify_matrix: Whether to sparsify the matrix or not.
     """
 
-    instruments = [instrument] if isinstance(instrument, ObsConfiguration) else instrument
+    obsconf_list = [obsconfs] if isinstance(obsconfs, ObsConfiguration) else obsconfs
     fakeits = []
 
-    for i, obs in enumerate(instruments):
-        countrate = jax.vmap(lambda p: forward_model(model, p, instrument, sparse=sparsify_matrix))(
-            parameters
+    for i, obsconf in enumerate(obsconf_list):
+        countrate = forward_model_with_multiple_inputs(
+            model, parameters, obsconf, sparse=sparsify_matrix
         )
 
         if apply_stat:

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/src/jaxspec/fit.py

@@ -13,7 +13,9 @@ import matplotlib.pyplot as plt
 import numpyro
 
 from jax import random
+from jax.experimental import mesh_utils
 from jax.random import PRNGKey
+from jax.sharding import PositionalSharding
 from numpyro.contrib.nested_sampling import NestedSampler
 from numpyro.distributions import Poisson, TransformedDistribution
 from numpyro.infer import AIES, ESS, MCMC, NUTS, Predictive
@@ -312,14 +314,27 @@ class BayesianModel:
         Check if the prior distribution include the observed data.
         """
         key_prior, key_posterior = jax.random.split(key, 2)
+        n_devices = len(jax.local_devices())
+        sharding = PositionalSharding(mesh_utils.create_device_mesh((n_devices,)))
+
+        # Sample from prior and correct if the number of samples is not a multiple of the number of devices
+        if num_samples % n_devices != 0:
+            num_samples = num_samples + n_devices - (num_samples % n_devices)
+
         prior_params = self.prior_samples(key=key_prior, num_samples=num_samples)
-        posterior_observations = self.mock_observations(prior_params, key=key_posterior)
+
+        # Split the parameters on every device
+        sharded_parameters = jax.device_put(prior_params, sharding)
+        posterior_observations = self.mock_observations(sharded_parameters, key=key_posterior)
 
         for key, value in self.observation_container.items():
             fig, ax = plt.subplots(
                 nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
             )
 
+            legend_plots = []
+            legend_labels = []
+
             y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
                 value.folded_counts.values, 1.0, "ct"
             )
@@ -337,6 +352,11 @@ class BayesianModel:
                 ax[0], value.out_energies, posterior_observations["obs_" + key], n_sigmas=3
             )
 
+            legend_plots.append((true_data_plot,))
+            legend_labels.append("Observed")
+            legend_plots += prior_plot
+            legend_labels.append("Prior Predictive")
+
             # 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
@@ -363,7 +383,7 @@ class BayesianModel:
             ax[1].set_ylim(0, 100)
             ax[0].set_xlim(value.out_energies.min(), value.out_energies.max())
             ax[0].loglog()
-            ax[0].legend(loc="upper right")
+            ax[0].legend(legend_plots, legend_labels)
             plt.suptitle(f"Prior Predictive coverage for {key}")
             plt.tight_layout()
             plt.show()

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/src/jaxspec/model/additive.py

@@ -156,13 +156,13 @@ class Gauss(AdditiveComponent):
     $$\mathcal{M}\left( E \right) = \frac{K}{\sigma \sqrt{2 \pi}}\exp\left(\frac{-(E-E_L)^2}{2\sigma^2}\right)$$
 
     !!! abstract "Parameters"
-        * $E_L$ (`E_l`) $\left[\text{keV}\right]$ : Energy of the line
+        * $E_L$ (`El`) $\left[\text{keV}\right]$ : Energy of the line
         * $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : Width of the line
         * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
     def __init__(self):
-        self.E_l = nnx.Param(2.0)
+        self.El = nnx.Param(2.0)
         self.sigma = nnx.Param(1e-2)
         self.norm = nnx.Param(1.0)
 
@@ -170,12 +170,12 @@ class Gauss(AdditiveComponent):
         return self.norm * (
             jsp.stats.norm.cdf(
                 e_high,
-                loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+                loc=jnp.asarray(self.El, dtype=jnp.float64),
                 scale=jnp.asarray(self.sigma, dtype=jnp.float64),
             )
             - jsp.stats.norm.cdf(
                 e_low,
-                loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+                loc=jnp.asarray(self.El, dtype=jnp.float64),
                 scale=jnp.asarray(self.sigma, dtype=jnp.float64),
             )
         )
@@ -246,13 +246,13 @@ class Agauss(AdditiveComponent):
     \frac{K}{\sigma \sqrt{2 \pi}} \exp\left(\frac{-(\lambda - \lambda_L)^2}{2 \sigma^2}\right)$$
 
     !!! abstract "Parameters"
-        * $\lambda_L$ (`lambda_l`) $\left[\unicode{x212B}\right]$ : Wavelength of the line in Angström
+        * $\lambda_L$ (`lambdal`) $\left[\unicode{x212B}\right]$ : Wavelength of the line in Angström
         * $\sigma$ (`sigma`) $\left[\unicode{x212B}\right]$ : Width of the line width in Angström
         * $K$ (`norm`) $\left[\frac{\unicode{x212B}~\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$: Normalization
     """
 
     def __init__(self):
-        self.lambda_l = nnx.Param(12.0)
+        self.lambdal = nnx.Param(12.0)
         self.sigma = nnx.Param(1e-2)
         self.norm = nnx.Param(1.0)
 
@@ -261,7 +261,7 @@ class Agauss(AdditiveComponent):
 
         return self.norm * jsp.stats.norm.pdf(
             hc / energy,
-            loc=jnp.asarray(self.lambda_l, dtype=jnp.float64),
+            loc=jnp.asarray(self.lambdal, dtype=jnp.float64),
             scale=jnp.asarray(self.sigma, dtype=jnp.float64),
         )
 
@@ -275,14 +275,14 @@ class Zagauss(AdditiveComponent):
     \frac{K (1+z)}{\sigma \sqrt{2 \pi}} \exp\left(\frac{-(\lambda/(1+z) - \lambda_L)^2}{2 \sigma^2}\right)$$
 
     !!! abstract "Parameters"
-        * $\lambda_L$ (`lambda_l`) $\left[\unicode{x212B}\right]$  : Wavelength of the line in Angström
+        * $\lambda_L$ (`lambdal`) $\left[\unicode{x212B}\right]$  : Wavelength of the line in Angström
         * $\sigma$ (`sigma`) $\left[\unicode{x212B}\right]$  : Width of the line width in Angström
         * $z$ (`redshift`) $\left[\text{dimensionless}\right]$ : Redshift
         * $K$ (`norm`) $\left[\frac{\unicode{x212B}~\text{photons}}{\text{keV}\text{cm}^2\text{s}}\right]$ : Normalization
     """
 
     def __init__(self):
-        self.lambda_l = nnx.Param(12.0)
+        self.lambdal = nnx.Param(12.0)
         self.sigma = nnx.Param(1e-2)
         self.redshift = nnx.Param(0.0)
         self.norm = nnx.Param(1.0)
@@ -297,7 +297,7 @@ class Zagauss(AdditiveComponent):
             * (1 + redshift)
             * jsp.stats.norm.pdf(
                 (hc / energy) / (1 + redshift),
-                loc=jnp.asarray(self.lambda_l, dtype=jnp.float64),
+                loc=jnp.asarray(self.lambdal, dtype=jnp.float64),
                 scale=jnp.asarray(self.sigma, dtype=jnp.float64),
             )
         )
@@ -311,14 +311,14 @@ class Zgauss(AdditiveComponent):
     \frac{K}{(1+z) \sigma \sqrt{2 \pi}}\exp\left(\frac{-(E(1+z)-E_L)^2}{2\sigma^2}\right)$$
 
     !!! abstract "Parameters"
-        * $E_L$ (`E_l`) $\left[\text{keV}\right]$ : Energy of the line
+        * $E_L$ (`El`) $\left[\text{keV}\right]$ : Energy of the line
         * $\sigma$ (`sigma`) $\left[\text{keV}\right]$ : Width of the line
         * $K$ (`norm`) $\left[\frac{\text{photons}}{\text{cm}^2\text{s}}\right]$ : Normalization
         * $z$ (`redshift`) $\left[\text{dimensionless}\right]$ : Redshift
     """
 
     def __init__(self):
-        self.E_l = nnx.Param(2.0)
+        self.El = nnx.Param(2.0)
         self.sigma = nnx.Param(1e-2)
         self.redshift = nnx.Param(0.0)
         self.norm = nnx.Param(1.0)
@@ -326,7 +326,7 @@ class Zgauss(AdditiveComponent):
     def continuum(self, energy) -> (jax.Array, jax.Array):
         return (self.norm / (1 + self.redshift)) * jsp.stats.norm.pdf(
             energy * (1 + self.redshift),
-            loc=jnp.asarray(self.E_l, dtype=jnp.float64),
+            loc=jnp.asarray(self.El, dtype=jnp.float64),
             scale=jnp.asarray(self.sigma, dtype=jnp.float64),
         )
 

{jaxspec-0.2.0 → jaxspec-0.2.1.dev2}/src/jaxspec/model/multiplicative.py

@@ -228,9 +228,9 @@ class Tbpcf(MultiplicativeComponent):
         self.nh = nnx.Param(1.0)
         self.f = nnx.Param(0.2)
 
-    def continuum(self, energy):
+    def factor(self, energy):
         sigma = jnp.interp(energy, self.energy, self.sigma, left=1e9, right=0.0)
-        return self.f * jnp.exp(-self.nh * sigma) + (1 - self.f)
+        return self.f * jnp.exp(-self.nh * sigma) + (1.0 - self.f)
 
 
 class FDcut(MultiplicativeComponent):
@@ -250,5 +250,5 @@ class FDcut(MultiplicativeComponent):
         self.Ec = nnx.Param(1.0)
         self.Ef = nnx.Param(3.0)
 
-    def continuum(self, energy):
+    def factor(self, energy):
         return (1 + jnp.exp((energy - self.Ec) / self.Ef)) ** -1