jaxspec 0.1.4__py3-none-any.whl → 0.2.1__py3-none-any.whl

jaxspec/_fit/_build_model.py

@@ -1,62 +1,25 @@
-import jax
-import numpyro
-import haiku as hk
-import numpy as np
+from typing import TYPE_CHECKING
+
 import jax.numpy as jnp
-from typing import Callable
+import numpy as np
+import numpyro
+
 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)
+    from ..model.abc import SpectralModel
+    from ..util.typing import PriorDictType
 
 
 def forward_model(
-        model: 'SpectralModel',
-        parameters,
-        obs_configuration: 'ObsConfiguration',
-        sparse=False,
-    ):
-
+    model: "SpectralModel",
+    parameters,
+    obs_configuration: "ObsConfiguration",
+    sparse=False,
+):
     energies = np.asarray(obs_configuration.in_energies)
 
     if sparse:
@@ -74,67 +37,27 @@ def forward_model(
     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=""):
+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)
+    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(sample, ArrayLike):
-            parameters[m][n] = jnp.ones(expand_shape) * sample
+        elif isinstance(value, ArrayLike):
+            parameters[key] = jnp.ones(expand_shape) * value
 
         else:
             raise ValueError(
-                f"Invalid prior type {type(sample)} for parameter {prefix}{m}_{n} : {sample}"
+                f"Invalid prior type {type(value)} for parameter {prefix}{module_name}_{param_name} : {value}"
             )
 
-    return parameters
+    return parameters

jaxspec/analysis/_plot.py

@@ -1,22 +1,58 @@
+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.stats import nbinom
+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,
-    percentiles: tuple = (16, 84),
+    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.
     """
-    y_low = nbinom.ppf(percentiles[0] / 100, y, 0.5)
-    y_high = nbinom.ppf(percentiles[1] / 100, y, 0.5)
 
     ax_to_plot = ax.errorbar(
         np.sqrt(x_bins[0] * x_bins[1]),
@@ -26,10 +62,133 @@ def _plot_poisson_data_with_error(
             y - y_low,
             y_high - y,
         ],
-        color="black",
-        linestyle="none",
-        alpha=0.3,
+        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