diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index 65688a620..644e2e79a 100644 --- a/bilby/core/prior/dict.py +++ b/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) diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py index cdd996f19..87ce4a633 100644 --- a/test/core/prior/dict_test.py +++ b/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):