2 | 
3 |
4 | 5 | []( 6 | https://pypi.org/project/koopmanlab/) 7 | []( 8 | https://github.com/Koopman-Laboratory/KoopmanLab) 9 | []( 10 | https://github.com/Koopman-Laboratory/KoopmanLab/blob/main/LICENSE) 11 | 12 | KoopmanLab is a package for Koopman Neural Operator with Pytorch. 13 | 14 | For more information, please refer to the following paper, where we provid detailed mathematical derivations, computational designs, and code explanations. 15 | - "[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). 16 | - "[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). 17 | 18 | # Installation 19 | KoopmanLab requires the following dependencies to be installed: 20 | - PyTorch >= 1.10 21 | - Numpy >= 1.23.2 22 | - Matplotlib >= 3.3.2 23 | 24 | You can install KoopmanLab package via the following approaches: 25 | 26 | - Install the stable version with `pip`: 27 | 28 | ``` 29 | $ pip install koopmanlab 30 | ``` 31 | 32 | - Install the current version by source code with `pip`: 33 | ``` 34 | $ git clone https://github.com/Koopman-Laboratory/KoopmanLab.git 35 | $ cd KoopmanLab 36 | $ pip install -e . 37 | ``` 38 | # Quick Start 39 | 40 | If you install KoopmanLab successfully, you can use our model directly by: 41 | 42 | ``` python 43 | import koopmanlab as kp 44 | encoder = kp.models.encoder_mlp(t_in, operator_size) 45 | decoder = kp.models.decoder_mlp(t_in, operator_size) 46 | KNO1d_model = kp.models.KNO1d(encoder, decoder, operator_size, modes_x = 16, decompose = 6) 47 | # Input size [batch, x, t_in] Output size [batch, x, t_in] for once iteration 48 | KNO2d_model = kp.models.KNO2d(encoder, decoder, operator_size, modes_x = 10, modes_y = 10, decompose = 6) 49 | # Input size [batch, x, t_in] Output size [batch, x, t_in] for once iteration 50 | ``` 51 | 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. 52 | 53 | # Usage 54 | 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. 55 | - [Dataset](https://drive.google.com/drive/folders/1UnbQh2WWc6knEHbLn-ZaXrKUZhp7pjt-) 56 | 57 | If you want to generate Navier-Stokes Equation data by yourself, the data generation configuration file can be found in the following link. 58 | 59 | - [File](https://github.com/zongyi-li/fourier_neural_operator/tree/master/data_generation/navier_stokes) 60 | 61 | Our package provides an easy way to create a Koopman neural operator model. 62 | ``` python 63 | import koopmanlab as kp 64 | MLP_KNO_2D = kp.model.koopman(backbone = "KNO2d", autoencoder = "MLP", device = device) 65 | MLP_KNO_2D = kp.model.koopman(backbone = "KNO2d", autoencoder = "MLP", o = o, m = m, r = r, t_in = 10, device = device) 66 | MLP_KNO_2D.compile() 67 | ## Parameter definitions: 68 | # o: the dimension of the learned Koopman operator 69 | # f: the number of frequency modes below frequency truncation threshold 70 | # r: the power of the Koopman operator 71 | # T_in: the duration length of input data 72 | # device : if CPU or GPU is used for calculating 73 | 74 | ViT_KNO = kp.model.koopman_vit(decoder = "MLP", resolution=(64, 64), patch_size=(2, 2), 75 | in_chans=1, out_chans=1, head_num=16, embed_dim=768, depth = 16, parallel = True, high_freq = True, device=device) 76 | ViT_KNO.compile() 77 | ## Parameter definitions: 78 | # depth: the depth of each head 79 | # head_num: the number of heads 80 | # resolution: the spatial resolution of input data 81 | # patch_size: the size of each patch (i.e., token) 82 | # in_chans: the number of target variables in the data set 83 | # out_chans: the number of predicted variables by ViT-KNO , which is usually same as in_chans 84 | # embed_dim: the embeding dimension 85 | # parallel: if data parallel is applied 86 | # high_freq: if high-frequency information complement is applied 87 | ``` 88 | 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`. 89 | ``` python 90 | MLP_KNO_2D.opt_init("Adam", lr = 0.005, step_size=100, gamma=0.5) 91 | ``` 92 | 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. 93 | ``` python 94 | train_loader, test_loader = kp.data.burgers(path, batch_size = 64, sub = 32) 95 | train_loader, test_loader = kp.data.shallow_water(path, batch_size = 5, T_in = 10, T_out = 40, sub = 1) 96 | train_loader, test_loader = kp.data.navier_stokes(path, batch_size = 10, T_in = 10, T_out = 40, type = "1e-3", sub = 1) 97 | ## Parameter definitions: 98 | # path: the file path of the downloaded data set 99 | # T_in: the duration length of input data 100 | # T_out: the duration length required to predict 101 | # Type: the viscosity coefficient of navier-stokes equation data set. 102 | # 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. 103 | ``` 104 | 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]`. 105 | 106 | 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`. 107 | ``` python 108 | MLP_KNO_2D.train_single(epochs=ep, trainloader = train_loader, evalloader = eval_loader) 109 | MLP_KNO_2D.train(epochs=ep, trainloader = train_loader, evalloader = eval_loader, T_out = T) 110 | MLP_KNO_2D.test_single(test_loader) 111 | MLP_KNO_2D.test(test_loader, T_out = T, path = "./fig/ns_time_error_1e-4/", is_save = True, is_plot = True) 112 | ``` 113 | 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. 114 | ``` python 115 | ViT_KNO.train_single(epochs=ep, trainloader = train_loader, evalloader = eval_loader) 116 | ViT_KNO.test_single(test_loader) 117 | ViT_KNO.train_multi(epochs=ep, trainloader = train_loader, evalloader = eval_loader, T_out = T_out) 118 | ViT_KNO.test_multi(test_loader) 119 | ## Parameter definitions: 120 | # epoch: epoch number of training 121 | # trainloader: dataloader of training, which is returning variable from torch.utils.data.DataLoader 122 | # evalloader: dataloader of evaluating, which is returning variable from torch.utils.data.DataLoader 123 | # test_loader: dataloader of testing, which is returning variable from torch.utils.data.DataLoader 124 | # T_out: the duration length required to predict 125 | ``` 126 | 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). 127 | ``` python 128 | MLP_KNO_2D.save(save_path) 129 | ## Parameter definitions: 130 | # save_path: the file path of the result saving 131 | ``` 132 | # Citation 133 | If you use KoopmanLab package for academic research, you are encouraged to cite the following paper: 134 | ``` 135 | @article{xiong2024koopman, 136 | title={Koopman neural operator as a mesh-free solver of non-linear partial differential equations}, 137 | author={Xiong, Wei and Huang, Xiaomeng and Zhang, Ziyang and Deng, Ruixuan and Sun, Pei and Tian, Yang}, 138 | journal={Journal of Computational Physics}, 139 | pages={113194}, 140 | year={2024}, 141 | publisher={Elsevier} 142 | } 143 | 144 | @article{xiong2023koopmanlab, 145 | title={Koopmanlab: machine learning for solving complex physics equations}, 146 | author={Xiong, Wei and Ma, Muyuan and Huang, Xiaomeng and Zhang, Ziyang and Sun, Pei and Tian, Yang}, 147 | journal={APL Machine Learning}, 148 | volume={1}, 149 | number={3}, 150 | year={2023}, 151 | publisher={AIP Publishing} 152 | } 153 | ``` 154 | # Acknowledgement 155 | Authors appreciate Abby, a talented artist, for designing the logo of KoopmanLab. 156 | 157 | # License 158 | [GPL-3.0 License](https://github.com/Koopman-Laboratory/KoopmanLab/blob/main/LICENSE) 159 | -------------------------------------------------------------------------------- /demo_ns.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import koopmanlab as kp 3 | # Setting your computing device 4 | torch.cuda.set_device(0) 5 | device = torch.device("cuda") 6 | 7 | # Path 8 | data_path = "./data/ns_V1e-3_N5000_T50.mat" 9 | fig_path = "./demo/fig/" 10 | save_path = "./demo/result/" 11 | 12 | # Loading Data 13 | train_loader, test_loader = kp.data.navier_stokes(data_path, batch_size = 10, T_in = 10, T_out = 40, type = "1e-3", sub = 1) 14 | 15 | # Hyper parameters 16 | ep = 1 # Training Epoch 17 | o = 32 # Koopman Operator Size 18 | m = 16 # Modes 19 | r = 8 # Power of Koopman Matrix 20 | 21 | # Model 22 | koopman_model = kp.model.koopman(backbone = "KNO2d", autoencoder = "MLP", o = o, m = m, r = r, t_in = 10, device = device) 23 | koopman_model.compile() 24 | koopman_model.opt_init("Adam", lr = 0.005, step_size=100, gamma=0.5) 25 | koopman_model.train(epochs=ep, trainloader = train_loader, evalloader = test_loader) 26 | 27 | # Result and Saving 28 | time_error = koopman_model.test(test_loader, path = fig_path, is_save = True, is_plot = True) 29 | filename = "ns_time_error_op" + str(o) + "m" + str(m) + "r" +str(r) + ".pt" 30 | torch.save({"time_error":time_error,"params":koopman_model.params}, save_path + filename) 31 | -------------------------------------------------------------------------------- /koopmanlab/__init__.py: -------------------------------------------------------------------------------- 1 | __all__ = ["data","model"] 2 | 3 | from . import data 4 | from . import model 5 | from . import func 6 | from .model import koopman, koopman_vit 7 | from .models import KNO1d,KNO2d 8 | from .models import encoder_mlp, decoder_mlp, encoder_conv1d, decoder_conv1d, encoder_conv2d, decoder_conv2d 9 | from .models import ViT 10 | -------------------------------------------------------------------------------- /koopmanlab/__pycache__/__init__.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/__init__.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/__init__.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/__init__.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/data.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/data.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/data.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/data.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/func.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/func.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/func.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/func.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/kno.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/kno.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/kno.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/kno.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/koopman_vit.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/koopman_vit.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/model.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/model.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/model.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/model.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/utilities.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/utilities.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/utils.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/utils.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/utils.cpython-38.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/utils.cpython-38.pyc -------------------------------------------------------------------------------- /koopmanlab/__pycache__/vit.cpython-310.pyc: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/koopmanlab/__pycache__/vit.cpython-310.pyc -------------------------------------------------------------------------------- /koopmanlab/data.py: -------------------------------------------------------------------------------- 1 | import h5py 2 | import torch 3 | import scipy.io 4 | import numpy as np 5 | 6 | def burgers(path, batch_size = 64, sub = 32): 7 | f = scipy.io.loadmat(path) 8 | x_data = f['a'][:,::sub] 9 | y_data = f['u'][:,::sub] 10 | 11 | x_train = torch.tensor(x_data[:1000,:],dtype=torch.float32) 12 | y_train = torch.tensor(y_data[:1000,:],dtype=torch.float32) 13 | x_test = torch.tensor(x_data[-200:,:],dtype=torch.float32) 14 | y_test = torch.tensor(y_data[-200:,:],dtype=torch.float32) 15 | 16 | S = x_train.shape[1] 17 | 18 | x_train = x_train.reshape(1000,S,1) 19 | x_test = x_test.reshape(200,S,1) 20 | x_test = x_test 21 | y_test = y_test 22 | 23 | print("Burgers Dataset has been loaded successfully!") 24 | print("X train shape:", x_train.shape, "Y train shape:", y_train.shape) 25 | print("X test shape:", x_test.shape, "Y test shape:", y_test.shape) 26 | 27 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size, shuffle=True) 28 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size, shuffle=False) 29 | return train_loader, test_loader 30 | 31 | def shallow_water(path, batch_size = 20, T_in = 10, T_out = 20, sub = 1): 32 | ntrain = 900 33 | ntest = 100 34 | total = ntrain + ntest 35 | f = h5py.File(path) 36 | data = f['data'][0:total] 37 | data = torch.tensor(data,dtype=torch.float32) 38 | # Traning data 39 | train_a = data[:ntrain,::sub,::sub,:T_in] 40 | train_u = data[:ntrain,::sub,::sub,T_in:T_out+T_in] 41 | # Testing data 42 | test_a = data[-ntest:,::sub,::sub,:T_in] 43 | test_u = data[-ntest:,::sub,::sub,T_in:T_out+T_in] 44 | 45 | print("Shallow Water Equations Dataset has been loaded successfully!") 46 | print("X train shape:", train_a.shape, "Y train shape:", train_u.shape) 47 | print("X test shape:", test_a.shape, "Y test shape:", test_u.shape) 48 | 49 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u), batch_size=batch_size, shuffle=True) 50 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=False) 51 | 52 | return train_loader, test_loader 53 | 54 | def navier_stokes(path, batch_size = 20, T_in = 10, T_out = 40, type = "1e-3", sub = 1,reshape = False): 55 | if type == "1e-3": 56 | ntrain = 1000 57 | ntest = 200 58 | total = ntrain + ntest 59 | f = h5py.File(path) 60 | data = f['u'][...,0:total] 61 | print("dataset shape : ", data.shape) # Print original shape of the data 62 | data = torch.tensor(data,dtype=torch.float32) 63 | data = data.permute(3,1,2,0) # The dimension of the data shape is [B, X, Y, T] 64 | 65 | # Traning data 66 | train_a = data[:ntrain,::sub,::sub,:T_in] 67 | train_u = data[:ntrain,::sub,::sub,T_in:T_out+T_in] 68 | # Testing data 69 | test_a = data[-ntest:,::sub,::sub,:T_in] 70 | test_u = data[-ntest:,::sub,::sub,T_in:T_out+T_in] 71 | 72 | if reshape: 73 | train_a = train_a.permute(reshape) 74 | train_u = train_u.permute(reshape) 75 | test_a = test_a.permute(reshape) 76 | test_u = test_u.permute(reshape) 77 | 78 | print("Navier-Stokes (vis = 1e-3) Dataset has been loaded successfully!") 79 | print("X train shape:", train_a.shape, "Y train shape:", train_u.shape) 80 | print("X test shape:", test_a.shape, "Y test shape:", test_u.shape) 81 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u), batch_size=batch_size, shuffle=True) 82 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=False) 83 | elif type == "1e-4": 84 | ntrain = 8000 85 | ntest = 200 86 | total = ntrain + ntest 87 | f = h5py.File(path) 88 | data = f['u'][...,0:8200] 89 | data = torch.tensor(data,dtype=torch.float32) 90 | data = data.permute(3,1,2,0) 91 | 92 | # Traning data 93 | train_a = data[:ntrain,::sub,::sub,:T_in] 94 | train_u = data[:ntrain,::sub,::sub,T_in:T_out+T_in] 95 | # Testing data 96 | test_a = data[-ntest:,::sub,::sub,:T_in] 97 | test_u = data[-ntest:,::sub,::sub,T_in:T_out+T_in] 98 | 99 | if reshape: 100 | train_a = train_a.permute(reshape) 101 | train_u = train_u.permute(reshape) 102 | test_a = test_a.permute(reshape) 103 | test_u = test_u.permute(reshape) 104 | 105 | print("Navier-Stokes (vis = 1e-4) Dataset has been loaded successfully!") 106 | print("X train shape:", train_a.shape, "Y train shape:", train_u.shape) 107 | print("X test shape:", test_a.shape, "Y test shape:", test_u.shape) 108 | 109 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u), batch_size=batch_size, shuffle=True) 110 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=False) 111 | elif type == "1e-5": 112 | ntrain = 1100 113 | ntest = 100 114 | total = ntrain + ntest 115 | f = scipy.io.loadmat(path) 116 | data = f['u'][...,0:total] 117 | data = torch.tensor(data,dtype=torch.float32) 118 | print(data.shape) 119 | 120 | 121 | # Traning data 122 | train_a = data[:ntrain,::sub,::sub,:T_in] 123 | train_u = data[:ntrain,::sub,::sub,T_in:T_out+T_in] 124 | # Testing data 125 | test_a = data[-ntest:,::sub,::sub,:T_in] 126 | test_u = data[-ntest:,::sub,::sub,T_in:T_out+T_in] 127 | 128 | if reshape: 129 | train_a = train_a.permute(reshape) 130 | train_u = train_u.permute(reshape) 131 | test_a = test_a.permute(reshape) 132 | test_u = test_u.permute(reshape) 133 | 134 | print("Navier-Stokes (vis = 1e-5) Dataset has been loaded successfully!") 135 | print("X train shape:", train_a.shape, "Y train shape:", train_u.shape) 136 | print("X test shape:", test_a.shape, "Y test shape:", test_u.shape) 137 | 138 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(train_a, train_u), batch_size=batch_size, shuffle=True) 139 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(test_a, test_u), batch_size=batch_size, shuffle=False) 140 | 141 | else: 142 | print("The type is unclaimed. Data loading is failed.") 143 | return 144 | 145 | return train_loader, test_loader 146 | 147 | def navier_stokes_single(path, batch_size = 64, T_in = 10, T_out = 40, type = "1e-4", sub = 1,reshape = False): 148 | f = scipy.io.loadmat(path) 149 | print(f["a"].shape) 150 | print(f["u"].shape) 151 | data = f["u"] 152 | loc = 0 153 | x = np.zeros([3980,256,256,1]) 154 | y = np.zeros([3980,256,256,1]) 155 | for i in range(20): 156 | for j in range(199): 157 | x[i,:,:,0:1] = data[i,:,:,j:j+1] 158 | y[i,:,:,0:1] = data[i,:,:,j+1:j+2] 159 | 160 | 161 | ntrain = 3600 162 | ntest = 200 163 | total = ntrain + ntest 164 | 165 | x = torch.tensor(x,dtype=torch.float32) 166 | y = torch.tensor(y,dtype=torch.float32) 167 | 168 | if reshape: 169 | x = x.permute(reshape) 170 | y = y.permute(reshape) 171 | 172 | # Traning data 173 | x_train = x[:ntrain] 174 | y_train = y[:ntrain] 175 | # Testing data 176 | x_test = x[-ntest:] 177 | y_test = y[-ntest:] 178 | 179 | 180 | print("Navier-Stokes (vis = 1e-4) Dataset has been loaded successfully!") 181 | print("X train shape:", x_train.shape, "Y train shape:", y_train.shape) 182 | print("X test shape:", x_test.shape, "Y test shape:", y_test.shape) 183 | 184 | train_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_train, y_train), batch_size=batch_size, shuffle=True) 185 | test_loader = torch.utils.data.DataLoader(torch.utils.data.TensorDataset(x_test, y_test), batch_size=batch_size, shuffle=False) 186 | 187 | return train_loader, test_loader 188 | -------------------------------------------------------------------------------- /koopmanlab/func.py: -------------------------------------------------------------------------------- 1 | import os 2 | import torch 3 | 4 | def save(obj, path): 5 | if type(obj) != "koopmanlab.model.koopman": 6 | print("Save function only for saving Koopman solver!") 7 | return 8 | (fpath,_) = os.path.split(path) 9 | if not os.path.isfile(path): 10 | os.makedirs(fpath) 11 | torch.save({"koopman":obj},path) 12 | print("Koopman solver has been saved successfully!") 13 | 14 | def load(path): 15 | f = torch.load(path) 16 | return f['koopman'] -------------------------------------------------------------------------------- /koopmanlab/model.py: -------------------------------------------------------------------------------- 1 | from koopmanlab.models import kno 2 | from koopmanlab import utils 3 | from koopmanlab.models import koopmanViT 4 | 5 | import os 6 | import torch 7 | import numpy as np 8 | import matplotlib.pyplot as plt 9 | from timeit import default_timer 10 | 11 | class koopman: 12 | def __init__(self, backbone = "KNO1d", autoencoder = "MLP", o = 16, m = 16, r = 8, t_in = 1, device = False): 13 | self.backbone = backbone 14 | self.autoencoder = autoencoder 15 | self.operator_size = o 16 | self.modes = m 17 | self.decompose = r 18 | self.device = device 19 | self.t_in = t_in 20 | # Core Model 21 | self.params = 0 22 | self.kernel = False 23 | # Opt Setting 24 | self.optimizer = False 25 | self.scheduler = False 26 | self.loss = torch.nn.MSELoss() 27 | def compile(self): 28 | if self.autoencoder == "MLP": 29 | encoder = kno.encoder_mlp(self.t_in, self.operator_size) 30 | decoder = kno.decoder_mlp(self.t_in, self.operator_size) 31 | print("The autoencoder type is MLP.") 32 | elif self.autoencoder == "Conv1d": 33 | encoder = kno.encoder_conv1d(self.t_in, self.operator_size) 34 | decoder = kno.decoder_conv1d(self.t_in, self.operator_size) 35 | print("The autoencoder type is Conv1d.") 36 | elif self.autoencoder == "Conv2d": 37 | encoder = kno.encoder_conv2d(self.t_in, self.operator_size) 38 | decoder = kno.decoder_conv2d(self.t_in, self.operator_size) 39 | print("The autoencoder type is Conv2d.") 40 | else: 41 | # encoder = kno.encoder_mlp(self.t_in, self.operator_size) 42 | # decoder = kno.decoder_mlp(self.t_in, self.operator_size) 43 | # print("The autoencoder type is MLP.") 44 | print("Wrong!") 45 | if self.backbone == "KNO1d": 46 | self.kernel = kno.KNO1d(encoder, decoder, self.operator_size, modes_x = self.modes, decompose = self.decompose).to(self.device) 47 | print("KNO1d model is completed.") 48 | 49 | elif self.backbone == "KNO2d": 50 | self.kernel = kno.KNO2d(encoder, decoder, self.operator_size, modes_x = self.modes, modes_y = self.modes,decompose = self.decompose).to(self.device) 51 | print("KNO2d model is completed.") 52 | 53 | if not self.kernel == False: 54 | self.params = utils.count_params(self.kernel) 55 | print("Koopman Model has been compiled!") 56 | print("The Model Parameters Number is ",self.params) 57 | def opt_init(self, opt, lr, step_size, gamma): 58 | if opt == "Adam": 59 | self.optimizer = utils.Adam(self.kernel.parameters(), lr= lr, weight_decay=1e-4) 60 | if not step_size == False: 61 | self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=step_size, gamma=gamma) 62 | 63 | def train_single(self, epochs, trainloader, evalloader = False): 64 | for ep in range(epochs): 65 | # Train 66 | self.kernel.train() 67 | t1 = default_timer() 68 | train_recons_full = 0 69 | train_pred_full = 0 70 | for xx, yy in trainloader: 71 | l_recons = 0 72 | bs = xx.shape[0] 73 | xx = xx.to(self.device) 74 | yy = yy.to(self.device) 75 | pred,im_re = self.kernel(xx) 76 | 77 | l_recons = self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 78 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 79 | 80 | train_pred_full += l_pred.item() 81 | train_recons_full += l_recons.item() 82 | 83 | loss = 5*l_pred + 0.5*l_recons 84 | self.optimizer.zero_grad() 85 | loss.backward() 86 | self.optimizer.step() 87 | train_pred_full = train_pred_full / len(trainloader) 88 | train_recons_full = train_recons_full / len(trainloader) 89 | t2 = default_timer() 90 | test_pred_full = 0 91 | test_recons_full = 0 92 | mse_test = 0 93 | # Test 94 | if evalloader: 95 | with torch.no_grad(): 96 | for xx, yy in evalloader: 97 | bs = xx.shape[0] 98 | loss = 0 99 | xx = xx.to(self.device) 100 | yy = yy.to(self.device) 101 | 102 | pred,im_re = self.kernel(xx) 103 | 104 | 105 | l_recons = self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 106 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 107 | 108 | 109 | test_pred_full += l_pred.item() 110 | test_recons_full += l_recons.item() 111 | 112 | test_pred_full = test_pred_full/len(evalloader) 113 | test_recons_full = test_recons_full/len(evalloader) 114 | 115 | self.scheduler.step() 116 | 117 | if evalloader: 118 | if ep == 0: 119 | print("Epoch","Time","[Train Recons MSE]","[Train Pred MSE]","[Eval Recons MSE]","[Eval Pred MSE]") 120 | print(ep, t2 - t1, train_recons_full, train_pred_full, test_recons_full, test_pred_full) 121 | else: 122 | if ep == 0: 123 | print("Epoch","Time","Train Recons MSE","Train Pred MSE") 124 | print(ep, t2 - t1, train_recons_full, train_pred_full) 125 | 126 | def test_single(self, testloader): 127 | test_pred_full = 0 128 | test_recons_full = 0 129 | with torch.no_grad(): 130 | for xx, yy in testloader: 131 | bs = xx.shape[0] 132 | loss = 0 133 | xx = xx.to(self.device) 134 | yy = yy.to(self.device) 135 | 136 | pred,im_re = self.kernel(xx) 137 | 138 | l_recons = self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 139 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 140 | 141 | test_pred_full += l_pred.item() 142 | test_recons_full += l_recons.item() 143 | test_pred_full = test_pred_full/len(testloader) 144 | test_recons_full = test_recons_full/len(testloader) 145 | print("Total prediction test mse error is ",test_pred_full) 146 | print("Total reconstruction test mse error is ",test_recons_full) 147 | return test_pred_full 148 | 149 | 150 | def train(self, epochs, trainloader, step = 1, T_out = 40, evalloader = False): 151 | T_eval = T_out 152 | for ep in range(epochs): 153 | self.kernel.train() 154 | t1 = default_timer() 155 | train_recons_full = 0 156 | train_pred_full = 0 157 | for xx, yy in trainloader: 158 | l_recons = 0 159 | xx = xx.to(self.device) 160 | yy = yy.to(self.device) 161 | bs = xx.shape[0] 162 | for t in range(0, T_out): 163 | y = yy[..., t:t + 1] 164 | 165 | im,im_re = self.kernel(xx) 166 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 167 | if t == 0: 168 | pred = im[...,-1:] 169 | else: 170 | pred = torch.cat((pred, im[...,-1:]), -1) 171 | 172 | xx = torch.cat((xx[..., step:], im[...,-1:]), dim=-1) 173 | 174 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 175 | loss = 5 * l_pred + 0.5 * l_recons 176 | 177 | train_pred_full += l_pred.item() 178 | train_recons_full += l_recons.item()/T_out 179 | 180 | self.optimizer.zero_grad() 181 | loss.backward() 182 | self.optimizer.step() 183 | train_pred_full = train_pred_full / len(trainloader) 184 | train_recons_full = train_recons_full / len(trainloader) 185 | t2 = default_timer() 186 | test_pred_full = 0 187 | test_recons_full = 0 188 | loc = 0 189 | mse_error = 0 190 | if evalloader: 191 | with torch.no_grad(): 192 | for xx, yy in evalloader: 193 | loss = 0 194 | xx = xx.to(self.device) 195 | yy = yy.to(self.device) 196 | 197 | for t in range(0, T_eval): 198 | y = yy[..., t:t + 1] 199 | im, im_re = self.kernel(xx) 200 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 201 | if t == 0: 202 | pred = im[...,-1:] 203 | else: 204 | pred = torch.cat((pred, im[...,-1:]), -1) 205 | xx = torch.cat((xx[..., 1:], im[...,-1:]), dim=-1) 206 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 207 | 208 | test_recons_full += l_recons.item() / T_eval 209 | test_pred_full += l_pred.item() 210 | 211 | loc = loc + 1 212 | mse_error = mse_error / loc 213 | test_recons_full = test_recons_full / len(evalloader) 214 | test_pred_full = test_pred_full / len(evalloader) 215 | self.scheduler.step() 216 | 217 | if evalloader: 218 | if ep == 0: 219 | print("Epoch","Time","[Train Recons MSE]","[Train Pred MSE]","[Eval Recons MSE]","[Eval Pred MSE]") 220 | print(ep, t2 - t1, train_recons_full, train_pred_full, test_recons_full, test_pred_full) 221 | else: 222 | if ep == 0: 223 | print("Epoch","Time","Train Recons MSE","Train Pred MSE") 224 | print(ep, t2 - t1, train_recons_full, train_pred_full) 225 | def test(self, testloader, step = 1, T_out = 40, path = False, is_save = False, is_plot = False): 226 | time_error = torch.zeros([T_out,1]) 227 | test_pred_full = 0 228 | test_recons_full = 0 229 | loc = 0 230 | with torch.no_grad(): 231 | for xx, yy in testloader: 232 | loss = 0 233 | bs = xx.shape[0] 234 | xx = xx.to(self.device) 235 | yy = yy.to(self.device) 236 | l_recons = 0 237 | for t in range(0, T_out): 238 | y = yy[..., t:t + 1] 239 | im, im_re = self.kernel(xx) 240 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 241 | t_error = self.loss(im[...,-1:],y) 242 | if t == 0: 243 | pred = im[...,-1:] 244 | else: 245 | pred = torch.cat((pred, im[...,-1:]), -1) 246 | time_error[t] = time_error[t] + t_error.item() 247 | xx = torch.cat((xx[..., 1:], im[...,-1:]), dim=-1) 248 | 249 | test_recons_full += l_recons.item() / T_out 250 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 251 | test_pred_full += l_pred.item() 252 | if(loc == 0 & is_save): 253 | torch.save({"pred":pred, "yy":yy}, path+ "pred_yy.pt") 254 | 255 | if(loc == 0 & is_plot): 256 | for i in range(T_out): 257 | plt.subplot(1,3,1) 258 | plt.title("Predict") 259 | plt.imshow(pred[0,...,i].cpu().detach().numpy()) 260 | plt.subplot(1,3,2) 261 | plt.imshow(yy[0,...,i].cpu().detach().numpy()) 262 | plt.title("Label") 263 | plt.subplot(1,3,3) 264 | plt.imshow(pred[0,...,i].cpu().detach().numpy()-yy[0,...,i].cpu().detach().numpy()) 265 | plt.title("Error") 266 | plt.show() 267 | plt.savefig(path + "time_"+str(i)+".png") 268 | plt.close() 269 | loc = loc + 1 270 | test_pred_full = test_pred_full / loc 271 | test_recons_full = test_recons_full / loc 272 | time_error = time_error / len(testloader) 273 | print("Total prediction test mse error is ",test_pred_full) 274 | print("Total reconstruction test mse error is ",test_recons_full) 275 | return time_error 276 | 277 | def save(self, path): 278 | (fpath,_) = os.path.split(path) 279 | if not os.path.isfile(fpath): 280 | os.makedirs(fpath) 281 | torch.save({"koopman":self,"model":self.kernel,"model_params":self.kernel.state_dict()}, path) 282 | 283 | class koopman_vit: 284 | def __init__(self, decoder = "Conv2d", depth = 16, resolution=(256, 256), patch_size=(4, 4), 285 | in_chans=1, out_chans=1, embed_dim=768, parallel = False, device = False): 286 | # Model Hyper-parameters 287 | self.decoder = decoder 288 | self.resolution = resolution 289 | self.patch_size = patch_size 290 | self.in_chans = in_chans 291 | self.out_chans = out_chans 292 | self.embed_dim = embed_dim 293 | self.depth = depth 294 | # Core Model 295 | self.params = 0 296 | self.kernel = False 297 | # Opt Setting 298 | self.optimizer = False 299 | self.scheduler = False 300 | self.device = device 301 | self.parallel = parallel 302 | self.loss = torch.nn.MSELoss() 303 | def compile(self): 304 | self.kernel = koopmanViT.ViT(img_size=self.resolution, patch_size=self.patch_size, in_chans=self.in_chans, out_chans=self.out_chans, num_blocks=self.num_blocks, embed_dim = self.embed_dim, depth=self.depth, settings = self.decoder).to(self.device) 305 | if self.parallel: 306 | self.kernel = torch.nn.DataParallel(self.kernel) 307 | self.params = utils.count_params(self.kernel) 308 | 309 | print("Koopman Fourier Vision Transformer has been compiled!") 310 | print("The Model Parameters Number is ",self.params) 311 | 312 | def opt_init(self, opt, lr, step_size, gamma): 313 | if opt == "Adam": 314 | self.optimizer = utils.Adam(self.kernel.parameters(), lr= lr, weight_decay=1e-4) 315 | if not step_size == False: 316 | self.scheduler = torch.optim.lr_scheduler.StepLR(self.optimizer, step_size=step_size, gamma=gamma) 317 | 318 | def train_multi(self, epochs, trainloader, T_out = 10, evalloader = False): 319 | T_eval = T_out 320 | for ep in range(epochs): 321 | self.kernel.train() 322 | t1 = default_timer() 323 | train_recons_full = 0 324 | train_pred_full = 0 325 | for xx, yy in trainloader: 326 | l_recons = 0 327 | xx = xx.to(self.device) # [batchsize,1,x,y] 328 | yy = yy.to(self.device) # [batchsize,T,x,y] 329 | bs = xx.shape[0] 330 | for t in range(0, T_out): 331 | y = yy[:, t:t + 1] 332 | im,im_re = self.kernel(xx) 333 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 334 | 335 | if t == 0: 336 | pred = im[:, -1:] 337 | else: 338 | pred = torch.cat((pred, im[:, -1:]), -1) 339 | 340 | xx = im 341 | 342 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 343 | loss = 5 * l_pred + 0.5 * l_recons 344 | 345 | train_pred_full += l_pred.item() 346 | train_recons_full += l_recons.item()/T_out 347 | 348 | self.optimizer.zero_grad() 349 | loss.backward() 350 | self.optimizer.step() 351 | train_pred_full = train_pred_full / len(trainloader) 352 | train_recons_full = train_recons_full / len(trainloader) 353 | t2 = default_timer() 354 | test_pred_full = 0 355 | test_recons_full = 0 356 | loc = 0 357 | mse_error = 0 358 | if evalloader: 359 | with torch.no_grad(): 360 | for xx, yy in evalloader: 361 | loss = 0 362 | xx = xx.to(self.device) 363 | yy = yy.to(self.device) 364 | 365 | for t in range(0, T_eval): 366 | y = yy[:, t:t + 1] 367 | im, im_re = self.kernel(xx) 368 | 369 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 370 | 371 | if t == 0: 372 | pred = im 373 | else: 374 | pred = torch.cat((pred, im), 1) 375 | 376 | xx = im 377 | 378 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 379 | 380 | test_recons_full += l_recons.item() / T_eval 381 | test_pred_full += l_pred.item() 382 | 383 | test_recons_full = test_recons_full / len(evalloader) 384 | test_pred_full = test_pred_full / len(evalloader) 385 | self.scheduler.step() 386 | 387 | if evalloader: 388 | if ep == 0: 389 | print("Epoch","Time","[Train Recons MSE]","[Train Pred MSE]","[Eval Recons MSE]","[Eval Pred MSE]") 390 | print(ep, t2 - t1, train_recons_full, train_pred_full, test_recons_full, test_pred_full) 391 | else: 392 | if ep == 0: 393 | print("Epoch","Time","Train Recons MSE","Train Pred MSE") 394 | print(ep, t2 - t1, train_recons_full, train_pred_full) 395 | 396 | def test_multi(self, testloader, step = 1, T_out = 5, path = False, is_save = False, is_plot = False): 397 | time_error = torch.zeros([T_out,1]) 398 | test_pred_full = 0 399 | test_recons_full = 0 400 | loc = 0 401 | with torch.no_grad(): 402 | for xx, yy in testloader: 403 | loss = 0 404 | bs = xx.shape[0] 405 | xx = xx.to(self.device) 406 | yy = yy.to(self.device) 407 | l_recons = 0 408 | for t in range(0, T_out): 409 | y = yy[:, t:t + 1] 410 | im, im_re = self.kernel(xx) 411 | 412 | 413 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 414 | t_error = self.loss(im, y) 415 | 416 | xx = im 417 | 418 | if t == 0: 419 | pred = im 420 | else: 421 | pred = torch.cat((pred, im), 1) 422 | time_error[t] = time_error[t] + t_error.item() 423 | 424 | test_recons_full += l_recons.item() / T_out 425 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 426 | test_pred_full += l_pred.item() 427 | 428 | if(loc == 0 & is_save): 429 | torch.save({"pred":pred, "yy":yy}, path+ "pred_yy.pt") 430 | 431 | if(loc == 0 & is_plot): 432 | for i in range(T_out): 433 | plt.subplot(1,3,1) 434 | plt.title("Predict") 435 | plt.imshow(pred[0,i].cpu().detach().numpy()) 436 | plt.subplot(1,3,2) 437 | plt.imshow(yy[0,i].cpu().detach().numpy()) 438 | plt.title("Label") 439 | plt.subplot(1,3,3) 440 | plt.imshow(pred[0,i].cpu().detach().numpy()-yy[0,i].cpu().detach().numpy()) 441 | plt.title("Error") 442 | plt.show() 443 | plt.savefig(path + "time_"+str(i)+".png") 444 | plt.close() 445 | 446 | loc = loc + 1 447 | test_pred_full = test_pred_full / loc 448 | test_recons_full = test_recons_full / loc 449 | time_error = time_error / len(testloader) 450 | print("Total prediction test mse error is ",test_pred_full) 451 | print("Total reconstruction test mse error is ",test_recons_full) 452 | return time_error 453 | 454 | 455 | def train_single(self, epochs, trainloader, evalloader = False): 456 | for ep in range(epochs): 457 | self.kernel.train() 458 | t1 = default_timer() 459 | train_recons_full = 0 460 | train_pred_full = 0 461 | for x, y in trainloader: 462 | l_recons = 0 463 | x = x.to(self.device) # [batchsize,1,64,64] 464 | y = y.to(self.device) # [batchsize,1,64,64] 465 | bs = x.shape[0] 466 | 467 | im,im_re = self.kernel(x) 468 | 469 | l_recons = self.loss(im_re.reshape(bs, -1), x.reshape(bs, -1)) 470 | l_pred = self.loss(im.reshape(bs, -1), y.reshape(bs, -1)) 471 | 472 | loss = 5 * l_pred + 0.5 * l_recons 473 | 474 | train_pred_full += l_pred.item() 475 | train_recons_full += l_recons.item() 476 | 477 | self.optimizer.zero_grad() 478 | loss.backward() 479 | self.optimizer.step() 480 | train_pred_full = train_pred_full / len(trainloader) 481 | train_recons_full = train_recons_full / len(trainloader) 482 | t2 = default_timer() 483 | test_pred_full = 0 484 | test_recons_full = 0 485 | loc = 0 486 | mse_error = 0 487 | if evalloader: 488 | with torch.no_grad(): 489 | for x, y in evalloader: 490 | loss = 0 491 | x = x.to(self.device) 492 | y = y.to(self.device) 493 | 494 | im, im_re = self.kernel(x) 495 | 496 | l_recons = self.loss(im_re.reshape(bs, -1), x.reshape(bs, -1)) 497 | l_pred = self.loss(im.reshape(bs, -1), y.reshape(bs, -1)) 498 | 499 | test_recons_full += l_recons.item() 500 | test_pred_full += l_pred.item() 501 | 502 | test_recons_full = test_recons_full / len(evalloader) 503 | test_pred_full = test_pred_full / len(evalloader) 504 | self.scheduler.step() 505 | 506 | if evalloader: 507 | if ep == 0: 508 | print("Epoch","Time","[Train Recons MSE]","[Train Pred MSE]","[Eval Recons MSE]","[Eval Pred MSE]") 509 | print(ep, t2 - t1, train_recons_full, train_pred_full, test_recons_full, test_pred_full) 510 | else: 511 | if ep == 0: 512 | print("Epoch","Time","Train Recons MSE","Train Pred MSE") 513 | print(ep, t2 - t1, train_recons_full, train_pred_full) 514 | 515 | def test_single(self, testloader, T_out = 1, path = False, is_save = False, is_plot = False): 516 | time_error = torch.zeros([T_out,1]) 517 | test_pred_full = 0 518 | test_recons_full = 0 519 | loc = 0 520 | idx = np.random.randint(0,len(testloader)) 521 | with torch.no_grad(): 522 | for xx, yy in testloader: 523 | loss = 0 524 | bs = xx.shape[0] 525 | xx = xx.to(self.device) 526 | yy = yy.to(self.device) 527 | l_recons = 0 528 | for t in range(0, T_out): 529 | y = yy[:, t:t + 1] 530 | im, im_re = self.kernel(xx) 531 | 532 | l_recons += self.loss(im_re.reshape(bs, -1), xx.reshape(bs, -1)) 533 | t_error = self.loss(im, y) 534 | 535 | xx = im 536 | 537 | if t == 0: 538 | pred = im 539 | else: 540 | pred = torch.cat((pred, im), 1) 541 | time_error[t] = time_error[t] + t_error.item() 542 | 543 | test_recons_full += l_recons.item() / T_out 544 | l_pred = self.loss(pred.reshape(bs, -1), yy.reshape(bs, -1)) 545 | test_pred_full += l_pred.item() 546 | 547 | if(loc == 0 & is_save): 548 | torch.save({"pred":pred, "yy":yy}, path+ "pred_yy.pt") 549 | 550 | if(loc == 0 & is_plot): 551 | for i in range(T_out): 552 | plt.subplot(1,3,1) 553 | plt.title("Predict") 554 | plt.imshow(pred[0,i].cpu().detach().numpy()) 555 | plt.subplot(1,3,2) 556 | plt.imshow(yy[0,i].cpu().detach().numpy()) 557 | plt.title("Label") 558 | plt.subplot(1,3,3) 559 | plt.imshow(pred[0,i].cpu().detach().numpy()-yy[0,i].cpu().detach().numpy()) 560 | plt.title("Error") 561 | plt.show() 562 | plt.savefig(path + "time_"+str(i)+".png") 563 | plt.close() 564 | loc = loc + 1 565 | 566 | test_pred_full = test_pred_full / len(testloader) 567 | test_recons_full = test_recons_full / len(testloader) 568 | time_error = time_error / len(testloader) 569 | print("Total prediction test mse error is ",test_pred_full) 570 | print("Total reconstruction test mse error is ",test_recons_full) 571 | 572 | return time_error 573 | 574 | def save(self, path): 575 | # (fpath,_) = os.path.split(path) 576 | # print(fpath, os.path.isfile(fpath)) 577 | # if not os.path.isfile(fpath): 578 | # os.makedirs(fpath) 579 | torch.save({"koopman":self,"model":self.kernel,"model_params":self.kernel.state_dict()}, path) 580 | -------------------------------------------------------------------------------- /koopmanlab/models/__init__.py: -------------------------------------------------------------------------------- 1 | from .kno import KNO1d,KNO2d 2 | from .kno import encoder_mlp, decoder_mlp, encoder_conv1d, decoder_conv1d, encoder_conv2d, decoder_conv2d 3 | from .koopmanViT import ViT -------------------------------------------------------------------------------- /koopmanlab/models/kno.py: -------------------------------------------------------------------------------- 1 | import torch 2 | import numpy as np 3 | import torch.nn as nn 4 | import torch.nn.functional as F 5 | 6 | torch.manual_seed(0) 7 | 8 | # The structure of Auto-Encoder 9 | class encoder_mlp(nn.Module): 10 | def __init__(self, t_len, op_size): 11 | super(encoder_mlp, self).__init__() 12 | self.layer = nn.Linear(t_len, op_size) 13 | def forward(self, x): 14 | x = self.layer(x) 15 | return x 16 | 17 | class decoder_mlp(nn.Module): 18 | def __init__(self, t_len, op_size): 19 | super(decoder_mlp, self).__init__() 20 | self.layer = nn.Linear(op_size, t_len) 21 | def forward(self, x): 22 | x = self.layer(x) 23 | return x 24 | 25 | class encoder_conv1d(nn.Module): 26 | def __init__(self, t_len, op_size): 27 | super(encoder_conv1d, self).__init__() 28 | self.layer = nn.Conv1d(t_len, op_size,1) 29 | def forward(self, x): 30 | x = x.permute([0,2,1]) 31 | x = self.layer(x) 32 | x = x.permute([0,2,1]) 33 | return x 34 | 35 | class decoder_conv1d(nn.Module): 36 | def __init__(self, t_len, op_size): 37 | super(decoder_conv1d, self).__init__() 38 | self.layer = nn.Conv1d(op_size, t_len,1) 39 | def forward(self, x): 40 | x = x.permute([0,2,1]) 41 | x = self.layer(x) 42 | x = x.permute([0,2,1]) 43 | return x 44 | 45 | class encoder_conv2d(nn.Module): 46 | def __init__(self, t_len, op_size): 47 | super(encoder_conv2d, self).__init__() 48 | self.layer = nn.Conv2d(t_len, op_size,1) 49 | def forward(self, x): 50 | x = x.permute([0,3,1,2]) 51 | x = self.layer(x) 52 | x = x.permute([0,2,3,1]) 53 | return x 54 | 55 | class decoder_conv2d(nn.Module): 56 | def __init__(self, t_len, op_size): 57 | super(decoder_conv2d, self).__init__() 58 | self.layer = nn.Conv2d(op_size, t_len,1) 59 | def forward(self, x): 60 | x = x.permute([0,3,1,2]) 61 | x = self.layer(x) 62 | x = x.permute([0,2,3,1]) 63 | return x 64 | 65 | # Koopman 1D structure 66 | class Koopman_Operator1D(nn.Module): 67 | def __init__(self, op_size, modes_x = 16): 68 | super(Koopman_Operator1D, self).__init__() 69 | self.op_size = op_size 70 | self.scale = (1 / (op_size * op_size)) 71 | self.modes_x = modes_x 72 | self.koopman_matrix = nn.Parameter(self.scale * torch.rand(op_size, op_size, self.modes_x, dtype=torch.cfloat)) 73 | # Complex multiplication 74 | def time_marching(self, input, weights): 75 | # (batch, t, x), (t, t+1, x) -> (batch, t+1, x) 76 | return torch.einsum("btx,tfx->bfx", input, weights) 77 | def forward(self, x): 78 | batchsize = x.shape[0] 79 | # Fourier Transform 80 | x_ft = torch.fft.rfft(x) 81 | # Koopman Operator Time Marching 82 | out_ft = torch.zeros(x_ft.shape, dtype=torch.cfloat, device = x.device) 83 | out_ft[:, :, :self.modes_x] = self.time_marching(x_ft[:, :, :self.modes_x], self.koopman_matrix) 84 | #Inverse Fourier Transform 85 | x = torch.fft.irfft(out_ft, n=x.size(-1)) 86 | return x 87 | 88 | class KNO1d(nn.Module): 89 | def __init__(self, encoder, decoder, op_size, modes_x = 16, decompose = 4, linear_type = True, normalization = False): 90 | super(KNO1d, self).__init__() 91 | # Parameter 92 | self.op_size = op_size 93 | self.decompose = decompose 94 | # Layer Structure 95 | self.enc = encoder 96 | self.dec = decoder 97 | self.koopman_layer = Koopman_Operator1D(self.op_size, modes_x = modes_x) 98 | self.w0 = nn.Conv1d(op_size, op_size, 1) 99 | self.linear_type = linear_type # If this variable is False, activate function will be worked after Koopman Matrix 100 | self.normalization = normalization 101 | if self.normalization: 102 | self.norm_layer = torch.nn.BatchNorm2d(op_size) 103 | def forward(self, x): 104 | # Reconstruct 105 | x_reconstruct = self.enc(x) 106 | x_reconstruct = torch.tanh(x_reconstruct) 107 | x_reconstruct = self.dec(x_reconstruct) 108 | # Predict 109 | x = self.enc(x) # Encoder 110 | x = torch.tanh(x) 111 | x = x.permute(0, 2, 1) 112 | x_w = x 113 | for i in range(self.decompose): 114 | x1 = self.koopman_layer(x) # Koopman Operator 115 | if self.linear_type: 116 | x = x + x1 117 | else: 118 | x = torch.tanh(x + x1) 119 | if self.normalization: 120 | x = torch.tanh(self.norm_layer(self.w0(x_w)) + x) 121 | else: 122 | x = torch.tanh(self.w0(x_w) + x) 123 | x = x.permute(0, 2, 1) 124 | x = self.dec(x) # Decoder 125 | return x, x_reconstruct 126 | 127 | # Koopman 2D structure 128 | class Koopman_Operator2D(nn.Module): 129 | def __init__(self, op_size, modes_x, modes_y): 130 | super(Koopman_Operator2D, self).__init__() 131 | self.op_size = op_size 132 | self.scale = (1 / (op_size * op_size)) 133 | self.modes_x = modes_x 134 | self.modes_y = modes_y 135 | self.koopman_matrix = nn.Parameter(self.scale * torch.rand(op_size, op_size, self.modes_x, self.modes_y, dtype=torch.cfloat)) 136 | 137 | # Complex multiplication 138 | def time_marching(self, input, weights): 139 | # (batch, t, x,y ), (t, t+1, x,y) -> (batch, t+1, x,y) 140 | return torch.einsum("btxy,tfxy->bfxy", input, weights) 141 | 142 | def forward(self, x): 143 | batchsize = x.shape[0] 144 | # Fourier Transform 145 | x_ft = torch.fft.rfft2(x) 146 | # Koopman Operator Time Marching 147 | out_ft = torch.zeros(x_ft.shape, dtype=torch.cfloat, device = x.device) 148 | out_ft[:, :, :self.modes_x, :self.modes_y] = self.time_marching(x_ft[:, :, :self.modes_x, :self.modes_y], self.koopman_matrix) 149 | out_ft[:, :, -self.modes_x:, :self.modes_y] = self.time_marching(x_ft[:, :, -self.modes_x:, :self.modes_y], self.koopman_matrix) 150 | #Inverse Fourier Transform 151 | x = torch.fft.irfft2(out_ft, s=(x.size(-2), x.size(-1))) 152 | return x 153 | 154 | class KNO2d(nn.Module): 155 | def __init__(self, encoder, decoder, op_size, modes_x = 10, modes_y = 10, decompose = 6, linear_type = True, normalization = False): 156 | super(KNO2d, self).__init__() 157 | # Parameter 158 | self.op_size = op_size 159 | self.decompose = decompose 160 | self.modes_x = modes_x 161 | self.modes_y = modes_y 162 | # Layer Structure 163 | self.enc = encoder 164 | self.dec = decoder 165 | self.koopman_layer = Koopman_Operator2D(self.op_size, self.modes_x, self.modes_y) 166 | self.w0 = nn.Conv2d(op_size, op_size, 1) 167 | self.linear_type = linear_type # If this variable is False, activate function will be worked after Koopman Matrix 168 | self.normalization = normalization 169 | if self.normalization: 170 | self.norm_layer = torch.nn.BatchNorm2d(op_size) 171 | def forward(self, x): 172 | # Reconstruct 173 | x_reconstruct = self.enc(x) 174 | x_reconstruct = torch.tanh(x_reconstruct) 175 | x_reconstruct = self.dec(x_reconstruct) 176 | # Predict 177 | x = self.enc(x) # Encoder 178 | x = torch.tanh(x) 179 | x = x.permute(0, 3, 1, 2) 180 | x_w = x 181 | for i in range(self.decompose): 182 | x1 = self.koopman_layer(x) # Koopman Operator 183 | if self.linear_type: 184 | x = x + x1 185 | else: 186 | x = torch.tanh(x + x1) 187 | if self.normalization: 188 | x = torch.tanh(self.norm_layer(self.w0(x_w)) + x) 189 | else: 190 | x = torch.tanh(self.w0(x_w) + x) 191 | x = x.permute(0, 2, 3, 1) 192 | x = self.dec(x) # Decoder 193 | return x, x_reconstruct 194 | -------------------------------------------------------------------------------- /koopmanlab/models/koopmanViT.py: -------------------------------------------------------------------------------- 1 | import math 2 | from functools import partial 3 | from collections import OrderedDict 4 | from copy import Error, deepcopy 5 | from re import S 6 | from numpy.lib.arraypad import pad 7 | import numpy as np 8 | import torch 9 | import torch.nn as nn 10 | import torch.nn.functional as F 11 | from timm.models.layers import DropPath, trunc_normal_ 12 | import torch.fft 13 | from torch.nn.modules.container import Sequential 14 | from torch.utils.checkpoint import checkpoint_sequential 15 | from einops import rearrange, repeat 16 | from einops.layers.torch import Rearrange 17 | 18 | class PatchEmbed(nn.Module): 19 | def __init__(self, img_size=(224, 224), patch_size=(16, 16), in_chans=3, embed_dim=768): 20 | super().__init__() 21 | num_patches = (img_size[1] // patch_size[1]) * (img_size[0] // patch_size[0]) 22 | self.img_size = img_size 23 | self.patch_size = patch_size 24 | self.num_patches = num_patches 25 | self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size) 26 | 27 | def forward(self, x): 28 | B, C, H, W = x.shape 29 | assert H == self.img_size[0] and W == self.img_size[1], f"Input image size ({H}*{W}) doesn't match model ({self.img_size[0]}*{self.img_size[1]})." 30 | x = self.proj(x).flatten(2).transpose(1, 2) 31 | return x 32 | 33 | class Mlp(nn.Module): 34 | def __init__(self, in_features, hidden_features=None, out_features=None, act_layer=nn.GELU, drop=0.): 35 | super().__init__() 36 | out_features = out_features or in_features 37 | hidden_features = hidden_features or in_features 38 | self.fc1 = nn.Linear(in_features, hidden_features) 39 | self.act = act_layer() 40 | self.fc2 = nn.Linear(hidden_features, out_features) 41 | self.drop = nn.Dropout(drop) 42 | 43 | def forward(self, x): 44 | x = self.fc1(x) 45 | x = self.act(x) 46 | x = self.drop(x) 47 | x = self.fc2(x) 48 | x = self.drop(x) 49 | return x 50 | 51 | # Use Fourier-Transformer structure to approximate linear Koopman Operator 52 | class Attention(nn.Module): 53 | def __init__( 54 | self, 55 | dim, 56 | num_heads=8, 57 | qkv_bias=False, 58 | qk_scale=None, 59 | attn_drop=0.0, 60 | proj_drop=0.0, 61 | input_size=(4, 14, 14), 62 | ): 63 | super().__init__() 64 | assert dim % num_heads == 0, "dim should be divisible by num_heads" 65 | self.num_heads = num_heads 66 | head_dim = dim // num_heads 67 | self.scale = qk_scale or head_dim**-0.5 68 | 69 | self.q = nn.Linear(dim, dim, bias=qkv_bias) 70 | self.k = nn.Linear(dim, dim, bias=qkv_bias) 71 | self.v = nn.Linear(dim, dim, bias=qkv_bias) 72 | assert attn_drop == 0.0 # do not use 73 | self.proj = nn.Linear(dim, dim) 74 | self.proj_drop = nn.Dropout(proj_drop) 75 | self.input_size = input_size 76 | assert input_size[1] == input_size[2] 77 | 78 | def forward(self, x): 79 | B, N, C = x.shape 80 | q = ( 81 | self.q(x) 82 | .reshape(B, N, self.num_heads, C // self.num_heads) 83 | .permute(0, 2, 1, 3) 84 | ) 85 | k = ( 86 | self.k(x) 87 | .reshape(B, N, self.num_heads, C // self.num_heads) 88 | .permute(0, 2, 1, 3) 89 | ) 90 | v = ( 91 | self.v(x) 92 | .reshape(B, N, self.num_heads, C // self.num_heads) 93 | .permute(0, 2, 1, 3) 94 | ) 95 | 96 | attn = (q @ k.transpose(-2, -1)) * self.scale 97 | 98 | attn = attn.softmax(dim=-1) 99 | 100 | x = (attn @ v).transpose(1, 2).reshape(B, N, C) 101 | x = self.proj(x) 102 | x = self.proj_drop(x) 103 | x = x.view(B, -1, C) 104 | return x 105 | 106 | 107 | class At_Block(nn.Module): 108 | """ 109 | Transformer Block in Fourier Domain 110 | """ 111 | def __init__( 112 | self, 113 | dim=768, 114 | num_heads=8, 115 | mlp_ratio=4.0, 116 | qkv_bias=False, 117 | qk_scale=None, 118 | drop=0.0, 119 | attn_drop=0.0, 120 | drop_path=0.0, 121 | act_layer=nn.GELU, 122 | norm_layer=nn.LayerNorm, 123 | attn_func=Attention, 124 | ): 125 | super().__init__() 126 | self.norm1 = norm_layer(dim) 127 | self.attn = attn_func( 128 | dim, 129 | num_heads=num_heads, 130 | qkv_bias=qkv_bias, 131 | qk_scale=qk_scale, 132 | attn_drop=attn_drop, 133 | proj_drop=drop, 134 | ) 135 | 136 | self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity() 137 | self.norm2 = norm_layer(dim) 138 | mlp_hidden_dim = int(dim * mlp_ratio) 139 | self.mlp = Mlp( 140 | in_features=dim, 141 | hidden_features=mlp_hidden_dim, 142 | act_layer=act_layer, 143 | drop=drop, 144 | ) 145 | 146 | def forward(self, x): 147 | x_ft = torch.fft.fft(x, dim=-1, norm="ortho") 148 | x_r = x_ft.real 149 | x_i = x_ft.imag 150 | # Calculate Real Part 151 | x_r = x_r + self.drop_path(self.attn(self.norm1(x_r))) 152 | x_r = x_r + self.drop_path(self.mlp(self.norm2(x_r))) 153 | # Calculate Imaginary Part 154 | x_i = x_i + self.drop_path(self.attn(self.norm1(x_i))) 155 | x_i = x_i + self.drop_path(self.mlp(self.norm2(x_i))) 156 | # Merge 157 | x_ft.real = x_r 158 | x_ft.imag = x_i 159 | x = torch.fft.ifft(x_ft, dim=-1, norm="ortho") 160 | return x 161 | 162 | # Use linear AFNO1D structure to approximate linear Koopman Operator 163 | class AFNO1D(nn.Module): 164 | def __init__(self, hidden_size, num_blocks=8, sparsity_threshold=0.01, hard_thresholding_fraction=1, hidden_size_factor=1): 165 | super().__init__() 166 | assert hidden_size % num_blocks == 0, f"hidden_size {hidden_size} should be divisble by num_blocks {num_blocks}" 167 | 168 | self.hidden_size = hidden_size 169 | self.sparsity_threshold = sparsity_threshold 170 | self.num_blocks = num_blocks 171 | self.block_size = self.hidden_size // self.num_blocks 172 | self.hard_thresholding_fraction = hard_thresholding_fraction 173 | self.hidden_size_factor = hidden_size_factor 174 | self.scale = 0.02 175 | 176 | self.w1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size, self.block_size * self.hidden_size_factor)) 177 | self.b1 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor)) 178 | self.w2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size * self.hidden_size_factor, self.block_size)) 179 | self.b2 = nn.Parameter(self.scale * torch.randn(2, self.num_blocks, self.block_size)) 180 | 181 | def forward(self, x): 182 | bias = x 183 | 184 | dtype = x.dtype 185 | x = x.float() 186 | B, N, C = x.shape 187 | 188 | x = torch.fft.rfft(x, dim=1, norm="ortho") 189 | x = x.reshape(B, N // 2 + 1, self.num_blocks, self.block_size) 190 | 191 | o1_real = torch.zeros([B, N // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) 192 | o1_imag = torch.zeros([B, N // 2 + 1, self.num_blocks, self.block_size * self.hidden_size_factor], device=x.device) 193 | o2_real = torch.zeros(x.shape, device=x.device) 194 | o2_imag = torch.zeros(x.shape, device=x.device) 195 | 196 | total_modes = N // 2 + 1 197 | kept_modes = int(total_modes * self.hard_thresholding_fraction) 198 | 199 | o1_real[:, :kept_modes] = F.relu( 200 | torch.einsum('...bi,bio->...bo', x[:, :kept_modes].real, self.w1[0]) - \ 201 | torch.einsum('...bi,bio->...bo', x[:, :kept_modes].imag, self.w1[1]) + \ 202 | self.b1[0] 203 | ) 204 | 205 | o1_imag[:, :kept_modes] = F.relu( 206 | torch.einsum('...bi,bio->...bo', x[:, :kept_modes].imag, self.w1[0]) + \ 207 | torch.einsum('...bi,bio->...bo', x[:, :kept_modes].real, self.w1[1]) + \ 208 | self.b1[1] 209 | ) 210 | 211 | o2_real[:, :kept_modes] = ( 212 | torch.einsum('...bi,bio->...bo', o1_real[:, :kept_modes], self.w2[0]) - \ 213 | torch.einsum('...bi,bio->...bo', o1_imag[:, :kept_modes], self.w2[1]) + \ 214 | self.b2[0] 215 | ) 216 | 217 | o2_imag[:, :kept_modes] = ( 218 | torch.einsum('...bi,bio->...bo', o1_imag[:, :kept_modes], self.w2[0]) + \ 219 | torch.einsum('...bi,bio->...bo', o1_real[:, :kept_modes], self.w2[1]) + \ 220 | self.b2[1] 221 | ) 222 | 223 | x = torch.stack([o2_real, o2_imag], dim=-1) 224 | x = F.softshrink(x, lambd=self.sparsity_threshold) 225 | x = torch.view_as_complex(x) 226 | x = x.reshape(B, N // 2 + 1, C) 227 | x = torch.fft.irfft(x, n=N, dim=1, norm="ortho") 228 | x = x.type(dtype) 229 | return x + bias 230 | 231 | class Af_Block(nn.Module): 232 | """ 233 | AdaptiveFNO Block 234 | """ 235 | def __init__( 236 | self, 237 | dim, 238 | mlp_ratio=4., 239 | drop=0., 240 | drop_path=0., 241 | act_layer=nn.GELU, 242 | norm_layer=nn.LayerNorm, 243 | double_skip=True, 244 | num_blocks=8, 245 | sparsity_threshold=0.01, 246 | hard_thresholding_fraction=1.0, 247 | ): 248 | super().__init__() 249 | self.norm1 = norm_layer(dim) 250 | self.filter = AFNO1D(dim, num_blocks, sparsity_threshold, hard_thresholding_fraction) 251 | self.drop_path = DropPath(drop_path) if drop_path > 0. else nn.Identity() 252 | #self.drop_path = nn.Identity() 253 | self.norm2 = norm_layer(dim) 254 | mlp_hidden_dim = int(dim * mlp_ratio) 255 | self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden_dim, act_layer=act_layer, drop=drop) 256 | self.double_skip = double_skip 257 | 258 | def forward(self, x): 259 | residual = x 260 | x = self.norm1(x) 261 | x = self.filter(x) 262 | 263 | if self.double_skip: 264 | x = x + residual 265 | residual = x 266 | 267 | x = self.norm2(x) 268 | x = self.mlp(x) 269 | x = self.drop_path(x) 270 | x = x + residual 271 | return x 272 | 273 | 274 | class ViT(nn.Module): 275 | def __init__( 276 | self, 277 | img_size=(720, 1440), 278 | patch_size=(8, 8), 279 | in_chans=20, 280 | out_chans=20, 281 | embed_dim=768, 282 | encoder_depth = 2, 283 | depth=10, 284 | mlp_ratio=4., 285 | drop_rate=0., 286 | drop_path_rate=0., 287 | num_blocks=16, 288 | sparsity_threshold=0.01, 289 | hard_thresholding_fraction=1.0, 290 | settings = "Conv2d", 291 | encoder_network = False 292 | ): 293 | super().__init__() 294 | self.encoder_network = encoder_network 295 | self.img_size = img_size 296 | self.patch_size = patch_size 297 | self.in_chans = in_chans 298 | self.out_chans = out_chans 299 | 300 | self.num_features = self.embed_dim = embed_dim 301 | self.num_blocks = num_blocks 302 | self.depth = depth 303 | self.settings = settings 304 | norm_layer = partial(nn.LayerNorm, eps=1e-6) 305 | 306 | self.patch_embed = PatchEmbed(img_size=self.img_size, patch_size=self.patch_size, in_chans=self.in_chans, embed_dim=embed_dim) 307 | num_patches = self.patch_embed.num_patches 308 | 309 | self.pos_embed = nn.Parameter(torch.zeros(1, num_patches, embed_dim)) 310 | self.pos_drop = nn.Dropout(p=drop_rate) 311 | 312 | dpr = [x.item() for x in torch.linspace(0, drop_path_rate, self.depth)] 313 | self.dpr = dpr 314 | 315 | self.h = img_size[0] // self.patch_size[0] 316 | self.w = img_size[1] // self.patch_size[1] 317 | 318 | 319 | # Encoder Settings 320 | self.encoder_depth = encoder_depth 321 | # There are two options. Af_Block represents using the AdaptiveFNO blocks, and At_Block represents using the Fourier-Transformer blocks. 322 | self.encoder_blocks = nn.ModuleList([ 323 | Af_Block(dim=embed_dim, mlp_ratio=mlp_ratio, drop=drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 324 | num_blocks=self.num_blocks, sparsity_threshold=sparsity_threshold, hard_thresholding_fraction=hard_thresholding_fraction) 325 | for i in range(encoder_depth)]) 326 | 327 | # Koopman Layers 328 | self.core_blocks = nn.ModuleList([ 329 | Af_Block(dim=embed_dim, mlp_ratio=mlp_ratio, drop=drop_rate, drop_path=dpr[i], norm_layer=norm_layer, 330 | num_blocks=self.num_blocks, sparsity_threshold=sparsity_threshold, hard_thresholding_fraction=hard_thresholding_fraction) 331 | for i in range(self.depth)]) 332 | 333 | self.norm = norm_layer(embed_dim) 334 | 335 | # High-frequency component 336 | self.w0 = nn.Conv1d(embed_dim, embed_dim, 1) # or user-defined more complicated convolutional structure 337 | 338 | # Decoder Settings 339 | if self.settings == "MLP": 340 | self.decoder_pred_mlp = nn.Linear(self.embed_dim, self.out_chans*self.patch_size[0]*self.patch_size[1], bias=False) 341 | elif self.settings == "Conv2d": 342 | self.decoder_pred_conv2d = nn.ConvTranspose2d(self.embed_dim, self.out_chans, kernel_size=self.patch_size, stride=self.patch_size) 343 | 344 | trunc_normal_(self.pos_embed, std=.02) 345 | self.apply(self._init_weights) 346 | 347 | def _init_weights(self, m): 348 | if isinstance(m, nn.Linear): 349 | trunc_normal_(m.weight, std=.02) 350 | if isinstance(m, nn.Linear) and m.bias is not None: 351 | nn.init.constant_(m.bias, 0) 352 | elif isinstance(m, nn.LayerNorm): 353 | nn.init.constant_(m.bias, 0) 354 | nn.init.constant_(m.weight, 1.0) 355 | 356 | @torch.jit.ignore 357 | def no_weight_decay(self): 358 | return {'pos_embed', 'cls_token'} 359 | 360 | def encoder(self, x): 361 | # Position Encoder 362 | B = x.shape[0] 363 | x = self.patch_embed(x) 364 | x = x + self.pos_embed 365 | x = self.pos_drop(x) 366 | # Encoder Network (if reconstruction task is hard, please use more complicated structure) 367 | for blk in self.encoder_blocks: 368 | x = blk(x) 369 | 370 | if self.encoder_network: 371 | x = self.encoder_network(x) 372 | 373 | return x 374 | 375 | def decoder(self, x): 376 | B = x.shape[0] 377 | x = x.reshape(B, self.h, self.w, self.embed_dim) 378 | if self.settings == "MLP": 379 | x = self.decoder_pred_mlp(x) 380 | x = rearrange( 381 | x, 382 | "b h w (p1 p2 c_out) -> b c_out (h p1) (w p2)", 383 | p1=self.patch_size[0], 384 | p2=self.patch_size[1], 385 | h=self.img_size[0] // self.patch_size[0], 386 | w=self.img_size[1] // self.patch_size[1], 387 | ) 388 | elif self.settings == "Conv2d": 389 | x = rearrange(x, "B H W C -> B C H W") 390 | x = self.decoder_pred_conv2d(x) 391 | return x 392 | 393 | def forward(self, x): 394 | x = self.encoder(x) 395 | # Reconstruction 396 | x_recons = self.decoder(x) 397 | # Prediction 398 | x_w = self.w0(x.permute(0,2,1)).permute(0,2,1) 399 | for blk in self.core_blocks: 400 | x = blk(x) 401 | x = blk(x) 402 | x = x + x_w 403 | x = self.decoder(x) 404 | return x, x_recons 405 | -------------------------------------------------------------------------------- /koopmanlab/utils.py: -------------------------------------------------------------------------------- 1 | import math 2 | import torch 3 | from torch import Tensor 4 | from typing import List, Optional 5 | from torch.optim.optimizer import Optimizer 6 | import operator 7 | from functools import reduce 8 | 9 | # print the number of parameters 10 | def count_params(model): 11 | c = 0 12 | for p in list(model.parameters()): 13 | c += reduce(operator.mul, 14 | list(p.size()+(2,) if p.is_complex() else p.size())) 15 | return c 16 | 17 | 18 | def adam(params: List[Tensor], 19 | grads: List[Tensor], 20 | exp_avgs: List[Tensor], 21 | exp_avg_sqs: List[Tensor], 22 | max_exp_avg_sqs: List[Tensor], 23 | state_steps: List[int], 24 | *, 25 | amsgrad: bool, 26 | beta1: float, 27 | beta2: float, 28 | lr: float, 29 | weight_decay: float, 30 | eps: float): 31 | r"""Functional API that performs Adam algorithm computation. 32 | See :class:`~torch.optim.Adam` for details. 33 | """ 34 | 35 | for i, param in enumerate(params): 36 | 37 | grad = grads[i] 38 | exp_avg = exp_avgs[i] 39 | exp_avg_sq = exp_avg_sqs[i] 40 | step = state_steps[i] 41 | 42 | bias_correction1 = 1 - beta1 ** step 43 | bias_correction2 = 1 - beta2 ** step 44 | 45 | if weight_decay != 0: 46 | grad = grad.add(param, alpha=weight_decay) 47 | 48 | # Decay the first and second moment running average coefficient 49 | exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1) 50 | exp_avg_sq.mul_(beta2).addcmul_(grad, grad.conj(), value=1 - beta2) 51 | if amsgrad: 52 | # Maintains the maximum of all 2nd moment running avg. till now 53 | torch.maximum(max_exp_avg_sqs[i], exp_avg_sq, out=max_exp_avg_sqs[i]) 54 | # Use the max. for normalizing running avg. of gradient 55 | denom = (max_exp_avg_sqs[i].sqrt() / math.sqrt(bias_correction2)).add_(eps) 56 | else: 57 | denom = (exp_avg_sq.sqrt() / math.sqrt(bias_correction2)).add_(eps) 58 | 59 | step_size = lr / bias_correction1 60 | 61 | param.addcdiv_(exp_avg, denom, value=-step_size) 62 | 63 | 64 | class Adam(Optimizer): 65 | r"""Implements Adam algorithm. 66 | It has been proposed in `Adam: A Method for Stochastic Optimization`_. 67 | The implementation of the L2 penalty follows changes proposed in 68 | `Decoupled Weight Decay Regularization`_. 69 | Args: 70 | params (iterable): iterable of parameters to optimize or dicts defining 71 | parameter groups 72 | lr (float, optional): learning rate (default: 1e-3) 73 | betas (Tuple[float, float], optional): coefficients used for computing 74 | running averages of gradient and its square (default: (0.9, 0.999)) 75 | eps (float, optional): term added to the denominator to improve 76 | numerical stability (default: 1e-8) 77 | weight_decay (float, optional): weight decay (L2 penalty) (default: 0) 78 | amsgrad (boolean, optional): whether to use the AMSGrad variant of this 79 | algorithm from the paper `On the Convergence of Adam and Beyond`_ 80 | (default: False) 81 | .. _Adam\: A Method for Stochastic Optimization: 82 | https://arxiv.org/abs/1412.6980 83 | .. _Decoupled Weight Decay Regularization: 84 | https://arxiv.org/abs/1711.05101 85 | .. _On the Convergence of Adam and Beyond: 86 | https://openreview.net/forum?id=ryQu7f-RZ 87 | """ 88 | 89 | def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8, 90 | weight_decay=0, amsgrad=False): 91 | if not 0.0 <= lr: 92 | raise ValueError("Invalid learning rate: {}".format(lr)) 93 | if not 0.0 <= eps: 94 | raise ValueError("Invalid epsilon value: {}".format(eps)) 95 | if not 0.0 <= betas[0] < 1.0: 96 | raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0])) 97 | if not 0.0 <= betas[1] < 1.0: 98 | raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1])) 99 | if not 0.0 <= weight_decay: 100 | raise ValueError("Invalid weight_decay value: {}".format(weight_decay)) 101 | defaults = dict(lr=lr, betas=betas, eps=eps, 102 | weight_decay=weight_decay, amsgrad=amsgrad) 103 | super(Adam, self).__init__(params, defaults) 104 | 105 | def __setstate__(self, state): 106 | super(Adam, self).__setstate__(state) 107 | for group in self.param_groups: 108 | group.setdefault('amsgrad', False) 109 | 110 | @torch.no_grad() 111 | def step(self, closure=None): 112 | """Performs a single optimization step. 113 | Args: 114 | closure (callable, optional): A closure that reevaluates the model 115 | and returns the loss. 116 | """ 117 | loss = None 118 | if closure is not None: 119 | with torch.enable_grad(): 120 | loss = closure() 121 | 122 | for group in self.param_groups: 123 | params_with_grad = [] 124 | grads = [] 125 | exp_avgs = [] 126 | exp_avg_sqs = [] 127 | max_exp_avg_sqs = [] 128 | state_steps = [] 129 | beta1, beta2 = group['betas'] 130 | 131 | for p in group['params']: 132 | if p.grad is not None: 133 | params_with_grad.append(p) 134 | if p.grad.is_sparse: 135 | raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead') 136 | grads.append(p.grad) 137 | 138 | state = self.state[p] 139 | # Lazy state initialization 140 | if len(state) == 0: 141 | state['step'] = 0 142 | # Exponential moving average of gradient values 143 | state['exp_avg'] = torch.zeros_like(p, memory_format=torch.preserve_format) 144 | # Exponential moving average of squared gradient values 145 | state['exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) 146 | if group['amsgrad']: 147 | # Maintains max of all exp. moving avg. of sq. grad. values 148 | state['max_exp_avg_sq'] = torch.zeros_like(p, memory_format=torch.preserve_format) 149 | 150 | exp_avgs.append(state['exp_avg']) 151 | exp_avg_sqs.append(state['exp_avg_sq']) 152 | 153 | if group['amsgrad']: 154 | max_exp_avg_sqs.append(state['max_exp_avg_sq']) 155 | 156 | # update the steps for each param group update 157 | state['step'] += 1 158 | # record the step after step update 159 | state_steps.append(state['step']) 160 | 161 | adam(params_with_grad, 162 | grads, 163 | exp_avgs, 164 | exp_avg_sqs, 165 | max_exp_avg_sqs, 166 | state_steps, 167 | amsgrad=group['amsgrad'], 168 | beta1=beta1, 169 | beta2=beta2, 170 | lr=group['lr'], 171 | weight_decay=group['weight_decay'], 172 | eps=group['eps']) -------------------------------------------------------------------------------- /logo.png: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/Koopman-Laboratory/KoopmanLab/c9e347a9df50103d308235148132ce7ad1c850a4/logo.png -------------------------------------------------------------------------------- /setup.py: -------------------------------------------------------------------------------- 1 | import io 2 | import setuptools 3 | 4 | with io.open("README.md", "r", encoding="utf-8") as f: 5 | long_description = f.read() 6 | 7 | setuptools.setup( 8 | name="koopmanlab", 9 | version="1.0.4", 10 | author="Wei Xiong, Tian Yang", 11 | author_email="xiongw21@mails.tsinghua.edu.cn", 12 | description="A library for Koopman Neural Operator with Pytorch", 13 | long_description=long_description, 14 | long_description_content_type="text/markdown", 15 | url="https://github.com/Koopman-Laboratory/KoopmanLab", 16 | packages=setuptools.find_packages(), 17 | classifiers=[ 18 | "Programming Language :: Python :: 3", 19 | "License :: OSI Approved :: GNU General Public License v3 (GPLv3)", 20 | "Operating System :: OS Independent", 21 | ], 22 | python_requires='>=3.8.5', 23 | install_requires = [ 24 | 'torch>=1.10', 25 | 'torchvision>=0.13.1', 26 | 'matplotlib>=3.3.2', 27 | 'numpy>=1.14.5', 28 | 'einops==0.5.0', 29 | 'timm==0.6.11', 30 | 'scipy==1.7.3', 31 | 'h5py==3.7.0', 32 | ] 33 | ) 34 | --------------------------------------------------------------------------------