Feature/update examples

conftest.py

@@ -161,7 +161,7 @@ def add_X_msd_no_input(doctest_namespace):
         damping=0.6,
     )
     # Specify initial conditions
-    x0 = msd.x0(np.array([1, -1]))
+    x0 = np.array([1, -1])
     # Simulate ODE
     t, x = msd.simulate(
         t_range,
@@ -201,7 +201,7 @@ def u(t):
         return 0.1 * np.sin(t)
 
     # Specify initial conditions
-    x0 = msd.x0(np.array([0, 0]))
+    x0 = np.array([0, 0])
     # Simulate ODE
     t, x = msd.simulate(t_range, t_step, x0, u, rtol=1e-8, atol=1e-8)
     # Format the data
@@ -235,7 +235,7 @@ def u(t):
         return 0.1 * np.sin(t)
 
     # Specify initial conditions
-    x0 = pend.x0(np.array([0, 0]))
+    x0 = np.array([0, 0])
     # Simulate ODE
     t, x = pend.simulate(t_range, t_step, x0, u, rtol=1e-8, atol=1e-8)
     # Format the data

doc/pykoop.rst

@@ -247,13 +247,13 @@ Examples
 Simple Koopman pipeline
 -----------------------
 
-.. include:: ../examples/example_pipeline_simple.py
-   :literal:
+.. plot:: ../examples/1_example_pipeline_simple.py
+   :include-source:
 
 Van der Pol Oscillator
 -----------------------
 
-.. plot:: ../examples/example_pipeline_vdp.py
+.. plot:: ../examples/2_example_pipeline_vdp.py
    :include-source:
 
 Cross-validation with ``scikit-learn``
@@ -269,8 +269,19 @@ necessary.
 Regressor parameters and lifting functions can easily be cross-validated using
 ``scikit-learn``:
 
-.. include:: ../examples/example_pipeline_cv.py
-   :literal:
+.. plot:: ../examples/3_example_pipeline_cv.py
+   :include-source:
+
+Asymptotic stability constraint
+-------------------------------
+
+In this example, three experimental EDMD-based regressors are compared to EDMD.
+Specifically, EDMD is compared to the asymptotic stability constraint and the
+H-infinity norm regularizer from [DF22]_ and [DF21]_, and the dissipativity
+constraint from [HIS19]_.
+
+.. plot:: ../examples/4_example_eigenvalue_comparison.py
+   :include-source:
 
 
 References

examples/1_example_pipeline_simple.py

@@ -0,0 +1,67 @@
+"""Example of how to use the Koopman pipeline."""
+
+from matplotlib import pyplot as plt
+from sklearn.preprocessing import MaxAbsScaler, StandardScaler
+
+import pykoop
+
+plt.rc('lines', linewidth=2)
+plt.rc('axes', grid=True)
+plt.rc('grid', linestyle='--')
+
+
+def example_pipeline_simple() -> None:
+    """Demonstrate how to use the Koopman pipeline."""
+    # Get example mass-spring-damper data
+    eg = pykoop.example_data_msd()
+
+    # Create pipeline
+    kp = pykoop.KoopmanPipeline(
+        lifting_functions=[
+            ('ma', pykoop.SkLearnLiftingFn(MaxAbsScaler())),
+            ('pl', pykoop.PolynomialLiftingFn(order=2)),
+            ('ss', pykoop.SkLearnLiftingFn(StandardScaler())),
+        ],
+        regressor=pykoop.Edmd(alpha=1),
+    )
+
+    # Fit the pipeline
+    kp.fit(
+        eg['X_train'],
+        n_inputs=eg['n_inputs'],
+        episode_feature=eg['episode_feature'],
+    )
+
+    # Predict using the pipeline
+    X_pred = kp.predict_trajectory(eg['x0_valid'], eg['u_valid'])
+
+    # Score using the pipeline
+    score = kp.score(eg['X_valid'])
+
+    # Plot results
+    fig, ax = plt.subplots(
+        kp.n_states_in_ + kp.n_inputs_in_,
+        1,
+        constrained_layout=True,
+        sharex=True,
+        figsize=(6, 6),
+    )
+    # Plot true trajectory
+    ax[0].plot(eg['t'], eg['X_valid'][:, 1], label='True trajectory')
+    ax[1].plot(eg['t'], eg['X_valid'][:, 2])
+    ax[2].plot(eg['t'], eg['X_valid'][:, 3])
+    # Plot predicted trajectory
+    ax[0].plot(eg['t'], X_pred[:, 1], '--', label='Predicted trajectory')
+    ax[1].plot(eg['t'], X_pred[:, 2], '--')
+    # Add labels
+    ax[-1].set_xlabel('$t$')
+    ax[0].set_ylabel('$x(t)$')
+    ax[1].set_ylabel(r'$\dot{x}(t)$')
+    ax[2].set_ylabel('$u$')
+    ax[0].set_title(f'True and predicted states; MSE={-1 * score:.2e}')
+    ax[0].legend(loc='upper right')
+
+
+if __name__ == '__main__':
+    example_pipeline_simple()
+    plt.show()

examples/example_pipeline_vdp.py → examples/2_example_pipeline_vdp.py

@@ -6,11 +6,15 @@
 import pykoop
 import pykoop.dynamic_models
 
+plt.rc('lines', linewidth=2)
+plt.rc('axes', grid=True)
+plt.rc('grid', linestyle='--')
+
 
 def example_pipeline_vdp() -> None:
     """Demonstrate how to use the Koopman pipeline."""
-    # Get sample data
-    X_vdp = pykoop.example_data_vdp()
+    # Get example Van der Pol data
+    eg = pykoop.example_data_vdp()
 
     # Create pipeline
     kp = pykoop.KoopmanPipeline(
@@ -26,82 +30,70 @@ def example_pipeline_vdp() -> None:
         regressor=pykoop.Edmd(),
     )
 
-    # Take last episode for validation
-    X_train = X_vdp[X_vdp[:, 0] < 4]
-    X_valid = X_vdp[X_vdp[:, 0] == 4]
-
     # Fit the pipeline
-    kp.fit(X_train, n_inputs=1, episode_feature=True)
-
-    # Extract initial conditions and input from validation episode
-    x0 = X_valid[[0], 1:3]
-    u = X_valid[:, 3:]
+    kp.fit(
+        eg['X_train'],
+        n_inputs=eg['n_inputs'],
+        episode_feature=eg['episode_feature'],
+    )
 
     # Predict with re-lifting between timesteps (default)
     X_pred_local = kp.predict_trajectory(
-        x0,
-        u,
+        eg['x0_valid'],
+        eg['u_valid'],
         relift_state=True,
-        episode_feature=False,
     )
 
     # Predict without re-lifting between timesteps
     X_pred_global = kp.predict_trajectory(
-        x0,
-        u,
+        eg['x0_valid'],
+        eg['u_valid'],
         relift_state=False,
-        episode_feature=False,
     )
 
     # Plot trajectories in phase space
-    fig, ax = plt.subplots(
-        constrained_layout=True,
-        figsize=(6, 6),
-    )
+    fig, ax = plt.subplots(constrained_layout=True, figsize=(6, 6))
     ax.plot(
-        X_valid[:, 1],
-        X_valid[:, 2],
+        eg['X_valid'][:, 1],
+        eg['X_valid'][:, 2],
         label='True trajectory',
     )
     ax.plot(
-        X_pred_local[:, 0],
         X_pred_local[:, 1],
+        X_pred_local[:, 2],
         '--',
         label='Local prediction',
     )
     ax.plot(
-        X_pred_global[:, 0],
         X_pred_global[:, 1],
+        X_pred_global[:, 2],
         '--',
         label='Global prediction',
     )
     ax.set_xlabel('$x_1[k]$')
     ax.set_ylabel('$x_2[k]$')
     ax.legend()
-    ax.grid(linestyle='--')
     ax.set_title('True and predicted phase-space trajectories')
 
     # Lift validation set
-    Psi_valid = kp.lift(X_valid[:, 1:], episode_feature=False)
+    Psi_valid = kp.lift(eg['X_valid'])
 
     # Predict lifted state with re-lifting between timesteps (default)
     Psi_pred_local = kp.predict_trajectory(
-        x0,
-        u,
+        eg['x0_valid'],
+        eg['u_valid'],
         relift_state=True,
         return_lifted=True,
         return_input=True,
-        episode_feature=False,
     )
 
     # Predict lifted state without re-lifting between timesteps
     Psi_pred_global = kp.predict_trajectory(
-        x0,
-        u,
+        eg['x0_valid'],
+        eg['u_valid'],
         relift_state=False,
         return_lifted=True,
         return_input=True,
-        episode_feature=False,
     )
 
     fig, ax = plt.subplots(
@@ -113,34 +105,27 @@ def example_pipeline_vdp() -> None:
         figsize=(6, 12),
     )
     for i in range(ax.shape[0]):
-        ax[i, 0].plot(Psi_valid[:, i], label='True trajectory')
-        ax[i, 0].plot(Psi_pred_local[:, i], '--', label='Local prediction')
-        ax[i, 0].plot(Psi_pred_global[:, i], '--', label='Global prediction')
-        ax[i, 0].grid(linestyle='--')
+        ax[i, 0].plot(
+            Psi_valid[:, i + 1],
+            label='True trajectory',
+        )
+        ax[i, 0].plot(
+            Psi_pred_local[:, i + 1],
+            '--',
+            label='Local prediction',
+        )
+        ax[i, 0].plot(
+            Psi_pred_global[:, i + 1],
+            '--',
+            label='Global prediction',
+        )
         ax[i, 0].set_ylabel(rf'$\vartheta_{i + 1}$')
 
     ax[-1, 0].set_xlabel('$k$')
     ax[0, 0].set_title('True and predicted lifted states')
     ax[-1, -1].legend(loc='lower right')
     fig.align_ylabels()
 
-    fig, ax = plt.subplots(
-        kp.n_inputs_out_,
-        1,
-        constrained_layout=True,
-        sharex=True,
-        squeeze=False,
-        figsize=(6, 6),
-    )
-    for i in range(ax.shape[0]):
-        j = kp.n_states_out_ + i
-        ax[i, 0].plot(Psi_valid[:, j], label='True trajectory')
-        ax[i, 0].grid(linestyle='--')
-
-    ax[-1, 0].set_xlabel('$k$')
-    ax[0, 0].set_ylabel('$u$')
-    ax[0, 0].set_title('Exogenous input')
-
 
 if __name__ == '__main__':
     example_pipeline_vdp()