Update topology/system particle bookkeeping in OpenMM

docs/releasehistory.md

@@ -11,6 +11,10 @@ Dates are given in YYYY-MM-DD format.
 
 Please note that all releases prior to a version 1.0.0 are considered pre-releases and many API changes will come before a stable release.
 
+## Current development
+
+* #1051 Updates `to_openmm_topology` to place virtual sites at the end, matching the behavior of `to_openmm_system`. The old behavior (collating within molecules) is available with `collate=True`.
+
 ## 0.4.0 - 2024
 
 * Pydantic v2 is now used at runtime. As a consequence, models containing `Interchange`s cannot also use models from the v1 API.

docs/using/output.md

@@ -69,6 +69,8 @@ openmm_positions: openmm.unit.Quantity = interchange.positions.to_openmm()
 openmm_box: openmm.unit.Quantity = interchange.box.to_openmm()
 ```
 
+If virtual sites are present in the system, they will all be placed at the end of the system and the topology. In this case, all virtual sites are placed in a new "VS" residue at the end of the topology. Optionally, virtual sites can collated within molecules in the topology, associated with the last residue in each molecule. In this case, all virtual sites are still placed at the end of the system. To do this, use `collate=True` as an argument to `Interchange.to_openmm_topology`. For discussion, see [issue #1049](https://github.com/openforcefield/openff-interchange/issues/1049).
+
 ## Amber
 
 An `Interchange` object can be written to Amber parameter/topology, coordinate, and SANDER run input files with [`Interchange.to_prmtop()`], [`Interchange.to_inpcrd()`], and [`Interchange.to_sander_input()`]:

examples/virtual_sites/virtual_sites.ipynb

@@ -88,7 +88,7 @@
     "        mdtraj.load(\n",
     "            \"trajectory.dcd\",\n",
     "            top=mdtraj.Topology.from_openmm(\n",
-    "                interchange.to_openmm_topology(),\n",
+    "                interchange.to_openmm_topology(collate=True),\n",
     "            ),\n",
     "        )\n",
     "    )"

openff/interchange/_tests/interoperability_tests/test_openmm.py

@@ -622,10 +622,11 @@ def test_valence_term_paticle_index_offsets(self, water, tip5p):
 
 @skip_if_missing("openmm")
 class TestToOpenMMTopology:
-    def test_num_virtual_sites(self, water, tip4p):
+    @pytest.mark.parametrize("collate", [True, False])
+    def test_num_virtual_sites(self, water, tip4p, collate):
         out = Interchange.from_smirnoff(tip4p, [water])
 
-        assert _get_num_virtual_sites(to_openmm_topology(out)) == 1
+        assert _get_num_virtual_sites(to_openmm_topology(out, collate=collate)) == 1
 
         # TODO: Monkeypatch Topology.to_openmm() and emit a warning when it seems
         #       to be used while virtual sites are present in a handler

openff/interchange/_tests/test_issues.py

@@ -112,3 +112,26 @@ def test_issue_1031(monkeypatch):
     # check a few atom names to ensure these didn't end up being empty sets
     for atom_name in ("NE2", "H3", "HA", "CH3", "CA", "CB", "CE1"):
         assert atom_name in openff_atom_names
+
+
+def test_issue_1049():
+    pytest.importorskip("openmm")
+
+    topology = Topology.from_molecules(
+        [
+            Molecule.from_smiles("C"),
+            Molecule.from_smiles("O"),
+            Molecule.from_smiles("O"),
+        ],
+    )
+
+    interchange = ForceField("openff-2.2.0.offxml", "opc.offxml").create_interchange(topology)
+
+    openmm_topology = interchange.to_openmm_topology()
+    openmm_system = interchange.to_openmm_system()
+
+    # the same index in system should also be a virtual site in the topology
+    for particle_index, particle in enumerate(openmm_topology.atoms()):
+        assert openmm_system.isVirtualSite(particle_index) == (
+            particle.element is None
+        ), f"particle index {particle_index} is a virtual site in the system OR topology but not both"

openff/interchange/_tests/unit_tests/interop/openmm/test_topology.py

@@ -0,0 +1,30 @@
+import pytest
+from openff.toolkit import Topology
+
+
+@pytest.mark.parametrize("collate", [True, False])
+def test_collate_or_not_virtual_site_ordering(
+    tip4p,
+    water,
+    collate,
+):
+    pytest.importorskip("openmm")
+
+    openmm_topology = tip4p.create_interchange(
+        Topology.from_molecules([water, water]),
+    ).to_openmm_topology(
+        collate=collate,
+    )
+
+    if collate:
+        # particles should be O H H VS O H H VS
+        #                     0 1 2 3  4 5 6 7
+        virtual_site_indices = (3, 7)
+
+    else:
+        # particles should be O H H O H H VS VS
+        #                     0 1 2 3 4 5 6  7
+        virtual_site_indices = (6, 7)
+
+    for particle_index, particle in enumerate(openmm_topology.atoms()):
+        assert (particle.element is None) == (particle_index in virtual_site_indices)

openff/interchange/components/interchange.py

@@ -539,13 +539,25 @@ def to_openmm(self, *args, **kwargs):
 
     def to_openmm_topology(
         self,
+        collate: bool = False,
         ensure_unique_atom_names: str | bool = "residues",
     ):
-        """Export components of this Interchange to an OpenMM Topology."""
+        """
+        Export components of this Interchange to an OpenMM Topology.
+
+        Parameters
+        ----------
+        collate
+            If False, the default, all virtual sites will be added to a single residue at the end of the topology.
+            If True, virtual sites will be collated with their associated molecule and added to the residue of the last
+            atom in the molecule they belong to.
+
+        """
         from openff.interchange.interop.openmm._topology import to_openmm_topology
 
         return to_openmm_topology(
             self,
+            collate=collate,
             ensure_unique_atom_names=ensure_unique_atom_names,
         )
 
@@ -646,7 +658,12 @@ def to_openmm_simulation(
 
     @requires_package("openmm")
     def to_pdb(self, file_path: Path | str, include_virtual_sites: bool = False):
-        """Export this Interchange to a .pdb file."""
+        """
+        Export this Interchange to a .pdb file.
+
+        Note that virtual sites are collated into each molecule, which differs from the default
+        behavior of Interchange.to_openmm_topology.
+        """
         from openff.interchange.interop.openmm import _to_pdb
 
         if self.positions is None:
@@ -661,6 +678,7 @@ def to_pdb(self, file_path: Path | str, include_virtual_sites: bool = False):
             )
 
             openmm_topology = self.to_openmm_topology(
+                collate=True,
                 ensure_unique_atom_names=False,
             )
             positions = get_positions_with_virtual_sites(self)

openff/interchange/interop/openmm/_topology.py

@@ -15,9 +15,22 @@
 
 def to_openmm_topology(
     interchange: "Interchange",
+    collate: bool = False,
     ensure_unique_atom_names: str | bool = "residues",
 ) -> "openmm.app.Topology":
-    """Create an OpenMM Topology containing some virtual site information (if appropriate)."""
+    """
+    Create an OpenMM Topology containing some virtual site information (if appropriate).
+
+    Parameters
+    ----------
+    interchange
+        The Interchange object to convert to an OpenMM Topology.
+    collate
+        If False, the default, all virtual sites will be added to a single residue at the end of the topology.
+        If True, virtual sites will be collated with their associated molecule and added to the residue of the last
+        atom in the molecule they belong to.
+
+    """
     # Heavily cribbed from the toolkit
     # https://github.com/openforcefield/openff-toolkit/blob/0.11.0rc2/openff/toolkit/topology/topology.py
 
@@ -126,7 +139,7 @@ def to_openmm_topology(
             last_chain = chain
             last_residue = residue
 
-        if has_virtual_sites:
+        if has_virtual_sites and collate:
             virtual_sites_in_this_molecule: list[VirtualSiteKey] = molecule_virtual_site_map[molecule_index]
             for this_virtual_site in virtual_sites_in_this_molecule:
                 virtual_site_name = this_virtual_site.name
@@ -173,6 +186,21 @@ def to_openmm_topology(
                 order=bond_order,
             )
 
+    if has_virtual_sites and not collate:
+        virtual_site_chain = openmm_topology.addChain(atom_chain_id)
+        virtual_site_residue = openmm_topology.addResidue("VS", virtual_site_chain)
+
+        for molecule_index, molecule in enumerate(topology.molecules):
+            virtual_sites_in_this_molecule: list[VirtualSiteKey] = molecule_virtual_site_map[molecule_index]
+            for this_virtual_site in virtual_sites_in_this_molecule:
+                virtual_site_name = this_virtual_site.name
+
+                openmm_topology.addAtom(
+                    virtual_site_name,
+                    virtual_site_element,
+                    virtual_site_residue,
+                )
+
     if interchange.box is not None:
         from openff.units.openmm import to_openmm