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/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..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))