Fix adaptive drift error

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,

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)

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