update PhyGeoNet (PaddlePaddle#731)

* Change the directory structure

* update phygeonet code

* update heat equation code

* update heat equation with bc code

* refine phygeonet code

* add callback for seed and logger

examples/phygeonet/heat_equation/conf/conf.yaml → examples/phygeonet/conf/heat_equation.yaml

@@ -14,6 +14,9 @@ hydra:
           - mode
           - output_dir
           - log_freq
+  callbacks:
+    init_callback:
+      _target_: ppsci.utils.callbacks.InitCallback
   sweep:
     # output directory for multirun
     dir: ${hydra.run.dir}
@@ -26,11 +29,11 @@ data_dir: ./data/heat_equation.npz
 output_dir: ${hydra:run.dir}
 dtheta: 0
 epochs: 1500
-len_data: 1
+
 # model settings
 MODEL:
   input_keys: [ 'coords' ]
-  output_keys: [ 'outputV' ]
+  output_keys: [ 'output_v' ]
   hidden_size: [16, 32, 16]
   h: 0.01
   ny: 19
@@ -41,8 +44,8 @@ MODEL:
 
 # training settings
 TRAIN:
-  lr: 1e-3
-  batchsize: 1
+  learning_rate: 1e-3
+  batch_size: 1
 
 # evaluation settings
 EVAL:

examples/phygeonet/heat_equation_parameterized_BC/conf/conf.yaml → examples/phygeonet/conf/heat_equation_with_bc.yaml

@@ -14,8 +14,10 @@ hydra:
           - mode
           - output_dir
           - log_freq
+  callbacks:
+    init_callback:
+      _target_: ppsci.utils.callbacks.InitCallback
   sweep:
-
     # output directory for multirun
     dir: ${hydra.run.dir}
     subdir: ./
@@ -29,12 +31,11 @@ output_dir: ${hydra:run.dir}
 r: 0.5
 R: 1
 dtheta: 0
-len_data: 2
 epochs: 2000
 # model settings
 MODEL:
   input_keys: [ 'coords' ]
-  output_keys: [ 'outputV' ]
+  output_keys: [ 'output_v' ]
   hidden_size: [16, 32, 16]
   h: 0.01
   ny: 276
@@ -45,8 +46,8 @@ MODEL:
 
 # training settings
 TRAIN:
-  lr: 1e-3
-  batchsize: 2
+  learning_rate: 1e-3
+  batch_size: 1
 
 # evaluation settings
 EVAL:

examples/phygeonet/heat_equation.py

@@ -0,0 +1,165 @@
+from typing import Dict
+
+import hydra
+import matplotlib.pyplot as plt
+import numpy as np
+import paddle
+import utils
+from omegaconf import DictConfig
+
+import ppsci
+from ppsci.utils import logger
+
+
+def train(cfg: DictConfig):
+    data = np.load(cfg.data_dir)
+    coords = data["coords"]
+    jinvs = data["jinvs"]
+    dxdxis = data["dxdxis"]
+    dydxis = data["dydxis"]
+    dxdetas = data["dxdetas"]
+    dydetas = data["dydetas"]
+
+    model = ppsci.arch.USCNN(**cfg.MODEL)
+
+    optimizer = ppsci.optimizer.Adam(cfg.TRAIN.learning_rate)(model)
+
+    iters_per_epoch = coords.shape[0]
+    sup_constraint_res = ppsci.constraint.SupervisedConstraint(
+        {
+            "dataset": {
+                "name": "NamedArrayDataset",
+                "input": {
+                    "coords": coords,
+                    "jinvs": jinvs,
+                    "dxdxis": dxdxis,
+                    "dydxis": dydxis,
+                    "dxdetas": dxdetas,
+                    "dydetas": dydetas,
+                },
+            },
+            "batch_size": cfg.TRAIN.batch_size,
+            "iters_per_epoch": iters_per_epoch,
+            "num_workers": 0,
+        },
+        ppsci.loss.FunctionalLoss(lambda out, label, weight: out["residual"]),
+        name="residual",
+    )
+    sup_constraint = {sup_constraint_res.name: sup_constraint_res}
+
+    def _transform_out(
+        _input: Dict[str, paddle.Tensor],
+        _output: Dict[str, paddle.Tensor],
+        pad_singleside: int = cfg.MODEL.pad_singleside,
+    ):
+        """Calculation residual.
+
+        Args:
+            _input (Dict[str, paddle.Tensor]): The input of the model.
+            _output (Dict[str, paddle.Tensor]): The output of the model.
+            pad_singleside (int, optional): Pad size. Defaults to cfg.MODEL.pad_singleside.
+        """
+        output_v = _output["output_v"]
+        jinv = _input["jinvs"]
+        dxdxi = _input["dxdxis"]
+        dydxi = _input["dydxis"]
+        dxdeta = _input["dxdetas"]
+        dydeta = _input["dydetas"]
+        output_v[:, 0, -pad_singleside:, pad_singleside:-pad_singleside] = 0
+        output_v[:, 0, :pad_singleside, pad_singleside:-pad_singleside] = 1
+        output_v[:, 0, pad_singleside:-pad_singleside, -pad_singleside:] = 1
+        output_v[:, 0, pad_singleside:-pad_singleside, 0:pad_singleside] = 1
+        output_v[:, 0, 0, 0] = 0.5 * (output_v[:, 0, 0, 1] + output_v[:, 0, 1, 0])
+        output_v[:, 0, 0, -1] = 0.5 * (output_v[:, 0, 0, -2] + output_v[:, 0, 1, -1])
+        dvdx = utils.dfdx(output_v, dydeta, dydxi, jinv)
+        d2vdx2 = utils.dfdx(dvdx, dydeta, dydxi, jinv)
+        dvdy = utils.dfdy(output_v, dxdxi, dxdeta, jinv)
+        d2vdy2 = utils.dfdy(dvdy, dxdxi, dxdeta, jinv)
+        continuity = d2vdy2 + d2vdx2
+        return {"residual": paddle.mean(continuity**2)}
+
+    model.register_output_transform(_transform_out)
+    solver = ppsci.solver.Solver(
+        model,
+        sup_constraint,
+        cfg.output_dir,
+        optimizer,
+        epochs=cfg.epochs,
+        iters_per_epoch=iters_per_epoch,
+    )
+    solver.train()
+    solver.plot_loss_history()
+
+
+def evaluate(cfg: DictConfig):
+    data = np.load(cfg.data_dir)
+    coords = data["coords"]
+
+    ofv_sb = paddle.to_tensor(data["OFV_sb"])
+
+    ## create model
+    pad_singleside = cfg.MODEL.pad_singleside
+    model = ppsci.arch.USCNN(**cfg.MODEL)
+    solver = ppsci.solver.Solver(
+        model,
+        pretrained_model_path=cfg.EVAL.pretrained_model_path,  ### the path of the model
+    )
+    output_v = solver.predict({"coords": paddle.to_tensor(coords)})
+    output_v = output_v["output_v"]
+
+    output_v[0, 0, -pad_singleside:, pad_singleside:-pad_singleside] = 0
+    output_v[0, 0, :pad_singleside, pad_singleside:-pad_singleside] = 1
+    output_v[0, 0, pad_singleside:-pad_singleside, -pad_singleside:] = 1
+    output_v[0, 0, pad_singleside:-pad_singleside, 0:pad_singleside] = 1
+    output_v[0, 0, 0, 0] = 0.5 * (output_v[0, 0, 0, 1] + output_v[0, 0, 1, 0])
+    output_v[0, 0, 0, -1] = 0.5 * (output_v[0, 0, 0, -2] + output_v[0, 0, 1, -1])
+
+    ev = paddle.sqrt(
+        paddle.mean((ofv_sb - output_v[0, 0]) ** 2) / paddle.mean(ofv_sb**2)
+    ).item()
+    logger.info(f"ev: {ev}")
+
+    output_v = output_v.numpy()
+    ofv_sb = ofv_sb.numpy()
+    fig = plt.figure()
+    ax = plt.subplot(1, 2, 1)
+    utils.visualize(
+        ax,
+        coords[0, 0, 1:-1, 1:-1],
+        coords[0, 1, 1:-1, 1:-1],
+        output_v[0, 0, 1:-1, 1:-1],
+        "horizontal",
+        [0, 1],
+    )
+    utils.set_axis_label(ax, "p")
+    ax.set_title("CNN " + r"$T$")
+    ax.set_aspect("equal")
+    ax = plt.subplot(1, 2, 2)
+    utils.visualize(
+        ax,
+        coords[0, 0, 1:-1, 1:-1],
+        coords[0, 1, 1:-1, 1:-1],
+        ofv_sb[1:-1, 1:-1],
+        "horizontal",
+        [0, 1],
+    )
+    utils.set_axis_label(ax, "p")
+    ax.set_aspect("equal")
+    ax.set_title("FV " + r"$T$")
+    fig.tight_layout(pad=1)
+    fig.savefig(f"{cfg.output_dir}/result.png", bbox_inches="tight")
+    plt.close(fig)
+
+
+@hydra.main(version_base=None, config_path="./conf", config_name="heat_equation.yaml")
+def main(cfg: DictConfig):
+    if cfg.mode == "train":
+        train(cfg)
+    elif cfg.mode == "eval":
+        evaluate(cfg)
+    else:
+        raise ValueError(f"cfg.mode should in ['train', 'eval'], but got '{cfg.mode}'")
+
+
+if __name__ == "__main__":
+    main()