jaxspec 0.1.3__tar.gz → 0.1.4__tar.gz

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: jaxspec
-Version: 0.1.3
+Version: 0.1.4
 Summary: jaxspec is a bayesian spectral fitting library for X-ray astronomy.
 Home-page: https://github.com/renecotyfanboy/jaxspec
 License: MIT
@@ -11,7 +11,7 @@ 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)
+Requires-Dist: arviz (>=0.17.1,<0.21.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)
@@ -23,17 +23,17 @@ Requires-Dist: jaxlib (>=0.4.30,<0.5.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: mendeleev (>=0.15,<0.19)
 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: 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: simpleeval (>=0.9.13,<1.1.0)
 Requires-Dist: sparse (>=0.15.1,<0.16.0)
 Requires-Dist: tinygp (>=0.3.0,<0.4.0)
 Requires-Dist: watermark (>=2.4.3,<3.0.0)
@@ -78,3 +78,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.1.4}/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.1.4}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "jaxspec"
-version = "0.1.3"
+version = "0.1.4"
 description = "jaxspec is a bayesian spectral fitting library for X-ray astronomy."
 authors = ["sdupourque <sdupourque@irap.omp.eu>"]
 license = "MIT"
@@ -20,18 +20,18 @@ numpyro = "^0.15.3"
 dm-haiku = "^0.0.12"
 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"
+optimistix = ">=0.0.7,<0.0.10"
 scipy = "<1.15"
-mendeleev = ">=0.15,<0.18"
+mendeleev = ">=0.15,<0.19"
 pyzmq = "<27"
 jaxns = "<2.6"
 pooch = "^1.8.2"
@@ -57,8 +57,8 @@ 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.8.0"
 jupyterlab = "^4.0.7"
 notebook = "^7.0.6"
 ipywidgets = "^8.1.1"

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

@@ -0,0 +1,140 @@
+import jax
+import numpyro
+import haiku as hk
+import numpy as np
+import jax.numpy as jnp
+from typing import Callable
+from jax.experimental.sparse import BCOO
+from typing import TYPE_CHECKING
+from numpyro.distributions import Poisson
+from jax.typing import ArrayLike
+from numpyro.distributions import Distribution
+
+
+if TYPE_CHECKING:
+    from ..model.abc import SpectralModel
+    from ..data import ObsConfiguration
+    from ..util.typing import PriorDictModel, PriorDictType
+
+
+
+class CountForwardModel(hk.Module):
+    """
+    A haiku module which allows to build the function that simulates the measured counts
+    """
+
+    # TODO: It has no point of being a haiku module, it should be a simple function
+
+    def __init__(self, model: 'SpectralModel', folding: 'ObsConfiguration', sparse=False):
+        super().__init__()
+        self.model = model
+        self.energies = jnp.asarray(folding.in_energies)
+
+        if (
+            sparse
+        ):  # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
+            self.transfer_matrix = BCOO.from_scipy_sparse(
+                folding.transfer_matrix.data.to_scipy_sparse().tocsr()
+            )
+
+        else:
+            self.transfer_matrix = jnp.asarray(folding.transfer_matrix.data.todense())
+
+    def __call__(self, parameters):
+        """
+        Compute the count functions for a given observation.
+        """
+
+        expected_counts = self.transfer_matrix @ self.model.photon_flux(parameters, *self.energies)
+
+        return jnp.clip(expected_counts, a_min=1e-6)
+
+
+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_numpyro_model_for_single_obs(
+    obs,
+    model,
+    background_model,
+    name: str = "",
+    sparse: bool = False,
+) -> Callable:
+    """
+    Build a numpyro model for a given observation and spectral model.
+    """
+
+    def numpyro_model(prior_params, observed=True):
+
+        # Return the expected countrate for a set of parameters
+        obs_model = jax.jit(lambda par: forward_model(model, par, obs, sparse=sparse))
+        countrate = obs_model(prior_params)
+
+        # Handle the background model
+        if (getattr(obs, "folded_background", None) is not None) and (background_model is not None):
+            bkg_countrate = background_model.numpyro_model(
+                obs, model, name="bkg_" + name, observed=observed
+            )
+
+        elif (getattr(obs, "folded_background", None) is None) and (background_model is not None):
+            raise ValueError(
+                "Trying to fit a background model but no background is linked to this observation"
+            )
+
+        else:
+            bkg_countrate = 0.0
+
+
+        # Register the observed value
+        # This is the case where we fit a model to a TOTAL spectrum as defined in OGIP standard
+        with numpyro.plate("obs_plate_" + name, len(obs.folded_counts)):
+            numpyro.sample(
+                "obs_" + name,
+                Poisson(countrate + bkg_countrate / obs.folded_backratio.data),
+                obs=obs.folded_counts.data if observed else None,
+            )
+
+    return numpyro_model
+
+
+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 = dict(hk.data_structures.to_haiku_dict(prior))
+
+    for i, (m, n, sample) in enumerate(hk.data_structures.traverse(prior)):
+        if isinstance(sample, Distribution):
+            parameters[m][n] = jnp.ones(expand_shape) * numpyro.sample(f"{prefix}{m}_{n}", sample)
+
+        elif isinstance(sample, ArrayLike):
+            parameters[m][n] = jnp.ones(expand_shape) * sample
+
+        else:
+            raise ValueError(
+                f"Invalid prior type {type(sample)} for parameter {prefix}{m}_{n} : {sample}"
+            )
+
+    return parameters

