Fix pertfile intrascale

README.md

@@ -388,13 +388,6 @@ If you want to load an existing system from a perturbation file and use the
 new `somd2` [ghost atom bonded-term modifications](https://github.com/OpenBioSim/ghostly),
 then simply omit the `--somd1-compatibility` option.
 
-> [!NOTE]
-> Using a ``pertfile`` as input is only supported when end states have the same
-> connectivity, i.e. it can't be used for ring-breaking perturbations, or
-> when non-bonded scale factors differ between the end states. For these cases,
-> you will need to generate a new perturbation stream file using `prepareFEP.py`
-> with the `--somd2` option as described above.
-
 ## GPU oversubscription
 
 If you have an NVIDIA GPU that supports the multi-process service (MPS), you can

src/somd2/_utils/_somd1.py

@@ -19,7 +19,12 @@
 # along with SOMD2. If not, see <http://www.gnu.org/licenses/>.
 #####################################################################
 
-__all__ = ["apply_pert", "make_compatible", "reconstruct_system"]
+__all__ = [
+    "apply_pert",
+    "make_compatible",
+    "reconstruct_intrascale",
+    "reconstruct_system",
+]
 
 from sire.system import System as _System
 from sire.legacy.System import System as _LegacySystem
@@ -611,6 +616,90 @@ def make_compatible(system, fix_perturbable_zero_sigmas=False):
     return system
 
 
+def reconstruct_intrascale(system):
+    """
+    Reconstruct end-state connectivity and intrascale matrices for perturbable
+    molecules from their bonded terms. This is required when a perturbation
+    file is used with AMBER topology/coordinate input, since the pertfile
+    assumes unchanged connectivity, which does not hold for ring-breaking or
+    ring-size-changing perturbations.
+
+    Parameters
+    ----------
+
+    system : sire.system.System, sire.legacy.System.System
+        The system containing the perturbable molecules.
+
+    Returns
+    -------
+
+    system : sire.system.System
+        The updated system with corrected connectivity0, connectivity1,
+        intrascale0, and intrascale1 properties on each perturbable molecule.
+    """
+
+    import sire.legacy.CAS as _SireCAS
+
+    if not isinstance(system, (_System, _LegacySystem)):
+        raise TypeError(
+            "'system' must of type 'sire.system.System' or 'sire.legacy.System.System'"
+        )
+
+    if isinstance(system, _LegacySystem):
+        system = _System(system)
+
+    system = system.clone()
+
+    try:
+        pert_mols = system.molecules("property is_perturbable")
+    except KeyError:
+        raise KeyError("No perturbable molecules in the system")
+
+    r = _SireCAS.Symbol("r")
+
+    for mol in pert_mols:
+        info = mol.info()
+
+        # Build connectivity from bond0 potentials, skipping zero-k bonds.
+        conn0 = _SireMol.Connectivity(info).edit()
+        for bond in mol.property("bond0").potentials():
+            if _SireMM.AmberBond(bond.function(), r).k() != 0.0:
+                conn0.connect(bond.atom0(), bond.atom1())
+        conn0 = conn0.commit()
+
+        # Build connectivity from bond1 potentials, skipping zero-k bonds.
+        conn1 = _SireMol.Connectivity(info).edit()
+        for bond in mol.property("bond1").potentials():
+            if _SireMM.AmberBond(bond.function(), r).k() != 0.0:
+                conn1.connect(bond.atom0(), bond.atom1())
+        conn1 = conn1.commit()
+
+        # Get the 1-4 scale factors from the lambda=0 forcefield.
+        ff = mol.property("forcefield0")
+        sf14 = _SireMM.CLJScaleFactor(
+            ff.electrostatic14_scale_factor(), ff.vdw14_scale_factor()
+        )
+
+        # Build intrascale matrices from the per-state connectivity.
+        intra0 = _SireMM.CLJNBPairs(conn0, sf14)
+        intra1 = _SireMM.CLJNBPairs(conn1, sf14)
+
+        edit_mol = mol.edit()
+
+        if conn0 == conn1:
+            edit_mol = edit_mol.set_property("connectivity", conn0).molecule()
+        else:
+            edit_mol = edit_mol.set_property("connectivity0", conn0).molecule()
+            edit_mol = edit_mol.set_property("connectivity1", conn1).molecule()
+
+        edit_mol = edit_mol.set_property("intrascale0", intra0).molecule()
+        edit_mol = edit_mol.set_property("intrascale1", intra1).molecule()
+
+        system.update(edit_mol.commit())
+
+    return system
+
+
 def reconstruct_system(system):
     """
     Reconstruct a perturbable system to its original state, i.e. extract the

src/somd2/runner/_base.py

@@ -147,6 +147,14 @@ def __init__(self, system, config):
                 _logger.error(msg)
                 raise IOError(msg)
 
+            # Reconstruct end-state connectivity and intrascale matrices from
+            # the bonded terms. The lambda=0 reference topology is used as the
+            # starting point and the pertfile does not express changes in
+            # connectivity or intrascale directly.
+            from .._utils._somd1 import reconstruct_intrascale
+
+            self._system = reconstruct_intrascale(self._system)
+
             # If we're not using SOMD1 compatibility, then reconstruct the original
             # perturbable system. We only need to do this if applying modifications
             # to ghost atom bonded terms.

src/somd2/runner/_repex.py

@@ -1657,6 +1657,10 @@ def _equilibrate(self, index):
             else:
                 system.set_time(_sr.u("0ps"))
 
+            # Resolve the water count while the context is still alive.
+            if gcmc_sampler is not None:
+                gcmc_sampler.num_waters()
+
             # Delete the dynamics object.
             self._dynamics_cache.delete(index)
 
@@ -1992,13 +1996,6 @@ def _reset_gcmc_sampler(gcmc_sampler, dynamics):
         dynamics: sire.mol.Dynamics
             The dynamics object associated with the GCMC sampler.
         """
-        # Ensure the water count is up to date before resetting. If the last
-        # move was a bulk sampling move, _is_bulk is True and num_waters()
-        # needs the stored context to recompute _N. Calling it here, while
-        # the context is still available, clears _is_bulk so that reset()
-        # can safely null the context.
-        gcmc_sampler.num_waters()
-
         # Reset the GCMC sampler. This resets the sampling statistics and
         # clears the associated OpenMM forces.
         gcmc_sampler.reset()

src/somd2/runner/_runner.py

@@ -628,6 +628,13 @@ def generate_lam_vals(lambda_base, increment=0.001):
             }
         )
 
