diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index 3ac54622e..050ee7a01 100644 --- a/bilby/core/prior/dict.py +++ b/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) diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py index 089611aee..6aa4d1a00 100644 --- a/test/core/prior/dict_test.py +++ b/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):