diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py index be3d543a9..943fb698a 100644 --- a/bilby/core/prior/dict.py +++ b/bilby/core/prior/dict.py @@ -598,20 +598,26 @@ def rescale(self, keys, theta): Parameters ========== - keys: list + keys: array-like List of prior keys to be rescaled - theta: list - List of randomly drawn values on a unit cube associated with the prior keys + theta: array-like + Randomly drawn values on a unit cube associated with the prior keys Returns ======= - list: List of floats containing the rescaled sample + list: + If theta is 1D, returns list of floats containing the rescaled sample. + If theta is 2D, returns list of lists containing the rescaled samples. """ - theta = list(theta) + if isinstance(theta, {}.values().__class__): + theta = list(theta) samples = [] for key, units in zip(keys, theta): samps = self[key].rescale(units) - samples += list(np.asarray(samps).flatten()) + samples.append(samps) + for i, samps in enumerate(samples): + # turns 0d-arrays into scalars + samples[i] = np.squeeze(samps).tolist() return samples def test_redundancy(self, key, disable_logging=False): @@ -830,30 +836,33 @@ def rescale(self, keys, theta): Parameters ========== - keys: list + keys: array-like List of prior keys to be rescaled - theta: list - List of randomly drawn values on a unit cube associated with the prior keys + theta: array-like + Randomly drawn values on a unit cube associated with the prior keys Returns ======= - list: List of floats containing the rescaled sample + list: + If theta is float for each key, returns list of floats containing the rescaled sample. + If theta is array-like for each key, returns list of lists containing the rescaled samples. """ - keys = list(keys) - theta = list(theta) + if isinstance(theta, {}.values().__class__): + theta = list(theta) + if isinstance(keys, {}.keys().__class__): + keys = list(keys) + self._check_resolved() self._update_rescale_keys(keys) result = dict() - for key, index in zip( - self.sorted_keys_without_fixed_parameters, self._rescale_indexes - ): - result[key] = self[key].rescale( - theta[index], **self.get_required_variables(key) - ) + for key, index in zip(self.sorted_keys_without_fixed_parameters, self._rescale_indexes): + result[key] = self[key].rescale(theta[index], **self.get_required_variables(key)) self[key].least_recently_sampled = result[key] samples = [] for key in keys: - samples += list(np.asarray(result[key]).flatten()) + # turns 0d-arrays into scalars + res = np.squeeze(result[key]).tolist() + samples.append(res) return samples def _update_rescale_keys(self, keys): diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py index 43c8913e3..1087754b3 100644 --- a/bilby/core/prior/joint.py +++ b/bilby/core/prior/joint.py @@ -63,8 +63,11 @@ def __init__(self, names, bounds=None): self.requested_parameters = dict() self.reset_request() - # a dictionary of the rescaled parameters - self.rescale_parameters = dict() + # a dictionary that stores the unit-cube values of parameters for later rescaling + self._current_unit_cube_parameter_values = dict() + # a dictionary of arrays that are used as intermediate return values of JointPrior.rescale() + # and updated in-place once all parameters have been requested + self._current_rescaled_parameter_values = dict() self.reset_rescale() # a list of sampled parameters @@ -94,15 +97,24 @@ def filled_rescale(self): Check if all the rescaled parameters have been filled. """ - return not np.any([val is None for val in self.rescale_parameters.values()]) + return not np.any([val is None for val in self._current_unit_cube_parameter_values.values()]) + + def set_rescale(self, key, values): + self._current_unit_cube_parameter_values[key] = np.array(values) + self._current_rescaled_parameter_values[key] = np.full_like(values, np.nan, dtype=float) def reset_rescale(self): """ Reset the rescaled parameters to None. """ - for name in self.names: - self.rescale_parameters[name] = None + self._current_unit_cube_parameter_values[name] = None + self._current_rescaled_parameter_values[name] = None + + def get_rescaled(self, key): + """Return an array that will be updated in-place once the rescale-operation + has been performed.""" + return self._current_rescaled_parameter_values[key] def get_instantiation_dict(self): subclass_args = infer_args_from_method(self.__init__) @@ -303,10 +315,12 @@ def rescale(self, value, **kwargs): Parameters ========== - value: array - A 1d vector sample (one for each parameter) drawn from a uniform + value: array or None + If given, a 1d vector sample (one for each parameter) drawn from a uniform distribution between 0 and 1, or a 2d NxM array of samples where N is the number of samples and M is the number of parameters. + If None, the values previously set using BaseJointPriorDist.set_rescale() are used, + the result is stored and can be accessed using get_rescaled(). kwargs: dict All keyword args that need to be passed to _rescale method, these keyword args are called in the JointPrior rescale methods for each parameter @@ -317,16 +331,30 @@ def rescale(self, value, **kwargs): An vector sample drawn from the multivariate Gaussian distribution. """ - samp = np.array(value) - if len(samp.shape) == 1: - samp = samp.reshape(1, self.num_vars) - - if len(samp.shape) != 2: - raise ValueError("Array is the wrong shape") - elif samp.shape[1] != self.num_vars: - raise ValueError("Array is the wrong shape") + if value is None: + if not self.filled_rescale(): + raise ValueError("Attempting to rescale from stored values without having set all required values.") + samp = np.array([self._current_unit_cube_parameter_values[key] for key in self.names]).T + if len(samp.shape) == 1: + samp = samp.reshape(1, self.num_vars) + else: + samp = np.asarray(value) + if len(samp.shape) == 1: + samp = samp.reshape(1, self.num_vars) + if len(samp.shape) != 2: + raise ValueError("Array is the wrong shape") + elif samp.shape[1] != self.num_vars: + raise ValueError("Array is the wrong shape") samp = self._rescale(samp, **kwargs) + # only store result if the rescale was done with saved unit-cube values + if value is None: + for i, key in enumerate(self.names): + # get the numpy array used for indermediate outputs + # prior to a full rescale-operation + output = self.get_rescaled(key) + # update the array in-place + output[...] = samp[:, i] return np.squeeze(samp) def _rescale(self, samp, **kwargs): @@ -790,19 +818,41 @@ def rescale(self, val, **kwargs): all kwargs passed to the dist.rescale method Returns ======= - float: - A sample from the prior parameter. - """ - - self.dist.rescale_parameters[self.name] = val + np.ndarray: + The samples from the prior parameter. If not all names in "dist" have been filled, + the array contains only np.nan. *This* specific array instance will be filled with + the rescaled value once all parameters have been requested + """ + + if self.dist.get_rescaled(self.name) is not None: + import warnings + warnings.warn( + f"Rescale values for {self.name} in {self.dist} have already been set, " + "indicating that another rescale-operation is in progress.\n" + "Call dist.reset_rescale() on the joint prior distribution associated " + "with this prior after using dist.rescale() and make sure all parameters " + "necessary for rescaling have been requested in PriorDict.rescale().\n" + "Resetting now.", + RuntimeWarning + ) + self.dist.reset_rescale() + self.dist.set_rescale(self.name, val) if self.dist.filled_rescale(): - values = np.array(list(self.dist.rescale_parameters.values())).T - samples = self.dist.rescale(values, **kwargs) + # If all names have been filled, perform rescale operation + self.dist.rescale(value=None, **kwargs) + # get the rescaled values for the requested parameter + output = self.dist.get_rescaled(self.name) + # reset the rescale operation self.dist.reset_rescale() - return samples else: - return [] # return empty list + # If not all names have been filled, return a *numpy array* + # filled only with `np.nan`. Once all names have been requested, + # this array is updated *in-place* with the rescaled values. + output = self.dist.get_rescaled(self.name) + + # have to return raw output to conserve in-place modifications + return output def sample(self, size=1, **kwargs): """ diff --git a/test/core/prior/conditional_test.py b/test/core/prior/conditional_test.py index 20c0cda93..abd83430f 100644 --- a/test/core/prior/conditional_test.py +++ b/test/core/prior/conditional_test.py @@ -334,7 +334,7 @@ def test_rescale_with_joint_prior(self): # set multivariate Gaussian distribution names = ["mvgvar_0", "mvgvar_1"] - mu = [[0.79, -0.83]] + mu = [[1, 1]] cov = [[[0.03, 0.], [0., 0.04]]] mvg = bilby.core.prior.MultivariateGaussianDist(names, mus=mu, covs=cov) @@ -349,7 +349,7 @@ def test_rescale_with_joint_prior(self): ) ) - ref_variables = list(self.test_sample.values()) + [0.4, 0.1] + ref_variables = list(self.test_sample.values()) + [0.5, 0.5] keys = list(self.test_sample.keys()) + names res = priordict.rescale(keys=keys, theta=ref_variables) @@ -359,9 +359,11 @@ def test_rescale_with_joint_prior(self): # check conditional values are still as expected expected = [self.test_sample["var_0"]] + self.assertFalse(np.any(np.isnan(res))) for ii in range(1, 4): expected.append(expected[-1] * self.test_sample[f"var_{ii}"]) - self.assertListEqual(expected, res[0:4]) + expected.extend([1, 1]) + self.assertListEqual(expected, res) def test_cdf(self): """ diff --git a/test/core/prior/dict_test.py b/test/core/prior/dict_test.py index 08e730bbe..a3c8e1cd6 100644 --- a/test/core/prior/dict_test.py +++ b/test/core/prior/dict_test.py @@ -33,8 +33,17 @@ def setUp(self): name="b", alpha=3, minimum=1, maximum=2, unit="m/s", boundary=None ) self.third_prior = bilby.core.prior.DeltaFunction(name="c", peak=42, unit="m") + + mvg = bilby.core.prior.MultivariateGaussianDist( + names=["testa", "testb"], + mus=[1, 1], + covs=np.array([[2.0, 0.5], [0.5, 2.0]]), + weights=1.0, + ) + self.testa = bilby.core.prior.MultivariateGaussian(dist=mvg, name="testa", unit="unit") + self.testb = bilby.core.prior.MultivariateGaussian(dist=mvg, name="testb", unit="unit") self.priors = dict( - mass=self.first_prior, speed=self.second_prior, length=self.third_prior + mass=self.first_prior, speed=self.second_prior, length=self.third_prior, testa=self.testa, testb=self.testb ) self.prior_set_from_dict = bilby.core.prior.PriorDict(dictionary=self.priors) self.default_prior_file = os.path.join( @@ -70,7 +79,7 @@ def test_prior_set_is_dict(self): self.assertIsInstance(self.prior_set_from_dict, dict) def test_prior_set_has_correct_length(self): - self.assertEqual(3, len(self.prior_set_from_dict)) + self.assertEqual(5, len(self.prior_set_from_dict)) def test_prior_set_has_expected_priors(self): self.assertDictEqual(self.priors, dict(self.prior_set_from_dict)) @@ -160,6 +169,12 @@ def test_to_file(self): "unit='m/s', boundary=None)\n", "mass = Uniform(minimum=0, maximum=1, name='a', latex_label='a', " "unit='kg', boundary=None)\n", + "testa_testb_mvg = MultivariateGaussianDist(names=['testa', 'testb'], nmodes=1, mus=[[1, 1]], " + "sigmas=[[1.4142135623730951, 1.4142135623730951]], " + "corrcoefs=[[[0.9999999999999998, 0.24999999999999994], [0.24999999999999994, 0.9999999999999998]]], " + "covs=[[[2.0, 0.5], [0.5, 2.0]]], weights=[1.0], bounds={'testa': (-inf, inf), 'testb': (-inf, inf)})\n", + "testa = MultivariateGaussian(dist=testa_testb_mvg, name='testa', latex_label='testa', unit='unit')\n", + "testb = MultivariateGaussian(dist=testa_testb_mvg, name='testb', latex_label='testb', unit='unit')\n", ] self.prior_set_from_dict.to_file(outdir="prior_files", label="to_file_test") with open("prior_files/to_file_test.prior") as f: @@ -178,6 +193,13 @@ def test_from_dict_with_string(self): self.assertDictEqual(self.prior_set_from_dict, from_dict) def test_convert_floats_to_delta_functions(self): + mvg = bilby.core.prior.MultivariateGaussianDist( + names=["testa", "testb"], + mus=[1, 1], + covs=np.array([[2.0, 0.5], [0.5, 2.0]]), + weights=1.0, + ) + self.prior_set_from_dict["d"] = 5 self.prior_set_from_dict["e"] = 7.3 self.prior_set_from_dict["f"] = "unconvertable" @@ -190,6 +212,8 @@ def test_convert_floats_to_delta_functions(self): name="b", alpha=3, minimum=1, maximum=2, unit="m/s", boundary=None ), length=bilby.core.prior.DeltaFunction(name="c", peak=42, unit="m"), + testa=bilby.core.prior.MultivariateGaussian(dist=mvg, name="testa", unit="unit"), + testb=bilby.core.prior.MultivariateGaussian(dist=mvg, name="testb", unit="unit"), d=bilby.core.prior.DeltaFunction(peak=5), e=bilby.core.prior.DeltaFunction(peak=7.3), f="unconvertable", @@ -321,12 +345,15 @@ def test_ln_prob(self): self.assertEqual(expected, self.prior_set_from_dict.ln_prob(samples)) def test_rescale(self): - theta = [0.5, 0.5, 0.5] + theta = [0.5, 0.5, 0.5, 0.5, 0.5] expected = [ self.first_prior.rescale(0.5), self.second_prior.rescale(0.5), self.third_prior.rescale(0.5), + self.testa.rescale(0.5), + self.testb.rescale(0.5) ] + assert not np.any(np.isnan(expected)) self.assertListEqual( sorted(expected), sorted( @@ -342,7 +369,7 @@ def test_cdf(self): Note that the format of inputs/outputs is different between the two methods. """ - sample = self.prior_set_from_dict.sample() + sample = self.prior_set_from_dict.sample_subset(keys=["length", "speed", "mass"]) original = np.array(list(sample.values())) new = np.array(self.prior_set_from_dict.rescale( sample.keys(),