Massive test update

Author: sichinaga

pydmd/havok.py

@@ -300,7 +300,7 @@ def hankel(self, X):
             )
 
         Hm = m - ((self._delays - 1) * self._lag)
-        H = np.empty((n * self._delays, Hm))
+        H = np.empty((n * self._delays, Hm), dtype="complex")
         for i in range(self._delays):
             H[i * n : (i + 1) * n] = X[:, i * self._lag : i * self._lag + Hm]
 
@@ -321,7 +321,7 @@ def dehankel(self, H):
         Hn, Hm = H.shape
         n = int(Hn / self._delays)
         m = int(Hm + ((self._delays - 1) * self._lag))
-        X = np.empty((n, m))
+        X = np.empty((n, m), dtype="complex")
         for i in range(self._delays):
             X[:, i * self._lag : i * self._lag + Hm] = H[i * n : (i + 1) * n]
         return X

tests/test_havok.py

@@ -1,9 +1,10 @@
 import numpy as np
 from pytest import raises
 from scipy.integrate import solve_ivp
-
 from numpy.testing import assert_equal
 
+from pydmd import DMD
+from pydmd import PiDMD
 from pydmd import HAVOK
 
 
@@ -12,8 +13,9 @@ def generate_lorenz_data(t_eval):
     Given a time vector t_eval = t1, t2, ..., evaluates and returns
     the snapshots of the Lorenz system as columns of the matrix X.
     """
+
     def lorenz_system(t, state):
-        sigma, rho, beta = 10, 28, 8/3 # chaotic parameters
+        sigma, rho, beta = 10, 28, 8 / 3  # chaotic parameters
         x, y, z = state
         x_dot = sigma * (y - x)
         y_dot = (x * (rho - z)) - y
@@ -38,7 +40,7 @@ def lorenz_system(t, state):
 
 
 # Generate Lorenz system data.
-dt = 0.001 # time step
+dt = 0.001  # time step
 m = 50000  # number of data samples
 t = np.arange(m) * dt
 X = generate_lorenz_data(t)
@@ -47,17 +49,11 @@ def lorenz_system(t, state):
 
 def test_error_fitted():
     """
-    Ensure that attempting to get attributes results
+    Ensure that attempting to get HAVOK attributes results
     in an error if fit() has not yet been called.
     """
     havok = HAVOK()
 
-    with raises(ValueError):
-        _ = havok.snapshots
-    with raises(ValueError):
-        _ = havok.ho_snapshots
-    with raises(ValueError):
-        _ = havok.time
     with raises(ValueError):
         _ = havok.modes
     with raises(ValueError):
@@ -82,12 +78,20 @@ def test_error_fitted():
 
 def test_hankel_1():
     """
-    Test
+    Test that the hankel and dehankel functions work as intended.
+    Use 1-D data, lag = 1, and various delay values.
     """
     dummy_data = np.array([[1, 2, 3, 4]])
 
     havok = HAVOK(delays=1)
-    assert_equal(havok.hankel(dummy_data), np.array([[1, 2, 3, 4],]))
+    assert_equal(
+        havok.hankel(dummy_data),
+        np.array(
+            [
+                [1, 2, 3, 4],
+            ]
+        ),
+    )
     assert_equal(havok.dehankel(havok.hankel(dummy_data)), dummy_data)
 
     havok = HAVOK(delays=2)
@@ -99,7 +103,25 @@ def test_hankel_1():
     assert_equal(havok.dehankel(havok.hankel(dummy_data)), dummy_data)
 
     havok = HAVOK(delays=4)
-    assert_equal(havok.hankel(dummy_data), np.array([[1,], [2,], [3,], [4,]]))
+    assert_equal(
+        havok.hankel(dummy_data),
+        np.array(
+            [
+                [
+                    1,
+                ],
+                [
+                    2,
+                ],
+                [
+                    3,
+                ],
+                [
+                    4,
+                ],
+            ]
+        ),
+    )
     assert_equal(havok.dehankel(havok.hankel(dummy_data)), dummy_data)
 
     with raises(ValueError):
@@ -109,7 +131,8 @@ def test_hankel_1():
 
 def test_hankel_2():
     """
-    Test
+    Test that the hankel and dehankel functions work as intended.
+    Use 2-D data, lag = 1, and various delay values.
     """
     dummy_data = np.array([[1, 2, 3], [4, 5, 6]])
     H2 = np.array([[1, 2], [4, 5], [2, 3], [5, 6]])
@@ -134,7 +157,8 @@ def test_hankel_2():
 
 def test_hankel_3():
     """
