def __init__(self, psrs, params=None): print('Initializing the model...') efac = parameter.Constant() equad = parameter.Constant() ef = white_signals.MeasurementNoise(efac=efac) eq = white_signals.EquadNoise(log10_equad=equad) tm = gp_signals.TimingModel(use_svd=True) s = eq + ef + tm model = [] for p in psrs: model.append(s(p)) self.pta = signal_base.PTA(model) # set white noise parameters if params is None: print('No noise dictionary provided!...') else: self.pta.set_default_params(params) self.psrs = psrs self.params = params self.Nmats = None
def white_noise_block(vary=False): """ Returns the white noise block of the model: 1. EFAC per backend/receiver system 2. EQUAD per backend/receiver system 3. ECORR per backend/receiver system :param vary: If set to true we vary these parameters with uniform priors. Otherwise they are set to constants with values to be set later. """ # define selection by observing backend selection = selections.Selection(selections.by_backend) # white noise parameters if vary: efac = parameter.Uniform(0.01, 10.0) equad = parameter.Uniform(-8.5, -5) ecorr = parameter.Uniform(-8.5, -5) else: efac = parameter.Constant() equad = parameter.Constant() ecorr = parameter.Constant() # white noise signals ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection) # combine signals s = ef + eq + ec return s
def test_equad_backend(self): """Test that backend-equad signal returns correct covariance.""" # set up signal and parameters equad = parameter.Uniform(-10, -5) selection = Selection(selections.by_backend) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) eqm = eq(self.psr) # parameters equads = [-6.1, -6.2, -6.3, -6.4] params = { "B1855+09_430_ASP_log10_equad": equads[0], "B1855+09_430_PUPPI_log10_equad": equads[1], "B1855+09_L-wide_ASP_log10_equad": equads[2], "B1855+09_L-wide_PUPPI_log10_equad": equads[3], } # correct value flags = ["430_ASP", "430_PUPPI", "L-wide_ASP", "L-wide_PUPPI"] nvec0 = np.zeros_like(self.psr.toas) for ct, flag in enumerate(np.unique(flags)): ind = flag == self.psr.backend_flags nvec0[ind] = 10**(2 * equads[ct]) * np.ones(np.sum(ind)) # test msg = "EQUAD covariance incorrect." assert np.all(eqm.get_ndiag(params) == nvec0), msg
def add_equad(self): """Return dictionary containing EQUAD white noise signal attributes :return: OrderedDict of EQUAD signal """ equad_dct = dict() equad = parameter.Uniform(-10.0, -4.0) eq = white_signals.EquadNoise(log10_equad=equad, selection=self.selection) self.equad_sig = eq(self.psr) for ii, param in enumerate(self.equad_sig.param_names): Nvec = np.ones_like(self.psr.toaerrs) * self.equad_sig._masks[ii] newsignal = OrderedDict({ 'type': 'equad', 'name': param, 'pmin': [-10.0], 'pmax': [-4.0], 'pstart': [-6.5], 'interval': [True], 'numpars': 1, 'Nvec': Nvec }) equad_dct.update({param: newsignal}) return equad_dct
def equad(self,option="by_backend"): """ EQUAD signal: adds EQUAD**2 to the ToA variance, where ToA variance are diagonal components of the Likelihood covariance matrix. """ if option not in selections.__dict__.keys(): raise ValueError('EQUAD option must be Enterprise selection function \ name') se=selections.Selection(selections.__dict__[option]) equadpr = interpret_white_noise_prior(self.params.equad) eqs = white_signals.EquadNoise(log10_equad=equadpr,selection=se) return eqs
def test_equad(self): """Test that equad signal returns correct covariance.""" # set up signal and parameters equad = parameter.Uniform(-10, -5) eq = white_signals.EquadNoise(log10_equad=equad) eqm = eq(self.psr) # parameters equad = -6.4 params = {"B1855+09_log10_equad": equad} # correct value nvec0 = 10**(2 * equad) * np.ones_like(self.psr.toas) # test msg = "EQUAD covariance incorrect." assert np.all(eqm.get_ndiag(params) == nvec0), msg
def test_powerlaw_equad(self): def powerlaw(f, log10_A=-15): return (10**log10_A) * f**2 def log10(A=10**-16): return np.log10(A) fquad = white_signals.EquadNoise(log10_equad=parameter.Function( log10, A=parameter.Uniform(10**-17, 10**-14))) fquad1 = fquad(self.psr) repr_A = '[B1855+09_log10_equad_A:Uniform(pmin=1e-17, pmax=1e-14)]' repr_B = '[B1855+09_log10_equad_A:Uniform(pmax=1e-14, pmin=1e-17)]' assert str(fquad1.params) == repr_A or str(fquad1.params) == repr_B assert np.allclose( fquad1.get_ndiag(params={'B1855+09_log10_equad_A': 10**-14})[:3], np.array([1e-28, 1e-28, 1e-28]))
def initialize_pta_sim(psrs, fgw, inc_efac=True, inc_equad=False, inc_ecorr=False, selection=None, inc_red_noise=False, noisedict=None): # continuous GW signal s = models.cw_block_circ(log10_fgw=np.log10(fgw), psrTerm=True) # linearized timing model s += gp_signals.TimingModel(use_svd=True) # white noise if selection == 'backend': selection = selections.Selection(selections.by_backend) if inc_efac: efac = parameter.Constant() s += white_signals.MeasurementNoise(efac=efac, selection=selection) if inc_equad: equad = parameter.Constant() s += white_signals.EquadNoise(log10_equad=equad, selection=selection) if inc_ecorr: ecorr = parameter.Constant() s += gp_signals.EcorrBasisModel(log10_ecorr=ecorr, selection=selection) if inc_red_noise: log10_A = parameter.Constant() gamma = parameter.Constant() pl = utils.powerlaw(log10_A=log10_A, gamma=gamma) s += gp_signals.FourierBasisGP(pl, components=30) model = [s(psr) for psr in psrs] pta = signal_base.PTA(model) # set white noise parameters if noisedict is None: print('No noise dictionary provided!...') else: pta.set_default_params(noisedict) return pta
def test_add_efac_equad(self): """Test that addition of efac and equad signal returns correct covariance. """ # set up signals efac = parameter.Uniform(0.1, 5) ef = white_signals.MeasurementNoise(efac=efac) equad = parameter.Uniform(-10, -5) eq = white_signals.EquadNoise(log10_equad=equad) s = ef + eq m = s(self.psr) # set parameters efac = 1.5 equad = -6.4 params = {"B1855+09_efac": efac, "B1855+09_log10_equad": equad} # correct value nvec0 = efac**2 * self.psr.toaerrs**2 nvec0 += 10**(2 * equad) * np.ones_like(self.psr.toas) # test msg = "EFAC/EQUAD covariance incorrect." assert np.all(m.get_ndiag(params) == nvec0), msg
# GW parameters (initialize with names here to use parameters in common across pulsars) log10_A_gw_1 = parameter.Uniform(-18, -13)('zlog10_A_gw') gamma_gw_1 = parameter.Constant(13 / 3)('zgamma_gw') # Second GW parameters log10_A_gw_2 = parameter.Uniform(-18, -13)('zlog10_A_other_gw') gamma_gw_2 = parameter.Constant(7 / 3)('zgamma_other_gw') ##### Set up signals ##### # timing model tm = gp_signals.TimingModel() # white noise ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=log10_equad, selection=selection) ec = white_signals.EcorrKernelNoise(log10_ecorr=log10_ecorr, selection=selection) # red noise (powerlaw with 30 frequencies) pl = utils.powerlaw(log10_A=red_noise_log10_A, gamma=red_noise_gamma) rn = gp_signals.FourierBasisGP(spectrum=pl, components=30, Tspan=Tspan) cpl_1 = utils.powerlaw(log10_A=log10_A_gw_1, gamma=gamma_gw_1) cpl_2 = utils.powerlaw(log10_A=log10_A_gw_2, gamma=gamma_gw_2) #Common red noise process with no correlations crn_1 = gp_signals.FourierBasisGP(spectrum=cpl_1, components=30, Tspan=Tspan, name='gw')
rn_log10_A = parameter.Uniform(-20, -11) rn_gamma = parameter.Uniform(0, 7) rn_pl = utils.powerlaw(log10_A=rn_log10_A, gamma=rn_gamma) # GWB parameters and powerlaw orf = utils.hd_orf() if args.ul: gwb_log10_A = parameter.LinearExp(-18, -12)('gwb_log10_A') else: gwb_log10_A = parameter.Uniform(-18, -12)('gwb_log10_A') gwb_gamma = parameter.Constant(13 / 3)('gwb_gamma') gwb_pl = utils.powerlaw(log10_A=gwb_log10_A, gamma=gwb_gamma) # signals ef = white_signals.MeasurementNoise(efac=efac, selection=bkend) eq = white_signals.EquadNoise(log10_equad=equad, selection=bkend) ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=bkend_NG) dmgp = gp_signals.BasisGP(dm_pl, dm_basis, name='dm_gp') dmexp = models.dm_exponential_dip(tmin=54500, tmax=55000, name='dmexp_1') dm1yr = models.dm_annual_signal() rn = gp_signals.FourierBasisGP(spectrum=rn_pl, components=30, Tspan=Tspan) gwb = gp_signals.FourierBasisCommonGP(gwb_pl, orf, components=30, Tspan=Tspan, name='gw') tm = gp_signals.TimingModel(use_svd=True) be = deterministic_signals.PhysicalEphemerisSignal( # widen prior on jup orbit
def test_single_pulsar(self): # get parameters from PAL2 style noise files params = get_noise_from_pal2(datadir + '/B1855+09_noise.txt') # setup basic model efac = parameter.Constant() equad = parameter.Constant() ecorr = parameter.Constant() log10_A = parameter.Constant() gamma = parameter.Constant() selection = Selection(selections.by_backend) ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection) pl = utils.powerlaw(log10_A=log10_A, gamma=gamma) rn = gp_signals.FourierBasisGP(pl) s = ef + eq + ec + rn m = s(self.psrs[0]) # set parameters m.set_default_params(params) # get parameters efacs = [params[key] for key in sorted(params.keys()) if 'efac' in key] equads = [ params[key] for key in sorted(params.keys()) if 'equad' in key ] ecorrs = [ params[key] for key in sorted(params.keys()) if 'ecorr' in key ] log10_A = params['B1855+09_red_noise_log10_A'] gamma = params['B1855+09_red_noise_gamma'] # correct value flags = ['430_ASP', '430_PUPPI', 'L-wide_ASP', 'L-wide_PUPPI'] nvec0 = np.zeros_like(self.psrs[0].toas) for ct, flag in enumerate(np.unique(flags)): ind = flag == self.psrs[0].backend_flags nvec0[ind] = efacs[ct]**2 * self.psrs[0].toaerrs[ind]**2 nvec0[ind] += 10**(2 * equads[ct]) * np.ones(np.sum(ind)) # get the basis bflags = self.psrs[0].backend_flags Umats = [] for flag in np.unique(bflags): mask = bflags == flag Umats.append( utils.create_quantization_matrix(self.psrs[0].toas[mask])[0]) nepoch = sum(U.shape[1] for U in Umats) U = np.zeros((len(self.psrs[0].toas), nepoch)) jvec = np.zeros(nepoch) netot = 0 for ct, flag in enumerate(np.unique(bflags)): mask = bflags == flag nn = Umats[ct].shape[1] U[mask, netot:nn + netot] = Umats[ct] jvec[netot:nn + netot] = 10**(2 * ecorrs[ct]) netot += nn # get covariance matrix cov = np.diag(nvec0) + np.dot(U * jvec[None, :], U.T) cf = sl.cho_factor(cov) logdet = np.sum(2 * np.log(np.diag(cf[0]))) # test msg = 'EFAC/ECORR logdet incorrect.' N = m.get_ndiag(params) assert np.allclose(N.solve(self.psrs[0].residuals, logdet=True)[1], logdet, rtol=1e-10), msg msg = 'EFAC/ECORR D1 solve incorrect.' assert np.allclose(N.solve(self.psrs[0].residuals), sl.cho_solve(cf, self.psrs[0].residuals), rtol=1e-10), msg msg = 'EFAC/ECORR 1D1 solve incorrect.' assert np.allclose(N.solve(self.psrs[0].residuals, left_array=self.psrs[0].residuals), np.dot(self.psrs[0].residuals, sl.cho_solve(cf, self.psrs[0].residuals)), rtol=1e-10), msg msg = 'EFAC/ECORR 2D1 solve incorrect.' T = m.get_basis(params) assert np.allclose(N.solve(self.psrs[0].residuals, left_array=T), np.dot(T.T, sl.cho_solve(cf, self.psrs[0].residuals)), rtol=1e-10), msg msg = 'EFAC/ECORR 2D2 solve incorrect.' assert np.allclose(N.solve(T, left_array=T), np.dot(T.T, sl.cho_solve(cf, T)), rtol=1e-10), msg F, f2 = utils.createfourierdesignmatrix_red(self.psrs[0].toas, nmodes=20) # spectrum test phi = utils.powerlaw(f2, log10_A=log10_A, gamma=gamma) msg = 'Spectrum incorrect for GP Fourier signal.' assert np.all(m.get_phi(params) == phi), msg # inverse spectrum test msg = 'Spectrum inverse incorrect for GP Fourier signal.' assert np.all(m.get_phiinv(params) == 1 / phi), msg
##### parameters and priors ##### # Uniform prior on EFAC and EQUAD efac = parameter.Uniform(0.1, 5.0) log10_equad = parameter.Uniform(-10.0, -4.0) # red noise parameters # Uniform in log10 Amplitude and in spectral index log10_A = parameter.Uniform(-18, -12) gamma = parameter.Uniform(0, 7) ##### Set up signals ##### # white noise ef = white_signals.MeasurementNoise(efac=efac) eq = white_signals.EquadNoise(log10_equad=log10_equad) # red noise (powerlaw with 30 frequencies) pl = utils.powerlaw(log10_A=log10_A, gamma=gamma) rn = gp_signals.FourierBasisGP(spectrum=pl, components=30) # timing model tm = gp_signals.TimingModel() # full model is sum of components model = ef + rn + tm + eq psr_dict = {} for psr in psrs: print('Working on ' + psr.name)
def white_noise_block(vary=False, inc_ecorr=False, gp_ecorr=False, efac1=False, select='backend', name=None): """ Returns the white noise block of the model: 1. EFAC per backend/receiver system 2. EQUAD per backend/receiver system 3. ECORR per backend/receiver system :param vary: If set to true we vary these parameters with uniform priors. Otherwise they are set to constants with values to be set later. :param inc_ecorr: include ECORR, needed for NANOGrav channelized TOAs :param gp_ecorr: whether to use the Gaussian process model for ECORR :param efac1: use a strong prior on EFAC = Normal(mu=1, stdev=0.1) """ if select == 'backend': # define selection by observing backend backend = selections.Selection(selections.by_backend) # define selection by nanograv backends backend_ng = selections.Selection(selections.nanograv_backends) else: # define no selection backend = selections.Selection(selections.no_selection) # white noise parameters if vary: if efac1: efac = parameter.Normal(1.0, 0.1) else: efac = parameter.Uniform(0.01, 10.0) equad = parameter.Uniform(-8.5, -5) if inc_ecorr: ecorr = parameter.Uniform(-8.5, -5) else: efac = parameter.Constant() equad = parameter.Constant() if inc_ecorr: ecorr = parameter.Constant() # white noise signals ef = white_signals.MeasurementNoise(efac=efac, selection=backend, name=name) eq = white_signals.EquadNoise(log10_equad=equad, selection=backend, name=name) if inc_ecorr: if gp_ecorr: 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, name=name) # combine signals if inc_ecorr: s = ef + eq + ec elif not inc_ecorr: s = ef + eq return s
def compute_like(self, npsrs=1, inc_corr=False, inc_kernel=False): # get parameters from PAL2 style noise files params = get_noise_from_pal2(datadir + "/B1855+09_noise.txt") params2 = get_noise_from_pal2(datadir + "/J1909-3744_noise.txt") params.update(params2) psrs = self.psrs if npsrs == 2 else [self.psrs[0]] if inc_corr: params.update({"GW_gamma": 4.33, "GW_log10_A": -15.0}) # find the maximum time span to set GW frequency sampling tmin = [p.toas.min() for p in psrs] tmax = [p.toas.max() for p in psrs] Tspan = np.max(tmax) - np.min(tmin) # setup basic model efac = parameter.Constant() equad = parameter.Constant() ecorr = parameter.Constant() log10_A = parameter.Constant() gamma = parameter.Constant() selection = Selection(selections.by_backend) ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection) pl = utils.powerlaw(log10_A=log10_A, gamma=gamma) rn = gp_signals.FourierBasisGP(pl) orf = utils.hd_orf() crn = gp_signals.FourierBasisCommonGP(pl, orf, components=20, name="GW", Tspan=Tspan) tm = gp_signals.TimingModel() log10_sigma = parameter.Uniform(-10, -5) log10_lam = parameter.Uniform(np.log10(86400), np.log10(1500 * 86400)) basis = create_quant_matrix(dt=7 * 86400) prior = se_kernel(log10_sigma=log10_sigma, log10_lam=log10_lam) se = gp_signals.BasisGP(prior, basis, name="se") # set up kernel stuff if isinstance(inc_kernel, bool): inc_kernel = [inc_kernel] * npsrs if inc_corr: s = ef + eq + ec + rn + crn + tm else: s = ef + eq + ec + rn + tm models = [] for ik, psr in zip(inc_kernel, psrs): snew = s + se if ik else s models.append(snew(psr)) pta = signal_base.PTA(models) # set parameters pta.set_default_params(params) # SE kernel parameters log10_sigmas, log10_lams = [-7.0, -6.5], [7.0, 6.5] params.update({ "B1855+09_se_log10_lam": log10_lams[0], "B1855+09_se_log10_sigma": log10_sigmas[0], "J1909-3744_se_log10_lam": log10_lams[1], "J1909-3744_se_log10_sigma": log10_sigmas[1], }) # get parameters efacs, equads, ecorrs, log10_A, gamma = [], [], [], [], [] lsig, llam = [], [] for pname in [p.name for p in psrs]: efacs.append([ params[key] for key in sorted(params.keys()) if "efac" in key and pname in key ]) equads.append([ params[key] for key in sorted(params.keys()) if "equad" in key and pname in key ]) ecorrs.append([ params[key] for key in sorted(params.keys()) if "ecorr" in key and pname in key ]) log10_A.append(params["{}_red_noise_log10_A".format(pname)]) gamma.append(params["{}_red_noise_gamma".format(pname)]) lsig.append(params["{}_se_log10_sigma".format(pname)]) llam.append(params["{}_se_log10_lam".format(pname)]) GW_gamma = 4.33 GW_log10_A = -15.0 # correct value tflags = [sorted(list(np.unique(p.backend_flags))) for p in psrs] cfs, logdets, phis, Ts = [], [], [], [] for ii, (ik, psr, flags) in enumerate(zip(inc_kernel, psrs, tflags)): nvec0 = np.zeros_like(psr.toas) for ct, flag in enumerate(flags): ind = psr.backend_flags == flag nvec0[ind] = efacs[ii][ct]**2 * psr.toaerrs[ind]**2 nvec0[ind] += 10**(2 * equads[ii][ct]) * np.ones(np.sum(ind)) # get the basis bflags = psr.backend_flags Umats = [] for flag in np.unique(bflags): mask = bflags == flag Umats.append( utils.create_quantization_matrix(psr.toas[mask])[0]) nepoch = sum(U.shape[1] for U in Umats) U = np.zeros((len(psr.toas), nepoch)) jvec = np.zeros(nepoch) netot = 0 for ct, flag in enumerate(np.unique(bflags)): mask = bflags == flag nn = Umats[ct].shape[1] U[mask, netot:nn + netot] = Umats[ct] jvec[netot:nn + netot] = 10**(2 * ecorrs[ii][ct]) netot += nn # get covariance matrix cov = np.diag(nvec0) + np.dot(U * jvec[None, :], U.T) cf = sl.cho_factor(cov) logdet = np.sum(2 * np.log(np.diag(cf[0]))) cfs.append(cf) logdets.append(logdet) F, f2 = utils.createfourierdesignmatrix_red(psr.toas, nmodes=20, Tspan=Tspan) Mmat = psr.Mmat.copy() norm = np.sqrt(np.sum(Mmat**2, axis=0)) Mmat /= norm U2, avetoas = create_quant_matrix(psr.toas, dt=7 * 86400) if ik: T = np.hstack((F, Mmat, U2)) else: T = np.hstack((F, Mmat)) Ts.append(T) phi = utils.powerlaw(f2, log10_A=log10_A[ii], gamma=gamma[ii]) if inc_corr: phigw = utils.powerlaw(f2, log10_A=GW_log10_A, gamma=GW_gamma) else: phigw = np.zeros(40) K = se_kernel(avetoas, log10_sigma=log10_sigmas[ii], log10_lam=log10_lams[ii]) k = np.diag( np.concatenate((phi + phigw, np.ones(Mmat.shape[1]) * 1e40))) if ik: k = sl.block_diag(k, K) phis.append(k) # manually compute loglike loglike = 0 TNrs, TNTs = [], [] for ct, psr in enumerate(psrs): TNrs.append(np.dot(Ts[ct].T, sl.cho_solve(cfs[ct], psr.residuals))) TNTs.append(np.dot(Ts[ct].T, sl.cho_solve(cfs[ct], Ts[ct]))) loglike += -0.5 * ( np.dot(psr.residuals, sl.cho_solve(cfs[ct], psr.residuals)) + logdets[ct]) TNr = np.concatenate(TNrs) phi = sl.block_diag(*phis) if inc_corr: hd = utils.hd_orf(psrs[0].pos, psrs[1].pos) phi[len(phis[0]):len(phis[0]) + 40, :40] = np.diag(phigw * hd) phi[:40, len(phis[0]):len(phis[0]) + 40] = np.diag(phigw * hd) cf = sl.cho_factor(phi) phiinv = sl.cho_solve(cf, np.eye(phi.shape[0])) logdetphi = np.sum(2 * np.log(np.diag(cf[0]))) Sigma = sl.block_diag(*TNTs) + phiinv cf = sl.cho_factor(Sigma) expval = sl.cho_solve(cf, TNr) logdetsigma = np.sum(2 * np.log(np.diag(cf[0]))) loglike -= 0.5 * (logdetphi + logdetsigma) loglike += 0.5 * np.dot(TNr, expval) method = ["partition", "sparse", "cliques"] for mth in method: eloglike = pta.get_lnlikelihood(params, phiinv_method=mth) msg = "Incorrect like for npsr={}, phiinv={}".format(npsrs, mth) assert np.allclose(eloglike, loglike), msg
def test_add_efac_equad_backend(self): """Test that addition of efac-backend and equad-backend signal returns correct covariance. """ selection = Selection(selections.by_backend) efac = parameter.Uniform(0.1, 5) equad = parameter.Uniform(-10, -5) ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) s = ef + eq m = s(self.psr) # set parameters efacs = [1.3, 1.4, 1.5, 1.6] equads = [-6.1, -6.2, -6.3, -6.4] params = { "B1855+09_430_ASP_efac": efacs[0], "B1855+09_430_PUPPI_efac": efacs[1], "B1855+09_L-wide_ASP_efac": efacs[2], "B1855+09_L-wide_PUPPI_efac": efacs[3], "B1855+09_430_ASP_log10_equad": equads[0], "B1855+09_430_PUPPI_log10_equad": equads[1], "B1855+09_L-wide_ASP_log10_equad": equads[2], "B1855+09_L-wide_PUPPI_log10_equad": equads[3], } # correct value flags = ["430_ASP", "430_PUPPI", "L-wide_ASP", "L-wide_PUPPI"] nvec0 = np.zeros_like(self.psr.toas) for ct, flag in enumerate(np.unique(flags)): ind = flag == self.psr.backend_flags nvec0[ind] = efacs[ct]**2 * self.psr.toaerrs[ind]**2 nvec0[ind] += 10**(2 * equads[ct]) * np.ones(np.sum(ind)) logdet = np.sum(np.log(nvec0)) # test msg = "EFAC/EQUAD covariance incorrect." assert np.all(m.get_ndiag(params) == nvec0), msg msg = "EFAC/EQUAD logdet incorrect." N = m.get_ndiag(params) assert np.allclose(N.solve(self.psr.residuals, logdet=True)[1], logdet, rtol=1e-10), msg msg = "EFAC/EQUAD D1 solve incorrect." assert np.allclose(N.solve(self.psr.residuals), self.psr.residuals / nvec0, rtol=1e-10), msg msg = "EFAC/EQUAD 1D1 solve incorrect." assert np.allclose( N.solve(self.psr.residuals, left_array=self.psr.residuals), np.dot(self.psr.residuals / nvec0, self.psr.residuals), rtol=1e-10, ), msg msg = "EFAC/EQUAD 2D1 solve incorrect." T = self.psr.Mmat assert np.allclose(N.solve(self.psr.residuals, left_array=T), np.dot(T.T, self.psr.residuals / nvec0), rtol=1e-10), msg msg = "EFAC/EQUAD 2D2 solve incorrect." assert np.allclose(N.solve(T, left_array=T), np.dot(T.T, T / nvec0[:, None]), rtol=1e-10), msg
def test_pta(self): # get parameters from PAL2 style noise files params = get_noise_from_pal2(datadir + '/B1855+09_noise.txt') params2 = get_noise_from_pal2(datadir + '/J1909-3744_noise.txt') params.update(params2) # setup basic model efac = parameter.Constant() equad = parameter.Constant() ecorr = parameter.Constant() log10_A = parameter.Constant() gamma = parameter.Constant() selection = Selection(selections.by_backend) ef = white_signals.MeasurementNoise(efac=efac, selection=selection) eq = white_signals.EquadNoise(log10_equad=equad, selection=selection) ec = white_signals.EcorrKernelNoise(log10_ecorr=ecorr, selection=selection) pl = utils.powerlaw(log10_A=log10_A, gamma=gamma) rn = gp_signals.FourierBasisGP(pl) s = ef + eq + ec + rn pta = s(self.psrs[0]) + s(self.psrs[1]) # set parameters pta.set_default_params(params) # get parameters efacs, equads, ecorrs, log10_A, gamma = [], [], [], [], [] for pname in [p.name for p in self.psrs]: efacs.append([ params[key] for key in sorted(params.keys()) if 'efac' in key and pname in key ]) equads.append([ params[key] for key in sorted(params.keys()) if 'equad' in key and pname in key ]) ecorrs.append([ params[key] for key in sorted(params.keys()) if 'ecorr' in key and pname in key ]) log10_A.append(params['{}_red_noise_log10_A'.format(pname)]) gamma.append(params['{}_red_noise_gamma'.format(pname)]) # correct value tflags = [sorted(list(np.unique(p.backend_flags))) for p in self.psrs] cfs, logdets, phis = [], [], [] for ii, (psr, flags) in enumerate(zip(self.psrs, tflags)): nvec0 = np.zeros_like(psr.toas) for ct, flag in enumerate(flags): ind = psr.backend_flags == flag nvec0[ind] = efacs[ii][ct]**2 * psr.toaerrs[ind]**2 nvec0[ind] += 10**(2 * equads[ii][ct]) * np.ones(np.sum(ind)) # get the basis bflags = psr.backend_flags Umats = [] for flag in np.unique(bflags): mask = bflags == flag Umats.append( utils.create_quantization_matrix(psr.toas[mask])[0]) nepoch = sum(U.shape[1] for U in Umats) U = np.zeros((len(psr.toas), nepoch)) jvec = np.zeros(nepoch) netot = 0 for ct, flag in enumerate(np.unique(bflags)): mask = bflags == flag nn = Umats[ct].shape[1] U[mask, netot:nn + netot] = Umats[ct] jvec[netot:nn + netot] = 10**(2 * ecorrs[ii][ct]) netot += nn # get covariance matrix cov = np.diag(nvec0) + np.dot(U * jvec[None, :], U.T) cf = sl.cho_factor(cov) logdet = np.sum(2 * np.log(np.diag(cf[0]))) cfs.append(cf) logdets.append(logdet) F, f2 = utils.createfourierdesignmatrix_red(psr.toas, nmodes=20) phi = utils.powerlaw(f2, log10_A=log10_A[ii], gamma=gamma[ii]) phis.append(phi) # tests Ns = pta.get_ndiag(params) pphis = pta.get_phi(params) pphiinvs = pta.get_phiinv(params) Ts = pta.get_basis(params) zipped = zip(logdets, cfs, phis, self.psrs, Ns, pphis, pphiinvs, Ts) for logdet, cf, phi, psr, N, pphi, pphiinv, T in zipped: msg = 'EFAC/ECORR logdet incorrect.' assert np.allclose(N.solve(psr.residuals, logdet=True)[1], logdet, rtol=1e-10), msg msg = 'EFAC/ECORR D1 solve incorrect.' assert np.allclose(N.solve(psr.residuals), sl.cho_solve(cf, psr.residuals), rtol=1e-10), msg msg = 'EFAC/ECORR 1D1 solve incorrect.' assert np.allclose(N.solve(psr.residuals, left_array=psr.residuals), np.dot(psr.residuals, sl.cho_solve(cf, psr.residuals)), rtol=1e-10), msg msg = 'EFAC/ECORR 2D1 solve incorrect.' assert np.allclose(N.solve(psr.residuals, left_array=T), np.dot(T.T, sl.cho_solve(cf, psr.residuals)), rtol=1e-10), msg msg = 'EFAC/ECORR 2D2 solve incorrect.' assert np.allclose(N.solve(T, left_array=T), np.dot(T.T, sl.cho_solve(cf, T)), rtol=1e-10), msg # spectrum test msg = 'Spectrum incorrect for GP Fourier signal.' assert np.all(pphi == phi), msg # inverse spectrum test msg = 'Spectrum inverse incorrect for GP Fourier signal.' assert np.all(pphiinv == 1 / phi), msg