Skip to content

BUG: plot_corner(priors=True) fails with ConditionalPrior #955

@fgittins

Description

@fgittins

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.

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions