Repository: StingraySoftware/stingray
Pull request #177 (2017-03-24 16:27:39)

Adding support for astropy.modeling.models
1. Support for astropy.modeling: now an user can choose to simulate one of the models defined in astropy.modeling.models, a library function that uses the astropy.modeling.Model syntax (see point 2 below), an ex-novo user defined astropy.modeling.Model function or an arbitrary combination of these.
Example:
`smoothbknpo = models.SmoothBrokenPowerLaw(norm=1., gamma_low=1., gamma_high=2., break_freq=1.)`
`lorentzian = models.GeneralizedLorentz1D(x_0=10, fwhm=1., value=10., power_coeff=2.)`
`myModel = smoothbknpo + lorentzian`
`sim = simulator.Simulator(N=50000, mean=30000., dt=0.01)`
`lc = sim.simulate(myModel)`

2. Refactored generalised Lorentzian and smooth broken power law functions (in stingray.simulator.models) using astropy.modeling.Model class

3. Code cleanup

<!-- 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/177)
<!-- Reviewable:end -->



Patch:
diff --git a/stingray/simulator/models.py b/stingray/simulator/models.py
index 875200da..23b876d4 100644
--- a/stingray/simulator/models.py
+++ b/stingray/simulator/models.py
@@ -1,4 +1,55 @@
 from __future__ import division, print_function, absolute_import
+from astropy.modeling.models import custom_model
+
+@custom_model
+def GeneralizedLorentz1D(x, x_0=1., fwhm=1., value=1., power_coeff=1.):
+    """
+    Generalized lorenzian function, 
+    implemented using astropy.modeling.models custom model
+
+    Parameters
+    ----------
+
+    x: numpy.ndarray
+        non-zero frequencies
+
+    x_0 : peak centeral frequency
+    fwhm : FWHM of the peak (gamma)
+    value : peak value at x=x0
+    power_coeff : power coefficient [n]
+
+    Returns
+    -------
+    model: astropy.modeling.Model
+        generalized Lorentzian psd model
+    """
+    assert power_coeff > 0., "The power coefficient should be greater than zero."
+    return value * (fwhm / 2)**power_coeff * 1./(abs(x - x_0)**power_coeff + (fwhm / 2)**power_coeff) 
+    
+@custom_model
+def SmoothBrokenPowerLaw(x, norm=1., gamma_low=1., gamma_high=1., break_freq=1.):
+    """
+    Smooth broken power law function,
+    implemented using astropy.modeling.models custom model
+
+    Parameters
+    ----------
+
+    x: numpy.ndarray
+        non-zero frequencies
+
+    norm: normalization frequency
+    gamma_low: power law index for f --> zero
+    gamma_high: power law index for f --> infinity
+    break_freq: break frequency
+
+    Returns
+    -------
+    model: astropy.modeling.Model
+        generalized smooth broken power law psd model
+    """
+    return norm * x**(-gamma_low) / (1. + (x / break_freq)**2)**(-(gamma_low - gamma_high) / 2)
+    
 
 def lorenzian(x, p):
     """
diff --git a/stingray/simulator/simulator.py b/stingray/simulator/simulator.py
index 3264180f..17731da3 100644
--- a/stingray/simulator/simulator.py
+++ b/stingray/simulator/simulator.py
@@ -2,6 +2,7 @@ from __future__ import division, print_function, absolute_import
 import numpy as np
 import numbers
 from scipy import signal
+import astropy.modeling.models
 
 from stingray import Lightcurve, AveragedPowerspectrum, io, utils
 import stingray.simulator.models as models
@@ -53,15 +54,13 @@ class Simulator(object):
 
         Examples
         --------
-        - x = simulate(2)
+        - x = simulate(beta)
             For generating a light curve using power law spectrum.
 
             Parameters
             ----------
-            Beta: int
+            beta: int
                 Defines the shape of spectrum
-            N: int
-                Number of samples
 
             Returns
             -------
@@ -79,17 +78,27 @@ class Simulator(object):
             -------
             lightCurve: `LightCurve` object
 
-        - x = simulate('lorenzian', p)
+        - x = simulate(model)
             For generating a light curve from pre-defined model
 
             Parameters
             ----------
-            model: str
-                name of the pre-defined model
+            model: astropy.modeling.Model
+                the pre-defined model
 
-            p: list iterable
-                model parameters. For details, see model definitions
-                in model.py
+            Returns
+            -------
+            lightCurve: `LightCurve` object
+        
+        - x = simulate('model', params)
+            For generating a light curve from pre-defined model
+
+            Parameters
+            ----------
+            model: string
+                the pre-defined model
+            params: list iterable or dict
+                the parameters for the pre-defined model
 
             Returns
             -------
@@ -136,12 +145,15 @@ class Simulator(object):
         if isinstance(args[0], numbers.Integral) and len(args) == 1: 
             return  self._simulate_power_law(args[0])
 
+        elif isinstance(args[0], astropy.modeling.Model) and len(args) == 1:
+            return self._simulate_model(args[0])
+        
+        elif utils.is_string(args[0]) and len(args) == 2:
+            return self._simulate_model_string(args[0], args[1])
+        
         elif len(args) == 1:
             return self._simulate_power_spectrum(args[0])
 
-        elif utils.is_string(args[0]) and len(args) == 2:
-            return self._simulate_model(args[0], args[1])
-
         elif len(args) == 2:
             return self._simulate_impulse_response(args[0], args[1])
 
@@ -151,6 +163,7 @@ class Simulator(object):
         else:
             raise ValueError("Length of arguments must be 1, 2 or 3.")
 
+
     def simulate_channel(self, channel, *args):
         """
         Simulate a lightcurve and add it to corresponding energy 
