BUG: Inconsistent return type between MultivariateGaussian and analytical priors

[hank-hua]

Describe the bug
PriorDict.ln_prob fails when the dictionary includes both MultivariateGaussian and Uniform/Gaussian priors. MVG priors return a scalar float for a single sample, whereas Uniform and Gaussian priors produce a numpy array with one element. The resulting mixed scalar/array outputs crashes np.sum.

To Reproduce

from bilby.core.prior import PriorDict, Uniform, Gaussian, MultivariateGaussian, MultivariateGaussianDist    
import bilby
print(bilby.__version__)
bilby.utils.random.seed(123)
prior_dict = PriorDict()                                                 
prior_dict["uniform"] = Uniform(name="uniform", minimum=0, maximum=1)
prior_dict["gaussian"] = Gaussian(name="gaussian", mu=0, sigma=1)
mean = [0, 0]                 
cov = [[1, 0.5], [0.5, 1]]              
names = ["mvg_param1", "mvg_param2"]                         
mvg = MultivariateGaussianDist(names, mus=mean, covs=cov)                       
prior_dict["mvg_param1"] = MultivariateGaussian(name="mvg_param1", dist=mvg)    
prior_dict["mvg_param2"] = MultivariateGaussian(name="mvg_param2", dist=mvg)
samples = prior_dict.sample(1)
print(f"{samples=}")   
for key in samples:                                                                         
    print(f"ln_prob for {key}: {samples[key]}, {prior_dict[key].ln_prob(samples[key])}")
print(f"{prior_dict.ln_prob(samples)=}")

Error message

2.7.1
samples={'uniform': array([0.68235186]), 'gaussian': array([-1.6088825]), 'mvg_param1': -0.21827224300008818, 'mvg_param2': -1.11710148572024}
ln_prob for uniform: [0.68235186], [0.]
ln_prob for gaussian: [-1.6088825], [-2.21318998]
ln_prob for mvg_param1: -0.21827224300008818, 0.0
ln_prob for mvg_param2: -1.11710148572024, -2.3951868665273013
Traceback (most recent call last):
  File "/home/h/Documents/test.py", line 18, in <module>
    print(f"{prior_dict.ln_prob(samples)=}")
  File "/home/h/.local/lib/python3.10/site-packages/bilby/core/prior/dict.py", line 579, in ln_prob
    ln_prob = np.sum([self[key].ln_prob(sample[key]) for key in sample], axis=axis)
  File "/home/h/.local/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 2313, in sum
    return _wrapreduction(a, np.add, 'sum', axis, dtype, out, keepdims=keepdims,
  File "/home/h/.local/lib/python3.10/site-packages/numpy/core/fromnumeric.py", line 88, in _wrapreduction
    return ufunc.reduce(obj, axis, dtype, out, **passkwargs)
ValueError: setting an array element with a sequence. The requested array has an inhomogeneous shape after 1 dimensions. The detected shape was (4,) + inhomogeneous part.

Expected behavior
The expected ln_prob should be around -2.2 + -2.4 = -4.6.

Suggested solution (optional)
A temporary workaround would be to make the return of ln_prob and sample in MultivariateGaussian be arrays.

from bilby.core.prior import PriorDict, Uniform, Gaussian, MultivariateGaussian, MultivariateGaussianDist
import bilby
print(bilby.__version__)
bilby.utils.random.seed(123)
import numpy as np

class MultivariateGaussianArray(MultivariateGaussian):
    def ln_prob(self, val):
        ln_prob = super().ln_prob(val)
        if isinstance(ln_prob, float):
            ln_prob = np.array([ln_prob])
        return ln_prob

    def sample(self, size=1, **kwargs):
        samps = super().sample(size=size, **kwargs)
        if size == 1:
            samps = np.array([samps])
        return samps

prior_dict = PriorDict()
prior_dict["uniform"] = Uniform(name="uniform", minimum=0, maximum=1)
prior_dict["gaussian"] = Gaussian(name="gaussian", mu=0, sigma=1)
mean = [0, 0]
cov = [[1, 0.5], [0.5, 1]]
names = ["mvg_param1", "mvg_param2"]
mvg = MultivariateGaussianDist(names, mus=mean, covs=cov)
prior_dict["mvg_param1"] = MultivariateGaussianArray(name="mvg_param1", dist=mvg)
prior_dict["mvg_param2"] = MultivariateGaussianArray(name="mvg_param2", dist=mvg)
samples = prior_dict.sample(1)
print(f"{samples=}")
for key in samples:
    print(f"ln_prob for {key}: {samples[key]}, {prior_dict[key].ln_prob(samples[key])}")
print(f"{prior_dict.ln_prob(samples)=}")

ln_prob then works as expected.

2.7.1
samples={'uniform': array([0.68235186]), 'gaussian': array([-1.6088825]), 'mvg_param1': array([-0.21827224]), 'mvg_param2': array([-1.11710149])}
ln_prob for uniform: [0.68235186], [0.]
ln_prob for gaussian: [-1.6088825], [-2.21318998]
ln_prob for mvg_param1: [-0.21827224], [0.]
ln_prob for mvg_param2: [-1.11710149], [-2.39518687]
prior_dict.ln_prob(samples)=-4.608376842798096

More permanently, the checks in ln_prob and sample could be removed to match the behaviour of the base class (or alternatively a check in the base class could be added to return a float for a single sample).

Environment (please complete the following information):

• Bilby version: 2.7.1
• Installation method [e.g. conda, pip, source]: pip