Merge pull request #58 from svigeland/master

changed how GWB frequencies are found in _get_freqs() and get_Fmats()

enterprise_extensions/blocks.py

@@ -75,8 +75,13 @@ def white_noise_block(vary=False, inc_ecorr=False, gp_ecorr=False,
                                   selection=backend, name=name)
     if inc_ecorr:
         if gp_ecorr:
-            ec = gp_signals.EcorrBasisModel(log10_ecorr=ecorr,
-                                            selection=backend_ng, name=name)
+            if name is None:
+                ec = gp_signals.EcorrBasisModel(log10_ecorr=ecorr,
+                                                selection=backend_ng)
+            else:
+                ec = gp_signals.EcorrBasisModel(log10_ecorr=ecorr,
+                                                selection=backend_ng, name=name)
+
         else:
             ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr,
                                                 selection=backend_ng,

enterprise_extensions/frequentist/F_statistic.py

@@ -4,30 +4,71 @@
 import scipy.linalg as sl
 import scipy.special
 
-from enterprise_extensions import models
+from enterprise.signals import signal_base
+from enterprise.signals import gp_signals
+from enterprise.signals import deterministic_signals
+from enterprise_extensions import blocks
+from enterprise_extensions import deterministic
 
 
 class FpStat(object):
     """
     Class for the Fp-statistic.
 
     :param psrs: List of `enterprise` Pulsar instances.
+    :param noisedict: Dictionary of white noise parameter values. Default=None
     :param psrTerm: Include the pulsar term in the CW signal model. Default=True
     :param bayesephem: Include BayesEphem model. Default=True
     """
 
     def __init__(self, psrs, params=None,
-                 psrTerm=True, bayesephem=True, wideband=False, pta=None):
+                 psrTerm=True, bayesephem=True, pta=None):
 
         if pta is None:
 
-            # initialize standard model with fixed white noise and powerlaw red noise
+            # initialize standard model with fixed white noise
+            # and powerlaw red noise
+            # uses the implementation of ECORR in gp_signals
             print('Initializing the model...')
-            self.pta = models.model_cw(psrs, noisedict=params, rn_psd='powerlaw',
-                                       ecc=False, psrTerm=psrTerm,
-                                       bayesephem=bayesephem, wideband=wideband)
+
+            tmin = np.min([p.toas.min() for p in psrs])
+            tmax = np.max([p.toas.max() for p in psrs])
+            Tspan = tmax - tmin
+
+            s = deterministic.cw_block_circ(amp_prior='log-uniform',
+                                            psrTerm=psrTerm, tref=tmin, name='cw')
+            s += gp_signals.TimingModel()
+            s += blocks.red_noise_block(prior='log-uniform', psd='powerlaw',
+                                        Tspan=Tspan, components=30)
+
+            if bayesephem:
+                s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True)
+
+            # adding white-noise, and acting on psr objects
+            models = []
+            for p in psrs:
+                if 'NANOGrav' in p.flags['pta']:
+                    s2 = s + blocks.white_noise_block(vary=False, inc_ecorr=True,
+                                                      gp_ecorr=True)
+                    models.append(s2(p))
+                else:
+                    s3 = s + blocks.white_noise_block(vary=False, inc_ecorr=False)
+                    models.append(s3(p))
+
+            pta = signal_base.PTA(models)
+
+            # set white noise parameters
+            if params is None:
+                print('No noise dictionary provided!')
+            else:
+                pta.set_default_params(params)
+
+            self.pta = pta
 
         else:
+
+            # user can specify their own pta object
+            # if ECORR is included, use the implementation in gp_signals
             self.pta = pta
 
         self.psrs = psrs

enterprise_extensions/frequentist/optimal_statistic.py

@@ -38,7 +38,7 @@ def __init__(self, psrs, bayesephem=True, gamma_common=4.33, orf='hd',
                                        bayesephem=bayesephem,
                                        gamma_common=gamma_common,
                                        wideband=wideband,
-                                       select=select, noisedict=noisedict)
+                                       select='backend', noisedict=noisedict)
         else:
             self.pta = pta
 
@@ -187,7 +187,7 @@ def get_Fmats(self, params={}):
         for sc in self.pta._signalcollections:
             ind = []
             for signal, idx in sc._idx.items():
-                if signal.signal_name == 'red noise':
+                if signal.signal_name == 'red noise' and signal.signal_id =='gw':
                     ind.append(idx)
             ix = np.unique(np.concatenate(ind))
             Fmats.append(sc.get_basis(params=params)[:, ix])
@@ -197,7 +197,7 @@ def get_Fmats(self, params={}):
     def _get_freqs(self,psrs):
         """ Hackish way to get frequency vector."""
         for sig in self.pta._signalcollections[0]._signals:
-            if sig.signal_name == 'red noise':
+            if sig.signal_name == 'red noise' and sig.signal_id == 'gw':
                 sig._construct_basis()
                 freqs = np.array(sig._labels[''])
                 break