From 0eb46ee91c91ed38ed50e44187831c09ea125edc Mon Sep 17 00:00:00 2001
From: Leon Wehrhan <mail@leonwehrhan.de>
Date: Fri, 22 Mar 2024 18:17:34 +0100
Subject: [PATCH] Implement iterative average structure (Issue#2039)

---
 package/AUTHORS                               |  1 +
 package/CHANGELOG                             |  6 ++
 package/MDAnalysis/analysis/rms.py            | 60 +++++++++++++++++++
 .../MDAnalysisTests/analysis/test_rms.py      | 36 +++++++++++
 4 files changed, 103 insertions(+)

diff --git a/package/AUTHORS b/package/AUTHORS
index 35e4029f776..3d6b749d152 100644
--- a/package/AUTHORS
+++ b/package/AUTHORS
@@ -236,6 +236,7 @@ Chronological list of authors
   2024
   - Aditya Keshari
   - Philipp Stärk
+  - Leon Wehrhan
 
 External code
 -------------
diff --git a/package/CHANGELOG b/package/CHANGELOG
index 78139e965e1..9d510401682 100644
--- a/package/CHANGELOG
+++ b/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
diff --git a/package/MDAnalysis/analysis/rms.py b/package/MDAnalysis/analysis/rms.py
index 36067d76b6d..767c2821144 100644
--- a/package/MDAnalysis/analysis/rms.py
+++ b/package/MDAnalysis/analysis/rms.py
@@ -167,6 +167,7 @@
 import warnings
 
 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 +333,65 @@ def process_selection(select):
     return select
 
 
+def iterative_average(
+    mobile, ref, weights=None, niter=100, eps=1e-8, verbose=False, **kwargs
+):
+    """Calculate an optimal reference that is also the average structure
+    after an RMSD alignment.
+    The optional 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.AtomGroup
+        mobile atomgroup to find the average for
+    ref : mda.AtomGroup
+        initial reference structure. Positions are changed by this function!
+    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
+
+    """
+    drmsd = np.inf
+    for i in range(niter):
+        # found a converged structure
+        if drmsd < eps:
+            break
+
+        res = align.AverageStructure(
+            mobile, reference=ref, weights=weights, **kwargs
+        ).run()
+        ref_ag = ref.atoms
+        drmsd = rmsd(ref_ag.positions, res.positions, weights=weights)
+        ref.positions = res.positions
+
+        if verbose:
+            print(
+                "i = {}, rmsd-change = {:.2f}, ave-rmsd = {:.2f}".format(
+                    i, drmsd, res.results.rmsd
+                )
+            )
+
+    return res
+
+
 class RMSD(AnalysisBase):
     r"""Class to perform RMSD analysis on a trajectory.
 
diff --git a/testsuite/MDAnalysisTests/analysis/test_rms.py b/testsuite/MDAnalysisTests/analysis/test_rms.py
index 56d2e459e94..a70f391fb6f 100644
--- a/testsuite/MDAnalysisTests/analysis/test_rms.py
+++ b/testsuite/MDAnalysisTests/analysis/test_rms.py
@@ -420,3 +420,39 @@ 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)
+        u.transfer_to_memory()
+        u1 = mda.Merge(u.select_atoms("bynum 1:10"))
+        return u1
+
+    @pytest.fixture()
+    def reference(self):
+        u = mda.Universe(PSF, DCD)
+        u1 = mda.Merge(u.select_atoms("bynum 1:10"))
+        return u1
+
+    def test_iterative_average(self, mobile, reference):
+
+        res = rms.iterative_average(mobile, reference, niter=100, verbose=True)
+        res = rms.iterative_average(mobile, reference, niter=10, verbose=False)
+        assert_almost_equal(
+            res.positions,
+            [
+                [11.73604393, 8.50079727, -10.44528103],
+                [12.36511898, 7.83993578, -10.83484173],
+                [12.09194851, 9.441535, -10.72461128],
+                [10.83195686, 8.30861378, -10.96393108],
+                [11.66462231, 8.39347267, -8.98323059],
+                [12.67261219, 8.5604887, -8.6807642],
+                [10.73963833, 9.56992817, -8.59075069],
+                [9.63879013, 9.27048492, -8.96795273],
+                [11.09461403, 10.51888371, -9.14199638],
+                [10.55657101, 9.93423939, -7.11236286],
+            ],
+            decimal=5,
+        )