From 0bf33f06cf3fdb3e5c134e79b351450206e69576 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 12:25:00 +0100 Subject: [PATCH 1/3] Use jackknife errors for drift and skip fwd/rev check for wide pairs. --- src/somd2/runner/_adaptive.py | 61 ++++++++++++++++++++++++----------- tests/runner/test_repex.py | 14 ++++++++ 2 files changed, 56 insertions(+), 19 deletions(-) diff --git a/src/somd2/runner/_adaptive.py b/src/somd2/runner/_adaptive.py index d86ce95..9d39123 100644 --- a/src/somd2/runner/_adaptive.py +++ b/src/somd2/runner/_adaptive.py @@ -96,7 +96,33 @@ 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. + """ + 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 +138,30 @@ 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. Only checked + # when both spreads are at most 2 kT, since beyond that the bias of the + # exponential average alone exceeds 1 kT. + 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 = [ + (_np.array_split(fwd, 2)[half], _np.array_split(rev, 2)[half]) + for half in range(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 +510,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..66123b0 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -2522,6 +2522,20 @@ 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 whose spread is wide enough for the bias of + the exponential averages to exceed 1 kT isn't flagged as unconverged. + """ + 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 From 4964b01c16b0b61ca64b887d144562aad295a8b5 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 12:33:52 +0100 Subject: [PATCH 2/3] Fix jackknife on short series and tidy convergence check docs. --- README.md | 5 +++-- src/somd2/runner/_adaptive.py | 15 ++++++++------- tests/runner/test_repex.py | 5 +++-- 3 files changed, 14 insertions(+), 11 deletions(-) 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 9d39123..63c2646 100644 --- a/src/somd2/runner/_adaptive.py +++ b/src/somd2/runner/_adaptive.py @@ -109,6 +109,11 @@ 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( @@ -144,9 +149,8 @@ def is_unconverged(fwd, rev): bool Whether the pair is unconverged. """ - # The forward and reverse estimates disagree by more than 1 kT. Only checked - # when both spreads are at most 2 kT, since beyond that the bias of the - # exponential average alone exceeds 1 kT. + # 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 @@ -155,10 +159,7 @@ def is_unconverged(fwd, rev): # 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. - halves = [ - (_np.array_split(fwd, 2)[half], _np.array_split(rev, 2)[half]) - for half in range(2) - ] + 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 diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index 66123b0..d67dae8 100644 --- a/tests/runner/test_repex.py +++ b/tests/runner/test_repex.py @@ -2524,8 +2524,9 @@ def test_rank_windows_flags_unconverged_pairs(): def test_wide_converged_pair_is_not_flagged(): """ - Validate that a converged pair whose spread is wide enough for the bias of - the exponential averages to exceed 1 kT isn't flagged as unconverged. + 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 From 06bf026957288fd4eb33fa5403cab620c3085172 Mon Sep 17 00:00:00 2001 From: Lester Hedges Date: Thu, 22 Oct 2026 12:40:08 +0100 Subject: [PATCH 3/3] Make the ranking tests hard pair consistent so it is not flagged. --- tests/runner/test_repex.py | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/tests/runner/test_repex.py b/tests/runner/test_repex.py index d67dae8..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)