From 2d259caf6b47d5e81ac9542bb5d62f4de592172d Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 20 Nov 2024 18:01:36 +0100
Subject: [PATCH 01/13] Changed rescale-method of JointPrior to always return
 correct-size array and update in-place once all keys are requested. Changed
 (Conditional)PriorDict.rescale to always return samples in right shape.

---
 bilby/core/prior/dict.py  | 35 ++++++++++++++------------
 bilby/core/prior/joint.py | 52 ++++++++++++++++++++++++++++-----------
 2 files changed, 57 insertions(+), 30 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index be3d543a9..2908bd083 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -600,18 +600,21 @@ def rescale(self, keys, theta):
         ==========
         keys: list
             List of prior keys to be rescaled
-        theta: list
-            List of randomly drawn values on a unit cube associated with the prior keys
+        theta: dict or array-like
+            Randomly drawn values on a unit cube associated with the prior keys
 
         Returns
         =======
-        list: List of floats containing the rescaled sample
+        list:
+            If theta is 1D, returns list of floats containing the rescaled sample.
+            If theta is 2D, returns list of lists containing the rescaled samples.
         """
-        theta = list(theta)
+        theta = [theta[key] for key in keys] if isinstance(theta, dict) else list(theta)
         samples = []
         for key, units in zip(keys, theta):
             samps = self[key].rescale(units)
-            samples += list(np.asarray(samps).flatten())
+            # turns 0d-arrays into scalars
+            samples.append(np.squeeze(samps).tolist())
         return samples
 
     def test_redundancy(self, key, disable_logging=False):
@@ -832,28 +835,28 @@ def rescale(self, keys, theta):
         ==========
         keys: list
             List of prior keys to be rescaled
-        theta: list
-            List of randomly drawn values on a unit cube associated with the prior keys
+        theta: dict or array-like
+            Randomly drawn values on a unit cube associated with the prior keys
 
         Returns
         =======
-        list: List of floats containing the rescaled sample
+        list:
+            If theta is float for each key, returns list of floats containing the rescaled sample.
+            If theta is array-like for each key, returns list of lists containing the rescaled samples.
         """
         keys = list(keys)
-        theta = list(theta)
+        theta = [theta[key] for key in keys] if isinstance(theta, dict) else list(theta)
         self._check_resolved()
         self._update_rescale_keys(keys)
         result = dict()
-        for key, index in zip(
-            self.sorted_keys_without_fixed_parameters, self._rescale_indexes
-        ):
-            result[key] = self[key].rescale(
-                theta[index], **self.get_required_variables(key)
-            )
+        for key, index in zip(self.sorted_keys_without_fixed_parameters, self._rescale_indexes):
+            result[key] = self[key].rescale(theta[index], **self.get_required_variables(key))
             self[key].least_recently_sampled = result[key]
         samples = []
         for key in keys:
-            samples += list(np.asarray(result[key]).flatten())
+            # turns 0d-arrays into scalars
+            res = np.squeeze(result[key]).tolist()
+            samples.append(res)
         return samples
 
     def _update_rescale_keys(self, keys):
diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index 43c8913e3..b088b15ba 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -63,8 +63,9 @@ def __init__(self, names, bounds=None):
         self.requested_parameters = dict()
         self.reset_request()
 
-        # a dictionary of the rescaled parameters
-        self.rescale_parameters = dict()
+        # a dictionary of the rescale(d) parameters
+        self._rescale_parameters = dict()
+        self._rescaled_parameters = dict()
         self.reset_rescale()
 
         # a list of sampled parameters
@@ -94,7 +95,12 @@ def filled_rescale(self):
         Check if all the rescaled parameters have been filled.
         """
 
-        return not np.any([val is None for val in self.rescale_parameters.values()])
+        return not np.any([val is None for val in self._rescale_parameters.values()])
+
+    def set_rescale(self, key, values):
+        values = np.array(values)
+        self._rescale_parameters[key] = values
+        self._rescaled_parameters[key] = np.atleast_1d(np.ones_like(values)) * np.nan
 
     def reset_rescale(self):
         """
@@ -102,7 +108,11 @@ def reset_rescale(self):
         """
 
         for name in self.names:
-            self.rescale_parameters[name] = None
+            self._rescale_parameters[name] = None
+            self._rescaled_parameters[name] = None
+
+    def get_rescaled(self, key):
+        return self._rescaled_parameters[key]
 
     def get_instantiation_dict(self):
         subclass_args = infer_args_from_method(self.__init__)
@@ -303,10 +313,11 @@ def rescale(self, value, **kwargs):
 
         Parameters
         ==========
-        value: array
-            A 1d vector sample (one for each parameter) drawn from a uniform
+        value: array or None
+            If given, a 1d vector sample (one for each parameter) drawn from a uniform
             distribution between 0 and 1, or a 2d NxM array of samples where
             N is the number of samples and M is the number of parameters.
+            If None, values previously set using BaseJointPriorDist.set_rescale() are used.
         kwargs: dict
             All keyword args that need to be passed to _rescale method, these keyword
             args are called in the JointPrior rescale methods for each parameter
@@ -317,7 +328,11 @@ def rescale(self, value, **kwargs):
             An vector sample drawn from the multivariate Gaussian
             distribution.
         """
-        samp = np.array(value)
+        if value is None:
+            samp = np.array(list(self._rescale_parameters.values())).T
+        else:
+            samp = np.array(value)
+
         if len(samp.shape) == 1:
             samp = samp.reshape(1, self.num_vars)
 
