Implement iterative average structure (Issue#2039)

package/AUTHORS

@@ -237,6 +237,8 @@ Chronological list of authors
   - Aditya Keshari
   - Philipp Stärk
   - Sampurna Mukherjee
+  - Leon Wehrhan
+
 
 External code
 -------------

package/CHANGELOG

@@ -16,7 +16,8 @@ The rules for this file:
 -------------------------------------------------------------------------------
 ??/??/?? IAlibay, HeetVekariya, marinegor, lilyminium, RMeli,
          ljwoods2, aditya292002, pstaerk, PicoCentauri, BFedder,
-         tyler.je.reddy, SampurnaM
+         tyler.je.reddy, SampurnaM, leonwehrhan
+
 
  * 2.8.0
 
@@ -40,6 +41,8 @@ Enhancements
  * Added parsing of arbitrary columns of the LAMMPS dump parser. (Issue #3504)
  * Documented the r0 attribute in the `Contacts` class and added the 
    `n_initial_contacts` attribute, with documentation. (Issue #2604, PR #4415)
+ * Implement average structures with iterative algorithm from 
+   DOI 10.1021/acs.jpcb.7b11988. (Issue #2039, PR #4524)
 
 Changes
  * As per SPEC0 the minimum supported Python version has been raised

package/MDAnalysis/analysis/align.py

@@ -174,6 +174,7 @@
 .. autoclass:: AlignTraj
 .. autoclass:: AverageStructure
 .. autofunction:: rotation_matrix
+.. autofunction:: iterative_average
 
 
 Helper functions
@@ -212,6 +213,8 @@
 from MDAnalysis.coordinates.memory import MemoryReader
 from MDAnalysis.lib.util import get_weights
 from MDAnalysis.lib.util import deprecate   # remove 3.0
+from MDAnalysis.lib.log import ProgressBar
+from ..due import due, Doi
 
 from .base import AnalysisBase
 
@@ -541,6 +544,120 @@
     return old_rmsd, new_rmsd
 
 
+@due.dcite(
+        Doi("10.1021/acs.jpcb.7b11988"),
+        description="Iterative Calculation of Opimal Reference",
+        path="MDAnalysis.analysis.align.iterative_average"
+)
+def iterative_average(
+    mobile, reference=None, select='all', weights=None, niter=100,
+    eps=1e-6, verbose=False, **kwargs
+):
+    """Iteratively calculate an optimal reference that is also the average
+    structure after an RMSD alignment.
+
+    The optimal reference is defined as average
+    structure of a trajectory, with the optimal reference used as input.
+    This function computes the optimal reference by using a starting
+    reference for the average structure, which is used as the reference
+    to calculate the average structure again. This is repeated until the
+    reference structure has converged. :footcite:p:`Linke2018`
+
+    Parameters
+    ----------
+    mobile : mda.Universe
+        Universe containing trajectory to be fitted to reference.
+    reference : mda.Universe (optional)
+        Universe containing the initial reference structure.
+    select : str or tuple or dict (optional)
+        Atom selection for fitting a substructue. Default is set to all.
+        Can be tuple or dict to define different selection strings for
+        mobile and target.
+    weights : str, array_like (optional)
+        Weights that can be used. If `None` use equal weights, if `'mass'`
+        use masses of ref as weights or give an array of arbitrary weights.
+    niter : int (optional)
+        Maximum number of iterations.
+    eps : float (optional)
+        RMSD distance at which reference and average are assumed to be
+        equal.
+    verbose : bool (optional)
+        Verbosity.
+    **kwargs : dict (optional)
+        AverageStructure kwargs.
+
+    Returns
+    -------
+    avg_struc : AverageStructure
+        AverageStructure result from the last iteration.
+
+    Example
+    -------
+    `iterative_average` can be used to obtain a :class:`MDAnalysis.Universe`
+    with the optimal reference structure.
+
+    ::
+
+        import MDAnalysis as mda
+        from MDAnalysis.analysis import align
+        from MDAnalysisTests.datafiles import PSF, DCD
+
+        u = mda.Universe(PSF, DCD)
+        av = align.iterative_average(u, u, verbose=True)
+
+        averaged_universe = av.results.universe
+
+    .. versionadded:: 2.8.0
+    """
+    if not reference:
+        reference = mobile
+
+    select = rms.process_selection(select)
+    ref = mda.Merge(reference.select_atoms(*select['reference']))
+    sel_mobile = select['mobile'][0]
+
+    weights = get_weights(ref.atoms, weights)
+
+    drmsd = np.inf
+    for i in ProgressBar(range(niter)):
+        # found a converged structure
+        if drmsd < eps:
+            break
+
+        avg_struc = AverageStructure(
+            mobile, reference=ref, select={
+                'mobile': sel_mobile, 'reference': 'all'
+                },
+            weights=weights, **kwargs
+        ).run()
+        drmsd = rms.rmsd(ref.atoms.positions, avg_struc.results.positions,
+                         weights=weights)
+        ref = avg_struc.results.universe
+
+        if verbose:
+            logger.debug(
+                f"iterative_average(): i = {i}, "
+                f"rmsd-change = {drmsd:.5f}, "
+                f"ave-rmsd = {avg_struc.results.rmsd:.5f}"
+            )
+
+    else:
+        errmsg = (
+            "iterative_average(): Did not converge in "
+            f"{niter} iterations to DRMSD < {eps}. "
+            f"Final average RMSD = {avg_struc.results.rmsd:.5f}"
+        )
+        logger.error(errmsg)
+        raise RuntimeError(errmsg)
+
+    logger.info(
+        f"iterative_average(): Converged to DRMSD < {eps}. "
+        f"Final average RMSD = {avg_struc.results.rmsd:.5f}"
+    )
+
+    return avg_struc
+
+
 class AlignTraj(AnalysisBase):
     """RMS-align trajectory to a reference structure using a selection.
 

package/MDAnalysis/analysis/rms.py

@@ -166,7 +166,9 @@
 import logging
 import warnings
 
+import MDAnalysis as mda
 import MDAnalysis.lib.qcprot as qcp
+import MDAnalysis.analysis.align as align
 from MDAnalysis.analysis.base import AnalysisBase
 from MDAnalysis.exceptions import SelectionError, NoDataError
 from MDAnalysis.lib.util import asiterable, iterable, get_weights

package/doc/sphinx/source/references.bib

@@ -770,3 +770,14 @@ @article{Kulke2022
   pages = {6161--6171},
   doi = {10.1021/acs.jctc.2c00327}
 }
+
+@article{Linke2018,
+  title = {Fully Anisotropic Rotational Diffusion Tensor from Molecular Dynamics Simulations},
+  author = {Linke, Max and Köfinger, Jürgen and Hummer, Gerhard},
+  year = {2018},
+  journal = {The Journal of Physical Chemistry B},
+  volume = {122},
+  number = {21},
+  pages = {5630--5639},
+  doi = {10.1021/acs.jpcb.7b11988}
+}

testsuite/MDAnalysisTests/analysis/test_align.py

@@ -594,6 +594,83 @@ def test_sequence_alignment_deprecation(self, atomgroups):
             align.sequence_alignment(mobile, reference)
 
 
+class TestIterativeAverage(object):
+    @pytest.fixture()
+    def mobile(self):
+        u = mda.Universe(PSF, DCD)
+        return u
+
+    @pytest.fixture()
+    def reference(self):
+        u = mda.Universe(PSF, DCD)
+        return u
+
+    def test_iterative_average_default(self, mobile):
+        res = align.iterative_average(mobile, select="bynum 1:10")
+        assert_allclose(
+            res.results.positions,
+            [
+                [11.93075595, 8.6729893, -10.49887605],
+                [12.60587898, 7.91673117, -10.73327464],
+                [12.45662411, 9.51900517, -10.35551193],
+                [11.27452274, 8.83003843, -11.2619057],
+                [11.25808119, 8.26794477, -9.23340715],
+                [12.02767222, 7.95332228, -8.57249317],
+                [10.54679871, 9.49505306, -8.61215292],
+                [9.99500556, 9.16624224, -7.75231192],
+                [9.83897407, 9.93134598, -9.29541129],
+                [11.45760169, 10.5857071, -8.13037669]
+            ],
+            atol=1e-5,
+        )
+
+    def test_iterative_average_eps_high(self, mobile):
+        res = align.iterative_average(mobile, select="bynum 1:10",
+                                      eps=1e-5)
+        assert_allclose(
+            res.results.positions,
+            [
+                [11.93075595, 8.6729893, -10.49887605],
+                [12.60587898, 7.91673117, -10.73327464],
+                [12.45662411, 9.51900517, -10.35551193],
+                [11.27452274, 8.83003843, -11.2619057],
+                [11.25808119, 8.26794477, -9.23340715],
+                [12.02767222, 7.95332228, -8.57249317],
+                [10.54679871, 9.49505306, -8.61215292],
+                [9.99500556, 9.16624224, -7.75231192],
+                [9.83897407, 9.93134598, -9.29541129],
+                [11.45760169, 10.5857071, -8.13037669]
+            ],
+            atol=1e-5,
+        )
+
+    def test_iterative_average_weights_mass(self, mobile, reference):
+        res = align.iterative_average(mobile, reference,
+                                      select="bynum 1:10",
+                                      niter=10, weights="mass")
+        assert_allclose(
+            res.results.positions,
+            [
+                [11.96438784, 8.85426235, -10.24735737],
+                [12.75920431, 8.27294545, -10.54295766],
+                [12.3285704, 9.72083717, -9.9419435],
+                [11.33941507, 9.03249423, -11.01306158],
+                [11.30988499, 8.14958885, -9.1205501],
+                [12.09108655, 7.85155906, -8.46681943],
+                [10.37499697, 9.13535837, -8.3732586],
+                [9.83883314, 8.57939098, -7.6195549],
+                [9.64405257, 9.55924307, -9.04315991],
+                [11.0678934, 10.27798773, -7.64881842]
+            ],
+            atol=1e-5,
+        )
+
+    def test_iterative_average_convergence_failure(self, mobile, reference):
+        with pytest.raises(RuntimeError):
+            _ = align.iterative_average(mobile, reference,
+                                        niter=1, eps=0)
+
+
 def test_alignto_reorder_atomgroups():
     # Issue 2977
     u = mda.Universe(PDB_helix)