Merge pull request #445 from CCPBioSim/441-all-atom-MOI-refactored

Residue level moment of inertia is now calculated using all atom positions + masses to match polymer and UA levels.

Author: ioanaapapa

CodeEntropy/levels/axes.py

@@ -82,15 +82,14 @@ def get_residue_axes(
             * Translational axes are principal axes of data_container.
             * Find edge heavy atom (i.e. heavy atoms bonded to neighbour residue).
               Compute rotation centre and axes as in get_terminal_axes.
-              Compute custom MOI, using heavy atom positions and
-              heavy atom + hydrogen masses.
+              Compute custom MOI, using all atoms positions and masses.
 
         - If bonded to at least two other residues:
             * Translational axes are principal axes of data_container.
             * Find edge heavy atoms (i.e. heavy atoms bonded to neighbour residues).
               Compute rotation centre and axes as in get_non_terminal_axes.
-              Compute a custom MOI, using heavy atom positions and
-              heavy atom + hydrogen masses.
+              Compute custom MOI, using all atoms positions and masses.
+
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -134,19 +133,8 @@ def get_residue_axes(
             f"resindex {index_next})"
         )
 
-        uas = residue.select_atoms("mass 2 to 999")
-        ua_masses = self.get_UA_masses(residue)
-
         if len(edge_atom_set) == 0:
-            moi_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=np.array(residue.center_of_mass()),
-                positions=uas.positions,
-                masses=ua_masses,
-                dimensions=data_container.dimensions[:3],
-            )
-            rot_axes, moment_of_inertia = self.get_principal_axes_from_tensor(
-                moi_tensor
-            )
+            rot_axes, moment_of_inertia = self.get_molecule_axes(residue)
             trans_axes = rot_axes  # per original convention
             rot_center = np.array(residue.center_of_mass())
         else:
@@ -167,8 +155,8 @@ def get_residue_axes(
                 )
             moment_of_inertia = self.get_custom_residue_moment_of_inertia(
                 center_of_mass=rot_center,
-                positions=uas.positions,
-                masses=ua_masses,
+                positions=residue.atoms.positions,
+                masses=residue.atoms.masses,
                 custom_rot_axes=rot_axes,
                 dimensions=data_container.dimensions[:3],
             )
@@ -181,7 +169,6 @@ def get_residue_axes_from_topology(
         mol,
         residue_atoms,
         topology,
-        box: np.ndarray | None,
     ):
         """Compute residue axes using cached static topology.
 
@@ -195,7 +182,6 @@ def get_residue_axes_from_topology(
             mol: Current-frame molecule fragment.
             residue_atoms: AtomGroup for the residue in the current frame.
             topology: Cached ``ResidueAxesTopology`` for this residue.
-            box: Current periodic box lengths. If omitted, ``u.dimensions`` is used.
 
         Returns:
             Tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
@@ -204,24 +190,10 @@ def get_residue_axes_from_topology(
                 - center: Residue centre, shape ``(3,)``.
                 - moment_of_inertia: Principal moments, shape ``(3,)``.
         """
-        dimensions = (
-            np.asarray(box, dtype=float)
-            if box is not None
-            else np.asarray(u.dimensions[:3], dtype=float)
-        )
         center = residue_atoms.center_of_mass(unwrap=True)
 
         if not topology.has_neighbor_bonds:
-            heavy_atoms = u.atoms[topology.residue_heavy_indices]
-            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=center,
-                positions=heavy_atoms.positions,
-                masses=topology.residue_ua_masses,
-                dimensions=dimensions,
-            )
-            rot_axes, moment_of_inertia = self.get_principal_axes_from_tensor(
-                moment_of_inertia_tensor
-            )
+            rot_axes, moment_of_inertia = self.get_molecule_axes(residue_atoms)
             trans_axes = rot_axes
         else:
             make_whole(mol.atoms)
@@ -277,7 +249,7 @@ def get_UA_axes(self, data_container, index: int, res_position):
 
         index = int(index)  # UA bead index
         heavy_atoms = data_container.select_atoms("mass 2 to 999")
-        # use the same customPI trans axes as the residue level
+        # use the same trans axes as the residue level
         if len(heavy_atoms) > 1:
             if len(data_container.residues) == 1:
                 # only the one residue => use principal axes