-    Test
+    Test that the hankel and dehankel functions work as intended.
+    Use 1-D data, lag = 2, and various delay values.
     """
     dummy_data = np.array([[1, 2, 3, 4, 5, 6]])
     H2 = np.array([[1, 2, 3, 4], [3, 4, 5, 6]])
@@ -149,12 +173,12 @@ def test_hankel_3():
     assert_equal(havok.hankel(dummy_data), H2)
     assert_equal(havok.dehankel(havok.hankel(dummy_data)), dummy_data)
 
-    havok = HAVOK(delays=3, lag=3)
+    havok = HAVOK(delays=3, lag=2)
     assert_equal(havok.hankel(dummy_data), H3)
     assert_equal(havok.dehankel(havok.hankel(dummy_data)), dummy_data)
 
     with raises(ValueError):
-        havok = HAVOK(delays=2, lag=2)
+        havok = HAVOK(delays=4, lag=2)
         havok.hankel(dummy_data)
 
 
@@ -186,43 +210,215 @@ def test_shape_2():
     assert havok.B.shape == (havok.r - 2, 2)
 
 
-def test_r():
+def test_snapshots_1():
     """
-    Ensure the accuracy of the r property in the following situations:
+    Test the stored snapshots and ho_snapshots.
+    Ensure that they are accurate in the default case with 1-D data.
     """
-    # If no svd truncation, r is the min of the dimensions of the hankel matrix
-    havok = HAVOK(svd_rank=-1)
+    havok = HAVOK()
+    with raises(ValueError):
+        _ = havok.snapshots
+    with raises(ValueError):
+        _ = havok.ho_snapshots
     havok.fit(x, t)
-    assert havok.r == min(havok.delays, len(t) - havok.delays + 1)
+    assert_equal(havok.snapshots, x)
+    assert_equal(havok.ho_snapshots, havok.hankel(x))
+
+
+def test_snapshots_2():
+    """
+    Test the stored snapshots and ho_snapshots.
+    Ensure that they are accurate in the default case with 2-D data.
+    """
+    havok = HAVOK().fit(X, t)
+    assert_equal(havok.snapshots, X)
+    assert_equal(havok.ho_snapshots, havok.hankel(X))
+
 
-    # Test the above case, but for a larger d value
-    havok = HAVOK(svd_rank=-1, delays=500)
+def test_time_1():
+    """
+    Test the stored time attribute.
+    Ensure that it is accurate in the default case.
+    """
+    havok = HAVOK()
+    with raises(ValueError):
+        _ = havok.time
     havok.fit(x, t)
-    assert havok.r == min(havok.delays, len(t) - havok.delays + 1)
+    assert_equal(havok.time, t)
+
+
+def test_time_2():
+    """
+    Test that a HAVOK model fitted with a time vector is essentially
+    the same as a HAVOK model fitted with the time-step dt. Check the
+    stored time vector and the computed HAVOK operator.
+    """
+    havok_1 = HAVOK().fit(x, t)
+    havok_2 = HAVOK().fit(x, dt)
+    assert_equal(havok_1.time, havok_2.time)
+    assert_equal(havok_1.operator, havok_2.operator)
+
 
-    # Test the above case, but for an even larger d value
-    havok = HAVOK(svd_rank=-1, delays=len(t) - 20)
+def test_plot_summary_1():
+    """
+    Test that everything is fine if we fit a HAVOK
+    model and ask for the default summary plot.
+    """
+    havok = HAVOK(svd_rank=16, delays=100)
     havok.fit(x, t)
-    assert havok.r == min(havok.delays, len(t) - havok.delays + 1)
+    havok.plot_summary()
 
 
-def test_reconstruction():
+def test_plot_summary_2():
     """