@@ -327,6 +342,11 @@ def rescale(self, value, **kwargs):
             raise ValueError("Array is the wrong shape")
 
         samp = self._rescale(samp, **kwargs)
+        if value is None:
+            for i, key in enumerate(self.names):
+                output = self.get_rescaled(key)
+                # update in-place for proper handling in PriorDict-instances
+                output[:] = samp[:, i]
         return np.squeeze(samp)
 
     def _rescale(self, samp, **kwargs):
@@ -790,19 +810,23 @@ def rescale(self, val, **kwargs):
             all kwargs passed to the dist.rescale method
         Returns
         =======
-        float:
-            A sample from the prior parameter.
+        np.ndarray:
+            The samples from the prior parameter. If not all names in "dist" have been filled,
+            the array contains only np.nan. *This* specific array instance will be filled with
+            the rescaled value once all parameters have been requested
         """
 
-        self.dist.rescale_parameters[self.name] = val
+        self.dist.set_rescale(self.name, val)
 
         if self.dist.filled_rescale():
-            values = np.array(list(self.dist.rescale_parameters.values())).T
-            samples = self.dist.rescale(values, **kwargs)
+            self.dist.rescale(values=None, **kwargs)
+            output = self.dist.get_rescaled(self.name)
             self.dist.reset_rescale()
-            return samples
         else:
-            return []  # return empty list
+            output = self.dist.get_rescaled(self.name)
+
+        # have to return raw output to conserve in-place modifications
+        return output
 
     def sample(self, size=1, **kwargs):
         """

From 5ce0bbe40a7f169181c4676611b58015bd026eb8 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Fri, 22 Nov 2024 09:55:06 +0100
Subject: [PATCH 02/13] Small fix to rescale of JointPrior

---
 bilby/core/prior/joint.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index b088b15ba..7e6655ee3 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -819,7 +819,7 @@ def rescale(self, val, **kwargs):
         self.dist.set_rescale(self.name, val)
 
         if self.dist.filled_rescale():
-            self.dist.rescale(values=None, **kwargs)
+            self.dist.rescale(value=None, **kwargs)
             output = self.dist.get_rescaled(self.name)
             self.dist.reset_rescale()
         else:

From 4e1a3d2d50f5abc97a3cb95153b7dac1eefa410b Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Mon, 25 Nov 2024 16:50:02 +0100
Subject: [PATCH 03/13] For jointprior rescale, only cast to list once its save
 to loose mutability

---
 bilby/core/prior/dict.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 2908bd083..0490d194e 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -613,8 +613,10 @@ def rescale(self, keys, theta):
         samples = []
         for key, units in zip(keys, theta):
             samps = self[key].rescale(units)
+            samples.append(samps)
+        for i, samps in enumerate(samples):
             # turns 0d-arrays into scalars
-            samples.append(np.squeeze(samps).tolist())
+            samples[i] = np.squeeze(samps).tolist()
         return samples
 
     def test_redundancy(self, key, disable_logging=False):

From 9abc7d81f31104ebc6d67568633d010a99fd9919 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Fri, 25 Oct 2024 16:23:00 +0200
Subject: [PATCH 04/13] Added qmc_quad based method for estimation of the
 constrained normalization factor

---
 bilby/core/prior/dict.py | 39 +++++++++++++++++++++++++++++++++++++++
 1 file changed, 39 insertions(+)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 0490d194e..ef58df8b2 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -5,6 +5,7 @@
 from io import open as ioopen
 
 import numpy as np
+from scipy.integrate import qmc_quad
 
 from .analytical import DeltaFunction
 from .base import Prior, Constraint
@@ -465,6 +466,44 @@ def sample_subset_constrained(self, keys=iter([]), size=None):
             }
             return all_samples
 
+    def normalize_constraint_factor_qmc_quad(self, keys, nrepeats=10, sampling_chunk=50000,
+                                             rel_error_target=0.01, max_trials=False):
+
+        def integrand(theta):
+            samples = np.array(self.rescale(keys=keys, theta=theta)).reshape((len(keys), -1))
+            samples = {key: samps for key, samps in zip(keys, samples)}
+            probs = self.evaluate_constraints(samples)
+            return probs
+
+        trials = 0
+        estimates = []
+        weights = []
+        while True:
+            res = qmc_quad(func=integrand, a=np.zeros(len(keys)), b=np.ones(len(keys)),
+                           n_estimates=nrepeats, n_points=sampling_chunk)
+            trials += sampling_chunk * nrepeats
+
+            # compute a weighted mean of the integral measurements and the resulting standard deviation
+            estimates.append(res.integral)
+            weights.append(1 / res.standard_error)
+
+            cumulative_weighted_error = 1 / np.sum(weights)
+            cumulative_weighted_integral = cumulative_weighted_error * np.sum(np.array(estimates) * np.array(weights))
+
+            # compute the rounded factor and relative error (as given by the standard deviation)
+            factor = 1 / cumulative_weighted_integral
+            rel_error = factor * cumulative_weighted_error
+
+            if rel_error > rel_error_target:
+                if max_trials and (trials + sampling_chunk * nrepeats > max_trials):
+                    break
+            else:
+                break
+
+        decimals = int(-np.floor(np.log10(3 * rel_error)))
+        factor_rounded = np.round(factor, decimals)
+        return factor_rounded
+
     def normalize_constraint_factor(
         self, keys, min_accept=10000, sampling_chunk=50000, nrepeats=10
     ):

