Repository: StingraySoftware/stingray
Pull request #184 (2017-04-13 15:00:17)

Generic base class for Variability vs energy spectra
I'm closing #183 and reopening here because there was a problem with the history following a rebase. 

Anyways, this will be the base for lag/energy, rms/energy, etc.

I added an implementation of `RMSEnergySpectrum` to show how such spectra should be implemented using this base class

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/stingraysoftware/stingray/184)
<!-- Reviewable:end -->



Patch:
diff --git a/stingray/events.py b/stingray/events.py
index 132580af..8faeb79d 100644
--- a/stingray/events.py
+++ b/stingray/events.py
@@ -145,7 +145,7 @@ class EventList(object):
         # Concatenate all lists
         times = [i for j in times for i in j]
 
-        return EventList(time=times)
+        return EventList(time=times, gti=lc.gti)
 
     def simulate_times(self, lc, use_spline=False, bin_time=None):
         """
@@ -171,10 +171,10 @@ class EventList(object):
         events_per_bin_predict = n_events_predict / n_bin
         
         if use_spline:
-            max_bin = np.max([4, 1000000 / events_per_bin_predict])
+            max_bin = np.long(np.max([4, 1000000 / events_per_bin_predict]))
             
         else:
-            max_bin = np.max([4, 5000000 / events_per_bin_predict])
+            max_bin = np.long(np.max([4, 5000000 / events_per_bin_predict]))
 
         ev_list = np.zeros(n_events_predict)
         nev = 0
@@ -232,6 +232,7 @@ class EventList(object):
         
         self.time = EventList(time).time
         self.ncounts = len(self.time)
+        self.gti = lc.gti
 
     def simulate_energies(self, spectrum):
         """
