diff --git a/enterprise_extensions/blocks.py b/enterprise_extensions/blocks.py index 1cb67978..796ecdf9 100644 --- a/enterprise_extensions/blocks.py +++ b/enterprise_extensions/blocks.py @@ -334,7 +334,7 @@ def bwm_sglpsr_block(Tmin, Tmax, amp_prior='log-uniform', def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', prior='log-uniform', dt=15, df=200, Tspan=None, components=30, - gamma_val=None, coefficients=False): + gamma_val=None, coefficients=False, vary=True): """ Returns DM noise model: @@ -357,12 +357,16 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', :param gamma_val: If given, this is the fixed slope of the power-law for powerlaw, turnover, or tprocess DM-variations + :param vary: + Whether to vary the parameters or use constant values. """ # dm noise parameters that are common if gp_kernel == 'diag': if psd in ['powerlaw', 'turnover', 'tprocess', 'tprocess_adapt']: # parameters shared by PSD functions - if prior == 'uniform': + if not vary: + log10_A_dm = parameter.Constant() + elif prior == 'uniform': log10_A_dm = parameter.LinearExp(-20, -11) elif prior == 'log-uniform' and gamma_val is not None: if np.abs(gamma_val - 4.33) < 0.1: @@ -374,6 +378,8 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', if gamma_val is not None: gamma_dm = parameter.Constant(gamma_val) + elif not vary: + gamma_dm = parameter.Constant() else: gamma_dm = parameter.Uniform(0, 7) @@ -381,26 +387,40 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', if psd == 'powerlaw': dm_prior = utils.powerlaw(log10_A=log10_A_dm, gamma=gamma_dm) elif psd == 'turnover': - kappa_dm = parameter.Uniform(0, 7) - lf0_dm = parameter.Uniform(-9, -7) + if vary: + kappa_dm = parameter.Uniform(0, 7) + lf0_dm = parameter.Uniform(-9, -7) + else: + kappa_dm = parameter.Constant() + lf0_dm = parameter.Constant() dm_prior = utils.turnover(log10_A=log10_A_dm, gamma=gamma_dm, lf0=lf0_dm, kappa=kappa_dm) elif psd == 'tprocess': - df = 2 - alphas_dm = gpp.InvGamma(df/2, df/2, size=components) + if vary: + df = 2 + alphas_dm = gpp.InvGamma(df/2, df/2, size=components) + else: + alphas_dm = parameter.Constant() dm_prior = gpp.t_process(log10_A=log10_A_dm, gamma=gamma_dm, alphas=alphas_dm) elif psd == 'tprocess_adapt': - df = 2 - alpha_adapt_dm = gpp.InvGamma(df/2, df/2, size=1) - nfreq_dm = parameter.Uniform(-0.5, 10-0.5) + if vary: + df = 2 + alpha_adapt_dm = gpp.InvGamma(df/2, df/2, size=1) + nfreq_dm = parameter.Uniform(-0.5, 10-0.5) + else: + alpha_adapt_dm = gpp.Constant() + nfreq_dm = parameter.Constant() + dm_prior = gpp.t_process_adapt(log10_A=log10_A_dm, gamma=gamma_dm, alphas_adapt=alpha_adapt_dm, nfreq=nfreq_dm) if psd == 'spectrum': - if prior == 'uniform': + if not vary: + log10_rho_dm = parameter.Constant() + elif prior == 'uniform': log10_rho_dm = parameter.LinearExp(-10, -4, size=components) elif prior == 'log-uniform': log10_rho_dm = parameter.Uniform(-10, -4, size=components) @@ -413,10 +433,16 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', elif gp_kernel == 'nondiag': if nondiag_kernel == 'periodic': # Periodic GP kernel for DM - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_p = parameter.Uniform(-4, 1) - log10_gam_p = parameter.Uniform(-3, 2) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + log10_p = parameter.Uniform(-2, 3) + log10_gam_p = parameter.Uniform(-4, 4) + else: + log10_sigma = parameter.Constant() + log10_ell = parameter.Constant() + log10_p = parameter.Constant() + log10_gam_p = parameter.Constant() dm_basis = gpk.linear_interp_basis_dm(dt=dt*const.day) dm_prior = gpk.periodic_kernel(log10_sigma=log10_sigma, @@ -425,12 +451,20 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', log10_p=log10_p) elif nondiag_kernel == 'periodic_rfband': # Periodic GP kernel for DM with RQ radio-frequency dependence - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_ell2 = parameter.Uniform(2, 7) - log10_alpha_wgt = parameter.Uniform(-4, 1) - log10_p = parameter.Uniform(-4, 1) - log10_gam_p = parameter.Uniform(-3, 2) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + log10_ell2 = parameter.Uniform(0, 7) + log10_alpha_wgt = parameter.Uniform(-4, 3) + log10_p = parameter.Uniform(-2, 3) + log10_gam_p = parameter.Uniform(-4, 4) + else: + log10_sigma = parameter.Constant() + log10_ell = parameter.Constant() + log10_ell2 = parameter.Constant() + log10_alpha_wgt = parameter.Constant() + log10_p = parameter.Constant() + log10_gam_p = parameter.Constant() dm_basis = gpk.get_tf_quantization_matrix(df=df, dt=dt*const.day, dm=True) @@ -441,18 +475,28 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', log10_ell2=log10_ell2) elif nondiag_kernel == 'sq_exp': # squared-exponential GP kernel for DM - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + else: + log10_sigma = parameter.Constant() + log10_ell = parameter.Constant() dm_basis = gpk.linear_interp_basis_dm(dt=dt*const.day) dm_prior = gpk.se_dm_kernel(log10_sigma=log10_sigma, log10_ell=log10_ell) elif nondiag_kernel == 'sq_exp_rfband': # Sq-Exp GP kernel for DM with RQ radio-frequency dependence - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_ell2 = parameter.Uniform(2, 7) - log10_alpha_wgt = parameter.Uniform(-4, 1) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + log10_ell2 = parameter.Uniform(0, 7) + log10_alpha_wgt = parameter.Uniform(-4, 3) + else: + log10_sigma = parameter.Consant() + log10_ell = parameter.Consant() + log10_ell2 = parameter.Consant() + log10_alpha_wgt = parameter.Consant() dm_basis = gpk.get_tf_quantization_matrix(df=df, dt=dt*const.day, dm=True) @@ -462,7 +506,10 @@ def dm_noise_block(gp_kernel='diag', psd='powerlaw', nondiag_kernel='periodic', log10_ell2=log10_ell2) elif nondiag_kernel == 'dmx_like': # DMX-like signal - log10_sigma = parameter.Uniform(-10, -4) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + else: + log10_sigma = parameter.Constant() dm_basis = gpk.linear_interp_basis_dm(dt=dt*const.day) dm_prior = gpk.dmx_ridge_prior(log10_sigma=log10_sigma) @@ -478,7 +525,7 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw', prior='log-uniform', dt=15, df=200, idx=4, include_quadratic=False, Tspan=None, name='chrom', components=30, - coefficients=False): + coefficients=False, vary=True): """ Returns GP chromatic noise model : @@ -490,7 +537,7 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw', Whether to use a diagonal kernel for the GP. ['diag','nondiag'] :param nondiag_kernel: Which nondiagonal kernel to use for the GP. - ['periodic','sq_exp','periodic_rfband','sq_exp_rfband'] + ['periodic','sq_exp','dmx_like'] :param psd: PSD to use for common red noise signal. Available options are ['powerlaw', 'turnover' 'spectrum'] @@ -501,7 +548,8 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw', :param df: frequency-scale for linear interpolation basis (MHz) :param idx: - Index of radio frequency dependence (i.e. DM is 2). Any float will work. + Index of radio frequency dependence (i.e. DM is 2). Any float will + work or 'vary' will vary between [2.5,5] :param include_quadratic: Whether to include a quadratic fit. :param name: Name of signal @@ -511,29 +559,49 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw', Number of frequencies to use in 'diag' GPs. :param coefficients: Whether to keep coefficients of the GP. + :param vary: + Whether to vary the parameters or use constant values. """ + if idx == 'vary': + if not vary: + idx = parameter.Constant() + else: + idx = parameter.Uniform(2.5, 5) + if gp_kernel == 'diag': chm_basis = gpb.createfourierdesignmatrix_chromatic(nmodes=components, - Tspan=Tspan) + Tspan=Tspan, + idx=idx) if psd in ['powerlaw', 'turnover']: - if prior == 'uniform': + if not vary: + log10_A = parameter.Constant() + elif prior == 'uniform': log10_A = parameter.LinearExp(-18, -11) elif prior == 'log-uniform': log10_A = parameter.Uniform(-18, -11) - gamma = parameter.Uniform(0, 7) + if vary: + gamma = parameter.Uniform(0, 7) + else: + gamma = parameter.Constant() # PSD if psd == 'powerlaw': chm_prior = utils.powerlaw(log10_A=log10_A, gamma=gamma) elif psd == 'turnover': - kappa = parameter.Uniform(0, 7) - lf0 = parameter.Uniform(-9, -7) + if vary: + kappa = parameter.Uniform(0, 7) + lf0 = parameter.Uniform(-9, -7) + else: + kappa = parameter.Constant() + lf0 = parameter.Constant() chm_prior = utils.turnover(log10_A=log10_A, gamma=gamma, lf0=lf0, kappa=kappa) if psd == 'spectrum': - if prior == 'uniform': + if not vary: + log10_rho = parameter.Constant() + elif prior == 'uniform': log10_rho = parameter.LinearExp(-10, -4, size=components) elif prior == 'log-uniform': log10_rho = parameter.Uniform(-10, -4, size=components) @@ -542,56 +610,46 @@ def chromatic_noise_block(gp_kernel='nondiag', psd='powerlaw', elif gp_kernel == 'nondiag': if nondiag_kernel == 'periodic': # Periodic GP kernel for DM - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_p = parameter.Uniform(-4, 1) - log10_gam_p = parameter.Uniform(-3, 2) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + log10_p = parameter.Uniform(-2, 3) + log10_gam_p = parameter.Uniform(-4, 4) + else: + log10_sigma = parameter.Constant() + log10_ell = parameter.Constant() + log10_p = parameter.Constant() + log10_gam_p = parameter.Constant() - chm_basis = gpk.linear_interp_basis_chromatic(dt=dt*const.day) + chm_basis = gpk.linear_interp_basis_chromatic(dt=dt*const.day,idx=idx) chm_prior = gpk.periodic_kernel(log10_sigma=log10_sigma, log10_ell=log10_ell, log10_gam_p=log10_gam_p, log10_p=log10_p) - elif nondiag_kernel == 'periodic_rfband': - # Periodic GP kernel for DM with RQ radio-frequency dependence - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_ell2 = parameter.Uniform(2, 7) - log10_alpha_wgt = parameter.Uniform(-4, 1) - log10_p = parameter.Uniform(-4, 1) - log10_gam_p = parameter.Uniform(-3, 2) - - chm_basis = gpk.get_tf_quantization_matrix(df=df, dt=dt*const.day, - dm=True, dm_idx=idx) - chm_prior = gpk.tf_kernel(log10_sigma=log10_sigma, - log10_ell=log10_ell, - log10_gam_p=log10_gam_p, - log10_p=log10_p, - log10_alpha_wgt=log10_alpha_wgt, - log10_ell2=log10_ell2) elif nondiag_kernel == 'sq_exp': # squared-exponential kernel for DM - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) + if vary: + log10_sigma = parameter.Uniform(-10, -4) + log10_ell = parameter.Uniform(0, 5) + else: + log10_sigma = parameter.Constant() + log10_ell = parameter.Constant() chm_basis = gpk.linear_interp_basis_chromatic(dt=dt*const.day, idx=idx) chm_prior = gpk.se_dm_kernel(log10_sigma=log10_sigma, log10_ell=log10_ell) - elif nondiag_kernel == 'sq_exp_rfband': - # Sq-Exp GP kernel for Chrom with RQ radio-frequency dependence - log10_sigma = parameter.Uniform(-10, -4) - log10_ell = parameter.Uniform(1, 4) - log10_ell2 = parameter.Uniform(2, 7) - log10_alpha_wgt = parameter.Uniform(-4, 1) - - chm_basis = gpk.get_tf_quantization_matrix(df=df, dt=dt*const.day, - dm=True, dm_idx=idx) - chm_prior = gpk.sf_kernel(log10_sigma=log10_sigma, - log10_ell=log10_ell, - log10_alpha_wgt=log10_alpha_wgt, - log10_ell2=log10_ell2) + + elif nondiag_kernel == 'dmx_like': + # DMX-like signal + if vary: + log10_sigma = parameter.Uniform(-10, -4) + else: + log10_sigma = parameter.Constant() + + chm_basis = gpk.linear_interp_basis_chromatic(dt=dt*const.day, idx=idx) + chm_prior = gpk.dmx_ridge_prior(log10_sigma=log10_sigma) cgp = gp_signals.BasisGP(chm_prior, chm_basis, name=name+'_gp', coefficients=coefficients) diff --git a/enterprise_extensions/chromatic/chromatic.py b/enterprise_extensions/chromatic/chromatic.py index 2da45a55..defcdafb 100644 --- a/enterprise_extensions/chromatic/chromatic.py +++ b/enterprise_extensions/chromatic/chromatic.py @@ -204,7 +204,7 @@ def dmx_delay(toas, freqs, dmx_ids, **kwargs): return wf -def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'): +def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp', vary=True): """ Returns chromatic exponential dip (i.e. TOA advance): @@ -216,15 +216,24 @@ def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'): :param sign: set sign of dip: 'positive', 'negative', or 'vary' :param name: Name of signal + :param vary: Whether to vary the parameters or use constant values. :return dmexp: chromatic exponential dip waveform. """ - t0_dmexp = parameter.Uniform(tmin, tmax) - log10_Amp_dmexp = parameter.Uniform(-10, -2) - log10_tau_dmexp = parameter.Uniform(0, 2.5) - if sign == 'vary': + if vary: + t0_dmexp = parameter.Uniform(tmin, tmax) + log10_Amp_dmexp = parameter.Uniform(-10, -2) + log10_tau_dmexp = parameter.Uniform(0, 2.5) + else: + t0_dmexp = parameter.Constant() + log10_Amp_dmexp = parameter.Constant() + log10_tau_dmexp = parameter.Constant() + + if sign == 'vary' and vary: sign_param = parameter.Uniform(-1.0, 1.0) + elif sign == 'vary' and not vary: + sign_param = parameter.Constant() elif sign == 'positive': sign_param = 1.0 else: @@ -238,7 +247,7 @@ def dm_exponential_dip(tmin, tmax, idx=2, sign='negative', name='dmexp'): def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative', - symmetric=False, name='dm_cusp'): + symmetric=False, name='dm_cusp', vary=True): """ Returns chromatic exponential cusp (i.e. TOA advance): @@ -250,16 +259,24 @@ def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative', :param sign: set sign of dip: 'positive', 'negative', or 'vary' :param name: Name of signal + :param vary: Whether to vary the parameters or use constant values. :return dmexp: chromatic exponential dip waveform. """ - t0_dm_cusp = parameter.Uniform(tmin, tmax) - log10_Amp_dm_cusp = parameter.Uniform(-10, -2) - log10_tau_dm_cusp_pre = parameter.Uniform(0, 2.5) + if vary: + t0_dm_cusp = parameter.Uniform(tmin, tmax) + log10_Amp_dm_cusp = parameter.Uniform(-10, -2) + log10_tau_dm_cusp_pre = parameter.Uniform(0, 2.5) + else: + t0_dm_cusp = parameter.Constant() + log10_Amp_dm_cusp = parameter.Constant() + log10_tau_dm_cusp_pre = parameter.Constant() - if sign == 'vary': + if sign == 'vary' and vary: sign_param = parameter.Uniform(-1.0, 1.0) + elif sign == 'vary' and not vary: + sign_param = parameter.Constant() elif sign == 'positive': sign_param = 1.0 else: @@ -267,8 +284,10 @@ def dm_exponential_cusp(tmin, tmax, idx=2, sign='negative', if symmetric: log10_tau_dm_cusp_post = 1 - else: + elif vary: log10_tau_dm_cusp_post = parameter.Uniform(0, 2.5) + else: + log10_tau_dm_cusp_post = parameter.Constant() wf = chrom_exp_cusp(log10_Amp=log10_Amp_dm_cusp, sign_param=sign_param, t0=t0_dm_cusp, log10_tau_pre=log10_tau_dm_cusp_pre, @@ -292,18 +311,28 @@ def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative', :param sign: set sign of dip: 'positive', 'negative', or 'vary' :param name: Name of signal + :param vary: Whether to vary the parameters or use constant values. :return dmexp: chromatic exponential dip waveform. """ - t0_dual_cusp = parameter.Uniform(tmin, tmax) - log10_Amp_dual_cusp_1 = parameter.Uniform(-10, -2) - log10_Amp_dual_cusp_2 = parameter.Uniform(-10, -2) - log10_tau_dual_cusp_pre_1 = parameter.Uniform(0, 2.5) - log10_tau_dual_cusp_pre_2 = parameter.Uniform(0, 2.5) + if vary: + t0_dual_cusp = parameter.Uniform(tmin, tmax) + log10_Amp_dual_cusp_1 = parameter.Uniform(-10, -2) + log10_Amp_dual_cusp_2 = parameter.Uniform(-10, -2) + log10_tau_dual_cusp_pre_1 = parameter.Uniform(0, 2.5) + log10_tau_dual_cusp_pre_2 = parameter.Uniform(0, 2.5) + else: + t0_dual_cusp = parameter.Constant() + log10_Amp_dual_cusp_1 = parameter.Constant() + log10_Amp_dual_cusp_2 = parameter.Constant() + log10_tau_dual_cusp_pre_1 = parameter.Constant() + log10_tau_dual_cusp_pre_2 = parameter.Constant() - if sign == 'vary': + if sign == 'vary' and vary: sign_param = parameter.Uniform(-1.0, 1.0) + elif sign == 'vary' and not vary: + sign_param = parameter.Constant() elif sign == 'positive': sign_param = 1.0 else: @@ -312,9 +341,12 @@ def dm_dual_exp_cusp(tmin, tmax, idx1=2, idx2=4, sign='negative', if symmetric: log10_tau_dual_cusp_post_1 = 1 log10_tau_dual_cusp_post_2 = 1 - else: + elif vary: log10_tau_dual_cusp_post_1 = parameter.Uniform(0, 2.5) log10_tau_dual_cusp_post_2 = parameter.Uniform(0, 2.5) + else: + log10_tau_dual_cusp_post_1 = parameter.Constant() + log10_tau_dual_cusp_post_2 = parameter.Constant() wf = chrom_dual_exp_cusp(t0=t0_dual_cusp, sign_param=sign_param, symmetric=symmetric, @@ -336,22 +368,29 @@ def dmx_signal(dmx_data, name='dmx_signal'): :param dmx_data: dictionary of DMX data for each pulsar from parfile. :param name: Name of signal. + :param vary: Whether to vary the parameters or use constant values. :return dmx_sig: dmx signal waveform. """ dmx = {} - for dmx_id in sorted(dmx_data): - dmx_data_tmp = dmx_data[dmx_id] - dmx.update({dmx_id: parameter.Normal(mu=dmx_data_tmp['DMX_VAL'], - sigma=dmx_data_tmp['DMX_ERR'])}) + if vary: + for dmx_id in sorted(dmx_data): + dmx_data_tmp = dmx_data[dmx_id] + dmx.update({dmx_id: parameter.Normal(mu=dmx_data_tmp['DMX_VAL'], + sigma=dmx_data_tmp['DMX_ERR'])}) + else: + for dmx_id in sorted(dmx_data): + dmx_data_tmp = dmx_data[dmx_id] + dmx.update({dmx_id: parameter.Constant()}) + wf = dmx_delay(dmx_ids=dmx_data, **dmx) dmx_sig = deterministic_signals.Deterministic(wf, name=name) return dmx_sig -def dm_annual_signal(idx=2, name='dm_s1yr'): +def dm_annual_signal(idx=2, name='dm_s1yr', vary=True): """ Returns chromatic annual signal (i.e. TOA advance): @@ -359,12 +398,17 @@ def dm_annual_signal(idx=2, name='dm_s1yr'): index of radio frequency dependence (i.e. DM is 2). If this is set to 'vary' then the index will vary from 1 - 6 :param name: Name of signal + :param vary: Whether to vary the parameters or use constant values. :return dm1yr: chromatic annual waveform. """ - log10_Amp_dm1yr = parameter.Uniform(-10, -2) - phase_dm1yr = parameter.Uniform(0, 2*np.pi) + if vary: + log10_Amp_dm1yr = parameter.Uniform(-10, -2) + phase_dm1yr = parameter.Uniform(0, 2*np.pi) + else: + log10_Amp_dm1yr = parameter.Constant() + phase_dm1yr = parameter.Constant() wf = chrom_yearly_sinusoid(log10_Amp=log10_Amp_dm1yr, phase=phase_dm1yr, idx=idx) diff --git a/enterprise_extensions/hypermodel.py b/enterprise_extensions/hypermodel.py index 1b5ed036..74a669a0 100644 --- a/enterprise_extensions/hypermodel.py +++ b/enterprise_extensions/hypermodel.py @@ -229,45 +229,41 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, sampler.addProposalToCycle(jp.draw_from_red_prior, 10) # DM GP noise prior draw - if 'dm_gp' in self.snames: + if 'dm_gp' in self.snames and len(jp.snames['dm_gp'])!=0: print('Adding DM GP noise prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dm_gp_prior, 10) # DM annual prior draw - if 'dm_s1yr' in jp.snames: + if 'dm_s1yr' in jp.snames and len(jp.snames['dm_s1yr'])!=0: print('Adding DM annual prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dm1yr_prior, 10) # DM dip prior draw - if 'dmexp' in '\t'.join(jp.snames): + dmexp_nm = [nm for nm in jp.snames if 'dmexp' in nm] + if len(dmexp_nm)!=0 and all([len(jp.snames[nm])!=0 for nm in dmexp_nm]): print('Adding DM exponential dip prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmexpdip_prior, 10) # DM cusp prior draw - if 'dm_cusp' in jp.snames: + if 'dm_cusp' in jp.snames and len(jp.snames['dm_cusp'])!=0: print('Adding DM exponential cusp prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmexpcusp_prior, 10) # DMX prior draw - if 'dmx_signal' in jp.snames: + if 'dmx_signal' in jp.snames and len(jp.snames['dmx_signal'])!=0: print('Adding DMX prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmx_prior, 10) # Chromatic GP noise prior draw - if 'chrom_gp' in self.snames: + if 'chrom_gp' in self.snames and len(jp.snames['chrom_gp'])!=0: print('Adding Chromatic GP noise prior draws...\n') sampler.addProposalToCycle(jp.draw_from_chrom_gp_prior, 10) # SW prior draw - if 'gp_sw' in jp.snames: + if 'gp_sw' in jp.snames and len(jp.snames['gp_sw'])!=0: print('Adding Solar Wind DM GP prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dm_sw_prior, 10) - # Chromatic GP noise prior draw - if 'chrom_gp' in self.snames: - print('Adding Chromatic GP noise prior draws...\n') - sampler.addProposalToCycle(jp.draw_from_chrom_gp_prior, 10) - # Ephemeris prior draw if 'd_jupiter_mass' in self.param_names: print('Adding ephemeris model prior draws...\n') diff --git a/enterprise_extensions/models.py b/enterprise_extensions/models.py index 5e5e5268..22fb8df1 100644 --- a/enterprise_extensions/models.py +++ b/enterprise_extensions/models.py @@ -34,8 +34,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, dm_type='gp', dmgp_kernel='diag', dm_psd='powerlaw', dm_nondiag_kernel='periodic', dmx_data=None, dm_annual=False, gamma_dm_val=None, - dm_dt=15, dm_df=200, - chrom_gp=False, chrom_gp_kernel='nondiag', + dm_dt=15, dm_df=200, dm_Nfreqs=30, + chrom_Nfreqs=30, chrom_gp=False, chrom_gp_kernel='nondiag', chrom_psd='powerlaw', chrom_idx=4, chrom_quad=False, chrom_kernel='periodic', chrom_dt=15, chrom_df=200, @@ -58,7 +58,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, psr_model=False, factorized_like=False, Tspan=None, fact_like_gamma=13./3, gw_components=10, fact_like_logmin=None, fact_like_logmax=None, - select='backend', tm_marg=False, dense_like=False, ng_twg_setup=False, wb_efac_sigma=0.25): + select='backend', tm_marg=False, dense_like=False, ng_twg_setup=False, wb_efac_sigma=0.25, + vary_dm=True, vary_chrom=True): """ Single pulsar noise model. @@ -94,7 +95,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, :param chrom_gp_kernel: GP kernel type to use in chrom ['diag','nondiag'] :param chrom_psd: power-spectral density of chromatic noise ['powerlaw','tprocess','free_spectrum'] - :param chrom_idx: frequency scaling of chromatic noise + :param chrom_idx: frequency scaling of chromatic noise. use 'vary' to vary + between [2.5,5]. :param chrom_kernel: Type of 'nondiag' time-domain chrom GP kernel to use ['periodic', 'sq_exp','periodic_rfband', 'sq_exp_rfband'] :param chrom_quad: Whether to add a quadratic chromatic term. Boolean @@ -148,6 +150,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, :param tm_marg: Use marginalized timing model. In many cases this will speed up the likelihood calculation significantly. :param dense_like: Use dense or sparse functions to evalute lnlikelihood + :param vary_dm: Whether to vary the DM model or use constant values + :param vary_chrom: Whether to vary the chromatic model or use constant values :return s: single pulsar noise model @@ -220,26 +224,20 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, if dm_type == 'gp': if dmgp_kernel == 'diag': s += dm_noise_block(gp_kernel=dmgp_kernel, psd=dm_psd, - prior=amp_prior, components=components, - gamma_val=gamma_dm_val, - coefficients=coefficients) + prior=amp_prior, components=dm_Nfreqs, + gamma_val=gamma_dm_val, Tspan=Tspan, + coefficients=coefficients, + vary=vary_dm) elif dmgp_kernel == 'nondiag': s += dm_noise_block(gp_kernel=dmgp_kernel, nondiag_kernel=dm_nondiag_kernel, dt=dm_dt, df=dm_df, - coefficients=coefficients) + coefficients=coefficients, + vary=vary_dm) elif dm_type == 'dmx': - s += chrom.dmx_signal(dmx_data=dmx_data[psr.name]) + s += chrom.dmx_signal(dmx_data=dmx_data[psr.name],vary=vary_dm) if dm_annual: - s += chrom.dm_annual_signal() - if chrom_gp: - s += chromatic_noise_block(gp_kernel=chrom_gp_kernel, - psd=chrom_psd, idx=chrom_idx, - components=components, - nondiag_kernel=chrom_kernel, - dt=chrom_dt, df=chrom_df, - include_quadratic=chrom_quad, - coefficients=coefficients) + s += chrom.dm_annual_signal(vary=vary_dm) if dm_expdip: if dm_expdip_tmin is None and dm_expdip_tmax is None: @@ -264,7 +262,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, s += chrom.dm_exponential_dip(tmin=tmin[dd], tmax=tmax[dd], idx=dm_expdip_idx[dd], sign=dmexp_sign, - name=dmdipname_base[dd]) + name=dmdipname_base[dd], + vary=vary_dm) if dm_cusp: if dm_cusp_tmin is None and dm_cusp_tmax is None: tmin = [psr.toas.min() / const.day for ii in range(num_dm_cusps)] @@ -288,7 +287,8 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, idx=dm_cusp_idx[dd-1], sign=dm_cusp_sign[dd-1], symmetric=dm_cusp_sym, - name=cusp_name_base+str(dd)) + name=cusp_name_base+str(dd), + vary=vary_dm) if dm_dual_cusp: if dm_dual_cusp_tmin is None and dm_cusp_tmax is None: tmin = psr.toas.min() / const.day @@ -306,13 +306,24 @@ def model_singlepsr_noise(psr, tm_var=False, tm_linear=False, idx2=dm_dual_cusp_idx2, sign=dm_dual_cusp_sign, symmetric=dm_dual_cusp_sym, - name=dual_cusp_name_base+str(dd)) + name=dual_cusp_name_base+str(dd), + vary=vary_dm) if dm_sw_deter: Tspan = psr.toas.max() - psr.toas.min() s += solar_wind_block(ACE_prior=True, include_swgp=dm_sw_gp, swgp_prior=swgp_prior, swgp_basis=swgp_basis, Tspan=Tspan) - + + if chrom_gp: + s += chromatic_noise_block(gp_kernel=chrom_gp_kernel, + psd=chrom_psd, idx=chrom_idx, + components=chrom_Nfreqs, + nondiag_kernel=chrom_kernel, + dt=chrom_dt, df=chrom_df, + include_quadratic=chrom_quad, + coefficients=coefficients, + Tspan=Tspan, + vary=vary_chrom) if extra_sigs is not None: s += extra_sigs diff --git a/enterprise_extensions/sampler.py b/enterprise_extensions/sampler.py index 700caece..1da68195 100644 --- a/enterprise_extensions/sampler.py +++ b/enterprise_extensions/sampler.py @@ -530,7 +530,9 @@ def draw_from_gwb_log_uniform_distribution(self, x, iter, beta): # draw parameter from signal model signal_name = [par for par in self.pnames if ('gw' in par and 'log10_A' in par)][0] - idx = list(self.pnames).index(signal_name) + + param_names = [par.name for par in self.params] + idx = list(param_names).index(signal_name) param = self.params[idx] q[self.pmap[str(param)]] = np.random.uniform(param.prior._defaults['pmin'], param.prior._defaults['pmax']) @@ -1160,28 +1162,33 @@ def setup_sampler(pta, outdir='chains', resume=False, print('Adding red noise prior draws...\n') sampler.addProposalToCycle(jp.draw_from_red_prior, 10) + # Chromatic GP noise prior draw + if 'chrom_gp' in jp.snames and len(jp.snames['chrom_gp'])!=0: + print('Adding Chromatic GP noise prior draws...\n') + sampler.addProposalToCycle(jp.draw_from_chrom_gp_prior, 10) + # DM GP noise prior draw - if 'dm_gp' in jp.snames: + if 'dm_gp' in jp.snames and len(jp.snames['dm_gp'])!=0: print('Adding DM GP noise prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dm_gp_prior, 10) # DM annual prior draw - if 'dm_s1yr' in jp.snames: + if 'dm_s1yr' in jp.snames and len(jp.snames['dm_s1yr'])!=0: print('Adding DM annual prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dm1yr_prior, 10) - + # DM dip prior draw - if 'dmexp' in jp.snames: + if 'dmexp' in jp.snames and len(jp.snames['dmexp'])!=0: print('Adding DM exponential dip prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmexpdip_prior, 10) # DM cusp prior draw - if 'dm_cusp' in jp.snames: + if 'dm_cusp' in jp.snames and len(jp.snames['dm_cusp'])!=0: print('Adding DM exponential cusp prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmexpcusp_prior, 10) # DMX prior draw - if 'dmx_signal' in jp.snames: + if 'dmx_signal' in jp.snames and len(jp.snames['dmx_signal'])!=0: print('Adding DMX prior draws...\n') sampler.addProposalToCycle(jp.draw_from_dmx_prior, 10)