jaxspec 0.0.6__py3-none-any.whl → 0.0.8__py3-none-any.whl

jaxspec/data/__init__.py

@@ -1,9 +1,9 @@
-# precommit is suppressing these imports
-from .obsconf import ObsConfiguration  # noqa: F401
-from .instrument import Instrument  # noqa: F401
-from .observation import Observation  # noqa: F401
 import astropy.units as u
 
+from .instrument import Instrument
+from .obsconf import ObsConfiguration
+from .observation import Observation
+
 u.add_enabled_aliases({"counts": u.count})
 u.add_enabled_aliases({"channel": u.dimensionless_unscaled})
 # Arbitrary units are found in .rsp files , let's hope it is compatible with what we would expect as the rmf x arf

jaxspec/data/obsconf.py

@@ -1,7 +1,8 @@
 import numpy as np
-import xarray as xr
-import sparse
 import scipy
+import sparse
+import xarray as xr
+
 from .instrument import Instrument
 from .observation import Observation
 
@@ -62,8 +63,26 @@ class ObsConfiguration(xr.Dataset):
 
     @classmethod
     def from_pha_file(
-        cls, pha_path, rmf_path=None, arf_path=None, bkg_path=None, low_energy: float = 1e-20, high_energy: float = 1e20
+        cls,
+        pha_path,
+        rmf_path: str | None = None,
+        arf_path: str | None = None,
+        bkg_path: str | None = None,
+        low_energy: float = 1e-20,
+        high_energy: float = 1e20,
     ):
+        r"""
+        Build the observation configuration from a PHA file.
+
+        Parameters:
+            pha_path: The path to the PHA file.
+            rmf_path: The path to the RMF file.
+            arf_path: The path to the ARF file.
+            bkg_path: The path to the background file.
+            low_energy: The lower bound of the energy range to consider.
+            high_energy: The upper bound of the energy range to consider.
+        """
+
         from .util import data_path_finder
 
         arf_path_default, rmf_path_default, bkg_path_default = data_path_finder(pha_path)
@@ -75,16 +94,34 @@ class ObsConfiguration(xr.Dataset):
         instrument = Instrument.from_ogip_file(rmf_path, arf_path=arf_path)
         observation = Observation.from_pha_file(pha_path, bkg_path=bkg_path)
 
-        return cls.from_instrument(instrument, observation, low_energy=low_energy, high_energy=high_energy)
+        return cls.from_instrument(
+            instrument, observation, low_energy=low_energy, high_energy=high_energy
+        )
 
     @classmethod
     def from_instrument(
-        cls, instrument: Instrument, observation: Observation, low_energy: float = 1e-20, high_energy: float = 1e20
+        cls,
+        instrument: Instrument,
+        observation: Observation,
+        low_energy: float = 1e-20,
+        high_energy: float = 1e20,
     ):
+        r"""
+        Build the observation configuration from an [`Instrument`][jaxspec.data.Instrument] and an [`Observation`][jaxspec.data.Observation] object.
+
+        Parameters:
+            instrument: The instrument object.
+            observation: The observation object.
+            low_energy: The lower bound of the energy range to consider.
+            high_energy: The upper bound of the energy range to consider.
+
+        """
         # First we unpack all the xarray data to classical np array for efficiency
         # We also exclude the bins that are flagged with bad quality on the instrument
         quality_filter = observation.quality.data == 0
-        grouping = scipy.sparse.csr_array(observation.grouping.data.to_scipy_sparse()) * quality_filter
+        grouping = (
+            scipy.sparse.csr_array(observation.grouping.data.to_scipy_sparse()) * quality_filter
+        )
         e_min_channel = instrument.coords["e_min_channel"].data
         e_max_channel = instrument.coords["e_max_channel"].data
         e_min_unfolded = instrument.coords["e_min_unfolded"].data
