ENH: Improve strain data read in and downsampling

bilby/gw/detector/strain_data.py

@@ -869,3 +869,166 @@ def check_frequency(self, freq):
             if notch.check_frequency(freq):
                 return True
         return False
+
+
+def resample_with_gwpy(data, sampling_frequency, **kwargs):
+    """
+    Resample a GWPy TimeSeries to a new sampling frequency.
+
+    Parameters:
+    ----------
+    data : gwpy.timeseries.TimeSeries
+        The input time series data to be resampled.
+    sampling_frequency : float
+        The target sampling frequency (Hz) for resampling.
+    **kwargs:
+        Additional keyword arguments passed to `gwpy.timeseries.TimeSeries.resample`.
+
+    Returns:
+    -------
+    gwpy.timeseries.TimeSeries
+        A new TimeSeries object resampled to the desired frequency.
+
+    """
+    return data.resample(sampling_frequency, **kwargs)
+
+
+def resample_with_lal(data, sampling_frequency, **kwargs):
+    """
+    Resample a GWPy TimeSeries using LAL's ResampleREAL8TimeSeries function.
+
+    Parameters:
+    ----------
+    data : gwpy.timeseries.TimeSeries
+        The input time series data to be resampled.
+    sampling_frequency : float
+        The target sampling frequency (Hz) for resampling.
+    **kwargs:
+        Not used, present for interface compatibility with `resample_with_gwpy`.
+
+    Returns:
+    -------
+    gwpy.timeseries.TimeSeries
+        A new TimeSeries object resampled to the desired frequency.
+
+    """
+    if kwargs:
+        raise ValueError(f"resample_with_lal does not support additional kwargs: {kwargs}")
+    try:
+        import lal
+        from gwpy.timeseries import TimeSeries
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("Cannot resample with lal: lal and gwpy are required")
+    lal_timeseries = data.to_lal()
+    lal.ResampleREAL8TimeSeries(
+        lal_timeseries, float(1 / sampling_frequency)
+    )
+    return TimeSeries(
+        lal_timeseries.data.data,
+        epoch=lal_timeseries.epoch,
+        dt=lal_timeseries.deltaT,
+    )
+
+
+RESAMPLING_FUNCTIONS = dict(
+    gwpy=resample_with_gwpy,
+    lal=resample_with_lal
+)
+
+
+def resample_timeseries(data, sampling_frequency, resampling_method="lal", **kwargs):
+    """
+    Resample a time series to a specified sampling frequency using a chosen method.
+
+    Parameters:
+    ----------
+    data : gwpy.timeseries.TimeSeries
+        The input time series data to be resampled.
+    sampling_frequency : float
+        The target sampling frequency (Hz) for resampling.
+    resampling_method : str, optional
+        The resampling method to use. Defaults to "lal".
+        Must be one of the methods defined in
+        `bilby.gw.detector.strain_data.RESAMPLING_FUNCTIONS`.
+    **kwargs:
+        Additional keyword arguments passed to the chosen resampling function.
+
+    Returns:
+    -------
+    gwpy.timeseries.TimeSeries
+        A new TimeSeries object resampled to the desired frequency.
+
+    Raises:
+    ------
+    ValueError
+        If the specified resampling method is not implemented.
+
+    """
+    if data.sample_rate.value == sampling_frequency:
+        logger.info("Sample rate matches data no resampling")
+        return data.copy()
+    elif resampling_method in RESAMPLING_FUNCTIONS:
+        logger.info(f"Resampling data to sampling_frequency {sampling_frequency} using {resampling_method}")
+        return RESAMPLING_FUNCTIONS[resampling_method](data, sampling_frequency, **kwargs)
+    else:
+        raise ValueError(f"Resampling method {resampling_method} not implemented")
+
+
+def find_and_read_data(start, end, ifo, frametype, channel, find_url_kwargs=None, read_kwargs=None,
+                       sampling_frequency=None, resampling_method="lal", resample_kwargs=None,
+                       dtype="float64"):
+    """
+    Find data using gw_data_find and then read it with gwpy. This assumes the data
+    exists on the provided host. Usually this is locally.
+
+    Parameters
+    ----------
+    start, end: float
+        The GPS start and end time
+    ifo: str
+        The detector to use, e.g. "H1", "L1", "V1", "K1", "A1"
+    frametype: str
+        The frametype to search for: not including the prepended detector name, e.g. "HOFT_C00_AR"
+    channel: str
+        The channel within the frame to read: not including the detector name, e.g. "GDS-CALIB_STRAIN_AR"
+    find_url_kwargs: dict
+        A dictionary of kwargs to pass to `gwdatafind.find_urls()`
+    read_kwargs: dict
+        A dictionary of kwargs to pass to `gwpy.timeseries.TimeSeries.read()`
+    sampling_frequency: float, optional
+        The target sampling frequency (Hz) to resample the data to. If not
+        given (default), no resampling is performed.
+    resample_kwargs: dict
+        A dictionary of kwargs to pass to the resampling function, see
+        `bilby.gw.detector.strain_data.resample_timeseries`.
+
+    Returns
+    -------
+    data: gwpy.timeseries.TimeSeries
+        The gwpy timeseries of the data
+
+    """
+    try:
+        from gwpy.timeseries import TimeSeries
+        from gwdatafind import find_urls
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("Cannot find and read data: gwpy and gwdatafind are required")
+
+    find_url_kwargs = find_url_kwargs or dict()
+    read_kwargs = read_kwargs or dict()
+    resample_kwargs = resample_kwargs or dict()
+
+    frametype_with_ifo = f"{ifo}_{frametype}"
+    channel_with_ifo = f"{ifo}:{channel}"
+    single_letter_ifo = ifo[0]
+
+    urls = find_urls(single_letter_ifo, frametype_with_ifo, start, end, **find_url_kwargs)
+
+    if len(urls) == 0:
+        raise ValueError(f"No data found for {ifo} {frametype} {start} {end}")
+
+    type_kwargs = dict(dtype=dtype, subok=True, copy=False)
+    data = TimeSeries.read(urls, channel_with_ifo, start=start, end=end, **read_kwargs).astype(**type_kwargs)
+    if sampling_frequency is not None:
+        data = resample_timeseries(data, sampling_frequency, resampling_method, **resample_kwargs)
+    return data

