diff --git a/enterprise_extensions/models.py b/enterprise_extensions/models.py index e66e2a75..c5a1df35 100644 --- a/enterprise_extensions/models.py +++ b/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()