Change residue level MOI calculation to all atom

CodeEntropy/levels/axes.py

@@ -82,15 +82,13 @@ 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 atom 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 atom positions and masses.
 
         Args:
             data_container (MDAnalysis.Universe or AtomGroup):
@@ -134,17 +132,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_custom_principal_axes(moi_tensor)
+            rot_axes, moi = self.get_vanilla_axes(residue)
             trans_axes = rot_axes  # per original convention
             rot_center = np.array(residue.center_of_mass())
         else:
@@ -165,8 +154,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],
             )
@@ -202,24 +191,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_custom_principal_axes(
-                moment_of_inertia_tensor
-            )
+            rot_axes, moment_of_inertia = self.get_vanilla_axes(residue_atoms)
             trans_axes = rot_axes
         else:
             make_whole(mol.atoms)
@@ -275,7 +250,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
@@ -925,8 +900,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
 

tests/regression/baselines/benzaldehyde/axes_off.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.08982962903796131,
-        "united_atom:Rovibrational": 32.16018134884085,
-        "residue:FTmat-Transvibrational": 88.7671666695003,
+        "united_atom:Transvibrational": 0.08982962903796153,
+        "united_atom:Rovibrational": 32.160181348840844,
+        "residue:FTmat-Transvibrational": 88.76716666950027,
         "residue:FTmat-Rovibrational": 61.61036267672132,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 203.1091118167158
+      "total": 203.10911181671574
     }
   }
 }

tests/regression/baselines/benzaldehyde/combined_forcetorque_false.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.07119323721997475,
-        "united_atom:Rovibrational": 49.68669738152346,
-        "residue:Transvibrational": 69.48692941204929,
-        "residue:Rovibrational": 68.46147102540942,
+        "united_atom:Transvibrational": 0.07119323721997481,
+        "united_atom:Rovibrational": 49.686697381523466,
+        "residue:Transvibrational": 69.80330954121192,
+        "residue:Rovibrational": 73.59711177182774,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 208.1878625488175
+      "total": 213.63988342439845
     }
   }
 }

tests/regression/baselines/benzaldehyde/frame_window.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 40.30267601961045,
-        "united_atom:Rovibrational": 38.21906858443615,
-        "residue:FTmat-Transvibrational": 73.41098578352612,
-        "residue:FTmat-Rovibrational": 57.881504393660364,
+        "united_atom:Transvibrational": 40.30267601961048,
+        "united_atom:Rovibrational": 38.219068584436165,
+        "residue:FTmat-Transvibrational": 74.59522673101434,
+        "residue:FTmat-Rovibrational": 58.48870270355095,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 230.29580627384846
+      "total": 232.0872455312273
     }
   }
 }

tests/regression/baselines/benzaldehyde/grouping_each.json

@@ -4,121 +4,121 @@
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.062499436482478535,
-        "residue:FTmat-Transvibrational": 16.83857036769833,
+        "residue:FTmat-Transvibrational": 16.92340904741005,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 31.909265242174982
       },
-      "total": 48.81033504635579
+      "total": 48.89517372606751
     },
     "1": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.04214874235319735,
-        "residue:FTmat-Transvibrational": 9.439228173427328,
+        "residue:FTmat-Transvibrational": 9.50856420599859,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 9.481376915780524
+      "total": 9.550712948351787
     },
     "2": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.0022147884278620493,
-        "residue:FTmat-Transvibrational": 17.449372254125855,
+        "residue:FTmat-Transvibrational": 17.765096867249817,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 31.909265242174982
       },
-      "total": 49.3608522847287
+      "total": 49.67657689785266
     },
     "3": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.001983835543157824,
-        "residue:FTmat-Transvibrational": 15.803976930845373,
+        "residue:FTmat-Transvibrational": 15.979343677823543,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 15.80596076638853
+      "total": 15.9813275133667
     },
     "4": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.003874883925216773,
-        "residue:FTmat-Transvibrational": 23.899916844406114,
+        "residue:FTmat-Transvibrational": 24.03232930250616,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 23.90379172833133
+      "total": 24.036204186431377
     },
     "5": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.010824819910336868,
