Repository: StingraySoftware/stingray
Pull request #213 (2017-07-18 18:55:32)

AutoCorrelation


Patch:
diff --git a/stingray/crosscorrelation.py b/stingray/crosscorrelation.py
index a4323407..f36e41da 100644
--- a/stingray/crosscorrelation.py
+++ b/stingray/crosscorrelation.py
@@ -6,6 +6,8 @@ from stingray import lightcurve
 from stingray.exceptions import StingrayError
 import stingray.utils as utils
 
+__all__ = ['CrossCorrelation', 'AutoCorrelation']
+
 
 class CrossCorrelation(object):
     def __init__(self, lc1=None, lc2=None, mode='same'):
@@ -42,8 +44,11 @@ class CrossCorrelation(object):
              There will be maximum correlation between lightcurves if one of the lightcurve is shifted by time_shift.
         n: int
              Number of points in self.corr(Length of cross-correlation data) 
+        auto: Boolean
+            A flag to indicate whether correlation is auto or cross.
         """
 
+        self.auto = False
         if isinstance(mode, str) is False:
             raise TypeError("mode must be a string")
 
@@ -207,3 +212,39 @@ class CrossCorrelation(object):
                 plt.savefig('corr.png')
             else:
                 plt.savefig(filename)
+        return plt
+
+
+class AutoCorrelation(CrossCorrelation):
+    def __init__(self, lc=None, mode='same'):
+        """
+        Make an auto-correlation from a light curve.
+        You can also make an empty Autocorrelation object to populate with your
+        own auto-correlation data.
+        Parameters
+        ----------
+        lc: lightcurve.Lightcurve object, optional, default None
+            The light curve data for correlation calculations.
+        mode: {'full', 'valid', 'same'}, optional, default 'same'
+            A string indicating the size of the correlation output.
+            See https://docs.scipy.org/doc/scipy-0.19.0/reference/generated/scipy.signal.correlate.html
+            for more details.
+
+        Attributes
+        ----------
+        lc1, lc2: lightcurve.Lightcurve 
+            The light curve data for correlation calculations.
+        corr: numpy.ndarray
+             An array of correlation data calculated from lightcurve data
+        time_lags: numpy.ndarray
+             An array of all possible time lags against which each point in corr is calculated 
+        dt: float
+             The time resolution of each lightcurve (used in time_lag calculations)
+        time_shift: float, zero
+             Max. Value of AutoCorrelation is always at zero lag.           
+        n: int
+             Number of points in self.corr(Length of auto-correlation data) 
+        """
+
+        CrossCorrelation.__init__(self, lc1=lc, lc2=lc, mode=mode)
+        self.auto = True


Test patch:
diff --git a/stingray/tests/test_crosscorrelation.py b/stingray/tests/test_crosscorrelation.py
index fb98d232..60a93e6f 100644
--- a/stingray/tests/test_crosscorrelation.py
+++ b/stingray/tests/test_crosscorrelation.py
@@ -5,7 +5,7 @@ import warnings
 import os
 
 from stingray import Lightcurve
-from stingray.crosscorrelation import CrossCorrelation
+from stingray.crosscorrelation import CrossCorrelation, AutoCorrelation
 from stingray.exceptions import StingrayError
 
 try:
@@ -35,6 +35,7 @@ class TestCrossCorrelation(object):
         assert cr.time_lags is None
         assert cr.dt is None
         assert cr.n is None
+        assert cr.auto is False
 
     def test_empty_cross_correlation_with_dt(self):
         cr = CrossCorrelation()
@@ -72,6 +73,7 @@ class TestCrossCorrelation(object):
         assert np.allclose(cr.time_lags, lags_result)
         assert np.isclose(cr.time_shift, 2.0)
         assert cr.mode == 'same'
+        assert cr.auto is False
 
     def test_cross_correlation_with_unequal_lc(self):
         result = np.array([-0.66666667, -0.33333333, -1., 0.66666667, 3.13333333])
@@ -85,6 +87,7 @@ class TestCrossCorrelation(object):
         assert np.allclose(cr.time_lags, lags_result)
         assert np.isclose(cr.time_shift, 2.0)
         assert cr.mode == 'same'
+        assert cr.auto is False
 
     def test_mode_with_bad_input(self):
         with pytest.raises(TypeError):
@@ -105,6 +108,7 @@ class TestCrossCorrelation(object):
         assert np.allclose(cr.corr, result)
         assert np.allclose(cr.time_lags, lags_result)
         assert np.isclose(cr.time_shift, 2.0)
+        assert cr.auto is False
 
     def test_timeshift_with_no_corr_but_lc_assigned(self):
         result = np.array([1.92, 2.16, 1.8, -14.44, 11.12])
@@ -121,6 +125,7 @@ class TestCrossCorrelation(object):
         assert np.allclose(cr.time_lags, lags_result)
         assert np.isclose(cr.time_shift, 2.0)
         assert cr.mode == 'same'
+        assert cr.auto is False
 
     def test_timeshift_with_corr_and_lc_assigned(self):
         result = np.array([1.92, 2.16, 1.8, -14.44, 11.12])
@@ -138,6 +143,7 @@ class TestCrossCorrelation(object):
         assert np.allclose(cr.time_lags, lags_result)
         assert np.isclose(cr.time_shift, 2.0)
         assert cr.mode == 'same'
+        assert cr.auto is False
 
     @pytest.mark.skipif(HAS_MPL, reason='Matplotlib is already installed if condition is met')
     def test_plot_matplotlib_not_installed(self):
@@ -193,3 +199,31 @@ class TestCrossCorrelation(object):
         cr.plot(save=True, filename='cr.png')
         assert os.path.isfile('cr.png')
         os.unlink('cr.png')
+
+    def test_auto_correlation(self):
+        result = np.array([1.68, -3.36, 5.2, -3.36, 1.68])
+        lags_result = np.array([-2, -1, 0, 1, 2])
+        ac = AutoCorrelation(self.lc1)
+        assert np.allclose(ac.lc1, self.lc1)
+        assert np.allclose(ac.lc2, self.lc1)
+        assert np.allclose(ac.corr, result)
+        assert np.isclose(ac.dt, self.lc1.dt)
+        assert ac.n == 5
+        assert np.allclose(ac.time_lags, lags_result)
+        assert np.isclose(ac.time_shift, 0.0)
+        assert ac.mode == 'same'
+        assert ac.auto is True
+
+    def test_auto_correlation_with_full_mode(self):
+        result = np.array([0.56, -1.48, 1.68, -3.36, 5.2, -3.36, 1.68, -1.48, 0.56])
+        lags_result = np.array([-4, -3, -2, -1, 0, 1, 2, 3, 4])
+        ac = AutoCorrelation(self.lc1, mode='full')
+        assert np.allclose(ac.lc1, self.lc1)
+        assert np.allclose(ac.lc2, self.lc1)
+        assert np.allclose(ac.corr, result)
+        assert np.isclose(ac.dt, self.lc1.dt)
+        assert ac.n == 9
+        assert np.allclose(ac.time_lags, lags_result)
+        assert np.isclose(ac.time_shift, 0.0)
+        assert ac.mode == 'full'
+        assert ac.auto is True