Repository: StingraySoftware/stingray
Pull request #222 (2017-07-25 12:42:46)

Simulation of event lists from lightcurve moved to simulator submodule
Move the core of the acceptance-rejection method to the `simulator` module where it belongs.

Comment for reviewers: the changes to `simulator.py` are only PEP8-related!


Patch:
diff --git a/stingray/events.py b/stingray/events.py
index 00b03d67..af4c5594 100644
--- a/stingray/events.py
+++ b/stingray/events.py
@@ -10,10 +10,10 @@ from .utils import simon, assign_value_if_none
 from .gti import cross_gtis, append_gtis, check_separate
 
 from .lightcurve import Lightcurve
+from stingray.simulator.base import simulate_times
 
 import numpy as np
 import numpy.random as ra
-import scipy.interpolate as sci
 
 __all__ = ['EventList']
 
@@ -150,90 +150,30 @@ class EventList(object):
 
     def simulate_times(self, lc, use_spline=False, bin_time=None):
         """
-        Assign (simulate) photon arrival times to event list, using acceptance-rejection
-        method.
+        Assign (simulate) photon arrival times to event list, using the
+        acceptance-rejection method.
 
         Parameters
         ----------
         lc: `Lightcurve` object
-        """ 
-       
-        times = lc.time
-        counts = lc.counts
-
-        bin_time = assign_value_if_none(bin_time, np.median(np.diff(times)))
-        n_bin = len(counts)
-        bin_start = 0
-        maxlc = np.max(counts)
-        intlc = maxlc * n_bin
-        n_events_predict = int(intlc + 10 * np.sqrt(intlc))
-
-        # Max number of events per chunk must be < 100000
-        events_per_bin_predict = n_events_predict / n_bin
-        
-        if use_spline:
-            max_bin = np.long(np.max([4, 1000000 / events_per_bin_predict]))
-            
-        else:
-            max_bin = np.long(np.max([4, 5000000 / events_per_bin_predict]))
-
-        ev_list = np.zeros(n_events_predict)
-        nev = 0
-
-        while bin_start < n_bin:
-            
-            t0 = times[bin_start]
-            bin_stop = min([bin_start + max_bin, n_bin + 1])
-            
-            lc_filt = counts[bin_start:bin_stop]
-            t_filt = times[bin_start:bin_stop]
-            length = t_filt[-1] - t_filt[0]
-            
-            n_bin_filt = len(lc_filt)
-            n_to_simulate = n_bin_filt * max(lc_filt)
-            safety_factor = 10
-
-            if n_to_simulate > 10000:
-                safety_factor = 4.
 
-            n_to_simulate += safety_factor * np.sqrt(n_to_simulate)
-            n_to_simulate = int(np.ceil(n_to_simulate))
-            n_predict = ra.poisson(np.sum(lc_filt))
-
-            random_ts = ra.uniform(t_filt[0] - bin_time / 2,
-                                   t_filt[-1] + bin_time / 2, n_to_simulate)
-
-            random_amps = ra.uniform(0, max(lc_filt), n_to_simulate)
-            
-            if use_spline:
-                lc_spl = sci.splrep(t_filt, lc_filt, s=np.longdouble(0), k=1)
-                pts = sci.splev(random_ts, lc_spl)
-            
-            else:
-                rough_bins = np.rint((random_ts - t0) / bin_time)
-                rough_bins = rough_bins.astype(int)
-                pts = lc_filt[rough_bins]
-            
-            good = random_amps < pts
-            len1 = len(random_ts)
-            random_ts = random_ts[good]
-            
-            len2 = len(random_ts)
-            random_ts = random_ts[:n_predict]
-            random_ts.sort()
-            
-            new_nev = len(random_ts)
-            ev_list[nev:nev + new_nev] = random_ts[:]
-            nev += new_nev
-            bin_start += max_bin
-
-        # Discard all zero entries at the end
-        time = ev_list[:nev]
-        time.sort()
-        
-        self.time = EventList(time).time
-        self.ncounts = len(self.time)
+        Other Parameters
+        ----------------
+        use_spline : bool
+            Approximate the light curve with a spline to avoid binning effects
+        bin_time : float
+            The bin time of the light curve, if it needs to be specified for
+            improved precision
+
+        Return
+        ------
+        times : array-like
+            Simulated photon arrival times
+        """
+        self.time = simulate_times(lc, use_spline=use_spline,
+                                   bin_time=bin_time)
         self.gti = lc.gti
+        self.ncounts = len(self.time)
 
     def simulate_energies(self, spectrum):
         """
