From f4cfa17659f8a7b874484c86fde070ac3a8ad525 Mon Sep 17 00:00:00 2001 From: Hazboun6 Date: Mon, 19 Aug 2024 18:03:42 -0700 Subject: [PATCH] feat(hypermodel.py): add nonlinear timing to hypermodels Co-authored-by: Andrew Kaiser --- enterprise_extensions/hypermodel.py | 117 ++++++++++++++++++++++++++-- 1 file changed, 111 insertions(+), 6 deletions(-) diff --git a/enterprise_extensions/hypermodel.py b/enterprise_extensions/hypermodel.py index 1b5ed036..c4e043f0 100644 --- a/enterprise_extensions/hypermodel.py +++ b/enterprise_extensions/hypermodel.py @@ -8,7 +8,7 @@ from enterprise import constants as const from PTMCMCSampler.PTMCMCSampler import PTSampler as ptmcmc -from .sampler import JumpProposal, get_parameter_groups, save_runtime_info +from .sampler import JumpProposal, get_parameter_groups, save_runtime_info, get_timing_groups, group_from_params class HyperModel(object): @@ -71,6 +71,27 @@ def __init__(self, models, log_weights=None): for q in uniq_params]].tolist() ######### + ######### + # Timing Model + self.tm_groups = [] + self.special_idxs = [] + for i, x in enumerate(self.params): + if "timing_model" in str(x): + self.tm_groups.append(i) + if "Uniform" in str(x): + pmin = float(str(x).split("Uniform")[-1].split("pmin=")[1].split(",")[0]) + pmax = float(str(x).split("Uniform")[-1].split("pmax=")[-1].split(")")[0]) + if pmin + pmax != 0.0: + self.special_idxs.append(i) + elif "BoundedNormal" in str(x): + pmin = float(str(x).split("BoundedNormal")[-1].split("[")[-1].split(",")[0]) + pmax = float(str(x).split("BoundedNormal")[-1].split("[")[-1].split(",")[1].split(']')[0]) + if pmin + pmax != 0.0: + self.special_idxs.append(i) + else: + self.special_idxs.append(i) + ######### + def get_lnlikelihood(self, x): # find model index variable @@ -115,6 +136,18 @@ def get_parameter_groups(self): unique_groups = [] for p in self.models.values(): groups = get_parameter_groups(p) + if self.tm_groups: + groups.extend(get_timing_groups(p)) + groups.append( + group_from_params( + p, + [ + x + for x in p.param_names + if any(y in x for y in ["timing_model", "ecorr"]) + ], + ) + ) # check for any duplicate groups # e.g. the GWB may have different indices in model 1 and model 2 for group in groups: @@ -127,12 +160,39 @@ def get_parameter_groups(self): unique_groups.extend([[len(self.param_names) - 1]]) return unique_groups - def initial_sample(self): + def initial_sample(self, tm_params_orig=None, tm_param_dict=None, zero_start=True): """ Draw an initial sample from within the hyper-model prior space. + :param tm_params_orig: dictionary of timing model parameter tuples, (val, err) + :param tm_param_dict: a nested dictionary of parameters to vary in the model and their user defined values and priors + :param zero_start: start all timing parameters at their parfile value (in tm_params_orig), or their refit values (tm_param_dict) """ - x0 = [np.array(p.sample()).ravel().tolist() for p in self.models[0].params] + if zero_start and tm_params_orig: + x0 = [] + for xx, p in enumerate(self.models[0].params): + if "timing" in p.name: + if "DMX" in p.name: + p_name = ("_").join(p.name.split("_")[-2:]) + else: + p_name = p.name.split("_")[-1] + if tm_params_orig[p_name][-1] == "normalized": + x0.append([np.double(0.0)]) + else: + if p_name in tm_param_dict.keys(): + x0.append([np.double(tm_param_dict[p_name]["prior_mu"])]) + else: + x0.append([np.double(tm_params_orig[p_name][0])]) + elif "dm_model" in p.name: + if "mu" in str(p): + x0.append([float(str(p).split("(")[1].split(",")[0].split("=")[-1])]) + else: + x0.append(np.array(p.sample()).ravel().tolist()) + else: + x0.append(np.array(p.sample()).ravel().tolist()) + else: + x0 = [np.array(p.sample()).ravel().tolist() for p in self.models[0].params] + uniq_params = [str(p) for p in self.models[0].params] for model in self.models.values(): @@ -161,8 +221,9 @@ def draw_from_nmodel_prior(self, x, iter, beta): return q, float(lqxy) def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, - empirical_distr=None, groups=None, human=None, - loglkwargs={}, logpkwargs={}): + empirical_distr=None, groups=None, timing=False, psr=None, human=None, + restrict_mass=True, + loglkwargs=None, logpkwargs=None): """ Sets up an instance of PTMCMC sampler. @@ -182,6 +243,12 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, draws from uniform distributions. """ + if loglkwargs is None: + loglkwargs = {} + + if logpkwargs is None: + logpkwargs = {} + # dimension of parameter space ndim = len(self.param_names) @@ -212,7 +279,9 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, save_runtime_info(self, sampler.outDir, human) # additional jump proposals - jp = JumpProposal(self, self.snames, empirical_distr=empirical_distr) + jp = JumpProposal(self, self.snames, empirical_distr=empirical_distr, + timing=timing, psr=psr, sampler=sampler, + restrict_mass=restrict_mass) sampler.jp = jp # always add draw from prior @@ -268,6 +337,16 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, 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: + 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') @@ -316,6 +395,32 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True, p in list(self.params) if 'gw' in str(p)]), 10) + # Non Linear Timing Draws + if "timing_model" in jp.snames: + print("Adding timing model jump proposal...\n") + sampler.addProposalToCycle(jp.draw_from_timing_model, 25) + if "timing_model" in jp.snames: + print("Adding timing model prior draw...\n") + sampler.addProposalToCycle(jp.draw_from_timing_model_prior, 25) + + # DM Model Draws + if "dm_model" in jp.snames and len(jp.snames["dm_model"]): + print("Adding dm model prior draw...\n") + sampler.addProposalToCycle(jp.draw_from_signal("dm_model"), 10) + + if timing: + if jp.restrict_mass: + # SCAM and AM Draws + # add SCAM + print("Adding SCAM Jump Proposal...\n") + sampler.addProposalToCycle(jp.covarianceJumpProposalSCAM, 20) + + # add AM + print("Adding AM Jump Proposal...\n") + sampler.addProposalToCycle(jp.covarianceJumpProposalAM, 20) + + # DE does not work well with restricting the pulsar mass + # Model index distribution draw if sample_nmodel: if 'nmodel' in self.param_names: