Add torch option for linalg for NMA and transform

prody/dynamics/pca.py

@@ -269,8 +269,18 @@ def performSVD(self, coordsets):
         dof = n_atoms * 3
         deviations = deviations.reshape((n_confs, dof)).T
 
-        vectors, values, self._temp = linalg.svd(deviations,
-                                                 full_matrices=False)
+        if linalg.__package__.startswith('torch'):
+            import torch
+            dev = torch.from_numpy(np.asarray(deviations))
+            if torch.cuda.is_available():
+                dev = dev.cuda()
+            vectors, values, self._temp = linalg.svd(dev, full_matrices=False)
+            vectors = vectors.detach().cpu().numpy()
+            values = values.detach().cpu().numpy()
+            self._temp = self._temp.detach().cpu().numpy()
+        else:
+            vectors, values, self._temp = linalg.svd(deviations,
+                                                     full_matrices=False)
         values = (values ** 2) / n_confs
         self._dof = dof
         self._n_atoms = n_atoms

prody/ensemble/ensemble.py

@@ -640,6 +640,10 @@ def _superpose(self, **kwargs):
                 matrix = dot((tar_org * weights).T,
                              (mob_org * weights)) / weights_dot
 
+            if linalg.__package__.startswith('torch'):
+                import torch
+                matrix = torch.from_numpy(matrix)
+
             U, s, Vh = svd(matrix)
             Id = array([[1, 0, 0], [0, 1, 0], [0, 0, sign(det(matrix))]])
             rotation = dot(Vh.T, dot(Id, U.T))

prody/measure/transform.py

@@ -152,8 +152,16 @@ def getTransformation(mob, tar, weights=None):
         tar = tar - tar_com
         matrix = np.dot((mob * weights).T, (tar * weights)) / weights_dot
 
+    if linalg.__package__.startswith('torch'):
+        import torch
+        matrix = torch.from_numpy(matrix)
+
     U, _, Vh = linalg.svd(matrix)
-    d = np.sign(linalg.det(np.dot(U, Vh)))
+    dot_U_Vh = np.dot(U, Vh)
+    if linalg.__package__.startswith('torch'):
+        dot_U_Vh = torch.from_numpy(dot_U_Vh)
+
+    d = np.sign(linalg.det(dot_U_Vh))
     Id = np.array([[1, 0, 0],
                    [0, 1, 0],
                    [0, 0, d]])

prody/utilities/eigtools.py

@@ -33,7 +33,7 @@ def solveEig(M, n_modes=None, zeros=False, turbo=True, expct_n_zeros=None, rever
             else:
                 eigvals = (0, n_modes+expct_n_zeros-1)
 
-    def _eigh(M, eigvals=None, turbo=True):
+    def _eigh(M, eigvals=None, turbo=True, n_modes=n_modes):
         if linalg.__package__.startswith('scipy'):
             from scipy.sparse import issparse
 
@@ -70,11 +70,25 @@ def _eigh(M, eigvals=None, turbo=True):
                 values = values[j:k]
                 vectors = vectors[:, j:k]
         else:
-            if n_modes is not None:
-                LOGGER.info('Scipy is not found, all modes were calculated.')
+            if linalg.__package__.startswith('torch'):
+                import torch
+                Mt = torch.from_numpy(np.asarray(M))
+                if torch.cuda.is_available():
+                    Mt = Mt.cuda()
+                values, vectors = linalg.eigh(Mt)
+                values = values.detach().cpu().numpy()
+                vectors = vectors.detach().cpu().numpy()
             else:
-                n_modes = dof
-            values, vectors = linalg.eigh(M)
+                values, vectors = linalg.eigh(M)
+
+            # numpy/torch eigh return the FULL ascending spectrum, whereas the scipy branch returns
+            # only the requested `eigvals` index range. Apply that same subset here, otherwise the
+            # caller's downstream slicing selects the wrong (smallest) modes.
+            if eigvals is not None:
+                lo, hi = int(eigvals[0]), int(eigvals[1])
+                values = values[lo:hi + 1]
+                vectors = vectors[:, lo:hi + 1]
+
         return values, vectors
 
     def _calc_n_zero_modes(M):

prody/utilities/misctools.py

@@ -15,7 +15,8 @@
 
 from xml.etree.ElementTree import Element
 
-__all__ = ['Everything', 'Cursor', 'ImageCursor', 'rangeString', 'alnum', 'importLA', 'dictElement',
+__all__ = ['Everything', 'Cursor', 'ImageCursor', 'rangeString', 'alnum', 'importLA',
+           'setLinalgBackend', 'getLinalgBackend', 'dictElement',
            'intorfloat', 'startswith', 'showFigure', 'countBytes', 'sqrtm',
            'saxsWater', 'count', 'addEnds', 'copy', 'dictElementLoop', 'index',
            'getDataPath', 'openData', 'chr2', 'toChararray', 'interpY', 'cmp', 'pystr',
@@ -224,17 +225,59 @@ def alnum(string, alt='_', trim=False, single=False):
     return result
 
 
+LINALG_BACKEND = None  # None -> auto (scipy > numpy > torch); or 'scipy'/'numpy'/'torch' to force
+
+
+def setLinalgBackend(backend):
+    """Force the linear-algebra backend used for NMA/PCA eigen/SVD calculations.
+
+    :arg backend: one of ``'scipy'``, ``'numpy'``, ``'torch'``, or **None** for automatic
+        selection (prefer scipy, then numpy, then torch). ``'torch'`` selects the experimental
+        GPU-capable :mod:`torch.linalg` path (uses CUDA when available)."""
+
+    global LINALG_BACKEND
+    if backend not in (None, 'scipy', 'numpy', 'torch'):
+        raise ValueError("backend must be None, 'scipy', 'numpy', or 'torch'")
+    LINALG_BACKEND = backend
+
+
+def getLinalgBackend():
+    """Returns the forced linear-algebra backend, or **None** if selection is automatic."""
+
+    return LINALG_BACKEND
+
+
 def importLA():
-    """Returns one of :mod:`scipy.linalg` or :mod:`numpy.linalg`."""
+    """Returns one of :mod:`scipy.linalg`, :mod:`numpy.linalg`, or :mod:`torch.linalg`.
+
+    Selection honours :func:`setLinalgBackend` if set; otherwise it prefers :mod:`scipy.linalg`
+    (the reference path, with subset eigensolvers), then :mod:`numpy.linalg`, and finally
+    :mod:`torch.linalg` (experimental GPU path). torch is NOT auto-preferred because its ``eigh``/
+    ``svd`` return the full spectrum / require tensors -- callers handle that, but scipy is faster
+    and better-tested for typical sizes; force torch with ``setLinalgBackend('torch')`` for GPU."""
+
+    backend = LINALG_BACKEND
+    if backend == 'torch':
+        import torch.linalg as linalg
+        return linalg
+    if backend == 'numpy':
+        import numpy.linalg as linalg
+        return linalg
+    if backend == 'scipy':
+        import scipy.linalg as linalg
+        return linalg
 
     try:
         import scipy.linalg as linalg
     except ImportError:
         try:
             import numpy.linalg as linalg
-        except:
-            raise ImportError('scipy.linalg or numpy.linalg is required for '
-                              'NMA and structure alignment calculations')
+        except ImportError:
+            try:
+                import torch.linalg as linalg
+            except ImportError:
+                raise ImportError('scipy.linalg or numpy.linalg or torch.linalg is required for '
+                                  'NMA and structure alignment calculations')
     return linalg
 
 def createStringIO():