Add control example

examples/8_control_vdp.py

@@ -0,0 +1,110 @@
+"""Example of how to use the Koopman operator for control."""
+
+import control
+import numpy as np
+from matplotlib import pyplot as plt
+
+import pykoop
+
+plt.rc('lines', linewidth=2)
+plt.rc('axes', grid=True)
+plt.rc('grid', linestyle='--')
+
+
+def example_control_vdp() -> None:
+    """Demonstrate how to use the Koopman operator for control."""
+    # Get example Van der Pol data
+    eg = pykoop.example_data_vdp()
+
+    # Create and fit linear Koopman pipeline
+    kp_lin = pykoop.KoopmanPipeline(
+        lifting_functions=None,
+        regressor=pykoop.Edmd(),
+    ).fit(
+        eg['X_train'],
+        n_inputs=eg['n_inputs'],
+        episode_feature=eg['episode_feature'],
+    )
+
+    # Create and fit polynomial Koopman pipeline
+    kp_poly = pykoop.KoopmanPipeline(
+        lifting_functions=[(
+            'sp',
+            pykoop.SplitPipeline(
+                lifting_functions_state=[
+                    ('pl', pykoop.PolynomialLiftingFn(order=3))
+                ],
+                lifting_functions_input=None,
+            ),
+        )],
+        regressor=pykoop.Edmd(),
+    ).fit(
+        eg['X_train'],
+        n_inputs=eg['n_inputs'],
+        episode_feature=eg['episode_feature'],
+    )
+
+    # Extract state-space matrices
+    U_lin = kp_lin.regressor_.coef_.T
+    A_lin = U_lin[:, :U_lin.shape[0]]
+    B_lin = U_lin[:, U_lin.shape[0]:]
+    U_poly = kp_poly.regressor_.coef_.T
+    A_poly = U_poly[:, :U_poly.shape[0]]
+    B_poly = U_poly[:, U_poly.shape[0]:]
+
+    # Set weighting matrices for LQR. Lifted state errors are not weighted
+    Q_lin = np.eye(2)
+    R_lin = np.eye(1)
+    Q_poly = np.diag([1, 1, 0, 0, 0, 0, 0, 0, 0])
+    R_poly = np.eye(1)
+
+    # Synthesize discrete LQR controllers using the ``control`` package
+    K_lin, _, _ = control.dlqr(A_lin, B_lin, Q_lin, R_lin, method='scipy')
+    K_poly, _, _ = control.dlqr(A_poly, B_poly, Q_poly, R_poly, method='scipy')
+
+    # Set number of timesteps to predict
+    N = 1000
+    # Get dynamic model
+    vdp = eg['dynamic_model']
+
+    # Predict ground truth dynamics without control
+    Xc_gt = np.zeros((N, 2))
+    Xc_gt[0, :] = np.array([1, 1])
+    for k in range(1, N):
+        Xc_gt[k, :] = vdp.f(0, Xc_gt[k - 1, :], 0)
+
+    # Predict dynamics with linear controller
+    Xc_lin = np.zeros((N, 2))
+    Uc_lin = np.zeros((N, 1))
+    Xc_lin[0, :] = np.array([1, 1])
+    for k in range(1, N):
+        Uc_lin[[k - 1], :] = (-K_lin @ Xc_lin[[k - 1], :].T).T
+        Xc_lin[k, :] = vdp.f(0, Xc_lin[k - 1, :], Uc_lin[k - 1, :].item())
+
+    # Predict dynamics with Koopman controller
+    Xc_poly = np.zeros((N, 2))
+    Uc_poly = np.zeros((N, 1))
+    Xc_poly[0, :] = np.array([1, 1])
+    for k in range(1, N):
+        lifted_state = kp_poly.lift_state(
+            Xc_poly[[k - 1], :],
+            episode_feature=False,
+        )
+        Uc_poly[[k - 1], :] = (-K_poly @ lifted_state.T).T
+        Xc_poly[k, :] = vdp.f(0, Xc_poly[k - 1, :], Uc_poly[k - 1, :].item())
+
+    # Plot trajectories
+    fig, ax = plt.subplots()
+    ax.plot(Xc_gt[:, 0], Xc_gt[:, 1], label='Uncontrolled')
+    ax.plot(Xc_lin[:, 0], Xc_lin[:, 1], label='Linear LQR')
+    ax.plot(Xc_poly[:, 0], Xc_poly[:, 1], label='Koopman LQR')
+    ax.scatter([0], [0], s=50, c='C3', zorder=2, marker='x', label='Target')
+    ax.set_xlabel('$x_1[k]$')
+    ax.set_ylabel('$x_2[k]$')
+    ax.legend(loc='lower right', ncol=2)
+    ax.set_title('Comparison of linear and Koopman LQR controllers')
+
+
+if __name__ == '__main__':
+    example_control_vdp()
+    plt.show()