Fixing redundant SVD computation

Author: greinerth

pydmd/utils.py

@@ -1,12 +1,51 @@
 """Utilities module."""
 
 import warnings
+from typing import Union
 
 import numpy as np
 from numpy.lib.stride_tricks import sliding_window_view
 
 
-def compute_rank(X, svd_rank=0):
+def __svht(sigma_svd: np.ndarray, rows: int, cols: int) -> int:
+    """
+    Singular Value Hard Threshold.
+
+    :param sigma_svd: Singual values computed by SVD
+    :type sigma_svd: np.ndarray
+    :param rows: Number of rows of original data matrix.
+    :type rows: int
+    :param cols: Number of columns of original data matrix.
+    :type cols: int
+    :return: Computed rank.
+    :rtype: int
+
+    References:
+    Gavish, Matan, and David L. Donoho, The optimal hard threshold for
+    singular values is, IEEE Transactions on Information Theory 60.8
+    (2014): 5040-5053.
+    https://ieeexplore.ieee.org/document/6846297
+    """
+    beta = np.divide(*sorted((rows, cols)))
+    beta_square = beta * beta
+    omega = 0.56 * beta_square * beta - 0.95 * beta_square + 1.82 * beta + 1.43
+    tau = np.median(sigma_svd) * omega
+    rank = np.sum(sigma_svd > tau)
+
+    if rank == 0:
+        warnings.warn(
+            "SVD optimal rank is 0. The largest singular values are "
+            "indistinguishable from noise. Setting rank truncation to 1.",
+            RuntimeWarning,
+        )
+        rank = 1
+
+    return rank
+
+
+def __compute_rank(
+    sigma_svd: np.ndarray, rows: int, cols: int, svd_rank: Union[float, int]
+) -> int:
     """
     Rank computation for the truncated Singular Value Decomposition.
     :param numpy.ndarray X: the matrix to decompose.
@@ -24,33 +63,47 @@ def compute_rank(X, svd_rank=0):
     singular values is, IEEE Transactions on Information Theory 60.8
     (2014): 5040-5053.
     """
-    U, s, _ = np.linalg.svd(X, full_matrices=False)
-
-    def omega(x):
-        return 0.56 * x**3 - 0.95 * x**2 + 1.82 * x + 1.43
-
     if svd_rank == 0:
-        beta = np.divide(*sorted(X.shape))
-        tau = np.median(s) * omega(beta)
-        rank = np.sum(s > tau)
-        if rank == 0:
-            warnings.warn(
-                "SVD optimal rank is 0. The largest singular values are "
-                "indistinguishable from noise. Setting rank truncation to 1.",
-                RuntimeWarning,
-            )
-            rank = 1
+        rank = __svht(sigma_svd, rows, cols)
+
     elif 0 < svd_rank < 1:
-        cumulative_energy = np.cumsum(s**2 / (s**2).sum())
+        sigma_svd_squared = np.square(sigma_svd)
+        cumulative_energy = np.cumsum(
+            sigma_svd_squared / sigma_svd_squared.sum()
+        )
         rank = np.searchsorted(cumulative_energy, svd_rank) + 1
+
     elif svd_rank >= 1 and isinstance(svd_rank, int):
-        rank = min(svd_rank, U.shape[1])
+        rank = min(svd_rank, sigma_svd.size)
+
     else:
-        rank = min(X.shape)
+        rank = min(rows, cols)
 
     return rank
 
 
+def compute_rank(X: np.ndarray, svd_rank=0):
+    """
+    Rank computation for the truncated Singular Value Decomposition.
+    :param numpy.ndarray X: the matrix to decompose.
+    :param svd_rank: the rank for the truncation; If 0, the method computes
+        the optimal rank and uses it for truncation; if positive interger,
+        the method uses the argument for the truncation; if float between 0
+        and 1, the rank is the number of the biggest singular values that
+        are needed to reach the 'energy' specified by `svd_rank`; if -1,
+        the method does not compute truncation. Default is 0.
+    :type svd_rank: int or float
+    :return: the computed rank truncation.
+    :rtype: int
+    References:
+    Gavish, Matan, and David L. Donoho, The optimal hard threshold for
+    singular values is, IEEE Transactions on Information Theory 60.8
+    (2014): 5040-5053.
+    """
+    _, s, _ = np.linalg.svd(X, full_matrices=False)
+    return __compute_rank(s, X.shape[0], X.shape[1], svd_rank)
+
+
 def compute_tlsq(X, Y, tlsq_rank):
     """
     Compute Total Least Square.
@@ -100,8 +153,8 @@ def compute_svd(X, svd_rank=0):
     singular values is, IEEE Transactions on Information Theory 60.8
     (2014): 5040-5053.
     """
-    rank = compute_rank(X, svd_rank)
     U, s, V = np.linalg.svd(X, full_matrices=False)
+    rank = __compute_rank(s, X.shape[0], X.shape[1], svd_rank)
     V = V.conj().T
 
     U = U[:, :rank]