Add analysis for symmetry corrected RMSD

docs/environment.yaml

@@ -7,3 +7,5 @@ dependencies:
 - python=3.12
 - sphinx
 - furo
+- spyrmsd
+- rdkit

environment.yml

@@ -18,3 +18,5 @@ dependencies:
   - pytest-cov
   - pytest-xdist
   - pytest-rerunfailures
+  - spyrmsd
+  - rdkit

pyproject.toml

@@ -18,6 +18,8 @@ dependencies = [
     'netCDF4',
     'openff-units',
     'pyyaml',
+    'spyrmsd',
+    'rdkit',
     'matplotlib',
 ]
 description="Trajectory analysis of free energy calculations."

src/openfe_analysis/rmsd.py

@@ -5,9 +5,10 @@
 import MDAnalysis as mda
 import netCDF4 as nc
 import numpy as np
+import spyrmsd.rmsd as srmsd
 from MDAnalysis.analysis import rms
 from MDAnalysis.analysis.base import AnalysisBase
-from MDAnalysis.transformations import unwrap
+from rdkit import Chem
 
 from .utils.apply_transformations import apply_alignment_transformations
 from .utils.universe_utils import create_universe_single_state
@@ -130,6 +131,85 @@ def _single_frame(self) -> None:
         )
 
 
