Apply image shifting by fragment

src/openfe_analysis/rmsd.py

@@ -328,12 +328,14 @@ def gather_rms_data(
             skip = max(n_frames // 500, 1)
 
         u_top = mda.Universe(pdb_topology)
+        protein = u_top.select_atoms("protein")
+        if protein:
+            protein.guess_bonds()
 
         for state_idx in range(n_lambda):
             # cheeky, but we can read the PDB topology once and reuse per universe
             # this then only hits the PDB file once for all replicas
             universe = create_universe_single_state(u_top._topology, ds, state_idx)
-
             prot = universe.select_atoms(protein_selection)
             ligand = universe.select_atoms(ligand_selection)
 
@@ -343,25 +345,28 @@ def gather_rms_data(
                     protein=prot,
                     ligands=[ligand] if ligand.n_atoms > 0 else None,
                 )
+
             elif ligand.n_atoms > 0:
                 apply_ligand_alignment_transformations(universe, ligand=ligand)
 
+            output["time(ps)"] = (
+                np.arange(len(universe.trajectory))[::skip] * universe.trajectory.dt
+            )
+            # unwrap/shift/align run once per frame, here
+            universe.transfer_to_memory(step=skip)
+
             if prot:
-                prot_rmsd = RMSDAnalysis(prot).run(step=skip)
+                prot_rmsd = RMSDAnalysis(prot).run()
                 output["protein_RMSD"].append(prot_rmsd.results.rmsd)
 
-                prot_rmsd2d = Protein2DRMSD(prot).run(step=skip)
+                prot_rmsd2d = Protein2DRMSD(prot).run()
                 output["protein_2D_RMSD"].append(prot_rmsd2d.results.rmsd2d)
 
             if ligand:
-                lig_rmsd = RMSDAnalysis(ligand, mass_weighted=True).run(step=skip)
+                lig_rmsd = RMSDAnalysis(ligand, mass_weighted=True).run()
                 output["ligand_RMSD"].append(lig_rmsd.results.rmsd)
 
-                lig_com_drift = LigandCOMDrift(ligand).run(step=skip)
+                lig_com_drift = LigandCOMDrift(ligand).run()
                 output["ligand_wander"].append(lig_com_drift.results.com_drift)
 
-            output["time(ps)"] = (
-                np.arange(len(universe.trajectory))[::skip] * universe.trajectory.dt
-            )
-
     return output

src/openfe_analysis/tests/test_rmsd_classes.py

@@ -28,6 +28,7 @@ def ligand(hybrid_system_skipped_pdb, simulation_skipped_nc):
     universe = universe_utils.create_universe_single_state(
         hybrid_system_skipped_pdb, simulation_skipped_nc, state=0
     )
+    universe.select_atoms("protein").guess_bonds()
     prot = universe.select_atoms("protein and name CA")
     ligand = universe.select_atoms("resname UNK")
     apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand])