-        "residue:FTmat-Transvibrational": 11.05479779559611,
+        "residue:FTmat-Transvibrational": 11.176470961462543,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 31.909265242174982
       },
-      "total": 42.97488785768143
+      "total": 43.096561023547864
     },
     "6": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.00018982532508837537,
-        "residue:FTmat-Transvibrational": 19.723925127490933,
+        "residue:FTmat-Transvibrational": 19.838516143816783,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 19.72411495281602
+      "total": 19.83870596914187
     },
     "7": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.004033941970593011,
-        "residue:FTmat-Transvibrational": 18.90676882142135,
+        "residue:FTmat-Transvibrational": 19.06667852893644,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 18.910802763391942
+      "total": 19.070712470907033
     },
     "8": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.020362757918507036,
-        "residue:FTmat-Transvibrational": 16.13204209116967,
+        "residue:FTmat-Transvibrational": 16.57161042803149,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 31.909265242174982
       },
-      "total": 48.06167009126315
+      "total": 48.50123842812498
     },
     "9": {
       "components": {
         "united_atom:Transvibrational": 0.0,
         "united_atom:Rovibrational": 0.00023139743934636996,
-        "residue:FTmat-Transvibrational": 21.42548529236268,
+        "residue:FTmat-Transvibrational": 21.777605911862764,
         "residue:FTmat-Rovibrational": 0.0,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 21.425716689802027
+      "total": 21.77783730930211
     }
   }
 }

tests/regression/baselines/benzaldehyde/selection_subset.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.07119323721997475,
-        "united_atom:Rovibrational": 49.68669738152346,
-        "residue:FTmat-Transvibrational": 87.43527331108173,
-        "residue:FTmat-Rovibrational": 61.67126452972779,
+        "united_atom:Transvibrational": 0.07119323721997481,
+        "united_atom:Rovibrational": 49.686697381523466,
+        "residue:FTmat-Transvibrational": 88.76716666950027,
+        "residue:FTmat-Rovibrational": 61.61036267672132,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 20.481571492615355
       },
-      "total": 219.34599995216834
+      "total": 220.61699145758038
     }
   }
 }

tests/regression/baselines/benzene/axes_off.json

@@ -2,9 +2,9 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.17824017761231503,
-        "united_atom:Rovibrational": 37.12706256545503,
-        "residue:FTmat-Transvibrational": 76.8200637641939,
+        "united_atom:Transvibrational": 0.17824017761231412,
+        "united_atom:Rovibrational": 37.12706256545502,
+        "residue:FTmat-Transvibrational": 76.82006376419392,
         "residue:FTmat-Rovibrational": 47.171663134129304,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,

tests/regression/baselines/benzene/combined_forcetorque_off.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 0.16311383522946216,
-        "united_atom:Rovibrational": 40.695812407899396,
-        "residue:Transvibrational": 57.59675197041957,
-        "residue:Rovibrational": 52.85496348419672,
+        "united_atom:Transvibrational": 0.16311383522946185,
+        "united_atom:Rovibrational": 40.69581240789943,
+        "residue:Transvibrational": 58.30403178667887,
+        "residue:Rovibrational": 53.84947651174329,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 151.31064169774515
+      "total": 153.01243454155104
     }
   }
 }

tests/regression/baselines/benzene/frame_window.json

@@ -2,15 +2,15 @@
   "groups": {
     "0": {
       "components": {
-        "united_atom:Transvibrational": 11.46552036084548,
-        "united_atom:Rovibrational": 34.98492056748493,
-        "residue:FTmat-Transvibrational": 71.83491878606151,
-        "residue:FTmat-Rovibrational": 55.6098400281744,
+        "united_atom:Transvibrational": 11.465520360845467,
+        "united_atom:Rovibrational": 34.98492056748487,
+        "residue:FTmat-Transvibrational": 74.31826351535895,
+        "residue:FTmat-Rovibrational": 55.98936607558771,
         "united_atom:Conformational": 0.0,
         "residue:Conformational": 0.0,
         "residue:Orientational": 0.0
       },
-      "total": 173.89519974256632
+      "total": 176.758070519277
     }
   }
 }