added kwargs for logL, logp in sampler (#170)

nanograv · Feb 10, 2022 · b5a2005 · b5a2005
1 parent 0983084
commit b5a2005
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 4 deletions.
diff --git a/enterprise_extensions/hypermodel.py b/enterprise_extensions/hypermodel.py
@@ -155,7 +155,8 @@ 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):
+                      empirical_distr=None, groups=None, human=None,
+                      loglkwargs={}, logpkwargs={}):
         """
         Sets up an instance of PTMCMC sampler.
 
@@ -189,7 +190,9 @@ def setup_sampler(self, outdir='chains', resume=False, sample_nmodel=True,
             groups = self.get_parameter_groups()
 
         sampler = ptmcmc(ndim, self.get_lnlikelihood, self.get_lnprior, cov,
-                         groups=groups, outDir=outdir, resume=resume)
+                         groups=groups, outDir=outdir, resume=resume,
+                         loglkwargs=loglkwargs, logpkwargs=logpkwargs)
+
         save_runtime_info(self, sampler.outDir, human)
 
         # additional jump proposals

diff --git a/enterprise_extensions/sampler.py b/enterprise_extensions/sampler.py
@@ -1022,7 +1022,8 @@ def save_runtime_info(pta, outdir='chains', human=None):
 
 
 def setup_sampler(pta, outdir='chains', resume=False,
-                  empirical_distr=None, groups=None, human=None, save_ext_dists=False):
+                  empirical_distr=None, groups=None, human=None,
+                  save_ext_dists=False, loglkwargs={}, logpkwargs={}):
     """
     Sets up an instance of PTMCMC sampler.
 
@@ -1062,7 +1063,8 @@ def setup_sampler(pta, outdir='chains', resume=False,
         groups = get_parameter_groups(pta)
 
     sampler = ptmcmc(ndim, pta.get_lnlikelihood, pta.get_lnprior, cov, groups=groups,
-                     outDir=outdir, resume=resume)
+                     outDir=outdir, resume=resume, loglkwargs=loglkwargs,
+                     logpkwargs=logpkwargs)
 
     save_runtime_info(pta, sampler.outDir, human)