From 9158a7f7672370bbd2427f3c35731f360c6a40a6 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 30 Oct 2024 14:45:42 +0100
Subject: [PATCH 05/13] added unit test

---
 test/core/prior/dict_test.py | 25 +++++++++++++++++++++++++
 1 file changed, 25 insertions(+)

diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py
index d6e6239f2..198333458 100644
--- a/test/core/prior/dict_test.py
+++ b/test/core/prior/dict_test.py
@@ -294,6 +294,31 @@ def test_sample_subset_constrained_as_array(self):
         self.assertTrue(isinstance(out, np.ndarray))
         self.assertTrue(out.shape == (len(keys), size))
 
+    def test_constrained_normalization(self):
+        prior_x = bilby.prior.Uniform(-1, 1, name="x")
+        prior_y = bilby.prior.Uniform(-1, 1, name="y")
+
+        r = 1
+
+        def radius_squared(sample_dict):
+            sample_dict["constraint"] = sample_dict["x"] ** 2 + sample_dict["y"] ** 2
+            return sample_dict
+
+        prior_constraint = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
+
+        prior_dict = bilby.prior.PriorDict(
+            {
+                "x": prior_x,
+                "y": prior_y,
+                "constraint": prior_constraint,
+            },
+            conversion_function=radius_squared,
+        )
+
+        factor = prior_dict.normalize_constraint_factor(keys=("x", "y"))
+        truth = 4 / (np.pi * r**2)
+        self.assertAlmostEqual(truth, factor, places=2)
+
     def test_sample(self):
         size = 7
         bilby.core.utils.random.seed(42)

From 5e16450a9fb023ef9f8b64ba9cded565819a3ecd Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 30 Oct 2024 14:46:52 +0100
Subject: [PATCH 06/13] Refactor normalization factor estimation and make it
 more robust against bugs

---
 bilby/core/prior/dict.py | 158 +++++++++++++++++++++++++--------------
 1 file changed, 103 insertions(+), 55 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index ef58df8b2..e27cab0a9 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -466,77 +466,122 @@ def sample_subset_constrained(self, keys=iter([]), size=None):
             }
             return all_samples
 
-    def normalize_constraint_factor_qmc_quad(self, keys, nrepeats=10, sampling_chunk=50000,
-                                             rel_error_target=0.01, max_trials=False):
+    def _integrate_normalization_factor_from_samples(self, keys, sampling_chunk):
+        samples = self.sample_subset(keys=keys, size=sampling_chunk)
+        count = np.count_nonzero(np.array(self.evaluate_constraints(samples), dtype=bool))
+        factor = count / sampling_chunk
+        return factor
+
+    def _integrate_normalization_factor_from_qmc_quad(self, keys, sampling_chunk):
 
         def integrand(theta):
-            samples = np.array(self.rescale(keys=keys, theta=theta)).reshape((len(keys), -1))
-            samples = {key: samps for key, samps in zip(keys, samples)}
+            samples = {key: self[key].rescale(units) for key, units in zip(keys, theta)}
             probs = self.evaluate_constraints(samples)
             return probs
 
+        res = qmc_quad(func=integrand, a=np.zeros(len(keys)), b=np.ones(len(keys)),
+                       n_estimates=1, n_points=sampling_chunk)
+
+        return res.integral
+
+    def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
+                                    rel_error_target=0.01, max_trials=False, method="qmc_quad", **kwargs):
+        """Estimates the probality normalization factor for constrained priors from (quasi) Monte Carlo integration.
+
+        Parameters
+        ==========
+        keys : list, tuple
+            The set of keys in the prior dict to perform the integration for. Must contain all keys that the constraint
+            depends on. For joint priors, the full distribution is sampled. Joint prior keys not present in 'keys' are
+            marginalized.
+        nrepeats: int
+            Number of repeated Monte Carlo integrations before convergence is checked. Higher numbers improve the
+            estimation of the sampling error.
+        sampling_chunk: int
+            Number of samples drawn per Monte Carlo integration. Higher numbers improve the integral estimation.
+        rel_error_target: float
+            The relative error targeted by the Monte Carlo integration. The relative error of the integral is estimated
+            from the standard deviation between repeated runs. The algorithm repeats 'nrepeats' iterations until the
+            rel_error_target is reached.
+        max_trials: False, int
+            Second termination criterion. The integration stops when the number of samples evaluated in the next
+            iteration would exceed 'max_trials' even if rel_error_target was not reached.
+        method: ["qmc_quad", "from_samples"]
+            Method to use for the Monte Carlo integration, effectively choosing between quasi Monte Carlo integration
+            and "normal" Monte Carlo integration. Due to computational overhead, the latter method will be faster in
+            most cases. However, 'qmc_quad' is expected to yield lower errors against the ground truth for cases with
+            few repeats, but high sampling_chunk.
+
+        Returns
+        =======
+        factor_rounded : float
+            The normalization factor, rounded to the number of significant digits based on the standard deviation of
+            the integration estimate.
+
+        """
+        keys = tuple(keys)
+        if keys in self._cached_normalizations.keys():
+            return self._cached_normalizations[keys]
+
+        sample_keys = [key for key in keys]
+
+        # check if the constraint can be applied to the selection of keys
+        try:
+            sample = self.sample_subset(keys, size=1)
+            self.conversion_function(sample)
+        except KeyError:
+            raise ValueError("'keys' does not contain all parameters needed to evaluate the constraint.")
+
+        for key in sample_keys:
+            if isinstance(self[key], JointPrior):
+                dist_keys = set(self[key].dist.names)
+                missing_keys = list(dist_keys - set(keys))
+                sample_keys.extend(missing_keys)
+
+        if method == "qmc_quad":
+            integrator = self._integrate_normalization_factor_from_qmc_quad
+        elif method == "from_samples":
+            integrator = self._integrate_normalization_factor_from_samples
+        else:
+            raise ValueError(f"Integration method {method} not understood")
+        try:
+            theta = np.random.uniform(0, 1, size=(len(keys), 1))
+            samples = {key: self[key].rescale(units) for key, units in zip(keys, theta)}
+            for key in keys:
+                if np.any(np.isnan(samples[key])):
+                    print("The rescale method appears to be not working. Switching to 'sample_based'.")
+                    integrator = self._integrate_normalization_factor_from_samples
+        except Exception:
+            print("The rescale method appears to be not working. Switching to 'sample_based'.")
+            integrator = self._integrate_normalization_factor_from_samples
+
         trials = 0
         estimates = []
-        weights = []
         while True:
-            res = qmc_quad(func=integrand, a=np.zeros(len(keys)), b=np.ones(len(keys)),
-                           n_estimates=nrepeats, n_points=sampling_chunk)
-            trials += sampling_chunk * nrepeats
+            for i in range(nrepeats):
+                integral = integrator(sample_keys, sampling_chunk, **kwargs)
+                trials += sampling_chunk
 
-            # compute a weighted mean of the integral measurements and the resulting standard deviation
-            estimates.append(res.integral)
-            weights.append(1 / res.standard_error)
+                # compute a weighted mean of the integral measurements and the resulting standard deviation
+                estimates.append(integral)
 
-            cumulative_weighted_error = 1 / np.sum(weights)
-            cumulative_weighted_integral = cumulative_weighted_error * np.sum(np.array(estimates) * np.array(weights))
+            standard_error = np.std(estimates, ddof=1) / np.sqrt(len(estimates))
+            cumulative_integral = np.sum(estimates) / len(estimates)
 
             # compute the rounded factor and relative error (as given by the standard deviation)
-            factor = 1 / cumulative_weighted_integral
-            rel_error = factor * cumulative_weighted_error
+            factor = 1 / cumulative_integral
+            rel_error = factor * standard_error
 
-            if rel_error > rel_error_target:
-                if max_trials and (trials + sampling_chunk * nrepeats > max_trials):
-                    break
-            else:
+            if rel_error < rel_error_target:
+                break
+            elif max_trials and ((trials + sampling_chunk * nrepeats) > max_trials):
                 break
 
         decimals = int(-np.floor(np.log10(3 * rel_error)))
         factor_rounded = np.round(factor, decimals)
-        return factor_rounded
-
-    def normalize_constraint_factor(
-        self, keys, min_accept=10000, sampling_chunk=50000, nrepeats=10
-    ):
-        if keys in self._cached_normalizations.keys():
-            return self._cached_normalizations[keys]
-        else:
-            factor_estimates = [
-                self._estimate_normalization(keys, min_accept, sampling_chunk)
-                for _ in range(nrepeats)
-            ]
-            factor = np.mean(factor_estimates)
-            if np.std(factor_estimates) > 0:
-                decimals = int(-np.floor(np.log10(3 * np.std(factor_estimates))))
-                factor_rounded = np.round(factor, decimals)
-            else:
-                factor_rounded = factor
-            self._cached_normalizations[keys] = factor_rounded
-            return factor_rounded
 
