Added flag to skip constraint normalization during sampling. This avoids...

This MR is a proposed solution to the mattermost discussion regarding intractably long initialization times with BilbyMCMC and BilbyPTMCMC when sampling in EoS model parameters and applying the standard EoS constraints to the standard parameter priors.

The long initialization times are a result of the method by which the constraint normalization factor is computed: samples are drawn from the prior after which constraints are applied. This process is repeated many times at which point the ratio of samples that passed the constraint requirements to the total number of samples are drawn is computed and used as the constraint normalization factor. The issue with this approach for EoS sampling is that EoS constraints require numerical integration and have an exceptionally low acceptance rate for samples drawn from the prior. As a result, the initialization process takes on the order of days and currently cannot be checkpointed. During sampling this is an unnecessary computation as the normalization factor, which appears in the numerator and denominator of Bayes' theorem, cancels in the MH ratio.

Therefore, this MR implements a flag that can be activated to disable computation of the constraint normalization factor, reenabling EoS parameter estimation when using the BilbyMCMC and BilbyPTMCMC samplers.

bilby/bilby_mcmc/sampler.py

@@ -127,6 +127,9 @@ class Bilby_MCMC(MCMCSampler):
     initial_sample_dict: dict
         A dictionary of the initial sample value. If incomplete, will overwrite
         the initial_sample drawn using initial_sample_method.
+    normalize_prior: bool
+        When False, disables calculation of constraint normalization factor
+        during prior probability computation. Default value is True.
     verbose: bool
         Whether to print diagnostic output during the run.
 
@@ -175,6 +178,7 @@ class Bilby_MCMC(MCMCSampler):
         resume=True,
         exit_code=130,
         verbose=True,
+        normalize_prior=True,
         **kwargs,
     ):
 
@@ -194,6 +198,7 @@ class Bilby_MCMC(MCMCSampler):
         self.kwargs["target_nsamples"] = self.kwargs["nsamples"]
         self.L1steps = self.kwargs["L1steps"]
         self.L2steps = self.kwargs["L2steps"]
+        self.normalize_prior = normalize_prior
         self.pt_inputs = ParallelTemperingInputs(
             **{key: self.kwargs[key] for key in ParallelTemperingInputs._fields}
         )
@@ -309,6 +314,7 @@ class Bilby_MCMC(MCMCSampler):
             evidence_method=self.evidence_method,
             initial_sample_method=self.initial_sample_method,
             initial_sample_dict=self.initial_sample_dict,
+            normalize_prior=self.normalize_prior,
         )
 
     def get_setup_string(self):
@@ -583,11 +589,13 @@ class BilbyPTMCMCSampler(object):
         evidence_method,
         initial_sample_method,
         initial_sample_dict,
+        normalize_prior=True,
     ):
         self.set_pt_inputs(pt_inputs)
         self.use_ratio = use_ratio
         self.initial_sample_method = initial_sample_method
         self.initial_sample_dict = initial_sample_dict
+        self.normalize_prior = normalize_prior
         self.setup_sampler_dictionary(convergence_inputs, proposal_cycle)
         self.set_convergence_inputs(convergence_inputs)
         self.pt_rejection_sample = pt_rejection_sample
@@ -658,6 +666,7 @@ class BilbyPTMCMCSampler(object):
                     use_ratio=self.use_ratio,
                     initial_sample_method=self.initial_sample_method,
                     initial_sample_dict=self.initial_sample_dict,
+                    normalize_prior=self.normalize_prior,
                 )
                 for Eindex in range(n)
             ]
@@ -1152,12 +1161,13 @@ class BilbyMCMCSampler(object):
         use_ratio=False,
         initial_sample_method="prior",
         initial_sample_dict=None,
+        normalize_prior=True,
     ):
         self.beta = beta
         self.Tindex = Tindex
         self.Eindex = Eindex
         self.use_ratio = use_ratio
-
+        self.normalize_prior = normalize_prior
         self.parameters = _sampling_convenience_dump.priors.non_fixed_keys
         self.ndim = len(self.parameters)
 
@@ -1232,7 +1242,10 @@ class BilbyMCMCSampler(object):
         return logl
 
     def log_prior(self, sample):
-        return _sampling_convenience_dump.priors.ln_prob(sample.parameter_only_dict)
+        return _sampling_convenience_dump.priors.ln_prob(
+            sample.parameter_only_dict,
+            normalized=self.normalize_prior,
+        )
 
     def accept_proposal(self, prop, proposal):
         self.chain.append(prop)