安装

KoopmanLab 需要安装以下依赖项：

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

你可以通过以下方式安装 KoopmanLab 包：

使用 pip 安装稳定版本：

$ pip install koopmanlab

通过源代码使用 pip 安装当前版本：

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

快速开始

如果你成功安装了 KoopmanLab，你可以直接使用我们的模型：

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)
# 输入大小 [batch, x, t_in] 输出大小 [batch, x, t_in]，一次迭代
KNO2d_model = kp.models.KNO2d(encoder, decoder, operator_size, modes_x = 10, modes_y = 10, decompose = 6)
# 输入大小 [batch, x, t_in] 输出大小 [batch, x, t_in]，一次迭代

如果你不想自定义训练、测试和绘图的算法，我们强烈建议你使用我们的基本 API 来构建 Koopman 模型。

使用方法

你可以阅读 demo_ns.py 来了解 KoopmanLab 的一些基本 API 和工作流程。如果你想运行 demo_ns.py，需要在你的计算资源中准备以下数据。

数据集

如果你想自己生成 Navier-Stokes 方程数据，可以在以下链接找到数据生成配置文件。

文件

我们的包提供了一种简单的方法来创建 Koopman 神经算子模型。

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()
## 参数定义：
# o：学习的 Koopman 算子的维度
# f：频率截断阈值以下的频率模式数
# r：Koopman 算子的幂
# T_in：输入数据的持续时间长度
# device：用于计算的设备是 CPU 还是 GPU

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()
## 参数定义：
# depth：每个头的深度
# head_num：头的数量
# resolution：输入数据的空间分辨率
# patch_size：每个补丁（即标记）的大小
# in_chans：数据集中目标变量的数量
# out_chans：ViT-KNO 预测的变量数量，通常与 in_chans 相同
# embed_dim：嵌入维度
# parallel：是否应用数据并行
# high_freq：是否应用高频信息补充

模型编译完成后，需要设置优化器以运行你自己的实验。如果你想对优化器和调度器进行更自定义的设置，可以使用任何 PyTorch 方法创建它们，并将它们分配给 Koopman 神经算子对象，例如 MLP_KNO_2D.optimizer 和 MLP_KNO_2D.scheduler。

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

如果你使用 Burgers 方程和 Navier-Stokes 方程数据或 PDEBench 提供的浅水方程数据，有三个特定的数据接口可供考虑。

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)
## 参数定义：
# path：下载的数据集的文件路径
# T_in：输入数据的持续时间长度
# T_out：需要预测的持续时间长度
# Type：Navier-Stokes 方程数据集的粘性系数
# sub：下采样缩放因子。例如，对空间分辨率为 64*64 的二维数据应用缩放因子 sub=2 将创建一个 32*32 的下采样空间。对空间分辨率为 1*64 的一维数据应用相同的因子意味着一个 1*32 的下采样空间。

我们建议你使用 PyTorch 方法 torch.utils.data.DataLoader 处理你的数据。在 KNO 模型中，2D 输入数据的形状是 [batchsize, x, y, t_len]，输出数据和标签的形状是 [batchsize, x, y, T]，其中 t_len 在 kp.model.koopman 中定义，T 在训练模块中定义。在 Koopman-ViT 模型中，2D 输入数据的形状是 [batchsize, in_chans, x, y]，输出数据和标签的形状是 [batchsize, out_chans, x, y]。 KoopmanLab提供了紧凑型KNO子系列的两种训练和两种测试方法。如果您的场景是单步预测，可以考虑使用train_single方法或使用train并设置T_out = 1。我们的软件包在test中提供了保存和可视化预测结果的方法。

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

对于ViT-KNO子系列，train和test方法设置为单步预测场景。具体来说，train_multi和test_multi方法提供多步迭代预测，模型在训练和测试中迭代T_out次。

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)
## 参数定义：
# epoch: 训练的轮数
# trainloader: 训练数据加载器，是torch.utils.data.DataLoader返回的变量
# evalloader: 评估数据加载器，是torch.utils.data.DataLoader返回的变量
# test_loader: 测试数据加载器，是torch.utils.data.DataLoader返回的变量
# T_out: 需要预测的持续时间长度

一旦您的模型训练完成，您可以使用KoopmanLab提供的保存模块来保存您的模型。保存的变量有三个属性：koopman是模型类变量（即保存的kno_model变量），model是训练好的模型变量（即保存的kno_model.kernel变量），model_params是训练好的模型变量的参数字典（即保存的kno_model.kernel.state_dict()变量）。

MLP_KNO_2D.save(save_path)
## 参数定义：
# save_path: 结果保存的文件路径

引用

如果您将KoopmanLab软件包用于学术研究，我们鼓励您引用以下论文：

@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}
}

致谢

作者感谢才华横溢的艺术家Abby设计了KoopmanLab的logo。

许可证

GPL-3.0 许可证