Merge pull request #36 from stevertaylor/master

adding different BayesEphem model options

enterprise_extensions/models.py

@@ -17,7 +17,7 @@
 
 from enterprise_extensions import model_utils
 
-#### Extra model components not part of base enterprise ####
+# Extra model components not part of base enterprise ####
 
 # timing model delay
 @signal_base.function
@@ -34,24 +34,26 @@ def tm_delay(residuals, t2pulsar, tmparams_orig, tmparams, which='all'):
     :return: difference between new and old residuals in seconds
     """
 
-    if which == 'all': keys = tmparams_orig.keys()
-    else: keys = which
+    if which == 'all':
+        keys = tmparams_orig.keys()
+    else:
+        keys = which
 
     # grab original timing model parameters and errors in dictionary
     orig_params = np.array([tmparams_orig[key] for key in keys])
 
     # put varying parameters into dictionary
-    tmparams_rescaled = np.atleast_1d(np.double(orig_params[:,0] +
-                                                tmparams * orig_params[:,1]))
+    tmparams_rescaled = np.atleast_1d(np.double(orig_params[:, 0] +
+                                                tmparams * orig_params[:, 1]))
     tmparams_vary = OrderedDict(zip(keys, tmparams_rescaled))
 
     # set to new values
     t2pulsar.vals(tmparams_vary)
     new_res = np.double(t2pulsar.residuals().copy())
 
-    # remmeber to set values back to originals
+    # remember to set values back to originals
     t2pulsar.vals(OrderedDict(zip(keys,
-                                  np.atleast_1d(np.double(orig_params[:,0])))))
+                                  np.atleast_1d(np.double(orig_params[:, 0])))))
 
     # Return the time-series for the pulsar
     return new_res - residuals
@@ -209,9 +211,9 @@ def chrom_exp_cusp(toas, freqs, log10_Amp=-7, sign_param=-1.0,
     """
     t0 *= const.day
     tau = 10**log10_tau * const.day
-    wf = (10**log10_Amp * np.heaviside(toas - t0, 1) * \
-        np.exp(- (toas - t0) / tau)) + (10**log10_Amp * \
-        (1 - np.heaviside(toas - t0, 1)) * np.exp(- (t0 - toas) / tau))
+    wf = (10**log10_Amp * np.heaviside(toas - t0, 1) *
+          np.exp(- (toas - t0) / tau)) + (10**log10_Amp *
+                                          (1 - np.heaviside(toas - t0, 1)) * np.exp(- (t0 - toas) / tau))
 
     return np.sign(sign_param) * wf * (1400 / freqs) ** idx
 
@@ -285,7 +287,7 @@ def createfourierdesignmatrix_chromatic(toas, freqs, nmodes=30, Tspan=None,
     :param toas: vector of time series in seconds
     :param freqs: radio frequencies of observations [MHz]
     :param nmodes: number of fourier coefficients to use
-    :param freq: option to output frequencies
+    :param freqs: option to output frequencies
     :param Tspan: option to some other Tspan
     :param logf: use log frequency spacing
     :param fmin: lower sampling frequency
@@ -381,7 +383,7 @@ def turnover_knee(f, log10_A, gamma, lfb, lfk, kappa, delta):
     """
     Generic turnover spectrum with a high-frequency knee.
     :param f: sampling frequencies of GWB
-    :param A: characteristic strain amplitude at f=1/yr
+    :param log10_A: characteristic strain amplitude at f=1/yr
     :param gamma: negative slope of PSD around f=1/yr (usually 13/3)
     :param lfb: log10 transition frequency at which environment dominates GWs
     :param lfk: log10 knee frequency due to population finiteness
@@ -421,7 +423,7 @@ def generalized_gwpol_psd(f, log10_A_tt=-15, log10_A_st=-15,
                               orf_aa_sl*gwpol_amps[3]])
 
     S_psd = prefactor * (gwpol_factors[0,:] * (f / const.fyr)**(-4/3) +
-                         np.sum(gwpol_factors[1:,:],axis=0) * \
+                         np.sum(gwpol_factors[1:,:],axis=0) *
                          (f / const.fyr)**(-2)) / \
     (8*np.pi**2*f**3)
 
@@ -436,8 +438,10 @@ def dropout_powerlaw(f, log10_A=-16, gamma=5, k_drop=0.5, k_threshold=0.5):
 
     df = np.diff(np.concatenate((np.array([0]), f[::2])))
 
-    if k_drop >= k_threshold: k_switch = 1.0
-    elif k_drop < k_threshold: k_switch = 0.0
+    if k_drop >= k_threshold:
+        k_switch = 1.0
+    elif k_drop < k_threshold:
+        k_switch = 0.0
 
     return k_switch * ((10**log10_A)**2 / 12.0 / np.pi**2 *
                        const.fyr**(gamma-3) * f**(-gamma) * np.repeat(df, 2))
@@ -456,8 +460,10 @@ def dropout_physical_ephem_delay(toas, planetssb, pos_t, frame_drift_rate=0,
     """
 
     # get dropout switch
-    if k_drop >= k_threshold: k_switch = 1.0
-    elif k_drop < k_threshold: k_switch = 0.0
+    if k_drop >= k_threshold:
+        k_switch = 1.0
+    elif k_drop < k_threshold:
+        k_switch = 0.0
 
     # convert toas to MJD
     mjd = toas / 86400
@@ -776,7 +782,7 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi,
     WARNING: This residual waveform is only accurate if the
     GW frequency is not significantly evolving over the
     observation time of the pulsar.
-    :param toa: pulsar observation times
+    :param toas: pulsar observation times
     :param theta: polar coordinate of pulsar
     :param phi: azimuthal coordinate of pulsar
     :param gwtheta: Polar angle of GW source in celestial coords [radians]
@@ -822,8 +828,8 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi,
     omhat = np.array([-singwtheta*cosgwphi, -singwtheta*singwphi, -cosgwtheta])
 
     # pulsar position vector
-    phat = np.array([np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi),\
-            np.cos(theta)])
+    phat = np.array([np.sin(theta)*np.cos(phi), np.sin(theta)*np.sin(phi),
+                     np.cos(theta)])
 
     fplus = 0.5 * (np.dot(m, phat)**2 - np.dot(n, phat)**2) / (1+np.dot(omhat, phat))
     fcross = (np.dot(m, phat)*np.dot(n, phat)) / (1 + np.dot(omhat, phat))