diff --git a/stingray/simulator/base.py b/stingray/simulator/base.py
new file mode 100644
index 00000000..7273e550
--- /dev/null
+++ b/stingray/simulator/base.py
@@ -0,0 +1,104 @@
+from __future__ import division, print_function
+import numpy as np
+import numpy.random as ra
+import scipy.interpolate as sci
+from ..utils import assign_value_if_none
+
+
+def simulate_times(lc, use_spline=False, bin_time=None):
+    """
+    Assign (simulate) photon arrival times to event list, using the
+    acceptance-rejection method.
+
+    Parameters
+    ----------
+    lc: `Lightcurve` object
+
+    Other Parameters
+    ----------------
+    use_spline : bool
+        Approximate the light curve with a spline to avoid binning effects
+    bin_time : float
+        The bin time of the light curve, if it needs to be specified for
+        improved precision
+
+    Return
+    ------
+    times : array-like
+        Simulated photon arrival times
+    """
+
+    times = lc.time
+    counts = lc.counts
+
+    bin_time = assign_value_if_none(bin_time, lc.dt)
+    n_bin = len(counts)
+    bin_start = 0
+    maxlc = np.max(counts)
+    intlc = maxlc * n_bin
+    n_events_predict = int(intlc + 10 * np.sqrt(intlc))
+
+    # Max number of events per chunk must be < 100000
+    events_per_bin_predict = n_events_predict / n_bin
+
+    if use_spline:
+        max_bin = np.long(np.max([4, 1000000 / events_per_bin_predict]))
+
+    else:
+        max_bin = np.long(np.max([4, 5000000 / events_per_bin_predict]))
+
+    ev_list = np.zeros(n_events_predict)
+    nev = 0
+
+    while bin_start < n_bin:
+
+        t0 = times[bin_start]
+        bin_stop = min([bin_start + max_bin, n_bin + 1])
+
+        lc_filt = counts[bin_start:bin_stop]
+        t_filt = times[bin_start:bin_stop]
+        length = t_filt[-1] - t_filt[0]
+
+        n_bin_filt = len(lc_filt)
+        n_to_simulate = n_bin_filt * max(lc_filt)
+        safety_factor = 10
+
+        if n_to_simulate > 10000:
+            safety_factor = 4.
+
+        n_to_simulate += safety_factor * np.sqrt(n_to_simulate)
+        n_to_simulate = int(np.ceil(n_to_simulate))
+        n_predict = ra.poisson(np.sum(lc_filt))
+
+        random_ts = ra.uniform(t_filt[0] - bin_time / 2,
+                               t_filt[-1] + bin_time / 2, n_to_simulate)
+
+        random_amps = ra.uniform(0, max(lc_filt), n_to_simulate)
+
+        if use_spline:
+            lc_spl = sci.splrep(t_filt, lc_filt, s=np.longdouble(0), k=1)
+            pts = sci.splev(random_ts, lc_spl)
+
+        else:
+            rough_bins = np.rint((random_ts - t0) / bin_time)
+            rough_bins = rough_bins.astype(int)
+            pts = lc_filt[rough_bins]
+
+        good = random_amps < pts
+        len1 = len(random_ts)
+        random_ts = random_ts[good]
+
+        len2 = len(random_ts)
+        random_ts = random_ts[:n_predict]
+        random_ts.sort()
+
+        new_nev = len(random_ts)
+        ev_list[nev:nev + new_nev] = random_ts[:]
+        nev += new_nev
+        bin_start += max_bin
+
+    # Discard all zero entries at the end
+    time = ev_list[:nev]
+    time.sort()
+
+    return time
diff --git a/stingray/simulator/simulator.py b/stingray/simulator/simulator.py
index 72b5788c..5cf4c7e7 100644
--- a/stingray/simulator/simulator.py
+++ b/stingray/simulator/simulator.py
@@ -7,6 +7,7 @@ import astropy.modeling.models
 from stingray import Lightcurve, AveragedPowerspectrum, io, utils
 import stingray.simulator.models as models
 
