KoopmanLab is a package for Koopman Neural Operator with Pytorch.

For more information, please refer to the following paper, where we provid detailed mathematical derivations, computational designs, and code explanations.

"Koopman neural operator as a mesh-free solver of non-linear partial differential equations." Journal of Computational Physics (2024). See also the arXiv preprint arXiv:2301.10022 (2023).
"KoopmanLab: machine learning for solving complex physics equations." APL Machine Learning (2023).

Installation

KoopmanLab requires the following dependencies to be installed:

PyTorch >= 1.10
Numpy >= 1.23.2
Matplotlib >= 3.3.2

You can install KoopmanLab package via the following approaches:

Install the stable version with pip:

$ pip install koopmanlab

Install the current version by source code with pip:

$ git clone https://github.com/Koopman-Laboratory/KoopmanLab.git
$ cd KoopmanLab
$ pip install -e .

Quick Start

If you install KoopmanLab successfully, you can use our model directly by:

import koopmanlab as kp
encoder = kp.models.encoder_mlp(t_in, operator_size)
decoder = kp.models.decoder_mlp(t_in, operator_size)
KNO1d_model = kp.models.KNO1d(encoder, decoder, operator_size, modes_x = 16, decompose = 6)
# Input size [batch, x, t_in] Output size [batch, x, t_in] for once iteration
KNO2d_model = kp.models.KNO2d(encoder, decoder, operator_size, modes_x = 10, modes_y = 10, decompose = 6)
# Input size [batch, x, t_in] Output size [batch, x, t_in] for once iteration

If you do not want to customize the algorithms for training, testing and plotting, we highly recommend that you use our basic APIs to build a Koopman model.

Usage

You can read demo_ns.py to learn about some basic APIs and workflow of KoopmanLab. If you want to run demo_ns.py, the following data need to be prepared in your computing resource.

Dataset

If you want to generate Navier-Stokes Equation data by yourself, the data generation configuration file can be found in the following link.

File

Our package provides an easy way to create a Koopman neural operator model.

import koopmanlab as kp
MLP_KNO_2D = kp.model.koopman(backbone = "KNO2d", autoencoder = "MLP", device = device)
MLP_KNO_2D = kp.model.koopman(backbone = "KNO2d", autoencoder = "MLP", o = o, m = m, r = r, t_in = 10, device = device)
MLP_KNO_2D.compile()
## Parameter definitions:
# o: the dimension of the learned Koopman operator
# f: the number of frequency modes below frequency truncation threshold
# r: the power of the Koopman operator
# T_in: the duration length of input data
# device : if CPU or GPU is used for calculating

ViT_KNO = kp.model.koopman_vit(decoder = "MLP", resolution=(64, 64), patch_size=(2, 2),
            in_chans=1, out_chans=1, head_num=16, embed_dim=768, depth = 16, parallel = True, high_freq = True, device=device)
ViT_KNO.compile()
## Parameter definitions:
# depth: the depth of each head 
# head_num: the number of heads
# resolution: the spatial resolution of input data
# patch_size: the size of each patch (i.e., token)
# in_chans: the number of target variables in the data set
# out_chans: the number of predicted variables by ViT-KNO , which is usually same as in_chans
# embed_dim: the embeding dimension
# parallel: if data parallel is applied
# high_freq: if high-frequency information complement is applied

Once the model is compiled, an optimizer setting is required to run your own experiments. If you want a more customized setting of optimizer and scheduler, you could use any PyTorch method to create them and assign them to Koopman neural operator object, eg. MLP_KNO_2D.optimizer and MLP_KNO_2D.scheduler.

MLP_KNO_2D.opt_init("Adam", lr = 0.005, step_size=100, gamma=0.5)

If you use Burgers equation and Navier-Stokes equation data or the shallow water data provided by PDEBench, there are three specifc data interfaces that you can consider.

