Repository: StingraySoftware/stingray
Pull request #231 (2017-07-27 22:09:48)

Bugfix + Add method to analyze regular chunks of a lightcurve with generic functions
This comes from the need of calculating excess variance in regular chunks of light curves (AGN stuff).

I made it a little more generic, so that people can pass their own statistics to be calculated on the light curve chunks. Excess variance is the example in `test_analyze_lc_chunks_fvar`.

Plus I fixed Fvar, that was missing a square root (screaming face)

EDIT: Plus I allowed to use rolling windows (steps of fractions of `chunk_length` in starting times), which involved also modifying slightly `bin_intervals_from_gtis`.

EDIT 2: Notebook here: https://github.com/StingraySoftware/notebooks/pull/16


Patch:
diff --git a/stingray/gti.py b/stingray/gti.py
index 306cc6c3..e8d82e6d 100644
--- a/stingray/gti.py
+++ b/stingray/gti.py
@@ -501,7 +501,8 @@ def join_gtis(gti0, gti1):
     return np.sort(final_gti, axis=0)
 
 
-def time_intervals_from_gtis(gtis, chunk_length):
+def time_intervals_from_gtis(gtis, chunk_length, fraction_step=1,
+                             epsilon=1e-5):
     """Returns equal time intervals compatible with GTIs.
 
     Used to start each FFT/PDS/cospectrum from the start of a GTI,
@@ -513,7 +514,10 @@ def time_intervals_from_gtis(gtis, chunk_length):
         [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
     chunk_length : float
         Length of the chunks
-
+    fraction_step : float
+        If the step is not a full chunk_length but less (e.g. a moving window),
+        this indicates the ratio between step step and `chunk_length` (e.g.
+        0.5 means that the window shifts of half chunk_length)
     Returns
     -------
     spectrum_start_times : array-like
@@ -525,11 +529,11 @@ def time_intervals_from_gtis(gtis, chunk_length):
     """
     spectrum_start_times = np.array([], dtype=np.longdouble)
     for g in gtis:
-        if g[1] - g[0] < chunk_length:
+        if g[1] - g[0] + epsilon < chunk_length:
             continue
 
