Merge pull request #940 from StingraySoftware/allow_remote

Allow loading of remote files, avoid extra loading of FITS files

Author: matteobachetti

docs/changes/940.feature.rst

@@ -0,0 +1 @@
+use the power of fits.open to load remote datasets

pyproject.toml

@@ -77,6 +77,7 @@ all = [
     "etils",
     "tfp-nightly",
     "typing_extensions",
+    "fsspec",
 ]
 docs = [
     "tomli>=1.1.0; python_version < '3.11'",

stingray/gti.py

@@ -4,10 +4,11 @@
 from collections.abc import Iterable
 import copy
 
-from astropy.io import fits
 from .utils import contiguous_regions, jit, HAS_NUMBA
 from .utils import assign_value_if_none, apply_function_if_none
 from .utils import check_iterables_close, is_sorted
+from .utils import fits_open_including_remote
+
 from stingray.exceptions import StingrayError
 from stingray.loggingconfig import setup_logger
 
@@ -122,7 +123,7 @@ def load_gtis(fits_file, gtistring=None):
 
     gtistring = assign_value_if_none(gtistring, "GTI")
     logger.info("Loading GTIS from file %s" % fits_file)
-    lchdulist = fits.open(fits_file, checksum=True, ignore_missing_end=True)
+    lchdulist = fits_open_including_remote(fits_file, checksum=True, ignore_missing_end=True)
     lchdulist.verify("warn")
 
     gtitable = lchdulist[gtistring].data
@@ -274,7 +275,7 @@ def get_gti_from_all_extensions(lchdulist, accepted_gtistrings=["GTI"], det_numb
     >>> assert np.allclose(gti, [[200, 250]])
     """
     if isinstance(lchdulist, str):
-        lchdulist = fits.open(lchdulist)
+        lchdulist = fits_open_including_remote(lchdulist)
     acc_gti_strs = copy.deepcopy(accepted_gtistrings)
     if det_numbers is not None:
         for i in det_numbers:

stingray/io.py

@@ -7,7 +7,6 @@
 from collections.abc import Iterable
 
 import numpy as np
-from astropy.io import fits
 from astropy.table import Table
 from astropy.logger import AstropyUserWarning
 import matplotlib.pyplot as plt
@@ -16,6 +15,7 @@
 
 import stingray.utils as utils
 from stingray.loggingconfig import setup_logger
+from stingray.utils import fits_open_including_remote
 
 
 from .utils import (
@@ -765,7 +765,7 @@ 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_hdus[self.hduname]
         self.gti_file = gti_file
         self._read_gtis(self.gti_file)
 
@@ -914,8 +914,8 @@ 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 = fits_open_including_remote(fname)
+        self.data_hdus = hdulist
         if not force_hduname:
             for hdu in hdulist:
                 if "TELESCOP" in hdu.header or "MISSION" in hdu.header:
@@ -1052,11 +1052,11 @@ def _read_gtis(self, gti_file=None, det_numbers=None):
         # So, here I'm reading a bunch of rows hoping that they represent the
         # detector number population
         if self.detector_key is not None:
-            with fits.open(self.fname) as hdul:
-                data = hdul[self.hduname].data
-                if self.detector_key in data.dtype.names:
-                    probe_vals = data[:100][self.detector_key]
-                    det_numbers = list(set(probe_vals))
+            hdul = self.data_hdus
+            data = hdul[self.hduname].data
+            if self.detector_key in data.dtype.names:
+                probe_vals = data[:100][self.detector_key]
+                det_numbers = list(set(probe_vals))
             del hdul
 
         accepted_gtistrings = self.gtistring.split(",")
@@ -1127,7 +1127,7 @@ def _get_idx_from_time_range(self, start, stop):
             time_edges[raw_max_idx] - stop >= 0
         ), f"Stop: {stop}; {time_edges[raw_max_idx] - stop} < 0"
 
-        with fits.open(self.fname) as hdulist:
+        with fits_open_including_remote(self.fname) as hdulist:
             filtered_times = hdulist[self.hduname].data[self.time_column][
                 raw_lower_edge : raw_upper_edge + 1
             ]
@@ -1219,7 +1219,7 @@ def _trace_nphots_in_file(self, nedges=1001):
 
         fname = self.fname
 
-        with fits.open(fname) as hdul:
+        with fits_open_including_remote(fname) as hdul:
             size = hdul[1].header["NAXIS2"]
             nedges = min(nedges, size // 10 + 2)
 
@@ -1335,7 +1335,7 @@ def read_header_key(fits_file, key, hdu=1):
         The value stored under ``key`` in ``fits_file``
     """
 
-    hdulist = fits.open(fits_file, ignore_missing_end=True)
+    hdulist = fits_open_including_remote(fits_file, ignore_missing_end=True)
     try:
         value = hdulist[hdu].header[key]
     except KeyError:  # pragma: no cover
@@ -1367,7 +1367,7 @@ def ref_mjd(fits_file, hdu=1):
         fits_file = fits_file[0]
         logger.info("opening %s" % fits_file)
 
-    hdulist = fits.open(fits_file, ignore_missing_end=True)
+    hdulist = fits_open_including_remote(fits_file, ignore_missing_end=True)
 
     ref_mjd_val = high_precision_keyword_read(hdulist[hdu].header, "MJDREF")
 

stingray/mission_support/rxte.py

@@ -5,6 +5,7 @@
 from astropy.table import Table
 
 from astropy.io import fits
+from stingray.utils import fits_open_including_remote
 
 c_match = re.compile(r"C\[(.*)\]")
 
@@ -343,7 +344,7 @@ def rxte_pca_event_file_interpretation(input_data, header=None, hduname=None):
 
     if isinstance(input_data, str):
         return rxte_pca_event_file_interpretation(
-            fits.open(input_data), header=header, hduname=hduname
+            fits_open_including_remote(input_data), header=header, hduname=hduname
         )
 
     if isinstance(input_data, fits.HDUList):

stingray/tests/test_events.py

@@ -15,6 +15,7 @@
 _HAS_PANDAS = importlib.util.find_spec("pandas") is not None
 _HAS_H5PY = importlib.util.find_spec("h5py") is not None
 _HAS_YAML = importlib.util.find_spec("yaml") is not None
+_HAS_FSSPEC = importlib.util.find_spec("fsspec") is not None
 
 
 class TestEvents(object):
@@ -248,6 +249,20 @@ def test_fits_standard(self):
         ev = ev.read(fname, fmt="hea")
         assert np.isclose(ev.mjdref, 55197.00076601852)
 
+    @pytest.mark.skipif("not _HAS_FSSPEC")
+    @pytest.mark.remote_data
+    def test_fits_standard_remote(self):
+        """Test that fits works with a standard event list
+        file.
+        """
+        fname = (
+            "s3://nasa-heasarc/swift/data/obs/2015_12/00037258040/xrt/event/"
+            "sw00037258040xwtw2st_cl.evt.gz"
+        )
+
+        ev = EventList.read(fname, fmt="hea")
+        assert np.isclose(ev.mjdref, 51910.00074287037)
+
     def test_fits_with_standard_file_and_calibrate_directly(self):
         """Test that fits works and calibration works."""
         fname = os.path.join(datadir, "monol_testA.evt")

stingray/utils.py

@@ -14,6 +14,7 @@
 from numpy import histogram as histogram_np
 from numpy import histogram2d as histogram2d_np
 from numpy import histogramdd as histogramdd_np
+from astropy.io import fits
 from .loggingconfig import setup_logger
 
 logger = setup_logger()
@@ -2410,3 +2411,69 @@ def _int_sum_non_zero(array):
         if a > 0:
             sum += int(a)
     return sum
+
+
+def fits_open_remote(filename, **kwargs):
+    """Open a remote FITS file.
+
+    This function attempts to open a FITS file using `astropy.io.fits.open`. If a
+    `PermissionError` is raised and the filename appears to be a remote URL,
+    it retries opening the file with fsspec.
+
+    Requires the `botocore` package to be installed.
+
+    Parameters
+    ----------
+    filename : str
+        The path or URL to the FITS file to open. Can be a local file path or a remote URL.
+    **kwargs
+        Additional keyword arguments passed to `astropy.io.fits.open`.
+
+    Returns
+    -------
+    hdulist : astropy.io.fits.HDUList
+        The opened FITS file as an HDUList object.
+
+    Raises
+    ------
+    PermissionError
+        If the file cannot be opened and anonymous access is not possible or fails.
+
+    """
+    import botocore
+
+    try:
+        # This will work for local files and remote files with proper permissions
+        return fits.open(filename, **kwargs)
+    except (PermissionError, botocore.exceptions.NoCredentialsError) as e:
+        if "://" in filename:
+            logger.info(f"Permission denied for {filename}, trying with fsspec.")
+            return fits.open(filename, use_fsspec=True, fsspec_kwargs={"anon": True}, **kwargs)
+        raise e
+
+
+def fits_open_including_remote(filename, **kwargs):
+    """Open a FITS file, including remote files with anonymous access if needed.
+
+    If the filename appears to be a remote URL, it calls `fits_open_remote` to handle
+    potential permission issues. Otherwise, it opens the file directly with
+    `astropy.io.fits.open`.
+
+    Parameters
+    ----------
+    filename : str
+        The path or URL to the FITS file to open. Can be a local file path or a remote URL.
+    **kwargs
+        Additional keyword arguments passed to `astropy.io.fits.open`.
+
+    Returns
+    -------
+    hdulist : astropy.io.fits.HDUList
+        The opened FITS file as an HDUList object.
+
+    """
+
+    if "://" in filename:
+        return fits_open_remote(filename, **kwargs)
+
+    return fits.open(filename, **kwargs)