Move structural analysis methods to child classes in SepTop Protocol

news/septop-abstract-analysis.rst

@@ -0,0 +1,27 @@
+**Added:**
+
+* <news item>
+
+**Changed:**
+
+* The per-phase trajectory analysis in the SepTop protocol is now run
+  through an abstract ``_run_trajectory_analysis`` method on
+  ``BaseSepTopAnalysisUnit``, implemented separately by the complex and solvent
+  analysis units, rather than branching on the simulation type within
+  ``_structural_analysis`` (PR #2049).
+
+**Deprecated:**
+
+* <news item>
+
+**Removed:**
+
+* <news item>
+
+**Fixed:**
+
+* <news item>
+
+**Security:**
+
+* <news item>

src/openfe/protocols/openmm_septop/base_units.py

@@ -35,17 +35,7 @@
 )
 from gufe.components import Component
 from gufe.protocols.errors import ProtocolUnitExecutionError
-from openfe_analysis.rmsd import (
-    LigandCOMDrift,
-    Protein2DRMSD,
-    SymmetryCorrectedLigandRMSD,
-)
 from openfe_analysis.utils import plotting
-from openfe_analysis.utils.apply_transformations import (
-    apply_complex_alignment_transformations,
-    apply_ligand_alignment_transformations,
-)
-from openfe_analysis.utils.universe_utils import create_universe_single_state
 from openff.toolkit.topology import Molecule as OFFMolecule
 from openff.units import unit as offunit
 from openff.units.openmm import ensure_quantity, from_openmm, to_openmm
