Repository: StingraySoftware/stingray
Pull request #185 (2017-04-18 09:02:45)

Baseline of light curve improved
Adds missing tests for baseline function, accounts for GTIs.

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/stingraysoftware/stingray/185)
<!-- Reviewable:end -->



Patch:
diff --git a/stingray/lightcurve.py b/stingray/lightcurve.py
index e4a4f3f2..b2354751 100644
--- a/stingray/lightcurve.py
+++ b/stingray/lightcurve.py
@@ -9,9 +9,9 @@ import numpy as np
 import stingray.io as io
 import stingray.utils as utils
 from stingray.exceptions import StingrayError
-from stingray.utils import simon, assign_value_if_none
+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
+from stingray.gti import check_gtis, create_gti_mask
 from astropy.stats import poisson_conf_interval
 
 __all__ = ["Lightcurve"]
@@ -393,6 +393,25 @@ class Lightcurve(object):
             raise IndexError("The index must be either an integer or a slice "
                              "object !")
 
+    def baseline(self, lam, p, niter=10):
+        """Calculate the baseline of the light curve, accounting for GTIs.
+
+        Parameters
+        ----------
+        lam : float
+            "smoothness" parameter. Larger values make the baseline stiffer
+            Typically 1e2 < lam < 1e9
+        p : float
+            "asymmetry" parameter. Smaller values make the baseline more 
+            "horizontal". Typically 0.001 < p < 0.1, but not necessary.
+        """
+        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)
+
+        return baseline
+
     @staticmethod
     def make_lightcurve(toa, dt, tseg=None, tstart=None, gti=None):
 


Test patch:
diff --git a/stingray/tests/test_lightcurve.py b/stingray/tests/test_lightcurve.py
index 77093807..e7611480 100644
--- a/stingray/tests/test_lightcurve.py
+++ b/stingray/tests/test_lightcurve.py
@@ -7,6 +7,7 @@ import matplotlib.pyplot as plt
 
 from stingray import Lightcurve
 from stingray.exceptions import StingrayError
+from stingray.gti import create_gti_mask
 
 np.random.seed(20150907)
 
@@ -566,9 +567,6 @@ class TestLightcurveRebin(object):
         assert np.allclose(lc_binned.counts, counts_test)
 
     def rebin_several(self, dt):
-        """
-        TODO: Not sure how to write tests for the rebin method!
-        """
         lc_binned = self.lc.rebin(dt)
         assert len(lc_binned.time) == len(lc_binned.counts)
 
@@ -576,3 +574,14 @@ class TestLightcurveRebin(object):
         dt_all = [2, 3, np.pi, 5]
         for dt in dt_all:
             self.rebin_several(dt)
+
+    def test_lc_baseline(self):
+        times = np.arange(0, 100, 0.01)
+        counts = np.random.normal(100, 0.1, 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)
+        assert np.all(lc.counts - baseline < 1)