+class SymmetryCorrectedLigandRMSD(AnalysisBase):
+    """
+    Symmetry-corrected 1D RMSD time series for a ligand AtomGroup.
+
+    Parameters
+    ----------
+    atomgroup : mda.AtomGroup
+        Ligand atoms to compute RMSD for. If ``rdmol`` is not provided,
+        bonds must be guessed on the atomgroup before instantiating this
+        class; use :func:`guess_ligand_bonds` for this purpose.
+    rdmol : Chem.Mol, optional
+        RDKit molecule corresponding to ``atomgroup``. If provided, it is
+        used directly and ``guess_ligand_bonds`` does not need to be called.
+        If ``None``, the RDKit molecule is derived from ``atomgroup`` via
+        ``convert_to("RDKIT")``.
+
+    Raises
+    ------
+    ValueError
+        If ``rdmol`` is ``None`` and no bonds are found on the atomgroup.
+    ValueError
+        If the number of atoms in ``atomgroup`` and ``rdmol`` do not match.
+    """
+
+    _analysis_algorithm_is_parallelizable = False
+
+    def __init__(
+        self,
+        atomgroup: mda.AtomGroup,
+        rdmol: Optional[Chem.Mol] = None,
+        **kwargs,
+    ):
+        super().__init__(atomgroup.universe.trajectory, **kwargs)
+        self._ag = atomgroup
+        if rdmol is None:
+            try:
+                has_bonds = len(atomgroup.bonds) > 0
+            except mda.exceptions.NoDataError:
+                has_bonds = False
+            if not has_bonds:
+                raise ValueError(
+                    "No bonds found on atomgroup. Call guess_ligand_bonds() "
+                    "before instantiating SymmetryCorrectedLigandRMSD, or "
+                    "pass an rdmol directly."
+                )
+        else:
+            if len(atomgroup) != rdmol.GetNumAtoms():
+                raise ValueError(
+                    f"atomgroup has {len(atomgroup)} atoms but rdmol has "
+                    f"{rdmol.GetNumAtoms()} atoms."
+                )
+        self._mol = rdmol if rdmol is not None else atomgroup.convert_to("RDKIT")
+        self._aprops = np.array([atom.GetAtomicNum() for atom in self._mol.GetAtoms()])
+        self._am = Chem.rdmolops.GetAdjacencyMatrix(self._mol)
+
+    def _prepare(self):
+        self.results.rmsd = np.zeros(self.n_frames, dtype=np.float64)
+        # reference is taken from the first analyzed frame, not necessarily frame 0
+        self._reference = self._ag.positions.copy()
+        self._isomorphisms: list | None = None
+
+    def _single_frame(self) -> None:
+        frame_rmsd, isomorphisms, _ = srmsd._rmsd_isomorphic_core(
+            coords1=self._ag.positions.copy(),
+            coords2=self._reference,
+            aprops1=self._aprops,
+            aprops2=self._aprops,
+            am1=self._am,
+            am2=self._am,
+            center=False,
+            minimize=False,
+            isomorphisms=self._isomorphisms,
+        )
+        self.results.rmsd[self._frame_index] = frame_rmsd
+        # cache isomorphisms after first frame to avoid redundant graph matching
+        if self._isomorphisms is None:
+            self._isomorphisms = isomorphisms
+
+
 class LigandCOMDrift(AnalysisBase):
     """
     Ligand center-of-mass displacement from initial position.

src/openfe_analysis/tests/test_rmsd_classes.py

@@ -6,8 +6,14 @@
 from MDAnalysis.analysis import diffusionmap, rms
 from MDAnalysisTests.datafiles import DCD, PSF
 from numpy.testing import assert_allclose, assert_almost_equal
-
-from openfe_analysis.rmsd import LigandCOMDrift, Protein2DRMSD, RMSDAnalysis
+from rdkit import Chem
+
+from openfe_analysis.rmsd import (
+    LigandCOMDrift,
+    Protein2DRMSD,
+    RMSDAnalysis,
+    SymmetryCorrectedLigandRMSD,
+)
 from openfe_analysis.utils import apply_transformations, universe_utils
 
 
@@ -28,6 +34,21 @@ def ligand(hybrid_system_skipped_pdb, simulation_skipped_nc):
     universe.trajectory.close()
 
 
+@pytest.fixture
+def minimal_universe():
+    """Minimal 3-atom water-like universe without bonds."""
+    u = mda.Universe.empty(3, n_residues=1, trajectory=True)
+    u.add_TopologyAttr("elements", ["O", "H", "H"])
+    u.add_TopologyAttr("names", ["O", "H1", "H2"])
+    u.add_TopologyAttr("resnames", ["UNK"])
+    u.add_TopologyAttr("resids", [1])
+    u.load_new(
+        np.array([[[0.0, 0.0, 0.0], [1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]]),
+        order="fac",
+    )
+    return u.select_atoms("all")
+
+
 @pytest.fixture()
 def correct_values():
     return [0, 4.68953]
@@ -133,3 +154,71 @@ def test_ligand_com_drift(ligand):
     expected = [0.0, 0.38549, 0.61483, 0.54140, 1.26861, 0.92772]
     assert len(result.results.com_drift) == len(ligand.universe.trajectory[::10])
     assert_allclose(result.results.com_drift[:6], expected, rtol=1e-3)
+
+
+class TestSymmetryCorrectedLigandRMSD:
+    def test_regression(self, ligand):
+        state_lig = universe_utils.select_state_atoms(ligand.universe, end_state="A").select_atoms(
+            "resname UNK"
+        )
+        result = SymmetryCorrectedLigandRMSD(state_lig).run(step=10)
+        expected = [0.0, 0.75138, 2.09003, 0.95125, 1.54566, 2.00029]
+        assert_allclose(result.results.rmsd[:6], expected, rtol=1e-3)
+
+    def test_zero_for_valid_swap(self):
+        """
+        For a water-like symmetric molecule, swapping the two equivalent H atoms
+        gives naive RMSD > 0 but SymmetryCorrectedLigandRMSD = 0.
+        """
+        # Build a minimal universe with two frames: reference and swapped
+        coords_ref = np.array(
+            [
+                [0.0, 0.0, 0.0],  # O
+                [1.0, 0.0, 0.0],  # H1
+                [0.0, 1.0, 0.0],  # H2
+            ]
+        )
+        coords_swapped = np.array(
+            [
+                [0.0, 0.0, 0.0],  # O
+                [0.0, 1.0, 0.0],  # H2 in H1's slot
+                [1.0, 0.0, 0.0],  # H1 in H2's slot
+            ]
+        )
+
+        u = mda.Universe.empty(3, trajectory=True)
+        u.add_TopologyAttr("elements", ["O", "H", "H"])
+        u.add_TopologyAttr("names", ["O", "H1", "H2"])
+        u.add_TopologyAttr("resnames", ["UNK"])
+        u.add_TopologyAttr("resids", [1])
+        u.add_TopologyAttr("bonds", [(0, 1), (0, 2)])
+        u.load_new(
+            np.array([coords_ref, coords_swapped]),
+            order="fac",
+        )
+
+        ag = u.select_atoms("all")
+
+        corrected = SymmetryCorrectedLigandRMSD(ag).run()
+        naive = RMSDAnalysis(ag).run()
+
+        # Frame 0 is reference — both should be 0
+        assert corrected.results.rmsd[0] == pytest.approx(0.0, abs=1e-5)
+        assert naive.results.rmsd[0] == pytest.approx(0.0, abs=1e-5)
+
+        # Frame 1 is the swap — naive sees displacement, corrected sees zero
+        assert naive.results.rmsd[1] > 0.0
+        assert corrected.results.rmsd[1] == pytest.approx(0.0, abs=1e-5)
+
+    def test_raises_on_missing_bonds(self, minimal_universe):
+        """Should raise ValueError if atomgroup has no bonds and no rdmol is provided."""
+        with pytest.raises(ValueError, match="No bonds found"):
+            SymmetryCorrectedLigandRMSD(minimal_universe)
+
+    def test_raises_on_atom_count_mismatch(self, minimal_universe):
+        """Should raise ValueError if atomgroup and rdmol have different atom counts."""
+        mol = Chem.RWMol()
+        mol.AddAtom(Chem.Atom(8))  # O
+        mol.AddAtom(Chem.Atom(1))  # H
+        with pytest.raises(ValueError, match="atomgroup has 3 atoms but rdmol has 2"):
+            SymmetryCorrectedLigandRMSD(minimal_universe, rdmol=mol.GetMol())

src/openfe_analysis/tests/utils/test_universe_utils.py

@@ -0,0 +1,79 @@
+import pytest
+
+from openfe_analysis.utils import apply_transformations
+from openfe_analysis.utils.universe_utils import (
+    create_universe_single_state,
+    guess_ligand_bonds,
+    select_state_atoms,
+)
+
+
+@pytest.fixture
+def universe(hybrid_system_skipped_pdb, simulation_skipped_nc):
+    universe = create_universe_single_state(
+        hybrid_system_skipped_pdb, simulation_skipped_nc, state=0
+    )
+    prot = universe.select_atoms("protein and name CA")
+    ligand = universe.select_atoms("resname UNK")
+    apply_transformations.apply_alignment_transformations(universe, prot, ligand)
+    yield universe
+    universe.trajectory.close()
+
+
+@pytest.fixture
+def ligand_ag(hybrid_system_skipped_pdb, simulation_skipped_nc):
+    universe = create_universe_single_state(
+        hybrid_system_skipped_pdb, simulation_skipped_nc, state=0
+    )
+    prot = universe.select_atoms("protein and name CA")
+    ligand = universe.select_atoms("resname UNK")
+    apply_transformations.apply_alignment_transformations(universe, prot, ligand)
+    ag = select_state_atoms(universe, end_state="A").select_atoms("resname UNK")
+    yield ag
+    universe.trajectory.close()
+
+
+def test_guess_ligand_bonds_adds_bonds(ligand_ag):
+    """Bonds should be present on the atomgroup after guess_ligand_bonds."""
+    original_count = len(ligand_ag.bonds)
+    # This also has stateB bond
+    assert original_count == 49
+    guess_ligand_bonds(ligand_ag, delete_existing=True)
+    # Now only 48 stateA bonds
+    assert len(ligand_ag.bonds) == 48
+
+
+def test_guess_ligand_bonds_modifies_universe_inplace(ligand_ag):
+    """Bond topology should be reflected on the parent universe after guessing."""
+    guess_ligand_bonds(ligand_ag)
+    universe_bonds = ligand_ag.universe.select_atoms("resname UNK").bonds
+    assert len(universe_bonds) > 0
+
+
+@pytest.mark.parametrize(
+    "end_state, expected_bfactors",
+    [
+        ("A", (0.25, 0.5)),
+        ("B", (0.75, 0.5)),
+    ],
+)
+def test_select_state_atoms(universe, end_state, expected_bfactors):
+    """State selection should include state-unique and shared atoms."""
+    state = select_state_atoms(universe, end_state=end_state)
+    assert len(state) > 0
+    assert all(atom.bfactor in expected_bfactors for atom in state)
+
+
+def test_select_state_atoms_invalid_state(universe):
+    """Invalid end_state should raise a ValueError."""
+    with pytest.raises(ValueError, match="end_state must be 'A' or 'B'"):
+        select_state_atoms(universe, end_state="C")
+
+
+def test_select_state_atoms_shared_atoms(universe):
+    """Shared atoms (bfactor 0.5) should appear in both state A and B selections."""
+    state_a = select_state_atoms(universe, end_state="A")
+    state_b = select_state_atoms(universe, end_state="B")
+    shared_a = set(atom.ix for atom in state_a if atom.bfactor == 0.5)
+    shared_b = set(atom.ix for atom in state_b if atom.bfactor == 0.5)
+    assert shared_a == shared_b