Fix forward/reverse convergence analysis: slice by time, not by replica

[aqemia-benedict-tan]

Summary

I recently benchmarked the default AbsoluteBindingProtocol for the TYK2 JACS dataset. The overall performance was strong ($R^2=0.50$, $RMSE=1.01$ kcal/mol), and the overlap matrices and replica exchange statistics looked healthy. However, the convergence behaviour appeared to be problematic, showing large fluctuations in both the forward and reverse estimates. The convergence plots generated for the solvent and complex leg for the ejm_47 ligand are shown below (input is ejm_47.json, raw results are given here: results.json).

Solvent leg

Complex leg

Interestingly, the real time analysis (complex_real_time_analysis.yaml, solvent_real_time_analysis.yaml) indicated satisfactory convergence for both the solvent and complex leg (within 1 kcal/mol at >0.5 fraction of uncorrelated samples):

I inspected how MultistateEquilFEAnalysis.get_forward_and_reverse_analysis generates the forward and reverse estimates, and it seems there is a bug in how the MBAR input matrix is sliced for the individual fractions < 1.0, that can result in large fluctuations in the partial-fraction $\Delta G$ estimates for calculations that are actually well converged.

The current code feeds slices with increasing number of columns/samples of the u_ln matrix obtained by self.analyzer._unbiased_decorrelated_u_ln to obtain DG/dDG estimates at increasing fraction of uncorrelated samples. This approach assumes that the columns are ordered per unit of time, so adding extra columns to your estimate corresponds to adding an increasing number of samples collected across all replicas.

self.analyzer is an openmmtools MultiStateSamplerAnalyzer; _unbiased_decorrelated_u_ln is one of its cached properties, ultimately built by _compute_mbar_decorrelated_energies. The 3D -> 2D reshape in openmmtools is handled by the reformat_energies_for_mbar method as shown below:

# see https://github.com/choderalab/openmmtools/blob/ba4031da49981fe9e8b0bcec010ab53e18a03ba2/openmmtools/multistate/multistateanalyzer.py#L994
k, l, n = u_kln.shape
if n_k is None:
    n_k = np.ones(k, dtype=np.int32)*n
u_ln = np.zeros([l, n_k.sum()])
n_counter = 0
for k_index in range(k):
    u_ln[:, n_counter:n_counter + n_k[k_index]] = u_kln[k_index, :, :n_k[k_index]]
    n_counter += n_k[k_index]

This flattening does not preserve the time-ordering of samples along the column axis; it groups them per replica (column = replica * M + iteration, where M is the number of decorrelated samples per replica). So at 50 % of the uncorrelated samples, u_ln[:, :u_ln.shape[1] // 2] selects all samples from half of the replicas and none from the other half. On its own that subset would still be usable — MBAR is permutation-invariant in its columns — but the function also passes MBAR a uniform per-state count (new_N_l = [chunk] * n_states) that the slice does not actually have: that equal split only holds once every replica is included (fraction 1.0). MBAR is therefore handed self-inconsistent inputs and returns a biased estimate for every fraction < 1.0. This also explains why the real-time MBAR estimate and the post-hoc convergence analysis agree only at the final estimate.

This PR reshapes the 2D u_ln matrix into its 3D (state, replica, iteration) form, slices the iteration (time) axis to take the first (forward) or last (reverse) chunk iterations of every replica, then flattens back to 2D for MBAR. Because each iteration contributes exactly one sample to every state, each slice now contains exactly chunk samples per state — consistent with the new_N_l = [chunk] * n_states passed to MBAR — and the "fraction" axis genuinely means fraction of simulation time. The reshape fix is self-contained within the _get_fraction_free_energy helper function, which was introduced by PR #1984. This means that the public get_forward_and_reverse_analysis is unchanged compared to the same PR.

I reran the forwards/reverse energy estimate for the TYK2 ejm_47 example with the fix implemented (generate_convergence.py), and obtained the convergence plots shown below (forwards estimate is now consistent with the real-time analysis yaml files).

Solvent leg

Complex leg

Tests

Both changes are in TestFEAnalysis (src/openfe/tests/protocols/test_openmmutils.py).

New — test_fraction_free_energy_slices_by_time. A fast, property-based regression test for the fix. It builds a small replica-major u_ln in which every column is tagged with its own flat index, mocks _get_free_energy to capture the matrices that _get_fraction_free_energy hands it, and asserts that the forward/reverse slices are the first/last chunk iterations of every replica (a slice in time), not the first/last chunk * n_states contiguous columns. It also guards explicitly against a regression to the old contiguous-column selection. Because it checks which samples are selected rather than the resulting MBAR numbers, it fails on the pre-fix code and passes on the fix, and needs no .nc fixture or MBAR run.

Modified — test_free_energies. This characterisation test pins the per-fraction forward_DGs / reverse_DGs values. Since the fix changes which samples enter each intermediate fraction, those values change for every fraction < 1.0. The final fraction (1.0) is unchanged — it uses the full data set either way — confirming the reported free energy is unaffected and only the intermediate path to it is redistributed. The forward_dDGs / reverse_dDGs arrays are left as-is (they still pass under the existing rtol=5e-01). The updated expected arrays were regenerated from the same fixture on this branch, and can be reproduced with:

from importlib import resources
from openff.units import unit
from openmmtools import multistate
from openfe.protocols.openmm_utils import multistate_analysis

with resources.as_file(resources.files("openfe.tests.data.openmm_rfe")) as d:
    nc, chk = str(d / "vacuum_nocoord.nc"), str(d / "vacuum_nocoord_checkpoint.nc")
reporter = multistate.MultiStateReporter(storage=nc, checkpoint_storage=chk)
analyzer = multistate.MultiStateSamplerAnalyzer(reporter)

# call the fixed method directly (skips the slow full-analysis __init__)
wrap = multistate_analysis.MultistateEquilFEAnalysis.__new__(
    multistate_analysis.MultistateEquilFEAnalysis)
wrap.analyzer, wrap.units = analyzer, unit.kilocalorie_per_mole
fr = wrap.get_forward_and_reverse_analysis(num_samples=10)
print(repr(fr["forward_DGs"].m), repr(fr["reverse_DGs"].m))

The existing MBAR-failure tests (test_forward_and_reverse_nan_on_mbar_failure, test_forward_and_reverse_none_on_final_fraction_failure) are unaffected by the fix and continue to pass.

LLM / AI generated code disclosure

LLMs or other AI-powered tools (beyond simple IDE use cases) were used in this contribution: yes

The new unit test test_fraction_free_energy_slices_by_time was generated by Claude Code (Opus 4.8) + manually reviewed and edited. The code shown above in the Tests section was also generated using Claude. Finally, the python script for local testing generate_convergence.py was also generated using Claude. All AI-generated code was manually reviewed + edited where necessary.

Checklist

• All new code is appropriately documented (user-facing code must have complete docstrings) - **still need to update docs wrt to return partial convergence analysis if mbar fails for individual fraction #1984 **
• Added a news entry, or the changes are not user-facing.
• Ran pre-commit: you can run pre-commit locally or comment on this PR with pre-commit.ci autofix.
• Filled in the AI generated code disclosure.

Developers certificate of origin

• I certify that this contribution is covered by the MIT License here and the Developer Certificate of Origin at https://developercertificate.org/.