From 498c71f41652186f90aa8f3edbe7986f74f5e12b Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 8 Jul 2026 13:28:51 +0200 Subject: [PATCH 1/8] Apply image shifting by fragment --- src/openfe_analysis/rmsd.py | 2 +- .../tests/utils/test_apply_transformations.py | 3 +++ src/openfe_analysis/utils/apply_transformations.py | 12 ++++++++---- 3 files changed, 12 insertions(+), 5 deletions(-) diff --git a/src/openfe_analysis/rmsd.py b/src/openfe_analysis/rmsd.py index 9c86200..9d79cc2 100644 --- a/src/openfe_analysis/rmsd.py +++ b/src/openfe_analysis/rmsd.py @@ -328,12 +328,12 @@ def gather_rms_data( skip = max(n_frames // 500, 1) u_top = mda.Universe(pdb_topology) + u_top.select_atoms("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) diff --git a/src/openfe_analysis/tests/utils/test_apply_transformations.py b/src/openfe_analysis/tests/utils/test_apply_transformations.py index 13f66ee..6217ddb 100644 --- a/src/openfe_analysis/tests/utils/test_apply_transformations.py +++ b/src/openfe_analysis/tests/utils/test_apply_transformations.py @@ -54,6 +54,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 +76,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) diff --git a/src/openfe_analysis/utils/apply_transformations.py b/src/openfe_analysis/utils/apply_transformations.py index 49f3e6f..a710df0 100644 --- a/src/openfe_analysis/utils/apply_transformations.py +++ b/src/openfe_analysis/utils/apply_transformations.py @@ -58,11 +58,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)) From 3fd980b2b1f1b8136a666c2b78ad42a34bba3f05 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 8 Jul 2026 15:56:08 +0200 Subject: [PATCH 2/8] Update tests --- src/openfe_analysis/tests/test_rmsd_classes.py | 3 +++ 1 file changed, 3 insertions(+) diff --git a/src/openfe_analysis/tests/test_rmsd_classes.py b/src/openfe_analysis/tests/test_rmsd_classes.py index 4bacb83..8050ac1 100644 --- a/src/openfe_analysis/tests/test_rmsd_classes.py +++ b/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"]] From e358ac13b1f659dd06dd44d8c05a5be4d61e02d2 Mon Sep 17 00:00:00 2001 From: Hannah Baumann <43765638+hannahbaumann@users.noreply.github.com> Date: Fri, 7 Aug 2026 11:21:49 +0200 Subject: [PATCH 3/8] Apply suggestion from @IAlibay Co-authored-by: Irfan Alibay --- src/openfe_analysis/rmsd.py | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/src/openfe_analysis/rmsd.py b/src/openfe_analysis/rmsd.py index 9d79cc2..6d64820 100644 --- a/src/openfe_analysis/rmsd.py +++ b/src/openfe_analysis/rmsd.py @@ -328,7 +328,9 @@ def gather_rms_data( skip = max(n_frames // 500, 1) u_top = mda.Universe(pdb_topology) - u_top.select_atoms("protein").guess_bonds() + 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 From 6b9ab583f2d37902f5680ce1a527f1cd432097eb Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 12 Aug 2026 14:05:39 +0200 Subject: [PATCH 4/8] Do alignment in memory --- src/openfe_analysis/rmsd.py | 19 +++++++++++-------- 1 file changed, 11 insertions(+), 8 deletions(-) diff --git a/src/openfe_analysis/rmsd.py b/src/openfe_analysis/rmsd.py index 9d79cc2..2c4544e 100644 --- a/src/openfe_analysis/rmsd.py +++ b/src/openfe_analysis/rmsd.py @@ -343,25 +343,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 From 37bcadc2d5c1b02ce0d3252a13519270fc504810 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 12 Aug 2026 14:23:15 +0200 Subject: [PATCH 5/8] Raise error if protein doesnt have bonds --- src/openfe_analysis/utils/apply_transformations.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/src/openfe_analysis/utils/apply_transformations.py b/src/openfe_analysis/utils/apply_transformations.py index a710df0..bfb2fb3 100644 --- a/src/openfe_analysis/utils/apply_transformations.py +++ b/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. " From 48dc08eda167dd7c54cf4b69d63e53700083fb98 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Wed, 12 Aug 2026 14:45:17 +0200 Subject: [PATCH 6/8] Update tests --- src/openfe_analysis/tests/utils/test_apply_transformations.py | 3 +++ src/openfe_analysis/tests/utils/test_universe_utils.py | 2 ++ 2 files changed, 5 insertions(+) diff --git a/src/openfe_analysis/tests/utils/test_apply_transformations.py b/src/openfe_analysis/tests/utils/test_apply_transformations.py index 6217ddb..3741c89 100644 --- a/src/openfe_analysis/tests/utils/test_apply_transformations.py +++ b/src/openfe_analysis/tests/utils/test_apply_transformations.py @@ -25,6 +25,7 @@ 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") + protein.guess_bonds() apply_transformations.apply_complex_alignment_transformations(universe_single_state, protein) chain = protein.segments[0].atoms @@ -91,6 +92,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.""" prot = universe_single_state.select_atoms("protein and name CA") + prot.guess_bonds() apply_transformations.apply_complex_alignment_transformations( universe_single_state, protein=prot ) @@ -113,6 +115,7 @@ def test_empty_protein_raises(universe_single_state): def test_atomgroup_as_ligands_raises(universe_single_state): prot = universe_single_state.select_atoms("protein and name CA") + prot.guess_bonds() lig = universe_single_state.select_atoms("resname UNK") with pytest.raises(TypeError, match="list of AtomGroups"): apply_transformations.apply_complex_alignment_transformations( diff --git a/src/openfe_analysis/tests/utils/test_universe_utils.py b/src/openfe_analysis/tests/utils/test_universe_utils.py index 519f64f..c2b7820 100644 --- a/src/openfe_analysis/tests/utils/test_universe_utils.py +++ b/src/openfe_analysis/tests/utils/test_universe_utils.py @@ -14,6 +14,7 @@ def universe(hybrid_system_skipped_pdb, simulation_skipped_nc): hybrid_system_skipped_pdb, simulation_skipped_nc, state=0 ) prot = universe.select_atoms("protein and name CA") + prot.guess_bonds() ligand = universe.select_atoms("resname UNK") apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand]) yield universe @@ -26,6 +27,7 @@ def ligand_ag(hybrid_system_skipped_pdb, simulation_skipped_nc): hybrid_system_skipped_pdb, simulation_skipped_nc, state=0 ) prot = universe.select_atoms("protein and name CA") + prot.guess_bonds() ligand = universe.select_atoms("resname UNK") apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand]) ag = select_state_atoms(universe, end_state="A").select_atoms("resname UNK") From 9d97900287104de0101230451f5f4fd4cb384a10 Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 13 Aug 2026 09:56:12 +0200 Subject: [PATCH 7/8] Fix tests --- .../tests/utils/test_apply_transformations.py | 12 ++++++------ .../tests/utils/test_universe_utils.py | 4 ++-- 2 files changed, 8 insertions(+), 8 deletions(-) diff --git a/src/openfe_analysis/tests/utils/test_apply_transformations.py b/src/openfe_analysis/tests/utils/test_apply_transformations.py index 3741c89..7d3f5ab 100644 --- a/src/openfe_analysis/tests/utils/test_apply_transformations.py +++ b/src/openfe_analysis/tests/utils/test_apply_transformations.py @@ -24,11 +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") - protein.guess_bonds() - 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()) @@ -91,8 +91,8 @@ 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") - prot.guess_bonds() apply_transformations.apply_complex_alignment_transformations( universe_single_state, protein=prot ) @@ -114,8 +114,8 @@ 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") - prot.guess_bonds() lig = universe_single_state.select_atoms("resname UNK") with pytest.raises(TypeError, match="list of AtomGroups"): apply_transformations.apply_complex_alignment_transformations( diff --git a/src/openfe_analysis/tests/utils/test_universe_utils.py b/src/openfe_analysis/tests/utils/test_universe_utils.py index c2b7820..4f2266a 100644 --- a/src/openfe_analysis/tests/utils/test_universe_utils.py +++ b/src/openfe_analysis/tests/utils/test_universe_utils.py @@ -13,8 +13,8 @@ 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") - prot.guess_bonds() ligand = universe.select_atoms("resname UNK") apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand]) yield universe @@ -26,8 +26,8 @@ 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") - prot.guess_bonds() ligand = universe.select_atoms("resname UNK") apply_transformations.apply_complex_alignment_transformations(universe, prot, [ligand]) ag = select_state_atoms(universe, end_state="A").select_atoms("resname UNK") From 94e56f15a3cd9ea9c161cddc901088d2c8b1757b Mon Sep 17 00:00:00 2001 From: hannahbaumann Date: Thu, 13 Aug 2026 10:14:34 +0200 Subject: [PATCH 8/8] Add test for missing bond error --- .../tests/utils/test_apply_transformations.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/src/openfe_analysis/tests/utils/test_apply_transformations.py b/src/openfe_analysis/tests/utils/test_apply_transformations.py index 7d3f5ab..d909930 100644 --- a/src/openfe_analysis/tests/utils/test_apply_transformations.py +++ b/src/openfe_analysis/tests/utils/test_apply_transformations.py @@ -129,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] + )