Repository: StingraySoftware/stingray
Pull request #256 (2017-09-07 22:22:45)

Allow folding searches and phaseogram from light curves
Basically, I added a `weights` option to a few functions, that allow to associate each time with a weight.

For example, if before we could only do something like:
```
event = EventList(...)
freq, stat = epoch_folding_search(event.time, frequencies)
```

now we can also do 
```
lc = Lightcurve(...)
freq, stat = epoch_folding_search(lc.time, frequencies, weights=lc.counts)
```

Same goes for `phaseogram`!


Patch:
diff --git a/stingray/modeling/parameterestimation.py b/stingray/modeling/parameterestimation.py
index 411fc2a9..394f8977 100644
--- a/stingray/modeling/parameterestimation.py
+++ b/stingray/modeling/parameterestimation.py
@@ -614,8 +614,7 @@ class ParameterEstimation(object):
                       nsim=1000, niter=200, nwalkers=500, burnin=200,
                       namestr="test", seed=None):
 
-        """
-        Calibrate the outcome of a Likelihood Ratio Test via MCMC.
+        """Calibrate the outcome of a Likelihood Ratio Test via MCMC.
 
         In order to compare models via likelihood ratio test, one generally
         aims to compute a p-value for the null hypothesis (generally the
@@ -623,15 +622,17 @@ class ParameterEstimation(object):
         distribution used to compute that p-value analytically given the
         observed likelihood ratio (a chi-square distribution) is not
         applicable:
+
         * the models are not nested (i.e. Model 1 is not a special, simpler
-        case of Model 2),
+          case of Model 2),
+
         * the parameter values fixed in Model 2 to retrieve Model 1 are at the
-        edges of parameter space (e.g. if one must set, say, an amplitude to
-        zero in order to remove a component in the more complex model, and
-        negative amplitudes are excluded a priori)
+          edges of parameter space (e.g. if one must set, say, an amplitude to
+          zero in order to remove a component in the more complex model, and
+          negative amplitudes are excluded a priori)
 
         In these cases, the observed likelihood ratio must be calibrated via
-         simulations of the simpler model (Model 1), using MCMC to take into
+        simulations of the simpler model (Model 1), using MCMC to take into
         account the uncertainty in the parameters. This function does
         exactly that: it computes the likelihood ratio for the observed data,
         and produces simulations to calibrate the likelihood ratio and
@@ -666,12 +667,10 @@ class ParameterEstimation(object):
             If True, set the internal state to do the optimization with the
             log-likelihood rather than the log-posterior.
 
-
         Returns
         -------
         pvalue : float [0,1]
             p-value 'n stuff
-
         """
 
         # compute the observed likelihood ratio
@@ -1020,7 +1019,7 @@ class PSDParEst(ParameterEstimation):
         s_all : numpy.ndarray of shape (nsamples, lpost1.npar)
             An array with MCMC samples derived from the null hypothesis model in
             `lpost1`. Its second dimension must match the number of free
-             parameters in `lpost1.model`.
+            parameters in `lpost1.model`.
 
         lpost1 : LogLikelihood or Posterior subclass object
             Object containing the null hypothesis model
@@ -1049,7 +1048,6 @@ class PSDParEst(ParameterEstimation):
         lrt_sim : numpy.ndarray
             An array with the simulated likelihood ratios for the simulated
             data
-
         """
 
         assert lpost1.__class__ == lpost2.__class__, "Both LogLikelihood or " \
diff --git a/stingray/modeling/posterior.py b/stingray/modeling/posterior.py
index bec3d66d..da937d2c 100644
--- a/stingray/modeling/posterior.py
+++ b/stingray/modeling/posterior.py
@@ -58,7 +58,7 @@ def set_logprior(lpost, priors):
         A dictionary containing the prior definitions. Keys are parameter
         names as defined by the model used in `lpost`. Items are functions
         that take a parameter as input and return the log-prior probability
-         of that parameter.
+        of that parameter.
 
     Returns
     -------