diff --git a/stingray/varenergyspectrum.py b/stingray/varenergyspectrum.py
new file mode 100644
index 00000000..79cc417f
--- /dev/null
+++ b/stingray/varenergyspectrum.py
@@ -0,0 +1,197 @@
+from __future__ import division
+import numpy as np
+from stingray.gti import check_separate, cross_two_gtis
+from stingray.lightcurve import Lightcurve
+from stingray.utils import assign_value_if_none, simon
+from stingray.crossspectrum import AveragedCrossspectrum
+from abc import ABCMeta, abstractmethod
+import six
+
+
+@six.add_metaclass(ABCMeta)
+class VarEnergySpectrum(object):
+    def __init__(self, events, freq_interval, energy_spec, ref_band=None,
+                 bin_time=1, use_pi=False, segment_size=None, events2=None):
+
+        """Base variability-energy spectrum.
+
+        This class is only a base for the various variability spectra, and it's
+        not to be instantiated by itself.
+
+        Parameters
+        ----------
+        events : stingray.events.EventList object
+            event list
+        freq_interval : [f0, f1], floats
+            the frequency range over which calculating the variability quantity
+        energy_spec : list or tuple (emin, emax, N, type)
+            if a list is specified, this is interpreted as a list of bin edges;
+            if a tuple is provided, this will encode the minimum and maximum
+            energies, the number of intervals, and "lin" or "log".
+
+        Other Parameters
+        ----------------
+        ref_band : [emin, emax], floats; default None
+            minimum and maximum energy of the reference band. If None, the
+            full band is used.
+        use_pi : boolean, default False
+            Use channel instead of energy
+        events2 : stingray.events.EventList object
+            event list for the second channel, if not the same. Useful if the
+            reference band has to be taken from another detector.
+            
+        Attributes
+        ----------
+        events1 : array-like
+            list of events used to produce the spectrum
+        events2 : array-like
+            if the spectrum requires it, second list of events
+        freq_interval : array-like
+            interval of frequencies used to calculate the spectrum
+        energy_intervals : [[e00, e01], [e10, e11], ...]
+            energy intervals used for the spectrum
+        spectrum : array-like
+            the spectral values, corresponding to each energy interval
+        spectrum_error : array-like
+            the errorbars corresponding to spectrum
+        
+        """
+        self.events1 = events
+        self.events2 = assign_value_if_none(events2, events)
+        self.freq_interval = freq_interval
+        self.use_pi = use_pi
+        self.bin_time = bin_time
+        if isinstance(energy_spec, tuple):
+            if energy_spec[-1].lower() not in ["lin", "log"]:
+                raise ValueError("Incorrect energy specification")
+
+            log_distr = True if energy_spec[-1].lower() == "log" else False
+
+            if log_distr:
+                energies = np.logspace(np.log10(energy_spec[0]),
+                                       np.log10(energy_spec[1]),
+                                       energy_spec[2] + 1)
+            else:
+                energies = np.linspace(energy_spec[0], energy_spec[1],
+                                       energy_spec[2] + 1)
+        else:
+            energies = np.asarray(energy_spec)
+
+        self.energy_intervals = list(zip(energies[0: -1], energies[1:]))
+
+        self.ref_band = assign_value_if_none(ref_band, [0, np.inf])
+
+        self.segment_size = segment_size
+        self.spectrum, self.spectrum_error = self._spectrum_function()
+
+    def _decide_ref_intervals(self, channel_band, ref_band):
+        """Eliminate channel_band from ref_band."""
+        if check_separate([ref_band], [channel_band]):
+            return np.asarray([ref_band])
+        not_channel_band = [[0, channel_band[0]],
+                            [channel_band[1], np.max([ref_band[-1],
+                                                      channel_band[1] + 1])]]
+        return cross_two_gtis([ref_band], not_channel_band)
+
+    def _construct_lightcurves(self, channel_band, tstart=None, tstop=None,
+                               exclude=True):
+        if self.use_pi:
+            energies1 = self.events1.pi
+            energies2 = self.events2.pi
+        else:
+            energies2 = self.events2.pha
+            energies1 = self.events1.pha
+
+        gti = cross_two_gtis(self.events1.gti, self.events2.gti)
+
+        tstart = assign_value_if_none(tstart,
+                                      np.max([self.events1.time[0],
+                                              self.events2.time[0]]))
+        tstop = assign_value_if_none(tstop,
+                                     np.min([self.events1.time[-1],
+                                             self.events2.time[-1]]))
+
+        good = (energies1 >= channel_band[0]) & (energies1 < channel_band[1])
+        base_lc = Lightcurve.make_lightcurve(self.events1.time[good],
+                                             self.bin_time,
+                                             tstart=tstart,
+                                             tseg=tstop - tstart,
+                                             gti=gti)
+
+        if exclude:
+            ref_intervals = self._decide_ref_intervals(channel_band,
+                                                       self.ref_band)
+        else:
+            ref_intervals = [self.ref_band]
+
+        ref_lc = Lightcurve(base_lc.time, np.zeros_like(base_lc.counts),
+                            gti=base_lc.gti)
+
+        for i in ref_intervals:
+            good = (energies2 >= i[0]) & (energies2 < i[1])
+            new_lc = Lightcurve.make_lightcurve(self.events2.time[good],
+                                                self.bin_time,
+                                                tstart=tstart,
+                                                tseg=tstop - tstart,
+                                                gti=gti)
+            ref_lc = ref_lc + new_lc
+
+        return base_lc, ref_lc
+
+    @abstractmethod
+    def _spectrum_function(self):
+        pass
+
+
+class RmsEnergySpectrum(VarEnergySpectrum):
+
+    def _spectrum_function(self):
+
+        rms_spec = np.zeros(len(self.energy_intervals))
+        rms_spec_err = np.zeros_like(rms_spec)
+        for i, eint in enumerate(self.energy_intervals):
+            base_lc, ref_lc = self._construct_lightcurves(eint,
+                                                          exclude=False)
+            xspect = AveragedCrossspectrum(base_lc, ref_lc,
+                                           segment_size=self.segment_size,
+                                           norm='frac')
+            good = (xspect.freq >= self.freq_interval[0]) & \
+                   (xspect.freq < self.freq_interval[1])
+            rms_spec[i] = np.sqrt(np.sum(xspect.power[good]*xspect.df))
+
+            # Root squared sum of errors of the spectrum
+            root_sq_err_sum = np.sqrt(np.sum(xspect.power[good]**2))*xspect.df
+            # But the rms is the squared root. So,
+            # Error propagation
+            rms_spec_err[i] = 1 / (2 * rms_spec[i]) * root_sq_err_sum
+
+        return rms_spec, rms_spec_err
+
+
+class LagEnergySpectrum(VarEnergySpectrum):
+
+    def _spectrum_function(self):
+
+        lag_spec = np.zeros(len(self.energy_intervals))
+        lag_spec_err = np.zeros_like(lag_spec)
+        for i, eint in enumerate(self.energy_intervals):
+            base_lc, ref_lc = self._construct_lightcurves(eint)
+            xspect = AveragedCrossspectrum(base_lc, ref_lc,
+                                           segment_size=self.segment_size)
+            good = (xspect.freq >= self.freq_interval[0]) & \
+                   (xspect.freq < self.freq_interval[1])
+            lag, lag_err = xspect.time_lag()
+            good_lag, good_lag_err = lag[good], lag_err[good]
+            coh, coh_err = xspect.coherence()
+            lag_spec[i] = np.mean(good_lag)
+            coh_check = coh > 1.2 / (1 + 0.2 * xspect.m)
+            if not np.all(coh_check[good]):
+                simon("Coherence is not ideal over the specified energy range."
+                      " Lag values and uncertainties might be underestimated. "
+                      "See Epitropakis and Papadakis, A\&A 591, 1113, 2016")
+
+            # Root squared sum of errors of the spectrum
+            # Verified!
+            lag_spec_err[i] = np.sqrt(np.sum(good_lag_err**2) / len(good_lag))
+
+        return lag_spec, lag_spec_err


