MAINT: tweaks to prior constraints

bilby/core/prior/dict.py

@@ -61,14 +61,46 @@ def __hash__(self):
         return hash(str(self))
 
     @xp_wrap
-    def evaluate_constraints(self, sample, *, xp=None):
-        out_sample = self.conversion_function(sample)
+    def evaluate_constraints(self, sample, *, strict=True, xp=None):
+        """Evaluate the constraints for a given sample.
+
+        Applies the conversion function to the sample and evaluates the
+        constraints on the converted sample.
+
+        Parameters
+        ==========
+        sample: dict
+            Dictionary of parameters used to evaluate the constraints.
+        strict: bool, optional
+            When True, raise if a constraint cannot be evaluated from the
+            provided sample. When False, skip constraints that cannot be
+            derived from a partial sample.
+
+        Raises
+        ======
+        ValueError:
+            If a constraint parameter is not present in the sample after
+            conversion and ``strict`` is True.
+        """
+        try:
+            out_sample = self.conversion_function(sample)
+        except KeyError:
+            if strict:
+                raise
+            out_sample = sample.copy()
         try:
             prob = xp.ones_like(next(iter(out_sample.values())), dtype=bool)
         except TypeError:
             prob = xp.ones_like(out_sample, dtype=bool)
         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:
+                    if not strict:
+                        continue
+                    raise ValueError(
+                        f"Constraint {key} is not present in the sample. "
+                        "Cannot evaluate constraints."
+                    )
                 prob *= self[key].prob(out_sample[key])
         return prob
 
@@ -440,6 +472,10 @@ def fixed_keys(self):
     def constraint_keys(self):
         return [k for k, p in self.items() if isinstance(p, Constraint)]
 
+    def _sample_has_all_constrained_keys(self, sample):
+        sampled_prior_keys = set(self.non_fixed_keys + self.fixed_keys)
+        return sampled_prior_keys.issubset(set(sample.keys()))
+
     def sample_subset_constrained(self, keys=iter([]), size=None, *, random_state=None):
         """
         Sample a subset of priors while ensuring constraints are satisfied.
@@ -475,7 +511,10 @@ def check_efficiency(n_tested, n_valid):
         if size is None or size == 1:
             while True:
                 sample = self.sample_subset(keys=keys, size=size, random_state=rng)
-                is_valid = self.evaluate_constraints(sample)
+                is_valid = self.evaluate_constraints(
+                    sample,
+                    strict=self._sample_has_all_constrained_keys(sample),
+                )
                 n_tested_samples += 1
                 n_valid_samples += int(is_valid.item())
                 check_efficiency(n_tested_samples, n_valid_samples)
@@ -489,9 +528,10 @@ def check_efficiency(n_tested, n_valid):
             xp = random_array_module(random_state)
             all_samples = {key: xp.asarray([]) for key in keys}
             _first_key = list(all_samples.keys())[0]
+            strict = self._sample_has_all_constrained_keys(all_samples)
             while len(all_samples[_first_key]) < needed:
                 samples = self.sample_subset(keys=keys, size=needed, random_state=rng)
-                keep = self.evaluate_constraints(samples)
+                keep = self.evaluate_constraints(samples, strict=strict)
                 for key in keys:
                     all_samples[key] = xp.hstack(
                         [all_samples[key], samples[key][keep].flatten()]
@@ -530,9 +570,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

@@ -347,6 +347,59 @@ def conversion_function(parameters):
                 self.assertEqual(N, len(samples2[key]))
             mock_warning.assert_not_called()
 
+    def test_sample_subset_constrained_with_partial_subset(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=-2, maximum=0)
+
+        samples = priors.sample_subset_constrained(
+            keys=["mass_1"], size=16, random_state=self.rng
+        )
+
+        self.assertListEqual(["mass_1"], list(samples.keys()))
+        self.assertEqual(16, len(samples["mass_1"]))
+
+    def test_sample_subset_constrained_full_sample_requires_constraints(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=-2, maximum=0)
+
+        with self.assertRaises(ValueError):
+            priors.sample_subset_constrained(
+                keys=list(priors.keys()), size=1, random_state=self.rng
+            )
+
+    def test_prob_on_partial_subset_requires_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=-2, maximum=0)
+
+        with self.assertRaises(KeyError):
+            priors.prob({"mass_1": 1.5})
+
     def test_sample_with_random_seed(self):
         """
         This test uses the default RNG, so don't specify random_state.
@@ -424,6 +477,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):