Merge pull request #298 from HydrogenSulfate/add_aistudio_run

update button for 'run on AIStudio'

Author: HydrogenSulfate

docs/zh/examples/ldc2d_steady.md

@@ -1,5 +1,7 @@
 # 2D-LDC(2D Lid Driven Cavity Flow)
 
+<a href="https://aistudio.baidu.com/aistudio/projectdetail/6137973" class="md-button md-button--primary" style>AI Studio快速体验</a>
+
 ## 1. 问题简介
 
 Lid Driven Cavity Flow，中文名称可译作“顶盖驱动方腔流”，一般指顶部平板以恒定速度驱动规则区域内封闭的不可压流体（例如水）的流动。在方腔流中可以观察到几乎所有可能发生在不可压流体中的流动现象，是一种典型的流体案例。
@@ -301,17 +303,17 @@ examples/ldc/ldc2d_steady_Re10.py:124:141
 
 ``` py linenums="143"
 --8<--
-examples/ldc/ldc2d_steady_Re10.py:143:162
+examples/ldc/ldc2d_steady_Re10.py:143:153
 --8<--
 ```
 
 ### 3.9 模型训练、评估与可视化
 
 完成上述设置之后，只需要将上述实例化的对象按顺序传递给 `ppsci.solver.Solver`，然后启动训练、评估、可视化。
 
-``` py linenums="164"
+``` py linenums="155"
 --8<--