+        # Resolve the water count while the old context is still alive. If the
+        # last equilibration move was a bulk sampling move, _is_bulk is True
+        # and num_waters() needs the stored context to recompute _N. Creating
+        # a new dynamics object below destroys that context.
+        if gcmc_sampler is not None:
+            gcmc_sampler.num_waters()
+
         # Create the dynamics object.
         dynamics = system.dynamics(**dynamics_kwargs)
 
@@ -639,12 +646,6 @@ def generate_lam_vals(lambda_base, increment=0.001):
         # Reset the GCMC sampler. This resets the sampling statistics and clears
         # the associated OpenMM forces.
         if gcmc_sampler is not None:
-            # Ensure the water count is up to date before resetting. If the
-            # last move was a bulk sampling move, _is_bulk is True and
-            # num_waters() needs the stored context to recompute _N. Calling
-            # it here, while the context is still available, clears _is_bulk
-            # so that reset() can safely null the context.
-            gcmc_sampler.num_waters()
             gcmc_sampler.reset()
 
             # Bind the GCMC sampler to the dynamics object.

tests/_utils/test_somd1.py

@@ -0,0 +1,32 @@
+import pytest
+import sire.legacy.Mol as _SireMol
+
+
+@pytest.fixture
+def mols(request):
+    return request.getfixturevalue(request.param)
+
+
+@pytest.mark.parametrize("mols", ["pert_fwd_mols", "pert_rev_mols"], indirect=True)
+def test_reconstruct_intrascale(mols):
+    """
+    Verify that reconstruct_intrascale correctly rebuilds end-state connectivity
+    and intrascale matrices from bond potentials.
+
+    The forward perturbation has two hydrogen atoms that are real at lambda=0
+    and become ghost atoms (du) at lambda=1; the reverse perturbation is the
+    mirror image. In both cases the reconstructed connectivity objects must
+    differ between the two end states.
+    """
+    from somd2._utils._somd1 import reconstruct_intrascale
+
+    system = reconstruct_intrascale(mols)
+
+    mol = system.molecules("property is_perturbable")[0]
+
+    conn0 = mol.property("connectivity0")
+    conn1 = mol.property("connectivity1")
+
+    assert isinstance(conn0, _SireMol.Connectivity)
+    assert isinstance(conn1, _SireMol.Connectivity)
+    assert conn0 != conn1

tests/conftest.py

@@ -71,7 +71,7 @@ def pert_fwd_mols():
     from somd2._utils._somd1 import apply_pert
 
     mols = sr.load_test_files("somd1_forward.prm7", "somd1_forward.rst7")
-    pert_file = str(Path(__file__).parent / "runner" / "inputs" / "forward.pert")
+    pert_file = str(Path(__file__).parent / "inputs" / "forward.pert")
     return apply_pert(mols, pert_file)
 
 
@@ -84,5 +84,5 @@ def pert_rev_mols():
     from somd2._utils._somd1 import apply_pert
 
     mols = sr.load_test_files("somd1_backward.prm7", "somd1_backward.rst7")
-    pert_file = str(Path(__file__).parent / "runner" / "inputs" / "backward.pert")
+    pert_file = str(Path(__file__).parent / "inputs" / "backward.pert")
     return apply_pert(mols, pert_file)