Repository: StingraySoftware/stingray
Pull request #263 (2017-09-26 14:31:50)

EF and Z search allow trial fdots
Another branch got mixed up with this, but the changes should be straightforward to get.
One is applying GTIs before calculating cross spectra, the other is adding trial dot to epoch folding-like searches


Patch:
diff --git a/stingray/crossspectrum.py b/stingray/crossspectrum.py
index 27db9673..188a0eab 100644
--- a/stingray/crossspectrum.py
+++ b/stingray/crossspectrum.py
@@ -603,6 +603,9 @@ class AveragedCrossspectrum(Crossspectrum):
 
         if self.gti is None:
             self.gti = cross_two_gtis(lc1.gti, lc2.gti)
+            lc1.gti = lc2.gti = self.gti
+            lc1._apply_gtis()
+            lc2._apply_gtis()
 
         check_gtis(self.gti)
 
@@ -621,15 +624,17 @@ class AveragedCrossspectrum(Crossspectrum):
             time_2 = lc2.time[start_ind:end_ind]
             counts_2 = lc2.counts[start_ind:end_ind]
             counts_2_err = lc2.counts_err[start_ind:end_ind]
+            gti1 = np.array([[time_1[0] - lc1.dt / 2,
+                             time_1[-1] + lc1.dt / 2]])
+            gti2 = np.array([[time_2[0] - lc2.dt / 2,
+                             time_2[-1] + lc2.dt / 2]])
             lc1_seg = Lightcurve(time_1, counts_1, err=counts_1_err,
                                  err_dist=lc1.err_dist,
-                                 gti=[[time_1[0] - lc1.dt / 2,
-                                       time_1[-1] + lc1.dt / 2]],
+                                 gti=gti1,
                                  dt=lc1.dt)
             lc2_seg = Lightcurve(time_2, counts_2, err=counts_2_err,
                                  err_dist=lc2.err_dist,
-                                 gti=[[time_2[0] - lc2.dt / 2,
-                                       time_2[-1] + lc2.dt / 2]],
+                                 gti=gti2,
                                  dt=lc2.dt)
             cs_seg = Crossspectrum(lc1_seg, lc2_seg, norm=self.norm)
             cs_all.append(cs_seg)
diff --git a/stingray/pulse/search.py b/stingray/pulse/search.py
index 4da3e99e..65fe4ebb 100644
--- a/stingray/pulse/search.py
+++ b/stingray/pulse/search.py
@@ -12,16 +12,18 @@ __all__ = ['epoch_folding_search', 'z_n_search', 'search_best_peaks',
 
 
 @jit(nopython=True)
-def _pulse_phase_fast(time, f, buffer_array):
+def _pulse_phase_fast(time, f, fdot, buffer_array):
     for i in range(len(time)):
-        buffer_array[i] = time[i] * f
+        buffer_array[i] = time[i] * f + 0.5 * time[i]**2 * fdot
         buffer_array[i] -= np.floor(buffer_array[i])
     return buffer_array
 
 
 def _folding_search(stat_func, times, frequencies, segment_size=5000,
-                    use_times=False, **kwargs):
-    stats = np.zeros_like(frequencies)
+                    use_times=False, fdots=0, **kwargs):
+
+    fgrid, fdgrid = np.meshgrid(frequencies, fdots)
+    stats = np.zeros_like(fgrid)
     times = (times - times[0]).astype(np.float64)
     length = times[-1]
     if length < segment_size:
@@ -33,23 +35,30 @@ def _folding_search(stat_func, times, frequencies, segment_size=5000,
         ts = times[good]
         if len(ts) < 1 or ts[-1] - ts[0] < 0.2 * segment_size:
             continue
-        for i, f in enumerate(frequencies):
-            if use_times:
-                kwargs_copy = {}
-                for key in kwargs.keys():
-                    if isinstance(kwargs[key], collections.Iterable) and \
-                        len(kwargs[key]) == len(times):
-
-                        kwargs_copy[key] = kwargs[key][good]
-                    else:
-                        kwargs_copy[key] = kwargs[key]
-                stats[i] += stat_func(ts, f, **kwargs_copy)
-            else:
-                buffer = np.zeros_like(ts)
-                phases = _pulse_phase_fast(ts, f, buffer)
-                stats[i] += stat_func(phases)
+        buffer = np.zeros_like(ts)
+        for i in range(stats.shape[0]):
+            for j in range(stats.shape[1]):
+                f = fgrid[i, j]
+                fd = fdgrid[i, j]
+                if use_times:
+                    kwargs_copy = {}
+                    for key in kwargs.keys():
+                        if isinstance(kwargs[key], collections.Iterable) and \
+                            len(kwargs[key]) == len(times):
+
+                            kwargs_copy[key] = kwargs[key][good]
+                        else:
+                            kwargs_copy[key] = kwargs[key]
+                    stats[i, j] += stat_func(ts, f, fd, **kwargs_copy)
+                else:
+                    phases = _pulse_phase_fast(ts, f, fd, buffer)
+                    stats[i, j] += stat_func(phases)
         count += 1
-    return frequencies, stats / count
+
+    if fgrid.shape[0] == 1:
+        return fgrid.flatten(), stats.flatten() / count
+    else:
+        return fgrid, fdgrid, stats / count
 
 
 @jit(nopython=True)
@@ -72,7 +81,7 @@ def _profile_fast(phase, nbin=128):
 
 
 def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
-                         expocorr=False, gti=None, weights=1):
+                         expocorr=False, gti=None, weights=1, fdots=0):
     """Performs epoch folding at trial frequencies in photon data.
 
     If no exposure correction is needed and numba is installed, it uses a fast
@@ -94,6 +103,8 @@ def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
         the number of bins of the folded profiles
     segment_size : float
         the length of the segments to be averaged in the periodogram
+    fdots : array-like
+        trial values of the first frequency derivative (optional)
     expocorr : bool
         correct for the exposure (Use it if the period is comparable to the
         length of the good time intervals). If True, GTIs have to be specified
@@ -109,20 +120,22 @@ def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
             raise ValueError('To calculate exposure correction, you need to'
                              ' specify the GTIs')
 
-        def fun(t, f, **kwargs):
-            return stat(fold_events(t, f, **kwargs)[1])
+        def stat_fun(t, f, fd=0, **kwargs):
+            return stat(fold_events(t, f, fd, **kwargs)[1])
 
         return \
-            _folding_search(fun, times, frequencies, segment_size=segment_size,
+            _folding_search(stat_fun, times, frequencies,
+                            segment_size=segment_size,
                             use_times=True, expocorr=expocorr, weights=weights,
-                            gti=gti, nbin=nbin)
+                            gti=gti, nbin=nbin, fdots=fdots)
 
     return _folding_search(lambda x: stat(_profile_fast(x, nbin=nbin)),
-                           times, frequencies, segment_size=segment_size)
+                           times, frequencies, segment_size=segment_size,
+                           fdots=fdots)
 
 
 def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
-               expocorr=False, weights=1, gti=None):
+               expocorr=False, weights=1, gti=None, fdots=0):
     """Calculates the Z^2_n statistics at trial frequencies in photon data.
 
     The "real" Z^2_n statistics is very slow. Therefore, in this function data
