Skip to content

Commit

Permalink
Merge pull request #182 from decargroup/add-control-example
Browse files Browse the repository at this point in the history
Add control example
  • Loading branch information
sdahdah authored Oct 4, 2024
2 parents 8283eb9 + 892f53d commit 1ba2728
Show file tree
Hide file tree
Showing 7 changed files with 421 additions and 3 deletions.
2 changes: 1 addition & 1 deletion CITATION.cff
Original file line number Diff line number Diff line change
Expand Up @@ -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"
2 changes: 1 addition & 1 deletion README.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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},
}
Expand Down
8 changes: 8 additions & 0 deletions doc/examples.rst
Original file line number Diff line number Diff line change
Expand Up @@ -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:
110 changes: 110 additions & 0 deletions examples/8_control_vdp.py
Original file line number Diff line number Diff line change
@@ -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()
299 changes: 299 additions & 0 deletions notebooks/8_control_vdp.ipynb

Large diffs are not rendered by default.

1 change: 1 addition & 0 deletions requirements-dev.txt
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
# Requirements to install ``pykoop``, run all unit tests, and generate docs
-r requirements.txt

control
pandas
pytest
pytest-regressions
Expand Down
2 changes: 1 addition & 1 deletion setup.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down

0 comments on commit 1ba2728

Please sign in to comment.