Implement iterative average structure (Issue#2039)

package/AUTHORS

@@ -236,6 +236,7 @@ Chronological list of authors
   2024
   - Aditya Keshari
   - Philipp Stärk
+  - Leon Wehrhan
 
 External code
 -------------

package/CHANGELOG

@@ -62,6 +62,12 @@ Deprecations
    (PR #4403)
 
 
+03/22/24 leonwehrhan
+
+Changes
+  * Implement average structures with iterative algorithm from DOI 10.1021/acs.jpcb.7b11988. (Issue #2039)
+
+
 12/26/23 IAlibay, ianmkenney, PicoCentauri, pgbarletta, p-j-smith, 
          richardjgowers, lilyminium, ALescoulie, hmacdope, HeetVekariya, 
          JoStoe, jennaswa, ljwoods2

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
@@ -332,6 +334,78 @@ def process_selection(select):
     return select
 
 
+def iterative_average(
+    mobile, ref, select='all', weights=None, niter=100, eps=1e-6,
+    verbose=False, **kwargs
+):
+    """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 to calculate the
+    average structure again. This is repeated until the reference
+    structure has converged.
+
+    Parameters
+    ----------
+    mobile : mda.Universe
+        Universe containing trajectory to be fitted to reference.
+    ref : mda.Universe
+        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
+    -------
+    res : AverageStructure
+        AverageStructure result from the last iteration.
+
+    """
+    select = process_selection(select)
+    ref = mda.Merge(ref.select_atoms(*select['reference']))
+    sel_mobile = select['mobile'][0]
+
+    weights = get_weights(ref.atoms, weights)
+
+    drmsd = np.inf
+    for i in range(niter):
+        # found a converged structure
+        if drmsd < eps:
+            break
+
+        res = align.AverageStructure(
+            mobile, reference=ref, select={
+                'mobile': sel_mobile, 'reference': 'all'
+                },
+            weights=weights, **kwargs
+        ).run()
+        drmsd = rmsd(ref.atoms.positions, res.results.positions, weights=weights)
+        ref = res.results.universe
+
+        if verbose:
+            print(
+                "i = {}, rmsd-change = {:.5f}, ave-rmsd = {:.5f}".format(
+                    i, drmsd, res.results.rmsd
+                )
+            )
+
+    return res
+
+
 class RMSD(AnalysisBase):
     r"""Class to perform RMSD analysis on a trajectory.
 

testsuite/MDAnalysisTests/analysis/test_rms.py

@@ -420,3 +420,59 @@ def test_rmsf_attr_warning(self, universe):
         wmsg = "The `rmsf` attribute was deprecated in MDAnalysis 2.0.0"
         with pytest.warns(DeprecationWarning, match=wmsg):
             assert_equal(rmsfs.rmsf, rmsfs.results.rmsf)
+
+
+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(self, mobile, reference):
+        res = rms.iterative_average(mobile, reference, niter=1, eps=0)
+        res = rms.iterative_average(mobile, reference, eps=1e-5)
+        res = rms.iterative_average(mobile, reference, select='bynum 1:10',
+                                    weights=np.ones(10))
+
+        res = rms.iterative_average(mobile, reference, select='bynum 1:10',
+                                    niter=10, verbose=True)
+        assert_almost_equal(
+            res.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]
+            ],
+            decimal=5,
+        )
+
+        res = rms.iterative_average(mobile, reference, select='bynum 1:10',
+                                    niter=10, weights='mass')
+        assert_almost_equal(
+            res.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]
+            ],
+            decimal=5,
+        )