Repository: StingraySoftware/stingray
Pull request #392 (2019-03-27 19:22:55)

New feature! Added lightcurve splitting
I keep running into the situation where I need to split light curves along data gaps, i.e. make a light curve of a contiguous light segment before a gap and a contiguous segment after a gap. I'm probably not the only one, so I wrote a method for `Lightcurve` that can do this.

Comments welcome! :) 


Patch:
diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index a640244e..68813a95 100644
--- a/stingray/lightcurve.py
+++ b/stingray/lightcurve.py
@@ -921,6 +921,59 @@ class Lightcurve(object):
 
         return self._truncate_by_index(start, stop)
 
+    def split(self, min_gap, min_points=1):
+        """
+        For data with gaps, it can sometimes be useful to be able to split
+        the light curve into separate, evenly sampled objects along those
+        data gaps. This method allows to do this: it finds data gaps of a
+        specified minimum size, and produces a list of new `Lightcurve`
+        objects for each contiguous segment.
+
+        Parameters
+        ----------
+        min_gap : float
+            The length of a data gap, in the same units as the `time` attribute
+            of the `Lightcurve` object. Any smaller gaps will be ignored, any
+            larger gaps will be identified and used to split the light curve.
+
+        min_points : int, default 1
+            The minimum number of data points in each light curve. Light
+            curves with fewer data points will be ignored.
+
+        Returns
+        -------
+        lc_split : iterable of `Lightcurve` objects
+            The list of all contiguous light curves
+
+        Example
+        -------
+        >>> time = np.array([1, 2, 3, 6, 7, 8, 11, 12, 13])
+        >>> counts = np.random.rand(time.shape[0])
+        >>> lc = Lightcurve(time, counts)
+        >>> split_lc = lc.split(1.5)
+
+        """
+        # calculate the difference between time bins
+        tdiff = np.diff(self.time)
+        # find all distances between time bins that are larger than `min_gap`
+        gap_idx = np.where(tdiff >= min_gap)[0]
+
+        # tolerance for the newly created GTIs: Note that this seems to work with a 
+        # tolerance of 2, but not if I substitute 10, and I don't know why
+        epsilon = np.min(tdiff)/2.0
+
+        # calculate new GTIs
+        gti_start = np.hstack([self.time[0]-epsilon, self.time[gap_idx+1]-epsilon])
+        gti_stop = np.hstack([self.time[gap_idx]+epsilon, self.time[-1]+epsilon])
+        
+        gti = np.vstack([gti_start, gti_stop]).T
+        if hasattr(self, 'gti') and self.gti is not None:
+            gti = cross_two_gtis(self.gti, gti)
+        self.gti = gti
+        
+        lc_split = self.split_by_gti(min_points=min_points)
+        return lc_split
+
     def sort(self, reverse=False):
         """
         Sort a Lightcurve object by time.
@@ -1262,11 +1315,17 @@ class Lightcurve(object):
         else:
             utils.simon("Format not understood.")
 
-    def split_by_gti(self):
+    def split_by_gti(self, min_points=2):
         """
         Split the current :class:`Lightcurve` object into a list of :class:`Lightcurve` objects, one
         for each continuous GTI segment as defined in the ``gti`` attribute.
 
+        Parameters
+        ----------
+        min_points : int, default 1
+            The minimum number of data points in each light curve. Light
+            curves with fewer data points will be ignored.
+
         Returns
         -------
         list_of_lcs : list
@@ -1278,6 +1337,13 @@ class Lightcurve(object):
         for i in range(len(start_bins)):
             start = start_bins[i]
             stop = stop_bins[i]
