jaxspec 0.1.4__py3-none-any.whl → 0.2.1__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

@@ -1,17 +1,16 @@
 from collections.abc import Mapping
 from pathlib import Path
-from typing import Literal, TypeVar
+from typing import TYPE_CHECKING, Literal, TypeVar
 
-import haiku as hk
 import jax
+import jax.numpy as jnp
 import numpy as np
 import numpyro
 
 from astropy.io import fits
-from numpy.typing import ArrayLike
+from jax.experimental.sparse import BCOO
 from numpyro import handlers
 
-from .._fit._build_model import CountForwardModel
 from ..model.abc import SpectralModel
 from ..util.online_storage import table_manager
 from . import Instrument, ObsConfiguration, Observation
@@ -19,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"],
@@ -127,69 +130,40 @@ 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 = []
+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
 
-    for i, instrument in enumerate(instruments):
-        transformed_model = hk.without_apply_rng(
-            hk.transform(
-                lambda par: CountForwardModel(model, instrument, sparse=sparsify_matrix)(par)
-            )
-        )
+    def flux_func(p):
+        return model.photon_flux(p, *energies)
 
-        def obs_model(p):
-            return transformed_model.apply(None, p)
+    for _ in parameter_dims:
+        flux_func = jax.vmap(flux_func)
 
-        with handlers.seed(rng_seed=rng_key):
-            counts = numpyro.sample(
-                f"likelihood_obs_{i}",
-                numpyro.distributions.Poisson(obs_model(parameters)),
-            )
+    flux_func = jax.jit(flux_func)
 
-        """
-        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
+    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()
         )
-        """
 
-        fakeits.append(np.array(counts, dtype=int))
+    else:
+        transfer_matrix = np.asarray(obs_configuration.transfer_matrix.data.todense())
 
-    return fakeits[0] if len(fakeits) == 1 else fakeits
+    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,
@@ -198,11 +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
 
-    TODO : avoid redundancy, better doc and type hints
+    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.
@@ -210,28 +205,23 @@ 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):
-        transformed_model = hk.without_apply_rng(
-            hk.transform(lambda par: CountForwardModel(model, obs, sparse=sparsify_matrix)(par))
+    for i, obsconf in enumerate(obsconf_list):
+        countrate = forward_model_with_multiple_inputs(
+            model, parameters, obsconf, sparse=sparsify_matrix
         )
 
-        @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)