LinearDensity now working with residues/segments (MDAnalysis#3572)

Fixes MDAnalysis#3571 : LinearDensity does not work on residues or segments

MDAnalysis/analysis/lineardensity.py:
  - in _prepare(self):
    - try/except statement replaced by if statements checking the grouping
      specified for LinearDensity.
    - masses and charges are obtained in a way appropriate for each case.

    - Floating point arithmetic issues could lead to negative radicands. A check
      was implemented to set numbers which are close to zero to precisely zero
      before the square root (for standard deviation) is calculated.

  - renamed previous test function and added expected masses and charges
  - added test functions for the other groupings with their respective expected
    masses and charges

Changed the if/if/if for reading masses and charges
to if/elif/else.
Will now raise an AttributeError if invalid grouping is chosen.
'residues', 'segments' and 'fragments' now use the same logic
(elem.atoms.total_mass()).

* Changed test of radicand to only convert negative numbers to 0
* Now uses CoM rather than CoG for density calculations
* Changed Class docstring to reflect change to CoM
* Now uses parameter for compound instead of list comprehension when
  calculating masses, charges, and CoM
* refactored test_lineardensity.py, make use of pytest functionality
  better
* Added test for case of invalid grouping selection


Co-authored-by: Irfan Alibay <IAlibay@users.noreply.github.com>

package/CHANGELOG

@@ -19,6 +19,8 @@ The rules for this file:
  * 2.2.0
 
 Fixes
+  * Fixed LinearDensity not working with grouping='residues' or 'segments'
+    (Issue #3571, PR #3572)
   * MOL2 can read files without providing optional columns:
     subst_id, subst_name and charge. (Issue #3385, PR #3598)
   * When converting an AtomGroup to an RDKit molecule (PR #3044):

package/MDAnalysis/analysis/lineardensity.py

@@ -43,8 +43,9 @@ class LinearDensity(AnalysisBase):
     select : AtomGroup
           any atomgroup
     grouping : str {'atoms', 'residues', 'segments', 'fragments'}
-          Density profiles will be computed on the center of geometry
-          of a selected group of atoms
+          Density profiles will be computed either on the atom positions (in
+          the case of 'atoms') or on the center of mass of the specified
+          grouping unit ('residues', 'segments', or 'fragments').
     binsize : float
           Bin width in Angstrom used to build linear density
           histograms. Defines the resolution of the resulting density
@@ -80,6 +81,15 @@ class LinearDensity(AnalysisBase):
        print(ldens.results.x.pos)
 
 
+    Alternatively, the other types of groupings can be selected using the ``grouping``
+    keyword. For example to calculated the density based on the :class:`ResidueGroup`s
+    of the system:
+    
+    .. code-block:: python
+       ldens = LinearDensity(selection, grouping='residues', binsize=1.0)
+       ldens.run()
+
+
     .. versionadded:: 0.14.0
 
     .. versionchanged:: 1.0.0
@@ -96,6 +106,12 @@ class LinearDensity(AnalysisBase):
        Results are now instances of
        :class:`~MDAnalysis.core.analysis.Results` allowing access
        via key and attribute.
+
+    .. versionchanged:: 2.2.0
+       Fixed a bug that caused LinearDensity to fail if grouping="residues"
+       or grouping="segments" were set.
+       Residues, segments, and fragments will be analysed based on their centre
+       of mass, not centre of geometry as previously stated.
     """
 
     def __init__(self, select, grouping='atoms', binsize=0.25, **kwargs):
@@ -141,18 +157,19 @@ def __init__(self, select, grouping='atoms', binsize=0.25, **kwargs):
         self.totalmass = None
 
     def _prepare(self):
-        # group must be a local variable, otherwise there will be
-        # issues with parallelization
-        group = getattr(self._ags[0], self.grouping)
-
         # Get masses and charges for the selection
-        try:  # in case it's not an atom
-            self.masses = np.array([elem.total_mass() for elem in group])
-            self.charges = np.array([elem.total_charge() for elem in group])
-        except AttributeError:  # much much faster for atoms
+        if self.grouping == "atoms":
             self.masses = self._ags[0].masses
             self.charges = self._ags[0].charges
 
+        elif self.grouping in ["residues", "segments", "fragments"]:
+            self.masses = self._ags[0].total_mass(compound=self.grouping)
+            self.charges = self._ags[0].total_charge(compound=self.grouping)
+
+        else:
+            raise AttributeError(
+                f"{self.grouping} is not a valid value for grouping.")
+
         self.totalmass = np.sum(self.masses)
 
     def _single_frame(self):
@@ -163,8 +180,8 @@ def _single_frame(self):
         if self.grouping == 'atoms':
             positions = self._ags[0].positions  # faster for atoms
         else:
-            # COM for res/frag/etc
-            positions = np.array([elem.centroid() for elem in self.group])
+            # Centre of mass for residues, segments, fragments
+            positions = self._ags[0].center_of_mass(compound=self.grouping)
 
         for dim in ['x', 'y', 'z']:
             idx = self.results[dim]['dim']
@@ -199,10 +216,20 @@ def _conclude(self):
             for key in ['pos', 'pos_std', 'char', 'char_std']:
                 self.results[dim][key] /= self.n_frames
             # Compute standard deviation for the error
-            self.results[dim]['pos_std'] = np.sqrt(self.results[dim][
-                'pos_std'] - np.square(self.results[dim]['pos']))
-            self.results[dim]['char_std'] = np.sqrt(self.results[dim][
-                'char_std'] - np.square(self.results[dim]['char']))
+            # For certain tests in testsuite, floating point imprecision
+            # can lead to negative radicands of tiny magnitude (yielding nan).
+            # radicand_pos and radicand_char are therefore calculated first and
+            # and negative values set to 0 before the square root
+            # is calculated.
+            radicand_pos = self.results[dim][
+                'pos_std'] - np.square(self.results[dim]['pos'])
+            radicand_pos[radicand_pos < 0] = 0
+            self.results[dim]['pos_std'] = np.sqrt(radicand_pos)
+
+            radicand_char = self.results[dim][
+                'char_std'] - np.square(self.results[dim]['char'])
+            radicand_char[radicand_char < 0] = 0
+            self.results[dim]['char_std'] = np.sqrt(radicand_char)
 
         for dim in ['x', 'y', 'z']:
             norm = k * self.results[dim]['slice_volume']

testsuite/MDAnalysisTests/analysis/test_lineardensity.py

@@ -22,18 +22,78 @@
 #
 import MDAnalysis as mda
 import numpy as np
+import pytest
 
 from MDAnalysisTests.datafiles import waterPSF, waterDCD
 from MDAnalysis.analysis.lineardensity import LinearDensity
-from numpy.testing import assert_almost_equal
+from numpy.testing import assert_allclose
 
 
-def test_serial():
+def test_invalid_grouping():
+    """Invalid groupings raise AttributeError"""
     universe = mda.Universe(waterPSF, waterDCD)
     sel_string = 'all'
     selection = universe.select_atoms(sel_string)
+    with pytest.raises(AttributeError):
+        # centroid is attribute of AtomGroup, but not valid here
+        ld = LinearDensity(selection, grouping="centroid", binsize=5)
+        ld.run()
 
-    xpos = np.array([0., 0., 0., 0.0072334, 0.00473299, 0.,
-                          0., 0., 0., 0.])
-    ld = LinearDensity(selection, binsize=5).run()
-    assert_almost_equal(xpos, ld.results['x']['pos'])
+
+# test data for grouping='atoms'
+expected_masses_atoms = np.array([15.9994, 1.008, 1.008, 15.9994, 1.008, 1.008,
+                                  15.9994, 1.008, 1.008, 15.9994, 1.008, 1.008,
+                                  15.9994, 1.008, 1.008])
+expected_charges_atoms = np.array([-0.834, 0.417, 0.417, -0.834, 0.417,
+                                   0.417, -0.834, 0.417, 0.417, -0.834,
+                                   0.417, 0.417, -0.834, 0.417, 0.417])
+expected_xpos_atoms = np.array([0., 0., 0., 0.0072334, 0.00473299, 0.,
+                                0., 0., 0., 0.])
+expected_xchar_atoms = np.array([0., 0., 0., 2.2158751e-05, -2.2158751e-05,
+                                 0., 0., 0., 0., 0.])
+
+# test data for grouping='residues'
+expected_masses_residues = np.array([18.0154, 18.0154, 18.0154, 18.0154,
+                                     18.0154])
+expected_charges_residues = np.array([0, 0, 0, 0, 0])
+expected_xpos_residues = np.array([0., 0., 0., 0.00717983, 0.00478656,
+                                   0., 0., 0., 0., 0.])
+expected_xchar_residues = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
+
+# test data for grouping='segments'
+expected_masses_segments = np.array([90.0770])
+expected_charges_segments = np.array([0])
+expected_xpos_segments = np.array([0., 0., 0., 0.01196639, 0.,
+                                   0., 0., 0., 0., 0.])
+expected_xchar_segments = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
+
+# test data for grouping='fragments'
+expected_masses_fragments = np.array([18.0154, 18.0154, 18.0154, 18.0154,
+                                      18.0154])
+expected_charges_fragments = np.array([0, 0, 0, 0, 0])
+expected_xpos_fragments = np.array([0., 0., 0., 0.00717983, 0.00478656,
+                                   0., 0., 0., 0., 0.])
+expected_xchar_fragments = np.array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
+
+
+@pytest.mark.parametrize("grouping, expected_masses, expected_charges, expected_xpos, expected_xchar", [
+    ("atoms", expected_masses_atoms, expected_charges_atoms,
+     expected_xpos_atoms, expected_xchar_atoms),
+    ("residues", expected_masses_residues, expected_charges_residues,
+     expected_xpos_residues, expected_xchar_residues),
+    ("segments", expected_masses_segments, expected_charges_segments,
+     expected_xpos_segments, expected_xchar_segments),
+    ("fragments", expected_masses_fragments, expected_charges_fragments,
+     expected_xpos_fragments, expected_xchar_fragments)
+])
+def test_lineardensity(grouping, expected_masses, expected_charges,
+                       expected_xpos, expected_xchar):
+    universe = mda.Universe(waterPSF, waterDCD)
+    sel_string = 'all'
+    selection = universe.select_atoms(sel_string)
+    ld = LinearDensity(selection, grouping, binsize=5).run()
+    assert_allclose(ld.masses, expected_masses)
+    assert_allclose(ld.charges, expected_charges)
+    # rtol changed here due to floating point imprecision
+    assert_allclose(ld.results['x']['pos'], expected_xpos, rtol=1e-06)
+    assert_allclose(ld.results['x']['char'], expected_xchar)