From 2ccfd5df6dacf5459c496c6af0005b9d6da93e33 Mon Sep 17 00:00:00 2001
From: lukashergt <lthergt@phas.ubc.ca>
Date: Mon, 20 Nov 2023 15:03:05 -0800
Subject: [PATCH] pcfix from #233

---
 cobaya/samplers/polychord/polychord.py | 41 ++++++++++----------------
 1 file changed, 16 insertions(+), 25 deletions(-)

diff --git a/cobaya/samplers/polychord/polychord.py b/cobaya/samplers/polychord/polychord.py
index cb51438c4..7f6132428 100644
--- a/cobaya/samplers/polychord/polychord.py
+++ b/cobaya/samplers/polychord/polychord.py
@@ -165,23 +165,6 @@ def initialize(self):
         self.pc_settings = settings.PolyChordSettings(
             self.nDims, self.nDerived, seed=(self.seed if self.seed is not None else -1),
             **{p: getattr(self, p) for p in pc_args if getattr(self, p) is not None})
-        # prior conversion from the hypercube
-        bounds = self.model.prior.bounds(
-            confidence_for_unbounded=self.confidence_for_unbounded)
-        # Check if priors are bounded (nan's to inf)
-        inf = np.where(np.isinf(bounds))
-        if len(inf[0]):
-            params_names = list(self.model.parameterization.sampled_params())
-            params = [params_names[i] for i in sorted(set(inf[0]))]
-            raise LoggedError(
-                self.log, "PolyChord needs bounded priors, but the parameter(s) '"
-                          "', '".join(params) + "' is(are) unbounded.")
-        locs = bounds[:, 0]
-        scales = bounds[:, 1] - bounds[:, 0]
-        # This function re-scales the parameters AND puts them in the right order
-        self.pc_prior = lambda x: (locs + np.array(x)[self.ordering] * scales).tolist()
-        # We will need the volume of the prior domain, since PolyChord divides by it
-        self.logvolume = np.log(np.prod(scales))
         # Prepare callback function
         if self.callback_function is not None:
             self.callback_function_callable = (
@@ -245,11 +228,9 @@ def run(self):
         Prepares the prior and likelihood functions, calls ``PolyChord``'s ``run``, and
         processes its output.
         """
-        # Prepare the posterior
-        # Don't forget to multiply by the volume of the physical hypercube,
-        # since PolyChord divides by it
 
-        def logpost(params_values):
+        # Prepare the PolyChord likelihood
+        def loglikelihood(params_values):
             result = self.model.logposterior(params_values)
             loglikes = result.loglikes
             if len(loglikes) != self.n_likes:
@@ -258,13 +239,23 @@ def logpost(params_values):
             if len(derived) != self.n_derived:
                 derived = np.full(self.n_derived, np.nan)
             derived = list(derived) + list(result.logpriors) + list(loglikes)
-            return (max(result.logpost + self.logvolume, self.pc_settings.logzero),
-                    derived)
+            if -np.inf in result.logpriors:
+                return self.pc_settings.logzero, derived
+            else:
+                return max(loglikes.sum(), self.pc_settings.logzero), derived
+
+        def prior(cube):
+            theta = np.empty_like(cube)
+            for i, xi in enumerate(np.array(cube)[self.ordering]):
+                theta[i] = self.model.prior.pdf[i].ppf(xi)
+            return theta
 
+        if is_main_process():
+            self.dump_paramnames(self.raw_prefix)
         sync_processes()
         self.mpi_info("Calling PolyChord...")
-        self.pc.run_polychord(logpost, self.nDims, self.nDerived, self.pc_settings,
-                              self.pc_prior, self.dumper)
+        self.pc.run_polychord(loglikelihood, self.nDims, self.nDerived, self.pc_settings,
+                              prior, self.dumper)
         self.process_raw_output()
 
     @property