-    def _estimate_normalization(self, keys, min_accept, sampling_chunk):
-        samples = self.sample_subset(keys=keys, size=sampling_chunk)
-        keep = np.atleast_1d(self.evaluate_constraints(samples))
-        if len(keep) == 1:
-            self._cached_normalizations[keys] = 1
-            return 1
-        all_samples = {key: np.array([]) for key in keys}
-        while np.count_nonzero(keep) < min_accept:
-            samples = self.sample_subset(keys=keys, size=sampling_chunk)
-            for key in samples:
-                all_samples[key] = np.hstack([all_samples[key], samples[key].flatten()])
-            keep = np.array(self.evaluate_constraints(all_samples), dtype=bool)
-        factor = len(keep) / np.count_nonzero(keep)
-        return factor
+        self._cached_normalizations[keys] = factor_rounded
+        return factor_rounded
 
     def prob(self, sample, **kwargs):
         """
@@ -558,6 +603,9 @@ def prob(self, sample, **kwargs):
         return self.check_prob(sample, prob)
 
     def check_prob(self, sample, prob):
+        if self.conversion_function == self.default_conversion_function:
+            return prob
+
         ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         if np.all(prob == 0.0):
             return prob * ratio
@@ -597,10 +645,10 @@ def ln_prob(self, sample, axis=None, normalized=True):
                                   normalized=normalized)
 
     def check_ln_prob(self, sample, ln_prob, normalized=True):
-        if normalized:
+        if normalized and (self.conversion_function != self.default_conversion_function):
             ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         else:
-            ratio = 1
+            return ln_prob
         if np.all(np.isinf(ln_prob)):
             return ln_prob
         else:

From dc8a5c8213671da4f143fdd73b7e293915def165 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Fri, 1 Nov 2024 17:48:48 +0100
Subject: [PATCH 07/13] Applied changes discussed on Github and updated
 qmc_quad for less computaional overhead and repeateability

---
 bilby/core/prior/dict.py     | 75 +++++++++++++++++-------------------
 test/core/prior/dict_test.py | 17 ++++++++
 2 files changed, 53 insertions(+), 39 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index e27cab0a9..146210abc 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -5,7 +5,7 @@
 from io import open as ioopen
 
 import numpy as np
-from scipy.integrate import qmc_quad
+from scipy.stats.qmc import Halton
 
 from .analytical import DeltaFunction
 from .base import Prior, Constraint
@@ -15,6 +15,7 @@
     check_directory_exists_and_if_not_mkdir,
     BilbyJsonEncoder,
     decode_bilby_json,
+    random
 )
 
 
@@ -468,24 +469,19 @@ def sample_subset_constrained(self, keys=iter([]), size=None):
 
     def _integrate_normalization_factor_from_samples(self, keys, sampling_chunk):
         samples = self.sample_subset(keys=keys, size=sampling_chunk)
-        count = np.count_nonzero(np.array(self.evaluate_constraints(samples), dtype=bool))
-        factor = count / sampling_chunk
+        factor = np.mean(self.evaluate_constraints(samples))
         return factor
 
     def _integrate_normalization_factor_from_qmc_quad(self, keys, sampling_chunk):
+        qrng = Halton(len(keys), seed=random.rng)
+        theta = qrng.random(sampling_chunk).T
+        samples = np.array(self.rescale(keys=keys, theta=theta)).reshape((len(keys), -1))
+        samples = {key: samps for key, samps in zip(keys, samples)}
+        factor = np.mean(self.evaluate_constraints(samples))
+        return factor
 
-        def integrand(theta):
-            samples = {key: self[key].rescale(units) for key, units in zip(keys, theta)}
-            probs = self.evaluate_constraints(samples)
-            return probs
-
-        res = qmc_quad(func=integrand, a=np.zeros(len(keys)), b=np.ones(len(keys)),
-                       n_estimates=1, n_points=sampling_chunk)
-
-        return res.integral
-
-    def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
-                                    rel_error_target=0.01, max_trials=False, method="qmc_quad", **kwargs):
+    def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, rel_error_target=0.01,
+                                    max_trials=False, method="qmc_quad", **kwargs):
         """Estimates the probality normalization factor for constrained priors from (quasi) Monte Carlo integration.
 
         Parameters
@@ -523,7 +519,7 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
         if keys in self._cached_normalizations.keys():
             return self._cached_normalizations[keys]
 
-        sample_keys = [key for key in keys]
+        sample_keys = tuple(keys)
 
         # check if the constraint can be applied to the selection of keys
         try:
@@ -538,22 +534,23 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
                 missing_keys = list(dist_keys - set(keys))
                 sample_keys.extend(missing_keys)
 
-        if method == "qmc_quad":
-            integrator = self._integrate_normalization_factor_from_qmc_quad
-        elif method == "from_samples":
+        if method == "from_samples":
             integrator = self._integrate_normalization_factor_from_samples
-        else:
-            raise ValueError(f"Integration method {method} not understood")
-        try:
+        elif method == "qmc_quad":
+            integrator = self._integrate_normalization_factor_from_qmc_quad
             theta = np.random.uniform(0, 1, size=(len(keys), 1))
-            samples = {key: self[key].rescale(units) for key, units in zip(keys, theta)}
-            for key in keys:
-                if np.any(np.isnan(samples[key])):
-                    print("The rescale method appears to be not working. Switching to 'sample_based'.")
-                    integrator = self._integrate_normalization_factor_from_samples
-        except Exception:
-            print("The rescale method appears to be not working. Switching to 'sample_based'.")
-            integrator = self._integrate_normalization_factor_from_samples
+            samples = self.rescale(keys, theta)
+            none_keys = []
+            for i, key in enumerate(keys):
+                if samples[i] is None:
+                    none_keys.append(key)
+            if len(none_keys) > 0:
+                print(f"The rescale method returns 'None', for key(s) {none_keys}. Switching to 'sample_based'.")
+                integrator = self._integrate_normalization_factor_from_samples
+                method = "from_samples"
+        else:
+            raise ValueError(f"Integration method {method} not understood.\n" +
+                             "Available options are ('from_samples','qmc_quad').")
 
         trials = 0
         estimates = []
@@ -566,7 +563,7 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
                 estimates.append(integral)
 
             standard_error = np.std(estimates, ddof=1) / np.sqrt(len(estimates))
-            cumulative_integral = np.sum(estimates) / len(estimates)
+            cumulative_integral = np.mean(estimates)
 
             # compute the rounded factor and relative error (as given by the standard deviation)
             factor = 1 / cumulative_integral
@@ -577,11 +574,14 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000,
             elif max_trials and ((trials + sampling_chunk * nrepeats) > max_trials):
                 break
 
-        decimals = int(-np.floor(np.log10(3 * rel_error)))
-        factor_rounded = np.round(factor, decimals)
+        if rel_error > 0:
+            decimals = int(-np.floor(np.log10(3 * rel_error)))
+            factor_rounded = np.round(factor, decimals)
+            self._cached_normalizations[keys] = factor_rounded
+        else:
+            self._cached_normalizations[keys] = factor
 
