Standarise the handling of dicts and data frames across prob methods

[GregoryAshton]

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

[Copilot]

Pull request overview

This pull request standardizes how PriorDict/ConditionalPriorDict probability methods infer the array backend (xp) when samples are provided as either dictionaries or DataFrame-like objects, aligning behavior across related probability APIs.

Changes:

• Updates prob (and conditional variants) to infer xp correctly when sample is a dict vs. a DataFrame-like object.
• Mirrors previously-added DataFrame/dict handling patterns across additional probability methods to keep the API consistent.

Suppressed comments (2)

bilby/core/prior/dict.py:848

• Same xp inference pattern as above: treating any non-dict sample as a DataFrame and reading sample.values can produce confusing errors for other mapping/array-like inputs. array_module(sample) already covers dicts and pandas objects, so it’s safer and simpler.

        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)

bilby/core/prior/dict.py:879

• Same xp inference issue here: the current branch assumes any non-dict sample is a DataFrame and accesses sample.values. Using array_module(sample) avoids relying on a specific attribute and matches array_module’s built-in handling for dict/pandas inputs.

        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)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[GregoryAshton]

Caude Summary

Fixes a regression introduced in 612c46b ("Support non-numpy array backends (#886)"), which added dict-vs-DataFrame handling to PriorDict.ln_prob but left three sibling methods on the old, dict-only code path. Any of them called with a pandas.DataFrame sample breaks, because DataFrame.values is a property, not a method like dict.values.

Fixed (bilby/core/prior/dict.py)

• PriorDict.prob — TypeError: 'numpy.ndarray' object is not callable
• ConditionalPriorDict.prob — same
• ConditionalPriorDict.ln_prob — same
• PriorDict.evaluate_constraints — different failure mode: wrapped in a try/except TypeError, so instead of crashing it silently falls back to ones_like(out_sample) on the whole DataFrame, producing shape (n_samples, n_keys) instead of (n_samples,). Depending on what's downstream this either raises a broadcasting ValueError or silently returns wrong-shaped output.

The first three get the same isinstance(sample, dict) / else-assume-DataFrame branch PriorDict.ln_prob already used. evaluate_constraints gets the equivalent branch, replacing the try/except with an explicit type check and taking one column (out_sample[next(iter(out_sample))]) instead of the whole frame as the shape template.

evaluate_constraints matters beyond .prob(): check_ln_prob also calls it, but only when some sample has non-finite ln_prob. So ConditionalPriorDict.ln_prob(dataframe) can look fixed on well-behaved input but still break the moment a batch contains a sample that violates a Constraint prior.

Tests added (test/core/prior/dict_test.py, test/core/prior/conditional_test.py)

test_prob_dataframe / test_ln_prob_dataframe on both TestPriorDict and TestConditionalPriorDict, calling each method with a pandas.DataFrame (including axis=0 for the batched case, matching ln_prob's existing convention). Each fails against main with the original TypeError (3 of 4 — PriorDict.ln_prob was already fixed pre-branch) and passes with the fix.

Verification

• Reproduced the evaluate_constraints shape bug directly: a 1000-row batch with half the rows violating a Constraint prior raised ValueError: operands could not be broadcast together with shapes (1000, 25) (1000,) before the fix; after, ln_prob returns the expected (1000,) shape and downstream masking/filtering behaves correctly.
• Full test/core/prior/ + test/gw/prior_test.py suite: 199 passed, plus 2 pre-existing failures unrelated to this change (missing healpy, missing jax in the test environment), confirmed present on stock main too.

[codecov]

Codecov Report

❌ Patch coverage is 72.72727% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 72.14%. Comparing base (a139afa) to head (d4ee17b).

Files with missing lines	Patch %	Lines
bilby/core/prior/dict.py	72.72%	0 Missing and 3 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #1128      +/-   ##
==========================================
+ Coverage   72.12%   72.14%   +0.02%     
==========================================
  Files          86       86              
  Lines       15371    15376       +5     
  Branches     2333     2337       +4     
==========================================
+ Hits        11086    11093       +7     
+ Misses       3578     3576       -2     
  Partials      707      707              

Flag	Coverage Δ
python311	72.07% <72.72%> (+0.02%)	⬆️
python312	72.07% <72.72%> (-0.01%)	⬇️
python313	72.09% <72.72%> (+0.01%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.