@@ -866,13 +872,13 @@ def compute_eccentric_residuals(toas, theta, phi, cos_gwtheta, gwphi,
     # no more than 100 harmonics
     nharm = min(nharm, 100)
 
-    ##### earth term #####
+    # earth term #
     splus, scross = utils.calculate_splus_scross(nmax=nharm, mc=mc, dl=dist,
                                                  h0=h0, F=F, e=e0, t=toas.copy(),
                                                  l0=l0, gamma=gamma0,
                                                  gammadot=gammadot, inc=inc)
 
-    ##### pulsar term #####
+    # pulsar term #
     if psrTerm:
         # pulsar distance
         pd = pdist
@@ -1957,7 +1963,7 @@ def model_1(psrs, psd='powerlaw', noisedict=None, components=30,
 
 def model_2a(psrs, psd='powerlaw', noisedict=None, components=30,
              gamma_common=None, upper_limit=False, bayesephem=False,
-             wideband=False, select='backend'):
+             be_type='orbel', wideband=False, select='backend'):
     """
     Reads in list of enterprise Pulsar instance and returns a PTA
     instantiated with model 2A from the analysis paper:
@@ -1989,6 +1995,8 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30,
         value.
     :param bayesephem:
         Include BayesEphem model. Set to False by default
+    :param be_type:
+        orbel, orbel-v2, setIII
     :param wideband:
         Use wideband par and tim files. Ignore ECORR. Set to False by default.
     """
@@ -2008,7 +2016,7 @@ def model_2a(psrs, psd='powerlaw', noisedict=None, components=30,
 
     # ephemeris model
     if bayesephem:
-        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True)
+        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True, model=be_type)
 
     # timing model
     s += gp_signals.TimingModel()
@@ -2044,7 +2052,7 @@ def model_general(psrs, tm_var=False, tm_linear=False, tmparam_list=None,
                   tm_svd=False, tm_norm=True, gamma_common=None,
                   upper_limit=False, upper_limit_red=None, upper_limit_dm=None,
                   upper_limit_common=None,
-                  bayesephem=False, wideband=False,
+                  bayesephem=False, be_type='orbel', wideband=False,
                   dm_var=False, dm_type='gp', dm_psd='powerlaw', dm_annual=False,
                   white_vary=False, gequad=False, dm_chrom=False,
                   dmchrom_psd='powerlaw', dmchrom_idx=4,
@@ -2087,6 +2095,8 @@ def model_general(psrs, tm_var=False, tm_linear=False, tmparam_list=None,
         value.
     :param bayesephem:
         Include BayesEphem model. Set to False by default
+    :param be_type:
+        orbel, orbel-v2, setIII
     :param wideband:
         Use wideband par and tim files. Ignore ECORR. Set to False by default.
     """
@@ -2156,12 +2166,12 @@ def model_general(psrs, tm_var=False, tm_linear=False, tmparam_list=None,
             s += dm_annual_signal()
         if dm_chrom:
             s += chromatic_noise_block(psd=dmchrom_psd, idx=dmchrom_idx,
-                                       name='chromatic', components=components,
+                                       name='chromatic', components=dm_components,
                                        coefficients=coefficients)
 
     # ephemeris model
     if bayesephem:
-        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True)
+        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True, model=be_type)
 
     # adding white-noise, and acting on psr objects
     models = []
@@ -2456,7 +2466,7 @@ def model_2d(psrs, psd='powerlaw', noisedict=None, components=30,
 
 def model_3a(psrs, psd='powerlaw', noisedict=None, components=30,
              gamma_common=None, upper_limit=False, bayesephem=False,
-             wideband=False):
+             be_type='orbel', wideband=False):
     """
     Reads in list of enterprise Pulsar instance and returns a PTA
     instantiated with model 3A from the analysis paper:
@@ -2488,6 +2498,8 @@ def model_3a(psrs, psd='powerlaw', noisedict=None, components=30,
         value.
     :param bayesephem:
         Include BayesEphem model. Set to False by default
+    :param be_type:
+        orbel, orbel-v2, setIII
     """
 
     amp_prior = 'uniform' if upper_limit else 'log-uniform'
@@ -2505,7 +2517,7 @@ def model_3a(psrs, psd='powerlaw', noisedict=None, components=30,
 
     # ephemeris model
     if bayesephem:
-        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True)
+        s += deterministic_signals.PhysicalEphemerisSignal(use_epoch_toas=True, model=be_type)
 
     # timing model
     s += gp_signals.TimingModel()