-        self._cached_normalizations[keys] = factor_rounded
-        return factor_rounded
+        return self._cached_normalizations[keys]
 
     def prob(self, sample, **kwargs):
         """
@@ -603,9 +603,6 @@ def prob(self, sample, **kwargs):
         return self.check_prob(sample, prob)
 
     def check_prob(self, sample, prob):
-        if self.conversion_function == self.default_conversion_function:
-            return prob
-
         ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         if np.all(prob == 0.0):
             return prob * ratio
@@ -645,7 +642,7 @@ def ln_prob(self, sample, axis=None, normalized=True):
                                   normalized=normalized)
 
     def check_ln_prob(self, sample, ln_prob, normalized=True):
-        if normalized and (self.conversion_function != self.default_conversion_function):
+        if normalized:
             ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         else:
             return ln_prob
diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py
index 198333458..8d44f4120 100644
--- a/test/core/prior/dict_test.py
+++ b/test/core/prior/dict_test.py
@@ -319,6 +319,23 @@ def radius_squared(sample_dict):
         truth = 4 / (np.pi * r**2)
         self.assertAlmostEqual(truth, factor, places=2)
 
+        r = 2
+
+        prior_constraint_2 = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
+
+        prior_dict_2 = bilby.prior.PriorDict(
+            {
+                "x": prior_x,
+                "y": prior_y,
+                "constraint": prior_constraint_2,
+            },
+            conversion_function=radius_squared,
+        )
+
+        factor_2 = prior_dict_2.normalize_constraint_factor(keys=("x", "y"))
+        truth_2 = 1
+        self.assertAlmostEqual(truth_2, factor_2, places=2)
+
     def test_sample(self):
         size = 7
         bilby.core.utils.random.seed(42)

From a63918d88970a0d60caf7232994c1e298ff339c2 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 12 Nov 2024 15:39:08 +0100
Subject: [PATCH 08/13] Added check if constraint estimation is necessary and
 updated check if rescale method is save

---
 bilby/core/prior/dict.py | 28 +++++++++++++++++++---------
 1 file changed, 19 insertions(+), 9 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 146210abc..cc888c742 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -442,6 +442,10 @@ def fixed_keys(self):
     def constraint_keys(self):
         return [k for k, p in self.items() if isinstance(p, Constraint)]
 
+    @property
+    def has_constraint(self):
+        return len(self.constraint_keys) > 0
+
     def sample_subset_constrained(self, keys=iter([]), size=None):
         if size is None or size == 1:
             while True:
@@ -538,14 +542,18 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
             integrator = self._integrate_normalization_factor_from_samples
         elif method == "qmc_quad":
             integrator = self._integrate_normalization_factor_from_qmc_quad
-            theta = np.random.uniform(0, 1, size=(len(keys), 1))
-            samples = self.rescale(keys, theta)
-            none_keys = []
-            for i, key in enumerate(keys):
-                if samples[i] is None:
-                    none_keys.append(key)
-            if len(none_keys) > 0:
-                print(f"The rescale method returns 'None', for key(s) {none_keys}. Switching to 'sample_based'.")
+            try:
+                theta = np.random.uniform(0, 1, size=(len(keys), 1))
+                samples = self.rescale(keys, theta)
+                none_keys = []
+                for i, key in enumerate(keys):
+                    if samples[i] is None:
+                        none_keys.append(key)
+                if len(none_keys) > 0:
+                    raise NotImplementedError(f"The rescale method returns 'None', for key(s) {none_keys}.")
+            except NotImplementedError as e:
+                print(f"The rescaling step fails with message:\n{e}")
+                print("Switching to method 'from_samples'")
                 integrator = self._integrate_normalization_factor_from_samples
                 method = "from_samples"
         else:
@@ -603,6 +611,8 @@ def prob(self, sample, **kwargs):
         return self.check_prob(sample, prob)
 
     def check_prob(self, sample, prob):
+        if not self.has_constraint:
+            return prob
         ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         if np.all(prob == 0.0):
             return prob * ratio
@@ -642,7 +652,7 @@ def ln_prob(self, sample, axis=None, normalized=True):
                                   normalized=normalized)
 
     def check_ln_prob(self, sample, ln_prob, normalized=True):
-        if normalized:
+        if normalized and self.has_constraint:
             ratio = self.normalize_constraint_factor(tuple(sample.keys()))
         else:
             return ln_prob

From 07a01c5a90fb999208c340d0dc88b440875c62a6 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 12 Nov 2024 17:23:58 +0100
Subject: [PATCH 09/13] switched to logger.info()

---
 bilby/core/prior/dict.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index cc888c742..30ad9a5a4 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -552,8 +552,8 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
                 if len(none_keys) > 0:
                     raise NotImplementedError(f"The rescale method returns 'None', for key(s) {none_keys}.")
             except NotImplementedError as e:
-                print(f"The rescaling step fails with message:\n{e}")
-                print("Switching to method 'from_samples'")
+                logger.info(f"The rescaling step fails with message:\n{e}")
+                logger.info("Switching to method 'from_samples'")
                 integrator = self._integrate_normalization_factor_from_samples
                 method = "from_samples"
         else:

From 9083bce859646426eb465735b506ef625d5818dd Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 12 Nov 2024 18:02:30 +0100
Subject: [PATCH 10/13] updated documentation to include versionchanged and
 seealso

---
 bilby/core/prior/dict.py | 21 +++++++++++++++++----
 1 file changed, 17 insertions(+), 4 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 30ad9a5a4..8b3c2d4ee 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -506,7 +506,7 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
         max_trials: False, int
             Second termination criterion. The integration stops when the number of samples evaluated in the next
             iteration would exceed 'max_trials' even if rel_error_target was not reached.
-        method: ["qmc_quad", "from_samples"]
+        method: ['qmc_quad', 'from_samples']
             Method to use for the Monte Carlo integration, effectively choosing between quasi Monte Carlo integration
             and "normal" Monte Carlo integration. Due to computational overhead, the latter method will be faster in
             most cases. However, 'qmc_quad' is expected to yield lower errors against the ground truth for cases with
@@ -514,10 +514,25 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
 
         Returns
         =======
-        factor_rounded : float
+        factor_rounded: float
             The normalization factor, rounded to the number of significant digits based on the standard deviation of
             the integration estimate.
 
+        Notes
+        =====
+        .. seealso::
+            Scipy's :scipy:stats:qmc:`Halton` class
+                Documentation of the quasi-random number generator :scipy:stats:qmc:`Halton` used for the quasi-monte
+                carlo integration of the normalization factor.
+            Scipy's :scipy:integrate:`qmc_quad` method
+                Documentation of the quasi-monte carlo integration scheme native to scipy.
+                The implementation, particularly the error estimate, motivates this implementation.
+                (The error estimate was re-implemented to also apply for the 'from_samples' method.)
+
+        .. versionchanged:: 4.0
+            The estimation of the normalization factor is now by default based on quasi-monte carlo integration.
+            Further, the default stopping criterion is now based on an target for the estimated relative integration
+            error.
         """
         keys = tuple(keys)
         if keys in self._cached_normalizations.keys():
