Package: koopmanlab
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." arXiv preprint arXiv:2301.10022 (2023)."KoopmanLab: machine learning for solving complex physics equations." arXiv preprint arXiv:2301.01104 (2023).InstallationKoopmanLab requires the following dependencies to be installed:PyTorch >= 1.10Numpy >= 1.23.2Matplotlib >= 3.3.2You can install KoopmanLab package via the following approaches:Install the stable version withpip:$ pip install koopmanlabInstall the current version by source code withpip:$ git clone https://github.com/Koopman-Laboratory/KoopmanLab.git
$ cd KoopmanLab
$ pip install -e .Quick StartIf you install KoopmanLab successfully, you can use our model directly by:importkoopmanlabaskpencoder=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 iterationKNO2d_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 iterationIf you do not want to customize the algorithms for training, testing and plotting, we highly recommend that you use our basic APIs to use Koopman model.UsageYou can readdemo_ns.pyto learn about some basic APIs and workflow of KoopmanLab. If you want to rundemo_ns.py, the following data need to be prepared in your computing resource.DatasetIf you want to generate Navier-Stokes Equation data by yourself, the data generation configuration file can be found in the following link.FileOur package provides an easy way to create a Koopman neural operator model.importkoopmanlabaskpMLP_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 calculatingViT_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 appliedOnce 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.optimizerandMLP_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 methodtorch.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 inkp.model.koopmanand 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 usetrain_singlemethod or usetrainwithT_out = 1. Our package provides a method to save and visualize your prediction results intest.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,trainandtestmethod is set with a single step predicition scenario. Specifically,train_multiandtest_multimethod provide multi-step iteration prediction, where the model iteratesT_outtimes 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 predictOnce your model has been trained, you can use the saving module provided in KoopmanLab to save your model. Saved variable has three attribute. wherekoopmanis the model class variable (i.e., the savedkno_modelvariable),modelis the trained model variable (i.e., the savedkno_model.kernelvariable), andmodel_paramsis the parameters dictionary of trained model variable (i.e., the savedkno_model.kernel.state_dict()variable).MLP_KNO_2D.save(save_path)## Parameter definitions:# save_path: the file path of the result savingCitationIf you use KoopmanLab package for academic research, you are encouraged to cite the following paper:@article{xiong2023kno,
  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={arXiv preprint arXiv:2301.10022},
  year={2023}
}

@article{xiong2023koopmanlab,
  title={KoopmanLab: A PyTorch module of Koopman neural operator family for solving partial differential equations},
  author={Xiong, Wei and Ma, Muyuan and Sun, Pei and Tian, Yang},
  journal={arXiv preprint arXiv:2301.01104},
  year={2023}
}AcknowledgementAuthors appreciate Abby, a talented artist, for designing the logo of KoopmanLab.LicenseGPL-3.0 License