@@ -376,42 +389,95 @@ class Simulator(object):
 
         return lc
 
-    def _simulate_model(self, mod, p):
+            
+    def _simulate_model(self, model):
         """
-        For generating a light curve from pre-defined model
+        For generating a light curve from a pre-defined model
 
         Parameters
         ----------
-        mod: str
-            name of the pre-defined model
-
-        p: list iterable
-            model parameters. For details, see model definitions
-            in model.py
-
+        model: astropy.modeling.Model derived function
+            the pre-defined model 
+            (library-based, available in astropy.modeling.models or custom-defined)
+        
         Returns
         -------
         lightCurve: `LightCurve` object
         """
 
-        # Define frequencies at which to compute PSD
-        w = np.fft.rfftfreq(self.red_noise*self.N, d=self.dt)[1:]
+        # Frequencies at which the PSD is to be computed 
+        # (only positive frequencies, since the signal is real)
+        nbins = self.red_noise*self.N
+        simfreq = np.fft.rfftfreq(nbins, d=self.dt)[1:]
+        
+        # Compute PSD from model
+        simpsd = model(simfreq)
+        
+        fac = np.sqrt(simpsd/2.)
+        pos_real   = np.random.normal(size=nbins//2)*fac
+        pos_imag   = np.random.normal(size=nbins//2)*fac
+        
+        pos_freq_transform = pos_real + 1j * pos_imag
 
-        if mod in dir(models):
-            mod = 'models.'+ mod
-            s = eval(mod +'(w, p)')
+        # Simulate light curve from its Fourier transform
+        arg  = np.concatenate(([self.mean], pos_freq_transform))
+    
+        # Inverse Fourier transform
+        long_lc = np.fft.irfft(arg)
+        
+        lc = Lightcurve(self.time, self._extract_and_scale(long_lc))
+        return lc
+        
+        
+    def _simulate_model_string(self, model_str, params):
+        """
+        For generating a light curve from a pre-defined model
 
-            # Draw two set of 'N' guassian distributed numbers
-            a1 = np.random.normal(size=len(s))
-            a2 = np.random.normal(size=len(s))
+        Parameters
+        ----------
+        model_str: string
+            name of the pre-defined model 
+        params: list or dictionary
+            parameters of the pre-defined model
+        
+        Returns
+        -------
+        lightCurve: `LightCurve` object
+        """
 
-            long_lc = self._find_inverse(a1*s, a2*s)
+        # Frequencies at which the PSD is to be computed 
+        # (only positive frequencies, since the signal is real)
+        nbins = self.red_noise*self.N
+        simfreq = np.fft.rfftfreq(nbins, d=self.dt)[1:]
+        
+        if model_str in dir(models):
+            if isinstance(params, dict):
+                model = eval('models.' + model_str + '(**params)')
+                # Compute PSD from model
+                simpsd = model(simfreq)
+            elif isinstance(params, list):
+                simpsd = eval('models.' + model_str + '(simfreq, params)')
+            else:
+                raise ValueError('Params should be list or dictionary!')    
+        
+            fac = np.sqrt(simpsd/2.)
+            pos_real   = np.random.normal(size=nbins//2)*fac
+            pos_imag   = np.random.normal(size=nbins//2)*fac
+            
+            pos_freq_transform = pos_real + 1j * pos_imag
+
+            # Simulate light curve from its Fourier transform
+            arg  = np.concatenate(([self.mean], pos_freq_transform))
+    
+            # Inverse Fourier transform
+            long_lc = np.fft.irfft(arg)
+        
             lc = Lightcurve(self.time, self._extract_and_scale(long_lc))
-
             return lc
-        
         else:
             raise ValueError('Model is not defined!')
+        
+        
 
     def _simulate_impulse_response(self, s, h, mode='same'):
         """


Test patch:
diff --git a/stingray/simulator/tests/test_simulator.py b/stingray/simulator/tests/test_simulator.py
index b1d09fd1..a54855f6 100644
--- a/stingray/simulator/tests/test_simulator.py
+++ b/stingray/simulator/tests/test_simulator.py
@@ -2,6 +2,7 @@ import numpy as np
 import os
 
 from astropy.tests.helper import pytest
+import astropy.modeling.models
 from stingray import Lightcurve, Crossspectrum, sampledata
 from stingray.simulator import simulator, models
 
@@ -202,6 +203,74 @@ class TestSimulator(object):
         assert np.all(
             np.abs(actual_prob - simulated_prob) < 3*np.sqrt(actual_prob))
 
+    def test_simulate_GeneralizedLorentz1D_str(self):
+        """
+        Simulate a light curve using the GeneralizedLorentz1D model
+        called as a string
+        """
+        assert len(self.simulator.simulate('GeneralizedLorentz1D', {'x_0':10, 'fwhm':1., 'value':10., 'power_coeff':2})), 1024
+        
+    def test_simulate_GeneralizedLorentz1D(self):
+        """
+        Simulate a light curve using the GeneralizedLorentz1D model
+        called as a astropy.modeling.Model class
+        """
+        mod = models.GeneralizedLorentz1D(x_0=10, fwhm=1., value=10., power_coeff=2)
+        assert len(self.simulator.simulate(mod)), 1024
+        
+    def test_simulate_SmoothBrokenPowerLaw_str(self):
+        """
+        Simulate a light curve using SmoothBrokenPowerLaw model
+        called as a string
+        """
+        assert len(self.simulator.simulate('SmoothBrokenPowerLaw', {'norm':1., 'gamma_low':1., 'gamma_high':2., 'break_freq':1.})), 1024
+    
+    def test_simulate_SmoothBrokenPowerLaw(self):
+        """
+        Simulate a light curve using SmoothBrokenPowerLaw model
+        called as a astropy.modeling.Model class
+        """
+        mod = models.SmoothBrokenPowerLaw(norm=1., gamma_low=1., gamma_high=2., break_freq=1.)
+        assert len(self.simulator.simulate(mod)), 1024
+        
+        
+    def test_simulate_generic_model(self):
+        """
+        Simulate a light curve using a generic model
+        called as a astropy.modeling.Model class
+        """
+        mod = astropy.modeling.models.Gaussian1D(amplitude=10., mean=1., stddev=2.)
+        assert len(self.simulator.simulate(mod)), 1024
+        
+        
+    def test_compare_composite(self):
+        """
+        Compare the PSD of a light curve simulated using a composite model
+        (using SmoothBrokenPowerLaw plus GeneralizedLorentz1D)
+        with the actual model
+        """
+        N = 50000
+        dt = 0.01
+        m = 30000.
+        
+        self.simulator = simulator.Simulator(N=N, mean=m, dt=dt)
+        smoothbknpo = models.SmoothBrokenPowerLaw(norm=1., gamma_low=1., gamma_high=2., break_freq=1.)
+        lorentzian = models.GeneralizedLorentz1D(x_0=10, fwhm=1., value=10., power_coeff=2.)
+        myModel = smoothbknpo + lorentzian
+        
+        lc = [self.simulator.simulate(myModel) for i in range(1, 50)]
+
+        simulated = self.simulator.powerspectrum(lc, lc[0].tseg)
+
+        w = np.fft.rfftfreq(N, d=dt)[1:]
+        actual = myModel(w)[:-1]
+
+        actual_prob = actual/float(sum(actual))
+        simulated_prob = simulated/float(sum(simulated))
+
+        assert np.all(np.abs(actual_prob - simulated_prob) < 3*np.sqrt(actual_prob))
+        
+        
     def test_simulate_wrong_model(self):
         """
         Simulate with a model that does not exist.