Features/1563 Add Dynamic Mode Decomposition (DMD)

ausprobieren.py

@@ -0,0 +1,44 @@
+"""
+for debugging will be removed later
+"""
+
+import heat as ht
+import torch
+import heat.decomposition as htd
+
+r = 6
+A_red = ht.array(
+    [
+        [0.0, -1.0, 0.0, 0.0, 0.0, 0.0],
+        [1.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+        [0.0, 0.0, 1.5, 0.0, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 0.5, 0.0, 0.0],
+        [0.0, 0.0, 0.0, 0.0, -1.5, 0.0],
+        [0.0, 0.0, 0.0, 0.0, 0.0, -0.5],
+    ],
+    split=None,
+    dtype=ht.float32,
+)
+x0_red = ht.random.randn(r, 1, split=None)
+m, n = 25 * ht.MPI_WORLD.size, 15
+X = ht.hstack(
+    [(ht.array(torch.linalg.matrix_power(A_red.larray, i) @ x0_red.larray)) for i in range(n)]
+)
+U = ht.random.randn(m, r, split=0)
+U, _ = ht.linalg.qr(U)
+X = U @ X
+
+
+dmd = htd.DMD(svd_solver="full", svd_rank=r)
+dmd.fit(X)
+print(dmd.rom_basis_.shape, dmd.rom_basis_.split)
+dmd.rom_basis_.resplit_(None)
+
+# check prediction of next states
+Y = dmd.predict_next(X)
+print(ht.allclose(Y[:, : n - 1], X[:, 1:], atol=1e-4, rtol=1e-4))
+
+# check batch prediction
+Y = dmd.predict(X[:, 0].resplit_(None), n)
+print(Y.shape)
+print(X - Y)

heat/core/linalg/svdtools.py

@@ -11,7 +11,7 @@
 from ..dndarray import DNDarray
 from .. import factories
 from .. import types
-from ..linalg import matmul, vector_norm
+from ..linalg import matmul, vector_norm, qr, svd
 from ..indexing import where
 from ..random import randn
 
@@ -21,7 +21,21 @@
 from math import log, ceil, floor, sqrt
 
 
-__all__ = ["hsvd_rank", "hsvd_rtol", "hsvd"]
+__all__ = ["hsvd_rank", "hsvd_rtol", "hsvd", "rsvd"]
+
+
+def _check_SVD_input(A):
+    if not isinstance(A, DNDarray):
+        raise TypeError(f"Argument needs to be a DNDarray but is {type(A)}.")
+    if not A.ndim == 2:
+        raise ValueError("A needs to be a 2D matrix")
+    if not A.dtype == types.float32 and not A.dtype == types.float64:
+        raise TypeError(
+            "Argument needs to be a DNDarray with datatype float32 or float64, but data type is {}.".format(
+                A.dtype
+            )
+        )
+    return None
 
 
 #######################################################################################
@@ -85,16 +99,7 @@ def hsvd_rank(
         [1] Iwen, Ong. A distributed and incremental SVD algorithm for agglomerative data analysis on large networks. SIAM J. Matrix Anal. Appl., 37(4), 2016.
         [2] Himpe, Leibner, Rave. Hierarchical approximate proper orthogonal decomposition. SIAM J. Sci. Comput., 40 (5), 2018.
     """
-    if not isinstance(A, DNDarray):
-        raise TypeError(f"Argument needs to be a DNDarray but is {type(A)}.")
-    if not A.ndim == 2:
-        raise ValueError("A needs to be a 2D matrix")
-    if not A.dtype == types.float32 and not A.dtype == types.float64:
-        raise TypeError(
-            "Argument needs to be a DNDarray with datatype float32 or float64, but data type is {}.".format(
-                A.dtype
-            )
-        )
+    _check_SVD_input(A)  # check if A is suitable input
     A_local_size = max(A.lshape_map[:, 1])
 
     if maxmergedim is not None and maxmergedim < 2 * (maxrank + safetyshift) + 1:
@@ -197,16 +202,7 @@ def hsvd_rtol(
         [1] Iwen, Ong. A distributed and incremental SVD algorithm for agglomerative data analysis on large networks. SIAM J. Matrix Anal. Appl., 37(4), 2016.
         [2] Himpe, Leibner, Rave. Hierarchical approximate proper orthogonal decomposition. SIAM J. Sci. Comput., 40 (5), 2018.
     """
-    if not isinstance(A, DNDarray):
-        raise TypeError(f"Argument needs to be a DNDarray but is {type(A)}.")
-    if not A.ndim == 2:
-        raise ValueError("A needs to be a 2D matrix")
-    if not A.dtype == types.float32 and not A.dtype == types.float64:
-        raise TypeError(
-            "Argument needs to be a DNDarray with datatype float32 or float64, but data type is {}.".format(
-                A.dtype
-            )
-        )
+    _check_SVD_input(A)  # check if A is suitable input
     A_local_size = max(A.lshape_map[:, 1])
 
     if maxmergedim is not None and maxrank is None:
@@ -529,3 +525,91 @@ def compute_local_truncated_svd(
         sigma_loc = torch.zeros(1, dtype=U_loc.dtype, device=U_loc.device)
         U_loc = torch.zeros(U_loc.shape[0], 1, dtype=U_loc.dtype, device=U_loc.device)
         return U_loc, sigma_loc, err_squared_loc
+
+
+##############################################################################################
+# Randomized SVD
+##############################################################################################
+
+
+def rsvd(
+    A: DNDarray,
+    rank: int,
+    n_oversamples: int = 10,
+    power_iter: int = 0,
+    qr_procs_to_merge: int = 2,
+) -> Union[Tuple[DNDarray, DNDarray, DNDarray], Tuple[DNDarray, DNDarray]]:
+    """
+    Randomized SVD (rSVD) with prescribed truncation rank `rank`.
+    If A = U diag(sigma) V^T is the true SVD of A, this routine computes an approximation for U[:,:rank] (and sigma[:rank], V[:,:rank]).
+
+    The accuracy of this approximation depends on the structure of A ("low-rank" is best) and appropriate choice of parameters.
+
+    Parameters
+    ----------
+    A : DNDarray
+        2D-array (float32/64) of which the rSVD has to be computed.
+    rank : int
+        truncation rank. (This parameter corresponds to `n_components` in sci-kit learn's TruncatedSVD.)
+    n_oversamples : int, optional
+        number of oversamples. The default is 10.
+    power_iter : int, optional
+        number of power iterations. The default is 1.
+    qr_procs_to_merge : int, optional
+        number of processes to merge at each step of QR decomposition in the power iteration (if power_iter > 0). The default is 2. See the corresponding remarks for `heat.linalg.qr` for more details.
+
+    Notes
+    ------
+    Memory requirements: the SVD computation of a matrix of size (rank + n_oversamples) x (rank + n_oversamples) must fit into the memory of a single process.
+    The implementation follows Algorithm 4.4 (randomized range finder) and Algorithm 5.1 (direct SVD) in [1].
+
+    References
+    -----------
+    [1] Halko, N., Martinsson, P. G., & Tropp, J. A. (2011). Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions. SIAM review, 53(2), 217-288.
+    """
+    _check_SVD_input(A)  # check if A is suitable input
+    if not isinstance(rank, int):
+        raise TypeError(f"rank must be an integer, but is {type(rank)}.")
+    if rank < 1:
+        raise ValueError(f"rank must be positive, but is {rank}.")
+    if not isinstance(n_oversamples, int):
+        raise TypeError(
+            f"if provided, n_oversamples must be an integer, but is {type(n_oversamples)}."
+        )
+    if n_oversamples < 0:
+        raise ValueError(f"n_oversamples must be non-negative, but is {n_oversamples}.")
+    if not isinstance(power_iter, int):
+        raise TypeError(f"if provided, power_iter must be an integer, but is {type(power_iter)}.")
+    if power_iter < 0:
+        raise ValueError(f"power_iter must be non-negative, but is {power_iter}.")
+
+    ell = rank + n_oversamples
+    q = power_iter
+
+    # random matrix
+    splitOmega = 1 if A.split == 0 else 0
+    Omega = randn(A.shape[1], ell, dtype=A.dtype, device=A.device, split=splitOmega)
+
+    # compute the range of A
+    Y = matmul(A, Omega)
+    Q, _ = qr(Y, procs_to_merge=qr_procs_to_merge)
+
+    # power iterations
+    for _ in range(q):
+        Y = matmul(A.T, Q)
+        Q, _ = qr(Y, procs_to_merge=qr_procs_to_merge)
+        Y = matmul(A, Q)
+        Q, _ = qr(Y, procs_to_merge=qr_procs_to_merge)
+
+    # compute the SVD of the projected matrix
+    B = matmul(Q.T, A)
+    B.resplit_(
+        None
+    )  # B will be of size ell x ell and thus small enough to fit into memory of a single process
+    U, sigma, V = svd.svd(B)  # actually just torch svd as input is not split anymore
+    U = matmul(Q, U)[:, :rank]
+    U.balance_()
+    S = sigma[:rank]
+    V = V[:, :rank]
+    V.balance_()
+    return U, S, V

heat/core/linalg/tests/test_svdtools.py

@@ -193,3 +193,58 @@ def test_hsvd_rank_part2(self):
                 self.assertTrue(U_orth_err <= dtype_tol)
                 self.assertTrue(V_orth_err <= dtype_tol)
                 self.assertTrue(true_rel_err <= dtype_tol)
+
+
+class TestRSVD(TestCase):
+    def test_rsvd(self):
+        for dtype in [ht.float32, ht.float64]:
+            dtype_tol = 1e-4 if dtype == ht.float32 else 1e-10
+            for split in [0, 1, None]:
+                X = ht.random.randn(200, 200, dtype=dtype, split=split)
+                for rank in [ht.MPI_WORLD.size, 10]:
+                    for n_oversamples in [5, 10]:
+                        for power_iter in [0, 1, 2, 3]:
+                            U, S, V = ht.linalg.rsvd(
+                                X, rank, n_oversamples=n_oversamples, power_iter=power_iter
+                            )
+                            self.assertEqual(U.shape, (X.shape[0], rank))
+                            self.assertEqual(S.shape, (rank,))
+                            self.assertEqual(V.shape, (X.shape[1], rank))
+                            self.assertTrue(ht.all(S >= 0))
+                            self.assertTrue(
+                                ht.allclose(
+                                    U.T @ U,
+                                    ht.eye(rank, dtype=U.dtype, split=U.split),
+                                    rtol=dtype_tol,
+                                    atol=dtype_tol,
+                                )
+                            )
+                            self.assertTrue(
+                                ht.allclose(
+                                    V.T @ V,
+                                    ht.eye(rank, dtype=V.dtype, split=V.split),
+                                    rtol=dtype_tol,
+                                    atol=dtype_tol,
+                                )
+                            )
+
+    def test_rsvd_catch_wrong_inputs(self):
+        X = ht.random.randn(10, 10)
+        # wrong dtype for rank
+        with self.assertRaises(TypeError):
+            ht.linalg.rsvd(X, "a")
+        # rank zero
+        with self.assertRaises(ValueError):
+            ht.linalg.rsvd(X, 0)
+        # wrong dtype for n_oversamples
+        with self.assertRaises(TypeError):
+            ht.linalg.rsvd(X, 10, n_oversamples="a")
+        # n_oversamples negative
+        with self.assertRaises(ValueError):
+            ht.linalg.rsvd(X, 10, n_oversamples=-1)
+        # wrong dtype for power_iter
+        with self.assertRaises(TypeError):
+            ht.linalg.rsvd(X, 10, power_iter="a")
+        # power_iter negative
+        with self.assertRaises(ValueError):
+            ht.linalg.rsvd(X, 10, power_iter=-1)

heat/core/random.py

@@ -129,8 +129,8 @@ def __counter_sequence(
         c_0 = (__counter & (max_count << 64)) >> 64
     c_1 = __counter & max_count
     total_elements = torch.prod(torch.tensor(shape))
-    if total_elements.item() > 2 * max_count:
-        raise ValueError(f"Shape is to big with {total_elements} elements")
+    # if total_elements.item() > 2 * max_count:
+    #    raise ValueError(f"Shape is to big with {total_elements} elements")
 
     if split is None:
         values = total_elements.item() // 2 + total_elements.item() % 2
@@ -619,7 +619,6 @@ def randint(
             x_0, x_1 = __threefry32(x_0, x_1, seed=__seed)
         else:  # torch.int64
             x_0, x_1 = __threefry64(x_0, x_1, seed=__seed)
-
         # stack the resulting sequence and normalize to given range
         values = torch.stack([x_0, x_1], dim=1).flatten()[lslice].reshape(lshape)
         # ATTENTION: this is biased and known, bias-free rejection sampling is difficult to do in parallel

heat/core/tests/test_random.py

@@ -616,7 +616,7 @@ def test_rand(self):
         self.assertTrue(ht.equal(a, b))
 
         # Too big arrays cant be created
-        with self.assertRaises(ValueError):
+        with self.assertRaises(RuntimeError):
             ht.random.randn(0x7FFFFFFFFFFFFFFF)
         with self.assertRaises(ValueError):
             ht.random.rand(3, 2, -2, 5, split=1)

heat/decomposition/__init__.py

@@ -3,3 +3,4 @@
 """
 
 from .pca import *
+from .dmd import *