Repository: rsagroup/rsatoolbox

Issue:
rescale() custom threshold
Rescale doesn't always converge to 1e-8. Users should be able to set a custom threshold.

Discussion:

Patch:
diff --git a/src/rsatoolbox/rdm/combine.py b/src/rsatoolbox/rdm/combine.py
index 6a8a8247..e8fc4a8b 100644
--- a/src/rsatoolbox/rdm/combine.py
+++ b/src/rsatoolbox/rdm/combine.py
@@ -96,7 +96,7 @@ def from_partials(
     )
 
 
-def rescale(rdms, method: str = 'evidence'):
+def rescale(rdms, method: str = 'evidence', threshold=1e-8):
     """Bring RDMs closer together
 
     Iteratively scales RDMs based on pairs in-common.
@@ -105,11 +105,15 @@ def rescale(rdms, method: str = 'evidence'):
     Args:
         method (str, optional): One of 'evidence', 'setsize' or
             'simple'. Defaults to 'evidence'.
+        threshold (float): Stop iterating when the sum of squares 
+            difference between iterations is smaller than this value.
+            A smaller value means more iterations, but the algorithm
+            may not always converge.
 
     Returns:
         RDMs: RDMs object with the aligned RDMs
     """
-    aligned, weights = _rescale(rdms.dissimilarities, method)
+    aligned, weights = _rescale(rdms.dissimilarities, method, threshold)
     rdm_descriptors = deepcopy(rdms.rdm_descriptors)
     if weights is not None:
         rdm_descriptors['rescalingWeights'] = weights
@@ -166,7 +170,7 @@ def _scale(vectors: ndarray) -> ndarray:
     return vectors / sqrt(_ss(vectors))
 
 
-def _rescale(dissim: ndarray, method: str) -> Tuple[ndarray, ndarray]:
+def _rescale(dissim: ndarray, method: str, threshold=1e-8) -> Tuple[ndarray, ndarray]:
     """Rescale RDM vectors
 
     See :meth:`rsatoolbox.rdm.combine.rescale`
@@ -191,7 +195,7 @@ def _rescale(dissim: ndarray, method: str) -> Tuple[ndarray, ndarray]:
 
     current_estimate = _scale(_mean(dissim))
     prev_estimate = np.full([n_conds, ], -inf)
-    while _ss(current_estimate - prev_estimate) > 1e-8:
+    while _ss(current_estimate - prev_estimate) > threshold:
         prev_estimate = current_estimate.copy()
         tiled_estimate = np.tile(current_estimate, [n_rdms, 1])
         tiled_estimate[np.isnan(dissim)] = nan


Test patch:
diff --git a/tests/test_rdms_combine.py b/tests/test_rdms_combine.py
index 9dd79eae..99ad3dec 100644
--- a/tests/test_rdms_combine.py
+++ b/tests/test_rdms_combine.py
@@ -2,7 +2,7 @@
 """
 # pylint: disable=import-outside-toplevel, no-self-use
 from unittest import TestCase
-from numpy import array, nan, isnan
+from numpy import array, nan, isnan, mean, abs as _abs, diff
 from numpy.testing import assert_almost_equal, assert_array_equal
 from scipy.stats import pearsonr
 
@@ -70,6 +70,28 @@ class RdmsCombineTests(TestCase):
             decimal=4
         )
 
+    def test_rescale_threshold(self):
+        """The rescale function bring the RDMs as close together as possible
+        """
+        from rsatoolbox.rdm.rdms import RDMs
+        from rsatoolbox.rdm.combine import rescale
+        partial_rdms = RDMs(
+            dissimilarities=array([
+                [  1,   2, nan,   3, nan, nan],
+                [nan, nan, nan,   4,   5, nan],
+            ])
+        )
+        ## high threshold, fewer iterations, substantial difference remaining
+        rescaled_rdms = rescale(partial_rdms, method='simple', threshold=10)
+        common_pair = rescaled_rdms.dissimilarities[:, 3]
+        rel_diff = _abs(diff(common_pair)/mean(common_pair))
+        self.assertGreater(rel_diff[0], 0.1)
+        ## low threshold, more iterations, difference small
+        rescaled_rdms = rescale(partial_rdms, method='simple', threshold=0.00001)
+        common_pair = rescaled_rdms.dissimilarities[:, 3]
+        rel_diff = _abs(diff(common_pair)/mean(common_pair))
+        self.assertLess(rel_diff[0], 0.01)
+
     def test_mean_no_weights(self):
         """RDMs.mean() returns an RDMs with the nan omitted mean of the rdms
         """