Add ref_model option to calcSignatureOverlaps with tests

prody/dynamics/signature.py

@@ -1388,7 +1388,7 @@ def calcSignatureCollectivity(mode_ensemble, masses=None):
 
     return sig
 
-def calcSignatureOverlaps(mode_ensemble, diag=True, collapse=False):
+def calcSignatureOverlaps(mode_ensemble, diag=True, collapse=False, ref_model=None):
     """Calculate average mode-mode overlaps for a ModeEnsemble.
 
     If *diag* is **True** (default) then only diagonal values will be calculated. 
@@ -1398,7 +1398,12 @@ def calcSignatureOverlaps(mode_ensemble, diag=True, collapse=False):
     a 4-dimensional sdarray that is a matrix of overlap matrices. 
 
     If *collapse* is **True** then these will be collapsed together, giving a 2-dimensional 
-    array for full matrices. This operation is not defined for diagonal values."""
+    array for full matrices. This operation is not defined for diagonal values.
+
+    The overlaps can now be calculated relative to a *ref_model* instead of within the ensemble.
+    """
+    if ref_model is not None and not isinstance(ref_model, (NMA, ModeSet, Mode, Vector)):
+        raise TypeError('ref_model must be of type NMA, ModeSet, Mode or Vector')
 
     if isinstance(mode_ensemble, ModeEnsemble):
         if not mode_ensemble.isMatched():
@@ -1407,9 +1412,13 @@ def calcSignatureOverlaps(mode_ensemble, diag=True, collapse=False):
 
         n_sets = mode_ensemble.numModeSets()
         n_modes = mode_ensemble.numModes()
+        n_atoms = mode_ensemble.numAtoms()
     else:
         if not isListLike(mode_ensemble):
             raise TypeError('mode_ensemble should be list-like or an instance of ModeEnsemble')
+
+        if not np.all(np.array([isinstance(modeset, (ModeSet, NMA)) for modeset in mode_ensemble])):
+            raise TypeError("mode_ensemble should contain ModeSet or NMA objects")
 
         n_sets = len(mode_ensemble)
 
@@ -1418,39 +1427,84 @@ def calcSignatureOverlaps(mode_ensemble, diag=True, collapse=False):
             n_modes = n_modes[0]
         else:
             raise ValueError('all mode sets in mode_ensemble should have the same number of modes')
-
-    if diag:
-        if collapse:
-            LOGGER.warn('cannot collapse diagonal values')
-        overlaps = np.zeros((n_modes, n_sets, n_sets))
+
+        n_atoms = np.array([modeset.numAtoms() for modeset in mode_ensemble])
+        if np.all(n_atoms == n_atoms[0]):
+            n_atoms = n_atoms[0]
+        else:
+            raise ValueError('all mode sets in mode_ensemble should have the same number of atoms')
+
+    if ref_model is None:
+        if diag:
+            if collapse:
+                LOGGER.warn('cannot collapse diagonal values')
+            overlaps = np.zeros((n_modes, n_sets, n_sets))
+        else:
+            if collapse:
+                overlaps = np.zeros((n_modes*n_sets, n_modes*n_sets))
+            else:
+                overlaps = np.zeros((n_modes, n_modes, n_sets, n_sets))
+
+        for i, modeset_i in enumerate(mode_ensemble):
+            for j, modeset_j in enumerate(mode_ensemble):
+                if j >= i:
+                    if diag:
+                        overlaps[:,i,j] = overlaps[:,j,i] = abs(calcOverlap(modeset_i, 
+                                                                            modeset_j, 
+                                                                            diag=True))
+                    else:
+                        if collapse:
+                            overlaps[i*n_modes:(i+1)*n_modes,
+                                     j*n_modes:(j+1)*n_modes] = np.abs(calcOverlap(modeset_i,
+                                                                                   modeset_j))
+                            overlaps[j*n_modes:(j+1)*n_modes,
+                                     i*n_modes:(i+1)*n_modes] = np.abs(calcOverlap(modeset_j,
+                                                                                   modeset_i))
+                        else:
+                            overlaps[:, :, i, j] = abs(calcOverlap(modeset_i,
+                                                                   modeset_j,
+                                                                   diag=False))
+                            overlaps[:, :, j, i] = abs(calcOverlap(modeset_j,
+                                                                  modeset_i,
+                                                                  diag=False))
     else:
-        if collapse:
-            overlaps = np.zeros((n_modes*n_sets, n_modes*n_sets))
+        if ref_model.numAtoms() != n_atoms:
+            raise ValueError("ref_model must have the same number of atoms as the ensemble")
+
+        n_modes_ref = ref_model.numModes()
+
+        if diag:
+            if ref_model.numModes() != n_modes:
+                raise ValueError(
+                    "cannot take diagonal if ref_model does not have the same number of modes as the ensemble"
+                )
+            if collapse:
+                LOGGER.warn('cannot collapse diagonal values')
+            overlaps = np.zeros((n_modes, n_sets, 1))
         else:
-            overlaps = np.zeros((n_modes, n_modes, n_sets, n_sets))
-
-    for i, modeset_i in enumerate(mode_ensemble):
-        for j, modeset_j in enumerate(mode_ensemble):
-            if j >= i:
-                if diag:
-                    overlaps[:,i,j] = overlaps[:,j,i] = abs(calcOverlap(modeset_i, 
-                                                                        modeset_j, 
-                                                                        diag=True))
+            if collapse:
+                overlaps = np.zeros((n_modes*n_sets, n_modes_ref))
+            else:
+                overlaps = np.zeros((n_modes, n_modes_ref, n_sets, 1))
+
+        j = 0 # only one ref_model
+        modeset_j = ref_model
+        for i, modeset_i in enumerate(mode_ensemble):
+            if diag:
+                overlaps[:,i,j] = abs(calcOverlap(modeset_i,
+                                                  modeset_j,
+                                                  diag=True))
+            else:
+                if collapse:
+                    overlaps[i*n_modes:(i+1)*n_modes,
+                             j*n_modes_ref:(j+1)*n_modes_ref] = np.abs(
+                                 calcOverlap(modeset_i, modeset_j)
+                            ).reshape(n_modes, n_modes_ref)
                 else:
-                    if collapse:
-                        overlaps[i*n_modes:(i+1)*n_modes,
-                                 j*n_modes:(j+1)*n_modes] = np.abs(calcOverlap(modeset_i,
-                                                                               modeset_j))
-                        overlaps[j*n_modes:(j+1)*n_modes,
-                                 i*n_modes:(i+1)*n_modes] = np.abs(calcOverlap(modeset_j,
-                                                                               modeset_i))
-                    else:
-                        overlaps[:, :, i, j] = abs(calcOverlap(modeset_i,
-                                                               modeset_j,
-                                                               diag=False))
-                        overlaps[:, :, j, i] = abs(calcOverlap(modeset_j,
-                                                               modeset_i,
-                                                               diag=False))
+                    overlaps[:, :, i, j] = abs(
+                        calcOverlap(modeset_i,
+                                    modeset_j,
+                                    diag=False)).reshape(n_modes, n_modes_ref)
 
     return overlaps
 

prody/tests/dynamics/test_signature.py

@@ -1,14 +1,17 @@
 """This module contains unit tests for :mod:`~prody.KDTree` module."""
 
-from numpy.testing import assert_array_equal, assert_equal
+import numpy as np
+from numpy.testing import assert_array_equal, assert_equal, assert_allclose
 from numpy.random import rand, randint
 
 from prody.dynamics import sdarray
+from prody.dynamics.signature import ModeEnsemble
 
 from prody.tests import unittest
 from prody.tests.datafiles import parseDatafile
 
-from prody import _PY3K, LOGGER
+from prody import (_PY3K, LOGGER, ANM, PDBEnsemble, calcEnsembleENMs,
+                   calcSignatureModes, calcSignatureOverlaps, calcOverlap)
 
 LOGGER.verbosity = 'none'
 
@@ -43,3 +46,222 @@ def testSlicing(self):
 
         s = S[0, 0, 0]
         #assert_array_equal(s, A[0, 0, 0], 'failed at sdarray slicing')
+
+
+class TestSignatureOverlapsRefModel(unittest.TestCase):
+    """Tests for the *ref_model* option of :func:`.calcSignatureOverlaps`, which
+    overlaps each modeset of the ensemble against a single reference model
+    rather than within the ensemble."""
+
+    @classmethod
+    def setUpClass(cls):
+        ca = parseDatafile('pdb1ake').select('calpha')
+        cls.ca = ca
+        cls.n_atoms = ca.numAtoms()
+        cls.n_modes = 5
+        cls.n_sets = 3
+
+        # a small matched mode ensemble on a common set of atoms
+        me = ModeEnsemble('test')
+        for _ in range(cls.n_sets):
+            anm = ANM()
+            anm.buildHessian(ca)
+            anm.calcModes(n_modes=cls.n_modes, zeros=False)
+            me.addModeSet(anm[:cls.n_modes])
+        cls.me = me
+
+        # a multi-mode reference and a single-mode reference on the same atoms
+        ref = ANM()
+        ref.buildHessian(ca)
+        ref.calcModes(n_modes=cls.n_modes, zeros=False)
+        cls.ref_full = ref[:cls.n_modes]
+        cls.ref_single = ref[0]
+
+    # -- shapes -----------------------------------------------------------
+
+    def testRefModelFullShape(self):
+        """A multi-mode reference gives (n_modes, n_modes_ref, n_sets, 1)."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_full, diag=False)
+        self.assertEqual(ov.shape,
+                         (self.n_modes, self.n_modes, self.n_sets, 1))
+
+    def testRefModelCollapseShape(self):
+        """collapse stacks each set's block: (n_modes*n_sets, n_modes_ref)."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_full,
+                                   diag=False, collapse=True)
+        self.assertEqual(ov.shape,
+                         (self.n_modes * self.n_sets, self.n_modes))
+        self.assertTrue(np.all(np.isfinite(ov)))
+
+    def testRefModelSingleModeNonCollapse(self):
+        """A single-mode reference gives n_modes_ref == 1 (reshape path)."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_single,
+                                   diag=False)
+        self.assertEqual(ov.shape, (self.n_modes, 1, self.n_sets, 1))
+
+    def testRefModelSingleModeCollapse(self):
+        """Single-mode reference with collapse: (n_modes*n_sets, 1)."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_single,
+                                   diag=False, collapse=True)
+        self.assertEqual(ov.shape, (self.n_modes * self.n_sets, 1))
+        self.assertTrue(np.all(np.isfinite(ov)))
+
+    def testRefModelDiagShape(self):
+        """diag against an equal-mode reference: (n_modes, n_sets, 1)."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_full, diag=True)
+        self.assertEqual(ov.shape, (self.n_modes, self.n_sets, 1))
+
+    # -- values -----------------------------------------------------------
+
+    def testNonCollapseMatchesCalcOverlap(self):
+        """Each block equals abs(calcOverlap(modeset_i, ref))."""
+        ov = calcSignatureOverlaps(self.me, ref_model=self.ref_full, diag=False)
+        for i in range(self.n_sets):
+            expected = np.abs(calcOverlap(self.me[i], self.ref_full))
+            assert_allclose(ov[:, :, i, 0], expected, atol=1e-8)
+
+    def testCollapseMatchesNonCollapse(self):
+        """The collapsed stack holds the same blocks as the 4-D array."""
+        ov4 = calcSignatureOverlaps(self.me, ref_model=self.ref_full, diag=False)
+        ov2 = calcSignatureOverlaps(self.me, ref_model=self.ref_full,
+                                    diag=False, collapse=True)
+        for i in range(self.n_sets):
+            block = ov2[i * self.n_modes:(i + 1) * self.n_modes, :]
+            assert_allclose(block, ov4[:, :, i, 0], atol=1e-8)
+
+    def testSingleModeCollapseMatchesNonCollapse(self):
+        """Single-mode reference: collapse and non-collapse agree."""
+        ov4 = calcSignatureOverlaps(self.me, ref_model=self.ref_single,
+                                    diag=False)
+        ov2 = calcSignatureOverlaps(self.me, ref_model=self.ref_single,
+                                    diag=False, collapse=True)
+        for i in range(self.n_sets):
+            block = ov2[i * self.n_modes:(i + 1) * self.n_modes, :]
+            assert_allclose(block, ov4[:, :, i, 0], atol=1e-8)
+
+    def testListInputEqualsEnsembleInput(self):
+        """A plain list of modesets is accepted and matches ModeEnsemble input."""
+        modesets = [ms for ms in self.me]
+        ov_list = calcSignatureOverlaps(modesets, ref_model=self.ref_full,
+                                        diag=False)
+        ov_ens = calcSignatureOverlaps(self.me, ref_model=self.ref_full,
+                                       diag=False)
+        assert_allclose(ov_list, ov_ens, atol=1e-8)
+
+    # -- errors -----------------------------------------------------------
+
+    def testRefModelWrongTypeRaises(self):
+        self.assertRaises(TypeError, calcSignatureOverlaps, self.me,
+                          ref_model='not-a-model')
+
+    def testRefModelWrongNumAtomsRaises(self):
+        small = ANM()
+        small.buildHessian(self.ca[:self.n_atoms // 2])
+        small.calcModes(n_modes=self.n_modes, zeros=False)
+        self.assertRaises(ValueError, calcSignatureOverlaps, self.me,
+                          ref_model=small[:self.n_modes], diag=False)
+
+    def testRefModelDiagModeMismatchRaises(self):
+        """diag needs the reference to have as many modes as the ensemble."""
+        self.assertRaises(ValueError, calcSignatureOverlaps, self.me,
+                          ref_model=self.ref_single, diag=True)
+
+    # -- regression: within-ensemble path still works ---------------------
+
+    def testWithinEnsembleCollapseUnchanged(self):
+        """Without ref_model, collapse gives the full block matrix."""
+        ov = calcSignatureOverlaps(self.me, diag=False, collapse=True)
+        self.assertEqual(
+            ov.shape,
+            (self.n_modes * self.n_sets, self.n_modes * self.n_sets))
+
+
+class TestSignatureOverlapsSignDy(unittest.TestCase):
+    """ref_model overlaps on GNM and ANM mode ensembles built the SignDy way,
+    i.e. :func:`.calcEnsembleENMs` on a :class:`.PDBEnsemble`, overlapped against
+    the mean signature model from :func:`.calcSignatureModes` (as in the SignDy
+    tutorial). Exercised on the 2k39 ubiquitin NMR ensemble."""
+
+    @classmethod
+    def setUpClass(cls):
+        cls.n_modes = 10
+        cls.n_sets = 6
+
+        ags = parseDatafile('2k39_ca')
+        ens = PDBEnsemble('2k39')
+        ens.setAtoms(ags)
+        ens.setCoords(ags.getCoords())
+        ens.addCoordset(ags.getCoordsets()[:cls.n_sets])
+        ens.iterpose()
+        cls.n_atoms = ens.numAtoms()
+
+        # one GNM and one ANM mode ensemble, each with its mean-signature model
+        cls.mes = {}
+        for model in ('gnm', 'anm'):
+            me = calcEnsembleENMs(ens, model=model, trim='reduce',
+                                  n_modes=cls.n_modes, match=True)
+            cls.mes[model] = (me, calcSignatureModes(me))
+
+    def testEnsembleShape(self):
+        """Both ensembles are matched with the expected dimensions."""
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                self.assertEqual(me.numModeSets(), self.n_sets)
+                self.assertEqual(me.numModes(), self.n_modes)
+                self.assertEqual(me.numAtoms(), self.n_atoms)
+                self.assertTrue(me.isMatched())
+                self.assertEqual(ref.numModes(), self.n_modes)
+                self.assertEqual(ref.numAtoms(), self.n_atoms)
+
+    def testRefModelFullShapeAndValues(self):
+        """(n_modes, n_modes_ref, n_sets, 1) and each block == calcOverlap."""
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                ov = calcSignatureOverlaps(me, ref_model=ref, diag=False)
+                self.assertEqual(
+                    ov.shape, (self.n_modes, self.n_modes, self.n_sets, 1))
+                for i in range(self.n_sets):
+                    expected = np.abs(calcOverlap(me[i], ref))
+                    assert_allclose(ov[:, :, i, 0], expected, atol=1e-8)
+
+    def testRefModelCollapse(self):
+        """Collapsed stack matches the 4-D blocks for both models."""
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                ov4 = calcSignatureOverlaps(me, ref_model=ref, diag=False)
+                ov2 = calcSignatureOverlaps(me, ref_model=ref, diag=False,
+                                            collapse=True)
+                self.assertEqual(
+                    ov2.shape, (self.n_modes * self.n_sets, self.n_modes))
+                for i in range(self.n_sets):
+                    block = ov2[i * self.n_modes:(i + 1) * self.n_modes, :]
+                    assert_allclose(block, ov4[:, :, i, 0], atol=1e-8)
+
+    def testRefModelSingleMode(self):
+        """A single signature mode reference exercises the reshape path."""
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                ov = calcSignatureOverlaps(me, ref_model=ref[0], diag=False)
+                self.assertEqual(ov.shape, (self.n_modes, 1, self.n_sets, 1))
+                ovc = calcSignatureOverlaps(me, ref_model=ref[0], diag=False,
+                                            collapse=True)
+                self.assertEqual(ovc.shape, (self.n_modes * self.n_sets, 1))
+                for i in range(self.n_sets):
+                    block = ovc[i * self.n_modes:(i + 1) * self.n_modes, :]
+                    assert_allclose(block, ov[:, :, i, 0], atol=1e-8)
+
+    def testRefModelDiag(self):
+        """diag against the equal-mode signature model: (n_modes, n_sets, 1)."""
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                ov = calcSignatureOverlaps(me, ref_model=ref, diag=True)
+                self.assertEqual(ov.shape, (self.n_modes, self.n_sets, 1))
+
+    def testSlicedEnsembleTutorialStyle(self):
+        """Mode-sliced ensembles (gnms[:, :5]) overlap against a sliced ref."""
+        k = 5
+        for model, (me, ref) in self.mes.items():
+            with self.subTest(model=model):
+                ov = calcSignatureOverlaps(me[:, :k], ref_model=ref[:k],
+                                           diag=False)
+                self.assertEqual(ov.shape, (k, k, self.n_sets, 1))