train_loader, test_loader = kp.data.burgers(path, batch_size = 64, sub = 32)
train_loader, test_loader = kp.data.shallow_water(path, batch_size = 5, T_in = 10, T_out = 40, sub = 1)
train_loader, test_loader = kp.data.navier_stokes(path, batch_size = 10, T_in = 10, T_out = 40, type = "1e-3", sub = 1)
## Parameter definitions:
# path: the file path of the downloaded data set
# T_in: the duration length of input data
# T_out: the duration length required to predict
# Type: the viscosity coefficient of navier-stokes equation data set.
# sub: the down-sampling scaling factor. For instance, a scaling factor sub=2 acting on a 2-dimensional data with the spatial resoluion 64*64 will create a down-sampled space of 32*32. The same factor action on a 1 dimensional data with the spatial resoluion 1*64 implies a down-sampled space of 1*32.

We recommend that you process your data by PyTorch method torch.utils.data.DataLoader. In KNO model, the shape of 2D input data is [batchsize, x, y, t_len] and the shape of output data and label is [batchsize, x, y, T], where t_len is defined in kp.model.koopman and T is defined in train module. In Koopman-ViT model, the shape of 2D input data is [batchsize, in_chans, x, y] and the shape of output data and label is [batchsize, out_chans, x, y].

The KoopmanLab provides two training and two testing methods of the compact KNO sub-family. If your scenario is single step prediction, you can consider to use train_single method or use train with T_out = 1. Our package provides a method to save and visualize your prediction results in test.

MLP_KNO_2D.train_single(epochs=ep, trainloader = train_loader, evalloader = eval_loader)
MLP_KNO_2D.train(epochs=ep, trainloader = train_loader, evalloader = eval_loader, T_out = T)
MLP_KNO_2D.test_single(test_loader)
MLP_KNO_2D.test(test_loader, T_out = T, path = "./fig/ns_time_error_1e-4/", is_save = True, is_plot = True)

As for the ViT-KNO sub-family, train and test method is set with a single step predicition scenario. Specifically, train_multi and test_multi method provide multi-step iteration prediction, where the model iterates T_out times in training and testing.

ViT_KNO.train_single(epochs=ep, trainloader = train_loader, evalloader = eval_loader)
ViT_KNO.test_single(test_loader)
ViT_KNO.train_multi(epochs=ep, trainloader = train_loader, evalloader = eval_loader, T_out = T_out)
ViT_KNO.test_multi(test_loader)
## Parameter definitions:
# epoch: epoch number of training
# trainloader: dataloader of training, which is returning variable from torch.utils.data.DataLoader
# evalloader: dataloader of evaluating, which is returning variable from torch.utils.data.DataLoader
# test_loader: dataloader of testing, which is returning variable from torch.utils.data.DataLoader
# T_out: the duration length required to predict

Once your model has been trained, you can use the saving module provided in KoopmanLab to save your model. Saved variable has three attribute. where koopman is the model class variable (i.e., the saved kno_model variable), model is the trained model variable (i.e., the saved kno_model.kernel variable), and model_params is the parameters dictionary of trained model variable (i.e., the saved kno_model.kernel.state_dict() variable).

MLP_KNO_2D.save(save_path)
## Parameter definitions:
# save_path: the file path of the result saving

Citation

If you use KoopmanLab package for academic research, you are encouraged to cite the following paper:

@article{xiong2024koopman,
  title={Koopman neural operator as a mesh-free solver of non-linear partial differential equations},
  author={Xiong, Wei and Huang, Xiaomeng and Zhang, Ziyang and Deng, Ruixuan and Sun, Pei and Tian, Yang},
  journal={Journal of Computational Physics},
  pages={113194},
  year={2024},
  publisher={Elsevier}
}

@article{xiong2023koopmanlab,
  title={Koopmanlab: machine learning for solving complex physics equations},
  author={Xiong, Wei and Ma, Muyuan and Huang, Xiaomeng and Zhang, Ziyang and Sun, Pei and Tian, Yang},
  journal={APL Machine Learning},
  volume={1},
  number={3},
  year={2023},
  publisher={AIP Publishing}
}

Acknowledgement

Authors appreciate Abby, a talented artist, for designing the logo of KoopmanLab.

License

GPL-3.0 License