【PPSCI Export&Infer No.20】shock_wave (#890)

* add shock_wave export&infer

* Update examples/shock_wave/shock_wave.py

* Update examples/shock_wave/shock_wave.py

---------

Co-authored-by: HydrogenSulfate <[email protected]>

Author: smallpoxscattered

docs/zh/examples/shock_wave.md

@@ -28,6 +28,32 @@
         python shock_wave.py -cn=shock_wave_Ma0.728 mode=eval EVAL.pretrained_model_path=https://paddle-org.bj.bcebos.com/paddlescience/models/shockwave/shock_wave_Ma0728_pretrained.pdparams
         ```
 
+=== "模型导出命令"
+
+    === "Ma=2.0"
+
+        ``` sh
+        python shock_wave.py mode=export
+        ```
+    === "Ma=0.728"
+
+        ``` sh
+        python shock_wave.py -cn=shock_wave_Ma0.728 mode=export
+        ```
+
+=== "模型推理命令"
+
+    === "Ma=2.0"
+
+        ``` sh
+        python shock_wave.py mode=infer
+        ```
+    === "Ma=0.728"
+
+        ``` sh
+        python shock_wave.py -cn=shock_wave_Ma0.728 mode=infer
+        ```
+
 ## 1. 背景简介
 
 激波是自然界以及工程应用中经常发现的现象。它们不仅广泛地存在于航空航天领域的可压缩流动中，而且也表现在理论与应用物理以及工程应用等其它领域。在超声速与高超声速流动中，激波的出现对流体流动的整体特征会产生重要影响。激波捕捉问题已在CFD领域发展了数十年，以弱解的数学理论为基础的激波捕捉方法以其简单易实现的特点发展迅速，并在复杂超声速、高超声速流动数值模拟中得到了广泛应用。

examples/shock_wave/conf/shock_wave_Ma0.728.yaml

@@ -11,6 +11,8 @@ hydra:
           - TRAIN.checkpoint_path
           - TRAIN.pretrained_model_path
           - EVAL.pretrained_model_path
+          - INFER.pretrained_model_path
+          - INFER.export_path
           - mode
           - output_dir
           - log_freq
@@ -70,3 +72,20 @@ TRAIN:
 EVAL:
   pretrained_model_path: null
   eval_with_no_grad: true
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/shockwave/shock_wave_Ma0728_pretrained.pdparams
+  export_path: ./inference/shock_wave_Ma0.728
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  device: gpu
+  engine: native
+  precision: fp32
+  onnx_path: ${INFER.export_path}.onnx
+  ir_optim: true
+  min_subgraph_size: 10
+  gpu_mem: 4000
+  gpu_id: 0
+  max_batch_size: 256
+  num_cpu_threads: 4
+  batch_size: 256

examples/shock_wave/conf/shock_wave_Ma2.0.yaml

@@ -11,6 +11,8 @@ hydra:
           - TRAIN.checkpoint_path
           - TRAIN.pretrained_model_path
           - EVAL.pretrained_model_path
+          - INFER.pretrained_model_path
+          - INFER.export_path
           - mode
           - output_dir
           - log_freq
@@ -70,3 +72,20 @@ TRAIN:
 EVAL:
   pretrained_model_path: null
   eval_with_no_grad: true
+
+INFER:
+  pretrained_model_path: https://paddle-org.bj.bcebos.com/paddlescience/models/shockwave/shock_wave_Ma2_pretrained.pdparams
+  export_path: ./inference/shock_wave_Ma2.0
+  pdmodel_path: ${INFER.export_path}.pdmodel
+  pdiparams_path: ${INFER.export_path}.pdiparams
+  device: gpu
+  engine: native
+  precision: fp32
+  onnx_path: ${INFER.export_path}.onnx
+  ir_optim: true
+  min_subgraph_size: 10
+  gpu_mem: 4000
+  gpu_id: 0
+  max_batch_size: 256
+  num_cpu_threads: 4
+  batch_size: 256

examples/shock_wave/shock_wave.py

@@ -518,6 +518,112 @@ def evaluate(cfg: DictConfig):
     plt.savefig(osp.join(cfg.output_dir, f"shock_wave(Ma_{cfg.MA:.3f}).png"))
 
 
+def export(cfg: DictConfig):
+    from paddle.static import InputSpec
+
+    # set models
+    model = ppsci.arch.MLP(**cfg.MODEL)
+    solver = ppsci.solver.Solver(
+        model,
+        pretrained_model_path=cfg.INFER.pretrained_model_path,
+    )
+
+    # export models
+    input_spec = [
+        {key: InputSpec([None, 1], "float32", name=key) for key in model.input_keys},
+    ]
+    solver.export(input_spec, cfg.INFER.export_path)
+
+
+def inference(cfg: DictConfig):
+    from deploy.python_infer import pinn_predictor
+
+    # set model predictor
+    predictor = pinn_predictor.PINNPredictor(cfg)
+
+    # visualize prediction
+    t = np.linspace(cfg.T, cfg.T, 1, dtype=np.float32)
+    x = np.linspace(0.0, cfg.Lx, cfg.Nd, dtype=np.float32)
+    y = np.linspace(0.0, cfg.Ly, cfg.Nd, dtype=np.float32)
+    _, x_grid, y_grid = np.meshgrid(t, x, y)
+
+    x_test = misc.cartesian_product(t, x, y)
+    x_test_dict = misc.convert_to_dict(
+        x_test,
+        cfg.MODEL.input_keys,
+    )
+    output_dict = predictor.predict(
+        x_test_dict,
+        cfg.INFER.batch_size,
+    )
+
+    # mapping data to cfg.MODEL.output_keys
+    output_dict = {
+        store_key: output_dict[infer_key]
+        for store_key, infer_key in zip(cfg.MODEL.output_keys, output_dict.keys())
+    }
+
+    u, v, p, rho = (
+        output_dict["u"],
+        output_dict["v"],
+        output_dict["p"],
+        output_dict["rho"],
+    )
+
+    zero_mask = (
+        (x_test[:, 1] - cfg.rx) ** 2 + (x_test[:, 2] - cfg.ry) ** 2
+    ) < cfg.rd**2
+    u[zero_mask] = 0
+    v[zero_mask] = 0
+    p[zero_mask] = 0
+    rho[zero_mask] = 0
+
+    u = u.reshape(cfg.Nd, cfg.Nd)
+    v = v.reshape(cfg.Nd, cfg.Nd)
+    p = p.reshape(cfg.Nd, cfg.Nd)
+    rho = rho.reshape(cfg.Nd, cfg.Nd)
+
+    fig, ax = plt.subplots(2, 2, sharex=True, sharey=True, figsize=(15, 15))
+
+    plt.subplot(2, 2, 1)
+    plt.contourf(x_grid[:, 0, :], y_grid[:, 0, :], u * 241.315, 60)
+    plt.title("U m/s")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    axe = plt.gca()
+    axe.set_aspect(1)
+    plt.colorbar()
+
+    plt.subplot(2, 2, 2)
+    plt.contourf(x_grid[:, 0, :], y_grid[:, 0, :], v * 241.315, 60)
+    plt.title("V m/s")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    axe = plt.gca()
+    axe.set_aspect(1)
+    plt.colorbar()
+
+    plt.subplot(2, 2, 3)
+    plt.contourf(x_grid[:, 0, :], y_grid[:, 0, :], p * 33775, 60)
+    plt.title("P Pa")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    axe = plt.gca()
+    axe.set_aspect(1)
+    plt.colorbar()
+
+    plt.subplot(2, 2, 4)
+    plt.contourf(x_grid[:, 0, :], y_grid[:, 0, :], rho * 0.58, 60)
+    plt.title("Rho kg/m^3")
+    plt.xlabel("x")
+    plt.ylabel("y")
+    axe = plt.gca()
+    axe.set_aspect(1)
+    plt.colorbar()
+
+    plt.savefig(osp.join(cfg.output_dir, f"shock_wave(Ma_{cfg.MA:.3f}).png"))
+
+
 @hydra.main(
     version_base=None, config_path="./conf", config_name="shock_wave_Ma2.0.yaml"
 )
@@ -526,8 +632,14 @@ def main(cfg: DictConfig):
         train(cfg)
     elif cfg.mode == "eval":
         evaluate(cfg)
+    elif cfg.mode == "export":
+        export(cfg)
+    elif cfg.mode == "infer":
+        inference(cfg)
     else:
-        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")
+        raise ValueError(
+            f"cfg.mode should in ['train', 'eval', 'export', 'infer'], but got '{cfg.mode}'"
+        )
 
 
 if __name__ == "__main__":