diff --git a/docs/example-warmstart.ipynb b/docs/example-warmstart.ipynb index 666b3cb5..5117f14b 100644 --- a/docs/example-warmstart.ipynb +++ b/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", diff --git a/evaluate/evaluate_sampling.py b/evaluate/evaluate_sampling.py index 951f6e6b..5e8562db 100644 --- a/evaluate/evaluate_sampling.py +++ b/evaluate/evaluate_sampling.py @@ -3,7 +3,7 @@ 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 @@ -11,7 +11,7 @@ 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) - diff --git a/tests/test_run.py b/tests/test_run.py index e0d20298..8d91ddc3 100644 --- a/tests/test_run.py +++ b/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': diff --git a/tests/test_stepsampling.py b/tests/test_stepsampling.py index 4666c0bb..18e9fb69 100644 --- a/tests/test_stepsampling.py +++ b/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 diff --git a/tests/test_utils.py b/tests/test_utils.py index fecfde99..281edf15 100644 --- a/tests/test_utils.py +++ b/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 diff --git a/ultranest/integrator.py b/ultranest/integrator.py index 4f92d42c..30a7a3dd 100644 --- a/ultranest/integrator.py +++ b/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,16 +256,15 @@ 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 @@ -260,14 +272,7 @@ def pop(Lmin): # 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') diff --git a/ultranest/stepsampler.py b/ultranest/stepsampler.py index e298485e..baaed968 100644 --- a/ultranest/stepsampler.py +++ b/ultranest/stepsampler.py @@ -878,6 +878,198 @@ def crop_bracket_at_unit_cube(ui, v, left, right, epsilon=1e-6): assert left <= 0 <= right, (left, right) return left, right, cropped_left, cropped_right +def _prepare_steps( + nsteps_done, nsteps, directions, ndraw, + current_interval, loglike, transform, region, ndim, region_filter, + Lmin, verbose, +): + point_sequence = [] + point_expectation = [] + intervals = [] + nsteps_prepared = 0 + while nsteps_prepared + nsteps_done < nsteps and len(point_sequence) < ndraw: + if verbose: + print("loop:", nsteps_prepared, nsteps_done, 'of', nsteps) + v = directions[nsteps_done + nsteps_prepared] + if verbose: + print("direction:", v) + if len(point_sequence) == 0: + ucurrent, left, right = current_interval + assert (ucurrent >= 0).all(), ucurrent + assert (ucurrent <= 1).all(), ucurrent + assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ( + 'cannot start from outside ellipsoid!', region.inside_ellipsoid(ucurrent.reshape((1, ndim)))) + if region_filter: + assert region.inside(ucurrent.reshape((1, ndim))), ( + 'cannot start from outside region!', region.inside(ucurrent.reshape((1, ndim)))) + assert loglike(transform(ucurrent.reshape((1, ndim)))) >= Lmin, ( + 'cannot start from outside!', loglike(transform(ucurrent.reshape((1, ndim)))), Lmin) + else: + left, right = None, None + assert (ucurrent >= 0).all(), ucurrent + assert (ucurrent <= 1).all(), ucurrent + if verbose: + print("preparing step: %d from %s" % (nsteps_prepared + nsteps_done, ucurrent)) + + if left is None or right is None: + # in each, find the end points using the expanded ellipsoid + assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') + left, right = ellipsoid_bracket(ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.enlarge) + left, right, _, _ = crop_bracket_at_unit_cube(ucurrent, v, left, right) + assert (ucurrent + v * left <= 1).all(), ( + ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) + assert (ucurrent + v * right <= 1).all(), ( + ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) + assert (ucurrent + v * left >= 0).all(), ( + ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) + assert (ucurrent + v * right >= 0).all(), ( + ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) + + assert left <= 0 <= right, (left, right) + if verbose: + print(" ellipsoid bracket found:", left, right) + + while True: + # sample in each a point until presumed success: + assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') + t = np.random.uniform(left, right) + unext = ucurrent + v * t + assert (unext >= 0).all(), unext + assert (unext <= 1).all(), unext + assert region.inside_ellipsoid(unext.reshape((1, ndim))), ('proposal landed outside ellipsoid!', t, left, right) + + # compute distance vector to center + d = unext - region.ellipsoid_center + # distance in normalised coordates: vector . matrix . vector + # where the matrix is the ellipsoid inverse covariance + r = np.einsum('j,jk,k->', d, region.ellipsoid_invcov, d) + if verbose: + print(" proposed slice point", t, r) + + likely_inside = r <= 1 + if not likely_inside and r <= region.enlarge: + # The exception is, when a point is between projected ellipsoid center and current point + # then it is also likely inside (if still inside the ellipsoid) + + # project ellipsoid center onto line + # region.ellipsoid_center = ucurrent + tc * v + tc = np.dot(region.ellipsoid_center - ucurrent, v) + # current point is at 0 by definition + if 0 < t < tc or tc < t < 0: + if verbose: + print(" proposed point is further inside than current point") + likely_inside = True + # print(" proposed point %.3f is going towards center %.3f" % (t, tc)) + # else: + # print(" proposed point %.3f is going away from center %.3f" % (t, tc)) + else: + # another exception is that points very close to the current point + # are very likely also inside + # to find that out, project all live points on the line + tall = np.einsum('ij,j->i', region.u - ucurrent, v) + # find the range and identify a small part of it + epsilon_nearby = 1e-6 + if tc < (tall.max() - tall.min()) * epsilon_nearby: + likely_inside = True + if verbose: + print(" proposed point is very nearby") + + if verbose: + print(" proposed point %s (%f) is likely %s (r=%f)" % (unext, t, 'inside' if likely_inside else 'outside', r)) + intervals.append((nsteps_prepared, ucurrent, v, left, right, t)) + point_sequence.append(unext) + point_expectation.append(likely_inside) + # If point radius in ellipsoid is <1, presume that it will be successful + if likely_inside: + nsteps_prepared += 1 + ucurrent = unext + assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') + break + + # Else, presume it will be unsuccessful, and sample another point + # shrink interval + if t > 0: + right = t + else: + left = t + + assert len(point_sequence) == len(point_expectation) + assert len(point_sequence) == len(intervals) + assert nsteps_prepared <= len(point_sequence) + + assert len(point_sequence) > 0, (len(point_sequence), ndraw, nsteps_prepared, nsteps_done, nsteps) + + if verbose: + print("proposed sequence:", point_sequence) + print("expectations:", point_expectation) + + return np.array(point_sequence, dtype=float), np.array(point_expectation, dtype=bool), intervals, nsteps_prepared + + +def _evaluate_with_filter( + region_filter, loglike, transform, Lmin, region, tregion, + point_sequence, point_expectation, + verbose +): + truncated = False + # region-filter, transform, tregion-filter, and evaluate the likelihood + if region_filter: + mask_inside = region.inside(point_sequence) + # identify first point that was expected to be inside, but was marked outside-of-region + i = np.where(np.logical_and(point_expectation, ~mask_inside))[0] + if verbose: + print("region filter says:", mask_inside, i) + if len(i) > 0: + imax = i[0] + 1 + # truncate there + point_sequence = point_sequence[:imax] + point_expectation = point_expectation[:imax] + mask_inside = mask_inside[:imax] + truncated |= True + del imax + if not mask_inside.any(): + return None + else: + mask_inside = None + + t_point_sequence = transform(point_sequence) + if region_filter and tregion is not None: + tmask = tregion.inside(t_point_sequence) + # identify first point that was expected to be inside, but was marked outside-of-region + i = np.where(np.logical_and(point_expectation, ~tmask))[0] + if verbose: + print("tregion filter says:", tmask, i) + mask_inside[~tmask] = False + del tmask + if len(i) > 0: + imax = i[0] + 1 + # truncate there + point_sequence = point_sequence[:imax] + point_expectation = point_expectation[:imax] + t_point_sequence = t_point_sequence[:imax] + mask_inside = mask_inside[:imax] + truncated |= True + del imax + if not mask_inside.any(): + return None + + # we expect the last point to be an accept, otherwise we would not terminate the sequence + assert point_expectation[-1] + if region_filter: + # set filtered ones to -np.inf + L = np.ones(len(t_point_sequence)) * -np.inf + nc = mask_inside.sum() + L[mask_inside] = loglike(t_point_sequence[mask_inside,:]) + else: + nc = len(point_sequence) + L = loglike(t_point_sequence) + Lmask = L > Lmin + + i = np.where(point_expectation != Lmask)[0] + if verbose: + print("reality:", Lmask) + print("difference:", point_expectation == Lmask) + return point_sequence, t_point_sequence, L, Lmask, i, nc, truncated class AHARMSampler(StepSampler): """Accelerated hit-and-run/slice sampler, vectorised. @@ -1024,6 +1216,7 @@ def __next__(self, region, Lmin, us, Ls, transform, loglike, ndraw=1024, plot=Fa Li = Ls[i] self.history.append((ui.copy(), Li.copy())) del i + print("starting at", ui) # set initially nleft = nsteps self.nsteps_done = 0 @@ -1052,178 +1245,25 @@ def __next__(self, region, Lmin, us, Ls, transform, loglike, ndraw=1024, plot=Fa nc = 0 while True: # prepare a sequence of points until nsteps are reached - point_sequence = [] - point_expectation = [] - intervals = [] - nsteps_prepared = 0 - while nsteps_prepared + self.nsteps_done < self.nsteps and len(point_sequence) < ndraw: - v = self.directions[self.nsteps_done + nsteps_prepared] - if len(point_sequence) == 0: - ucurrent, left, right = self.current_interval - assert (ucurrent >= 0).all(), ucurrent - assert (ucurrent <= 1).all(), ucurrent - assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ( - 'cannot start from outside ellipsoid!', region.inside_ellipsoid(ucurrent.reshape((1, ndim)))) - if self.region_filter: - assert region.inside(ucurrent.reshape((1, ndim))), ( - 'cannot start from outside region!', region.inside(ucurrent.reshape((1, ndim)))) - assert loglike(transform(ucurrent.reshape((1, ndim)))) >= Lmin, ( - 'cannot start from outside!', loglike(transform(ucurrent.reshape((1, ndim)))), Lmin) - else: - left, right = None, None - assert (ucurrent >= 0).all(), ucurrent - assert (ucurrent <= 1).all(), ucurrent - if verbose: - print("preparing step: %d from %s" % (nsteps_prepared + self.nsteps_done, ucurrent)) - - if left is None or right is None: - # in each, find the end points using the expanded ellipsoid - assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') - left, right = ellipsoid_bracket(ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.enlarge) - left, right, _, _ = crop_bracket_at_unit_cube(ucurrent, v, left, right) - assert (ucurrent + v * left <= 1).all(), ( - ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) - assert (ucurrent + v * right <= 1).all(), ( - ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) - assert (ucurrent + v * left >= 0).all(), ( - ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) - assert (ucurrent + v * right >= 0).all(), ( - ucurrent, v, region.ellipsoid_center, region.ellipsoid_inv_axes, region.ellipsoid_invcov, region.enlarge) - - assert left <= 0 <= right, (left, right) - if verbose: - print(" ellipsoid bracket found:", left, right) - - while True: - # sample in each a point until presumed success: - assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') - t = np.random.uniform(left, right) - unext = ucurrent + v * t - assert (unext >= 0).all(), unext - assert (unext <= 1).all(), unext - assert region.inside_ellipsoid(unext.reshape((1, ndim))), ('proposal landed outside ellipsoid!', t, left, right) - - # compute distance vector to center - d = unext - region.ellipsoid_center - # distance in normalised coordates: vector . matrix . vector - # where the matrix is the ellipsoid inverse covariance - r = np.einsum('j,jk,k->', d, region.ellipsoid_invcov, d) - - likely_inside = r <= 1 - if not likely_inside and r <= region.enlarge: - # The exception is, when a point is between projected ellipsoid center and current point - # then it is also likely inside (if still inside the ellipsoid) - - # project ellipsoid center onto line - # region.ellipsoid_center = ucurrent + tc * v - tc = np.dot(region.ellipsoid_center - ucurrent, v) - # current point is at 0 by definition - if 0 < t < tc or tc < t < 0: - if verbose: - print(" proposed point is further inside than current point") - likely_inside = True - # print(" proposed point %.3f is going towards center %.3f" % (t, tc)) - # else: - # print(" proposed point %.3f is going away from center %.3f" % (t, tc)) - else: - # another exception is that points very close to the current point - # are very likely also inside - # to find that out, project all live points on the line - tall = np.einsum('ij,j->i', region.u - ucurrent, v) - # find the range and identify a small part of it - epsilon_nearby = 1e-6 - if tc < (tall.max() - tall.min()) * epsilon_nearby: - likely_inside = True - if verbose: - print(" proposed point is very nearby") + point_sequence, point_expectation, intervals, nsteps_prepared = _prepare_steps( + self.nsteps_done, self.nsteps, self.directions, ndraw, + self.current_interval, loglike, transform, region, ndim, self.region_filter, + Lmin, verbose + ) + point_sequence, t_point_sequence, L, Lmask, indices_deviating, nc_here, truncated = _evaluate_with_filter( + self.region_filter, loglike, transform, Lmin, region, tregion, + point_sequence, point_expectation, + verbose + ) + del point_expectation + nc += nc_here - if verbose: - print(" proposed point %s (%f) is likely %s (r=%f)" % (unext, t, 'inside' if likely_inside else 'outside', r)) - intervals.append((nsteps_prepared, ucurrent, v, left, right, t)) - point_sequence.append(unext) - point_expectation.append(likely_inside) - # If point radius in ellipsoid is <1, presume that it will be successful - if likely_inside: - nsteps_prepared += 1 - ucurrent = unext - assert region.inside_ellipsoid(ucurrent.reshape((1, ndim))), ('current point outside ellipsoid!') - break - - # Else, presume it will be unsuccessful, and sample another point - # shrink interval - if t > 0: - right = t - else: - left = t - - # the above defines a sequence of points (u) - - assert len(point_sequence) > 0, (len(point_sequence), ndraw, nsteps_prepared, self.nsteps_done, self.nsteps) - point_sequence = np.array(point_sequence, dtype=float) - point_expectation = np.array(point_expectation, dtype=bool) - if verbose: - print("proposed sequence:", point_sequence) - print("expectations:", point_expectation) - truncated = False - # region-filter, transform, tregion-filter, and evaluate the likelihood - if self.region_filter: - mask_inside = region.inside(point_sequence) - # identify first point that was expected to be inside, but was marked outside-of-region - i = np.where(np.logical_and(point_expectation, ~mask_inside))[0] - if verbose: - print("region filter says:", mask_inside, i) - if len(i) > 0: - imax = i[0] + 1 - # truncate there - point_sequence = point_sequence[:imax] - point_expectation = point_expectation[:imax] - mask_inside = mask_inside[:imax] - truncated |= True - del imax - if not mask_inside.any(): - continue - - t_point_sequence = transform(point_sequence) - if self.region_filter and tregion is not None: - tmask = tregion.inside(t_point_sequence) - # identify first point that was expected to be inside, but was marked outside-of-region - i = np.where(np.logical_and(point_expectation, ~tmask))[0] - if verbose: - print("tregion filter says:", tmask, i) - mask_inside[~tmask] = False - del tmask - if len(i) > 0: - imax = i[0] + 1 - # truncate there - point_sequence = point_sequence[:imax] - point_expectation = point_expectation[:imax] - t_point_sequence = t_point_sequence[:imax] - mask_inside = mask_inside[:imax] - truncated |= True - del imax - if not mask_inside.any(): - continue - - # we expect the last point to be an accept, otherwise we would not terminate the sequence - assert point_expectation[-1] - if self.region_filter: - # set filtered ones to -np.inf - L = np.ones(len(t_point_sequence)) * -np.inf - nc += mask_inside.sum() - L[mask_inside] = loglike(t_point_sequence[mask_inside,:]) - else: - nc += len(point_sequence) - L = loglike(t_point_sequence) - Lmask = L > Lmin - i = np.where(point_expectation != Lmask)[0] self.nrejects += (~Lmask).sum() - if verbose: - print("reality:", Lmask) - print("difference:", point_expectation == Lmask) - print("calling likelihood with %5d prepared points, accepted:" % ( - len(point_sequence)), '=' * (i[0] + Lmask[i[0]] * 1 if len(i) > 0 else len(Lmask))) + #print("calling likelihood with %5d prepared points, accepted:" % ( + # len(point_sequence)), '=' * (i[0] + Lmask[i[0]] * 1 if len(i) > 0 else len(Lmask))) # identify first point that was unexpected - if len(i) == 0 and nsteps_prepared + self.nsteps_done == self.nsteps: + any_deviating = len(indices_deviating) > 0 + if any_deviating and nsteps_prepared + self.nsteps_done == self.nsteps: # everything according to prediction. if verbose: print("everything according to prediction and done") @@ -1232,7 +1272,7 @@ def __next__(self, region, Lmin, us, Ls, transform, loglike, ndraw=1024, plot=Fa self.history.append((ui, Li)) self.finalize_chain(region=region, Lmin=Lmin, Ls=Ls) return point_sequence[-1], t_point_sequence[-1], L[-1], nc - elif len(i) == 0: + elif any_deviating: # everything according to prediction. if verbose: print("everything according to prediction") @@ -1249,7 +1289,7 @@ def __next__(self, region, Lmin, us, Ls, transform, loglike, ndraw=1024, plot=Fa assert region.inside(ucurrent.reshape((1, ndim))), ('suggested point outside region!', region.inside(ucurrent.reshape((1, ndim)))) else: # point i unexpectedly inside or outside - imax = i[0] + imax = indices_deviating[0] for ui, Li in zip(point_sequence[:imax][Lmask[:imax]], L[:imax][Lmask[:imax]]): self.history.append((ui, Li)) nsteps_prepared, ucurrent, v, left, right, t = intervals[imax] diff --git a/ultranest/utils.py b/ultranest/utils.py index e59ebff7..00e46ad2 100644 --- a/ultranest/utils.py +++ b/ultranest/utils.py @@ -276,6 +276,27 @@ def is_affine_transform(a, b): return True +def normalised_kendall_tau_distance(values1, values2, i=None, j=None): + """ + Normalised Kendall tau distance between two arrays, + *values1* and *values2*, both of length N. + + see https://en.wikipedia.org/wiki/Kendall_tau_distance + + You can optionally pass precomputed indices: + i, j = np.meshgrid(np.arange(N), np.arange(N)) + """ + N = len(values1) + assert len(values2) == N, "Both lists have to be of equal length" + if i is None or j is None: + i, j = np.meshgrid(np.arange(N), np.arange(N)) + a = np.argsort(values1) + b = np.argsort(values2) + ndisordered = np.logical_or(np.logical_and(a[i] < a[j], b[i] > b[j]), np.logical_and(a[i] > a[j], b[i] < b[j])).sum() + return ndisordered / (N * (N - 1)) + + + def _merge_transform_loglike_gradient_function(transform, loglike, gradient): def transform_loglike_gradient(u): """combine transform, likelihood and gradient function"""