make warmstart fuzziness more rigorous with kendall tau distances

docs/example-warmstart.ipynb

@@ -127,7 +127,7 @@
    "source": [
     "results = []\n",
     "\n",
-    "yseen = y[:i]\n",
+    "yseen = y[:]\n",
     "\n",
     "# delete any existing content:\n",
     "ReactiveNestedSampler(parameters, log_likelihood, prior_transform,\n",
@@ -139,7 +139,8 @@
     "    \n",
     "    yseen = y[:i]\n",
     "    sampler = ReactiveNestedSampler(parameters, log_likelihood, prior_transform,\n",
-    "                                              log_dir='warmstartdoc', resume='resume-similar')\n",
+    "                                    log_dir='warmstartdoc', resume='resume-similar',\n",
+    "                                    warmstart_max_tau=0.5)\n",
     "    ncall_initial = int(sampler.ncall)\n",
     "    res = sampler.run(frac_remain=0.5, viz_callback=None)\n",
     "    results.append((i, res, ncall_initial))\n",
@@ -188,7 +189,7 @@
    "source": [
     "## Conclusion\n",
     "\n",
-    "Notice the time saving in the bottom panel a cost reduction of ~30%. This benefit is *independent of problem dimension*. The cost savings are higher, the more similar the modified problem is.\n",
+    "Notice the time saving in the bottom panel by more thanhalf. This benefit is *independent of problem dimension*. The cost savings are higher, the more similar the modified problem is.\n",
     "\n",
     "This feature allows you to:\n",
     "\n",

evaluate/evaluate_sampling.py

@@ -3,15 +3,15 @@
 from ultranest.mlfriends import ScalingLayer, AffineLayer, MLFriends
 from ultranest.stepsampler import RegionMHSampler, CubeMHSampler
 from ultranest.stepsampler import CubeSliceSampler, RegionSliceSampler, RegionBallSliceSampler, RegionSequentialSliceSampler, SpeedVariableRegionSliceSampler
-#from ultranest.stepsampler import DESampler
+from ultranest.stepsampler import AHARMSampler
 #from ultranest.stepsampler import OtherSamplerProxy, SamplingPathSliceSampler, SamplingPathStepSampler
 #from ultranest.stepsampler import GeodesicSliceSampler, RegionGeodesicSliceSampler
 import tqdm
 import joblib
 import warnings
 from problems import transform, get_problem
 
-mem = joblib.Memory('.', verbose=False)
+#mem = joblib.Memory('.', verbose=False)
 
 def quantify_step(a, b):
     # euclidean step distance
@@ -28,7 +28,7 @@ def quantify_step(a, b):
     radial_step = np.abs(ra - rb)
     return [stepsize, angular_step, radial_step]
 
-@mem.cache
+#@mem.cache
 def evaluate_warmed_sampler(problemname, ndim, nlive, nsteps, sampler):
     loglike, grad, volume, warmup = get_problem(problemname, ndim=ndim)
     if hasattr(sampler, 'set_gradient'):
@@ -134,35 +134,22 @@ def main(args):
 
     samplers = [
         #CubeMHSampler(nsteps=16), #CubeMHSampler(nsteps=4), CubeMHSampler(nsteps=1),
-        RegionMHSampler(nsteps=16), #RegionMHSampler(nsteps=4), RegionMHSampler(nsteps=1),
+        #RegionMHSampler(nsteps=16), #RegionMHSampler(nsteps=4), RegionMHSampler(nsteps=1),
         ##DESampler(nsteps=16), DESampler(nsteps=4), #DESampler(nsteps=1),
         #CubeSliceSampler(nsteps=2*ndim), CubeSliceSampler(nsteps=ndim), CubeSliceSampler(nsteps=max(1, ndim//2)),
         #RegionSliceSampler(nsteps=ndim), RegionSliceSampler(nsteps=max(1, ndim//2)),
-        RegionSliceSampler(nsteps=2), RegionSliceSampler(nsteps=4), 
-        RegionSliceSampler(nsteps=ndim), RegionSliceSampler(nsteps=4*ndim),
+        #RegionSliceSampler(nsteps=2), RegionSliceSampler(nsteps=4), 
+        #RegionSliceSampler(nsteps=ndim), RegionSliceSampler(nsteps=4*ndim),
         #RegionBallSliceSampler(nsteps=2*ndim), RegionBallSliceSampler(nsteps=ndim), RegionBallSliceSampler(nsteps=max(1, ndim//2)),
        # RegionSequentialSliceSampler(nsteps=2*ndim), RegionSequentialSliceSampler(nsteps=ndim), RegionSequentialSliceSampler(nsteps=max(1, ndim//2)),
 
         #SpeedVariableRegionSliceSampler([Ellipsis]*ndim), SpeedVariableRegionSliceSampler([slice(i, ndim) for i in range(ndim)]),
         #SpeedVariableRegionSliceSampler([Ellipsis]*ndim + [slice(1 + ndim//2, None)]*ndim), 
 
-        RegionMHSampler(nsteps=2*ndim, adaptive_nsteps='move-distance'),
-        RegionMHSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances'),
-        RegionMHSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances-NN'),
-        RegionMHSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances'),
-        RegionMHSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances-NN'),
-
-        RegionSliceSampler(nsteps=2*ndim, adaptive_nsteps='move-distance'),
-        RegionSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances'),
-        RegionSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances-NN'),
-        RegionSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances'),
-        RegionSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances-NN'),
-
-        RegionBallSliceSampler(nsteps=2*ndim, adaptive_nsteps='move-distance'),
-        RegionBallSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances'),
-        RegionBallSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-total-distances-NN'),
-        RegionBallSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances'),
-        RegionBallSliceSampler(nsteps=2*ndim, adaptive_nsteps='proposal-summed-distances-NN'),
+        AHARMSampler(nsteps=64),
+        AHARMSampler(nsteps=64, adaptive_nsteps='move-distance'),
+        AHARMSampler(nsteps=64, region_filter=False),
+        AHARMSampler(nsteps=64, orthogonalise=True),
     ]
     if ndim < 14:
         samplers.insert(0, MLFriendsSampler())
@@ -311,4 +298,3 @@ def main(args):
 
     args = parser.parse_args()
     main(args)
-

tests/test_run.py

@@ -324,7 +324,8 @@ def transform(x):
             try:
                 sampler = ReactiveNestedSampler(paramnames,
                     loglike, transform=transform,
-                    log_dir=folder, resume=resume, vectorized=True, draw_multiple=False)
+                    log_dir=folder, resume=resume, vectorized=True, draw_multiple=False,
+                    warmstart_max_tau=0.5)
             except Exception as e:
                 # we expect an error for resuming with a changed likelihood
                 if resume != 'resume':

tests/test_stepsampling.py

@@ -1,7 +1,7 @@
 import numpy as np
 from ultranest.mlfriends import ScalingLayer, AffineLayer, MLFriends
 from ultranest import ReactiveNestedSampler
-from ultranest.stepsampler import RegionMHSampler, CubeMHSampler, CubeSliceSampler, RegionSliceSampler, SpeedVariableRegionSliceSampler, AHARMSampler
+from ultranest.stepsampler import RegionMHSampler, CubeMHSampler, CubeSliceSampler, RegionSliceSampler, SpeedVariableRegionSliceSampler, AHARMSampler, RegionBallSliceSampler
 from ultranest.stepsampler import generate_region_random_direction, ellipsoid_bracket, crop_bracket_at_unit_cube
 from ultranest.pathsampler import SamplingPathStepSampler
 from numpy.testing import assert_allclose

tests/test_utils.py

@@ -1,7 +1,7 @@
 import numpy as np
 import tempfile
 import os
-from ultranest.utils import vectorize, is_affine_transform
+from ultranest.utils import vectorize, is_affine_transform, normalised_kendall_tau_distance
 from numpy.testing import assert_allclose
 
 
@@ -29,3 +29,20 @@ def test_is_affine_transform():
 	assert is_affine_transform(a, a - 1)
 	assert is_affine_transform(a, a * 10000 - 5000.)
 	assert not is_affine_transform(a, a**2)
+
+def test_tau():
+
+	assert normalised_kendall_tau_distance(np.arange(400), np.arange(400)) == 0
+	assert normalised_kendall_tau_distance(np.arange(2000), np.arange(2000)) == 0
+	a = np.array([1, 2, 3, 4, 5])
+	b = np.array([3, 4, 1, 2, 5])
+	assert normalised_kendall_tau_distance(a, b) == 0.4
+	i, j = np.meshgrid(np.arange(len(a)), np.arange(len(b)))
+	assert normalised_kendall_tau_distance(a, b, i, j) == 0.4
+	assert normalised_kendall_tau_distance(a, a, i, j) == 0
+
+	try:
+		normalised_kendall_tau_distance(np.arange(5), np.arange(10))
+		raise Exception("expect error")
+	except AssertionError:
+		pass

ultranest/integrator.py

@@ -16,7 +16,7 @@
 from numpy import log, exp, logaddexp
 import numpy as np
 
-from .utils import create_logger, make_run_dir, resample_equal, vol_prefactor, vectorize, listify as _listify, is_affine_transform
+from .utils import create_logger, make_run_dir, resample_equal, vol_prefactor, vectorize, listify as _listify, is_affine_transform, normalised_kendall_tau_distance
 from ultranest.mlfriends import MLFriends, AffineLayer, ScalingLayer, find_nearby, WrappingEllipsoid
 from .store import HDF5PointStore, TextPointStore, NullPointStore
 from .viz import get_default_viz_callback, nicelogger
@@ -112,7 +112,9 @@ def _explore_iterator_batch(explorer, pop, x_dim, num_params, pointpile, batchsi
         yield batch
 
 
-def resume_from_similar_file(log_dir, x_dim, loglikelihood, transform, verbose=False, vectorized=False, ndraw=400):
+
+def resume_from_similar_file(log_dir, x_dim, loglikelihood, transform,
+    max_tau=0, verbose=False, ndraw=400):
     """
     Change a stored UltraNest run to a modified loglikelihood/transform.
 
@@ -128,6 +130,16 @@ def resume_from_similar_file(log_dir, x_dim, loglikelihood, transform, verbose=F
         new transform function
     verbose: bool
         show progress
+    ndraw: int
+        set to >1 if functions can take advantage of vectorized computations
+    max_tau: float
+        Allowed dissimilarity in the live point ordering, quantified as
+        normalised Kendall tau distance.
+        
+        max_tau=0 is the very conservative choice of stopping the warm start
+        when the live point order differs.
+        Near 1 are completely different live point orderings.
+        Values in between permit mild disorder.
 
     Returns
     ----------
@@ -201,13 +213,14 @@ def pop(Lmin):
         roots2.append(pointpile2.make_node(logl_new, u, v))
         pointstore2.add(_listify([-np.inf, logl_new, 0.0], u, v), 1)
 
-    batchsize = ndraw if vectorized else 1
+    batchsize = ndraw
     explorer = BreadthFirstIterator(roots)
     explorer2 = BreadthFirstIterator(roots2)
     main_iterator2 = SingleCounter()
     main_iterator2.Lmax = logls_new.max()
     good_state = True
 
+    indices1, indices2 = np.meshgrid(np.arange(len(logls_new)), np.arange(len(logls_new)))
     last_good_like = -1e300
     last_good_state = 0
     epsilon = 1 + 1e-6
@@ -243,31 +256,23 @@ def pop(Lmin):
             if len(active_values) != len(active_values2):
                 if verbose == 2:
                     print("stopping, number of live points differ (%d vs %d)" % (len(active_values), len(active_values2)))
+                    good_state = False
                 break
-            active_values_order = np.argsort(active_values)
-            active_values2_order = np.argsort(active_values2)
-            mask_order_different = active_values_order != active_values2_order
-            reorders = list(zip(active_values_order[active_values_order!=active_values2_order], active_values2_order[active_values_order!=active_values2_order]))
-            reorders_are_onlyswaps = all((b, a) in reorders for a, b in reorders)
-
-            order_consistent = (np.argmin(active_values) == np.argmin(active_values2))
-            order_fully_consistent = (np.argsort(active_values) == np.argsort(active_values2)).all()
-            if order_fully_consistent and len(active_values) > 10 or reorders_are_onlyswaps and len(active_values) > 10:
+
+
+            if len(active_values) != len(indices1):
+                indices1, indices2 = np.meshgrid(np.arange(len(active_values)), np.arange(len(active_values2)))
+            tau = normalised_kendall_tau_distance(active_values, active_values2, indices1, indices2)
+            order_consistent = tau <= max_tau
+            if order_consistent and len(active_values) > 10 and len(active_values) > 10:
                 good_state = True
             elif not order_consistent:
                 good_state = False
             else:
                 # maintain state
                 pass
             if verbose == 2:
-                print(
-                    niter, mask_order_different.sum(),
-                    'S' if reorders_are_onlyswaps else ' ',
-                    'C' if order_consistent else ' ',
-                    'F' if order_fully_consistent else ' ',
-                    'G' if good_state else ' ',
-                    '%5.2f' % (100 * (np.argsort(active_values) == np.argsort(active_values2)).mean())
-                )
+                print(niter, tau)
             if good_state:
                 #print("        (%.1e)   L=%.1f" % (last_good_like, Lmin2))
                 #assert last_good_like < Lmin2, (last_good_like, Lmin2)
@@ -286,14 +291,16 @@ def pop(Lmin):
                 logl_new = logls_new[j]
                 j += 1
 
-
                 #print(j, Lmin2, '->', logl_new, 'instead of', Lmin, '->', [c.value for c in node2.children])
+                child2 = pointpile2.make_node(logl_new, u, v)
+                node2.children.append(child2)
                 if logl_new > Lmin2:
+                    pointstore2.add(_listify([Lmin2, logl_new, 0.0], u, v), 1)
+                else:
                     if verbose == 2:
                         print("cannot use new point because it would decrease likelihood (%.1f->%.1f)" % (Lmin2, logl_new))
-                    child2 = pointpile2.make_node(logl_new, u, v)
-                    node2.children.append(child2)
-                    pointstore2.add(_listify([Lmin2, logl_new, 0.0], u, v), 1)
+                    #good_state = False
+                    #break
 
             main_iterator2.passing_node(node2, active_nodes2)
 
@@ -901,6 +908,7 @@ def __init__(self,
                  ndraw_min=128,
                  ndraw_max=65536,
                  storage_backend='hdf5',
+                 warmstart_max_tau=-1,
                  ):
         """Initialise nested sampler.
 
@@ -968,6 +976,12 @@ def __init__(self,
         storage_backend: str or class
             Class to use for storing the evaluated points (see ultranest.store)
             'hdf5' is strongly recommended. 'tsv' and 'csv' are also possible.
+
+        warmstart_max_tau: float
+            Maximum disorder to accept when resume='resume-similar';
+            Live points are reused as long as the live point order 
+            is below this normalised Kendall tau distance.
+            Values from 0 (highly conservative) to 1 (extremely negligent).
         """
         self.paramnames = param_names
         x_dim = len(self.paramnames)
@@ -1055,13 +1069,17 @@ def __init__(self,
             assert self.log_to_disk
             if resume_similar and self.log_to_disk:
                 assert storage_backend == 'hdf5', 'resume-similar is only supported for HDF5 files'
+                assert 0 <= warmstart_max_tau <= 1, 'warmstart_max_tau parameter needs to be set to a value between 0 and 1'
                 # close
                 self.pointstore.close()
                 del self.pointstore
                 # rewrite points file
                 if self.log:
                     self.logger.info('trying to salvage points from previous, different run ...')
-                resume_from_similar_file(log_dir, x_dim, loglike, transform, vectorized=vectorized, ndraw=ndraw_min)
+                resume_from_similar_file(
+                    log_dir, x_dim, loglike, transform,
+                    ndraw=ndraw_min if vectorized else 1,
+                    max_tau=warmstart_max_tau, verbose=False)
                 self.pointstore = HDF5PointStore(
                     os.path.join(self.logs['results'], 'points.hdf5'),
                     3 + self.x_dim + self.num_params, mode='a' if resume else 'w')