{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 用于时空PDE系统的物理编码消息传递图神经网络\n",
    "\n",
    "[![下载Notebook](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_notebook.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.5.0/mindflow/zh_cn/data_mechanism_fusion/mindspore_phympgn.ipynb) [![下载样例代码](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_download_code.svg)](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/notebook/r2.5.0/mindflow/zh_cn/data_mechanism_fusion/mindspore_phympgn.py) [![查看源文件](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/resource/_static/logo_source.svg)](https://gitee.com/mindspore/docs/blob/r2.5.0/docs/mindflow/docs/source_zh_cn/data_mechanism_fusion/phympgn.ipynb)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "偏微分方程（PDEs）控制的复杂动力系统广泛存在于各个学科当中。近年来，数据驱动的神经网络模型在预测时空动态上取得了极好的效果。\n",
    "\n",
    "物理编码的消息传递图网络（PhyMPGN: Physics-encoded Message Passing Graph Network for spatiotemporal PDE systems），可以使用少量训练数据在不规则计算域上建模时空PDE系统。具体来说，\n",
    "\n",
    "- 提出了一个使用消息传递机制的物理编码图学习模型，使用二阶龙格库塔（Runge-Kutta）数值方案进行时间步进\n",
    "- 考虑到物理现象中普遍存在扩散过程，设计了一个可学习的Laplace Block，编码了离散拉普拉斯-贝尔特拉米算子（Laplace-Beltrami Operator）\n",
    "- 提出了一个新颖的填充策略在模型中编码不同类型的边界条件\n",
    "\n",
    "论文链接: [https://arxiv.org/abs/2410.01337](https://gitee.com/link?target=https%3A%2F%2Farxiv.org%2Fabs%2F2410.01337)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 问题描述"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "考虑由如下形式控制的时空PDE系统：\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\dot {\\boldsymbol {u}}(\\boldsymbol x, t) = \\boldsymbol F (t, \\boldsymbol x, \\boldsymbol u, \\nabla \\boldsymbol u, \\Delta \\boldsymbol u, \\dots)\n",
    "\\tag{1}\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "其中$\\boldsymbol u(\\boldsymbol x, y) \\in \\mathbb{R}^m$是具有$m$个分量的状态变量向量，例如速度、温度或者压力等，它的定义在时空域$\\{ \\boldsymbol x, t \\} \\in \\Omega \\times [0, \\mathcal{T}]$上；$\\dot{\\boldsymbol u}$代表$\\boldsymbol u$对时间的导数，$\\boldsymbol F$是依赖于当前状态$\\boldsymbol u$和其空间导数的非线性算子。\n",
    "\n",
    "假设在空间域$\\Omega$上有着非均匀且稀疏的观测结点$\\{ \\boldsymbol x_0, \\dots, \\boldsymbol x_{N-1} \\}$（即，非结构化网格），在时刻$t_0, \\dots, t_{T-1}$，这些结点上的观测为$\\{ \\boldsymbol U(t_0), \\dots, \\boldsymbol U(t_{T-1}) \\}$，其中的$\\boldsymbol U(t_i) = \\{ \\boldsymbol u(\\boldsymbol x_0, t_i), \\dots, \\boldsymbol u (\\boldsymbol x_{N-1}, t_i) \\}$代表某些物理量。考虑到很多物理现象包含扩散过程，我们假设PDE中的扩散项是已知的先验知识。我们的目标是使用少量训练数据学习一个图神经网络模型，在稀疏非结构网格上预测不同的时空动态系统，处理不同的边界条件，为任意的初始条件产生后续动态轨迹。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "本案例展示PhyMPGN如何求解圆柱绕流（Cylinder Flow）问题。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "圆柱绕流 (Cylinder Flow) 动态系统由如下的Navier-Stokes方程控制\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\dot{\\boldsymbol{u}} = - \\boldsymbol{u} \\cdot \\nabla \\boldsymbol{u} -\\frac{1}{\\rho} \\nabla p + \\frac{\\mu}{\\rho} \\Delta \\boldsymbol{u} + \\boldsymbol{f}\n",
    "\\tag{2}\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "其中流体密度$\\rho=1$，流体粘度系数$\\mu=5.0\\times10^{-3}$，外力$f=0$。该圆柱绕流系统左边界为入口，右边界为出口，圆柱表面为无滑移边界条件，上下边界为对称边界条件。本案例关注于，在保持流体密度$\\rho$，圆柱大小$D=2$和流体粘度系数$\\mu$不变的情况下，泛化入射流速度$U_m$。因为雷诺数定义为$Re=\\rho U_m D / \\mu$，所以泛化入射流速度$U_m$也意味着泛化不同的雷诺数。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型方法"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "![PhyMPGN网络结构](https://mindspore-website.obs.cn-north-4.myhuaweicloud.com/website-images/r2.5.0/docs/mindflow/docs/source_zh_cn/data_mechanism_fusion/images/phympgn.png)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "对于式（1），可以使用二阶龙格库塔（Runge-Kutta, RK2）方案进行离散：\n",
    "\n",
    "$$\n",
    "\\begin{equation}\n",
    "\\boldsymbol u^{k+1} = \\boldsymbol u^k + \\frac{1}{2}(\\boldsymbol g_1 + \\boldsymbol g_2); \\quad \\boldsymbol g_1 = \\boldsymbol F(t^k, \\boldsymbol x, \\boldsymbol u^k, \\dots); \\quad \\boldsymbol g_2 = \\boldsymbol F(t^{k+1}, \\boldsymbol x, \\boldsymbol u^k + \\delta t \\boldsymbol g_1, \\dots)\n",
    "\\tag{3}\n",
    "\\end{equation}\n",
    "$$\n",
    "\n",
    "其中$\\boldsymbol u^k$为$t^k$时刻的状态变量，$\\delta t$为时刻$t^k$和$t^{k+1}$之间的时间间隔。根据式（3），我们构建一个GNN来学习非线性算子$\\boldsymbol F$.\n",
    "\n",
    "如图所示，我们使用NN block来学习非线性算子$\\boldsymbol F$。NN block又可以分为两部分：采用编码器-处理器-解码器架构的GNN block和可学习的Laplace block。因为物理现象中扩散过程的普遍存在性，我们设计了可学习的Laplace block，编码离散拉普拉斯贝尔特拉米算子（Laplace-Beltrami operator），来学习由PDE中扩散项导致的增量；而GNN block来学习PDE中其他项导致的增量。"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 准备环节"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "- 确保已安装相关依赖库，如MindSpore等\n",
    "- 确保已下载好[圆柱绕流数据](https://download-mindspore.osinfra.cn/mindscience/mindflow/dataset/applications/data_mechanism_fusion/PhyMPGN/)\n",
    "- 确保在[yamls/train.yaml](https://gitee.com/mindspore/mindscience/blob/r0.7/MindFlow/applications/data_mechanism_fusion/phympgn/yamls/train.yaml)配置文件中已配置好数据和模型权重等相关保存路径"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 代码执行步骤"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "代码执行流程如下步骤：\n",
    "\n",
    "1. 读取配置文件\n",
    "2. 构建数据集\n",
    "3. 构建模型\n",
    "4. 模型训练\n",
    "5. 模型推理\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 读取配置文件"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindflow.utils import log_config, load_yaml_config, print_log\n",
    "from easydict import EasyDict\n",
    "import os.path as osp\n",
    "from pathlib import Path\n",
    "\n",
    "\n",
    "def load_config(config_file_path, train):\n",
    "    config = load_yaml_config(config_file_path)\n",
    "    config['train'] = train\n",
    "    config = EasyDict(config)\n",
    "    log_dir = './logs'\n",
    "    if train:\n",
    "        log_file = f'phympgn-{config.experiment_name}'\n",
    "    else:\n",
    "        log_file = f'phympgn-{config.experiment_name}-te'\n",
    "    if not osp.exists(osp.join(log_dir, f'{log_file}.log')):\n",
    "        Path(osp.join(log_dir, f'{log_file}.log')).touch()\n",
    "    log_config(log_dir, log_file)\n",
    "    print_log(config)\n",
    "    return config"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file_path = 'yamls/train.yaml'\n",
    "config = load_config(config_file_path=config_file_path, train=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore as ms\n",
    "\n",
    "ms.set_device(device_target='Ascend', device_id=7)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建数据集"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import PDECFDataset, get_data_loader\n",
    "\n",
    "\n",
    "print_log('Train...')\n",
    "print_log('Loading training data...')\n",
    "tr_dataset = PDECFDataset(\n",
    "    root=config.path.data_root_dir,\n",
    "    raw_files=config.path.tr_raw_data,\n",
    "    dataset_start=config.data.dataset_start,\n",
    "    dataset_used=config.data.dataset_used,\n",
    "    time_start=config.data.time_start,\n",
    "    time_used=config.data.time_used,\n",
    "    window_size=config.data.tr_window_size,\n",
    "    training=True\n",
    ")\n",
    "tr_loader = get_data_loader(\n",
    "    dataset=tr_dataset,\n",
    "    batch_size=config.optim.batch_size\n",
    ")\n",
    "\n",
    "print_log('Loading validation data...')\n",
    "val_dataset = PDECFDataset(\n",
    "    root=config.path.data_root_dir,\n",
    "    raw_files=config.path.val_raw_data,\n",
    "    dataset_start=config.data.dataset_start,\n",
    "    dataset_used=config.data.dataset_used,\n",
    "    time_start=config.data.time_start,\n",
    "    time_used=config.data.time_used,\n",
    "    window_size=config.data.val_window_size\n",
    ")\n",
    "val_loader = get_data_loader(\n",
    "    dataset=val_dataset,\n",
    "    batch_size=config.optim.batch_size\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 构建模型"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import PhyMPGN\n",
    "\n",
    "print_log('Building model...')\n",
    "model = PhyMPGN(\n",
    "    encoder_config=config.network.encoder_config,\n",
    "    mpnn_block_config=config.network.mpnn_block_config,\n",
    "    decoder_config=config.network.decoder_config,\n",
    "    laplace_block_config=config.network.laplace_block_config,\n",
    "    integral=config.network.integral\n",
    ")\n",
    "print_log(f'Number of parameters: {model.num_params}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型训练"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mindflow import get_multi_step_lr\n",
    "from mindspore import nn\n",
    "import numpy as np\n",
    "\n",
    "from src import Trainer, TwoStepLoss\n",
    "\n",
    "\n",
    "lr_scheduler = get_multi_step_lr(\n",
    "    lr_init=config.optim.lr,\n",
    "    milestones=list(np.arange(0, config.optim.start_epoch+config.optim.epochs,\n",
    "                              step=config.optim.steplr_size)[1:]),\n",
    "    gamma=config.optim.steplr_gamma,\n",
    "    steps_per_epoch=len(tr_loader),\n",
    "    last_epoch=config.optim.start_epoch+config.optim.epochs-1\n",
    ")\n",
    "optimizer = nn.AdamWeightDecay(model.trainable_params(), learning_rate=lr_scheduler,\n",
    "                               eps=1.0e-8, weight_decay=1.0e-2)\n",
    "trainer = Trainer(\n",
    "    model=model, optimizer=optimizer, scheduler=lr_scheduler, config=config,\n",
    "    loss_func=TwoStepLoss()\n",
    ")\n",
    "trainer.train(tr_loader, val_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "[Epoch    1/1600] Batch Time: 2.907 (3.011)  Data Time: 0.021 (0.035)  Graph Time: 0.004 (0.004)  Grad Time: 2.863 (2.873)  Optim Time: 0.006 (0.022)\n",
    "\n",
    "[Epoch    1/1600] Batch Time: 1.766 (1.564)  Data Time: 0.022 (0.044)  Graph Time: 0.003 (0.004)\n",
    "\n",
    "[Epoch    1/1600] tr_loss: 1.36e-02        val_loss: 1.29e-02 [MIN]\n",
    "\n",
    "[Epoch    2/1600] Batch Time: 3.578 (3.181)  Data Time: 0.024 (0.038)  Graph Time: 0.004 (0.004)  Grad Time: 3.531 (3.081)  Optim Time: 0.004 (0.013)\n",
    "\n",
    "[Epoch    2/1600] Batch Time: 1.727 (1.664)  Data Time: 0.023 (0.042)  Graph Time: 0.003 (0.004)\n",
    "\n",
    "[Epoch    2/1600] tr_loss: 1.15e-02        val_loss: 9.55e-03 [MIN]\n",
    "\n",
    "...\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 模型推理"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "config_file_path = 'yamls/train.yaml'\n",
    "config = load_config(config_file_path=config_file_path, train=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import mindspore as ms\n",
    "\n",
    "ms.set_device(device_target='Ascend', device_id=7)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from src import PDECFDataset, get_data_loader, Trainer, PhyMPGN\n",
    "from mindspore import nn\n",
    "\n",
    "\n",
    "# test datasets\n",
    "te_dataset = PDECFDataset(\n",
    "    root=config.path.data_root_dir,\n",
    "    raw_files=config.path.te_raw_data,\n",
    "    dataset_start=config.data.te_dataset_start,\n",
    "    dataset_used=config.data.te_dataset_used,\n",
    "    time_start=config.data.time_start,\n",
    "    time_used=config.data.time_used,\n",
    "    window_size=config.data.te_window_size,\n",
    "    training=False\n",
    ")\n",
    "te_loader = get_data_loader(\n",
    "    dataset=te_dataset,\n",
    "    batch_size=1,\n",
    "    shuffle=False,\n",
    ")\n",
    "print_log('Building model...')\n",
    "model = PhyMPGN(\n",
    "    encoder_config=config.network.encoder_config,\n",
    "    mpnn_block_config=config.network.mpnn_block_config,\n",
    "    decoder_config=config.network.decoder_config,\n",
    "    laplace_block_config=config.network.laplace_block_config,\n",
    "    integral=config.network.integral\n",
    ")\n",
    "print_log(f'Number of parameters: {model.num_params}')\n",
    "trainer = Trainer(\n",
    "    model=model, optimizer=None, scheduler=None, config=config,\n",
    "    loss_func=nn.MSELoss()\n",
    ")\n",
    "print_log('Test...')\n",
    "trainer.test(te_loader)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    ".[TEST 0/9] MSE at 2000t: 5.06e-04, armse: 0.058, time: 185.3432s\n",
    "\n",
    "[TEST 1/9] MSE at 2000t: 4.83e-04, armse: 0.040, time: 186.3979s\n",
    "\n",
    "[TEST 2/9] MSE at 2000t: 1.95e-03, armse: 0.062, time: 177.0030s\n",
    "\n",
    "...\n",
    "\n",
    "[TEST 8/9] MSE at 2000t: 1.42e-01, armse: 0.188, time: 163.1219s\n",
    "\n",
    "[TEST 9] Mean Loss: 4.88e-02, Mean armse: 0.137, corre: 0.978, time: 173.3827s"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "zbc_ms2.5.0",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}