gyration_moments method / allow shape parameters to yield per residue quantities

package/CHANGELOG

@@ -15,10 +15,13 @@ The rules for this file:
 ------------------------------------------------------------------------------
 
 ??/??/?? IAlibay, pgbarletta, mglagolev, hmacdope, manuel.nuno.melo, chrispfae, 
-         ooprathamm, MeetB7, BFedder, v-parmar, MoSchaeffler, jbarnoud, jandom, xhgchen
+         ooprathamm, MeetB7, BFedder, v-parmar, MoSchaeffler, jbarnoud, jandom,
+         xhgchen, jaclark5
  * 2.5.0
 
 Fixes
+  * Allows shape_parameter and asphericity to yield per residue quantities
+    (Issue #3002, PR #3905)
   * Add tests for "No path data" exception raise in test_psa.py (Issue #4036)
   * Fix uninitialized `format` variable issue when calling `selections.get_writer` directly (PR #4043)
   * Fix tests should use results.rmsf to avoid DeprecationWarning (Issue #3695)
@@ -30,6 +33,8 @@ Fixes
     (Issue #3336)
 
 Enhancements
+  * Added AtomGroup TopologyAttr to calculate gyration moments (Issue #3904,
+    PR #3905)
   * Add support for TPR files produced by Gromacs 2023 (Issue #4047)
   * Add distopia distance calculation library bindings as a selectable backend
     for `calc_bonds` in `MDA.lib.distances`. (Issue #3783, PR #3914)

package/MDAnalysis/core/topologyattrs.py

@@ -1729,7 +1729,100 @@ def radius_of_gyration(group, wrap=False, **kwargs):
     @warn_if_not_unique
     @_pbc_to_wrap
     @check_atomgroup_not_empty
-    def shape_parameter(group, wrap=False):
+    def gyration_moments(group, wrap=False, unwrap=False, compound='group'):
+        r"""Moments of the gyration tensor.
+
+        The moments are defined as the eigenvalues of the gyration
+        tensor.
+
+        .. math::
+
+            \mathsf{T} = \frac{1}{N} \sum_{i=1}^{N} (\mathbf{r}_\mathrm{i} - 
+                \mathbf{r}_\mathrm{COM})(\mathbf{r}_\mathrm{i} - \mathbf{r}_\mathrm{COM})
+
+        Where :math:`\mathbf{r}_\mathrm{COM}` is the center of mass.
+
+        See [Dima2004a]_ for background information.
+
+        Parameters
+        ----------
+        wrap : bool, optional
+            If ``True``, move all atoms within the primary unit cell before
+            calculation. [``False``]
+        unwrap : bool, optional
+            If ``True``, compounds will be unwrapped before computing their centers.
+        compound : {'group', 'segments', 'residues', 'molecules', 'fragments'}, optional
+            Which type of component to keep together during unwrapping.
+
+        Returns
+        -------
+        principle_moments_of_gyration : numpy.ndarray
+            Gyration vector(s) of (compounds of) the group in :math:`Å^2`.
+            If `compound` was set to ``'group'``, the output will be a single
+            vector of length 3. Otherwise, the output will be a 2D array of shape
+            ``(n,3)`` where ``n`` is the number of compounds.
+
+
+        .. versionadded:: 2.5.0
+        """
+
+        def _gyration(recenteredpos, masses):
+            if len(masses.shape) > 1:
+                masses = np.squeeze(masses)
+            tensor = np.einsum( "ki,kj->ij",
+                                recenteredpos,
+                                np.einsum("ij,i->ij", recenteredpos, masses),
+                              )
+            return np.linalg.eigvalsh(tensor/np.sum(masses))
+
+        atomgroup = group.atoms
+        masses = atomgroup.masses
+
+        com = atomgroup.center_of_mass(
+            wrap=wrap, unwrap=unwrap, compound=compound)
+
+        if compound == 'group':
+             if wrap:
+                 recenteredpos = (atomgroup.pack_into_box(inplace=False) - com)
+             elif unwrap:
+                 recenteredpos = (atomgroup.unwrap(inplace=False,
+                                                   compound=compound, 
+                                                   reference=None, 
+                                                  ) - com)
+             else:
+                 recenteredpos = (atomgroup.positions - com)
+             eig_vals = _gyration(recenteredpos, masses)
+        else:
+             (atom_masks, 
+              compound_masks, 
+              n_compounds) = atomgroup._split_by_compound_indices(compound)
+
+             if unwrap:
+                 coords = atomgroup.unwrap(
+                     compound=compound, 
+                     reference=None, 
+                     inplace=False
+                 )
+             else:
+                 coords = atomgroup.positions
+
+             eig_vals = np.empty((n_compounds, 3), dtype=np.float64)
+             for compound_mask, atom_mask in zip(compound_masks, atom_masks):
+                 eig_vals[compound_mask, :] = [_gyration(
+                      coords[mask] - com[compound_mask][i],
+                      masses[mask][:, None]
+                     ) for i, mask in enumerate(atom_mask)]
+
+        return eig_vals
+
+    transplants[GroupBase].append(
+        ('gyration_moments', gyration_moments))
+
+
+    @warn_if_not_unique
+    @_pbc_to_wrap
+    @check_atomgroup_not_empty
+    def shape_parameter(group, wrap=False, unwrap=False, compound='group'):
         """Shape parameter.
 
         See [Dima2004a]_ for background information.
@@ -1739,6 +1832,10 @@ def shape_parameter(group, wrap=False):
         wrap : bool, optional
             If ``True``, move all atoms within the primary unit cell before
             calculation. [``False``]
+        unwrap : bool, optional
+            If ``True``, compounds will be unwrapped before computing their centers.
+        compound : {'group', 'segments', 'residues', 'molecules', 'fragments'}, optional
+            Which type of component to keep together during unwrapping.
 
 
         .. versionadded:: 0.7.7
@@ -1748,24 +1845,17 @@ def shape_parameter(group, wrap=False):
            Renamed `pbc` kwarg to `wrap`. `pbc` is still accepted but
            is deprecated and will be removed in version 3.0.
            Superfluous kwargs were removed.
+        .. versionchanged:: 2.5.0
+           Added calculation for any `compound` type
         """
         atomgroup = group.atoms
-        masses = atomgroup.masses
-
-        com = atomgroup.center_of_mass(wrap=wrap)
-        if wrap:
-            recenteredpos = atomgroup.pack_into_box(inplace=False) - com
+        eig_vals = atomgroup.gyration_moments(wrap=wrap, unwrap=unwrap, compound=compound)
+        if len(eig_vals.shape) > 1:
+            shape = 27.0 * np.prod(eig_vals - np.mean(eig_vals, axis=1), axis=1
+                                   ) / np.power(np.sum(eig_vals, axis=1), 3)
         else:
-            recenteredpos = atomgroup.positions - com
-        tensor = np.zeros((3, 3))
-
-        for x in range(recenteredpos.shape[0]):
-            tensor += masses[x] * np.outer(recenteredpos[x, :],
-                                           recenteredpos[x, :])
-        tensor /= atomgroup.total_mass()
-        eig_vals = np.linalg.eigvalsh(tensor)
-        shape = 27.0 * np.prod(eig_vals - np.mean(eig_vals)
-                               ) / np.power(np.sum(eig_vals), 3)
+            shape = 27.0 * np.prod(eig_vals - np.mean(eig_vals)
+                                   ) / np.power(np.sum(eig_vals), 3)
 
         return shape
 
@@ -1776,7 +1866,7 @@ def shape_parameter(group, wrap=False):
     @_pbc_to_wrap
     @check_wrap_and_unwrap
     @check_atomgroup_not_empty
-    def asphericity(group, wrap=False, unwrap=None, compound='group'):
+    def asphericity(group, wrap=False, unwrap=False, compound='group'):
         """Asphericity.
 
         See [Dima2004b]_ for background information.
@@ -1800,32 +1890,17 @@ def asphericity(group, wrap=False, unwrap=None, compound='group'):
         .. versionchanged:: 2.1.0
            Renamed `pbc` kwarg to `wrap`. `pbc` is still accepted but
            is deprecated and will be removed in version 3.0.
+        .. versionchanged:: 2.5.0
+           Added calculation for any `compound` type
         """
         atomgroup = group.atoms
-        masses = atomgroup.masses
-
-        com = atomgroup.center_of_mass(
-            wrap=wrap, unwrap=unwrap, compound=compound)
-        if compound != 'group':
-            com = (com * group.masses[:, None]
-                   ).sum(axis=0) / group.masses.sum()
-
-        if wrap:
-            recenteredpos = (atomgroup.pack_into_box(inplace=False) - com)
-        elif unwrap:
-            recenteredpos = (atomgroup.unwrap(inplace=False) - com)
+        eig_vals = atomgroup.gyration_moments(wrap=wrap, unwrap=unwrap, compound=compound)
+        if len(eig_vals.shape) > 1:
+            shape = (3.0 / 2.0) * (np.sum((eig_vals - np.mean(eig_vals, axis=1))**2, axis=1) /
+                                   np.sum(eig_vals, axis=1)**2)
         else:
-            recenteredpos = (atomgroup.positions - com)
-
-        tensor = np.zeros((3, 3))
-        for x in range(recenteredpos.shape[0]):
-            tensor += masses[x] * np.outer(recenteredpos[x],
-                                           recenteredpos[x])
-
-        tensor /= atomgroup.total_mass()
-        eig_vals = np.linalg.eigvalsh(tensor)
-        shape = (3.0 / 2.0) * (np.sum((eig_vals - np.mean(eig_vals))**2) /
-                               np.sum(eig_vals)**2)
+            shape = (3.0 / 2.0) * (np.sum((eig_vals - np.mean(eig_vals))**2) /
+                                   np.sum(eig_vals)**2)
 
         return shape
 

testsuite/MDAnalysisTests/core/test_atomgroup.py

@@ -994,7 +994,8 @@ class TestUnwrapFlag(object):
             np.array([[7333.79167791, -211.8997285, -721.50785456],
                       [-211.8997285, 7059.07470427, -91.32156884],
                       [-721.50785456, -91.32156884, 6509.31735029]]),
-        'asphericity': 0.02060121,
+        'asphericity': np.array([0.135, 0.047, 0.094]),
+        'shape_parameter': np.array([-0.112, -0.004,  0.02]),
     }
 
     ref_Unwrap_residues = {
@@ -1009,7 +1010,8 @@ class TestUnwrapFlag(object):
         'moment_of_inertia': np.array([[16687.941, -1330.617, 2925.883],
                                        [-1330.617, 19256.178, 3354.832],
                                        [2925.883,  3354.832, 8989.946]]),
-        'asphericity': 0.2969491080,
+        'asphericity': np.array([0.61 , 0.701, 0.381]),
+        'shape_parameter': np.array([-0.461,  0.35 ,  0.311]),
     }
 
     ref_noUnwrap = {
@@ -1019,6 +1021,7 @@ class TestUnwrapFlag(object):
                                        [0.0, 98.6542, 0.0],
                                        [0.0, 0.0, 98.65421327]]),
         'asphericity': 1.0,
+        'shape_parameter': 1.0,
     }
 
     ref_Unwrap = {
@@ -1028,6 +1031,7 @@ class TestUnwrapFlag(object):
                                        [0.0, 132.673, 0.0],
                                        [0.0, 0.0, 132.673]]),
         'asphericity': 1.0,
+        'shape_parameter': 1.0,
     }
 
     @pytest.fixture(params=[False, True])  # params indicate shuffling
@@ -1054,7 +1058,8 @@ def unwrap_group(self):
     @pytest.mark.parametrize('method_name', ('center_of_geometry',
                                              'center_of_mass',
                                              'moment_of_inertia',
-                                             'asphericity'))
+                                             'asphericity',
+                                             'shape_parameter'))
     def test_residues(self, ag, unwrap, ref, method_name):
         method = getattr(ag, method_name)
         if unwrap: