Repository: StingraySoftware/stingray
Pull request #268 (2017-10-15 17:10:57)

Add offset correction to baseline in light curve
This change in the methods allows to add an offset to the baseline calculated by `baseline_als`, so that it aligns with the mean of the light curve.
I changed the method in `utils.py` and its usage in `Lightcurve.baseline()`


Patch:
diff --git a/astropy_helpers b/astropy_helpers
index d23a53f4..14ca346b 160000
--- a/astropy_helpers
+++ b/astropy_helpers
@@ -1,1 +1,1 @@
-Subproject commit d23a53f46dd1c3703e5eee63dca3f53bd18a4e8b
+Subproject commit 14ca346b0da3e92e65bdc398cc8f726cbd7be57d
diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index ab661794..b230836c 100644
--- a/stingray/lightcurve.py
+++ b/stingray/lightcurve.py
@@ -434,28 +434,28 @@ class Lightcurve(object):
                              "object !")
 
     def __eq__(self,other_lc):
-	    """
-	        Compares two :class:`Lightcurve` objects.
-
-	        Light curves are equal only if their counts as well as times at which those counts occur equal.
-
-	        Examples
-	        --------
-	        >>> time = [1, 2, 3]
-	        >>> count1 = [100, 200, 300]
-	        >>> count2 = [100, 200, 300]
-	        >>> lc1 = Lightcurve(time, count1)
-	        >>> lc2 = Lightcurve(time, count2)
-	        >>> lc1 == lc2
-	        True
-	    """
-	    if not isinstance(other_lc, Lightcurve):
-	        raise ValueError('Lightcurve can only be compared with a Lightcurve Object')
-	    if (np.allclose(self.time, other_lc.time) and np.allclose(self.counts, other_lc.counts)):
-	        return True
-	    return False
-
-    def baseline(self, lam, p, niter=10):
+        """
+            Compares two :class:`Lightcurve` objects.
+
+            Light curves are equal only if their counts as well as times at which those counts occur equal.
+
+            Example
+            -------
+            >>> time = [1, 2, 3]
+            >>> count1 = [100, 200, 300]
+            >>> count2 = [100, 200, 300]
+            >>> lc1 = Lightcurve(time, count1)
+            >>> lc2 = Lightcurve(time, count2)
+            >>> lc1 == lc2
+            True
+        """
+        if not isinstance(other_lc, Lightcurve):
+            raise ValueError('Lightcurve can only be compared with a Lightcurve Object')
+        if (np.allclose(self.time, other_lc.time) and np.allclose(self.counts, other_lc.counts)):
+            return True
+        return False
+
+    def baseline(self, lam, p, niter=10, offset_correction=False):
         """Calculate the baseline of the light curve, accounting for GTIs.
 
         Parameters
@@ -466,11 +466,22 @@ class Lightcurve(object):
         p : float
             "asymmetry" parameter. Smaller values make the baseline more
             "horizontal". Typically 0.001 < p < 0.1, but not necessary.
+
+        Other parameters
+        ----------------
+        offset_correction : bool, default False
+            by default, this method does not align to the running mean of the
+            light curve, but it goes below the light curve. Setting align to
+            True, an additional step is done to shift the baseline so that it
+            is shifted to the middle of the light curve noise distribution.
         """
         baseline = np.zeros_like(self.time)
         for g in self.gti:
             good = create_gti_mask(self.time, [g])
-            baseline[good] = baseline_als(self.counts[good], lam, p, niter)
+            _, baseline[good] = \
+                baseline_als(self.time[good], self.counts[good], lam, p,
+                             niter, offset_correction=offset_correction,
+                             return_baseline=True)
 
         return baseline
 
diff --git a/stingray/utils.py b/stingray/utils.py
index e4ede6ab..4901540f 100644
--- a/stingray/utils.py
+++ b/stingray/utils.py
@@ -32,6 +32,19 @@ except ImportError:
             return wrapped_f
 
 
