Merge pull request #182 from decargroup/add-control-example

Add control example

CITATION.cff

@@ -10,6 +10,6 @@ authors:
     orcid: "https://orcid.org/0000-0002-1987-9268"
     affiliation: "McGill University"
 title: "decargroup/pykoop"
-version: v2.0.1
+version: v2.0.2
 doi: 10.5281/zenodo.5576490
 url: "https://github.com/decargroup/pykoop"

README.rst

@@ -221,7 +221,7 @@ If you use this software in your research, please cite it as below or see
         url={https://github.com/decargroup/pykoop},
         publisher={Zenodo},
         author={Steven Dahdah and James Richard Forbes},
-        version = {{v2.0.1}},
+        version = {{v2.0.2}},
         year={2024},
     }
 

doc/examples.rst

@@ -70,3 +70,11 @@ For more details on how these features are generated, see [RR07]_.
 
 .. plot:: ../examples/7_example_rff_duffing.py
    :include-source:
+
+Control using the Koopman operator
+----------------------------------
+
+This example demonstrates Koopman LQR control using a Van der Pol oscillator.
+
+.. plot:: ../examples/8_control_vdp.py
+   :include-source:

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()

requirements-dev.txt

@@ -1,6 +1,7 @@
 # Requirements to install ``pykoop``, run all unit tests, and generate docs
 -r requirements.txt
 
+control
 pandas
 pytest
 pytest-regressions

setup.py

@@ -7,7 +7,7 @@
 
 setuptools.setup(
     name='pykoop',
-    version='2.0.1',
+    version='2.0.2',
     description=('Koopman operator identification library in Python, '
                  'compatible with `scikit-learn`'),
     long_description=readme,