-examples/ldc/ldc2d_steady_Re10.py:164:
+examples/ldc/ldc2d_steady_Re10.py:155:
 --8<--
 ```
 

examples/ldc/ldc2d_steady_Re10.py

@@ -23,9 +23,9 @@
     # set random seed for reproducibility
     ppsci.utils.misc.set_random_seed(42)
     # set output directory
-    output_dir = "./ldc2d_steady_Re10" if not args.output_dir else args.output_dir
+    OUTPUT_DIR = "./ldc2d_steady_Re10" if not args.output_dir else args.output_dir
     # initialize logger
-    logger.init_logger("ppsci", f"{output_dir}/train.log", "info")
+    logger.init_logger("ppsci", f"{OUTPUT_DIR}/train.log", "info")
 
     # set model
     model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 9, 50, "tanh", False, False)
@@ -37,24 +37,24 @@
     geom = {"rect": ppsci.geometry.Rectangle((-0.05, -0.05), (0.05, 0.05))}
 
     # set dataloader config
-    iters_per_epoch = 1
+    ITERS_PER_EPOCH = 1
     train_dataloader_cfg = {
         "dataset": "IterableNamedArrayDataset",
-        "iters_per_epoch": iters_per_epoch,
+        "iters_per_epoch": ITERS_PER_EPOCH,
     }
 
-    npoint_pde = 99**2
-    npoint_top = 101
-    npoint_bottom = 101
-    npoint_left = 99
-    npoint_right = 99
+    NPOINT_PDE = 99**2
+    NPOINT_TOP = 101
+    NPOINT_BOTTOM = 101
+    NPOINT_LEFT = 99
+    NPOINT_RIGHT = 99
 
     # set constraint
     pde_constraint = ppsci.constraint.InteriorConstraint(
         equation["NavierStokes"].equations,
         {"continuity": 0, "momentum_x": 0, "momentum_y": 0},
         geom["rect"],
-        {**train_dataloader_cfg, "batch_size": npoint_pde},
+        {**train_dataloader_cfg, "batch_size": NPOINT_PDE},
         ppsci.loss.MSELoss("sum"),
         evenly=True,
         weight_dict={
@@ -68,7 +68,7 @@
         {"u": lambda out: out["u"], "v": lambda out: out["v"]},
         {"u": 1, "v": 0},
         geom["rect"],
-        {**train_dataloader_cfg, "batch_size": npoint_top},
+        {**train_dataloader_cfg, "batch_size": NPOINT_TOP},
         ppsci.loss.MSELoss("sum"),
         criteria=lambda x, y: np.isclose(y, 0.05),
         name="BC_top",
@@ -77,7 +77,7 @@
         {"u": lambda out: out["u"], "v": lambda out: out["v"]},
         {"u": 0, "v": 0},
         geom["rect"],
-        {**train_dataloader_cfg, "batch_size": npoint_bottom},
+        {**train_dataloader_cfg, "batch_size": NPOINT_BOTTOM},
         ppsci.loss.MSELoss("sum"),
         criteria=lambda x, y: np.isclose(y, -0.05),
         name="BC_bottom",
@@ -86,7 +86,7 @@
         {"u": lambda out: out["u"], "v": lambda out: out["v"]},
         {"u": 0, "v": 0},
         geom["rect"],
-        {**train_dataloader_cfg, "batch_size": npoint_left},
+        {**train_dataloader_cfg, "batch_size": NPOINT_LEFT},
         ppsci.loss.MSELoss("sum"),
         criteria=lambda x, y: np.isclose(x, -0.05),
         name="BC_left",
@@ -95,7 +95,7 @@
         {"u": lambda out: out["u"], "v": lambda out: out["v"]},
         {"u": 0, "v": 0},
         geom["rect"],
-        {**train_dataloader_cfg, "batch_size": npoint_right},
+        {**train_dataloader_cfg, "batch_size": NPOINT_RIGHT},
         ppsci.loss.MSELoss("sum"),
         criteria=lambda x, y: np.isclose(x, 0.05),
         name="BC_right",
@@ -110,26 +110,26 @@
     }
 
     # set training hyper-parameters
-    epochs = 20000 if not args.epochs else args.epochs
+    EPOCHS = 20000 if not args.epochs else args.epochs
     lr_scheduler = ppsci.optimizer.lr_scheduler.Cosine(
-        epochs,
-        iters_per_epoch,
+        EPOCHS,
+        ITERS_PER_EPOCH,
         0.001,
-        warmup_epoch=int(0.05 * epochs),
+        warmup_epoch=int(0.05 * EPOCHS),
     )()
 
     # set optimizer
     optimizer = ppsci.optimizer.Adam(lr_scheduler)((model,))
 
     # set validator
-    npoints_eval = npoint_pde
+    NPOINTS_EVAL = NPOINT_PDE
     residual_validator = ppsci.validate.GeometryValidator(
         equation["NavierStokes"].equations,
         {"momentum_x": 0, "continuity": 0, "momentum_y": 0},
         geom["rect"],
         {
             "dataset": "NamedArrayDataset",
-            "total_size": npoints_eval,
+            "total_size": NPOINTS_EVAL,
             "batch_size": 8192,
             "sampler": {"name": "BatchSampler"},
         },
@@ -141,18 +141,9 @@
     validator = {residual_validator.name: residual_validator}
 
     # set visualizer(optional)
-    npoint_bc = npoint_top + npoint_bottom + npoint_left + npoint_right
-    vis_interior_points = geom["rect"].sample_interior(npoint_pde, evenly=True)
-    vis_boundary_points = geom["rect"].sample_boundary(npoint_bc, evenly=True)
-
     # manually collate input data for visualization,
-    # interior+boundary
-    vis_points = {}
-    for key in vis_interior_points:
-        vis_points[key] = np.concatenate(
-            (vis_interior_points[key], vis_boundary_points[key])
-        )
-
+    NPOINT_BC = NPOINT_TOP + NPOINT_BOTTOM + NPOINT_LEFT + NPOINT_RIGHT
+    vis_points = geom["rect"].sample_interior(NPOINT_PDE + NPOINT_BC, evenly=True)
     visualizer = {
         "visulzie_u_v": ppsci.visualize.VisualizerVtu(
             vis_points,
@@ -165,11 +156,11 @@
     solver = ppsci.solver.Solver(
         model,
         constraint,
-        output_dir,
+        OUTPUT_DIR,
         optimizer,
         lr_scheduler,
-        epochs,
-        iters_per_epoch,
+        EPOCHS,
+        ITERS_PER_EPOCH,
         eval_during_train=True,
         eval_freq=200,
         equation=equation,
@@ -185,15 +176,16 @@
     solver.visualize()
 
     # directly evaluate pretrained model(optional)
+    logger.init_logger("ppsci", f"{OUTPUT_DIR}/eval.log", "info")
     solver = ppsci.solver.Solver(
         model,
         constraint,
-        output_dir,
+        OUTPUT_DIR,
         equation=equation,
         geom=geom,
         validator=validator,
         visualizer=visualizer,
-        pretrained_model_path=f"{output_dir}/checkpoints/latest",
+        pretrained_model_path=f"{OUTPUT_DIR}/checkpoints/latest",
     )
     solver.eval()
     # visualize prediction for pretrained model(optional)