+try:
+    from statsmodels.robust import mad as mad  # pylint: disable=unused-import
+except ImportError:
+    def mad(data, c=0.6745, axis=None):
+        """Straight from statsmodels's source code, adapted"""
+        data = np.asarray(data)
+        if axis is not None:
+            center = np.apply_over_axes(np.median, data, axis)
+        else:
+            center = np.median(data)
+        return np.median((np.fabs(data - center)) / c, axis=axis)
+
+
 __all__ = ['simon', 'rebin_data', 'rebin_data_log', 'look_for_array_in_array',
            'is_string', 'is_iterable', 'order_list_of_arrays',
            'optimal_bin_time', 'contiguous_regions', 'is_int',
@@ -346,30 +359,71 @@ def get_random_state(random_state=None):
     return random_state
 
 
-def baseline_als(y, lam, p, niter=10):
+def _offset(x, off):
+    """An offset."""
+    return off
+
+
+def offset_fit(x, y, offset_start=0):
+    """Fit a constant offset to the data.
+
+    Parameters
+    ----------
+    x : array-like
+    y : array-like
+    offset_start : float
+        Constant offset, initial value
+
+    Returns
+    -------
+    offset : float
+        Fitted offset
+    """
+    from scipy.optimize import curve_fit
+    par, _ = curve_fit(_offset, x, y, [offset_start],
+                       maxfev=6000)
+    return par[0]
+
+
+def _als(y, lam, p, niter=10):
     """Baseline Correction with Asymmetric Least Squares Smoothing.
 
     Modifications to the routine from Eilers & Boelens 2005
-    https://www.researchgate.net/publication/228961729_Technical_Report_Baseline_Correction_with_Asymmetric_Least_Squares_Smoothing
+    https://www.researchgate.net/publication/
+        228961729_Technical_Report_Baseline_Correction_with_
+        Asymmetric_Least_Squares_Smoothing
     The Python translation is partly from
-    http://stackoverflow.com/questions/29156532/python-baseline-correction-library
+    http://stackoverflow.com/questions/29156532/
+        python-baseline-correction-library
 
     Parameters
     ----------
-    y : array of floats
-        the "light curve". It assumes equal spacing.
+    y : array-like
+        the data series corresponding to x
     lam : float
-        "smoothness" parameter. Larger values make the baseline stiffer
-        Typically 1e2 < lam < 1e9
+        the lambda parameter of the ALS method. This control how much the
+        baseline can adapt to local changes. A higher value corresponds to a
+        stiffer baseline
     p : float
-        "asymmetry" parameter. Smaller values make the baseline more
-        "horizontal". Typically 0.001 < p < 0.1, but not necessary.
+        the asymmetry parameter of the ALS method. This controls the overall
+        slope tollerated for the baseline. A higher value correspond to a
+        higher possible slope
+
+    Other parameters
+    ----------------
+    niter : int
+        The number of iterations to perform
+
+    Returns
+    -------
+    z : array-like, same size as y
+        Fitted baseline.
     """
     from scipy import sparse
     L = len(y)
     D = sparse.csc_matrix(np.diff(np.eye(L), 2))
     w = np.ones(L)
-    for i in range(niter):
+    for _ in range(niter):
         W = sparse.spdiags(w, 0, L, L)
         Z = W + lam * D.dot(D.transpose())
         z = sparse.linalg.spsolve(Z, w * y)
@@ -377,6 +431,72 @@ def baseline_als(y, lam, p, niter=10):
     return z
 
 
+def baseline_als(x, y, lam=None, p=None, niter=10, return_baseline=False,
+                 offset_correction=False):
+    """Baseline Correction with Asymmetric Least Squares Smoothing.
+
+    Parameters
+    ----------
+    x : array-like
+        the sample time/number/position
+    y : array-like
+        the data series corresponding to x
+    lam : float
+        the lambda parameter of the ALS method. This control how much the
+        baseline can adapt to local changes. A higher value corresponds to a
+        stiffer baseline
+    p : float
+        the asymmetry parameter of the ALS method. This controls the overall
+        slope tollerated for the baseline. A higher value correspond to a
+        higher possible slope
+
+    Other Parameters
+    ----------------
+    niter : int
+        The number of iterations to perform
+    return_baseline : bool
+        return the baseline?
+    offset_correction : bool
+        also correct for an offset to align with the running mean of the scan
+
+    Returns
+    -------
+    y_subtracted : array-like, same size as y
+        The initial time series, subtracted from the trend
+    baseline : array-like, same size as y
+        Fitted baseline. Only returned if return_baseline is True
+
+    Examples
+    --------
+    >>> x = np.arange(0, 10, 0.01)
+    >>> y = np.zeros_like(x) + 10
+    >>> ysub = baseline_als(x, y)
+    >>> np.all(ysub < 0.001)
+    True
+    """
+
+    if lam is None:
+        lam = 1e11
+    if p is None:
+        p = 0.001
+
+    z = _als(y, lam, p, niter=niter)
+
+    ysub = y - z
+    offset = 0
+    if offset_correction:
+        std = mad(ysub)
+
+        good = np.abs(ysub) < 10 * std
+
+        offset = offset_fit(x[good], ysub[good], 0)
+
+    if return_baseline:
+        return ysub - offset, z + offset
+    else:
+        return ysub - offset
+
+
 def excess_variance(lc, normalization='fvar'):
     """Calculate the excess variance.
 


Test patch:
diff --git a/stingray/tests/test_lightcurve.py b/stingray/tests/test_lightcurve.py
index 5f5cde81..db7ca0c7 100644
--- a/stingray/tests/test_lightcurve.py
+++ b/stingray/tests/test_lightcurve.py
@@ -734,6 +734,18 @@ class TestLightcurveRebin(object):
         baseline = lc.baseline(10000, 0.01)
         assert np.all(lc.counts - baseline < 1)
 
+    def test_lc_baseline_offset(self):
+        times = np.arange(0, 100, 0.01)
+        input_stdev = 0.1
+        counts = np.random.normal(100, input_stdev, len(times)) + \
+            0.001 * times
+        gti = [[-0.005, 50.005], [59.005, 100.005]]
+        good = create_gti_mask(times, gti)
+        counts[np.logical_not(good)] = 0
+        lc = Lightcurve(times, counts, gti=gti)
+        baseline = lc.baseline(10000, 0.01, offset_correction=True)
+        assert np.isclose(np.std(lc.counts - baseline), input_stdev, rtol=0.1)
+
     def test_change_mjdref(self):
         lc_new = self.lc.change_mjdref(57000)
         assert lc_new.mjdref == 57000