diff --git a/README.md b/README.md index b8af331..a7ab166 100644 --- a/README.md +++ b/README.md @@ -447,8 +447,9 @@ accounting for the samples each window has already collected and for cold windows costing and sampling `--adaptive-mix-interval` times less than hot ones. The following are always hot: -- pairs whose forward and reverse free energy estimates disagree, or whose - estimate drifts during the simulation, +- pairs whose estimate drifts during the simulation or, for pairs with narrow + energy distributions, whose forward and reverse free energy estimates + disagree, - pairs that rarely exchange, - windows with fewer than `--adaptive-min-samples` statistically independent samples since the burn-in, diff --git a/src/somd2/runner/_adaptive.py b/src/somd2/runner/_adaptive.py index d86ce95..63c2646 100644 --- a/src/somd2/runner/_adaptive.py +++ b/src/somd2/runner/_adaptive.py @@ -96,7 +96,38 @@ def exp_estimate(du): return -(m + _np.log(_np.mean(_np.exp(-du - m)))) -def is_unconverged(fwd, rev, g_fwd, g_rev): +def _pair_estimate(fwd, rev): + """ + The reduced free energy difference between a pair of neighbouring windows, + averaged over the forward and reverse exponential estimates. + """ + return 0.5 * (exp_estimate(fwd) - exp_estimate(rev)) + + +def _jackknife_error(fwd, rev, num_batches=10): + """ + The standard error of the pair estimate, from a jackknife over contiguous + batches, which accounts for correlation and the nonlinear average. + """ + # Every batch needs a sample, and the error needs at least two batches. + num_batches = min(num_batches, len(fwd), len(rev)) + if num_batches < 2: + return _np.inf + + f = _np.array_split(fwd, num_batches) + r = _np.array_split(rev, num_batches) + estimates = _np.array( + [ + _pair_estimate( + _np.concatenate(f[:i] + f[i + 1 :]), _np.concatenate(r[:i] + r[i + 1 :]) + ) + for i in range(num_batches) + ] + ) + return _np.sqrt((num_batches - 1) * estimates.var()) + + +def is_unconverged(fwd, rev): """ Whether the free energy difference between a pair of neighbouring windows shows signs of poor sampling, regardless of its uncertainty. @@ -112,33 +143,26 @@ def is_unconverged(fwd, rev, g_fwd, g_rev): The reduced energy difference to the lower window, sampled in the upper window. - g_fwd, g_rev: float - The statistical inefficiencies of fwd and rev. - Returns ------- bool Whether the pair is unconverged. """ - # The forward and reverse estimates disagree by more than 1 kT. - if abs(exp_estimate(fwd) + exp_estimate(rev)) > 1.0: + # The forward and reverse estimates disagree by more than 1 kT, checked only + # for spreads up to 2 kT, beyond which the bias of the average can exceed that. + if ( + max(fwd.std(), rev.std()) <= 2.0 + and abs(exp_estimate(fwd) + exp_estimate(rev)) > 1.0 + ): return True # The estimates from the two halves of the data differ by more than 0.25 kT, # and by more than twice their combined standard error. - estimates = [] - errors = [] - for half in range(2): - f = _np.array_split(fwd, 2)[half] - r = _np.array_split(rev, 2)[half] - estimates.append(0.5 * (exp_estimate(f) - exp_estimate(r))) - errors.append( - 0.5 * _np.sqrt(f.var() * g_fwd / len(f) + r.var() * g_rev / len(r)) - ) - - drift = abs(estimates[0] - estimates[1]) - return drift > 0.25 and drift > 2.0 * _np.hypot(*errors) + halves = list(zip(_np.array_split(fwd, 2), _np.array_split(rev, 2))) + drift = abs(_pair_estimate(*halves[1]) - _pair_estimate(*halves[0])) + error = _np.hypot(*(_jackknife_error(f, r) for f, r in halves)) + return drift > 0.25 and drift > 2.0 * error def cold_rate(first_cycle, num_cycles, checkpoint_cycles, mix_interval): @@ -487,7 +511,7 @@ def rank( num_proposed = proposed[k, k + 1] if ( num_proposed > 0 and accepted[k, k + 1] / num_proposed < 0.05 - ) or is_unconverged(fwd, rev, g_upper[k], g_lower[k + 1]): + ) or is_unconverged(fwd, rev): flagged.append(k) weights = _np.sqrt(fwd_variance + rev_variance) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index d7c0a2a..4ee74a7 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -1697,10 +1697,12 @@ def _fill_ranking_data(runner, hard_pair, num_samples=200, seed=7): rng = np.random.default_rng(seed) num_lambda = runner._config.num_lambda + # Means of +/- dG + sigma^2 / 2 make the forward and reverse estimates agree. for k in range(num_lambda - 1): sigma = 1.5 if k == hard_pair else 0.1 - runner._ranker.upper[k] = list(rng.normal(0.5, sigma, num_samples)) - runner._ranker.lower[k + 1] = list(rng.normal(-0.5, sigma, num_samples)) + shift = 0.5 * sigma**2 + runner._ranker.upper[k] = list(rng.normal(0.5 + shift, sigma, num_samples)) + runner._ranker.lower[k + 1] = list(rng.normal(-0.5 + shift, sigma, num_samples)) def test_rank_windows(): @@ -1825,7 +1827,7 @@ def test_rank_windows_fill_windows_stay_cold_candidates(): """ runner = _ranking_runner(8, adaptive_min_samples=10) - # The flagged pair 0 gives windows 0, 1 and 2, so one window is filled. + # The hard pair 0 gives windows 0, 1 and 2, so one window is filled. _fill_ranking_data(runner, hard_pair=0) _rank_twice(runner, num_gpus=2) @@ -2162,7 +2164,7 @@ def test_rank_windows_fills_whole_batches(oversubscription_factor, capacity, num 8, adaptive_min_samples=10, oversubscription_factor=oversubscription_factor ) - # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. + # The hard pair 0 gives windows 0, 1 and 2, one on each of three GPUs. _fill_ranking_data(runner, hard_pair=0) _rank_twice(runner, num_gpus=4, capacity=capacity) @@ -2180,7 +2182,7 @@ def test_rank_windows_unfilled_with_shared_contexts(): """ runner = _ranking_runner(8, is_cached=True, adaptive_min_samples=10) - # The flagged pair 0 gives windows 0, 1 and 2, one on each of three GPUs. + # The hard pair 0 gives windows 0, 1 and 2, one on each of three GPUs. _fill_ranking_data(runner, hard_pair=0) _rank_twice(runner, num_gpus=4) @@ -2522,6 +2524,21 @@ def test_rank_windows_flags_unconverged_pairs(): assert hot is not None and {0, 1, 2}.issubset(hot) +def test_wide_converged_pair_is_not_flagged(): + """ + Validate that a converged pair with a wide spread isn't flagged as + unconverged, either by the bias of its forward and reverse exponential + averages or by drift between the two halves of its data. + """ + from somd2.runner._adaptive import is_unconverged + + rng = np.random.default_rng(0) + fwd = rng.normal(8.0, 4.0, 2000) + rev = rng.normal(8.0, 4.0, 2000) + + assert not is_unconverged(fwd, rev) + + def test_checkpoint_without_adaptive_state(): """ Validate that a checkpoint written before adaptive replica exchange was