test/gw/detector/strain_data_test.py

@@ -3,8 +3,15 @@
 
 import numpy as np
 import scipy.signal
+from gwpy.timeseries import TimeSeries
 
 import bilby
+from bilby.gw.detector.strain_data import (
+    find_and_read_data,
+    resample_timeseries,
+    resample_with_gwpy,
+    resample_with_lal,
+)
 
 
 class TestInterferometerStrainData(unittest.TestCase):
@@ -420,5 +427,132 @@ def test_init_fail(self):
             bilby.gw.detector.strain_data.NotchList([(30, 20, 20)])
 
 
+class TestResampling(unittest.TestCase):
+    def setUp(self):
+        self.sampling_frequency = 512
+        self.data = TimeSeries(
+            np.random.normal(0, 1, self.sampling_frequency * 4),
+            sample_rate=self.sampling_frequency,
+            epoch=0,
+        )
+
+    def test_resample_with_gwpy(self):
+        new_data = resample_with_gwpy(self.data, self.sampling_frequency / 2)
+        self.assertIsInstance(new_data, TimeSeries)
+        self.assertEqual(new_data.sample_rate.value, self.sampling_frequency / 2)
+
+    def test_resample_with_lal(self):
+        new_data = resample_with_lal(self.data, self.sampling_frequency / 2)
+        self.assertIsInstance(new_data, TimeSeries)
+        self.assertEqual(new_data.sample_rate.value, self.sampling_frequency / 2)
+
+    def test_resample_with_lal_unsupported_kwargs(self):
+        with self.assertRaises(ValueError):
+            resample_with_lal(self.data, self.sampling_frequency / 2, window="hann")
+
+    def test_resample_timeseries_same_rate(self):
+        new_data = resample_timeseries(self.data, self.sampling_frequency)
+        self.assertIsNot(new_data, self.data)
+        self.assertTrue(np.array_equal(new_data.value, self.data.value))
+
+    def test_resample_timeseries_lal(self):
+        new_data = resample_timeseries(
+            self.data, self.sampling_frequency / 2, resampling_method="lal"
+        )
+        self.assertEqual(new_data.sample_rate.value, self.sampling_frequency / 2)
+
+    def test_resample_timeseries_gwpy(self):
+        new_data = resample_timeseries(
+            self.data, self.sampling_frequency / 2, resampling_method="gwpy"
+        )
+        self.assertEqual(new_data.sample_rate.value, self.sampling_frequency / 2)
+
+    def test_resample_timeseries_kwargs_passed_through(self):
+        new_data = resample_timeseries(
+            self.data,
+            self.sampling_frequency / 2,
+            resampling_method="gwpy",
+            window="hann",
+        )
+        self.assertEqual(new_data.sample_rate.value, self.sampling_frequency / 2)
+
+    def test_resample_timeseries_invalid_method(self):
+        with self.assertRaises(ValueError):
+            resample_timeseries(
+                self.data, self.sampling_frequency / 2, resampling_method="not-a-method"
+            )
+
+
+class TestFindAndReadData(unittest.TestCase):
+    def setUp(self):
+        self.start = 0
+        self.end = 4
+        self.ifo = "H1"
+        self.frametype = "HOFT_C00_AR"
+        self.channel = "GDS-CALIB_STRAIN_AR"
+        self.sampling_frequency = 512
+        self.data = TimeSeries(
+            np.random.normal(0, 1, self.sampling_frequency * 4),
+            sample_rate=self.sampling_frequency,
+            epoch=self.start,
+        )
+
+    @mock.patch("gwdatafind.find_urls")
+    @mock.patch("gwpy.timeseries.TimeSeries.read")
+    def test_find_and_read_data(self, mock_read, mock_find_urls):
+        mock_find_urls.return_value = ["file://fake/H-H1_HOFT_C00_AR-0-4.gwf"]
+        mock_read.return_value = self.data
+
+        data = find_and_read_data(
+            self.start, self.end, self.ifo, self.frametype, self.channel,
+            sampling_frequency=self.sampling_frequency,
+        )
+
+        mock_find_urls.assert_called_once_with(
+            "H", f"{self.ifo}_{self.frametype}", self.start, self.end
+        )
+        mock_read.assert_called_once_with(
+            mock_find_urls.return_value,
+            f"{self.ifo}:{self.channel}",
+            start=self.start,
+            end=self.end,
+        )
+        self.assertEqual(data.sample_rate.value, self.sampling_frequency)
+
+    @mock.patch("gwdatafind.find_urls")
+    def test_find_and_read_data_no_urls_raises(self, mock_find_urls):
+        mock_find_urls.return_value = []
+        with self.assertRaises(ValueError):
+            find_and_read_data(
+                self.start, self.end, self.ifo, self.frametype, self.channel,
+                sampling_frequency=self.sampling_frequency,
+            )
+
+    @mock.patch("gwdatafind.find_urls")
+    @mock.patch("gwpy.timeseries.TimeSeries.read")
+    def test_find_and_read_data_resamples(self, mock_read, mock_find_urls):
+        mock_find_urls.return_value = ["file://fake/H-H1_HOFT_C00_AR-0-4.gwf"]
+        mock_read.return_value = self.data
+
+        data = find_and_read_data(
+            self.start, self.end, self.ifo, self.frametype, self.channel,
+            sampling_frequency=self.sampling_frequency / 2,
+        )
+
+        self.assertEqual(data.sample_rate.value, self.sampling_frequency / 2)
+
+    @mock.patch("gwdatafind.find_urls")
+    @mock.patch("gwpy.timeseries.TimeSeries.read")
+    def test_find_and_read_data_no_sampling_frequency_skips_resampling(self, mock_read, mock_find_urls):
+        mock_find_urls.return_value = ["file://fake/H-H1_HOFT_C00_AR-0-4.gwf"]
+        mock_read.return_value = self.data
+
+        data = find_and_read_data(
+            self.start, self.end, self.ifo, self.frametype, self.channel,
+        )
+
+        self.assertEqual(data.sample_rate.value, self.sampling_frequency)
+
+
 if __name__ == "__main__":
     unittest.main()