From 70d77dd1ff816b32f9e2484159edb0f11ccf4d5a Mon Sep 17 00:00:00 2001 From: Gregory Ashton Date: Mon, 17 Aug 2026 02:12:10 -0700 Subject: [PATCH 1/2] Standarise the handling of dicts and data frames across prob methods Previously, ln_prob adding conditional catches to handle dictionaries and data drames properly. But, this wasn't mirrored across related methods. This adds that mirroring --- bilby/core/prior/dict.py | 15 ++++++++++++--- 1 file changed, 12 insertions(+), 3 deletions(-) diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index 65688a620..4d44780cc 100644 --- a/bilby/core/prior/dict.py +++ b/bilby/core/prior/dict.py @@ -557,8 +557,11 @@ def prob(self, sample, *, normalized=True, xp=None, **kwargs): float: Joint probability of all individual sample probabilities """ - if xp is None: + if xp is None and isinstance(sample, dict): xp = array_module(sample.values()) + elif xp is None: + # assume input is a dataframe + xp = array_module(sample.values) prob = xp.prod(xp.stack([self[key].prob(sample[key], xp=xp) for key in sample]), **kwargs) return self.check_prob(sample, prob, normalized=normalized, xp=xp) @@ -838,8 +841,11 @@ def prob(self, sample, *, normalized=True, xp=None, **kwargs): """ self._prepare_evaluation(*zip(*sample.items())) - if xp is None: + if xp is None and isinstance(sample, dict): xp = array_module(sample.values()) + elif xp is None: + # assume input is a dataframe + xp = array_module(sample.values) res = xp.asarray([ self[key].prob(sample[key], **self.get_required_variables(key), xp=xp) for key in sample @@ -866,8 +872,11 @@ def ln_prob(self, sample, *, axis=None, normalized=True, xp=None): """ self._prepare_evaluation(*zip(*sample.items())) - if xp is None: + if xp is None and isinstance(sample, dict): xp = array_module(sample.values()) + elif xp is None: + # assume input is a dataframe + xp = array_module(sample.values) res = xp.asarray([ self[key].ln_prob(sample[key], **self.get_required_variables(key), xp=xp) for key in sample From d4ee17b17f964a4541e1e14597ba6d5f5b807348 Mon Sep 17 00:00:00 2001 From: Gregory Ashton Date: Mon, 17 Aug 2026 02:34:18 -0700 Subject: [PATCH 2/2] Add tests and fix another bug --- bilby/core/prior/dict.py | 9 ++++++--- test/core/prior/conditional_test.py | 25 +++++++++++++++++++++++++ test/core/prior/dict_test.py | 25 +++++++++++++++++++++++++ 3 files changed, 56 insertions(+), 3 deletions(-) diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index 4d44780cc..f947132ca 100644 --- a/bilby/core/prior/dict.py +++ b/bilby/core/prior/dict.py @@ -63,10 +63,13 @@ def __hash__(self): @xp_wrap def evaluate_constraints(self, sample, *, xp=None): out_sample = self.conversion_function(sample) - try: + if isinstance(out_sample, dict): prob = xp.ones_like(next(iter(out_sample.values())), dtype=bool) - except TypeError: - prob = xp.ones_like(out_sample, dtype=bool) + else: + # assume input is a dataframe; take a single column so the shape + # matches the number of samples, not the whole (n_samples, n_keys) + # frame. + prob = xp.ones_like(out_sample[next(iter(out_sample))], dtype=bool) for key in self: if isinstance(self[key], Constraint) and key in out_sample: prob *= self[key].prob(out_sample[key]) diff --git a/test/core/prior/conditional_test.py b/test/core/prior/conditional_test.py index 2d3a874f0..fea14e1b2 100644 --- a/test/core/prior/conditional_test.py +++ b/test/core/prior/conditional_test.py @@ -286,6 +286,31 @@ def test_ln_prob_illegal_conditions(self): with self.assertRaises(bilby.core.prior.IllegalConditionsException): self.conditional_priors.ln_prob(sample=self.test_sample) + def test_prob_dataframe(self): + """ + Regression test: :code:`ConditionalPriorDict.prob` used to assume + :code:`sample.values()` was callable, which fails for a + :code:`pandas.DataFrame` because :code:`DataFrame.values` is a + property, not a method. :code:`PriorDict.prob`/:code:`ln_prob` handle + this by branching on :code:`isinstance(sample, dict)`; the + :code:`ConditionalPriorDict` overrides need the same branch. + """ + sample = pd.DataFrame( + {key: [float(value), float(value)] for key, value in self.test_sample.items()} + ) + # axis=0 is needed to get one probability per row rather than the + # fully-reduced scalar product over the whole (rows, keys) array. + prob = self.conditional_priors.prob(sample=sample, axis=0) + np.testing.assert_allclose(np.asarray(prob), [float(self.test_value)] * 2) + + def test_ln_prob_dataframe(self): + """See :code:`test_prob_dataframe`; same issue affects :code:`ln_prob`.""" + sample = pd.DataFrame( + {key: [float(value), float(value)] for key, value in self.test_sample.items()} + ) + ln_prob = self.conditional_priors.ln_prob(sample=sample, axis=0) + np.testing.assert_allclose(np.asarray(ln_prob), [float(np.log(self.test_value))] * 2) + def test_sample_subset_all_keys(self): bilby.core.utils.random.seed(5) self.assertDictEqual( diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py index cdd996f19..ad3496d54 100644 --- a/test/core/prior/dict_test.py +++ b/test/core/prior/dict_test.py @@ -4,6 +4,7 @@ import array_api_compat as aac import numpy as np +import pandas as pd import pytest import bilby @@ -389,6 +390,30 @@ def test_ln_prob(self): self.assertEqual(expected, self.prior_set_from_dict.ln_prob(samples)) self.assertEqual(aac.get_namespace(expected), self.xp) + def test_prob_dataframe(self): + """ + Regression test: :code:`prob` used to assume :code:`sample.values()` + was callable, which fails for a :code:`pandas.DataFrame` because + :code:`DataFrame.values` is a property, not a method. + """ + sample = pd.DataFrame({"mass": [0.3, 0.6], "speed": [1.2, 1.5]}) + expected = np.asarray(self.first_prior.prob(sample["mass"].to_numpy())) * np.asarray( + self.second_prior.prob(sample["speed"].to_numpy()) + ) + # axis=0 is needed to get one probability per row rather than the + # fully-reduced scalar product over the whole (rows, keys) array. + prob = self.prior_set_from_dict.prob(sample, axis=0) + np.testing.assert_allclose(np.asarray(prob), expected) + + def test_ln_prob_dataframe(self): + """See :code:`test_prob_dataframe`; same issue affects :code:`ln_prob`.""" + sample = pd.DataFrame({"mass": [0.3, 0.6], "speed": [1.2, 1.5]}) + expected = np.asarray(self.first_prior.ln_prob(sample["mass"].to_numpy())) + np.asarray( + self.second_prior.ln_prob(sample["speed"].to_numpy()) + ) + ln_prob = self.prior_set_from_dict.ln_prob(sample, axis=0) + np.testing.assert_allclose(np.asarray(ln_prob), expected) + def test_rescale(self): theta = [0.5, 0.5, 0.5] expected = [