Test patch:
diff --git a/stingray/tests/test_varenergyspectrum.py b/stingray/tests/test_varenergyspectrum.py
new file mode 100644
index 00000000..d9c8fcc7
--- /dev/null
+++ b/stingray/tests/test_varenergyspectrum.py
@@ -0,0 +1,171 @@
+import numpy as np
+from stingray.events import EventList
+from stingray.varenergyspectrum import VarEnergySpectrum, RmsEnergySpectrum
+from stingray.varenergyspectrum import LagEnergySpectrum
+from stingray.lightcurve import Lightcurve
+
+from astropy.tests.helper import pytest
+np.random.seed(20150907)
+
+
+class DummyVarEnergy(VarEnergySpectrum):
+    def _spectrum_function(self):
+        return None, None
+
+
+class TestVarEnergySpectrum(object):
+
+    @classmethod
+    def setup_class(cls):
+        tstart = 0.0
+        tend = 100.0
+        nphot = 1000
+        alltimes = np.random.uniform(tstart, tend, nphot)
+        alltimes.sort()
+        cls.events = EventList(alltimes,
+                               pha=np.random.uniform(0.3, 12, nphot),
+                               gti = [[tstart, tend]])
+        cls.vespec = DummyVarEnergy(cls.events, [0., 10000],
+                                    (0.5, 5, 10, "lin"), [0.3, 10],
+                                    bin_time=0.1)
+        cls.vespeclog = \
+            DummyVarEnergy(cls.events, [0., 10000],
+                           (0.5, 5, 10, "log"), [0.3, 10])
+
+    def test_intervals_overlapping(self):
+        ref_int = self.vespec._decide_ref_intervals([0.5, 6], [0.3, 10])
+        np.testing.assert_allclose(ref_int, [[0.3, 0.5], [6, 10]])
+        ref_int = self.vespec._decide_ref_intervals([0.5, 11], [0.3, 10])
+        np.testing.assert_allclose(ref_int, [[0.3, 0.5]])
+
+    def test_intervals_non_overlapping(self):
+        ref_int = self.vespec._decide_ref_intervals([6, 11], [0.3, 5])
+        np.testing.assert_allclose(ref_int, [[0.3, 5]])
+
+    def test_ref_band_none(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pha=[0,0,0,0,1,1],
+                           gti=[[0, 0.65]])
+        vespec = DummyVarEnergy(events, [0., 10000],
+                                (0, 1, 2, "lin"),
+                                bin_time=0.1)
+
+    def test_energy_spec_wrong_list_not_tuple(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pha=[0, 0, 0, 0, 1, 1],
+                           gti=[[0, 0.65]])
+        # Test using a list instead of tuple
+        # with pytest.raises(ValueError):
+        vespec = DummyVarEnergy(events, [0., 10000],
+                                [0, 1, 2, "lin"],
+                                bin_time=0.1)
+
+    def test_energy_spec_wrong_str(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pha=[0, 0, 0, 0, 1, 1],
+                           gti=[[0, 0.65]])
+        # Test using a list instead of tuple
+        with pytest.raises(ValueError):
+            vespec = DummyVarEnergy(events, [0., 10000],
+                                    (0, 1, 2, "xxx"),
+                                    bin_time=0.1)
+
+    def test_construct_lightcurves(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pha=[0,0,0,0,1,1],
+                           gti=[[0, 0.65]])
+        vespec = DummyVarEnergy(events, [0., 10000],
+                                (0, 1, 2, "lin"), [0.5, 1.1],
+                                bin_time=0.1)
+        base_lc, ref_lc = \
+            vespec._construct_lightcurves([0, 0.5],
+                                          tstart=0, tstop=0.65)
+        np.testing.assert_allclose(base_lc.counts, [1, 0, 2, 1, 0, 0])
+        np.testing.assert_allclose(ref_lc.counts, [0, 0, 0, 1, 0, 1])
+
+    def test_construct_lightcurves_no_exclude(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pha=[0,0,0,0,1,1],
+                           gti=[[0, 0.65]])
+
+        vespec = DummyVarEnergy(events, [0., 10000],
+                                (0, 1, 2, "lin"), [0, 0.5],
+                                bin_time=0.1)
+        base_lc, ref_lc = \
+            vespec._construct_lightcurves([0, 0.5],
+                                          tstart=0, tstop=0.65,
+                                          exclude=False)
+        np.testing.assert_equal(base_lc.counts, ref_lc.counts)
+
+    def test_construct_lightcurves_pi(self):
+        events = EventList([0.09, 0.21, 0.23, 0.32, 0.4, 0.54],
+                           pi=np.asarray([0, 0, 0, 0, 1, 1]),
+                           gti=[[0, 0.65]])
+        vespec = DummyVarEnergy(events, [0., 10000],
+                                (0, 1, 2, "lin"), [0.5, 1.1], use_pi=True,
+                                   bin_time=0.1)
+        base_lc, ref_lc = \
+            vespec._construct_lightcurves([0, 0.5],
+                                          tstart=0, tstop=0.65)
+        np.testing.assert_allclose(base_lc.counts, [1, 0, 2, 1, 0, 0])
+        np.testing.assert_allclose(ref_lc.counts, [0, 0, 0, 1, 0, 1])
+
+
+class TestRMSEnergySpectrum(object):
+    @classmethod
+    def setup_class(cls):
+        from ..simulator.simulator import Simulator
+
+        simulator = Simulator(0.1, 10000, rms=0.4, mean=200)
+        test_lc = simulator.simulate(1)
+        test_ev1, test_ev2 = EventList(), EventList()
+        test_ev1.simulate_times(test_lc)
+        test_ev2.simulate_times(test_lc)
+        test_ev1.pha = np.random.uniform(0.3, 12, len(test_ev1.time))
+        test_ev2.pha = np.random.uniform(0.3, 12, len(test_ev2.time))
+
+        cls.rms = RmsEnergySpectrum(test_ev1, [0., 100],
+                                    (0.3, 12, 5, "lin"),
+                                    bin_time=0.01,
+                                    segment_size=100,
+                                    events2=test_ev2)
+
+    def test_correct_rms_values(self):
+        # Assert that it is close to 0.4 (since we don't have infinite spectral
+        # coverage, it will be a little less!)
+        assert np.allclose(self.rms.spectrum, 0.37, 0.05)
+
+    def test_correct_rms_errorbars(self):
+
+        # Assert that the rms measured at all energies is the same
+        assert np.all(
+            np.abs(self.rms.spectrum - self.rms.spectrum[0]) < \
+                self.rms.spectrum_error)
+
+
+class TestLagEnergySpectrum(object):
+    @classmethod
+    def setup_class(cls):
+        from ..simulator.simulator import Simulator
+        simulator = Simulator(0.1, 10000, rms=0.4, mean=2000)
+        test_lc1 = simulator.simulate(2)
+        test_lc2 = Lightcurve(test_lc1.time,
+                              np.array(np.roll(test_lc1.counts, 4)),
+                              err_dist=test_lc1.err_dist)
+
+        test_ev1, test_ev2 = EventList(), EventList()
+        test_ev1.simulate_times(test_lc1)
+        test_ev2.simulate_times(test_lc2)
+        test_ev1.pha = np.random.uniform(0.3, 9, len(test_ev1.time))
+        test_ev2.pha = np.random.uniform(9, 12, len(test_ev2.time))
+
+        cls.lag = LagEnergySpectrum(test_ev1, [0., 0.5],
+                                    (0.3, 9, 4, "lin"), [9, 12],
+                                    bin_time=0.1,
+                                    segment_size=30,
+                                    events2=test_ev2)
+
+    def test_lagspectrum_values_and_errors(self):
+
+        assert np.all(np.abs(self.lag.spectrum - 0.2) < \
+                      3 * self.lag.spectrum_error)