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

jaxspec/data/instrument.py

@@ -1,7 +1,10 @@
 import os
+
 import numpy as np
 import xarray as xr
+
 from matplotlib import colors
+
 from .ogip import DataARF, DataRMF
 
 
@@ -25,7 +28,15 @@ class Instrument(xr.Dataset):
     )
 
     @classmethod
-    def from_matrix(cls, redistribution_matrix, spectral_response, e_min_unfolded, e_max_unfolded, e_min_channel, e_max_channel):
+    def from_matrix(
+        cls,
+        redistribution_matrix,
+        spectral_response,
+        e_min_unfolded,
+        e_max_unfolded,
+        e_min_channel,
+        e_max_channel,
+    ):
         return cls(
             {
                 "redistribution": (
@@ -65,7 +76,7 @@ class Instrument(xr.Dataset):
         )
 
     @classmethod
-    def from_ogip_file(cls, rmf_path: str | os.PathLike, arf_path: str | os.PathLike = None):
+    def from_ogip_file(cls, rmf_path: str | os.PathLike, arf_path: str | os.PathLike | None = None):
         """
         Load the data from OGIP files.
 
@@ -80,37 +91,61 @@ class Instrument(xr.Dataset):
             specresp = DataARF.from_file(arf_path).specresp
 
         else:
-            specresp = np.ones(rmf.energ_lo.shape)
+            specresp = rmf.matrix.sum(axis=0)
+            rmf.sparse_matrix /= specresp
 
-        return cls.from_matrix(rmf.sparse_matrix, specresp, rmf.energ_lo, rmf.energ_hi, rmf.e_min, rmf.e_max)
+        return cls.from_matrix(
+            rmf.sparse_matrix, specresp, rmf.energ_lo, rmf.energ_hi, rmf.e_min, rmf.e_max
+        )
 
-    def plot_redistribution(self, **kwargs):
+    def plot_redistribution(
+        self,
+        xscale: str = "log",
+        yscale: str = "log",
+        cmap=None,
+        vmin: float = 1e-6,
+        vmax: float = 1e0,
+        add_labels: bool = True,
+        **kwargs,
+    ):
         """
         Plot the redistribution probability matrix
 
         Parameters:
+            xscale : The scale of the x-axis.
+            yscale : The scale of the y-axis.
+            cmap : The colormap to use.
+            vmin : The minimum value for the colormap.
+            vmax : The maximum value for the colormap.
+            add_labels : Whether to add labels to the plot.
             **kwargs : `kwargs` passed to https://docs.xarray.dev/en/latest/generated/xarray.plot.pcolormesh.html#xarray.plot.pcolormesh
         """
+
         import cmasher as cmr
 
         return xr.plot.pcolormesh(
             self.redistribution,
             x="e_max_unfolded",
             y="e_max_channel",
-            xscale="log",
-            yscale="log",
-            cmap=cmr.ember_r,
-            norm=colors.LogNorm(vmin=1e-6, vmax=1),
-            add_labels=True,
+            xscale=xscale,
+            yscale=yscale,
+            cmap=cmr.ember_r if cmap is None else cmap,
+            norm=colors.LogNorm(vmin=vmin, vmax=vmax),
+            add_labels=add_labels,
             **kwargs,
         )
 
-    def plot_area(self, **kwargs):
+    def plot_area(self, xscale: str = "log", yscale: str = "log", where: str = "post", **kwargs):
         """
         Plot the effective area
 
         Parameters:
+            xscale : The scale of the x-axis.
+            yscale : The scale of the y-axis.
+            where : The position of the steps.
             **kwargs : `kwargs` passed to https://docs.xarray.dev/en/latest/generated/xarray.DataArray.plot.line.html#xarray.DataArray.plot.line
         """
 
-        return self.area.plot.step(x="e_min_unfolded", xscale="log", yscale="log", where="post", **kwargs)
+        return self.area.plot.step(
+            x="e_min_unfolded", xscale=xscale, yscale=yscale, where=where, **kwargs
+        )

jaxspec/data/obsconf.py

@@ -157,8 +157,10 @@ 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": (
@@ -200,6 +202,14 @@ class ObsConfiguration(xr.Dataset):
                     "unit": "photons",
                 },
             ),
+            "folded_background_unscaled": (
+                ["folded_channel"],
+                folded_background_unscaled,
+                {
+                    "description": "To be done",
+                    "unit": "photons",
+                },
+            ),
         }
 
         return cls(
@@ -234,8 +244,8 @@ class ObsConfiguration(xr.Dataset):
         cls,
         instrument: Instrument,
         exposure: float,
-        low_energy: float = 1e-20,
-        high_energy: float = 1e20,
+        low_energy: float = 1e-300,
+        high_energy: float = 1e300,
     ):
         """
         Create a mock observation configuration from an instrument object. The fake observation will have zero counts.

jaxspec/data/observation.py

@@ -46,14 +46,16 @@ 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:
+        if background is None or background_unscaled is None:
             background = np.zeros_like(counts, dtype=np.int64)
+            background_unscaled = np.zeros_like(counts, dtype=np.int64)
 
         data_dict = {
             "counts": (
@@ -92,9 +94,14 @@ 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.int64),
+                np.asarray(np.ma.filled(grouping @ background), dtype=np.float64),
                 {"description": "Background counts", "unit": "photons"},
             ),
         }
@@ -128,14 +135,20 @@ class Observation(xr.Dataset):
         else:
             backratio = np.ones_like(pha.counts)
 
+        if (bkg is not None) and ("NET" in pha.flags):
+            counts = pha.counts + bkg.counts * backratio
+        else:
+            counts = pha.counts
+
         return cls.from_matrix(
-            pha.counts,
+            counts,
             pha.grouping,
             pha.channel,
             pha.quality,
             pha.exposure,
             backratio=backratio,
-            background=bkg.counts if bkg is not None else None,
+            background=bkg.counts * backratio if bkg is not None else None,
+            background_unscaled=bkg.counts if bkg is not None else None,
             attributes=metadata,
         )
 

jaxspec/data/ogip.py

@@ -1,9 +1,11 @@
-import numpy as np
 import os
+
 import astropy.units as u
+import numpy as np
 import sparse
-from astropy.table import QTable
+
 from astropy.io import fits
+from astropy.table import QTable
 
 
 class DataPHA:
@@ -25,6 +27,7 @@ class DataPHA:
         ancrfile=None,
         backscal=1.0,
         areascal=1.0,
+        flags=None,
     ):
         self.channel = np.asarray(channel, dtype=int)
         self.counts = np.asarray(counts, dtype=int)
