jaxspec 0.1.3__tar.gz → 0.2.0__tar.gz

{jaxspec-0.1.3 → jaxspec-0.2.0}/PKG-INFO

@@ -1,43 +1,42 @@
-Metadata-Version: 2.1
+Metadata-Version: 2.3
 Name: jaxspec
-Version: 0.1.3
+Version: 0.2.0
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
-Home-page: https://github.com/renecotyfanboy/jaxspec
 License: MIT
 Author: sdupourque
 Author-email: sdupourque@irap.omp.eu
-Requires-Python: >=3.10,<3.12
+Requires-Python: >=3.10,<3.13
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3
 Classifier: Programming Language :: Python :: 3.10
 Classifier: Programming Language :: Python :: 3.11
-Requires-Dist: arviz (>=0.17.1,<0.20.0)
+Classifier: Programming Language :: Python :: 3.12
+Requires-Dist: arviz (>=0.17.1,<0.21.0)
 Requires-Dist: astropy (>=6.0.0,<7.0.0)
+Requires-Dist: catppuccin (>=2.3.4,<3.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.12,<0.0.13)
-Requires-Dist: gpjax (>=0.8.0,<0.9.0)
+Requires-Dist: flax (>=0.10.1,<0.11.0)
 Requires-Dist: interpax (>=0.3.3,<0.4.0)
-Requires-Dist: jax (>=0.4.33,<0.5.0)
-Requires-Dist: jaxlib (>=0.4.30,<0.5.0)
-Requires-Dist: jaxns (<2.6)
+Requires-Dist: jax (>=0.4.37,<0.5.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)
-Requires-Dist: mendeleev (>=0.15,<0.18)
+Requires-Dist: mendeleev (>=0.15,<0.20)
 Requires-Dist: networkx (>=3.1,<4.0)
 Requires-Dist: numpy (<2.0.0)
-Requires-Dist: numpyro (>=0.15.3,<0.16.0)
-Requires-Dist: optimistix (>=0.0.7,<0.0.8)
+Requires-Dist: numpyro (>=0.16.1,<0.17.0)
+Requires-Dist: optimistix (>=0.0.7,<0.0.10)
 Requires-Dist: pandas (>=2.2.0,<3.0.0)
 Requires-Dist: pooch (>=1.8.2,<2.0.0)
-Requires-Dist: pyzmq (<27)
 Requires-Dist: scipy (<1.15)
 Requires-Dist: seaborn (>=0.13.1,<0.14.0)
-Requires-Dist: simpleeval (>=0.9.13,<0.10.0)
-Requires-Dist: sparse (>=0.15.1,<0.16.0)
+Requires-Dist: simpleeval (>=0.9.13,<1.1.0)
+Requires-Dist: sparse (>=0.15.4,<0.16.0)
 Requires-Dist: tinygp (>=0.3.0,<0.4.0)
 Requires-Dist: watermark (>=2.4.3,<3.0.0)
 Project-URL: Documentation, https://jaxspec.readthedocs.io/en/latest/
+Project-URL: Homepage, https://github.com/renecotyfanboy/jaxspec
 Description-Content-Type: text/markdown
 
 <p align="center">