-    Test the accuracy of the HAVOK reconstruction.
+    Test that everything is fine if we fit a HAVOK model and
+    ask for various (but still valid) plot modifications.
     """
     havok = HAVOK(svd_rank=16, delays=100)
     havok.fit(x, t)
-    error = x - havok.reconstructed_data.real
-    assert np.linalg.norm(error) / np.linalg.norm(x) < 0.2
+    havok.plot_summary(
+        num_plot=15000,
+        index_linear=(0, 1),
+        forcing_threshold=0.005,
+        min_jump_dist=200,
+        figsize=(15, 4),
+        dpi=100,
+    )
+
+
+def test_plot_summary_3():
+    """
+    Test that an error is thrown if we ask for
+    a summary from an unfitted HAVOK model.
+    """
+    havok = HAVOK()
+    with raises(ValueError):
+        havok.plot_summary()
+
+
+def test_reconstruction_1():
+    """
+    Test the accuracy of the HAVOK reconstruction.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+    error = x - havok.reconstructed_data
+    assert np.linalg.norm(error) / np.linalg.norm(x) < 0.05
+
+
+def test_reconstruction_2():
+    """
+    Test the accuracy of the sHAVOK reconstruction.
+    """
+    havok = HAVOK(svd_rank=4, delays=100, structured=True).fit(x, t)
+    error = x[:-1] - havok.reconstructed_data
+    assert np.linalg.norm(error) / np.linalg.norm(x[:-1]) < 0.07
+
+
+def test_predict_1():
+    """
+    Test the accuracy of the HAVOK prediction.
+    Ensure that prediction with the HAVOK forcing term and the times
+    of fitting simply yields the computed data reconstruction.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+    assert_equal(
+        havok.predict(havok.forcing, havok.time),
+        havok.reconstructed_data,
+    )
+
+
+def test_predict_2():
+    """
+    Test the accuracy of the HAVOK prediction.
+    Test that predicting beyond the training set isn't absurdly inaccurate.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+
+    # Build a longer data set and fit a HAVOK model to it.
+    t_long = np.arange(2 * m) * dt
+    x_long = generate_lorenz_data(t_long)[0]
+    havok_long = HAVOK(svd_rank=16, delays=100).fit(x_long, t_long)
+
+    # We only use the long HAVOK model to obtain a long forcing signal.
+    forcing_long = havok_long.forcing
+    time_long = t_long[: len(forcing_long)]
 
+    # Get the error of the full prediction.
+    error = x_long - havok.predict(forcing_long, time_long)
+    assert np.linalg.norm(error) / np.linalg.norm(x_long) < 0.45
 
-def test_predict():
+
+def test_predict_3():
     """
     Test the accuracy of the HAVOK prediction.
+    Test that predicting with V0 indices is functionally
+    the same as predicting with array-valued V0 inputs.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+    assert_equal(
+        havok.predict(havok.forcing, havok.time, V0=0),
+        havok.predict(havok.forcing, havok.time, V0=havok.linear_dynamics[0]),
+    )
+    assert_equal(
+        havok.predict(havok.forcing, havok.time, V0=1),
+        havok.predict(havok.forcing, havok.time, V0=havok.linear_dynamics[1]),
+    )
+    assert_equal(
+        havok.predict(havok.forcing, havok.time, V0=-1),
+        havok.predict(havok.forcing, havok.time, V0=havok.linear_dynamics[-1]),
+    )
+
+
+def test_threshold_1():
     """
+    Test compute_threshold function.
+    Test that threshold computation works as expected, whether you plot or not.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+    thres_1 = havok.compute_threshold(p=0.1, bins=100, plot=False)
+    thres_2 = havok.compute_threshold(p=0.1, bins=100, plot=True)
+    assert thres_1 == thres_2
+
+
+def test_threshold_2():
+    """
+    Test compute_threshold function.
+    Test that threshold computation works as expected, whether you index
+    the stored forcing term, or provide a custom array-valued forcing term.
+    """
+    havok = HAVOK(svd_rank=16, delays=100).fit(x, t)
+    vr = havok.forcing.flatten()
+    thres_1 = havok.compute_threshold(p=0.1, bins=100)
+    thres_2 = havok.compute_threshold(forcing=vr, p=0.1, bins=100)
+    assert thres_1 == thres_2
 
 
-def test_time():
+def test_dmd_1():
     """
-    Test that 
+    Test that HAVOK works when used with an externally-defined DMD model.
+    Test the basic exact DMD model - ensure that plot_summary works just
+    fine and that the data reconstruction is still accurate.
     """
+    dmd = DMD(svd_rank=-1)
+    havok = HAVOK(svd_rank=16, delays=100, dmd=dmd).fit(x, t)
+    havok.plot_summary()
+    error = x - havok.reconstructed_data
+    assert np.linalg.norm(error) / np.linalg.norm(x) < 0.05
+
+
+def test_dmd_2():
+    """
+    Test that HAVOK works when used with an externally-defined DMD model.
+    Test a physics-informed DMD model - ensure that plot_summary works just
+    fine and that the data reconstruction is still accurate.
+    """
+    pidmd = PiDMD(
+        svd_rank=-1,
+        manifold="diagonal",
+        manifold_opt=2,
+        compute_A=True,
+    )
+    havok = HAVOK(svd_rank=4, delays=100, dmd=pidmd).fit(x, t)
+    havok.plot_summary()
+    error = x - havok.reconstructed_data
+    assert np.linalg.norm(error) / np.linalg.norm(x) < 0.1