Conserve rescale shape

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):

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):
         """

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):
         """

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(),