{jaxspec-0.1.3 → jaxspec-0.1.4}/src/jaxspec/analysis/results.py

@@ -4,6 +4,7 @@ from collections.abc import Mapping
 from typing import TYPE_CHECKING, Any, Literal, TypeVar
 
 import arviz as az
+import astropy.cosmology.units as cu
 import astropy.units as u
 import jax
 import jax.numpy as jnp
@@ -287,7 +288,8 @@ class FitResult:
         self,
         e_min: float,
         e_max: float,
-        redshift: float | ArrayLike = 0.1,
+        redshift: float | ArrayLike = None,
+        distance: float | ArrayLike = None,
         observer_frame: bool = True,
         cosmology: Cosmology = Planck18,
         unit: Unit = u.erg / u.s,
@@ -310,6 +312,17 @@ class FitResult:
         if not observer_frame:
             raise NotImplementedError()
 
+        if redshift is None and distance is None:
+            raise ValueError("Either redshift or distance must be specified.")
+
+        if distance is not None:
+            if redshift is not None:
+                raise ValueError("Redshift must be None as a distance is specified.")
+            else:
+                redshift = distance.to(
+                    cu.redshift, cu.redshift_distance(cosmology, kind="luminosity")
+                ).value
+
         @jax.jit
         @jnp.vectorize
         def vectorized_flux(*pars):
@@ -489,13 +502,13 @@ class FitResult:
 
         match y_type:
             case "counts":
-                y_units = u.photon
+                y_units = u.ct
             case "countrate":
-                y_units = u.photon / u.s
+                y_units = u.ct / u.s
             case "photon_flux":
-                y_units = u.photon / u.cm**2 / u.s
+                y_units = u.ct / u.cm**2 / u.s
             case "photon_flux_density":
