Repository: StingraySoftware/stingray
Pull request #211 (2017-07-17 22:27:22)

Functions to search for pulsations, using epoch folding and/or Z^2_n stat
The basic functionality for pulsation searches with the Epoch Folding algorithm and the Z^2_n statistics is implemented. 
It also implements a way to use simple thresholding of peaks in the periodogram.

To speed up the processing, I added a dependency on the `numba` library. I made it in such a way that if it is not installed the processing happens anyway, just much slower. On travis, all tests but one run without it.

The full fit of the peak with a `sinc` function à la Leahy+87, initially planned for this PR, will be implemented later.


Patch:
diff --git a/.travis.yml b/.travis.yml
index f8847b62..654ec5a3 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -45,12 +45,12 @@ matrix:
     fast_finish: true
     include:
 
-        #Try without importing h5py
+        #Try without importing h5py but importing numba
         - python: 2.7
-          env: CONDA_DEPENDENCIES='Cython numpy scipy nose matplotlib'
+          env: CONDA_DEPENDENCIES='Cython numpy scipy nose matplotlib numba'
 
         - python: 3.5
-          env: CONDA_DEPENDENCIES='Cython numpy scipy nose matplotlib'
+          env: CONDA_DEPENDENCIES='Cython numpy scipy nose matplotlib numba' SETUP_CMD='test --coverage'
 
         # Try without using corner
         - python: 3.5
diff --git a/stingray/pulse/pulsar.py b/stingray/pulse/pulsar.py
index 684f8346..a13868ef 100644
--- a/stingray/pulse/pulsar.py
+++ b/stingray/pulse/pulsar.py
@@ -125,7 +125,7 @@ def phase_exposure(start_time, stop_time, period, nbin=16, gtis=None):
             l0 = l[0]
             l1 = l[1]
             # Discards bins untouched by these limits
-            goodbins = np.logical_and(phs[:, 0] < l1, phs[:, 1] > l0)
+            goodbins = np.logical_and(phs[:, 0] <= l1, phs[:, 1] >= l0)
             idxs = np.arange(len(phs), dtype=int)[goodbins]
             for i in idxs:
                 start = max([phs[i, 0], l0])
@@ -165,6 +165,9 @@ def fold_events(times, *frequency_derivatives, **opts):
         Good time intervals
     ref_time : float, optional, default 0
         Reference time for the timing solution
