MAINT: tweaks to prior constraints

bilby/core/prior/dict.py

@@ -55,13 +55,29 @@ def __init__(self, dictionary=None, filename=None, conversion_function=None):
             self.conversion_function = self.default_conversion_function
 
     def evaluate_constraints(self, sample):
+        """Evaluate the constraints for a given sample.
+
+        Applies the conversion function to the sample and evaluates the
+        constraints on the converted sample.
+
+        Raises
+        ======
+        ValueError:
+            If a constraint parameter is not present in the sample after
+            conversion.
+        """
         out_sample = self.conversion_function(sample)
         try:
             prob = np.ones_like(next(iter(out_sample.values())))
         except TypeError:
             prob = np.ones_like(out_sample)
         for key in self:
-            if isinstance(self[key], Constraint) and key in out_sample:
+            if isinstance(self[key], Constraint):
+                if key not in out_sample:
+                    raise ValueError(
+                        f"Constraint {key} is not present in the sample. "
+                        "Cannot evaluate constraints."
+                    )
                 prob *= self[key].prob(out_sample[key])
         return prob
 
@@ -511,9 +527,6 @@ def normalize_constraint_factor(
     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)

test/core/prior/dict_test.py

@@ -392,6 +392,79 @@ def test_redundancy(self):
         for key in self.prior_set_from_dict.keys():
             self.assertFalse(self.prior_set_from_dict.test_redundancy(key=key))
 
+    def test_evaluate_constraints(self):
+
+        def conversion_function(parameters):
+            converted_parameters = parameters.copy()
+            converted_parameters["delta_mass"] = (
+                parameters["mass_1"] - parameters["mass_2"]
+            )
+            return converted_parameters
+
+        priors = bilby.core.prior.PriorDict(conversion_function=conversion_function)
+        priors["mass_1"] = bilby.core.prior.Uniform(minimum=1.38, maximum=2)
+        priors["mass_2"] = bilby.core.prior.Uniform(minimum=1, maximum=1.4)
+        priors["delta_mass"] = bilby.core.prior.Constraint(minimum=0.4, maximum=1.4)
+
+        theta = {"mass_1": 1.7, "mass_2": 1.2}
+        self.assertTrue(priors.evaluate_constraints(theta))
+
+        theta = {"mass_1": 1.5, "mass_2": 1.2}
+        self.assertFalse(priors.evaluate_constraints(theta))
+
+    def test_evaluate_constraints_batches(self):
+
+        def conversion_function(parameters):
+            converted_parameters = parameters.copy()
+            converted_parameters["delta_mass"] = (
+                parameters["mass_1"] - parameters["mass_2"]
+            )
+            return converted_parameters
+
+        priors = bilby.core.prior.PriorDict(conversion_function=conversion_function)
+        priors["mass_1"] = bilby.core.prior.Uniform(minimum=1.38, maximum=2)
+        priors["mass_2"] = bilby.core.prior.Uniform(minimum=1, maximum=1.4)
+        priors["delta_mass"] = bilby.core.prior.Constraint(minimum=0.4, maximum=1.4)
+
+        theta = {"mass_1": np.array([1.7, 1.5]), "mass_2": np.array([1.2, 1.2])}
+        expected = np.array([True, False])
+        self.assertTrue(np.array_equal(expected, priors.evaluate_constraints(theta)))
+
+    def test_evaluate_constraints_missing_keys(self):
+
+        def conversion_function(parameters):
+            return parameters.copy()
+
+        priors = bilby.core.prior.PriorDict(conversion_function=conversion_function)
+        priors["mass_1"] = bilby.core.prior.Uniform(minimum=1.38, maximum=2)
+        priors["mass_2"] = bilby.core.prior.Uniform(minimum=1, maximum=1.4)
+        priors["delta_mass"] = bilby.core.prior.Constraint(minimum=0.4, maximum=1.4)
+
+        theta = {"mass_1": 1.5, "mass_2": 1.2}
+
+        with self.assertRaises(
+            ValueError,
+            msg="Constraint delta_mass is not present in the sample. Cannot evaluate constraints."
+        ):
+            priors.evaluate_constraints(theta)
+
+    def test_normalize_constraint_keys(self):
+
+        def conversion_function(parameters):
+            converted_parameters = parameters.copy()
+            converted_parameters["mass_ratio"] = parameters["mass_2"] / parameters["mass_1"]
+            return converted_parameters
+
+        priors = bilby.core.prior.PriorDict(conversion_function=conversion_function)
+        priors["mass_1"] = bilby.core.prior.Uniform(minimum=1, maximum=2)
+        priors["mass_2"] = bilby.core.prior.Uniform(minimum=1, maximum=2)
+        priors["mass_ratio"] = bilby.core.prior.Constraint(minimum=0.0, maximum=1.0)
+
+        # Factor should close to 2 since half the prior volume is removed by the constraint
+        keys = ("mass_1", "mass_2")
+        factor = priors.normalize_constraint_factor(keys)
+        self.assertAlmostEqual(factor, 2.0, delta=0.01)
+
 
 class TestJsonIO(unittest.TestCase):
     def setUp(self):