From 456576bf54ddab44650f3b618f4d2728fa5fa05a Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 20 Nov 2024 18:01:36 +0100
Subject: [PATCH 1/7] Changed rescale-method of JointPrior to always return
 correct-size array and update in-place once all keys are requested. Changed
 (Conditional)PriorDict.rescale to always return samples in right shape.

---
 bilby/core/prior/dict.py  | 35 ++++++++++++++------------
 bilby/core/prior/joint.py | 52 ++++++++++++++++++++++++++++-----------
 2 files changed, 57 insertions(+), 30 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index be3d543a9..2908bd083 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -600,18 +600,21 @@ def rescale(self, keys, theta):
         ==========
         keys: list
             List of prior keys to be rescaled
-        theta: list
-            List of randomly drawn values on a unit cube associated with the prior keys
+        theta: dict or 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)
+        theta = [theta[key] for key in keys] if isinstance(theta, dict) else list(theta)
         samples = []
         for key, units in zip(keys, theta):
             samps = self[key].rescale(units)
-            samples += list(np.asarray(samps).flatten())
+            # turns 0d-arrays into scalars
+            samples.append(np.squeeze(samps).tolist())
         return samples
 
     def test_redundancy(self, key, disable_logging=False):
@@ -832,28 +835,28 @@ def rescale(self, keys, theta):
         ==========
         keys: list
             List of prior keys to be rescaled
-        theta: list
-            List of randomly drawn values on a unit cube associated with the prior keys
+        theta: dict or 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)
+        theta = [theta[key] for key in keys] if isinstance(theta, dict) else list(theta)
         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..b088b15ba 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -63,8 +63,9 @@ 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 of the rescale(d) parameters
+        self._rescale_parameters = dict()
+        self._rescaled_parameters = dict()
         self.reset_rescale()
 
         # a list of sampled parameters
@@ -94,7 +95,12 @@ 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._rescale_parameters.values()])
+
+    def set_rescale(self, key, values):
+        values = np.array(values)
+        self._rescale_parameters[key] = values
+        self._rescaled_parameters[key] = np.atleast_1d(np.ones_like(values)) * np.nan
 
     def reset_rescale(self):
         """
@@ -102,7 +108,11 @@ def reset_rescale(self):
         """
 
         for name in self.names:
-            self.rescale_parameters[name] = None
+            self._rescale_parameters[name] = None
+            self._rescaled_parameters[name] = None
+
+    def get_rescaled(self, key):
+        return self._rescaled_parameters[key]
 
     def get_instantiation_dict(self):
         subclass_args = infer_args_from_method(self.__init__)
@@ -303,10 +313,11 @@ 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, values previously set using BaseJointPriorDist.set_rescale() are used.
         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,7 +328,11 @@ def rescale(self, value, **kwargs):
             An vector sample drawn from the multivariate Gaussian
             distribution.
         """
-        samp = np.array(value)
+        if value is None:
+            samp = np.array(list(self._rescale_parameters.values())).T
+        else:
+            samp = np.array(value)
+
         if len(samp.shape) == 1:
             samp = samp.reshape(1, self.num_vars)
 
@@ -327,6 +342,11 @@ def rescale(self, value, **kwargs):
             raise ValueError("Array is the wrong shape")
 
         samp = self._rescale(samp, **kwargs)
+        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)
 
     def _rescale(self, samp, **kwargs):
@@ -790,19 +810,23 @@ def rescale(self, val, **kwargs):
             all kwargs passed to the dist.rescale method
         Returns
         =======
-        float:
-            A sample from the prior parameter.
+        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
         """
 
-        self.dist.rescale_parameters[self.name] = val
+        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)
+            self.dist.rescale(values=None, **kwargs)
+            output = self.dist.get_rescaled(self.name)
             self.dist.reset_rescale()
-            return samples
         else:
