Conserve rescale shape

[JasperMartins]

This PR updates the rescale functions of PriorDict and ConditionalPriorDict to preserve the correct shape of the rescaled samples.

The PR also updates the behavior of JointPrior.rescale and BaseJointPriorDist.

JointPrior.rescale now returns a mutable numpy array. If not all required keys have been requested for rescaling, the arrays is filled with np.nan. BaseJointPriorDist keeps track of this array. Once all required keys are requested, the rescale operation is performed, and the returned arrays are populated with the rescaled values.
This change enables out-of-order sets of keys in rescale - previously priordict.rescale(keys=["a", "JointPrior_b", "c", "JointPrior_d"], theta=...) where "JointPrior_b" and "JointPrior_d" share the same dist would have resulted in an order [*samples_a, *samples_c, *samples_b, *samples_d].

Relates to the discussion in #850

[JasperMartins]

I wonder if it would not make more sense to preserve numpy arrays to make it easier to define conversion_functions and so on without needing to cast to array again

[ColmTalbot]

I think it is would be worth returning arrays, I think it's worth accepting that this is going to not be fully backward compatible.

[ColmTalbot]

Thanks for opening this, I think this will be a good change, and I'm in favour of also not pushing everything through lists.

I don't really follow the logic in the joint prior, so some additional docstrings/comments would probably help there.

[ColmTalbot]

I think it is would be worth returning arrays, I think it's worth accepting that this is going to not be fully backward compatible.

[ColmTalbot]

It might be worth returning the same type as the input, i.e., return a dict if a dict is passed.

Another option would be to add a specific method to handle dicts, I've thought about doing this a few times before, e.g., PriorDict.rescale_dict that just wraps PriorDict.rescale.

[JasperMartins]

This could also apply to the conservation of numpy arrays, ie return a dict, list or array depending on the input.

[ColmTalbot]

I can't see where these values are set to a non-trivial value, what is the distinction between these two?

[JasperMartins]

The values are updated in-place in the rescale function:

bilby/bilby/core/prior/joint.py

Lines 345 to 350 in 4e1a3d2

	if value is None:
	for i, key in enumerate(self.names):
	output = self.get_rescaled(key)
	# update in-place for proper handling in PriorDict-instances
	output[:] = samp[:, i]
	return np.squeeze(samp)

bilby/bilby/core/prior/joint.py

Lines 819 to 829 in 4e1a3d2

	self.dist.set_rescale(self.name, val)
	if self.dist.filled_rescale():
	self.dist.rescale(value=None, **kwargs)
	output = self.dist.get_rescaled(self.name)
	self.dist.reset_rescale()
	else:
	output = self.dist.get_rescaled(self.name)
	# have to return raw output to conserve in-place modifications
	return output

So, in self._rescaled_parameters, the values of the rescaled parameters are stored once the rescaling operation is performed. Prior to this, the array contains only np.NaN to clarify that the rescaling has not yet happened. The per-key numpy arrays stored in self._rescaled_parameters are returned as references in JointPrior.rescale. Once all keys have been requested, the values in the arrays are updated in place, meaning that they will also be updated in the array returned by JointPrior.rescale.

I agree that this should be documented better, particularly to ensure that PriorDict.rescale is not updated in a way that loses track of the reference to the arrays stored in BaseJointPriorDict._rescaled_parameters in the future.

[JasperMartins]

I have updated the changes to joint priors with more comments / docstrings and changed variable names to be more reflective of what happens. I have also noticed the implementation attempts to solve a similar issue as #848 when rescaling joint priors.

For now, I have left the return-values of PriorDict.rescale() and ConditionalPriorDict.rescale() as lists. Changing to arrays should, I think, be accompanied by a larger-scale switch to numpy arrays for the return values of methods of Prior() that I don't have time for right at the moment.