-                y_units = u.photon / u.cm**2 / u.s / x_unit
+                y_units = u.ct / u.cm**2 / u.s / x_unit
             case _:
                 raise ValueError(
                     f"Unknown y_type: {y_type}. Must be 'counts', 'countrate', 'photon_flux' or 'photon_flux_density'"
@@ -566,16 +579,16 @@ class FitResult:
                     case "photon_flux_density":
                         denominator = (xbins[1] - xbins[0]) * integrated_arf * exposure
 
-                y_samples = (count * u.photon / denominator).to(y_units)
-                y_observed = (obsconf.folded_counts.data * u.photon / denominator).to(y_units)
+                y_samples = (count * u.ct / denominator).to(y_units)
+                y_observed = (obsconf.folded_counts.data * u.ct / denominator).to(y_units)
                 y_observed_low = (
                     nbinom.ppf(percentile[0] / 100, obsconf.folded_counts.data, 0.5)
-                    * u.photon
+                    * u.ct
                     / denominator
                 ).to(y_units)
                 y_observed_high = (
                     nbinom.ppf(percentile[1] / 100, obsconf.folded_counts.data, 0.5)
-                    * u.photon
+                    * u.ct
                     / denominator
                 ).to(y_units)
 
@@ -611,18 +624,18 @@ class FitResult:
                 if self.background_model is not None:
                     # We plot the background only if it is included in the fit, i.e. by subtracting
                     ratio = obsconf.folded_backratio.data
-                    y_samples_bkg = (bkg_count * u.photon / (denominator * ratio)).to(y_units)
+                    y_samples_bkg = (bkg_count * u.ct / (denominator * ratio)).to(y_units)
                     y_observed_bkg = (
-                        obsconf.folded_background.data * u.photon / (denominator * ratio)
+                        obsconf.folded_background.data * u.ct / (denominator * ratio)
                     ).to(y_units)
                     y_observed_bkg_low = (
                         nbinom.ppf(percentile[0] / 100, obsconf.folded_background.data, 0.5)
-                        * u.photon
+                        * u.ct
                         / (denominator * ratio)
                     ).to(y_units)
                     y_observed_bkg_high = (
                         nbinom.ppf(percentile[1] / 100, obsconf.folded_background.data, 0.5)
-                        * u.photon
+                        * u.ct
                         / (denominator * ratio)
                     ).to(y_units)
 

{jaxspec-0.1.3 → jaxspec-0.1.4}/src/jaxspec/data/observation.py

@@ -1,5 +1,6 @@
 import numpy as np
 import xarray as xr
+
 from .ogip import DataPHA
 
 
@@ -23,7 +24,16 @@ class Observation(xr.Dataset):
     folded_background: xr.DataArray
     """The background counts, after grouping"""
 
-    __slots__ = ("grouping", "channel", "quality", "exposure", "background", "folded_background", "counts", "folded_counts")
+    __slots__ = (
+        "grouping",
+        "channel",
+        "quality",
+        "exposure",
+        "background",
+        "folded_background",
+        "counts",
+        "folded_counts",
+    )
 
     _default_attributes = {"description": "X-ray observation dataset"}
 
@@ -46,7 +56,11 @@ class Observation(xr.Dataset):
             background = np.zeros_like(counts, dtype=np.int64)
 
         data_dict = {
-            "counts": (["instrument_channel"], np.asarray(counts, dtype=np.int64), {"description": "Counts", "unit": "photons"}),
+            "counts": (
+                ["instrument_channel"],
+                np.asarray(counts, dtype=np.int64),
+                {"description": "Counts", "unit": "photons"},
+            ),
             "folded_counts": (
                 ["folded_channel"],
                 np.asarray(np.ma.filled(grouping @ counts), dtype=np.int64),
@@ -57,7 +71,11 @@ class Observation(xr.Dataset):
                 grouping,
                 {"description": "Grouping matrix."},
             ),
-            "quality": (["instrument_channel"], np.asarray(quality, dtype=np.int64), {"description": "Quality flag."}),
+            "quality": (
+                ["instrument_channel"],
+                np.asarray(quality, dtype=np.int64),
+                {"description": "Quality flag."},
+            ),
             "exposure": ([], float(exposure), {"description": "Total exposure", "unit": "s"}),
             "backratio": (
                 ["instrument_channel"],
@@ -84,20 +102,29 @@ class Observation(xr.Dataset):
         return cls(
             data_dict,
             coords={
-                "channel": (["instrument_channel"], np.asarray(channel, dtype=np.int64), {"description": "Channel number"}),
+                "channel": (
+                    ["instrument_channel"],
+                    np.asarray(channel, dtype=np.int64),
+                    {"description": "Channel number"},
+                ),
                 "grouped_channel": (
                     ["folded_channel"],
                     np.arange(len(grouping @ counts), dtype=np.int64),
                     {"description": "Channel number"},
                 ),
             },
-            attrs=cls._default_attributes if attributes is None else attributes | cls._default_attributes,
+            attrs=cls._default_attributes
+            if attributes is None
+            else attributes | cls._default_attributes,
         )
 
     @classmethod
     def from_ogip_container(cls, pha: DataPHA, bkg: DataPHA | None = None, **metadata):
         if bkg is not None:
-            backratio = np.nan_to_num((pha.backscal * pha.exposure * pha.areascal) / (bkg.backscal * bkg.exposure * bkg.areascal))
+            backratio = np.nan_to_num(
+                (pha.backscal * pha.exposure * pha.areascal)
+                / (bkg.backscal * bkg.exposure * bkg.areascal)
+            )
         else:
             backratio = np.ones_like(pha.counts)
 
@@ -114,6 +141,14 @@ class Observation(xr.Dataset):
 
     @classmethod
     def from_pha_file(cls, pha_path: str, bkg_path: str | None = None, **metadata):
+        """
+        Build an observation from a PHA file
+
+        Parameters:
+            pha_path : Path to the PHA file
+            bkg_path : Path to the background file
+            metadata : Additional metadata to add to the observation
+        """
         from .util import data_path_finder
 
         arf_path, rmf_path, bkg_path_default = data_path_finder(pha_path)
@@ -155,7 +190,16 @@ class Observation(xr.Dataset):
 
         fig = plt.figure(figsize=(6, 6))
         gs = fig.add_gridspec(
-            2, 2, width_ratios=(4, 1), height_ratios=(1, 4), left=0.1, right=0.9, bottom=0.1, top=0.9, wspace=0.05, hspace=0.05
+            2,
+            2,
+            width_ratios=(4, 1),
+            height_ratios=(1, 4),
+            left=0.1,
+            right=0.9,
+            bottom=0.1,
+            top=0.9,
+            wspace=0.05,
+            hspace=0.05,
         )
         ax = fig.add_subplot(gs[1, 0])
         ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)

{jaxspec-0.1.3 → jaxspec-0.1.4}/src/jaxspec/data/util.py

@@ -11,7 +11,7 @@ from astropy.io import fits
 from numpy.typing import ArrayLike
 from numpyro import handlers
 
-from ..fit import CountForwardModel
+from .._fit._build_model import CountForwardModel
 from ..model.abc import SpectralModel
 from ..util.online_storage import table_manager
 from . import Instrument, ObsConfiguration, Observation

{jaxspec-0.1.3 → jaxspec-0.1.4}/src/jaxspec/fit.py

@@ -7,7 +7,6 @@ from functools import cached_property
 from typing import Literal
 
 import arviz as az
-import haiku as hk
 import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
@@ -15,17 +14,15 @@ import numpy as np
 import numpyro
 
 from jax import random
-from jax.experimental.sparse import BCOO
 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.distributions import Poisson, TransformedDistribution
 from numpyro.infer import AIES, ESS, MCMC, NUTS, Predictive
 from numpyro.infer.inspect import get_model_relations
 from numpyro.infer.reparam import TransformReparam
 from numpyro.infer.util import log_density
 
+from ._fit._build_model import build_prior, forward_model
 from .analysis._plot import _plot_poisson_data_with_error
 from .analysis.results import FitResult
 from .data import ObsConfiguration
@@ -34,101 +31,6 @@ from .model.background import BackgroundModel
 from .util.typing import PriorDictModel, PriorDictType
 
 
-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 = dict(hk.data_structures.to_haiku_dict(prior))
-
-    for i, (m, n, sample) in enumerate(hk.data_structures.traverse(prior)):
-        if isinstance(sample, Distribution):
-            parameters[m][n] = jnp.ones(expand_shape) * numpyro.sample(f"{prefix}{m}_{n}", sample)
-
-        elif isinstance(sample, ArrayLike):
-            parameters[m][n] = jnp.ones(expand_shape) * sample
-
-        else:
-            raise ValueError(
-                f"Invalid prior type {type(sample)} for parameter {prefix}{m}_{n} : {sample}"
-            )
-
-    return parameters
-
-
-def build_numpyro_model_for_single_obs(
-    obs: ObsConfiguration,
-    model: SpectralModel,
-    background_model: BackgroundModel,
-    name: str = "",
-    sparse: bool = False,
-) -> Callable:
-    """
-    Build a numpyro model for a given observation and spectral model.
-    """
-
-    def numpyro_model(prior_params, observed=True):
-        # prior_params = build_prior(prior_distributions, name=name)
-        transformed_model = hk.without_apply_rng(
-            hk.transform(lambda par: CountForwardModel(model, obs, sparse=sparse)(par))
-        )
-
-        if (getattr(obs, "folded_background", None) is not None) and (background_model is not None):
-            bkg_countrate = background_model.numpyro_model(
-                obs, model, name="bkg_" + name, observed=observed
-            )
-        elif (getattr(obs, "folded_background", None) is None) and (background_model is not None):
-            raise ValueError(
-                "Trying to fit a background model but no background is linked to this observation"
-            )
-
-        else:
-            bkg_countrate = 0.0
-
-        obs_model = jax.jit(lambda p: transformed_model.apply(None, p))
-        countrate = obs_model(prior_params)
-
-        # This is the case where we fit a model to a TOTAL spectrum as defined in OGIP standard
-        with numpyro.plate("obs_plate_" + name, len(obs.folded_counts)):
-            numpyro.sample(
-                "obs_" + name,
-                Poisson(countrate + bkg_countrate / obs.folded_backratio.data),
-                obs=obs.folded_counts.data if observed else None,
-            )
-
-    return numpyro_model
-
-
-class CountForwardModel(hk.Module):
-    """
-    A haiku module which allows to build the function that simulates the measured counts
-    """
-
-    def __init__(self, model: SpectralModel, folding: ObsConfiguration, sparse=False):
-        super().__init__()
-        self.model = model
-        self.energies = jnp.asarray(folding.in_energies)
-
-        if (
-            sparse
-        ):  # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
-            self.transfer_matrix = BCOO.from_scipy_sparse(
-                folding.transfer_matrix.data.to_scipy_sparse().tocsr()
-            )
-
-        else:
-            self.transfer_matrix = jnp.asarray(folding.transfer_matrix.data.todense())
-
-    def __call__(self, parameters):
-        """
-        Compute the count functions for a given observation.
-        """
-
-        expected_counts = self.transfer_matrix @ self.model.photon_flux(parameters, *self.energies)
-
-        return jnp.clip(expected_counts, a_min=1e-6)
-
-
 class BayesianModel:
     """
     Base class for a Bayesian model. This class contains the necessary methods to build a model, sample from the prior
@@ -157,7 +59,6 @@ class BayesianModel:
         self.model = model
         self._observations = observations
         self.background_model = background_model
-        self.pars = tree_map(lambda x: jnp.float64(x), self.model.params)
         self.sparse = sparsify_matrix
 
         if not callable(prior_distributions):
@@ -197,22 +98,50 @@ class BayesianModel:
         Build the numpyro model using the observed data, the prior distributions and the spectral model.
         """
 
-        def model(observed=True):
+        def numpyro_model(observed=True):
+            # Instantiate and register the parameters of the spectral model and the background
             prior_params = self.prior_distributions_func()
 
             # Iterate over all the observations in our container and build a single numpyro model for each observation
-            for i, (key, observation) in enumerate(self.observation_container.items()):
+            for i, (name, observation) in enumerate(self.observation_container.items()):
+                # Check that we can indeed fit a background
+                if (getattr(observation, "folded_background", None) is not None) and (
+                    self.background_model is not None
+                ):
+                    # This call should register the parameter and observation of our background model
+                    bkg_countrate = self.background_model.numpyro_model(
+                        observation, name=name, observed=observed
+                    )
+
+                elif (getattr(observation, "folded_background", None) is None) and (
+                    self.background_model is not None
+                ):
+                    raise ValueError(
+                        "Trying to fit a background model but no background is linked to this observation"
+                    )
+
+                else:
+                    bkg_countrate = 0.0
+
                 # We expect that prior_params contains an array of parameters for each observation
                 # They can be identical or different for each observation
-                params = tree_map(lambda x: x[i], prior_params)
+                params = jax.tree.map(lambda x: x[i], prior_params)
 
-                obs_model = build_numpyro_model_for_single_obs(
-                    observation, self.model, self.background_model, name=key, sparse=self.sparse
+                # Forward model the observation and get the associated countrate
+                obs_model = jax.jit(
+                    lambda par: forward_model(self.model, par, observation, sparse=self.sparse)
                 )
+                obs_countrate = obs_model(params)
 
-                obs_model(params, observed=observed)
+                # Register the observation as an observed site
+                with numpyro.plate("obs_plate_" + name, len(observation.folded_counts)):
+                    numpyro.sample(
+                        "obs_" + name,
+                        Poisson(obs_countrate + bkg_countrate / observation.folded_backratio.data),
+                        obs=observation.folded_counts.data if observed else None,
+                    )
 
-        return model
+        return numpyro_model
 
     @cached_property
     def transformed_numpyro_model(self) -> Callable:
@@ -352,7 +281,9 @@ class BayesianModel:
         return fakeit(key, parameters)
 
     def prior_predictive_coverage(
-        self, key: PRNGKey = PRNGKey(0), num_samples: int = 1000, percentiles: tuple = (16, 84)
+        self,
+        key: PRNGKey = PRNGKey(0),
+        num_samples: int = 1000,
     ):
         """
         Check if the prior distribution include the observed data.
@@ -363,24 +294,36 @@ class BayesianModel:
 
         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]
+                nrows=2, ncols=1, sharex=True, figsize=(5, 6), height_ratios=[3, 1]
             )
 
             _plot_poisson_data_with_error(
                 axs[0],
                 value.out_energies,
                 value.folded_counts.values,
-                percentiles=percentiles,
+                percentiles=(16, 84),
             )
 
-            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),
-            )
+            for i, (envelop_percentiles, color, alpha) in enumerate(
+                zip(
+                    [(16, 86), (2.5, 97.5), (0.15, 99.85)],
+                    ["#03045e", "#0077b6", "#00b4d8"],
+                    [0.5, 0.4, 0.3],
+                )
+            ):
+                lower, upper = np.percentile(
+                    posterior_observations["obs_" + key], envelop_percentiles, axis=0
+                )
+
+                axs[0].stairs(
+                    upper,
+                    edges=[*list(value.out_energies[0]), value.out_energies[1][-1]],
+                    baseline=lower,
+                    alpha=alpha,
+                    fill=True,
+                    color=color,
+                    label=rf"${1+i}\sigma$",
+                )
 
             # rank = np.vstack((posterior_observations["obs_" + key], value.folded_counts.values)).argsort(axis=0)[-1] / (num_samples) * 100
             counts = posterior_observations["obs_" + key]
@@ -408,7 +351,9 @@ class BayesianModel:
             axs[1].set_ylim(0, 100)
             axs[0].set_xlim(value.out_energies.min(), value.out_energies.max())
             axs[0].loglog()
+            axs[0].legend(loc="upper right")
             plt.suptitle(f"Prior Predictive coverage for {key}")
+            plt.tight_layout()
             plt.show()
 
 
@@ -513,7 +458,11 @@ class BayesianModelFitter(BayesianModel, ABC):
             predictive_parameters
         ]
 
-        parameters = [x for x in inference_data.posterior.keys() if not x.endswith("_base")]
+        parameters = [
+            x
+            for x in inference_data.posterior.keys()
+            if not x.endswith("_base") or x.startswith("_")
+        ]
         inference_data.posterior = inference_data.posterior[parameters]
         inference_data.prior = inference_data.prior[parameters]
 

{jaxspec-0.1.3 → jaxspec-0.1.4}/src/jaxspec/model/background.py

@@ -1,22 +1,28 @@
 from abc import ABC, abstractmethod
 
+import jax
 import jax.numpy as jnp
 import numpyro
 import numpyro.distributions as dist
 
 from jax.scipy.integrate import trapezoid
+from numpyro.distributions import Poisson
 from tinygp import GaussianProcess, kernels
 
+from .._fit._build_model import build_prior, forward_model
+from ..util.typing import PriorDictModel
+from .abc import SpectralModel
+
 
 class BackgroundModel(ABC):
     """
-    This class handles the modelization of backgrounds in our spectra. This is handled in a separate class for now
+    Handles the background modelling in our spectra. This is handled in a separate class for now
     since backgrounds can be phenomenological models fitted directly on the folded spectrum. This is not the case for
     the source model, which is fitted on the unfolded spectrum. This might be changed later.
     """
 
     @abstractmethod
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+    def numpyro_model(self, observation, name: str = "", observed=True):
         """
         Build the model for the background.
         """
@@ -25,7 +31,7 @@ class BackgroundModel(ABC):
 
 class SubtractedBackground(BackgroundModel):
     """
-    This class is to use when implying that the observed background should be simply subtracted from the observed.
+    Define a model where the observed background is simply subtracted from the observed.
 
     !!! danger
 
@@ -35,93 +41,40 @@ class SubtractedBackground(BackgroundModel):
 
     """
 
-    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
-        _, observed_counts = obs.out_energies, obs.folded_background.data
-        numpyro.deterministic(f"{name}", observed_counts)
+    def numpyro_model(self, observation, name: str = "", observed=True):
+        _, observed_counts = observation.out_energies, observation.folded_background.data
+        numpyro.deterministic(f"bkg_{name}", observed_counts)
 
-        return jnp.zeros_like(observed_counts)
+        return observed_counts
 
 
 class BackgroundWithError(BackgroundModel):
     """
-    This class is to use when implying that the observed background should be simply subtracted from the observed. It
+    Define a model where the observed background is subtracted from the observed accounting for its intrinsic spread. It
     fits a countrate for each background bin assuming a Poisson distribution.
-
-    !!! warning
-        This is the same as [`ConjugateBackground`][jaxspec.model.background.ConjugateBackground]
-        but slower since it performs the fit using MCMC instead of analytical solution.
     """
 
-    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
+    def numpyro_model(self, obs, name: str = "", observed=True):
+        # We can't use the build_prior_function method here because the parameter size varies
+        # with the current observation. It must be instantiated in place.
         # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)
+
         _, observed_counts = obs.out_energies, obs.folded_background.data
         alpha = observed_counts + 1
         beta = 1
-        countrate = numpyro.sample(f"{name}_params", dist.Gamma(alpha, rate=beta))
+        countrate = numpyro.sample(f"_bkg_{name}_countrate", dist.Gamma(alpha, rate=beta))
 
-        with numpyro.plate(f"{name}_plate", len(observed_counts)):
+        with numpyro.plate(f"bkg_{name}_plate", len(observed_counts)):
             numpyro.sample(
-                f"{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
+                f"bkg_{name}", dist.Poisson(countrate), obs=observed_counts if observed else None
             )
 
         return countrate
 
 
-'''
-# TODO: Implement this class and sample it with Gibbs Sampling
-
-class ConjugateBackground(BackgroundModel):
-    r"""
-    This class fit an expected rate $\\lambda$ in each bin of the background spectrum. Assuming a Gamma prior
-    distribution, we can analytically derive the posterior as a Negative binomial distribution.
-
-    $$ p(\\lambda_{\text{Bkg}}) \\sim \\Gamma \\left( \alpha, \beta \right) \\implies
-    p\\left(\\lambda_{\text{Bkg}} | \text{Counts}_{\text{Bkg}}\right) \\sim \text{NB}\\left(\alpha, \frac{\beta}{\beta +1}
-    \right) $$
-
-    !!! info
-        Here, $\alpha$ and $\beta$ are set to $\alpha = \text{Counts}_{\text{Bkg}} + 1$ and $\beta = 1$. Doing so,
-        the prior distribution is such that $\\mathbb{E}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}} +1$ and
-        $\text{Var}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}}+1$. The +1 is to avoid numerical issues when the
-        counts are 0, and add a small scatter even if the measured background is effectively null.
-
-    ??? abstract "References"
-
-        - https://en.wikipedia.org/wiki/Conjugate_prior
-        - https://www.acsu.buffalo.edu/~adamcunn/probability/gamma.html
-        - https://bayesiancomputationbook.com/markdown/chp_01.html?highlight=conjugate#conjugate-priors
-        - https://vioshyvo.github.io/Bayesian_inference/conjugate-distributions.html
-
-    """
-
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
-        # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)
-        # alpha = observed_counts + 1
-        # beta = 1
-
-        with numpyro.plate(f"{name}_plate", len(observed_counts)):
-            countrate = numpyro.sample(f"{name}", dist.Gamma(2 * observed_counts + 1, 2), obs=None)
-
-        return countrate
-'''
-
-"""
-class SpectralBackgroundModel(BackgroundModel):
-    # I should pass the current spectral model as an argument to the background model
-    # In the numpyro model function
-    def __init__(self, model, prior):
-        self.model = model
-        self.prior = prior
-
-    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
-        #TODO : keep the sparsification from top model
-        transformed_model = hk.without_apply_rng(hk.transform(lambda par: CountForwardModel(model, obs, sparse=False)(par)))
-"""
-
-
 class GaussianProcessBackground(BackgroundModel):
     """
-    This class use a Gaussian Process to model the background. The GP is built using the
+    Define a Gaussian Process to model the background. The GP is built using the
     [`tinygp`](https://tinygp.readthedocs.io/en/stable/guide.html) library.
     """
 
@@ -146,16 +99,7 @@ class GaussianProcessBackground(BackgroundModel):
         self.n_nodes = n_nodes
         self.kernel = kernel
 
-    def numpyro_model(self, obs, spectral_model, name: str = "bkg", observed=True):
-        """
-        Build the model for the background.
-
-        Parameters:
-            energy: The energy bins lower and upper values (e_low, e_high).
-            observed_counts: The observed counts in each energy bin.
-            name: The name of the background model for parameters disambiguation.
-            observed: Whether the model is observed or not. Useful for `numpyro.infer.Predictive` calls.
-        """
+    def numpyro_model(self, obs, name: str = "", observed=True):
         energy, observed_counts = obs.out_energies, obs.folded_background.data
 
         if (observed_counts is not None) and (self.n_nodes >= len(observed_counts)):
@@ -163,28 +107,92 @@ class GaussianProcessBackground(BackgroundModel):
                 "More nodes than channels in the observation associated with GaussianProcessBackground."
             )
 
+        else:
+            observed_counts = jnp.asarray(observed_counts)
+
         # The parameters of the GP model
-        mean = numpyro.sample(f"{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0))
-        sigma = numpyro.sample(f"{name}_sigma", dist.HalfNormal(3.0))
-        rho = numpyro.sample(f"{name}_rho", dist.HalfNormal(10.0))
+        mean = numpyro.sample(
+            f"_bkg_{name}_mean", dist.Normal(jnp.log(jnp.mean(observed_counts)), 2.0)
+        )
+        sigma = numpyro.sample(f"_bkg_{name}_sigma", dist.HalfNormal(3.0))
+        rho = numpyro.sample(f"_bkg_{name}_rho", dist.HalfNormal(10.0))
 
         # Set up the kernel and GP objects
         kernel = sigma**2 * self.kernel(rho)
         nodes = jnp.linspace(0, 1, self.n_nodes)
         gp = GaussianProcess(kernel, nodes, diag=1e-5 * jnp.ones_like(nodes), mean=mean)
 
-        log_rate = numpyro.sample(f"_{name}_log_rate_nodes", gp.numpyro_dist())
+        log_rate = numpyro.sample(f"_bkg_{name}_log_rate_nodes", gp.numpyro_dist())
+
         interp_count_rate = jnp.exp(
             jnp.interp(energy, nodes * (self.e_max - self.e_min) + self.e_min, log_rate)
         )
         count_rate = trapezoid(interp_count_rate, energy, axis=0)
 
         # Finally, our observation model is Poisson
-        with numpyro.plate(f"{name}_plate", len(observed_counts)):
-            # TODO : change to Poisson Likelihood when there is no background model
-            # TODO : Otherwise clip the background model to 1e-6 to avoid numerical issues
+        with numpyro.plate("bkg_plate_" + name, len(observed_counts)):
             numpyro.sample(
-                f"{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
+                f"bkg_{name}", dist.Poisson(count_rate), obs=observed_counts if observed else None
             )
 
         return count_rate
+
+
+class SpectralBackgroundModel(BackgroundModel):
+    def __init__(self, spectral_model: "SpectralModel", prior_distributions, sparse=False):
+        self.spectral_model = spectral_model
+        self.prior = PriorDictModel.from_dict(prior_distributions).nested_dict
+        self.sparse = sparse
+
+    def numpyro_model(self, observation, name: str = "", observed=True):
+        params = build_prior(self.prior, prefix=f"_bkg_{name}")
+        bkg_model = jax.jit(
+            lambda par: forward_model(self.spectral_model, par, observation, sparse=self.sparse)
+        )
+        bkg_countrate = bkg_model(params)
+
+        with numpyro.plate("bkg_plate_" + name, len(observation.folded_background)):
+            numpyro.sample(
+                "bkg_" + name,
+                Poisson(bkg_countrate),
+                obs=observation.folded_background.data if observed else None,
+            )
+
+        return bkg_countrate
+
+
+'''
+class ConjugateBackground(BackgroundModel):
+    r"""
+    This class fit an expected rate $\\lambda$ in each bin of the background spectrum. Assuming a Gamma prior
+    distribution, we can analytically derive the posterior as a Negative binomial distribution.
+
+    $$ p(\\lambda_{\text{Bkg}}) \\sim \\Gamma \\left( \alpha, \beta \right) \\implies
+    p\\left(\\lambda_{\text{Bkg}} | \text{Counts}_{\text{Bkg}}\right) \\sim \text{NB}\\left(\alpha, \frac{\beta}{\beta +1}
+    \right) $$
+
+    !!! info
+        Here, $\alpha$ and $\beta$ are set to $\alpha = \text{Counts}_{\text{Bkg}} + 1$ and $\beta = 1$. Doing so,
+        the prior distribution is such that $\\mathbb{E}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}} +1$ and
+        $\text{Var}[\\lambda_{\text{Bkg}}] = \text{Counts}_{\text{Bkg}}+1$. The +1 is to avoid numerical issues when the
+        counts are 0, and add a small scatter even if the measured background is effectively null.
+
+    ??? abstract "References"
+
+        - https://en.wikipedia.org/wiki/Conjugate_prior
+        - https://www.acsu.buffalo.edu/~adamcunn/probability/gamma.html
+        - https://bayesiancomputationbook.com/markdown/chp_01.html?highlight=conjugate#conjugate-priors
+        - https://vioshyvo.github.io/Bayesian_inference/conjugate-distributions.html
+
+    """
+
+    def numpyro_model(self, energy, observed_counts, name: str = "bkg", observed=True):
+        # Gamma in numpyro is parameterized by concentration and rate (alpha/beta)
+        # alpha = observed_counts + 1
+        # beta = 1
+
+        with numpyro.plate(f"{name}_plate", len(observed_counts)):
+            countrate = numpyro.sample(f"{name}", dist.Gamma(2 * observed_counts + 1, 2), obs=None)
+
+        return countrate
+'''