@@ -134,7 +171,10 @@ class ObsConfiguration(xr.Dataset):
             "area": (
                 ["unfolded_channel"],
                 area,
-                {"description": "Effective area with the same restrictions as the transfer matrix.", "units": "cm^2"},
+                {
+                    "description": "Effective area with the same restrictions as the transfer matrix.",
+                    "units": "cm^2",
+                },
             ),
             "exposure": ([], exposure, {"description": "Total exposure", "unit": "s"}),
             "folded_counts": (
@@ -148,7 +188,9 @@ class ObsConfiguration(xr.Dataset):
             "folded_backratio": (
                 ["folded_channel"],
                 folded_backratio,
-                {"description": "Background scaling after grouping, with the same restrictions as the transfer matrix."},
+                {
+                    "description": "Background scaling after grouping, with the same restrictions as the transfer matrix."
+                },
             ),
             "folded_background": (
                 ["folded_channel"],
@@ -186,3 +228,6 @@ class ObsConfiguration(xr.Dataset):
             },
             attrs=observation.attrs | instrument.attrs,
         )
+
+    def plot_counts(self, **kwargs):
+        return self.folded_counts.plot.step(x="e_min_folded", where="post", **kwargs)

jaxspec/data/util.py

@@ -1,102 +1,130 @@
-import importlib.resources
-import numpyro
+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 haiku as hk
-from pathlib import Path
-from numpy.typing import ArrayLike
-from collections.abc import Mapping
-from typing import TypeVar, Tuple
+import numpyro
+
 from astropy.io import fits
+from numpy.typing import ArrayLike
+from numpyro import handlers
 
-from . import Observation, Instrument, ObsConfiguration
-from ..model.abc import SpectralModel
 from ..fit import CountForwardModel
-from numpyro import handlers
+from ..model.abc import SpectralModel
+from ..util.online_storage import table_manager
+from . import Instrument, ObsConfiguration, Observation
 
 K = TypeVar("K")
 V = TypeVar("V")
 
 
-def load_example_observations():
+def load_example_pha(
+    source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"],
+) -> (Observation, list[Observation] | dict[str, Observation]):
     """
     Load some example observations from the package data.
+
+    Parameters:
+        source: The source to be loaded. Can be either "NGC7793_ULX4_PN" or "NGC7793_ULX4_ALL".
     """
 
-    example_observations = {
-        "PN": Observation.from_pha_file(
-            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(
-            str(importlib.resources.files("jaxspec") / "data/example_data/MOS1_spectrum_grp.fits"),
-            low_energy=0.3,
-            high_energy=7,
-        ),
-        "MOS2": Observation.from_pha_file(
-            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2_spectrum_grp.fits"),
-            low_energy=0.3,
-            high_energy=7,
-        ),
-    }
-
-    return example_observations
-
-
-def load_example_instruments():
+    if source == "NGC7793_ULX4_PN":
+        return Observation.from_pha_file(
+            table_manager.fetch("example_data/NGC7793_ULX4/PN_spectrum_grp20.fits"),
+            bkg_path=table_manager.fetch("example_data/NGC7793_ULX4/PNbackground_spectrum.fits"),
+        )
+
+    elif source == "NGC7793_ULX4_ALL":
+        return {
+            "PN": Observation.from_pha_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/PN_spectrum_grp20.fits"),
+                bkg_path=table_manager.fetch(
+                    "example_data/NGC7793_ULX4/PNbackground_spectrum.fits"
+                ),
+            ),
+            "MOS1": Observation.from_pha_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS1_spectrum_grp.fits"),
+                bkg_path=table_manager.fetch(
+                    "example_data/NGC7793_ULX4/MOS1background_spectrum.fits"
+                ),
+            ),
+            "MOS2": Observation.from_pha_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS2_spectrum_grp.fits"),
+                bkg_path=table_manager.fetch(
+                    "example_data/NGC7793_ULX4/MOS2background_spectrum.fits"
+                ),
+            ),
+        }
+
+    else:
+        raise ValueError(f"{source} not recognized.")
+
+
+def load_example_instruments(source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"]):
     """
     Load some example instruments from the package data.
+
+    Parameters:
+        source: The source to be loaded. Can be either "NGC7793_ULX4_PN" or "NGC7793_ULX4_ALL".
+
     """
+    if source == "NGC7793_ULX4_PN":
+        return Instrument.from_ogip_file(
+            table_manager.fetch("example_data/NGC7793_ULX4/PN.rmf"),
+            table_manager.fetch("example_data/NGC7793_ULX4/PN.arf"),
+        )
 
-    example_instruments = {
-        "PN": Instrument.from_ogip_file(
-            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(
-            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(
-            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2.rmf"),
-            str(importlib.resources.files("jaxspec") / "data/example_data/MOS2.arf"),
-        ),
-    }
-
-    return example_instruments
-
-
-def load_example_foldings():
+    elif source == "NGC7793_ULX4_ALL":
+        return {
+            "PN": Instrument.from_ogip_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/PN.rmf"),
+                table_manager.fetch("example_data/NGC7793_ULX4/PN.arf"),
+            ),
+            "MOS1": Instrument.from_ogip_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS1.rmf"),
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS1.arf"),
+            ),
+            "MOS2": Instrument.from_ogip_file(
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS2.rmf"),
+                table_manager.fetch("example_data/NGC7793_ULX4/MOS2.arf"),
+            ),
+        }
+
+    else:
+        raise ValueError(f"{source} not recognized.")
+
+
+def load_example_obsconf(source: Literal["NGC7793_ULX4_PN", "NGC7793_ULX4_ALL"]):
     """
-    Load some example instruments from the package data.
+    Load some example ObsConfigurations.
+
+    Parameters:
+        source: The source to be loaded. Can be either "NGC7793_ULX4_PN" or "NGC7793_ULX4_ALL".
     """
 
-    example_instruments = load_example_instruments()
-    example_observations = load_example_observations()
-
-    example_foldings = {
-        "PN": ObsConfiguration.from_instrument(
-            example_instruments["PN"],
-            example_observations["PN"],
-            low_energy=0.3,
-            high_energy=7.5,
-        ),
-        "MOS1": ObsConfiguration.from_instrument(
-            example_instruments["MOS1"],
-            example_observations["MOS1"],
-            low_energy=0.3,
-            high_energy=7,
-        ),
-        "MOS2": ObsConfiguration.from_instrument(
-            example_instruments["MOS2"],
-            example_observations["MOS2"],
-            low_energy=0.3,
-            high_energy=7,
-        ),
-    }
-
-    return example_foldings
+    if source in "NGC7793_ULX4_PN":
+        instrument = load_example_instruments(source)
+        observation = load_example_pha(source)
+
+        return ObsConfiguration.from_instrument(
+            instrument, observation, low_energy=0.5, high_energy=8.0
+        )
+
+    elif source == "NGC7793_ULX4_ALL":
+        instruments_dict = load_example_instruments(source)
+        observations_dict = load_example_pha(source)
+
+        return {
+            key: ObsConfiguration.from_instrument(
+                instruments_dict[key], observations_dict[key], low_energy=0.5, high_energy=8.0
+            )
+            for key in instruments_dict.keys()
+        }
+
+    else:
+        raise ValueError(f"{source} not recognized.")
 
 
 def fakeit(
@@ -107,7 +135,7 @@ def fakeit(
     sparsify_matrix: bool = False,
 ) -> ArrayLike | list[ArrayLike]:
     """
-    This function is a convenience function that allows to simulate spectra from a given model and a set of parameters.
+    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.
@@ -125,7 +153,9 @@ def fakeit(
 
     for i, instrument in enumerate(instruments):
         transformed_model = hk.without_apply_rng(
-            hk.transform(lambda par: CountForwardModel(model, instrument, sparse=sparsify_matrix)(par))
+            hk.transform(
+                lambda par: CountForwardModel(model, instrument, sparse=sparsify_matrix)(par)
+            )
         )
 
         def obs_model(p):
@@ -167,8 +197,7 @@ def fakeit_for_multiple_parameters(
     sparsify_matrix: bool = False,
 ):
     """
-    This function is a convenience function that allows to simulate spectra multiple spectra from a given model and a
-    set of parameters.
+    Convenience function to simulate multiple spectra from a given model and a set of parameters.
 
     TODO : avoid redundancy, better doc and type hints
 
@@ -209,9 +238,10 @@ def fakeit_for_multiple_parameters(
     return fakeits[0] if len(fakeits) == 1 else fakeits
 
 
-def data_path_finder(pha_path: str) -> Tuple[str | None, str | None, str | None]:
+def data_path_finder(pha_path: str) -> tuple[str | None, str | None, str | None]:
     """
-    This function tries its best to find the ARF, RMF and BKG files associated with a given PHA file.
+    Function which tries its best to find the ARF, RMF and BKG files associated with a given PHA file.
+
     Parameters:
         pha_path: The PHA file path.