@@ -1520,18 +1510,19 @@ def _analyze_multistate_energies(
         return analyzer.unit_results_dict
 
     @staticmethod
-    def _run_complex_analysis(
-        ds,
+    @abc.abstractmethod
+    def _run_trajectory_analysis(
+        ds: nc.Dataset,
         pdb_file: pathlib.Path,
         skip: int,
         ligand_A_indices: list[int],
         ligand_B_indices: list[int],
         rdmol_A: Chem.Mol,
         rdmol_B: Chem.Mol,
-        protein_selection: str,
+        protein_selection: str | None,
     ) -> tuple[dict[str, list[np.ndarray]], np.ndarray | None]:
         """
-        Run structural analysis for the complex phase.
+        Analyze a trajectory.
 
         Parameters
         ----------
@@ -1549,8 +1540,8 @@ def _run_complex_analysis(
           RDKit molecule for ligand A, used for symmetry-corrected RMSD.
         rdmol_B : Chem.Mol
           RDKit molecule for ligand B, used for symmetry-corrected RMSD.
-        protein_selection : str
-          MDAnalysis selection string for the protein atoms used for
+        protein_selection : str | None
+          Optional MDAnalysis selection string for the protein atoms used for
           alignment and RMSD calculations.
 
         Returns
@@ -1561,115 +1552,13 @@ def _run_complex_analysis(
         time_ps : np.ndarray or None
           Time array in picoseconds corresponding to the analyzed frames.
         """
-        n_lambda = ds.dimensions["state"].size
-        per_state_data: dict[str, list[np.ndarray]] = {
-            "ligand_A_RMSD": [],
-            "ligand_B_RMSD": [],
-            "ligand_A_COM_drift": [],
-            "ligand_B_COM_drift": [],
-            "protein_2D_RMSD": [],
-        }
-        time_ps: np.ndarray | None = None
-        # Read the PDB topology once and reuse across all lambda states.
-        # Passing u_top._topology to create_universe_single_state avoids
-        # reading the PDB file from disk for every lambda.
-        u_top = mda.Universe(pdb_file)
-        prot_indices = u_top.select_atoms(protein_selection).indices
-        for state_idx in range(n_lambda):
-            universe = create_universe_single_state(u_top._topology, ds, state=state_idx)
-            prot = universe.atoms[prot_indices]
-            lig_A = universe.atoms[ligand_A_indices]
-            lig_B = universe.atoms[ligand_B_indices]
-            apply_complex_alignment_transformations(universe, protein=prot, ligands=[lig_A, lig_B])
-
-            if prot:
-                prot_rmsd2d = Protein2DRMSD(prot).run(step=skip)
-                per_state_data["protein_2D_RMSD"].append(prot_rmsd2d.results.rmsd2d)
-
-            for label, lig, rdmol in [
-                ("ligand_A", lig_A, rdmol_A),
-                ("ligand_B", lig_B, rdmol_B),
-            ]:
-                lig_rmsd = SymmetryCorrectedLigandRMSD(lig, rdmol=rdmol).run(step=skip)
-                per_state_data[f"{label}_RMSD"].append(lig_rmsd.results.rmsd)
-
-                lig_drift = LigandCOMDrift(lig).run(step=skip)
-                per_state_data[f"{label}_COM_drift"].append(lig_drift.results.com_drift)
-
-            if time_ps is None:
-                time_ps = np.arange(len(universe.trajectory))[::skip] * universe.trajectory.dt
-        return per_state_data, time_ps
-
-    @staticmethod
-    def _run_solvent_analysis(
-        ds,
-        pdb_file: pathlib.Path,
-        skip: int,
-        ligand_A_indices: list[int],
-        ligand_B_indices: list[int],
-        rdmol_A: Chem.Mol,
-        rdmol_B: Chem.Mol,
-    ) -> tuple[dict[str, list[np.ndarray]], np.ndarray | None]:
-        """
-        Run structural analysis for the solvent phase.
-
-        Parameters
-        ----------
-        ds : netCDF4.Dataset
-          Open NetCDF dataset for the multistate trajectory.
-        pdb_file : pathlib.Path
-          Path to the subsampled PDB file.
-        skip : int
-          Frame stride for analysis.
-        ligand_A_indices : list[int]
-          Atom indices of ligand A in the subsampled system.
-        ligand_B_indices : list[int]
-          Atom indices of ligand B in the subsampled system.
-        rdmol_A : Chem.Mol
-          RDKit molecule for ligand A, used for symmetry-corrected RMSD.
-        rdmol_B : Chem.Mol
-          RDKit molecule for ligand B, used for symmetry-corrected RMSD.
-
-        Returns
-        -------
-        per_state_data : dict[str, list[np.ndarray]]
-          Per-state analysis results for ``ligand_A_RMSD``, ``ligand_B_RMSD``.
-        time_ps : np.ndarray or None
-          Time array in picoseconds corresponding to the analyzed frames.
-        """
-        n_lambda = ds.dimensions["state"].size
-        per_state_data: dict[str, list[np.ndarray]] = {
-            "ligand_A_RMSD": [],
-            "ligand_B_RMSD": [],
-        }
-        time_ps: np.ndarray | None = None
-        # Read the PDB topology once and reuse across all lambda states.
-        # Passing u_top._topology to create_universe_single_state avoids
-        # reading the PDB file from disk for every lambda.
-        u_top = mda.Universe(pdb_file)
-        for state_idx in range(n_lambda):
-            for label, indices, rdmol in [
-                ("ligand_A", ligand_A_indices, rdmol_A),
-                ("ligand_B", ligand_B_indices, rdmol_B),
-            ]:
-                universe = create_universe_single_state(u_top._topology, ds, state=state_idx)
-                lig = universe.atoms[indices]
-                apply_ligand_alignment_transformations(universe, ligand=lig)
-
-                lig_rmsd = SymmetryCorrectedLigandRMSD(lig, rdmol=rdmol).run(step=skip)
-                per_state_data[f"{label}_RMSD"].append(lig_rmsd.results.rmsd)
-
-                if time_ps is None:
-                    time_ps = np.arange(len(universe.trajectory))[::skip] * universe.trajectory.dt
-
-        return per_state_data, time_ps
+        ...
 
-    @staticmethod
     def _structural_analysis(
+        self,
         pdb_file: pathlib.Path,
         trj_file: pathlib.Path,
         output_directory: pathlib.Path,
-        simtype: Literal["complex", "solvent"],
         ligand_A_indices: list[int],
         ligand_B_indices: list[int],
         rdmol_A: Chem.Mol,
@@ -1690,9 +1579,6 @@ def _structural_analysis(
         output_directory : pathlib.Path
             The output directory where plots and the data NPZ file
             will be stored.
-        simtype : Literal["complex", "solvent"]
-            Either ``"complex"`` or ``"solvent"``. Controls whether protein
-            analyses are run and how alignment is applied.
         ligand_A_indices : list[int]
             Atom indices of ligand A in the subsampled system.
         ligand_B_indices : list[int]
@@ -1740,28 +1626,17 @@ def _structural_analysis(
                     # max against 1 to avoid skip=0 case
                     skip = max(n_frames // 500, 1)
 
-                if simtype == "complex":
-                    data, time_ps = BaseSepTopAnalysisUnit._run_complex_analysis(
-                        ds=ds,
-                        pdb_file=pdb_file,
-                        skip=skip,
-                        ligand_A_indices=ligand_A_indices,
-                        ligand_B_indices=ligand_B_indices,
-                        rdmol_A=rdmol_A,
-                        rdmol_B=rdmol_B,
-                        protein_selection=protein_selection,
-                    )
+                data, time_ps = self._run_trajectory_analysis(
+                    ds=ds,
+                    pdb_file=pdb_file,
+                    skip=skip,
+                    ligand_A_indices=ligand_A_indices,
+                    ligand_B_indices=ligand_B_indices,
+                    rdmol_A=rdmol_A,
+                    rdmol_B=rdmol_B,
+                    protein_selection=protein_selection,
+                )
 
-                else:
-                    data, time_ps = BaseSepTopAnalysisUnit._run_solvent_analysis(
-                        ds=ds,
-                        pdb_file=pdb_file,
-                        skip=skip,
-                        ligand_A_indices=ligand_A_indices,
-                        ligand_B_indices=ligand_B_indices,
-                        rdmol_A=rdmol_A,
-                        rdmol_B=rdmol_B,
-                    )
             npz_data = {k: np.asarray(v, dtype=np.float32) for k, v in data.items()}
             npz_data["time_ps"] = np.asarray(time_ps, dtype=np.float32)
 
@@ -1873,7 +1748,6 @@ def run(
             pdb_file=pdb_file,
             trj_file=trajectory,
             output_directory=self.shared_basepath,
-            simtype=self.simtype,
             ligand_A_indices=ligand_A_indices,
             ligand_B_indices=ligand_B_indices,
             rdmol_A=smc_A.to_rdkit(),