EDMD and NNDMD for a simple linear system

In this tutorial, we use EDMD to learn approximate Koopman eigenfunctions for the following two-dimensional, linear system (example from Rowley, Williams & Kevrekidis, “IPAM MTWS4”, 2014):

\[\begin{split}{\bf x}_{k+1} = \begin{bmatrix} 0.8 & -0.05\\ 0 & 0.7\end{bmatrix}{\bf x}_k,\end{split}\]

where \({\bf x} = [x, y]^T\). Polynomial functions \(x^iy^j\) with \(i,j\in [0,3]\) up to order 3 are used as basis functions. The true eigenfunctions and eigenvalues can be approximated as follows:

\[\begin{split}\varphi_{ij}(x,y) \approx (0.894x - 0.447y)^i y^j\quad \text{for}\quad i, j\in\mathbb{N}\\ \lambda_{ij} = (0.8)^i(0.7)^j\end{split}\]

import sys
sys.path.append('../src')

import pykoopman as pk
from pykoopman.common import Linear2Ddynamics
import scipy
import numpy as np
import numpy.random as rnd
np.random.seed(42)  # for reproducibility
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings('ignore')

Collect training data

# Create instance of the dynamical system
sys = Linear2Ddynamics()

# Collect training data
n_pts = 51
n_int = 1
xx, yy = np.meshgrid(np.linspace(-1, 1, n_pts), np.linspace(-1, 1, n_pts))
x = np.vstack((xx.flatten(), yy.flatten()))
n_traj = x.shape[1]

X, Y = sys.collect_data(x, n_int, n_traj)

fig, axs = plt.subplots(1, 1, tight_layout=True, figsize=(4, 4))
for traj_idx in range(n_traj):
    axs.plot([X[0, traj_idx::n_traj], Y[0, traj_idx::n_traj]],
             [X[1, traj_idx::n_traj], Y[1, traj_idx::n_traj]], '-k',alpha=0.1)
axs.set(ylabel=r'$x_2$',
        xlabel=r'$x_1$')

[Text(0, 0.5, '$x_2$'), Text(0.5, 0, '$x_1$')]

Analytic Koopman eigenvalues and eigenfunctions

degree = 3
true_efuns = np.zeros((X.shape[1], 12))
true_evals = np.zeros(12)
counter = 0
for i in range(degree+1):
    for j in range(degree+1):
        if i*j <= degree:
            true_evals[counter] = (0.8**i) * (0.7**j)
            # tmp = ((0.894 * xx - 0.447 * yy)**i) * (yy**j)
            # true_efuns[:, counter] = tmp.flatten()
            true_efuns[:, counter] = ((0.894 * x[0, :] - 0.447 * x[1, :])**i) * \
                                     (x[1, :]**j)
            counter += 1

sort_idx = np.argsort(true_evals)
sort_idx = sort_idx[::-1]
true_evals = true_evals[sort_idx]
true_efuns = true_efuns[:, sort_idx]
sys.visualize_modes(X, true_efuns, eigvals=true_evals)
print(true_evals)

[1.     0.8    0.7    0.64   0.56   0.512  0.49   0.448  0.392  0.3584
 0.343  0.2744]

Data-driven approximation with EDMD and polynomial basis functions

EDMDc = pk.regression.EDMD()
obsv = pk.observables.Polynomial(degree=3)
model = pk.Koopman(observables=obsv, regressor=EDMDc)
model.fit(X.T, y=Y.T)

psi = model.psi(x_col=x)
eigenvalues = np.real(np.diag(model.lamda))
order = np.argsort(eigenvalues)[::-1]

print(eigenvalues[order])
sys.visualize_modes(X, psi.T, eigvals=eigenvalues, order=order)

[1.    0.8   0.7   0.64  0.56  0.512 0.49  0.448 0.392 0.343]

Data-driven approximation with NN-DMD

import lightning as L
class IncrementalSequenceLoss(L.Callback):
    def on_train_epoch_end(self, trainer, pl_module):
        max_look_forward = pl_module.look_forward
        if trainer.callback_metrics["loss"] < 1e-2 and \
                pl_module.masked_loss_metric.max_look_forward < max_look_forward:
            print("increase width size from {} to {}".format(
                pl_module.masked_loss_metric.max_look_forward,
                pl_module.masked_loss_metric.max_look_forward+1) )
            print("")
            pl_module.masked_loss_metric.max_look_forward += 1

look_forward = 1
dlk_regressor = pk.regression.NNDMD(look_forward=look_forward,
            config_encoder=dict(input_size=2, hidden_sizes=[32] * 2, output_size=2,
                                activations="linear"),
            config_decoder=dict(input_size=2, hidden_sizes=[32] * 2, output_size=2,
                                activations="linear"),
                                    batch_size=512, lbfgs=True, \
                                    normalize=True, normalize_mode='max',
                                    trainer_kwargs=dict(max_epochs=3, callbacks=[
        IncrementalSequenceLoss()]))

model = pk.Koopman(regressor=dlk_regressor)
model.fit(X.T, Y.T)

