Standarise the handling of dicts and data frames across prob methods

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])
@@ -557,8 +560,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 +844,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 +875,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

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(

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 = [