add support for fermi_gbm files

Author: ankitkhushwaha

stingray/io.py

@@ -743,6 +743,8 @@ def __init__(
         data_kind="events",
     ):
         self.fname = fname
+        self._data = fits.open(self.fname)
+
         self._meta_attrs = []
         self.gtistring = gtistring
         self.output_class = output_class
@@ -759,7 +761,10 @@ def __init__(
             additional_columns = [self.detector_key]
         elif self.detector_key != "NONE":
             additional_columns.append(self.detector_key)
-        self.data_hdu = fits.open(self.fname)[self.hduname]
+        self.data_hdu = self._data[self.hduname]
+
+        if self.energy_column == "EBOUNDS":
+            self.edata_hdu = self._data["EBOUNDS"]
         self.gti_file = gti_file
         self._read_gtis(self.gti_file)
 
@@ -793,7 +798,6 @@ def __getitem__(self, index):
         """Return an element or a slice of the object, e.g. ``ts[1]`` or ``ts[1:2]."""
 
         data = self.data_hdu.data[index]
-
         return self.transform_slice(data)
 
     def transform_slice(self, data):
@@ -854,18 +858,34 @@ def _transform_slice_into_events(self, data):
             )
         except ValueError:
             pi_energy_func = None
-        if self.energy_column in data.dtype.names:
-            conversion = 1
+
+        if self.energy_column is not None and self.energy_column.lower() == "ebounds":
+            channels = self.data_hdu.data["PHA"]
+            elower = self.edata_hdu.data["E_MIN"]
+            ehigher = self.edata_hdu.data["E_MAX"]
+            emid = elower + (ehigher - elower) / 2.0
+
+            energy = np.array([emid[c] for c in channels])
+
             if (
-                hasattr(data.columns[self.energy_column], "unit")
-                and (unit := data.columns[self.energy_column].unit) is not None
+                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 = data[self.energy_column] * conversion
-        elif self.pi_column.lower() in [col.lower() for col in data.dtype.names]:
-            new_ts.pi = data[self.pi_column]
-            if pi_energy_func is not None:
-                new_ts.energy = pi_energy_func(new_ts.pi)
+            new_ts.energy = energy * conversion
+        else:
+            if self.energy_column in data.dtype.names:
+                conversion = 1
+                if (
+                    hasattr(data.columns[self.energy_column], "unit")
+                    and (unit := data.columns[self.energy_column].unit) is not None
+                ):
+                    conversion = (1 * u.Unit(unit)).to(u.keV).value
+                new_ts.energy = data[self.energy_column] * conversion
+            elif self.pi_column.lower() in [col.lower() for col in data.dtype.names]:
+                new_ts.pi = data[self.pi_column]
+                if pi_energy_func is not None:
+                    new_ts.energy = pi_energy_func(new_ts.pi)
 
         det_numbers = None
         if self.detector_key is not None and self.detector_key in data.dtype.names:
@@ -906,15 +926,15 @@ def _initialize_header_events(self, fname, force_hduname=None):
             If not None, the name of the HDU to read. If None, an extension called
             EVENTS or the first extension.
         """
-        hdulist = fits.open(fname)
+        hdulist = self._data
 
         if not force_hduname:
-            for hdu in hdulist:
+            for hdu in self._data:
                 if "TELESCOP" in hdu.header or "MISSION" in hdu.header:
                     probe_header = hdu.header
                     break
         else:
-            probe_header = hdulist[force_hduname].header
+            probe_header = self._data[force_hduname].header
 
         # We need the minimal information to read the mission database.
         # That is, the name of the mission/telescope, the instrument and,
@@ -955,12 +975,11 @@ def _initialize_header_events(self, fname, force_hduname=None):
 
         # If the EVENT/``force_hduname`` extension is not found, try the first extension
         # which is usually the one containing the data
-        if hduname not in hdulist:
+        if hduname not in self._data:
             warnings.warn(f"HDU {hduname} not found. Trying first extension")
             hduname = 1
         self._add_meta_attr("hduname", hduname)
-
-        header = hdulist[hduname].header
+        header = self._data[hduname].header
         if "OBS_ID" in header:
             self._add_meta_attr("obsid", header["OBS_ID"])
         if "TIMEDEL" in header:
@@ -970,7 +989,7 @@ def _initialize_header_events(self, fname, force_hduname=None):
         # No need to cope with dicts working badly with Netcdf, for example. The header
         # can be saved back and forth to files and be interpreted through
         # fits.Header.fromstring(self.header) when needed.
-        self._add_meta_attr("header", hdulist[self.hduname].header.tostring())
+        self._add_meta_attr("header", self._data[self.hduname].header.tostring())
         self._add_meta_attr("nphot", header["NAXIS2"])
 
         # These are the important keywords for timing.
@@ -1016,16 +1035,28 @@ def _initialize_header_events(self, fname, force_hduname=None):
         default_pi_column = get_key_from_mission_info(db, "ecol", "PI", instr, self.mode)
         if isinstance(default_pi_column, str):
             default_pi_column = _case_insensitive_search_in_list(
-                default_pi_column, hdulist[self.hduname].data.columns.names
+                default_pi_column, self._data[self.hduname].data.columns.names
+            )
+        if default_pi_column is None:
+            default_pi_column = get_key_from_mission_info(db, "ecol", "PHA", instr, self.mode)
+
+        if isinstance(default_pi_column, str):
+            default_pi_column = _case_insensitive_search_in_list(
+                default_pi_column, self._data[self.hduname].data.columns.names
             )
 
         self._add_meta_attr("pi_column", default_pi_column)
 
         # If a column named "energy" is found, we read it and assume the energy conversion
         # is already done.
         energy_column = _case_insensitive_search_in_list(
-            "energy", hdulist[self.hduname].data.columns.names
+            "energy", self._data[self.hduname].data.columns.names
         )
+        if energy_column is None:
+            energy_column = _case_insensitive_search_in_list(
+                "ebounds", [hdu.name for hdu in self._data]
+            )
+
         self._add_meta_attr("energy_column", energy_column)
 
     def _read_gtis(self, gti_file=None, det_numbers=None):