jaxspec 0.0.2__py3-none-any.whl → 0.0.4__py3-none-any.whl

jaxspec/data/observation.py

@@ -1,6 +1,6 @@
-import os
 import numpy as np
 import xarray as xr
+from .ogip import DataPHA
 
 
 class Observation(xr.Dataset):
@@ -43,31 +43,40 @@ class Observation(xr.Dataset):
             attributes = {}
 
         if background is None:
-            background = np.zeros_like(counts, dtype=int)
+            background = np.zeros_like(counts, dtype=np.int64)
 
         data_dict = {
-            "counts": (["instrument_channel"], np.array(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.array(grouping @ counts, dtype=int),
+                np.asarray(np.ma.filled(grouping @ counts), dtype=np.int64),
                 {"description": "Folded counts, after grouping", "unit": "photons"},
             ),
             "grouping": (
                 ["folded_channel", "instrument_channel"],
-                np.array(grouping, dtype=bool),
+                grouping,
                 {"description": "Grouping matrix."},
             ),
-            "quality": (["instrument_channel"], np.array(quality, dtype=int), {"description": "Quality flag."}),
+            "quality": (["instrument_channel"], np.asarray(quality, dtype=np.int64), {"description": "Quality flag."}),
             "exposure": ([], float(exposure), {"description": "Total exposure", "unit": "s"}),
-            "backratio": ([], float(backratio), {"description": "Background scaling (SRC_BACKSCAL/BKG_BACKSCAL)"}),
+            "backratio": (
+                ["instrument_channel"],
+                np.asarray(backratio, dtype=float),
+                {"description": "Background scaling (SRC_BACKSCAL/BKG_BACKSCAL)"},
+            ),
+            "folded_backratio": (
+                ["folded_channel"],
+                np.asarray(np.ma.filled(grouping @ backratio), dtype=float),
+                {"description": "Background scaling after grouping"},
+            ),
             "background": (
                 ["instrument_channel"],
-                np.array(background, dtype=int),
+                np.asarray(background, dtype=np.int64),
                 {"description": "Background counts", "unit": "photons"},
             ),
             "folded_background": (
                 ["folded_channel"],
-                np.array(grouping @ background, dtype=int),
+                np.asarray(np.ma.filled(grouping @ background), dtype=np.int64),
                 {"description": "Background counts", "unit": "photons"},
             ),
         }
@@ -75,7 +84,7 @@ class Observation(xr.Dataset):
         return cls(
             data_dict,
             coords={
-                "channel": (["instrument_channel"], np.array(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),
@@ -86,15 +95,11 @@ class Observation(xr.Dataset):
         )
 
     @classmethod
-    def from_pha_file(cls, pha_file: str | os.PathLike, **kwargs):
-        from .util import data_loader
-
-        pha, arf, rmf, bkg, metadata = data_loader(pha_file)
-
+    def from_ogip_container(cls, pha: DataPHA, bkg: DataPHA | None = None, **metadata):
         if bkg is not None:
-            backratio = (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 = 1.0
+            backratio = np.ones_like(pha.counts)
 
         return cls.from_matrix(
             pha.counts,
@@ -107,6 +112,28 @@ class Observation(xr.Dataset):
             attributes=metadata,
         )
 
+    @classmethod
+    def from_pha_file(cls, pha_path: str, bkg_path: str | None = None, **metadata):
+        from .util import data_path_finder
+
+        arf_path, rmf_path, bkg_path_default = data_path_finder(pha_path)
+        bkg_path = bkg_path_default if bkg_path is None else bkg_path
+
+        pha = DataPHA.from_file(pha_path)
+        bkg = DataPHA.from_file(bkg_path) if bkg_path is not None else None
+
+        if metadata is None:
+            metadata = {}
+
+        metadata.update(
+            observation_file=pha_path,
+            background_file=bkg_path,
+            response_matrix_file=rmf_path,
+            ancillary_response_file=arf_path,
+        )
+
+        return cls.from_ogip_container(pha, bkg=bkg, **metadata)
+
     def plot_counts(self, **kwargs):
         """
         Plot the counts
@@ -133,7 +160,7 @@ class Observation(xr.Dataset):
         ax = fig.add_subplot(gs[1, 0])
         ax_histx = fig.add_subplot(gs[0, 0], sharex=ax)
         ax_histy = fig.add_subplot(gs[1, 1], sharey=ax)
-        sns.heatmap(self.grouping.T, ax=ax, cbar=False)
+        sns.heatmap(self.grouping.data.todense().T, ax=ax, cbar=False)
         ax_histx.step(np.arange(len(self.folded_counts)), self.folded_counts, where="post")
         ax_histy.step(self.counts, np.arange(len(self.counts)), where="post")
 

jaxspec/data/ogip.py

@@ -1,6 +1,7 @@
 import numpy as np
 import os
 import astropy.units as u
+import sparse
 from astropy.table import QTable
 from astropy.io import fits
 
@@ -8,7 +9,6 @@ from astropy.io import fits
 class DataPHA:
     r"""
     Class to handle PHA data defined with OGIP standards.
-
     ??? info "References"
         * [The OGIP standard PHA file format](https://heasarc.gsfc.nasa.gov/docs/heasarc/ofwg/docs/spectra/ogip_92_007/node5.html)
     """
@@ -26,30 +26,40 @@ class DataPHA:
         backscal=1.0,
         areascal=1.0,
     ):
-        self.channel = channel
-        self.counts = counts
-        self.exposure = exposure
+        self.channel = np.asarray(channel, dtype=int)
+        self.counts = np.asarray(counts, dtype=int)
+        self.exposure = float(exposure)
 
-        self.quality = quality
+        self.quality = np.asarray(quality, dtype=int)
         self.backfile = backfile
         self.respfile = respfile
         self.ancrfile = ancrfile
-        self.backscal = backscal
-        self.areascal = areascal
+        self.backscal = np.asarray(backscal, dtype=float)
+        self.areascal = np.asarray(areascal, dtype=float)
 
         if grouping is not None:
-            # Indices array of beginning of each group
+            # Indices array of the beginning of each group
             b_grp = np.where(grouping == 1)[0]
-            # Indices array of ending of each group
+            # Indices array of the ending of each group
             e_grp = np.hstack((b_grp[1:], len(channel)))
-            # Matrix to multiply with counts/channel to have counts/group
-            grp_matrix = np.zeros((len(b_grp), len(channel)), dtype=bool)
+
+            # Prepare data for sparse matrix
+            rows = []
+            cols = []
+            data = []
 
             for i in range(len(b_grp)):
-                grp_matrix[i, b_grp[i] : e_grp[i]] = 1
+                for j in range(b_grp[i], e_grp[i]):
+                    rows.append(i)
+                    cols.append(j)
+                    data.append(True)
+
+            # Create a COO sparse matrix
+            grp_matrix = sparse.COO((data, (rows, cols)), shape=(len(b_grp), len(channel)), fill_value=0)
 
         else:
-            grp_matrix = np.eye(len(channel))
+            # Identity matrix case, use sparse for efficiency
+            grp_matrix = sparse.eye(len(channel), format="coo", dtype=bool)
 
         self.grouping = grp_matrix
 
@@ -65,24 +75,44 @@ class DataPHA:
         data = QTable.read(pha_file, "SPECTRUM")
         header = fits.getheader(pha_file, "SPECTRUM")
 
-        if "QUALITY" in data.colnames:
+        if header.get("GROUPING") == 0:
+            grouping = None
+        elif "GROUPING" in data.colnames:
+            grouping = data["GROUPING"]
+        else:
+            raise ValueError("No grouping column found in the PHA file.")
+
+        if header.get("QUALITY") == 0:
+            quality = np.zeros(len(data["CHANNEL"]), dtype=bool)
+        elif "QUALITY" in data.colnames:
             quality = data["QUALITY"]
         else:
-            if header.get("QUALITY") == 0:
-                quality = np.ones(len(data["CHANNEL"]), dtype=bool)
-            else:
-                raise ValueError("No quality column found in the PHA file.")
+            raise ValueError("No QUALITY column found in the PHA file.")
+
+        if "BACKSCAL" in header:
+            backscal = header["BACKSCAL"] * np.ones_like(data["CHANNEL"])
+        elif "BACKSCAL" in data.colnames:
+            backscal = data["BACKSCAL"]
+        else:
+            raise ValueError("No BACKSCAL found in the PHA file.")
+
+        if "AREASCAL" in header:
+            areascal = header["AREASCAL"]
+        elif "AREASCAL" in data.colnames:
+            areascal = data["AREASCAL"]
+        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
         kwargs = {
-            "grouping": data["GROUPING"] if "GROUPING" in data.colnames else None,
+            "grouping": grouping,
             "quality": quality,
             "backfile": header.get("BACKFILE"),
             "respfile": header.get("RESPFILE"),
             "ancrfile": header.get("ANCRFILE"),
-            "backscal": header.get("BACKSCAL", 1.0),
-            "areascal": header.get("AREASCAL", 1.0),
+            "backscal": backscal,
+            "areascal": areascal,
         }
 
         return cls(data["CHANNEL"], data["COUNTS"], header["EXPOSURE"], **kwargs)
@@ -123,18 +153,16 @@ class DataARF:
 class DataRMF:
     r"""
     Class to handle RMF data defined with OGIP standards.
-
     ??? info "References"
         * [The Calibration Requirements for Spectral Analysis (Definition of RMF and ARF file formats)](https://heasarc.gsfc.nasa.gov/docs/heasarc/caldb/docs/memos/cal_gen_92_002/cal_gen_92_002.html)
         * [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):
+    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
-        self.n_grp = n_grp  # "Entry" energies
+        self.n_grp = n_grp
         self.f_chan = f_chan
         self.n_chan = n_chan
         self.matrix_entry = matrix
@@ -144,35 +172,51 @@ class DataRMF:
         self.e_min = e_min
         self.e_max = e_max
 
-        self.full_matrix = np.zeros(self.n_grp.shape + self.channel.shape)
+        # Prepare data for sparse matrix
+        rows = []
+        cols = []
+        data = []
 
-        for i, n_grp in enumerate(self.n_grp):
+        for i, n_grp_val in enumerate(self.n_grp):
             base = 0
 
             if np.size(self.f_chan[i]) == 1:
-                # ravel()[0] allows to get the value of the array without triggering numpy's conversion from
-                # multidimensional array to scalar
                 low = int(self.f_chan[i].ravel()[0])
                 high = min(
                     int(self.f_chan[i].ravel()[0] + self.n_chan[i].ravel()[0]),
-                    self.full_matrix.shape[1],
+                    len(self.channel),
                 )
 
-                self.full_matrix[i, low:high] = self.matrix_entry[i][0 : high - low]
+                rows.extend([i] * (high - low))
+                cols.extend(range(low, high))
+                data.extend(self.matrix_entry[i][0 : high - low])
 
             else:
-                for j in range(n_grp):
+                for j in range(n_grp_val):
                     low = self.f_chan[i][j]
-                    high = min(self.f_chan[i][j] + self.n_chan[i][j], self.full_matrix.shape[1])
+                    high = min(self.f_chan[i][j] + self.n_chan[i][j], len(self.channel))
 
-                    self.full_matrix[i, low:high] = self.matrix_entry[i][base : base + self.n_chan[i][j]]
+                    rows.extend([i] * (high - low))
+                    cols.extend(range(low, high))
+                    data.extend(self.matrix_entry[i][base : base + self.n_chan[i][j]])
 
                     base += self.n_chan[i][j]
 
-        # Transposed matrix so that we just have to multiply by the spectrum
-        self.matrix = self.full_matrix.T
-        # self.matrix = bsr_matrix(self.full_matrix.T).T
-        # self.sparse_matrix = sparse.BCOO.fromdense(jnp.copy(self.full_matrix))
+        # Convert lists to numpy arrays
+        rows = np.array(rows)
+        cols = np.array(cols)
+        data = np.array(data)
+
+        # Sometimes, zero elements are given in the matrix rows, so we get rid of them
+        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)))
+        self.sparse_matrix = coo.T  # .tocsr()
+
+    @property
+    def matrix(self):
+        return np.asarray(self.sparse_matrix.todense())
 
     @classmethod
     def from_file(cls, rmf_file: str | os.PathLike):
@@ -182,15 +226,31 @@ class DataRMF:
         Parameters:
             rmf_file: The RMF file path.
         """
+        extension_names = [hdu[1] for hdu in fits.info(rmf_file, output=False)]
+
+        if "MATRIX" in extension_names:
+            matrix_extension = "MATRIX"
 
-        matrix_table = QTable.read(rmf_file, "MATRIX")
+        elif "SPECRESP MATRIX" in extension_names:
+            matrix_extension = "SPECRESP MATRIX"
+            # raise NotImplementedError("SPECRESP MATRIX extension is not yet supported")
+
+        else:
+            raise ValueError("No MATRIX or SPECRESP MATRIX extension found in the RMF file")
+
+        matrix_table = QTable.read(rmf_file, matrix_extension)
         ebounds_table = QTable.read(rmf_file, "EBOUNDS")
 
+        matrix_header = fits.getheader(rmf_file, matrix_extension)
+
+        f_chan_column_pos = list(matrix_table.columns).index("F_CHAN") + 1
+        tlmin_fchan = int(matrix_header[f"TLMIN{f_chan_column_pos}"])
+
         return cls(
             matrix_table["ENERG_LO"],
             matrix_table["ENERG_HI"],
             matrix_table["N_GRP"],
-            matrix_table["F_CHAN"],
+            matrix_table["F_CHAN"] - tlmin_fchan,
             matrix_table["N_CHAN"],
             matrix_table["MATRIX"],
             ebounds_table["CHANNEL"],

jaxspec/data/util.py

@@ -1,15 +1,14 @@
 import importlib.resources
-import os
-import warnings
 import numpyro
 import jax
 import numpy as np
 import haiku as hk
+from pathlib import Path
 from numpy.typing import ArrayLike
 from collections.abc import Mapping
-from typing import TypeVar
+from typing import TypeVar, Tuple
+from astropy.io import fits
 
-from .ogip import DataPHA, DataARF, DataRMF
 from . import Observation, Instrument, ObsConfiguration
 from ..model.abc import SpectralModel
 from ..fit import CountForwardModel
@@ -26,17 +25,17 @@ def load_example_observations():
 
     example_observations = {
         "PN": Observation.from_pha_file(
-            importlib.resources.files("jaxspec") / "data/example_data/PN_spectrum_grp20.fits",
+            str(importlib.resources.files("jaxspec") / "data/example_data/PN_spectrum_grp20.fits"),
             low_energy=0.3,
             high_energy=7.5,
         ),
         "MOS1": Observation.from_pha_file(
-            importlib.resources.files("jaxspec") / "data/example_data/MOS1_spectrum_grp.fits",
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS1_spectrum_grp.fits"),
             low_energy=0.3,
             high_energy=7,
         ),
         "MOS2": Observation.from_pha_file(
-            importlib.resources.files("jaxspec") / "data/example_data/MOS2_spectrum_grp.fits",
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2_spectrum_grp.fits"),
             low_energy=0.3,
             high_energy=7,
         ),
@@ -52,16 +51,16 @@ def load_example_instruments():
 
     example_instruments = {
         "PN": Instrument.from_ogip_file(
-            importlib.resources.files("jaxspec") / "data/example_data/PN.arf",
-            importlib.resources.files("jaxspec") / "data/example_data/PN.rmf",
+            str(importlib.resources.files("jaxspec") / "data/example_data/PN.rmf"),
+            str(importlib.resources.files("jaxspec") / "data/example_data/PN.arf"),
         ),
         "MOS1": Instrument.from_ogip_file(
-            importlib.resources.files("jaxspec") / "data/example_data/MOS1.arf",
-            importlib.resources.files("jaxspec") / "data/example_data/MOS1.rmf",
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS1.rmf"),
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS1.arf"),
         ),
         "MOS2": Instrument.from_ogip_file(
-            importlib.resources.files("jaxspec") / "data/example_data/MOS2.arf",
-            importlib.resources.files("jaxspec") / "data/example_data/MOS2.rmf",
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2.rmf"),
+            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2.arf"),
         ),
     }
 
@@ -202,46 +201,51 @@ def fakeit_for_multiple_parameters(
     return fakeits[0] if len(fakeits) == 1 else fakeits
 
 
-def data_loader(pha_path, arf_path=None, rmf_path=None, bkg_path=None):
+def data_path_finder(pha_path: str) -> Tuple[str | None, str | None, str | None]:
     """
-    This function is a convenience function that allows to load PHA, ARF and RMF data
-    from a given PHA file, using either the ARF/RMF/BKG filenames in the header or the
-    specified filenames overwritten by the user.
-
+    This function tries its best to find the ARF, RMF and BKG files associated with a given PHA file.
     Parameters:
         pha_path: The PHA file path.
+
+    Returns:
         arf_path: The ARF file path.
         rmf_path: The RMF file path.
         bkg_path: The BKG file path.
     """
 
-    pha = DataPHA.from_file(pha_path)
-
-    if arf_path is None:
-        if pha.ancrfile != "none" and pha.ancrfile != "":
-            arf_path = os.path.join(os.path.dirname(pha_path), pha.ancrfile)
-    if rmf_path is None:
-        if pha.respfile != "none" and pha.respfile != "":
-            rmf_path = os.path.join(os.path.dirname(pha_path), pha.respfile)
-    if bkg_path is None:
-        if pha.backfile != "none" and pha.backfile != "":
-            bkg_path = os.path.join(os.path.dirname(pha_path), pha.backfile)
-
-    arf = DataARF.from_file(arf_path) if arf_path is not None else None
-    rmf = DataRMF.from_file(rmf_path) if rmf_path is not None else None
-
-    try:
-        bkg = DataPHA.from_file(bkg_path) if bkg_path is not None else None
-
-    except FileNotFoundError:
-        bkg = None
-        warnings.warn(f"Background file {bkg_path} is specified in header but not found.")
-
-    metadata = {
-        "observation_file": pha_path,
-        "background_file": bkg_path,
-        "response_matrix_file": rmf_path,
-        "ancillary_response_file": arf_path,
-    }
+    def find_path(file_name: str, directory: str) -> str | None:
+        if file_name.lower() != "none" and file_name != "":
+            return find_file_or_compressed_in_dir(file_name, directory)
+        else:
+            return None
+
+    header = fits.getheader(pha_path, "SPECTRUM")
+    directory = str(Path(pha_path).parent)
+
+    arf_path = find_path(header.get("ANCRFILE", "none"), directory)
+    rmf_path = find_path(header.get("RESPFILE", "none"), directory)
+    bkg_path = find_path(header.get("BACKFILE", "none"), directory)
+
+    return arf_path, rmf_path, bkg_path
+
+
+def find_file_or_compressed_in_dir(path: str | Path, directory: str | Path) -> str:
+    """
+    Try to find a file or its .gz compressed version in a given directory and return
+    the full path of the file.
+    """
+    path = Path(path) if isinstance(path, str) else path
+    directory = Path(directory) if isinstance(directory, str) else directory
+
+    if directory.joinpath(path).exists():
+        return str(directory.joinpath(path))
+
+    matching_files = list(directory.glob(str(path) + "*"))
+
+    if matching_files:
+        file = matching_files[0]
+        if file.suffix == ".gz":
+            return str(file)
 
-    return pha, arf, rmf, bkg, metadata
+    else:
+        raise FileNotFoundError(f"Can't find {path}(.gz) in {directory}.")

jaxspec/fit.py

@@ -7,6 +7,7 @@ from typing import Callable, TypeVar
 from abc import ABC
 from jax import random
 from jax.tree_util import tree_map
+from jax.experimental.sparse import BCSR
 from .analysis.results import ChainResult
 from .model.abc import SpectralModel
 from .data import ObsConfiguration
@@ -45,16 +46,19 @@ def build_numpyro_model(
     model: SpectralModel,
     background_model: BackgroundModel,
     name: str = "",
+    sparse: bool = False,
 ) -> Callable:
     def numpro_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)(par)))
+        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):
-            # TODO : Raise warning when setting a background model but there is no background spectra loaded
             bkg_countrate = background_model.numpyro_model(
                 obs.out_energies, obs.folded_background.data, name=name + "bkg", 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
 
@@ -65,7 +69,7 @@ def build_numpyro_model(
         with numpyro.plate(name + "obs_plate", len(obs.folded_counts)):
             numpyro.sample(
                 name + "obs",
-                Poisson(countrate + bkg_countrate * obs.backratio.data),
+                Poisson(countrate + bkg_countrate / obs.folded_backratio.data),
                 obs=obs.folded_counts.data if observed else None,
             )
 
@@ -77,11 +81,16 @@ 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):
+    def __init__(self, model: SpectralModel, folding: ObsConfiguration, sparse=False):
         super().__init__()
         self.model = model
         self.energies = jnp.asarray(folding.in_energies)
-        self.transfer_matrix = jnp.asarray(folding.transfer_matrix.data)
+
+        if sparse:  # folding.transfer_matrix.data.density > 0.015 is a good criterion to consider sparsify
+            self.transfer_matrix = BCSR.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):
         """
@@ -119,7 +128,7 @@ class BayesianModelAbstract(ABC):
         num_samples: int = 1000,
         max_tree_depth: int = 10,
         target_accept_prob: float = 0.8,
-        dense_mass=True,
+        dense_mass=False,
         mcmc_kwargs: dict = {},
     ) -> ChainResult:
         """
@@ -184,13 +193,14 @@ class BayesianModelAbstract(ABC):
 
 class BayesianModel(BayesianModelAbstract):
     """
-    Class to fit a model to a given set of observation using a Bayesian approach.
+    Class to fit a model to a given observation using a Bayesian approach.
     """
 
-    def __init__(self, model, observation, background_model: BackgroundModel = None):
+    def __init__(self, model, observation, background_model: BackgroundModel = None, sparsify_matrix: bool = False):
         super().__init__(model)
         self.observation = observation
         self.pars = tree_map(lambda x: jnp.float64(x), self.model.params)
+        self.sparse = sparsify_matrix
         self.background_model = background_model
 
     def numpyro_model(self, prior_distributions: HaikuDict[Distribution]) -> Callable:
@@ -204,7 +214,7 @@ class BayesianModel(BayesianModelAbstract):
 
         def model(observed=True):
             prior_params = build_prior(prior_distributions)
-            obs_model = build_numpyro_model(self.observation, self.model, self.background_model)
+            obs_model = build_numpyro_model(self.observation, self.model, self.background_model, sparse=self.sparse)
             obs_model(prior_params, observed=observed)
 
         return model