diff --git a/stingray/pulse/pulsar.py b/stingray/pulse/pulsar.py
index 1d05151a..8863a9bc 100644
--- a/stingray/pulse/pulsar.py
+++ b/stingray/pulse/pulsar.py
@@ -64,13 +64,12 @@ def pulse_phase(times, *frequency_derivatives, **opts):
 
 
 def phase_exposure(start_time, stop_time, period, nbin=16, gtis=None):
-    '''
-    Calculate the exposure on each phase of a pulse profile.
+    """Calculate the exposure on each phase of a pulse profile.
 
     Parameters
     ----------
     start_time, stop_time : float
-        Starting and stopping time (or phase if ``period``==1)
+        Starting and stopping time (or phase if ``period==1``)
     period : float
         The pulse period (if 1, equivalent to phases)
 
@@ -84,9 +83,9 @@ def phase_exposure(start_time, stop_time, period, nbin=16, gtis=None):
     ----------------
     nbin : int, optional, default 16
         The number of bins in the profile
-    gtis : [[gti0_0, gti0_1], [gti1_0, gti1_1], ...], optional, default None
+    gtis : [[gti00, gti01], [gti10, gti11], ...], optional, default None
         Good Time Intervals
-    '''
+    """
     if gtis is None:
         gtis = np.array([[start_time, stop_time]])
 
diff --git a/stingray/pulse/search.py b/stingray/pulse/search.py
index 2661fdb9..f93774bc 100644
--- a/stingray/pulse/search.py
+++ b/stingray/pulse/search.py
@@ -1,5 +1,6 @@
 from __future__ import division, print_function
 import numpy as np
+import collections
 from .pulsar import stat, fold_events, z_n, pulse_phase
 from ..utils import jit, HAS_NUMBA
 from ..utils import contiguous_regions
@@ -18,7 +19,8 @@ def _pulse_phase_fast(time, f, buffer_array):
     return buffer_array
 
 
-def _folding_search(stat_func, times, frequencies, segment_size=5000):
+def _folding_search(stat_func, times, frequencies, segment_size=5000,
+                    use_times=False, **kwargs):
     stats = np.zeros_like(frequencies)
     times = (times - times[0]).astype(np.float64)
     length = times[-1]
@@ -27,13 +29,25 @@ def _folding_search(stat_func, times, frequencies, segment_size=5000):
     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)
+        good = (times >=s) & (times < s + segment_size)
+        ts = times[good]
         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)
+            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)
         count += 1
     return frequencies, stats / count
 
@@ -58,7 +72,7 @@ def _profile_fast(phase, nbin=128):
 
 
 def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
-                         expocorr=False):
+                         expocorr=False, gti=None, weights=1):
     """Performs epoch folding at trial frequencies in photon data.
 
     If no exposure correction is needed and numba is installed, it uses a fast
@@ -82,21 +96,33 @@ def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
         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.)
+        length of the good time intervals). If True, GTIs have to be specified
+        via the ``gti`` keyword
+    gti : [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
+        Good time intervals
+    weights : array-like
+        weight for each time. This might be, for example, the number of counts
+        if the times array contains the time bins of a light curve
     """
-    if expocorr or not HAS_NUMBA:
+    if expocorr or not HAS_NUMBA or isinstance(weights, collections.Iterable):
+        if expocorr and gti is None:
+            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])
+
         return \
-            _folding_search(lambda x: stat(fold_events(np.sort(x), 1,
-                                                       nbin=nbin,
-                                                       expocorr=expocorr)[1]),
-                               times, frequencies, segment_size=segment_size)
+            _folding_search(fun, times, frequencies, segment_size=segment_size,
+                            use_times=True, expocorr=expocorr, weights=weights,
+                            gti=gti)
 
     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):
+               expocorr=False, weights=1, gti=None):
     """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
@@ -127,15 +153,24 @@ def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
     expocorr : bool
         correct for the exposure (Use it if the period is comparable to the
         length of the good time intervals.)
+    gti : [[gti0_0, gti0_1], [gti1_0, gti1_1], ...]
+        Good time intervals
+    weights : array-like
+        weight for each time. This might be, for example, the number of counts
+        if the times array contains the time bins of a light curve
     """
     phase = np.arange(0, 1, 1 / nbin)
-    if expocorr or not HAS_NUMBA:
+    if expocorr or not HAS_NUMBA or isinstance(weights, collections.Iterable):
+        if expocorr and gti is None:
+            raise ValueError('To calculate exposure correction, you need to'
+                             ' specify the GTIs')
+        def fun(t, f, **kwargs):
+            return z_n(phase, n=nharm,
+                       norm=fold_events(t, f, **kwargs)[1])
         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)
+            _folding_search(fun, times, frequencies, segment_size=segment_size,
+                            use_times=True, expocorr=expocorr, weights=weights,
+                            gti=gti)
 
     return _folding_search(lambda x: z_n(phase, n=nharm,
                                          norm=_profile_fast(x, nbin=nbin)),
@@ -299,7 +334,8 @@ def plot_phaseogram(phaseogram, phase_bins, time_bins, unit_str='s', ax=None,
 
 
 def phaseogram(times, f, nph=128, nt=32, ph0=0, mjdref=None, fdot=0, fddot=0,
-               pepoch=None, plot=False, phaseogram_ax=None, **plot_kwargs):
+               pepoch=None, plot=False, phaseogram_ax=None,
+               weights=None, **plot_kwargs):
     """
     Calculate and plot the phaseogram of a pulsar observation.
 
@@ -335,6 +371,8 @@ def phaseogram(times, f, nph=128, nt=32, ph0=0, mjdref=None, fdot=0, fddot=0,
         If the input pulse solution is referred to a given time, give it here.
         It has no effect (just a phase shift of the pulse) if `fdot` is zero.
         if `mjdref` is specified, pepoch MUST be in MJD
+    weights : array
+        Weight for each time
     plot : bool
         Return the axes in the additional_info, and don't close the plot, so
         that the user can add information to it.
@@ -373,6 +411,12 @@ def phaseogram(times, f, nph=128, nt=32, ph0=0, mjdref=None, fdot=0, fddot=0,
     allphases = np.concatenate([phases, phases + 1]).astype('float64')
     allts = \
         np.concatenate([times, times]).astype('float64')
+    if weights is not None and isinstance(weights, collections.Iterable):
+        if len(weights) != len(times):
+            raise ValueError('The length of weights must match the length of '
+                             'times')
+        weights = \
+            np.concatenate([weights, weights]).astype('float64')
 
     if use_mjdref:
         allts = allts / 86400 + mjdref
@@ -380,7 +424,8 @@ def phaseogram(times, f, nph=128, nt=32, ph0=0, mjdref=None, fdot=0, fddot=0,
     phas, binx, biny = np.histogram2d(allphases, allts,
                bins=(np.linspace(0, 2, nph * 2 + 1),
                      np.linspace(np.min(allts),
-                                 np.max(allts), nt + 1)))
+                                 np.max(allts), nt + 1)),
+               weights=weights)
 
     if plot:
         phaseogram_ax = plot_phaseogram(phas, binx, biny, ax=phaseogram_ax,


Test patch:
diff --git a/stingray/pulse/tests/test_search.py b/stingray/pulse/tests/test_search.py
index f12dde83..3771af59 100644
--- a/stingray/pulse/tests/test_search.py
+++ b/stingray/pulse/tests/test_search.py
@@ -6,9 +6,11 @@ from stingray.pulse.pulsar import fold_events
 import numpy as np
 from stingray import Lightcurve
 from stingray.events import EventList
+import pytest
 
 np.random.seed(20150907)
 
+
 class TestAll(object):
     """Unit tests for the stingray.pulse.search module."""
     @classmethod
@@ -21,7 +23,8 @@ class TestAll(object):
         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]])
+        cls.gti = [[cls.tstart, cls.tend]]
+        lc = Lightcurve(cls.times, cls.counts, gti=cls.gti)
         events = EventList()
         events.simulate_times(lc)
         cls.event_times = events.time