@@ -36,6 +39,7 @@ class DataPHA:
         self.ancrfile = ancrfile
         self.backscal = np.asarray(backscal, dtype=float)
         self.areascal = np.asarray(areascal, dtype=float)
+        self.flags = flags
 
         if grouping is not None:
             # Indices array of the beginning of each group
@@ -55,7 +59,9 @@ class DataPHA:
                     data.append(True)
 
             # Create a COO sparse matrix
-            grp_matrix = sparse.COO((data, (rows, cols)), shape=(len(b_grp), len(channel)), fill_value=0)
+            grp_matrix = sparse.COO(
+                (data, (rows, cols)), shape=(len(b_grp), len(channel)), fill_value=0
+            )
 
         else:
             # Identity matrix case, use sparse for efficiency
@@ -74,12 +80,7 @@ class DataPHA:
 
         data = QTable.read(pha_file, "SPECTRUM")
         header = fits.getheader(pha_file, "SPECTRUM")
-
-        if header.get("HDUCLAS2") == "NET":
-            raise ValueError(
-                f"The HDUCLAS2={header.get('HDUCLAS2')} keyword in the PHA file is not supported."
-                f"Please open an issue if this is required."
-            )
+        flags = []
 
         if header.get("HDUCLAS3") == "RATE":
             raise ValueError(
@@ -114,6 +115,8 @@ class DataPHA:
         else:
             raise ValueError("No BACKSCAL found in the PHA file.")
 
+        backscal = np.where(backscal == 0, 1.0, backscal)
+
         if "AREASCAL" in header:
             areascal = header["AREASCAL"]
         elif "AREASCAL" in data.colnames:
@@ -121,8 +124,9 @@ class DataPHA:
         else:
             raise ValueError("No AREASCAL found in the PHA file.")
 
-        # Grouping and quality parameters are in binned PHA dataset
-        # Backfile, respfile and ancrfile are in primary header
+        if header.get("HDUCLAS2") == "NET":
+            flags.append("NET")
+
         kwargs = {
             "grouping": grouping,
             "quality": quality,
@@ -131,6 +135,7 @@ class DataPHA:
             "ancrfile": header.get("ANCRFILE"),
             "backscal": backscal,
             "areascal": areascal,
+            "flags": flags,
         }
 
         return cls(data["CHANNEL"], data["COUNTS"], header["EXPOSURE"], **kwargs)
@@ -176,7 +181,19 @@ class DataRMF:
         * [The Calibration Requirements for Spectral Analysis Addendum: Changes log](https://heasarc.gsfc.nasa.gov/docs/heasarc/caldb/docs/memos/cal_gen_92_002a/cal_gen_92_002a.html)
     """
 
-    def __init__(self, energ_lo, energ_hi, n_grp, f_chan, n_chan, matrix, channel, e_min, e_max, low_threshold=0.0):
+    def __init__(
+        self,
+        energ_lo,
+        energ_hi,
+        n_grp,
+        f_chan,
+        n_chan,
+        matrix,
+        channel,
+        e_min,
+        e_max,
+        low_threshold=0.0,
+    ):
         # RMF stuff
         self.energ_lo = energ_lo  # "Entry" energies
         self.energ_hi = energ_hi  # "Entry" energies
@@ -229,7 +246,9 @@ class DataRMF:
         idxs = data > low_threshold
 
         # Create a COO sparse matrix and then convert to CSR for efficiency
-        coo = sparse.COO([rows[idxs], cols[idxs]], data[idxs], shape=(len(self.energ_lo), len(self.channel)))
+        coo = sparse.COO(
+            [rows[idxs], cols[idxs]], data[idxs], shape=(len(self.energ_lo), len(self.channel))
+        )
         self.sparse_matrix = coo.T  # .tocsr()
 
     @property

jaxspec/data/util.py

@@ -2,16 +2,13 @@ from collections.abc import Mapping
 from pathlib import Path
 from typing import Literal, TypeVar
 
-import haiku as hk
 import jax
-import numpy as np
 import numpyro
 
 from astropy.io import fits
-from numpy.typing import ArrayLike
 from numpyro import handlers
 
-from .._fit._build_model import CountForwardModel
+from .._fit._build_model import forward_model
 from ..model.abc import SpectralModel
 from ..util.online_storage import table_manager
 from . import Instrument, ObsConfiguration, Observation
@@ -127,67 +124,6 @@ def load_example_obsconf(source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"])
         raise ValueError(f"{source} not recognized.")
 
 
-def fakeit(
-    instrument: ObsConfiguration | list[ObsConfiguration],
-    model: SpectralModel,
-    parameters: Mapping[K, V],
-    rng_key: int = 0,
-    sparsify_matrix: bool = False,
-) -> ArrayLike | list[ArrayLike]:
-    """
-    Convenience function to simulate a spectrum from a given model and a set of parameters.
-    It requires an instrumental setup, and unlike in
-    [XSPEC's fakeit](https://heasarc.gsfc.nasa.gov/xanadu/xspec/manual/node72.html), the error on the counts is given
-    exclusively by Poisson statistics.
-
-    Parameters:
-        instrument: The instrumental setup.
-        model: The model to use.
-        parameters: The parameters of the model.
-        rng_key: The random number generator seed.
-        sparsify_matrix: Whether to sparsify the matrix or not.
-    """
-
-    instruments = [instrument] if isinstance(instrument, ObsConfiguration) else instrument
-    fakeits = []
-
-    for i, instrument in enumerate(instruments):
-        transformed_model = hk.without_apply_rng(
-            hk.transform(
-                lambda par: CountForwardModel(model, instrument, sparse=sparsify_matrix)(par)
-            )
-        )
-
-        def obs_model(p):
-            return transformed_model.apply(None, p)
-
-        with handlers.seed(rng_seed=rng_key):
-            counts = numpyro.sample(
-                f"likelihood_obs_{i}",
-                numpyro.distributions.Poisson(obs_model(parameters)),
-            )
-
-        """
-        pha = DataPHA(
-            instrument.rmf.channel,
-            np.array(counts, dtype=int)*u.ct,
-            instrument.exposure,
-            grouping=instrument.grouping)
-
-        observation = Observation(
-            pha=pha,
-            arf=instrument.arf,
-            rmf=instrument.rmf,
-            low_energy=instrument.low_energy,
-            high_energy=instrument.high_energy
-        )
-        """
-
-        fakeits.append(np.array(counts, dtype=int))
-
-    return fakeits[0] if len(fakeits) == 1 else fakeits
-
-
 def fakeit_for_multiple_parameters(
     instrument: ObsConfiguration | list[ObsConfiguration],
     model: SpectralModel,
@@ -199,7 +135,6 @@ def fakeit_for_multiple_parameters(
     """
     Convenience function to simulate multiple spectra from a given model and a set of parameters.
 
-    TODO : avoid redundancy, better doc and type hints
 
     Parameters:
         instrument: The instrumental setup.
@@ -214,24 +149,19 @@ def fakeit_for_multiple_parameters(
     fakeits = []
 
     for i, obs in enumerate(instruments):
-        transformed_model = hk.without_apply_rng(
-            hk.transform(lambda par: CountForwardModel(model, obs, sparse=sparsify_matrix)(par))
+        countrate = jax.vmap(lambda p: forward_model(model, p, instrument, sparse=sparsify_matrix))(
+            parameters
         )
 
-        @jax.jit
-        @jax.vmap
-        def obs_model(p):
-            return transformed_model.apply(None, p)
-
         if apply_stat:
             with handlers.seed(rng_seed=rng_key):
                 spectrum = numpyro.sample(
                     f"likelihood_obs_{i}",
-                    numpyro.distributions.Poisson(obs_model(parameters)),
+                    numpyro.distributions.Poisson(countrate),
                 )
 
         else:
-            spectrum = obs_model(parameters)
+            spectrum = countrate
 
         fakeits.append(spectrum)
 

jaxspec/fit.py

@@ -10,7 +10,6 @@ import arviz as az
 import jax
 import jax.numpy as jnp
 import matplotlib.pyplot as plt
-import numpy as np
 import numpyro
 
 from jax import random
@@ -23,12 +22,16 @@ 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._plot import (
+    _error_bars_for_observed_data,
+    _plot_binned_samples_with_error,
+    _plot_poisson_data_with_error,
+)
 from .analysis.results import FitResult
 from .data import ObsConfiguration
 from .model.abc import SpectralModel
 from .model.background import BackgroundModel
-from .util.typing import PriorDictModel, PriorDictType
+from .util.typing import PriorDictType
 
 
 class BayesianModel:
@@ -63,10 +66,12 @@ class BayesianModel:
 
         if not callable(prior_distributions):
             # Validate the entry with pydantic
-            prior = PriorDictModel.from_dict(prior_distributions).nested_dict
+            # prior = PriorDictModel.from_dict(prior_distributions).
 
             def prior_distributions_func():
-                return build_prior(prior, expand_shape=(len(self.observation_container),))
+                return build_prior(
+                    prior_distributions, expand_shape=(len(self.observation_container),)
+                )
 
         else:
             prior_distributions_func = prior_distributions
@@ -74,6 +79,22 @@ class BayesianModel:
         self.prior_distributions_func = prior_distributions_func
         self.init_params = self.prior_samples()
 
+        # Check the priors are suited for the observations
+        split_parameters = [
+            (param, shape[-1])
+            for param, shape in jax.tree.map(lambda x: x.shape, self.init_params).items()
+            if (len(shape) > 1)
+            and not param.startswith("_")
+            and not param.startswith("bkg")  # hardcoded for subtracted background
+        ]
+
+        for parameter, proposed_number_of_obs in split_parameters:
+            if proposed_number_of_obs != len(self.observation_container):
+                raise ValueError(
+                    f"Invalid splitting in the prior distribution. "
+                    f"Expected {len(self.observation_container)} but got {proposed_number_of_obs} for {parameter}"
+                )
+
     @cached_property
     def observation_container(self) -> dict[str, ObsConfiguration]:
         """
@@ -137,7 +158,9 @@ class BayesianModel:
                 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,
                     )
 
@@ -293,37 +316,26 @@ class BayesianModel:
         posterior_observations = self.mock_observations(prior_params, key=key_posterior)
 
         for key, value in self.observation_container.items():
-            fig, axs = plt.subplots(
+            fig, ax = plt.subplots(
                 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=(16, 84),
+            y_observed, y_observed_low, y_observed_high = _error_bars_for_observed_data(
+                value.folded_counts.values, 1.0, "ct"
             )
 
-            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
-                )
+            true_data_plot = _plot_poisson_data_with_error(
+                ax[0],
+                value.out_energies,
+                y_observed.value,
+                y_observed_low.value,
+                y_observed_high.value,
+                alpha=0.7,
+            )
 
-                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$",
-                )
+            prior_plot = _plot_binned_samples_with_error(
+                ax[0], value.out_energies, posterior_observations["obs_" + key], n_sigmas=3
+            )
 
             # rank = np.vstack((posterior_observations["obs_" + key], value.folded_counts.values)).argsort(axis=0)[-1] / (num_samples) * 100
             counts = posterior_observations["obs_" + key]
@@ -336,22 +348,22 @@ class BayesianModel:
 
             rank = (less_than_obs + 0.5 * equal_to_obs) / num_samples * 100
 
-            axs[1].stairs(rank, edges=[*list(value.out_energies[0]), value.out_energies[1][-1]])
+            ax[1].stairs(rank, edges=[*list(value.out_energies[0]), value.out_energies[1][-1]])
 
-            axs[1].plot(
+            ax[1].plot(
                 (value.out_energies.min(), value.out_energies.max()),
                 (50, 50),
                 color="black",
                 linestyle="--",
             )
 
-            axs[1].set_xlabel("Energy (keV)")
-            axs[0].set_ylabel("Counts")
-            axs[1].set_ylabel("Rank (%)")
-            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")
+            ax[1].set_xlabel("Energy (keV)")
+            ax[0].set_ylabel("Counts")
+            ax[1].set_ylabel("Rank (%)")
+            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")
             plt.suptitle(f"Prior Predictive coverage for {key}")
             plt.tight_layout()
             plt.show()
@@ -544,7 +556,6 @@ class MCMCFitter(BayesianModelFitter):
         return FitResult(
             self,
             inference_data,
-            self.model.params,
             background_model=self.background_model,
         )
 
@@ -590,11 +601,13 @@ class NSFitter(BayesianModelFitter):
         ns = NestedSampler(
             bayesian_model,
             constructor_kwargs=dict(
-                num_parallel_workers=1,
                 verbose=verbose,
                 difficult_model=True,
-                max_samples=1e6,
+                max_samples=1e5,
                 parameter_estimation=True,
+                gradient_guided=True,
+                devices=jax.devices(),
+                # init_efficiency_threshold=0.01,
                 num_live_points=num_live_points,
             ),
             termination_kwargs=termination_kwargs if termination_kwargs else dict(),
@@ -613,6 +626,5 @@ class NSFitter(BayesianModelFitter):
         return FitResult(
             self,
             inference_data,
-            self.model.params,
             background_model=self.background_model,
         )