https://github.com/Koopman-Laboratory/KoopmanLab
A library for Koopman Neural Operator with Pytorch.
https://github.com/Koopman-Laboratory/KoopmanLab
deep-learning koopman neural-network neural-operator pytorch scientific-machine-learning
Last synced: 2 days ago
JSON representation
A library for Koopman Neural Operator with Pytorch.
- Host: GitHub
- URL: https://github.com/Koopman-Laboratory/KoopmanLab
- Owner: Koopman-Laboratory
- License: gpl-3.0
- Created: 2022-12-21T09:58:28.000Z (over 3 years ago)
- Default Branch: main
- Last Pushed: 2024-10-05T03:08:27.000Z (over 1 year ago)
- Last Synced: 2026-01-28T22:56:43.658Z (5 months ago)
- Topics: deep-learning, koopman, neural-network, neural-operator, pytorch, scientific-machine-learning
- Language: Python
- Homepage:
- Size: 579 KB
- Stars: 314
- Watchers: 7
- Forks: 25
- Open Issues: 9
-
Metadata Files:
- Readme: README.md
- License: LICENSE
- Citation: CITATION.cff
Awesome Lists containing this project
- awesome-ai-cae - Koopman-Laboratory/KoopmanLab - Koopman Neural Operator for mesh-free nonlinear PDE solving. Multi-scale decomposition. (AI/ML for Simulation / Verified vs declared)
README
[](
https://pypi.org/project/koopmanlab/)
[](
https://github.com/Koopman-Laboratory/KoopmanLab)
[](
https://github.com/Koopman-Laboratory/KoopmanLab/blob/main/LICENSE)
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](https://www.sciencedirect.com/science/article/pii/S0021999124004431)." Journal of Computational Physics (2024). See also the arXiv preprint arXiv:2301.10022 (2023).
- "[KoopmanLab: machine learning for solving complex physics equations](https://pubs.aip.org/aip/aml/article/1/3/036110/2910717/KoopmanLab-Machine-learning-for-solving-complex)." 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:
``` python
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](https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-)
If you want to generate Navier-Stokes Equation data by yourself, the data generation configuration file can be found in the following link.
- [File](https://github.com/zongyi-li/fourier_neural_operator/tree/master/data_generation/navier_stokes)
Our package provides an easy way to create a Koopman neural operator model.
``` python
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`.
``` python
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.
``` python
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`.
``` python
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.
``` python
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).
``` python
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](https://github.com/Koopman-Laboratory/KoopmanLab/blob/main/LICENSE)