jaxspec 0.1.3__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

@@ -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"}
 
@@ -36,17 +46,23 @@ 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": (["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 +73,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"],
@@ -74,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"},
             ),
         }
@@ -84,36 +109,59 @@ 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)
 
+        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,
         )
 
     @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 +203,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/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 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)