Add dipole and quadrupole caclulation in AtomGroup

package/CHANGELOG

@@ -13,7 +13,8 @@ The rules for this file:
   * release numbers follow "Semantic Versioning" http://semver.org
 
 ------------------------------------------------------------------------------
-??/??/?? IAlibay, Luthaf, hmacdope, rafaelpap, jbarnoud, BFedder, aya9aladdin
+??/??/?? IAlibay, Luthaf, hmacdope, rafaelpap, jbarnoud, BFedder, aya9aladdin,
+         jaclark5
 
  * 2.4.0
 
@@ -25,6 +26,7 @@ Fixes
   * fixes guessed_attributes and read_attributes methods of the Topology class, as 
     those methods were giving unexpected results for some topology attributes
     (e.g. bonds, angles) (PR #3779).
+  * Added dipole and quadrupole moment topology methods to AtomGroup (#3841)
 
 
 

package/MDAnalysis/core/topologyattrs.py

@@ -202,6 +202,17 @@ def stub_method(self, *args, **kwargs):
         )
         raise NoDataError(message)
 
+    def stub_method(self, *args, **kwargs):
+        message = (
+            '{class_name}.{method_name}() '
+            'not available; this requires {attribute_name}'
+        ).format(
+            class_name=self.__class__.__name__,
+            method_name=method_name,
+            attribute_name=attribute_name,
+        )
+        raise NoDataError(message)
+
     annotation = textwrap.dedent("""\
         .. note::
 
@@ -1988,6 +1999,8 @@ def center_of_charge(group, wrap=False, unwrap=False, compound='group'):
         .. versionadded:: 2.2.0
         """
         atoms = group.atoms
+        if not len(atoms):
+            raise ValueError("There are no atoms provided group")
         return atoms.center(weights=atoms.charges.__abs__(),
                             wrap=wrap,
                             compound=compound,
@@ -2032,6 +2045,239 @@ def total_charge(group, compound='group'):
     transplants[GroupBase].append(
         ('total_charge', total_charge))
 
+    @warn_if_not_unique
+    @_pbc_to_wrap
+    @check_wrap_and_unwrap
+    def dipole_moment(group, wrap=False, unwrap=False, 
+        compound='group', center="mass"):
+        r"""Dipole vector of the group.
+
+        Computes the dipole moment of :class:`Atoms<Atom>` in the group.
+        Dipoleper :class:`Residue`, :class:`Segment`, molecule, or
+        fragment can be obtained by setting the `compound` parameter
+        accordingly.
+
+        Parameters
+        ----------
+        group : obj
+            AtomGroup
+        wrap : bool, optional
+            If ``True`` and `compound` is ``'group'``, move all atoms to the
+            primary unit cell before calculation.
+            If ``True`` and `compound` is not ``group``, the
+            centers of mass of each compound will be calculated without moving
+            any atoms to keep the compounds intact.
+        unwrap : bool, optional
+            If ``True``, compounds will be unwrapped before computing their
+            centers.
+        compound : {'group', 'segments', 'residues', 'molecules', \
+                    'fragments'}, optional
+            If ``'group'``, a single dipole vector returns. Otherwise, an
+            array of each :class:`Segment`, :class:`Residue`, molecule, or
+            fragment will be returned as an array of position vectors, i.e.
+            a 2d array.
+            Note that, in any case, *only* the positions of
+            :class:`Atoms<Atom>` *belonging to the group* will be taken into
+            account.
+        center : str, optional
+            Choose whether the dipole moment is calculated at the center of 
+            "mass" or the center of "charge"
+
+        Returns
+        -------
+        numpy.ndarray
+            Dipole moment of (compounds of) the group in e*angstroms.
+            If `compound` was set to ``'group'``, the output will be a single
+            value. Otherwise, the output will be a 1d array of shape ``(n,)``
+            where ``n`` is the number of compounds.
+
+        .. versionadded:: 2.4.0
+        """
+
+        compound = compound.lower()
+
+        atomgroup = group.atoms
+        charges = atomgroup.charges
+
+        if center == "mass":
+            masses = atomgroup.masses
+            ref = atomgroup.center_of_mass(
+                wrap=wrap, 
+                unwrap=unwrap, 
+                compound=compound
+            )
+        elif center == "charge":
+            ref = atomgroup.center_of_charge(
+                wrap=wrap,
+                unwrap=unwrap,
+                compound=compound
+            )
+        else:
+            choices = ["mass", "charge"]
+            raise ValueError(
+                "The dipole center, {}, is not supported. Choose: {}".format(
+                 center,choices))
+
+        if compound == 'group':
+            if wrap:
+                recenteredpos = (atomgroup.pack_into_box(inplace=False) - ref)
+            elif unwrap:
+                recenteredpos = (atomgroup.unwrap(inplace=False) - ref)
+            else:
+                recenteredpos = (atomgroup.positions - ref)
+            dipole_moment = np.sum(recenteredpos*charges[:,np.newaxis], axis=0)
+        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
+            chgs = atomgroup.charges
+
+            dipole_moment = np.empty((n_compounds, 3), dtype=np.float64)
+            for compound_mask, atom_mask in zip(compound_masks, atom_masks):
+                dipole_moment[compound_mask] = np.sum(
+                    (coords[atom_mask]-ref[compound_mask][:,None,:])
+                    *chgs[atom_mask][:,:,None], axis=1)
+
+        return dipole_moment
+
+    transplants[GroupBase].append(
+        ('dipole_moment', dipole_moment))
+
+    @warn_if_not_unique
+    @_pbc_to_wrap
+    @check_wrap_and_unwrap
+    def quadrupole_moment(group, wrap=False, unwrap=False, 
+        compound='group', center="mass"):
+        r"""Quadrupole vector of the group according to [Gray1984]_
+         
+        .. math::
+            Q = \sqrt{\frac{2}{3}\boldsymbol{\hat{Q}}:\boldsymbol{\hat{Q}}}
+
+        where :math:`\boldsymbol{\hat{Q}}` is the traceless quadrupole 
+        tensor.
+
+        Computes the quadrupole moment of :class:`Atoms<Atom>` in the group.
+        Dipoleper :class:`Residue`, :class:`Segment`, molecule, or
+        fragment can be obtained by setting the `compound` parameter
+        accordingly.
+
+        Parameters
+        ----------
+        group : obj
+            AtomGroup
+        wrap : bool, optional
+            If ``True`` and `compound` is ``'group'``, move all atoms to the
+            primary unit cell before calculation.
+            If ``True`` and `compound` is not ``group``, the
+            centers of mass of each compound will be calculated without moving
+            any atoms to keep the compounds intact.
+        unwrap : bool, optional
+            If ``True``, compounds will be unwrapped before computing their
+            centers.
+        compound : {'group', 'segments', 'residues', 'molecules', \
+                    'fragments'}, optional
+            If ``'group'``, a single quadrupole value returns. Otherwise, an
+            array of each :class:`Segment`, :class:`Residue`, molecule, or
+            fragment will be returned as an array of position vectors, i.e.
+            a 1d array.
+            Note that, in any case, *only* the positions of
+            :class:`Atoms<Atom>` *belonging to the group* will be taken into
+            account.
+        center : str, optional
+            Choose whether the quadrupole moment is calculated at the center of 
+            "mass" or the center of "charge"
+
+        Returns
+        -------
+        numpy.ndarray
+            Quadrupole moment of (compounds of) the group in e*angstroms**2.
+            If `compound` was set to ``'group'``, the output will be a single
+            value. Otherwise, the output will be a 1d array of shape ``(n,)``
+            where ``n`` is the number of compounds.
+
+        .. versionadded:: 2.4.0
+        """
+        def __quadrupole(recenteredpos, charges):
+            tensor = np.zeros((3,3))
+            for i,atom in enumerate(recenteredpos):
+                tensor += np.matmul(
+                    atom[:,np.newaxis],
+                    atom[np.newaxis,:])*charges[i]
+
+            quad_trace = np.sum(np.diag(tensor))
+            tensor = 3*tensor/2
+            for j in (0,1,2):
+                tensor[j][j] += -quad_trace/2
+            return np.sqrt(2*np.tensordot(tensor,tensor)/3)
+
+        compound = compound.lower()
+
+        atomgroup = group.atoms
+        charges = atomgroup.charges
+
+        if center == "mass":
+            masses = atomgroup.masses
+            ref = atomgroup.center_of_mass(
+                wrap=wrap,
+                unwrap=unwrap,
+                compound=compound
+            )
+        elif center == "charge":
+            ref = atomgroup.center_of_charge(
+                wrap=wrap,
+                unwrap=unwrap,
+                compound=compound
+            )
+        else:
+            choices = ["mass", "charge"]
+            raise ValueError(
+                "The dipole center, {}, is not supported. Choose: {}".format(
+                 center,choices))
+
+        if compound == 'group':
+            if wrap:
+                recenteredpos = (atomgroup.pack_into_box(inplace=False) - ref)
+            elif unwrap:
+                recenteredpos = (atomgroup.unwrap(inplace=False) - ref)
+            else:
+                recenteredpos = (atomgroup.positions - ref)
+            quad_moment = __quadrupole(recenteredpos, charges)
+        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
+            chgs = atomgroup.charges
+
+            quad_moment = np.empty(n_compounds, dtype=np.float64)
+            for compound_mask, atom_mask in zip(compound_masks, atom_masks):
+                quad_moment[compound_mask] = [__quadrupole(
+                     coords[mask]-ref[compound_mask][i],
+                     chgs[mask][:,None]
+                    ) for i,mask in enumerate(atom_mask)]
+
+        return quad_moment
+
+    transplants[GroupBase].append(
+        ('quadrupole_moment', quadrupole_moment))
+
 
 class FormalCharges(AtomAttr):
     """Formal charge on each atom"""

package/doc/sphinx/source/documentation_pages/references.rst

@@ -1,5 +1,5 @@
 .. -*- coding: utf-8 -*-
-.. note: make sure that no lines accidentaly start with a single character
+.. note: make sure that no lines accidentally start with a single character
 ..       followed by a period: reST interprets it as an enumerated list and
 ..       messes up the formatting
 
@@ -204,6 +204,17 @@ please cite [Dima2004b]_.
    6564-6570. doi:`10.1021/jp037128y
    <https://doi.org/10.1021/jp037128y>`_
 
+If you calculate shape parameters using
+:meth:`~MDAnalysis.core.group.AtomGroup.quadrupole_moment`,
+:meth:`~MDAnalysis.core.group.ResidueGroup.quadrupole_moment`,
+:meth:`~MDAnalysis.core.group.SegmentGroup.quadrupole_moment`
+please cite [Gray1984]_.
+
+.. [Gray1984] Gray, C. G., Gubbins, K. E., & Joslin C. G. (1984). 
+   Theory of molecular fluids. The International Series of 
+   Monographs on Chemistry 9. Oxford; New York: Oxford University
+   Press.
+
 If you use use the dielectric analysis code in
 :class:`~MDAnalysis.analysis.dielectric.DielectricConstant` please cite [Neumann1983]_.