@@ -36,6 +39,26 @@ class TestAll(object):
         assert np.any(times > 25)
         assert np.all((phases >= 0)&(phases <= 2))
 
+    def test_phaseogram_bad_weights(self):
+        with pytest.raises(ValueError) as excinfo:
+            phaseogr, phases, times, additional_info = \
+                phaseogram(self.event_times, self.pulse_frequency,
+                           weights=[0, 2])
+        assert 'must match' in str(excinfo)
+
+    def test_phaseogram_weights(self):
+        phaseogr, phases, times, additional_info = \
+            phaseogram(self.times, self.pulse_frequency, weights=self.counts,
+                       nph=16)
+        assert np.all(times < 25.6)
+        assert np.any(times > 25)
+        assert np.all((phases >= 0)&(phases <= 2))
+        import matplotlib.pyplot as plt
+        fig = plt.figure('Phaseogram direct weights')
+        plot_phaseogram(phaseogr, phases, times)
+        plt.savefig('phaseogram_weights.png')
+        plt.close(fig)
+
     def test_phaseogram_mjdref(self):
         phaseogr, phases, times, additional_info = \
             phaseogram(self.event_times, self.pulse_frequency,
@@ -72,7 +95,7 @@ class TestAll(object):
     def test_plot_profile(self):
         import matplotlib.pyplot as plt
         phase, prof, _ = fold_events(self.event_times,
-                                            self.pulse_frequency)
+                                     self.pulse_frequency)
         ax = plot_profile(phase, prof)
         plt.savefig('profile_direct.png')
         plt.close(plt.gcf())
@@ -105,40 +128,76 @@ class TestAll(object):
 
     def test_epoch_folding_search(self):
         """Test pulse phase calculation, frequency only."""
-        frequencies = np.arange(9.85, 9.95, 0.1/self.tseg)
+        frequencies = np.arange(9.8, 9.99, 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]
+        assert np.allclose(maxstatbin, frequencies[minbin], 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)
+        with pytest.raises(ValueError) as excinfo:
+            freq, stat = epoch_folding_search(self.event_times, frequencies,
+                                              nbin=16, 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.89, 9.91, 0.1/self.tseg)
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
         freq, stat = epoch_folding_search(self.event_times, frequencies,
-                                          nbin=16, expocorr=True)
+                                          nbin=16, expocorr=True, gti=self.gti)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
-        assert maxstatbin == frequencies[minbin]
+        assert np.allclose(maxstatbin, frequencies[minbin], 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)
+        freq, stat = epoch_folding_search(self.times, frequencies,
+                                          nbin=16, weights=self.counts)
+
+        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(self):
         """Test pulse phase calculation, frequency only."""
-        frequencies = np.arange(9.85, 9.95, 0.3/self.tseg)
+        frequencies = np.arange(9.8, 9.99, 0.1/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]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)
 
     def test_z_n_search_expocorr(self):
         """Test pulse phase calculation, frequency only."""
-        frequencies = np.arange(9.89, 9.91, 0.1/self.tseg)
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
         freq, stat = z_n_search(self.event_times, frequencies, nbin=16,
-                                nharm=1, expocorr=True)
+                                nharm=1, expocorr=True, gti=self.gti)
+
+        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_expocorr_fails(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        with pytest.raises(ValueError) as excinfo:
+            freq, stat = z_n_search(self.event_times, frequencies, nharm=1,
+                                    nbin=16, expocorr=True)
+        assert 'To calculate exposure correction' in str(excinfo)
+
+    def test_z_n_search_weights(self):
+        """Test pulse phase calculation, frequency only."""
+        frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
+        freq, stat = z_n_search(self.times, frequencies, nbin=16,
+                                nharm=1, weights=self.counts)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
-        assert maxstatbin == frequencies[minbin]
+        assert np.allclose(maxstatbin, frequencies[minbin], atol=0.1/self.tseg)