@@ -411,16 +383,7 @@ def get_UA_axes_from_topology(
         heavy_atom = u.atoms[int(topology.heavy_atom_index)]
 
         if len(heavy_atoms) > 1:
-            center = residue_atoms.center_of_mass(unwrap=True)
-            moment_of_inertia_tensor = self.get_moment_of_inertia_tensor(
-                center_of_mass=center,
-                positions=heavy_atoms.positions,
-                masses=topology.residue_ua_masses,
-                dimensions=dimensions,
-            )
-            trans_axes, _moment_of_inertia = self.get_principal_axes_from_tensor(
-                moment_of_inertia_tensor
-            )
+            trans_axes = self.get_principal_axes_from_group(residue_atoms)
         else:
             make_whole(residue_atoms)
             trans_axes = self.get_principal_axes_from_group(residue_atoms)
@@ -921,8 +884,8 @@ def get_custom_residue_moment_of_inertia(
 
         Args:
             center_of_mass: (3, ) COM for bead
-            positions: (N,3) positions of the UAs in the bead
-            masses: (N,) masses of the UAs in the bead
+            positions: (N,3) positions of the atoms in the bead
+            masses: (N,) masses of the atoms in the bead
             custom_rot_axes: (3,3) array of residue rotation axes
             dimensions: (3,) simulation_box_dimensions
 
@@ -1137,30 +1100,6 @@ def get_principal_axes_from_group(self, group) -> np.ndarray:
         )
         return principal_axes
 
-    def get_UA_masses(self, molecule) -> list[float]:
-        """Return united-atom (UA) masses for a molecule.
-
-        For a given molecule, return a list of masses of UAs (combination of the
-        heavy atoms + bonded hydrogen atoms). This list is used to get the moment
-        of inertia tensor for molecules larger than one UA.
-
-        Args:
-            molecule: MDAnalysis AtomGroup representing the molecule.
-
-        Returns:
-            list[float]: UA masses for each heavy atom.
-        """
-        ua_masses: list[float] = []
-        for atom in molecule:
-            if atom.mass > 1.1:
-                ua_mass = float(atom.mass)
-                bonded_atoms = molecule.select_atoms(f"bonded index {atom.index}")
-                bonded_h_atoms = bonded_atoms.select_atoms("mass 1 to 1.1")
-                for h in bonded_h_atoms:
-                    ua_mass += float(h.mass)
-                ua_masses.append(ua_mass)
-        return ua_masses
-
     def get_chain(self, residue, first, last):
         """
         For a given MDAnalysis AtomGroup and two given heavy atoms

CodeEntropy/levels/nodes/covariance.py

@@ -578,7 +578,6 @@ def _build_residue_vectors(
                 relative_res_i=relative_res_i,
                 axes_manager=axes_manager,
                 axes_topology=axes_topology,
-                box=box,
                 customised_axes=customised_axes,
             )
 
@@ -615,7 +614,6 @@ def _get_residue_axes(
         relative_res_i: int,
         axes_manager: Any,
         axes_topology: Any | None,
-        box: np.ndarray | None,
         customised_axes: bool,
     ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
         """Return axes, centre, and inertia data for a residue bead.
@@ -630,7 +628,6 @@ def _get_residue_axes(
             might differ from residue index of first residue in topology.
             axes_manager: Axes helper used to select axes, centres, and moments.
             axes_topology: Optional cached axes topology generated during static setup.
-            box: Optional periodic box vector.
             customised_axes: Whether customised residue axes should be used.
 
         Returns:
@@ -644,11 +641,7 @@ def _get_residue_axes(
 
             if residue_topology is not None:
                 return axes_manager.get_residue_axes_from_topology(
-                    u=u,
-                    mol=mol,
-                    residue_atoms=res.atoms,
-                    topology=residue_topology,
-                    box=box,
+                    u=u, mol=mol, residue_atoms=res.atoms, topology=residue_topology
                 )
 
             return axes_manager.get_residue_axes(

tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json

@@ -3,14 +3,14 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 0.06817494136874551,
-        "united_atom:Rovibrational": 49.686697381523516,
-        "residue:Transvibrational": 69.25273989458984,
-        "residue:Rovibrational": 68.6481013342223,
+        "united_atom:Rovibrational": 49.686697381523466,
+        "residue:Transvibrational": 69.24977659775257,
+        "residue:Rovibrational": 69.73922011617867,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 208.13728504431975
+      "total": 209.22544052943883
     }
   }
 }

tests/regression/baselines/benzaldehyde/default.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 24.92182210331405,
-        "united_atom:Rovibrational": 27.950376507672583,
-        "residue:FTmat-Transvibrational": 70.97090266225698,
-        "residue:FTmat-Rovibrational": 59.588669973122094,
+        "united_atom:Transvibrational": 24.921822103314035,
+        "united_atom:Rovibrational": 27.950376507672594,
+        "residue:FTmat-Transvibrational": 71.231562004914,
+        "residue:FTmat-Rovibrational": 60.43602521189648,
         "united_atom:Conformational": 7.522643899702263,
         "residue:Conformational": 0.0,
         "residue:Orientational": 59.43920971558428
       },
-      "total": 250.39362486165226
+      "total": 251.50163944308366
     }
   }
 }

tests/regression/baselines/benzaldehyde/frame_window.json

@@ -3,14 +3,14 @@
     "0": {
       "components": {
         "united_atom:Transvibrational": 37.90368186636453,
-        "united_atom:Rovibrational": 38.21906858443616,
-        "residue:FTmat-Transvibrational": 73.76730786578054,
-        "residue:FTmat-Rovibrational": 57.689489298548764,
+        "united_atom:Rovibrational": 38.219068584436165,
+        "residue:FTmat-Transvibrational": 74.05827280885852,
+        "residue:FTmat-Rovibrational": 58.53451127009362,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 228.06111910774533
+      "total": 229.1971060223682
     }
   }
 }