Conversion function keys added to samples when sampling with `size=1`
When sampling a prior with a conversion_function
with size=1
, any keys added inside the conversion_function
end up in the sampled parameters, while sampling with size > 1
keeps them out. Consider the following example:
from bilby.core.prior import Constraint, PriorDict, Uniform
def mass_constraint(samples):
samples["mass_ratio"] = samples["mass_1"] / samples["mass_2"]
return samples
prior = PriorDict(conversion_function=mass_constraint)
prior["mass_1"] = Uniform(name="mass_1", minimum=5, maximum=50)
prior["mass_2"] = Uniform(name="mass_2", minimum=5, maximum=50)
prior["mass_ratio"] = Constraint(0.02, 1)
print(prior.sample(1))
print(prior.sample(2))
Output:
{'mass_1': array([25.29565642]), 'mass_2': array([48.40744659]), 'mass_ratio': array([0.52255713])}
{'mass_1': array([13.46337225, 44.1621476 ]), 'mass_2': array([14.61401658, 45.24270904])}
Obviously you could avoid this altogether by just doing a samples = samples.copy()
inside the mass_constraint
function, but at least for our use case the number of samples we need to draw can vary by orders of magnitude, so for larger arrays this copy will be nontrivial. Plus it seems like consistency here would be useful.
I'm not sure which behavior would be considered more desirable, but I would assume it's the latter in the example above? The difference seems to come inside PriorDict.sample_subset_constrained
, where when size == 1
you return the samples as-is after calling self.evaluate_constraints
(which adds the erroneous keys), but in the second clause for size > 1
you explicitly loop through self.keys
to append to all_samples
.
I can make a quick MR which adopts the latter behavior for the size == 1
case, but it should be easy to change to adopt the former behavior if that's considered more desirable by the developers.