add ebound support and tests

Author: ankitkhushwaha

stingray/events.py

@@ -631,7 +631,7 @@ def read(cls, filename, fmt=None, rmf_file=None, **kwargs):
             additional_columns = kwargs.pop("additional_columns", None)
 
             evt = FITSTimeseriesReader(
-                filename, output_class=EventList, additional_columns=additional_columns
+                filename, output_class=EventList, additional_columns=additional_columns, **kwargs
             )[:]
 
             if rmf_file is not None:

stingray/io.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 from astropy.io import fits
+import astropy.units as u
 from astropy.table import Table
 from astropy.logger import AstropyUserWarning
 import matplotlib.pyplot as plt
@@ -741,6 +742,7 @@ def __init__(
         gtistring=None,
         additional_columns=None,
         data_kind="events",
+        **kwargs,
     ):
         self.fname = fname
         self._data = fits.open(self.fname)
@@ -757,6 +759,8 @@ def __init__(
                 f"{data_kind} is an unknown data kind."
             )
         self.data_kind = data_kind
+        self.high_precision = kwargs.pop("high_precision", None)
+
         if additional_columns is None and self.detector_key != "NONE":
             additional_columns = [self.detector_key]
         elif self.detector_key != "NONE":
@@ -765,8 +769,12 @@ def __init__(
 
         if self.energy_column == "EBOUNDS":
             self.edata_hdu = self._data["EBOUNDS"]
+            self.emin = kwargs.pop("emin", None)
+            self.emax = kwargs.pop("emax", None)
         self.gti_file = gti_file
         self._read_gtis(self.gti_file)
+        if kwargs != {}:
+            warnings.warn(f"Unrecognized keywords: {list(kwargs.keys())}")
 
     @property
     def time(self):
@@ -843,12 +851,6 @@ def _transform_slice_into_events(self, data):
         if self._mission_specific_processing is not None:
             data = self._mission_specific_processing(data, header=self.header, hduname=self.hduname)
 
-        # Set the times
-        setattr(
-            new_ts,
-            self.main_array_attr,
-            data[self.time_column][:] + self.timezero,
-        )
         # Get conversion function PI->Energy
         try:
             pi_energy_func = get_rough_conversion_function(
@@ -865,15 +867,32 @@ def _transform_slice_into_events(self, data):
             ehigher = self.edata_hdu.data["E_MAX"]
             emid = elower + (ehigher - elower) / 2.0
 
-            energy = np.array([emid[c] for c in channels])
+            self.emin = np.min(elower) if self.emin is None else self.emin
+            self.emax = np.max(ehigher) if self.emax is None else self.emax
 
+            energy = np.array([emid[c] for c in channels])
             if (
                 hasattr(self.edata_hdu.columns["E_MIN"], "unit")
                 and (unit := self.edata_hdu.columns["E_MIN"].unit) is not None
             ):
                 conversion = (1 * u.Unit(unit)).to(u.keV).value
-            new_ts.energy = energy * conversion
-            new_ts.pi = channels
+
+            if isinstance(self.emin, u.Quantity):
+                self.emin = self.emin.to(u.keV).value
+            if isinstance(self.emax, u.Quantity):
+                self.emax = self.emax.to(u.keV).value
+
+            if self.emin > self.emax:
+                self.emin, self.emax = self.emax, self.emin
+
+            mask = (energy > self.emin) & (energy < self.emax)
+            energy = energy[mask]
+            channels = channels[mask]
+
+            time_dtype = np.float128 if self.high_precision is True else np.float64
+
+            new_ts.energy = np.asanyarray(energy * conversion, dtype=np.float64)
+            new_ts.pi = np.asanyarray(channels, dtype=time_dtype)
         else:
             if self.energy_column in data.dtype.names:
                 conversion = 1
@@ -888,6 +907,16 @@ def _transform_slice_into_events(self, data):
                 if pi_energy_func is not None:
                     new_ts.energy = pi_energy_func(new_ts.pi)
 
+        if "mask" not in locals():
+            mask = np.ones(data[self.time_column].shape, dtype=bool)
+
+        # Set the times
+        setattr(
+            new_ts,
+            self.main_array_attr,
+            data[self.time_column][mask] + self.timezero,
+        )
+
         det_numbers = None
         if self.detector_key is not None and self.detector_key in data.dtype.names:
             new_ts.detector_id = data[self.detector_key]

stingray/tests/test_events.py

@@ -6,6 +6,7 @@
 from numpy.testing import assert_allclose, assert_array_equal, assert_equal
 
 from astropy.io import fits
+import astropy.units as u
 from astropy.time import Time
 
 from ..events import EventList
@@ -375,6 +376,7 @@ def setup_class(self):
             elower = hdul[1].data["E_MIN"]
             ehigher = hdul[1].data["E_MAX"]
             emid = elower + (ehigher - elower) / 2.0
+            self.emin, self.emax = np.min(elower), np.max(ehigher)
             self.energy = np.array([emid[c] for c in self.pi])
 
     def test_read_fermi(self):
@@ -386,6 +388,43 @@ def test_check_energy_pi(self):
         assert_allclose(evt.energy, self.energy, atol=1e-8)
         assert_array_equal(evt.pi, self.pi)
 
+    def test_high_precision(self):
+        evt = EventList.read(self.fermi_file, high_precision=True)
+
+        assert np.issubdtype(evt.time.dtype, np.float128)
+        assert np.issubdtype(evt.pi.dtype, np.float128)
+
+    @pytest.mark.parametrize(
+        "emin, emax",
+        [
+            (2, 45.2),
+            (2 * u.keV, 45.2 * u.keV),
+            (2000 * u.eV, 45.2 * u.keV),
+            (None, 45.2 * u.keV),
+            (2000 * u.eV, None),
+            (None, None),
+        ],
+    )
+    def test_mask_ebound(self, emin, emax):
+        evt = EventList.read(self.fermi_file, emin=emin, emax=emax)
+
+        emin = emin if emin is not None else self.emin
+        emax = emax if emax is not None else self.emax
+
+        if isinstance(emin, u.Quantity):
+            emin = emin.to(u.keV).value
+        if isinstance(emax, u.Quantity):
+            emax = emax.to(u.keV).value
+
+        mask = (self.energy > emin) & (self.energy < emax)
+        energy = self.energy[mask]
+        time = self.time[mask]
+        pi = self.pi[mask]
+
+        assert_array_equal(evt.time, time)
+        assert_array_equal(evt.pi, pi)
+        assert_array_equal(evt.energy, energy)
+
     def test_check_time_gti(self):
         evt = EventList.read(self.fermi_file)
 
@@ -394,6 +433,10 @@ def test_check_time_gti(self):
         assert_equal(evt.ncounts, self.time.shape[0])
         assert_allclose(evt.mjdref, self.header["MJDREFI"])
 
+    def test_kwargs(self):
+        with pytest.warns(UserWarning, match="Unrecognized keywords:"):
+            EventList.read(self.fermi_file, alpha=1)
+
 
 class TestJoinEvents:
     def test_join_without_times_simulated(self):