update example code for PINNPredictor (PaddlePaddle#918)

liaoxin2 · Jun 3, 2024 · 0780c19 · 0780c19
1 parent 2fdee45
commit 0780c19
Show file tree

Hide file tree

Showing 3 changed files with 69 additions and 7 deletions.
diff --git a/deploy/python_infer/__init__.py b/deploy/python_infer/__init__.py
@@ -15,8 +15,13 @@
 from deploy.python_infer.base import Predictor
 from deploy.python_infer.pinn_predictor import PINNPredictor
 
+
 # alias as PINNPredictor can be used in most cases
-GeneralPredictor = PINNPredictor
+class GeneralPredictor(PINNPredictor):
+    """Use PINNPredictor as GeneralPredictor."""
+
+    pass
+
 
 __all__ = [
     "Predictor",

diff --git a/deploy/python_infer/pinn_predictor.py b/deploy/python_infer/pinn_predictor.py
@@ -30,6 +30,57 @@ class PINNPredictor(base.Predictor):
 
     Args:
         cfg (DictConfig): Running configuration.
+
+    Examples:
+        >>> import numpy as np
+        >>> import paddle
+        >>> from omegaconf import DictConfig
+        >>> from paddle.static import InputSpec
+        >>> import ppsci
+        >>> from deploy.python_infer import pinn_predictor
+        >>> model = ppsci.arch.MLP(("x", "y"), ("u", "v", "p"), 3, 16)
+        >>> static_model = paddle.jit.to_static(
+        ...     model,
+        ...     input_spec=[
+        ...         {
+        ...             key: InputSpec([None, 1], "float32", name=key)
+        ...             for key in model.input_keys
+        ...         },
+        ...     ],
+        ... )
+        >>> paddle.jit.save(static_model, "./inference")
+        >>> cfg = DictConfig(
+        ...     {
+        ...         "log_freq": 10,
+        ...         "INFER": {
+        ...             "pdmodel_path": "./inference.pdmodel",
+        ...             "pdiparams_path": "./inference.pdiparams",
+        ...             "device": "cpu",
+        ...             "engine": "native",
+        ...             "precision": "fp32",
+        ...             "onnx_path": None,
+        ...             "ir_optim": True,
+        ...             "min_subgraph_size": 15,
+        ...             "gpu_mem": 500,
+        ...             "gpu_id": 0,
+        ...             "max_batch_size": 10,
+        ...             "num_cpu_threads": 10,
+        ...         }
+        ...     }
+        ... )
+        >>> predictor = pinn_predictor.PINNPredictor(cfg) # doctest: +SKIP
+        >>> pred = predictor.predict(
+        ...     {
+        ...         "x": np.random.randn(4, 1).astype("float32"),
+        ...         "y": np.random.randn(4, 1).astype("float32"),
+        ...     },
+        ...     batch_size=2,
+        ... ) # doctest: +SKIP
+        >>> for k, v in pred.items(): # doctest: +SKIP
+        ...     print(k, v.shape) # doctest: +SKIP
+        save_infer_model/scale_0.tmp_0 (4, 1)
+        save_infer_model/scale_1.tmp_0 (4, 1)
+        save_infer_model/scale_2.tmp_0 (4, 1)
     """
 
     def __init__(

diff --git a/ppsci/geometry/mesh.py b/ppsci/geometry/mesh.py
@@ -23,6 +23,7 @@
 
 import numpy as np
 import paddle
+from typing_extensions import Literal
 
 from ppsci.geometry import geometry
 from ppsci.geometry import geometry_3d
@@ -351,7 +352,7 @@ def inflated_random_points(self, n, distance, random="pseudo", criteria=None):
 
     def _approximate_area(
         self,
-        random: str = "pseudo",
+        random: Literal["pseudo"] = "pseudo",
         criteria: Optional[Callable] = None,
         n_appr: int = 10000,
     ) -> float:
@@ -444,7 +445,12 @@ def random_boundary_points(self, n, random="pseudo"):
         return points, normal, areas
 
     def sample_boundary(
-        self, n, random="pseudo", criteria=None, evenly=False, inflation_dist=None
+        self,
+        n: int,
+        random: Literal["pseudo"] = "pseudo",
+        criteria: Optional[Callable[..., np.ndarray]] = None,
+        evenly: bool = False,
+        inflation_dist: Union[float, Tuple[float, ...]] = None,
     ) -> Dict[str, np.ndarray]:
         # TODO(sensen): support for time-dependent points(repeat data in time)
         if inflation_dist is not None:
@@ -561,10 +567,10 @@ def random_points(self, n, random="pseudo", criteria=None):
 
     def sample_interior(
         self,
-        n,
-        random="pseudo",
-        criteria=None,
-        evenly=False,
+        n: int,
+        random: Literal["pseudo"] = "pseudo",
+        criteria: Optional[Callable[..., np.ndarray]] = None,
+        evenly: bool = False,
         compute_sdf_derivatives: bool = False,
     ):
         """Sample random points in the geometry and return those meet criteria."""