-        newtimes = np.arange(g[0], g[1] - chunk_length,
-                             np.longdouble(chunk_length),
+        newtimes = np.arange(g[0], g[1] - chunk_length + epsilon,
+                             np.longdouble(chunk_length) * fraction_step,
                              dtype=np.longdouble)
         spectrum_start_times = \
             np.append(spectrum_start_times,
@@ -540,7 +544,8 @@ def time_intervals_from_gtis(gtis, chunk_length):
     return spectrum_start_times, spectrum_start_times + chunk_length
 
 
-def bin_intervals_from_gtis(gtis, chunk_length, time, dt=None, epsilon=0.001):
+def bin_intervals_from_gtis(gtis, chunk_length, time, dt=None, fraction_step=1,
+                            epsilon=0.001):
     """Similar to intervals_from_gtis, but given an input time array.
 
     Used to start each FFT/PDS/cospectrum from the start of a GTI,
@@ -564,6 +569,10 @@ def bin_intervals_from_gtis(gtis, chunk_length, time, dt=None, epsilon=0.001):
         Time resolution of the light curve.
     epsilon : float, default 0.001
         The tolerance, in fraction of dt, for the comparisons at the borders
+    fraction_step : float
+        If the step is not a full chunk_length but less (e.g. a moving window),
+        this indicates the ratio between step step and `chunk_length` (e.g.
+        0.5 means that the window shifts of half chunk_length)
 
     Examples
     --------
@@ -621,8 +630,8 @@ def bin_intervals_from_gtis(gtis, chunk_length, time, dt=None, epsilon=0.001):
         if time[stopbin - 1] > g[1] - dt/2 + epsilon*dt:
             stopbin -= 1
 
-        newbins = np.arange(startbin, stopbin - nbin + 1, nbin,
-                            dtype=np.long)
+        newbins = np.arange(startbin, stopbin - nbin + 1,
+                            int(nbin * fraction_step), dtype=np.long)
         spectrum_start_bins = \
             np.append(spectrum_start_bins,
                       newbins)
diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index 6e98da9f..a4f1e517 100644
--- a/stingray/lightcurve.py
+++ b/stingray/lightcurve.py
@@ -4,6 +4,7 @@ Definition of :class:`Lightcurve`.
 :class:`Lightcurve` is used to create light curves out of photon counting data
 or to save existing light curves in a class that's easy to use.
 """
+from __future__ import division, print_function
 import logging
 import numpy as np
 import stingray.io as io
@@ -11,7 +12,7 @@ import stingray.utils as utils
 from stingray.exceptions import StingrayError
 from stingray.utils import simon, assign_value_if_none, baseline_als
 from stingray.gti import cross_two_gtis, join_gtis, gti_border_bins
-from stingray.gti import check_gtis, create_gti_mask
+from stingray.gti import check_gtis, create_gti_mask, bin_intervals_from_gtis
 from astropy.stats import poisson_conf_interval
 
 __all__ = ["Lightcurve"]
@@ -156,8 +157,9 @@ class Lightcurve(object):
                 err_low, err_high = poisson_conf_interval(np.asarray(counts),
                     interval='frequentist-confidence', sigma=1)
                 # calculate approximately symmetric uncertainties
-                err = (np.absolute(err_low) + np.absolute(err_high) -
-                       2 * np.asarray(counts))/2.0
+                err_low -= np.asarray(counts)
+                err_high -= np.asarray(counts)
+                err = (np.absolute(err_low) + np.absolute(err_high))/2.0
                 # other estimators can be implemented for other statistics
             else:
                 simon("Stingray only uses poisson err_dist at the moment, "
@@ -789,11 +791,129 @@ class Lightcurve(object):
 
         new_counts, new_time, new_counts_err = zip(*sorted(zip(self.counts, self.time, self.counts_err) , reverse=reverse))
 
-
         self.time = np.asarray(new_time)
         self.counts = np.asarray(new_counts)
         self.counts_err = np.asarray(new_counts_err)
 
+    def estimate_chunk_length(self, min_total_counts=100, min_time_bins=100):
+        """Choose a reasonable chunk length.
+
+        The user specifies a condition on the total counts in each chunk and
+        the minimum number of time bins.
+
+        Other Parameters
+        ----------------
+        min_total_counts : int
+            Minimum number of counts for each chunk
+        min_time_bins : int
+            Minimum number of time bins
+
+        Returns
+        -------
+        chunk_length : float
+            The length of the light curve chunks that satisfies the conditions
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> time = np.arange(150)
+        >>> count = np.zeros_like(time) + 3
+        >>> lc = Lightcurve(time, count)
+        >>> lc.estimate_chunk_length(min_total_counts=10, min_time_bins=3)
+        4.0
+        >>> lc.estimate_chunk_length(min_total_counts=10, min_time_bins=5)
+        5.0
+        >>> count[2:4] = 1
+        >>> lc = Lightcurve(time, count)
+        >>> lc.estimate_chunk_length(min_total_counts=3, min_time_bins=1)
+        4.0
+        """
+
+        rough_estimate = np.ceil(min_total_counts / self.meancounts) * self.dt
+
+        chunk_length = np.max([rough_estimate, min_time_bins * self.dt])
+
+        keep_searching = True
+        while keep_searching:
+            start_times, stop_times, results = \
+                self.analyze_lc_chunks(chunk_length, np.sum)
+            mincounts = np.min(results)
+            if mincounts >= min_total_counts:
+                keep_searching = False
+            else:
+                chunk_length *= np.ceil(min_total_counts / mincounts) * self.dt
+
+        return chunk_length
+
+    def analyze_lc_chunks(self, chunk_length, func, fraction_step=1, **kwargs):
+        """Analyze chunks of the light curve with any function.
+
+        Parameters
+        ----------
+        chunk_length : float
+            Length in seconds of the light curve chunks
+        func : function
+            Function accepting a `Lightcurve` object as single argument, plus
+            possible additional keyword arguments, and returning a number or a
+            tuple - e.g., (result, error) where both result and error are
+            numbers.
+
+        Other parameters
+        ----------------
+        fraction_step : float
+            If the step is not a full chunk_length but less (e.g. a moving window),
+            this indicates the ratio between step step and `chunk_length` (e.g.
+            0.5 means that the window shifts of half chunk_length)
+        kwargs : keyword arguments
+            These additional keyword arguments, if present, they will be passed
+            to `func`
+
+        Returns
+        -------
+        start_times : array
+            Lower time boundaries of all chunks.
+        stop_times : array
+            Higher time boundaries of all chunks.
+        result : array of N elements
+            The result of `func` for each chunk of the light curve
+
+        Examples
+        --------
+        >>> import numpy as np
+        >>> time = np.arange(0, 10, 0.1)
+        >>> counts = np.zeros_like(time) + 10
+        >>> lc = Lightcurve(time, counts)
+        >>> # Define a function that calculates the mean
+        >>> mean_func = lambda x: np.mean(x)
+        >>> # Calculate the mean in chunks of 5 seconds
+        >>> start, stop, res = lc.analyze_lc_chunks(5, mean_func)
+        >>> len(res) == 2
+        True
+        >>> np.all(res == 10)
+        True
+        """
+        start, stop = bin_intervals_from_gtis(self.gti, chunk_length,
+                                              self.time,
+                                              fraction_step=fraction_step,
+                                              dt=self.dt)
+        start_times = self.time[start] - self.dt * 0.5
+
+        # Remember that stop is one element above the last element, because
+        # it's defined to be used in intervals start:stop
+        stop_times = self.time[stop - 1] + self.dt * 1.5
+
+        results = []
+        for i, (st, sp) in enumerate(zip(start, stop)):
+            lc_filt = self[st:sp]
+            res = func(lc_filt, **kwargs)
+            results.append(res)
+
+        results = np.array(results)
+
+        if len(results.shape) == 2:
+            results = [results[:, i] for i in range(results.shape[1])]
+        return start_times, stop_times, results
+
     def plot(self, witherrors=False, labels=None, axis=None, title=None,
              marker='-', save=False, filename=None):
         """
@@ -941,6 +1061,18 @@ class Lightcurve(object):
 
         return list_of_lcs
 
-    def apply_gtis(self):
-        good = create_gti_mask(self.time, lc.gti)
+    def _apply_gtis(self):
+        """Apply GTIs to a light curve after modification."""
+        check_gtis(self.gti)
+
+        good = create_gti_mask(self.time, self.gti)
+
+        self.time = self.time[good]
+        self.counts = self.counts[good]
+        self.counts_err = self.counts_err[good]
+        self.countrate = self.countrate[good]
+        self.countrate_err = self.countrate_err[good]
 
+        self.meanrate = np.mean(self.countrate)
+        self.meancounts = np.mean(self.counts)
+        self.n = self.counts.shape[0]
diff --git a/stingray/utils.py b/stingray/utils.py
index 74868918..a9695b33 100644
--- a/stingray/utils.py
+++ b/stingray/utils.py
@@ -390,9 +390,9 @@ def excess_variance(lc, normalization='fvar'):
     lc_actual_var = np.var(lc.counts)
     var_xs = lc_actual_var - lc_mean_var
     mean_lc = np.mean(lc.counts)
-    mean_ctvar = np.mean(mean_lc ** 2)
+    mean_ctvar = mean_lc ** 2
 
-    fvar = var_xs / mean_ctvar
+    fvar = np.sqrt(var_xs / mean_ctvar)
 
     N = len(lc.counts)
     var_xs_err_A = np.sqrt(2 / N) * lc_mean_var / mean_lc ** 2


Test patch:
diff --git a/stingray/tests/test_gti.py b/stingray/tests/test_gti.py
index 81f86f2c..1166a2a3 100644
--- a/stingray/tests/test_gti.py
+++ b/stingray/tests/test_gti.py
@@ -128,6 +128,14 @@ class TestGTI(object):
         assert np.all(start_times == np.array([0, 128, 256, 1022]))
         assert np.all(stop_times == np.array([0, 128, 256, 1022]) + 128)
 
+    def test_time_intervals_from_gtis_frac(self):
+        """Test the division of start and end times to calculate spectra."""
+        start_times, stop_times = \
+            time_intervals_from_gtis([[0, 400], [1022, 1200],
+                                      [1210, 1220]], 128, fraction_step=0.5)
+        assert np.all(start_times == np.array([0, 64, 128, 192, 256, 1022]))
+        assert np.all(stop_times == start_times + 128)
+
     def test_bin_intervals_from_gtis(self):
         """Test the division of start and end times to calculate spectra."""
         times = np.arange(0.5, 13.5)
@@ -137,6 +145,16 @@ class TestGTI(object):
         assert np.all(start_bins == np.array([0, 2, 6]))
         assert np.all(stop_bins == np.array([2, 4, 8]))
 
+    def test_bin_intervals_from_gtis_frac(self):
+        """Test the division of start and end times to calculate spectra."""
+        times = np.arange(0.5, 13.5)
+        start_bins, stop_bins = \
+            bin_intervals_from_gtis([[0, 5], [6, 8]], 2, times,
+                                    fraction_step=0.5)
+
+        assert np.all(start_bins == np.array([0, 1, 2, 3, 6]))
+        assert np.all(stop_bins == np.array([2, 3, 4, 5, 8]))
+
     def test_gti_border_bins(self):
         times = np.arange(0.5, 2.5)
 
diff --git a/stingray/tests/test_lightcurve.py b/stingray/tests/test_lightcurve.py
index 47bbd00f..c493c433 100644
--- a/stingray/tests/test_lightcurve.py
+++ b/stingray/tests/test_lightcurve.py
@@ -89,6 +89,40 @@ class TestLightcurve(object):
         lc = Lightcurve(self.times, self.counts)
         assert lc.n == 4
 
+    def test_analyze_lc_chunks(self):
+        lc = Lightcurve(self.times, self.counts, gti=self.gti)
+
+        def func(lc):
+            return lc.time[0]
+        start, stop, res = lc.analyze_lc_chunks(2, func)
+        assert start[0] == 0.5
+        assert np.all(start + lc.dt / 2 == res)
+
+    def test_analyze_lc_chunks_fvar_fracstep(self):
+        dt = 0.1
+        tstart = 0
+        tstop = 100
+        times = np.arange(tstart, tstop, dt)
+        gti = np.array([[tstart - dt/2, tstop - dt/2]])
+        # Simulate something *clearly* non-constant
+        counts = np.random.poisson(
+            10000 + 2000 * np.sin(2 * np.pi * times))
+
+        lc = Lightcurve(times, counts, gti=gti)
+
+        def excvar(lc):
+            from stingray.utils import excess_variance
+            return excess_variance(lc, normalization='fvar')
+
+        start, stop, res = lc.analyze_lc_chunks(20, excvar, fraction_step=0.5)
+        # excess_variance returns fvar and fvar_err
+        res, res_err = res
+
+        assert np.allclose(start[0], gti[0, 0])
+        assert np.all(res > 0)
+        # This must be a clear measurement of fvar
+        assert np.all(res > res_err)
+
     def test_bin_edges(self):
         bin_lo = [0.5,  1.5,  2.5,  3.5]
         bin_hi = [1.5,  2.5,  3.5,  4.5]
@@ -660,4 +694,16 @@ class TestLightcurveRebin(object):
     def test_change_mjdref(self):
         lc_new = self.lc.change_mjdref(57000)
         assert lc_new.mjdref == 57000
-    
\ No newline at end of file
+
+    def test_apply_gtis(self):
+        time = np.arange(150)
+        count = np.zeros_like(time) + 3
+        lc = Lightcurve(time, count, gti=[[-0.5, 150.5]])
+        lc.gti = [[-0.5, 2.5], [12.5, 14.5]]
+        lc._apply_gtis()
+        assert lc.n == 5
+        assert np.all(lc.time == np.array([0, 1, 2, 13, 14]))
+        lc.gti = [[-0.5, 10.5]]
+        lc._apply_gtis()
+        assert lc.n == 3
+        assert np.all(lc.time == np.array([0, 1, 2]))