+
+            if np.isclose(stop-start, 1):
+                logging.warning("Segment with a single time bin! Ignoring this segment!")
+                continue
+            if (stop - start) < min_points:
+                continue
+
             # Note: GTIs are consistent with default in this case!
             new_lc = Lightcurve(self.time[start:stop], self.counts[start:stop],
                                 err=self.counts_err[start:stop],


Test patch:
diff --git a/stingray/tests/test_lightcurve.py b/stingray/tests/test_lightcurve.py
index 33113854..b5eabfa5 100644
--- a/stingray/tests/test_lightcurve.py
+++ b/stingray/tests/test_lightcurve.py
@@ -566,6 +566,69 @@ class TestLightcurve(object):
         np.testing.assert_almost_equal(lc2.gti[-1][-1], 2.5)
         assert lc2.mjdref == lc.mjdref
 
+    def test_split_with_two_segments(self):
+        test_time = np.array([1, 2, 3, 6, 7, 8])
+        test_counts = np.random.rand(len(test_time))
+        lc_test = Lightcurve(test_time, test_counts)
+        slc = lc_test.split(1.5)
+
+        assert len(slc) == 2
+
+    def test_split_has_correct_data_points(self):
+        test_time = np.array([1, 2, 3, 6, 7, 8])
+        test_counts = np.random.rand(len(test_time))
+        lc_test = Lightcurve(test_time, test_counts)
+        slc = lc_test.split(1.5)
+
+        assert np.all((slc[0].time == [1, 2, 3]))
+        assert np.all((slc[1].time == [6, 7 ,8]))
+        assert np.all((slc[0].counts == test_counts[:3]))
+        assert np.all((slc[1].counts == test_counts[3:]))
+
+    def test_split_with_three_segments(self):
+        test_time = np.array([1, 2, 3, 6, 7, 8, 10, 11, 12])
+        test_counts = np.random.rand(len(test_time))
+        lc_test = Lightcurve(test_time, test_counts)
+        slc = lc_test.split(1.5)
+
+        assert len(slc) == 3
+
+    def test_threeway_split_has_correct_data_points(self):
+        test_time = np.array([1, 2, 3, 6, 7, 8, 10, 11, 12])
+        test_counts = np.random.rand(len(test_time))
+        lc_test = Lightcurve(test_time, test_counts)
+        slc = lc_test.split(1.5)
+
+        assert np.all((slc[0].time == [1, 2, 3]))
+        assert np.all((slc[1].time == [6, 7 ,8]))
+        assert np.all((slc[2].time == [10, 11 ,12]))
+        assert np.all((slc[0].counts == test_counts[:3]))
+        assert np.all((slc[1].counts == test_counts[3:6]))
+        assert np.all((slc[2].counts == test_counts[6:]))
+
+    def test_split_with_gtis(self):
+        test_time = np.array([1, 2, 3, 6, 7, 8, 10, 11, 12])
+        test_counts = np.random.rand(len(test_time))
+        gti = [[0,4], [9, 13]]
+        lc_test = Lightcurve(test_time, test_counts, gti=gti)
+        slc = lc_test.split(1.5)
+
+        assert np.all((slc[0].time == [1, 2, 3]))
+        assert np.all((slc[1].time == [10, 11 ,12]))
+        assert np.all((slc[0].counts == test_counts[:3]))
+        assert np.all((slc[1].counts == test_counts[6:]))
+
+    def test_consecutive_gaps(self):
+        test_time = np.array([1, 2, 3, 6, 9, 10, 11])
+        test_counts = np.random.rand(len(test_time))
+        lc_test = Lightcurve(test_time, test_counts)
+        slc = lc_test.split(1.5)
+
+        assert np.all((slc[0].time == [1, 2, 3]))
+        assert np.all((slc[1].time == [9, 10, 11]))
+        assert np.all((slc[0].counts == test_counts[:3]))
+        assert np.all((slc[1].counts == test_counts[4:]))
+
     def test_sort(self):
         _times = [2, 1, 3, 4]
         _counts = [40, 10, 20, 5]
@@ -718,6 +781,21 @@ class TestLightcurve(object):
         assert np.all(lc0.gti == [[0.5, 4.5]])
         assert np.all(lc1.gti == [[5.5, 7.5]])
 
+    def test_split_lc_by_gtis_minpoints(self):
+        times = [1, 2, 3, 4, 5, 6, 7, 8]
+        counts = [1, 1, 1, 1, 2, 3, 3, 2]
+        gti = [[0.5, 3.5], [3.5, 5.5], [5.5, 8.5]]
+        min_points = 3
+
+        lc = Lightcurve(times, counts, gti=gti)
+        list_of_lcs = lc.split_by_gti(min_points=min_points)
+        lc0 = list_of_lcs[0]
+        lc1 = list_of_lcs[1]
+        assert np.all(lc0.time == [1, 2, 3])
+        assert np.all(lc1.time == [6, 7, 8])
+        assert np.all(lc0.counts == [1, 1, 1])
+        assert np.all(lc1.counts == [3, 3, 2])
+
     def test_shift(self):
         times = [1, 2, 3, 4, 5, 6, 7, 8]
         counts = [1, 1, 1, 1, 2, 3, 3, 2]