Repository: StingraySoftware/stingray
Pull request #260 (2017-09-14 08:42:21)

Fix psr search with zn
* Fix failing z_n^2 test when nbin is not 16; test properly

* Solve bug when calculating orbital delays


Patch:
diff --git a/stingray/pulse/pulsar.py b/stingray/pulse/pulsar.py
index 8863a9bc..322c3a71 100644
--- a/stingray/pulse/pulsar.py
+++ b/stingray/pulse/pulsar.py
@@ -9,6 +9,7 @@ from scipy.optimize import minimize, basinhopping, curve_fit
 try:
     import pint.toa as toa
     import pint
+    from pint.models import get_model
     HAS_PINT = True
 except ImportError:
     HAS_PINT = False
@@ -592,8 +593,22 @@ def get_TOA(prof, period, tstart, template=None, additional_phase=0,
     return toa, toaerr
 
 
+def _load_and_prepare_TOAs(mjds, ephem="DE405"):
+    toalist = [None] * len(mjds)
+    for i, m in enumerate(mjds):
+        toalist[i] = toa.TOA(m, obs='Barycenter', scale='tdb')
+
+    toalist = toa.TOAs(toalist = toalist)
+    if 'tdb' not in toalist.table.colnames:
+        toalist.compute_TDBs()
+    if 'ssb_obs_pos' not in toalist.table.colnames:
+        toalist.compute_posvels(ephem, False)
+    return toalist
+
+
 def get_orbital_correction_from_ephemeris_file(mjdstart, mjdstop, parfile,
-                                               ntimes=1000, ephem="DE405"):
+                                               ntimes=1000, ephem="DE405",
+                                               return_pint_model=False):
     """Get a correction for orbital motion from pulsar parameter file.
 
     Parameters
@@ -617,6 +632,7 @@ def get_orbital_correction_from_ephemeris_file(mjdstart, mjdstop, parfile,
         Function that accepts times in MJDs and returns the deorbited times.
     """
     from scipy.interpolate import interp1d
+    from astropy import units
     simon("Assuming events are already referred to the solar system "
           "barycenter (timescale is TDB)")
     if not HAS_PINT:
@@ -624,23 +640,17 @@ def get_orbital_correction_from_ephemeris_file(mjdstart, mjdstop, parfile,
                           "functionality: github.com/nanograv/pint")
 
     mjds = np.linspace(mjdstart, mjdstop, ntimes)
-    toalist = [None] * len(mjds)
-    for i, m in enumerate(mjds):
-        toalist[i] = toa.TOA(m, obs='Barycenter', scale='tdb')
-
-    toalist = toa.TOAs(toalist = toalist)
-    if 'tdb' not in toalist.table.colnames:
-        toalist.compute_TDBs()
-    if 'ssb_obs_pos' not in toalist.table.colnames:
-        toalist.compute_posvels(ephem, False)
-
-    m = pint.models.get_model(parfile)
-    tdbtimes = m.get_barycentric_toas(toalist.table)
-
-    correction_mjd = interp1d(mjds, tdbtimes)
+    toalist = _load_and_prepare_TOAs(mjds, ephem=ephem)
+    m = get_model(parfile)
+    delays = m.delay(toalist.table)
+    correction_mjd = \
+        interp1d(mjds, (toalist.table['tdbld'] * units.d - delays).to(units.d).value)
 
     def correction_sec(times, mjdref):
         deorb_mjds = correction_mjd(times / 86400 + mjdref)
         return np.array((deorb_mjds - mjdref) * 86400)
 
-    return correction_sec, correction_mjd
\ No newline at end of file
+    retvals = [correction_sec, correction_mjd]
+    if return_pint_model:
+        retvals.append(m)
+    return retvals
diff --git a/stingray/pulse/search.py b/stingray/pulse/search.py
index f93774bc..4da3e99e 100644
--- a/stingray/pulse/search.py
+++ b/stingray/pulse/search.py
@@ -115,7 +115,7 @@ def epoch_folding_search(times, frequencies, nbin=128, segment_size=5000,
         return \
             _folding_search(fun, times, frequencies, segment_size=segment_size,
                             use_times=True, expocorr=expocorr, weights=weights,
-                            gti=gti)
+                            gti=gti, nbin=nbin)
 
     return _folding_search(lambda x: stat(_profile_fast(x, nbin=nbin)),
                            times, frequencies, segment_size=segment_size)
@@ -166,7 +166,7 @@ def z_n_search(times, frequencies, nharm=4, nbin=128, segment_size=5000,
                              ' specify the GTIs')
         def fun(t, f, **kwargs):
             return z_n(phase, n=nharm,
-                       norm=fold_events(t, f, **kwargs)[1])
+                       norm=fold_events(t, f, nbin=nbin, **kwargs)[1])
         return \
             _folding_search(fun, times, frequencies, segment_size=segment_size,
                             use_times=True, expocorr=expocorr, weights=weights,


Test patch:
diff --git a/stingray/pulse/tests/test_pulse.py b/stingray/pulse/tests/test_pulse.py
index 48144d21..3332b45a 100644
--- a/stingray/pulse/tests/test_pulse.py
+++ b/stingray/pulse/tests/test_pulse.py
@@ -54,9 +54,10 @@ class TestAll(object):
 
         mjdstart, mjdstop = mjds[0] - 1, mjds[-1] + 1
 
-        correction_sec, correction_mjd = \
+        correction_sec, correction_mjd, model = \
             get_orbital_correction_from_ephemeris_file(mjdstart, mjdstop,
-                                                       parfile, ntimes=1000)
+                                                       parfile, ntimes=1000,
+                                                       return_pint_model=True)
 
         mjdref = 50000
         toa_sec = (mjds - mjdref) * 86400
diff --git a/stingray/pulse/tests/test_search.py b/stingray/pulse/tests/test_search.py
index 3771af59..2eeb766f 100644
--- a/stingray/pulse/tests/test_search.py
+++ b/stingray/pulse/tests/test_search.py
@@ -130,7 +130,7 @@ class TestAll(object):
         """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,
-                                          nbin=16)
+                                          nbin=43)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
@@ -141,14 +141,14 @@ class TestAll(object):
         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)
+                                              nbin=23, expocorr=True)
         assert 'To calculate exposure correction' in str(excinfo)
 
     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,
-                                          nbin=16, expocorr=True, gti=self.gti)
+                                          nbin=42, expocorr=True, gti=self.gti)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
@@ -167,8 +167,8 @@ class TestAll(object):
     def test_z_n_search(self):
         """Test pulse phase calculation, frequency only."""
         frequencies = np.arange(9.8, 9.99, 0.1/self.tseg)
-        freq, stat = z_n_search(self.event_times, frequencies, nbin=16,
-                                nharm=1)
+        freq, stat = z_n_search(self.event_times, frequencies, nbin=25,
+                                nharm=2)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
@@ -177,8 +177,8 @@ class TestAll(object):
     def test_z_n_search_expocorr(self):
         """Test pulse phase calculation, frequency only."""
         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, gti=self.gti)
+        freq, stat = z_n_search(self.event_times, frequencies, nbin=64,
+                                nharm=2, expocorr=True, gti=self.gti)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))
         maxstatbin = freq[np.argmax(stat)]
@@ -189,13 +189,13 @@ class TestAll(object):
         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)
+                                    nbin=35, 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,
+        freq, stat = z_n_search(self.times, frequencies, nbin=44,
                                 nharm=1, weights=self.counts)
 
         minbin = np.argmin(np.abs(frequencies - self.pulse_frequency))