@@ -150,6 +163,8 @@ def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
         the number of bins of the folded profiles
     segment_size : float
         the length of the segments to be averaged in the periodogram
+    fdots : array-like
+        trial values of the first frequency derivative (optional)
     expocorr : bool
         correct for the exposure (Use it if the period is comparable to the
         length of the good time intervals.)
@@ -164,17 +179,19 @@ def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
         if expocorr and gti is None:
             raise ValueError('To calculate exposure correction, you need to'
                              ' specify the GTIs')
-        def fun(t, f, **kwargs):
+        def stat_fun(t, f, fd=0, **kwargs):
             return z_n(phase, n=nharm,
-                       norm=fold_events(t, f, nbin=nbin, **kwargs)[1])
+                       norm=fold_events(t, f, fd, nbin=nbin, **kwargs)[1])
         return \
-            _folding_search(fun, times, frequencies, segment_size=segment_size,
+            _folding_search(stat_fun, times, frequencies,
+                            segment_size=segment_size,
                             use_times=True, expocorr=expocorr, weights=weights,
-                            gti=gti)
+                            gti=gti, fdots=fdots)
 
     return _folding_search(lambda x: z_n(phase, n=nharm,
                                          norm=_profile_fast(x, nbin=nbin)),
-                           times, frequencies, segment_size=segment_size)
+                           times, frequencies, segment_size=segment_size,
+                           fdots=fdots)
 
 
 def search_best_peaks(x, stat, threshold):


Test patch:
diff --git a/stingray/pulse/tests/test_search.py b/stingray/pulse/tests/test_search.py
index 2eeb766f..8ecb5a0b 100644
--- a/stingray/pulse/tests/test_search.py
+++ b/stingray/pulse/tests/test_search.py
@@ -136,6 +136,19 @@ class TestAll(object):
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
+    def test_epoch_folding_search_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = epoch_folding_search(self.event_times, frequencies,
+                                                nbin=43, fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)
+
     def test_epoch_folding_search_expocorr_fails(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
@@ -144,7 +157,7 @@ class TestAll(object):
                                               nbin=23, expocorr=True)
         assert 'To calculate exposure correction' in str(excinfo)
 
-    def test_epoch_folding_search_expocorr_(self):
+    def test_epoch_folding_search_expocorr(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
         freq, stat = epoch_folding_search(self.event_times, frequencies,
@@ -154,6 +167,21 @@ class TestAll(object):
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
+    def test_epoch_folding_search_expocorr_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = epoch_folding_search(self.event_times, frequencies,
+                                                nbin=42, expocorr=True, gti=self.gti,
+                                                fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)
+
     def test_epoch_folding_search_weights(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
@@ -164,6 +192,20 @@ class TestAll(object):
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
+    def test_epoch_folding_search_weights_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = epoch_folding_search(self.times, frequencies,
+                                                nbin=16, weights=self.counts,
+                                                fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)
+
     def test_z_n_search(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
@@ -174,6 +216,19 @@ class TestAll(object):
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
+    def test_z_n_search_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = z_n_search(self.event_times, frequencies, nbin=25,
+                                      nharm=2, fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)
+
     def test_z_n_search_expocorr(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
@@ -184,6 +239,20 @@ class TestAll(object):
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
+    def test_z_n_search_expocorr_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = z_n_search(self.event_times, frequencies, nbin=64,
+                                      nharm=2, expocorr=True, gti=self.gti,
+                                      fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)
+
     def test_z_n_search_expocorr_fails(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
@@ -201,3 +270,17 @@ class TestAll(object):
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
         assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+
+    def test_z_n_search_weights_fdot(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        fdots = [-0.1, 0, 0.1]
+        freq, fdot, stat = z_n_search(self.times, frequencies, nbin=44,
+                                      nharm=1, weights=self.counts,
+                                      fdots=fdots)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq.flatten()[np.argmax(stat)]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
+        maxfdot = fdot.flatten()[np.argmax(stat)]
+        assert np.allclose(maxfdot, 0.0, atol=0.1/self.tseg)