diff --git a/enterprise_extensions/blocks.py b/enterprise_extensions/blocks.py index 708cf902..0f3230f4 100644 --- a/enterprise_extensions/blocks.py +++ b/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, diff --git a/enterprise_extensions/frequentist/F_statistic.py b/enterprise_extensions/frequentist/F_statistic.py index 3cec41be..ac039bde 100644 --- a/enterprise_extensions/frequentist/F_statistic.py +++ b/enterprise_extensions/frequentist/F_statistic.py @@ -4,7 +4,11 @@ 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): @@ -12,22 +16,59 @@ 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 diff --git a/enterprise_extensions/frequentist/optimal_statistic.py b/enterprise_extensions/frequentist/optimal_statistic.py index e2171c30..395a3da2 100644 --- a/enterprise_extensions/frequentist/optimal_statistic.py +++ b/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