-
Notifications
You must be signed in to change notification settings - Fork 123
Description
Summary
Using the plot_corner method with priors=True raises an error when conditional priors are present in the prior dictionary. The root cause is that priors are treated as independent when being plotted, which is not valid for ConditionalPriors.
Error message
For the example I show below, I get the following error:
ValueError: ('Lengths must match to compare', (590,), (300,))
This occurs because the prob method is called independently on each prior, regardless of whether it is conditional (see here).
Suggested fix
An effective solution might be to modify plot_corner for ConditionalPriorDict to jointly sample the priors (e.g., using .sample(...)), rather than plotting them independently. This would respect the conditioning relationships.
Minimal reproducible example
A minimal modification of the linear-regression example demonstrates the issue:
"""
Modified from the `bilby` linear-regression example. The change is that the
prior on `c` becomes a conditional prior, which is conditioned on the value of
`m`.
"""
import bilby
import corner
import numpy
from bilby.core.utils import random
def model(time, m, c):
return time * m + c
def condition_func(reference_params, m):
"""
Functional form of the condition for the prior on `c`. Ensures that `c` is
less than or equal to `m`.
"""
condition = numpy.minimum(reference_params["maximum"], m)
return dict(minimum=reference_params["minimum"], maximum=condition)
def main():
seed = 20250520
random.seed(seed)
label = "linear_regression"
outdir = "data"
injection_parameters = dict(m=0.5, c=0.2)
sampling_frequency = 10
time_duration = 10
time = numpy.arange(0, time_duration, 1 / sampling_frequency)
N = len(time)
sigma = random.rng.normal(1, 0.01, N)
data = model(time, **injection_parameters) + random.rng.normal(0, sigma, N)
likelihood = bilby.likelihood.GaussianLikelihood(time, data, model, sigma)
# Here we define a conditional prior on `c`
priors = bilby.core.prior.ConditionalPriorDict()
priors["m"] = bilby.core.prior.Uniform(minimum=0, maximum=5, name=r"$m$")
priors["c"] = bilby.core.prior.ConditionalUniform(
minimum=-2, maximum=2, name=r"$c$", condition_func=condition_func
)
# Plot the priors
samples = priors.sample(1_000_000)
samples = numpy.vstack([samples["m"], samples["c"]]).T
fig = corner.corner(
samples,
labels=[r"$m$", r"$c$"],
color="tab:blue",
fill_contours=True,
)
fig.tight_layout()
fig.savefig(f"{outdir}/{label}_priors.png")
result = bilby.run_sampler(
likelihood=likelihood,
priors=priors,
sampler="dynesty",
nlive=250,
injection_parameters=injection_parameters,
outdir=outdir,
label=label,
rstate=random.rng,
)
# NOTE: Plotting the priors raises an error directly after sampling
# NOTE: If using the cached result, there is no error but the plot is wrong
result.plot_corner(priors=True)
if __name__ == "__main__":
main()For this example—with a uniform conditional prior—the plotting results in an error because the theta values (of size 300) do not match with the stored prior samples for self.maximum (of size 590). The reason for the 590 self.maximum values is because the sampler saves them during sampling and they are either 2 or the value for m for that prior sample (whichever is smaller).
Let me know if you think it's worth opening a PR to address this.