INFO: GPU available: True (cuda), used: True
[rank_zero.py:64 -                _info() ] GPU available: True (cuda), used: True
INFO: TPU available: False, using: 0 TPU cores
[rank_zero.py:64 -                _info() ] TPU available: False, using: 0 TPU cores
INFO: IPU available: False, using: 0 IPUs
[rank_zero.py:64 -                _info() ] IPU available: False, using: 0 IPUs
INFO: HPU available: False, using: 0 HPUs
[rank_zero.py:64 -                _info() ] HPU available: False, using: 0 HPUs
D:\work\pykoopman\venv\lib\site-packages\lightning\pytorch\trainer\connectors\logger_connector\logger_connector.py:67: UserWarning: Starting from v1.9.0, `tensorboardX` has been removed as a dependency of the `lightning.pytorch` package, due to potential conflicts with other packages in the ML ecosystem. For this reason, `logger=True` will use `CSVLogger` as the default logger, unless the `tensorboard` or `tensorboardX` packages are found. Please `pip install lightning[extra]` or one of them to enable TensorBoard support by default
  warning_cache.warn(
D:\work\pykoopman\venv\lib\site-packages\lightning\pytorch\trainer\configuration_validator.py:69: UserWarning: You passed in a `val_dataloader` but have no `validation_step`. Skipping val loop.
  rank_zero_warn("You passed in a `val_dataloader` but have no `validation_step`. Skipping val loop.")
D:\work\pykoopman\src\pykoopman\regression\_nndmd.py:898: UserWarning: Warning: no validation data prepared
  warn("Warning: no validation data prepared")
INFO: You are using a CUDA device ('NVIDIA GeForce RTX 3080 Ti Laptop GPU') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
[rank_zero.py:64 -                _info() ] You are using a CUDA device ('NVIDIA GeForce RTX 3080 Ti Laptop GPU') that has Tensor Cores. To properly utilize them, you should set `torch.set_float32_matmul_precision('medium' | 'high')` which will trade-off precision for performance. For more details, read https://pytorch.org/docs/stable/generated/torch.set_float32_matmul_precision.html#torch.set_float32_matmul_precision
INFO: LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
[cuda.py:58 -     set_nvidia_flags() ] LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]
INFO:
  | Name                | Type                    | Params
----------------------------------------------------------------
0 | _encoder            | FFNN                    | 1.2 K
1 | _decoder            | FFNN                    | 1.2 K
2 | _koopman_propagator | StandardKoopmanOperator | 4
3 | masked_loss_metric  | MaskedMSELoss           | 0
----------------------------------------------------------------
2.3 K     Trainable params
0         Non-trainable params
2.3 K     Total params
0.009     Total estimated model params size (MB)
[model_summary.py:83 -            summarize() ]
  | Name                | Type                    | Params
----------------------------------------------------------------
0 | _encoder            | FFNN                    | 1.2 K
1 | _decoder            | FFNN                    | 1.2 K
2 | _koopman_propagator | StandardKoopmanOperator | 4
3 | masked_loss_metric  | MaskedMSELoss           | 0
----------------------------------------------------------------
2.3 K     Trainable params
0         Non-trainable params
2.3 K     Total params
0.009     Total estimated model params size (MB)
D:\work\pykoopman\venv\lib\site-packages\lightning\pytorch\trainer\connectors\data_connector.py:442: PossibleUserWarning: The dataloader, train_dataloader, does not have many workers which may be a bottleneck. Consider increasing the value of the `num_workers` argument` (try 20 which is the number of cpus on this machine) in the `DataLoader` init to improve performance.
  rank_zero_warn(
D:\work\pykoopman\venv\lib\site-packages\lightning\pytorch\loops\fit_loop.py:281: PossibleUserWarning: The number of training batches (6) is smaller than the logging interval Trainer(log_every_n_steps=50). Set a lower value for log_every_n_steps if you want to see logs for the training epoch.
  rank_zero_warn(

INFO: `Trainer.fit` stopped: `max_epochs=3` reached.
[rank_zero.py:64 -                _info() ] `Trainer.fit` stopped: `max_epochs=3` reached.

Koopman(observables=Identity(),
        regressor=NNDMD(batch_size=512,
                        config_decoder={'activations': 'linear',
                                        'hidden_sizes': [32, 32],
                                        'input_size': 2, 'output_size': 2},
                        config_encoder={'activations': 'linear',
                                        'hidden_sizes': [32, 32],
                                        'input_size': 2, 'output_size': 2},
                        lbfgs=True, normalize_mode='max',
                        trainer_kwargs={'callbacks': [<__main__.IncrementalSequenceLoss object at 0x000002CC1356A410>],
                                        'max_epochs': 3}))

In a Jupyter environment, please rerun this cell to show the HTML representation or trust the notebook.
On GitHub, the HTML representation is unable to render, please try loading this page with nbviewer.org.

psi = model.psi(x_col=x)
eigenvalues = np.real(np.diag(model.lamda))
order = np.argsort(eigenvalues)[::-1]

print(eigenvalues[order])
sys.visualize_modes(X, psi.T, eigvals=eigenvalues, order=order)

[0.79999924 0.70000166]

Check shape of matrices related to Koopman operator

model.A.shape

(2, 2)

model.C.shape

(2, 2)

model.W.shape

(2, 2)

model.lamda.shape

(2, 2)