@@ -566,8 +581,6 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
             for i in range(nrepeats):
                 integral = integrator(sample_keys, sampling_chunk, **kwargs)
                 trials += sampling_chunk
-
-                # compute a weighted mean of the integral measurements and the resulting standard deviation
                 estimates.append(integral)
 
             standard_error = np.std(estimates, ddof=1) / np.sqrt(len(estimates))

From 64e85cca156647ea05dcf60290e81bdce9073e15 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 13 Nov 2024 18:35:10 +0100
Subject: [PATCH 11/13] Updated method documentation

---
 bilby/core/prior/dict.py | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 8b3c2d4ee..db76ab700 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -490,7 +490,7 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
 
         Parameters
         ==========
-        keys : list, tuple
+        keys: list, tuple
             The set of keys in the prior dict to perform the integration for. Must contain all keys that the constraint
             depends on. For joint priors, the full distribution is sampled. Joint prior keys not present in 'keys' are
             marginalized.
@@ -521,11 +521,11 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
         Notes
         =====
         .. seealso::
-            Scipy's :scipy:stats:qmc:`Halton` class
-                Documentation of the quasi-random number generator :scipy:stats:qmc:`Halton` used for the quasi-monte
-                carlo integration of the normalization factor.
-            Scipy's :scipy:integrate:`qmc_quad` method
-                Documentation of the quasi-monte carlo integration scheme native to scipy.
+            :py:class:`scipy.stats.qmc.Halton`
+                Documentation of the quasi-random number generator used for the quasi-monte carlo-based method
+                to integrate the normalization factor.
+            :py:func:`scipy.integrate.qmc_quad`
+                Documentation of the scipy-native quasi-monte carlo integration scheme.
                 The implementation, particularly the error estimate, motivates this implementation.
                 (The error estimate was re-implemented to also apply for the 'from_samples' method.)
 

From 88e9ee7bffe76ac2972eb604c06e95acef9aa08f Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Mon, 25 Nov 2024 16:37:56 +0100
Subject: [PATCH 12/13] Added test cases and and jointprior handling