+    expocorr : bool, default False
+        Correct each bin for exposure (use when the period of the pulsar is
+        comparable to that of GTIs)
 
     '''
     nbin = _default_value_if_no_key(opts, "nbin", 16)
@@ -179,7 +182,10 @@ def fold_events(times, *frequency_derivatives, **opts):
 
     gtis -= ref_time
     times -= ref_time
-    dt = times[1] - times[0]
+    # This dt has not the same meaning as in the Lightcurve case.
+    # it's just to define stop_time as a meaningful value after
+    # the last event.
+    dt = np.abs(times[1] - times[0])
     start_time = times[0]
     stop_time = times[-1] + dt
 
@@ -194,12 +200,14 @@ def fold_events(times, *frequency_derivatives, **opts):
                                      weights=weights)
 
     if expocorr:
-        expo_norm = phase_exposure(start_phase, stop_phase, 1, nbin)
+        expo_norm = phase_exposure(start_phase, stop_phase, 1, nbin,
+                                   gtis=gti_phases)
         simon("For exposure != 1, the uncertainty might be incorrect")
     else:
         expo_norm = 1
 
     # TODO: this is wrong. Need to extend this to non-1 weights
+
     raw_profile_err = np.sqrt(raw_profile)
 
     return bins[:-1] + np.diff(bins) / 2, raw_profile / expo_norm, \
diff --git a/stingray/pulse/search.py b/stingray/pulse/search.py
new file mode 100644
index 00000000..49d09c8e
--- /dev/null
+++ b/stingray/pulse/search.py
@@ -0,0 +1,189 @@
+from __future__ import division, print_function
+import numpy as np
+from .pulsar import stat, fold_events, z_n
+from ..utils import jit, HAS_NUMBA
+from ..utils import contiguous_regions
+
+
+@jit(nopython=True)
+def _pulse_phase_fast(time, f, buffer_array):
+    for i in range(len(time)):
+        buffer_array[i] = time[i] * f
+        buffer_array[i] -= np.floor(buffer_array[i])
+    return buffer_array
+
+
+def _folding_search(stat_func, times, frequencies, segment_size=5000):
+    stats = np.zeros_like(frequencies)
+    times = (times - times[0]).astype(np.float64)
+    length = times[-1]
+    if length < segment_size:
+        segment_size = length
+    start_times = np.arange(times[0], times[-1], segment_size)
+    count = 0
+    for s in start_times:
+        ts = times[(times >=s) & (times < s + segment_size)]
+        buffer = np.zeros_like(ts)
+        if len(ts) < 1 or ts[-1] - ts[0] < 0.2 * segment_size:
+            continue
+        for i, f in enumerate(frequencies):
+            phases = _pulse_phase_fast(ts, f, buffer)
+            stats[i] += stat_func(phases)
+        count += 1
+    return frequencies, stats / count
+
+
+@jit(nopython=True)
+def _bincount_fast(phase):
+    return np.bincount(phase)
+
+
+@jit(nopython=True)
+def _profile_fast(phase, nbin=128):
+    phase_bin = np.zeros(len(phase) + 2, dtype=np.int64)
+    # This is done to force bincount from 0 to nbin -1
+    phase_bin[-1] = nbin - 1
+    phase_bin[-2] = 0
+    for i in range(len(phase)):
+        phase_bin[i] = np.int64(np.floor(phase[i] * nbin))
+    bc = _bincount_fast(phase_bin)
+    bc[0] -= 1
+    bc[-1] -= 1
+    return bc
+
+
+def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
+                         expocorr=False):
+    """Performs epoch folding at trial frequencies in photon data.
+
+    If no exposure correction is needed and numba is installed, it uses a fast
+    algorithm to perform the folding. Otherwise, it runs a *much* slower
+    algorithm, which however yields a more precise result.
+    The search can be done in segments and the results averaged. Use
+    segment_size to control this
+
+    Parameters
+    ----------
+    times : array-like
+        the event arrival times
+    frequencies : array-like
+        the trial values for the frequencies
+
+    Other Parameters
+    ----------------
+    nbin : int
+        the number of bins of the folded profiles
+    segment_size : float
+        the length of the segments to be averaged in the periodogram
+    expocorr : bool
+        correct for the exposure (Use it if the period is comparable to the
+        length of the good time intervals.)
+    """
+    if expocorr or not HAS_NUMBA:
+        return \
+            _folding_search(lambda x: stat(fold_events(np.sort(x), 1,
+                                                       nbin=nbin,
+                                                       expocorr=expocorr)[1]),
+                               times, frequencies, segment_size=segment_size)
+
+    return _folding_search(lambda x: stat(_profile_fast(x, nbin=nbin)),
+                           times, frequencies, segment_size=segment_size)
+
+
+def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
+               expocorr=False):
+    """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
+    are folded first, and then the statistics is calculated using the value of
+    the profile as an additional normalization term.
+    The two methods are mostly equivalent. However, the number of bins has to
+    be chosen wisely: if the number of bins is too small, the search for high
+    harmonics is ineffective.
+    If no exposure correction is needed and numba is installed, it uses a fast
+    algorithm to perform the folding. Otherwise, it runs a *much* slower
+    algorithm, which however yields a more precise result.
+    The search can be done in segments and the results averaged. Use
+    segment_size to control this
+
+    Parameters
+    ----------
+    times : array-like
+        the event arrival times
+    frequencies : array-like
+        the trial values for the frequencies
+
+    Other Parameters
+    ----------------
+    nbin : int
+        the number of bins of the folded profiles
+    segment_size : float
+        the length of the segments to be averaged in the periodogram
+    expocorr : bool
+        correct for the exposure (Use it if the period is comparable to the
+        length of the good time intervals.)
+    """
+    phase = np.arange(0, 1, 1 / nbin)
+    if expocorr or not HAS_NUMBA:
+        return \
+            _folding_search(
+                lambda x: z_n(phase, n=nharm,
+                              norm=fold_events(np.sort(x), 1, nbin=nbin,
+                                               expocorr=expocorr)[1]),
+                               times, frequencies, segment_size=segment_size)
+
+    return _folding_search(lambda x: z_n(phase, n=nharm,
+                                         norm=_profile_fast(x, nbin=nbin)),
+                           times, frequencies, segment_size=segment_size)
+
+
+def search_best_peaks(x, stat, threshold):
+    """Search peaks above threshold in an epoch folding periodogram.
+
+    If more values of stat are above threshold and are contiguous, only the
+    largest one is returned (see Examples).
+
+    Parameters
+    ----------
+    x : array-like
+        The x axis of the periodogram (frequencies, periods, ...)
+    stat : array-like
+        The y axis. It must have the same shape as x
+    threshold : float
+        The threshold value over which we look for peaks in the stat array
+
+    Examples
+    --------
+    >>> # Test multiple peaks
+    >>> x = np.arange(10)
+    >>> stat = [0, 0, 0.5, 0, 0, 1, 1, 2, 1, 0]
+    >>> best = search_best_peaks(x, stat, 0.5)
+    >>> len(best)
+    2
+    >>> best[0]
+    2.0
+    >>> best[1]
+    7.0
+    >>> # Test no peak above threshold
+    >>> x = np.arange(10)
+    >>> stat = [0, 0, 0.4, 0, 0, 0, 0, 0, 0, 0]
+    >>> search_best_peaks(x, stat, 0.5)
+    []
+
+    Returns
+    -------
+    best_x : array-like
+        the array containing the peaks above threshold. If no peaks are above
+        threshold, an empty list is returned.
+    """
+    stat = np.asarray(stat)
+    x = np.asarray(x)
+    peaks = stat >= threshold
+    regions = contiguous_regions(peaks)
+    if len(regions) == 0:
+        return []
+    best_x = np.zeros(len(regions))
+    for i, r in enumerate(regions):
+        best_x[i] = x[r[0]:r[1]][np.argmax(stat[r[0]:r[1]])]
+
+    return best_x
diff --git a/stingray/utils.py b/stingray/utils.py
index 81431621..2dfdf50a 100644
--- a/stingray/utils.py
+++ b/stingray/utils.py
@@ -10,11 +10,23 @@ import numpy as np
 # If numba is installed, import jit. Otherwise, define an empty decorator with
 # the same name.
 
+HAS_NUMBA = False
 try:
     from numba import jit
-except:
-    def jit(fun):
-        return fun
+    HAS_NUMBA = True
+except ImportError:
+    warnings.warn("Numba not installed. Faking it")
+
+    class jit(object):
+
+        def __init__(self, *args, **kwargs):
+            pass
+
+        def __call__(self, func):
+            def wrapped_f(*args, **kwargs):
+                return func(*args, **kwargs)
+
+            return wrapped_f
 
 
 def simon(message, **kwargs):


Test patch:
diff --git a/stingray/pulse/tests/test_pulse.py b/stingray/pulse/tests/test_pulse.py
index 34a29e13..c0824461 100644
--- a/stingray/pulse/tests/test_pulse.py
+++ b/stingray/pulse/tests/test_pulse.py
@@ -119,17 +119,35 @@ class TestAll(object):
 
     def test_pulse_profile2(self):
         nbin = 16
-        times = np.arange(0, 1, 1/nbin)
+        dt = 1/nbin
+        times = np.arange(0, 2, dt)
+        gtis = np.array([[-0.5*dt, 2 + 0.5*dt]])
 
         period = 1
-        ph, p, pe = fold_events(times, 1, nbin=nbin, expocorr=True)
+        ph, p, pe = fold_events(times, 1, nbin=nbin, expocorr=True, gtis=gtis)
 
         np.testing.assert_array_almost_equal(ph, np.arange(nbin)/nbin +
                                              0.5/nbin)
-        np.testing.assert_array_almost_equal(p, np.ones(nbin))
-        np.testing.assert_array_almost_equal(pe, np.ones(nbin))
+        np.testing.assert_array_almost_equal(p, 2 * np.ones(nbin))
+        np.testing.assert_array_almost_equal(pe, 2**0.5 * np.ones(nbin))
 
-    def test_zn(self):
+    def test_pulse_profile3(self):
+        nbin = 16
+        dt = 1/nbin
+        times = np.arange(0, 2 - dt, dt)
+        gtis = np.array([[-0.5*dt, 2 - dt]])
+
+        ph, p, pe = fold_events(times, 1, nbin=nbin, expocorr=True,
+                                gtis=gtis)
+
+        np.testing.assert_array_almost_equal(ph, np.arange(nbin)/nbin +
+                                             0.5/nbin)
+        np.testing.assert_array_almost_equal(p, 2 * np.ones(nbin))
+        expected_err = 2**0.5 * np.ones(nbin)
+        expected_err[-1] = 2  # Because of the change of exposure
+        np.testing.assert_array_almost_equal(pe, expected_err)
+
+    def test_zn_2(self):
         np.testing.assert_almost_equal(z_n(np.arange(1), n=1, norm=1), 2)
         np.testing.assert_almost_equal(z_n(np.arange(1), n=2, norm=1), 4)
         np.testing.assert_almost_equal(z_n(np.arange(2), n=2, norm=1), 8)
diff --git a/stingray/pulse/tests/test_search.py b/stingray/pulse/tests/test_search.py
new file mode 100644
index 00000000..3266a19f
--- /dev/null
+++ b/stingray/pulse/tests/test_search.py
@@ -0,0 +1,72 @@
+from __future__ import division, print_function
+from stingray.pulse.search import epoch_folding_search, z_n_search, _profile_fast
+import numpy as np
+from stingray import Lightcurve
+from stingray.events import EventList
+
+np.random.seed(20150907)
+
+class TestAll(object):
+    """Unit tests for the stingray.pulse.search module."""
+    @classmethod
+    def setup_class(cls):
+        cls.pulse_frequency = 1/0.101
+        cls.tstart = 0
+        cls.tend = 25.25
+        cls.tseg = cls.tend - cls.tstart
+        cls.dt = 0.0202
+        cls.times = np.arange(cls.tstart, cls.tend, cls.dt) + cls.dt / 2
+        cls.counts = \
+            100 + 20 * np.cos(2 * np.pi * cls.times * cls.pulse_frequency)
+        lc = Lightcurve(cls.times, cls.counts, gti=[[cls.tstart, cls.tend]])
+        events = EventList()
+        events.simulate_times(lc)
+        cls.event_times = events.time
+
+    def test_prepare(self):
+        pass
+
+    def test_profile_fast(self):
+        test_phase = np.arange(0, 1, 1/16)
+        prof = _profile_fast(test_phase, nbin=16)
+        assert np.all(prof == np.ones(16))
+
+    def test_epoch_folding_search(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.85, 9.95, 0.1/self.tseg)
+        freq, stat = epoch_folding_search(self.event_times, frequencies,
+                                          nbin=16)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq[np.argmax(stat)]
+        assert maxstatbin == frequencies[minbin]
+
+    def test_epoch_folding_search_expocorr(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.89, 9.91, 0.1/self.tseg)
+        freq, stat = epoch_folding_search(self.event_times, frequencies,
+                                          nbin=16, expocorr=True)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq[np.argmax(stat)]
+        assert maxstatbin == frequencies[minbin]
+
+    def test_z_n_search(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.85, 9.95, 0.3/self.tseg)
+        freq, stat = z_n_search(self.event_times, frequencies, nbin=16,
+                                nharm=1)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq[np.argmax(stat)]
+        assert maxstatbin == frequencies[minbin]
+
+    def test_z_n_search_expocorr(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.89, 9.91, 0.1/self.tseg)
+        freq, stat = z_n_search(self.event_times, frequencies, nbin=16,
+                                nharm=1, expocorr=True)
+
+        minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
+        maxstatbin = freq[np.argmax(stat)]
+        assert maxstatbin == frequencies[minbin]