+
 class Simulator(object):
 
     def __init__(self, dt=1, N=1024, mean=0, rms=1, red_noise=1,
@@ -218,7 +219,8 @@ class Simulator(object):
         channel = [lc for lc in self.channels if lc[0] == channel]
 
         if len(channel) == 0:
-            raise KeyError('This channel does not exist or has already been deleted.')
+            raise KeyError('This channel does not exist or has already been '
+                           'deleted.')
         else:
             index = self.channels.index(channel[0])
             del self.channels[index]
@@ -231,8 +233,8 @@ class Simulator(object):
         channels = [lc for lc in self.channels if lc[0] in channels]
 
         if len(channels) != n:
-            raise KeyError('One of more of the channels do not exist or have already been'
-                'deleted.')
+            raise KeyError('One of more of the channels do not exist or have '
+                           'already been deleted.')
         else:
             indices = [self.channels.index(channel) for channel in channels]
             for i in sorted(indices, reverse=True):
@@ -275,7 +277,8 @@ class Simulator(object):
 
         return np.append(h_zeros, h_ones)
 
-    def relativistic_ir(self, t1=3, t2=4, t3=10, p1=1, p2=1.4, rise=0.6, decay=0.1):
+    def relativistic_ir(self, t1=3, t2=4, t3=10, p1=1, p2=1.4, rise=0.6,
+                        decay=0.1):
         """
         Construct a realistic impulse response considering the relativistic
         effects.
@@ -400,7 +403,8 @@ class Simulator(object):
         ----------
         model: astropy.modeling.Model derived function
             the pre-defined model
-            (library-based, available in astropy.modeling.models or custom-defined)
+            (library-based, available in astropy.modeling.models or
+            custom-defined)
 
         Returns
         -------
@@ -557,7 +561,9 @@ class Simulator(object):
             lc = long_lc
         else:
             # Make random cut and extract light curve of length 'N'
-            extract = self.random_state.randint(self.N-1, self.red_noise*self.N - self.N+1)
+            extract = \
+                self.random_state.randint(self.N-1,
+                                          self.red_noise*self.N - self.N+1)
             lc = np.take(long_lc, range(extract, extract + self.N))
 
         avg = np.mean(lc)


Test patch:
diff --git a/stingray/simulator/tests/test_base.py b/stingray/simulator/tests/test_base.py
new file mode 100644
index 00000000..a1c9c98c
--- /dev/null
+++ b/stingray/simulator/tests/test_base.py
@@ -0,0 +1,33 @@
+from __future__ import division
+from stingray.simulator.base import simulate_times
+from stingray.lightcurve import Lightcurve
+import numpy as np
+
+
+class TestSimulator(object):
+
+    @classmethod
+    def setup_class(self):
+        self.time = [0.5, 1.5, 2.5, 3.5]
+        self.counts_flat = [3000, 3000, 3000, 3000]
+        self.gti = [[0, 4]]
+
+    def test_simulate_times(self):
+        """Simulate photon arrival times for an event list
+        from light curve.
+        """
+        lc = Lightcurve(self.time, self.counts_flat, gti=self.gti)
+        times = simulate_times(lc)
+        lc_sim = Lightcurve.make_lightcurve(times, gti=lc.gti, dt=lc.dt,
+                                            tstart=lc.tstart, tseg=lc.tseg)
+        assert np.all((lc - lc_sim).counts < 3 * np.sqrt(lc.counts))
+
+    def test_simulate_times_with_spline(self):
+        """Simulate photon arrival times, with use_spline option
+        enabled.
+        """
+        lc = Lightcurve(self.time, self.counts_flat, gti=self.gti)
+        times = simulate_times(lc, use_spline=True)
+        lc_sim = Lightcurve.make_lightcurve(times, gti=lc.gti, dt=lc.dt,
+                                            tstart=lc.tstart, tseg=lc.tseg)
+        assert np.all((lc - lc_sim).counts < 3 * np.sqrt(lc.counts))
diff --git a/stingray/tests/test_events.py b/stingray/tests/test_events.py
index 13915309..4c3aabf7 100644
--- a/stingray/tests/test_events.py
+++ b/stingray/tests/test_events.py
@@ -23,7 +23,9 @@ class TestEvents(object):
     def setup_class(self):
         self.time = [0.5, 1.5, 2.5, 3.5]
         self.counts = [3000, 2000, 2200, 3600]
-        self.spectrum = [[1, 2, 3, 4, 5, 6],[1000, 2040, 1000, 3000, 4020, 2070]]
+        self.counts_flat = [3000, 3000, 3000, 3000]
+        self.spectrum = [[1, 2, 3, 4, 5, 6],
+                         [1000, 2040, 1000, 3000, 4020, 2070]]
         self.gti = [[0, 4]]
 
     def test_inequal_length(self):
@@ -52,28 +54,21 @@ class TestEvents(object):
         """Simulate photon arrival times for an event list 
         from light curve.
         """
-        lc = Lightcurve(self.time, self.counts)
+        lc = Lightcurve(self.time, self.counts_flat, gti=self.gti)
         ev = EventList()
         ev.simulate_times(lc)
+        lc_sim = ev.to_lc(dt=lc.dt, tstart=lc.tstart, tseg=lc.tseg)
+        assert np.all((lc - lc_sim).counts < 3 * np.sqrt(lc.counts))
 
     def test_simulate_times_with_spline(self):
         """Simulate photon arrival times, with use_spline option
         enabled.
         """
-        lc = Lightcurve(self.time, self.counts)
+        lc = Lightcurve(self.time, self.counts_flat, gti=self.gti)
         ev = EventList()
-        ev.simulate_times(lc, use_spline = True)
-
-    def test_recover_lc(self):
-        """Test that the original light curve can be recovered from 
-        event list, which was used to simulate it.
-        """
-        lc = Lightcurve(self.time, self.counts)
-        ev = EventList()
-        ev.simulate_times(lc)
-
-        lc_rcd = ev.to_lc(dt=1, tstart=0, tseg=4)
-        assert np.all(np.abs(lc_rcd.counts - lc.counts) < 3 * np.sqrt(lc.counts))
+        ev.simulate_times(lc, use_spline=True)
+        lc_sim = ev.to_lc(dt=lc.dt, tstart=lc.tstart, tseg=lc.tseg)
+        assert np.all((lc - lc_sim).counts < 3 * np.sqrt(lc.counts))
 
     def test_simulate_energies(self):
         """Assign photon energies to an event list.