@@ -150,6 +151,7 @@ def test_separate_ligands_fixes_pbc_spike(
 
             # Combined approach should produce a spike
             u_combined = universe_utils.create_universe_single_state(d["pdb"], ds, state=state_idx)
+            u_combined.select_atoms("protein").guess_bonds()
             prot = u_combined.select_atoms("protein and name CA")
             lig_A = u_combined.atoms[d["ligand_A_indices"]]
             lig_B = u_combined.atoms[d["ligand_B_indices"]]
@@ -164,6 +166,7 @@ def test_separate_ligands_fixes_pbc_spike(
 
             # Separate approach should fix it
             u_separate = universe_utils.create_universe_single_state(d["pdb"], ds, state=state_idx)
+            u_separate.select_atoms("protein").guess_bonds()
             prot = u_separate.select_atoms("protein and name CA")
             lig_A = u_separate.atoms[d["ligand_A_indices"]]
             lig_B = u_separate.atoms[d["ligand_B_indices"]]

src/openfe_analysis/tests/utils/test_apply_transformations.py

@@ -24,10 +24,11 @@ def universe_single_state(hybrid_system_skipped_pdb, simulation_skipped_nc):
 def test_chain_radius_of_gyration_stable(universe_single_state):
     """Protein chains should not explode or collapse due to PBC errors
     after applying alignment transformations."""
-    protein = universe_single_state.select_atoms("protein and name CA")
-    apply_transformations.apply_complex_alignment_transformations(universe_single_state, protein)
+    universe_single_state.select_atoms("protein").guess_bonds()
+    prot = universe_single_state.select_atoms("protein and name CA")
+    apply_transformations.apply_complex_alignment_transformations(universe_single_state, prot)
 
-    chain = protein.segments[0].atoms
+    chain = prot.segments[0].atoms
     rgs = []
     for ts in universe_single_state.trajectory[:50]:
         rgs.append(chain.radius_of_gyration())
@@ -54,6 +55,7 @@ def test_multichain_rmsd_shifting(simulation_skipped_nc, hybrid_system_skipped_p
     u = universe_utils.create_universe_single_state(
         hybrid_system_skipped_pdb, simulation_skipped_nc, 0
     )
+    u.select_atoms("protein").guess_bonds()
     prot = u.select_atoms("protein and name CA")
     # Do other transformations, but no shifting
     unwrap_tr = unwrap(prot)
@@ -75,7 +77,9 @@ def test_multichain_rmsd_shifting(simulation_skipped_nc, hybrid_system_skipped_p
     u2 = universe_utils.create_universe_single_state(
         hybrid_system_skipped_pdb, simulation_skipped_nc, 0
     )
+    u2.select_atoms("protein").guess_bonds()
     prot2 = u2.select_atoms("protein and name CA")
+    assert len(list(prot2.fragments)) == 2
     apply_transformations.apply_complex_alignment_transformations(u2, protein=prot2)
 
     R2 = rms.RMSD(prot2)
@@ -87,6 +91,7 @@ def test_multichain_rmsd_shifting(simulation_skipped_nc, hybrid_system_skipped_p
 
 def test_rmsd_reference_is_first_frame(universe_single_state):
     """After alignment, RMSD at the first frame should be zero."""
+    universe_single_state.select_atoms("protein").guess_bonds()
     prot = universe_single_state.select_atoms("protein and name CA")
     apply_transformations.apply_complex_alignment_transformations(
         universe_single_state, protein=prot
@@ -109,6 +114,7 @@ def test_empty_protein_raises(universe_single_state):
 
 
 def test_atomgroup_as_ligands_raises(universe_single_state):
+    universe_single_state.select_atoms("protein").guess_bonds()
     prot = universe_single_state.select_atoms("protein and name CA")
     lig = universe_single_state.select_atoms("resname UNK")
     with pytest.raises(TypeError, match="list of AtomGroups"):
@@ -123,3 +129,14 @@ def test_empty_ligand_raises(universe_single_state):
         apply_transformations.apply_ligand_alignment_transformations(
             universe_single_state, ligand=empty
         )
+
+
+def test_missing_protein_bond_raises(universe_single_state):
+    prot = universe_single_state.select_atoms("protein and name CA")
+    ligand = universe_single_state.select_atoms("resname UNK")
+    assert not prot.bonds  # precondition: the guard's trigger is actually present
+
+    with pytest.raises(ValueError, match="no bonds"):
+        apply_transformations.apply_complex_alignment_transformations(
+            universe_single_state, prot, [ligand]
+        )

src/openfe_analysis/tests/utils/test_universe_utils.py

@@ -13,6 +13,7 @@ def universe(hybrid_system_skipped_pdb, simulation_skipped_nc):
     universe = create_universe_single_state(
         hybrid_system_skipped_pdb, simulation_skipped_nc, state=0
     )
+    universe.select_atoms("protein").guess_bonds()
     prot = universe.select_atoms("protein and name CA")
     ligand = universe.select_atoms("resname UNK")
     apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand])
@@ -25,6 +26,7 @@ def ligand_ag(hybrid_system_skipped_pdb, simulation_skipped_nc):
     universe = create_universe_single_state(
         hybrid_system_skipped_pdb, simulation_skipped_nc, state=0
     )
+    universe.select_atoms("protein").guess_bonds()
     prot = universe.select_atoms("protein and name CA")
     ligand = universe.select_atoms("resname UNK")
     apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand])

src/openfe_analysis/utils/apply_transformations.py

@@ -44,6 +44,12 @@ def apply_complex_alignment_transformations(
     if protein is None or not protein:
         raise ValueError("protein AtomGroup is empty or None")
 
+    if not protein.bonds:
+        raise ValueError(
+            "protein AtomGroup has no bonds which would lead to wrong alignment. "
+            "Call guess_bonds() on the protein before applying these transformations."
+        )
+
     if isinstance(ligands, mda.AtomGroup):
         raise TypeError(
             "ligands must be a list of AtomGroups, not a single AtomGroup. "
@@ -58,11 +64,15 @@ def apply_complex_alignment_transformations(
     # 1. Make molecules whole (protein + optional ligand)
     transforms = [unwrap(group)]
 
-    # 2. Closest image shift for protein chains + ligand (if present)
-    chains = [seg.atoms for seg in protein.segments]
-    shift_targets = chains[1:] + ligands
+    # 2. Closest image shift for protein fragments + ligand (if present)
+    fragments = list(protein.fragments)
+    # Pick the largest fragment as the reference
+    ref_idx = max(range(len(fragments)), key=lambda i: fragments[i].n_atoms)
+    reference = fragments[ref_idx]
+    shift_targets = [f for i, f in enumerate(fragments) if i != ref_idx]
+    shift_targets += ligands
     if shift_targets:
-        transforms.append(ClosestImageShift(chains[0], shift_targets))
+        transforms.append(ClosestImageShift(reference=reference, targets=shift_targets))
 
     # 3. Align on protein backbone/atoms
     transforms.append(Aligner(protein))