@@ -78,3 +77,24 @@ Once the environment is set up, you can install jaxspec directly from pypi
 pip install jaxspec --upgrade
 ```
 
+## Citation 
+
+If you use `jaxspec` in your research, please consider citing the following article 
+
+```
+@ARTICLE{2024A&A...690A.317D,
+       author = {{Dupourqu{\'e}}, S. and {Barret}, D. and {Diez}, C.~M. and {Guillot}, S. and {Quintin}, E.},
+        title = "{jaxspec: A fast and robust Python library for X-ray spectral fitting}",
+      journal = {\aap},
+     keywords = {methods: data analysis, methods: statistical, X-rays: general},
+         year = 2024,
+        month = oct,
+       volume = {690},
+          eid = {A317},
+        pages = {A317},
+          doi = {10.1051/0004-6361/202451736},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2024A&A...690A.317D},
+      adsnote = {Provided by the SAO/NASA Astrophysics Data System}
+}
+```
+

{jaxspec-0.1.3 → jaxspec-0.2.0}/README.md

@@ -35,3 +35,24 @@ Once the environment is set up, you can install jaxspec directly from pypi
 ```
 pip install jaxspec --upgrade
 ```
+
+## Citation 
+
+If you use `jaxspec` in your research, please consider citing the following article 
+
+```
+@ARTICLE{2024A&A...690A.317D,
+       author = {{Dupourqu{\'e}}, S. and {Barret}, D. and {Diez}, C.~M. and {Guillot}, S. and {Quintin}, E.},
+        title = "{jaxspec: A fast and robust Python library for X-ray spectral fitting}",
+      journal = {\aap},
+     keywords = {methods: data analysis, methods: statistical, X-rays: general},
+         year = 2024,
+        month = oct,
+       volume = {690},
+          eid = {A317},
+        pages = {A317},
+          doi = {10.1051/0004-6361/202451736},
+       adsurl = {https://ui.adsabs.harvard.edu/abs/2024A&A...690A.317D},
+      adsnote = {Provided by the SAO/NASA Astrophysics Data System}
+}
+```

{jaxspec-0.1.3 → jaxspec-0.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jaxspec"
-version = "0.1.3"
+version = "0.2.0"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = ["sdupourque <sdupourque@irap.omp.eu>"]
 license = "MIT"
@@ -10,39 +10,37 @@ documentation = "https://jaxspec.readthedocs.io/en/latest/"
 
 
 [tool.poetry.dependencies]
-python = ">=3.10,<3.12"
-jax = "^0.4.33"
-jaxlib = "^0.4.30"
+python = ">=3.10,<3.13"
+jax = "^0.4.37"
 numpy = "<2.0.0"
 pandas = "^2.2.0"
 astropy = "^6.0.0"
-numpyro = "^0.15.3"
-dm-haiku = "^0.0.12"
+numpyro = "^0.16.1"
 networkx = "^3.1"
 matplotlib = "^3.8.0"
-arviz = ">=0.17.1,<0.20.0"
+arviz = ">=0.17.1,<0.21.0"
 chainconsumer = "^1.1.2"
-simpleeval = "^0.9.13"
+simpleeval = ">=0.9.13,<1.1.0"
 cmasher = "^1.6.3"
-gpjax = "^0.8.0"
 jaxopt = "^0.8.1"
 tinygp = "^0.3.0"
 seaborn = "^0.13.1"
-sparse = "^0.15.1"
-optimistix = "^0.0.7"
+sparse = "^0.15.4"
+optimistix = ">=0.0.7,<0.0.10"
 scipy = "<1.15"
-mendeleev = ">=0.15,<0.18"
-pyzmq = "<27"
-jaxns = "<2.6"
+mendeleev = ">=0.15,<0.20"
+jaxns = "^2.6.7"
 pooch = "^1.8.2"
 interpax = "^0.3.3"
 watermark = "^2.4.3"
+catppuccin = "^2.3.4"
+flax = "^0.10.1"
 
 
 [tool.poetry.group.docs.dependencies]
 mkdocs = "^1.6.1"
 mkdocs-material = "^9.4.6"
-mkdocstrings = {extras = ["python"], version = ">=0.24,<0.27"}
+mkdocstrings = {extras = ["python"], version = ">=0.24,<0.28"}
 mkdocs-jupyter = "^0.25.0"
 
 
@@ -50,15 +48,15 @@ mkdocs-jupyter = "^0.25.0"
 chex = "^0.1.83"
 mktestdocs = "^0.2.1"
 coverage = "^7.3.2"
-pytest-cov = ">=4.1,<6.0"
+pytest-cov = ">=4.1,<7.0"
 flake8 = "^7.0.0"
 pytest = "^8.0.0"
 testbook = "^0.4.2"
 
 
 [tool.poetry.group.dev.dependencies]
-pre-commit = "^3.5.0"
-ruff = ">=0.2.1,<0.7.0"
+pre-commit = ">=3.5,<5.0"
+ruff = ">=0.2.1,<0.9.0"
 jupyterlab = "^4.0.7"
 notebook = "^7.0.6"
 ipywidgets = "^8.1.1"

jaxspec-0.2.0/src/jaxspec/_fit/_build_model.py

@@ -0,0 +1,63 @@
+from typing import TYPE_CHECKING
+
+import jax.numpy as jnp
+import numpy as np
+import numpyro
+
+from jax.experimental.sparse import BCOO
+from jax.typing import ArrayLike
+from numpyro.distributions import Distribution
+
+if TYPE_CHECKING:
+    from ..data import ObsConfiguration
+    from ..model.abc import SpectralModel
+    from ..util.typing import PriorDictType
+
+
+def forward_model(
+    model: "SpectralModel",
+    parameters,
+    obs_configuration: "ObsConfiguration",
+    sparse=False,
+):
+    energies = np.asarray(obs_configuration.in_energies)
+
+    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 = transfer_matrix @ model.photon_flux(parameters, *energies)
+
+    # 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 build_prior(prior: "PriorDictType", expand_shape: tuple = (), prefix=""):
+    """
+    Transform a dictionary of prior distributions into a dictionary of parameters sampled from the prior.
+    Must be used within a numpyro model.
+    """
+    parameters = {}
+
+    for key, value in prior.items():
+        # Split the key to extract the module name and parameter name
+        module_name, param_name = key.rsplit("_", 1)
+        if isinstance(value, Distribution):
+            parameters[key] = jnp.ones(expand_shape) * numpyro.sample(
+                f"{prefix}{module_name}_{param_name}", value
+            )
+
+        elif isinstance(value, ArrayLike):
+            parameters[key] = jnp.ones(expand_shape) * value
+
+        else:
+            raise ValueError(
+                f"Invalid prior type {type(value)} for parameter {prefix}{module_name}_{param_name} : {value}"
+            )
+
+    return parameters

jaxspec-0.2.0/src/jaxspec/analysis/_plot.py

@@ -0,0 +1,194 @@
+from __future__ import annotations
+
+import catppuccin
+import matplotlib.pyplot as plt
+import numpy as np
+
+from astropy import units as u
+from catppuccin.extras.matplotlib import load_color
+from cycler import cycler
+from jax.typing import ArrayLike
+from scipy.integrate import trapezoid
+from scipy.stats import nbinom, norm
+
+from jaxspec.data import ObsConfiguration
+
+PALETTE = catppuccin.PALETTE.latte
+
+COLOR_CYCLE = [
+    load_color(PALETTE.identifier, color)
+    for color in ["sky", "teal", "green", "yellow", "peach", "maroon", "red", "pink", "mauve"][::-1]
+]
+
+LINESTYLE_CYCLE = ["dashed", "dotted", "dashdot", "solid"]
+
+SPECS_CYCLE = cycler(linestyle=LINESTYLE_CYCLE) * cycler(color=COLOR_CYCLE)
+
+SPECTRUM_COLOR = load_color(PALETTE.identifier, "blue")
+SPECTRUM_DATA_COLOR = load_color(PALETTE.identifier, "overlay2")
+BACKGROUND_COLOR = load_color(PALETTE.identifier, "sapphire")
+BACKGROUND_DATA_COLOR = load_color(PALETTE.identifier, "overlay0")
+
+
+def sigma_to_percentile_intervals(sigmas):
+    intervals = []
+    for sigma in sigmas:
+        lower_bound = 100 * norm.cdf(-sigma)
+        upper_bound = 100 * norm.cdf(sigma)
+        intervals.append((lower_bound, upper_bound))
+    return intervals
+
+
+def _plot_poisson_data_with_error(
+    ax: plt.Axes,
+    x_bins: ArrayLike,
+    y: ArrayLike,
+    y_low: ArrayLike,
+    y_high: ArrayLike,
+    color=SPECTRUM_DATA_COLOR,
+    linestyle="none",
+    alpha=0.3,
+):
+    """
+    Plot Poisson data with error bars. We extrapolate the intrinsic error of the observation assuming a prior rate
+    distributed according to a Gamma RV.
+    """
+
+    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=color,
+        linestyle=linestyle,
+        alpha=alpha,
+        capsize=2,
+    )
+
+    return ax_to_plot
+
+
+def _plot_binned_samples_with_error(
+    ax: plt.Axes,
+    x_bins: ArrayLike,
+    y_samples: ArrayLike,
+    color=SPECTRUM_COLOR,
+    alpha_median: float = 0.7,
+    alpha_envelope: (float, float) = (0.15, 0.25),
+    linestyle="solid",
+    n_sigmas=3,
+):
+    """
+    Helper function to plot the posterior predictive distribution of the model. The function
+    computes the percentiles of the posterior predictive distribution and plot them as a shaded
+    area. If the observed data is provided, it is also plotted as a step function.
+
+    Parameters:
+        x_bins: The bin edges of the data (2 x N).
+        y_samples: The samples of the posterior predictive distribution (Samples X N).
+        ax: The matplotlib axes object.
+        color: The color of the posterior predictive distribution.
+    """
+
+    median = ax.stairs(
+        list(np.median(y_samples, axis=0)),
+        edges=[*list(x_bins[0]), x_bins[1][-1]],
+        color=color,
+        alpha=alpha_median,
+        linestyle=linestyle,
+    )
+
+    # The legend cannot handle fill_between, so we pass a fill to get a fancy icon
+    (envelope,) = ax.fill(np.nan, np.nan, alpha=alpha_envelope[-1], facecolor=color)
+
+    if n_sigmas == 1:
+        alpha_envelope = (alpha_envelope[1], alpha_envelope[0])
+
+    for percentile, alpha in zip(
+        sigma_to_percentile_intervals(list(range(n_sigmas, 0, -1))),
+        np.linspace(*alpha_envelope, n_sigmas),
+    ):
+        percentiles = np.percentile(y_samples, percentile, axis=0)
+        ax.stairs(
+            percentiles[1],
+            edges=[*list(x_bins[0]), x_bins[1][-1]],
+            baseline=percentiles[0],
+            alpha=alpha,
+            fill=True,
+            color=color,
+        )
+
+    return [(median, envelope)]
+
+
+def _compute_effective_area(
+    obsconf: ObsConfiguration,
+    x_unit: str | u.Unit = "keV",
+):
+    """
+    Helper function to compute the bins and effective area of an observational configuration
+
+    Parameters:
+        obsconf: The observational configuration.
+        x_unit: The unit 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.
+    """
+
+    # Note to Simon : do not change xbins[1] - xbins[0] to
+    # np.diff, you already did this twice and forgot that it does not work since diff keeps the dimensions
+    # and enable weird broadcasting that makes the plot fail
+
+    xbins = obsconf.out_energies * u.keV
+    xbins = xbins.to(x_unit, u.spectral())
+
+    # This computes the total effective area within all bins
+    # This is a bit weird since the following computation is equivalent to ignoring the RMF
+    exposure = obsconf.exposure.data * u.s
+    mid_bins_arf = obsconf.in_energies.mean(axis=0) * u.keV
+    mid_bins_arf = mid_bins_arf.to(x_unit, u.spectral())
+    e_grid = np.linspace(*xbins, 10)
+    interpolated_arf = np.interp(e_grid, mid_bins_arf, obsconf.area)
+    integrated_arf = (
+        trapezoid(interpolated_arf, x=e_grid, axis=0)
+        / (
+            np.abs(
+                xbins[1] - xbins[0]
+            )  # Must fold in abs because some units reverse the ordering of the bins
+        )
+        * u.cm**2
+    )
+
+    return xbins, exposure, integrated_arf
+
+
+def _error_bars_for_observed_data(observed_counts, denominator, units, sigma=1):
+    r"""
+    Compute the error bars for the observed data assuming a prior Gamma distribution
+
+    Parameters:
+        observed_counts: array of integer counts
+        denominator: normalization factor (e.g. effective area)
+        units: unit to convert to
+        sigma: dispersion to use for quantiles computation
+
+    Returns:
+        y_observed: observed counts in the desired units
+        y_observed_low: lower bound of the error bars
+        y_observed_high: upper bound of the error bars
+    """
+
+    percentile = sigma_to_percentile_intervals([sigma])[0]
+
+    y_observed = (observed_counts * u.ct / denominator).to(units)
+
+    y_observed_low = (
+        nbinom.ppf(percentile[0] / 100, observed_counts, 0.5) * u.ct / denominator
+    ).to(units)
+
+    y_observed_high = (
+        nbinom.ppf(percentile[1] / 100, observed_counts, 0.5) * u.ct / denominator
+    ).to(units)
+
+    return y_observed, y_observed_low, y_observed_high