-            return []  # return empty list
+            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):
         """

From 8270402a4fde7639c9a3ce1dba28a30b20a3c022 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Fri, 22 Nov 2024 09:55:06 +0100
Subject: [PATCH 2/7] Small fix to rescale of JointPrior

---
 bilby/core/prior/joint.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index b088b15ba..7e6655ee3 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -819,7 +819,7 @@ def rescale(self, val, **kwargs):
         self.dist.set_rescale(self.name, val)
 
         if self.dist.filled_rescale():
-            self.dist.rescale(values=None, **kwargs)
+            self.dist.rescale(value=None, **kwargs)
             output = self.dist.get_rescaled(self.name)
             self.dist.reset_rescale()
         else:

From 27f4ef6b0de3b8a50cc897eddf39fd2f89e87a39 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Mon, 25 Nov 2024 16:50:02 +0100
Subject: [PATCH 3/7] For jointprior rescale, only cast to list once its save
 to loose mutability

---
 bilby/core/prior/dict.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 2908bd083..0490d194e 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -613,8 +613,10 @@ def rescale(self, keys, theta):
         samples = []
         for key, units in zip(keys, theta):
             samps = self[key].rescale(units)
+            samples.append(samps)
+        for i, samps in enumerate(samples):
             # turns 0d-arrays into scalars
-            samples.append(np.squeeze(samps).tolist())
+            samples[i] = np.squeeze(samps).tolist()
         return samples
 
     def test_redundancy(self, key, disable_logging=False):

From 60d7be10874a72399aaf43c602b43c52c599203f Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 21 Jan 2025 16:27:34 +0100
Subject: [PATCH 4/7] Improved nomenclature, added comments, new tests that
 include joint priors

---
 bilby/core/prior/joint.py           | 49 ++++++++++++++++++-----------
 test/core/prior/conditional_test.py |  8 +++--
 test/core/prior/dict_test.py        | 35 ++++++++++++++++++---
 3 files changed, 66 insertions(+), 26 deletions(-)

diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index 7e6655ee3..9d2349343 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -63,9 +63,11 @@ def __init__(self, names, bounds=None):
         self.requested_parameters = dict()
         self.reset_request()
 
-        # a dictionary of the rescale(d) parameters
-        self._rescale_parameters = dict()
-        self._rescaled_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
@@ -95,24 +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):
-        values = np.array(values)
-        self._rescale_parameters[key] = values
-        self._rescaled_parameters[key] = np.atleast_1d(np.ones_like(values)) * np.nan
+        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._rescaled_parameters[name] = None
+            self._current_unit_cube_parameter_values[name] = None
+            self._current_rescaled_parameter_values[name] = None
 
     def get_rescaled(self, key):
-        return self._rescaled_parameters[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__)
@@ -317,7 +319,7 @@ def rescale(self, value, **kwargs):
             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, values previously set using BaseJointPriorDist.set_rescale() are used.
+            If None, the values previously set using BaseJointPriorDist.set_rescale() are used.
         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
@@ -329,9 +331,11 @@ def rescale(self, value, **kwargs):
             distribution.
         """
         if value is None:
-            samp = np.array(list(self._rescale_parameters.values())).T
+            samp = np.array(list(self._current_unit_cube_parameter_values.values())).T
         else:
-            samp = np.array(value)
+            for key, val in zip(self.names, value):
+                self.set_rescale(key, val)
+            samp = np.asarray(value)
 
         if len(samp.shape) == 1:
             samp = samp.reshape(1, self.num_vars)
@@ -342,11 +346,12 @@ def rescale(self, value, **kwargs):
             raise ValueError("Array is the wrong shape")
 
         samp = self._rescale(samp, **kwargs)
-        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]
+        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):
@@ -819,10 +824,16 @@ def rescale(self, val, **kwargs):
         self.dist.set_rescale(self.name, val)
 
         if self.dist.filled_rescale():
+            # 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()
         else:
+            # 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
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(),

From 6c1bd66849ce8bceab25e6c38c0334b690b2eb22 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 21 Jan 2025 17:32:43 +0100
Subject: [PATCH 5/7] Added Warning for previously set rescale_parameters

