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

GPU available: True (cuda), used: True
TPU available: False, using: 0 TPU cores
IPU available: False, using: 0 IPUs
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:70: 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\pykoopman\regression\_nndmd.py:877: UserWarning: Warning: no validation data prepared
  warn("Warning: no validation data prepared")
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
LOCAL_RANK: 0 - CUDA_VISIBLE_DEVICES: [0]

  | 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:430: 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:280: 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(

`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 0x00000132194D9090>],
                                        'max_epochs': 3}))

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.7999816  0.70000285]

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)