---
 bilby/core/prior/dict.py     | 12 ++++----
 test/core/prior/dict_test.py | 57 ++++++++++++++++++++++++++++++++----
 2 files changed, 59 insertions(+), 10 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index db76ab700..91a1a1cb0 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -550,18 +550,20 @@ def normalize_constraint_factor(self, keys, nrepeats=10, sampling_chunk=10000, r
         for key in sample_keys:
             if isinstance(self[key], JointPrior):
                 dist_keys = set(self[key].dist.names)
-                missing_keys = list(dist_keys - set(keys))
-                sample_keys.extend(missing_keys)
+                missing_keys = tuple(dist_keys - set(sample_keys))
+                sample_keys += missing_keys
+                for key in missing_keys:
+                    self.sample_subset(missing_keys, 1)
 
         if method == "from_samples":
             integrator = self._integrate_normalization_factor_from_samples
         elif method == "qmc_quad":
             integrator = self._integrate_normalization_factor_from_qmc_quad
             try:
-                theta = np.random.uniform(0, 1, size=(len(keys), 1))
-                samples = self.rescale(keys, theta)
+                theta = np.random.uniform(0, 1, size=(len(sample_keys), 1))
+                samples = self.rescale(sample_keys, theta)
                 none_keys = []
-                for i, key in enumerate(keys):
+                for i, key in enumerate(sample_keys):
                     if samples[i] is None:
                         none_keys.append(key)
                 if len(none_keys) > 0:
diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py
index 8d44f4120..940856cbf 100644
--- a/test/core/prior/dict_test.py
+++ b/test/core/prior/dict_test.py
@@ -298,12 +298,11 @@ def test_constrained_normalization(self):
         prior_x = bilby.prior.Uniform(-1, 1, name="x")
         prior_y = bilby.prior.Uniform(-1, 1, name="y")
 
-        r = 1
-
         def radius_squared(sample_dict):
             sample_dict["constraint"] = sample_dict["x"] ** 2 + sample_dict["y"] ** 2
             return sample_dict
 
+        r = 1
         prior_constraint = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
 
         prior_dict = bilby.prior.PriorDict(
@@ -317,10 +316,9 @@ def radius_squared(sample_dict):
 
         factor = prior_dict.normalize_constraint_factor(keys=("x", "y"))
         truth = 4 / (np.pi * r**2)
-        self.assertAlmostEqual(truth, factor, places=2)
+        self.assertAlmostEqual(truth, factor, delta=truth * 0.01)
 
         r = 2
-
         prior_constraint_2 = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
 
         prior_dict_2 = bilby.prior.PriorDict(
@@ -334,7 +332,56 @@ def radius_squared(sample_dict):
 
         factor_2 = prior_dict_2.normalize_constraint_factor(keys=("x", "y"))
         truth_2 = 1
-        self.assertAlmostEqual(truth_2, factor_2, places=2)
+        self.assertAlmostEqual(truth_2, factor_2, delta=truth_2 * 0.01)
+
+        prior_dist_1 = bilby.prior.MultivariateGaussianDist(
+            names=["x", "y", "z"],
+            mus=[0, 0, 0],
+            sigmas=[1, 1, 1],
+        )
+
+        prior_dist_2 = bilby.prior.MultivariateGaussianDist(
+            names=["x", "y"],
+            mus=[0, 0],
+            sigmas=[1, 1],
+        )
+
+        joint_prior_x = bilby.prior.JointPrior(dist=prior_dist_1, name="x")
+        joint_prior_y = bilby.prior.JointPrior(dist=prior_dist_1, name="y")
+        joint_prior_z = bilby.prior.JointPrior(dist=prior_dist_1, name="z")
+
+        r = 1
+        prior_constraint_3 = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
+
+        joint_prior_x_2 = bilby.prior.JointPrior(dist=prior_dist_2, name="x")
+        joint_prior_y_2 = bilby.prior.JointPrior(dist=prior_dist_2, name="y")
+        prior_constraint_4 = bilby.prior.Constraint(minimum=0, maximum=r**2, name="constraint")
+        prior_dict_4 = bilby.prior.PriorDict(
+            {
+                "x": joint_prior_x_2,
+                "y": joint_prior_y_2,
+                "constraint": prior_constraint_4,
+            },
+            conversion_function=radius_squared,
+        )
+        prior_dict_3 = bilby.prior.PriorDict(
+            {
+                "x": joint_prior_x,
+                "y": joint_prior_y,
+                "z": joint_prior_z,
+                "constraint": prior_constraint_3,
+            },
+            conversion_function=radius_squared,
+        )
+
+        factor_3 = prior_dict_3.normalize_constraint_factor(keys=("x", "y"))
+        factor_4 = prior_dict_4.normalize_constraint_factor(keys=("x", "y"))
+
+        from scipy.stats import chi2
+
+        truth_3 = 1 / (chi2(df=2).cdf(1))
+        self.assertAlmostEqual(truth_3, factor_3, delta=truth_3 * 0.01)
+        self.assertAlmostEqual(truth_3, factor_4, delta=truth_3 * 0.01)
 
     def test_sample(self):
         size = 7

From 1b02af6147d3d2a11e2e6e15e4f2ee7a1741592a Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Mon, 25 Nov 2024 16:54:10 +0100
Subject: [PATCH 13/13] Adopt new rescale function that maintains correct
 output shape

---
 bilby/core/prior/dict.py     | 2 +-
 test/core/prior/dict_test.py | 2 +-
 2 files changed, 2 insertions(+), 2 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 91a1a1cb0..489b0655b 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -479,7 +479,7 @@ def _integrate_normalization_factor_from_samples(self, keys, sampling_chunk):
     def _integrate_normalization_factor_from_qmc_quad(self, keys, sampling_chunk):
         qrng = Halton(len(keys), seed=random.rng)
         theta = qrng.random(sampling_chunk).T
-        samples = np.array(self.rescale(keys=keys, theta=theta)).reshape((len(keys), -1))
+        samples = self.rescale(keys=keys, theta=theta)
         samples = {key: samps for key, samps in zip(keys, samples)}
         factor = np.mean(self.evaluate_constraints(samples))
         return factor
diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py
index 940856cbf..8c2f702e0 100644
--- a/test/core/prior/dict_test.py
+++ b/test/core/prior/dict_test.py
@@ -299,7 +299,7 @@ def test_constrained_normalization(self):
         prior_y = bilby.prior.Uniform(-1, 1, name="y")
 
         def radius_squared(sample_dict):
-            sample_dict["constraint"] = sample_dict["x"] ** 2 + sample_dict["y"] ** 2
+            sample_dict["constraint"] = np.array(sample_dict["x"]) ** 2 + np.array(sample_dict["y"]) ** 2
             return sample_dict
 
         r = 1