diff --git a/src/openfe_analysis/rmsd.py b/src/openfe_analysis/rmsd.py index 9c86200..edf22a1 100644 --- a/src/openfe_analysis/rmsd.py +++ b/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 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"]] diff --git a/src/openfe_analysis/tests/utils/test_apply_transformations.py b/src/openfe_analysis/tests/utils/test_apply_transformations.py index 13f66ee..d909930 100644 --- a/src/openfe_analysis/tests/utils/test_apply_transformations.py +++ b/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] + ) diff --git a/src/openfe_analysis/tests/utils/test_universe_utils.py b/src/openfe_analysis/tests/utils/test_universe_utils.py index 519f64f..4f2266a 100644 --- a/src/openfe_analysis/tests/utils/test_universe_utils.py +++ b/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]) diff --git a/src/openfe_analysis/utils/apply_transformations.py b/src/openfe_analysis/utils/apply_transformations.py index 49f3e6f..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. " @@ -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))