---
 bilby/core/prior/joint.py | 32 ++++++++++++++++++++++----------
 1 file changed, 22 insertions(+), 10 deletions(-)

diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index 9d2349343..cd2d0c622 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -332,18 +332,18 @@ def rescale(self, value, **kwargs):
         """
         if value is None:
             samp = np.array(list(self._current_unit_cube_parameter_values.values())).T
+            if len(samp.shape) == 1:
+                samp = samp.reshape(1, self.num_vars)
         else:
-            for key, val in zip(self.names, value):
-                self.set_rescale(key, val)
             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")
+            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")
+            for key, val in zip(self.names, samp.T):
+                self.set_rescale(key, val)
 
         samp = self._rescale(samp, **kwargs)
         for i, key in enumerate(self.names):
@@ -821,6 +821,18 @@ def rescale(self, val, **kwargs):
             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():

From 2ddf678007f231f7c6855f5b00b6719abaf943c3 Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Tue, 21 Jan 2025 18:22:18 +0100
Subject: [PATCH 6/7] removed dict as input type for rescale theta

---
 bilby/core/prior/dict.py | 18 +++++++++++-------
 1 file changed, 11 insertions(+), 7 deletions(-)

diff --git a/bilby/core/prior/dict.py b/bilby/core/prior/dict.py
index 0490d194e..943fb698a 100644
--- a/bilby/core/prior/dict.py
+++ b/bilby/core/prior/dict.py
@@ -598,9 +598,9 @@ def rescale(self, keys, theta):
 
         Parameters
         ==========
-        keys: list
+        keys: array-like
             List of prior keys to be rescaled
-        theta: dict or array-like
+        theta: array-like
             Randomly drawn values on a unit cube associated with the prior keys
 
         Returns
@@ -609,7 +609,8 @@ def rescale(self, keys, theta):
             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 = [theta[key] for key in keys] if isinstance(theta, dict) else list(theta)
+        if isinstance(theta, {}.values().__class__):
+            theta = list(theta)
         samples = []
         for key, units in zip(keys, theta):
             samps = self[key].rescale(units)
@@ -835,9 +836,9 @@ def rescale(self, keys, theta):
 
         Parameters
         ==========
-        keys: list
+        keys: array-like
             List of prior keys to be rescaled
-        theta: dict or array-like
+        theta: array-like
             Randomly drawn values on a unit cube associated with the prior keys
 
         Returns
@@ -846,8 +847,11 @@ def rescale(self, keys, theta):
             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 = [theta[key] for key in keys] if isinstance(theta, dict) else 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()

From 1a1ec8735b26f18b1e78bb6a52f59100dc79f67b Mon Sep 17 00:00:00 2001
From: Michael Jasper Martins <jasper.martins@aei.mpg.de>
Date: Wed, 22 Jan 2025 12:05:17 +0100
Subject: [PATCH 7/7] Better Error Handling

---
 bilby/core/prior/joint.py | 23 +++++++++++++----------
 1 file changed, 13 insertions(+), 10 deletions(-)

diff --git a/bilby/core/prior/joint.py b/bilby/core/prior/joint.py
index cd2d0c622..1087754b3 100644
--- a/bilby/core/prior/joint.py
+++ b/bilby/core/prior/joint.py
@@ -319,7 +319,8 @@ def rescale(self, value, **kwargs):
             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.
+            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
@@ -331,7 +332,9 @@ def rescale(self, value, **kwargs):
             distribution.
         """
         if value is None:
-            samp = np.array(list(self._current_unit_cube_parameter_values.values())).T
+            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:
@@ -342,16 +345,16 @@ def rescale(self, value, **kwargs):
                 raise ValueError("Array is the wrong shape")
             elif samp.shape[1] != self.num_vars:
                 raise ValueError("Array is the wrong shape")
-            for key, val in zip(self.names, samp.T):
-                self.set_rescale(key, val)
 
         samp = self._rescale(samp, **kwargs)
-        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]
+        # 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):