[hannahbaumann]

Description

Fixes

• NVT context created during minimization leads to GPU memory issues in the ABFE protocol. #817

The new test was written in a back and forth with Claude opus4.8

Todos

• Implement feature / fix bug
• Add tests
• Update documentation as needed
• Update changelog to summarize changes in behavior, enhancements, and bugfixes implemented in this PR

Status

• Ready to go

Changelog message

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 85.42%. Comparing base (f6ef22a) to head (d128493).

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[ijpulidos]

Thank you so much for this! I can validate the results with this. Without these changes the behavior is basically like what's shown in the following screen capture

whereas with these changes is more like

So we can definitely see the benefit from these changes. However, I think the tests need to be improved, we probably don't want to change the current test as much as this PR is suggesting, we should keep the test as untouched as possible, and maybe writing a new one based on something like the following script

import pathlib

import openmm
from openmm import unit
from openmmtools import alchemy, cache, multistate, states, testsystems

REPLICAS = 31  # an ABFE complex leg: 30 lambda windows, plus its propagation context
STORAGE = pathlib.Path(__file__).parent / "test_multistatesampler_context_cache.nc"
STORAGE.unlink(missing_ok=True)

# A solvated, alchemical, barostatted system, as an ABFE complex leg simulates
water = testsystems.WaterBox(box_edge=7.0 * unit.nanometer, cutoff=1.0 * unit.nanometer)
system = alchemy.AbsoluteAlchemicalFactory().create_alchemical_system(
    water.system, alchemy.AlchemicalRegion(alchemical_atoms=[0, 1, 2])
)
system.addForce(openmm.MonteCarloBarostat(1.0 * unit.bar, 298.15 * unit.kelvin))

# The thermodynamic states of one leg: one system, differing only in lambda
thermodynamic_states = states.create_thermodynamic_state_protocol(
    system,
    protocol={
        "lambda_electrostatics": [index / (REPLICAS - 1) for index in range(REPLICAS)],
        "lambda_sterics": [1.0] * REPLICAS,
    },
    constants={"temperature": 298.15 * unit.kelvin, "pressure": 1.0 * unit.bar},
    composable_states=[alchemy.AlchemicalState.from_system(system)],
)
print(
    "propagation shares",
    len({s._standard_system_hash for s in thermodynamic_states}),
    "context(s)",
)

# One sampler state per replica, all starting from the same box of water.
sampler_states = [
    states.SamplerState(
        positions=water.positions, box_vectors=system.getDefaultPeriodicBoxVectors()
    )
    for _ in range(REPLICAS)
]

sampler = multistate.MultiStateSampler(number_of_iterations=1)
# Same unbounded, CUDA-backed cache openfe configures.
sampler.energy_context_cache = cache.ContextCache(
    capacity=None, time_to_live=None, platform=openmm.Platform.getPlatformByName("CUDA")
)

sampler.create(thermodynamic_states, sampler_states, storage=str(STORAGE))
print("after create:", len(sampler.energy_context_cache), "contexts")

# Keep it very short: we only care about context sharing, not a real minimization.
sampler.minimize(max_iterations=1)
print("after minimize:", len(sampler.energy_context_cache), "contexts")

print("holding", len(sampler.energy_context_cache), "contexts")

which is inspired in the script in OpenFreeEnergy/openfe#2222 but using the MultiStateSampler minimize method directly instead. I'll be working on the test changes for this as soon as possible. Thanks for the contribution.