Unify api invoking and fix bugs&typos (#1246)

* fix math_module

* fix gammaln dtype

* unify invoking device API for different device and fix typos

* fix eval

* fix solver docstrings

Author: HydrogenSulfate

examples/biharmonic2d/biharmonic2d.py

@@ -290,14 +290,14 @@ def evaluate(cfg: DictConfig):
         np.linspace(start=0, stop=cfg.LENGTH, num=num_x, endpoint=True),
         np.linspace(start=0, stop=cfg.WIDTH, num=num_y, endpoint=True),
     )
-    x_faltten = paddle.to_tensor(
+    x_flatten = paddle.to_tensor(
         x_grad.flatten()[:, None], dtype=paddle.get_default_dtype(), stop_gradient=False
     )
-    y_faltten = paddle.to_tensor(
+    y_flatten = paddle.to_tensor(
         y_grad.flatten()[:, None], dtype=paddle.get_default_dtype(), stop_gradient=False
     )
     outs_pred = solver.predict(
-        {"x": x_faltten, "y": y_faltten}, batch_size=num_cords, no_grad=False
+        {"x": x_flatten, "y": y_flatten}, batch_size=num_cords, no_grad=False
     )
 
     # generate label
@@ -309,11 +309,11 @@ def evaluate(cfg: DictConfig):
         paddle.to_tensor(Q, dtype=paddle.get_default_dtype())
         * paddle.sin(
             paddle.to_tensor(np.pi / cfg.LENGTH, dtype=paddle.get_default_dtype())
-            * x_faltten,
+            * x_flatten,
         )
         * paddle.sin(
             paddle.to_tensor(np.pi / cfg.WIDTH, dtype=paddle.get_default_dtype())
-            * y_faltten,
+            * y_flatten,
         )
     )
 
@@ -334,10 +334,10 @@ def compute_outs(w, x, y):
         Q_y = -jacobian((w_x2 + w_y2), y) * D
         return {"Mx": M_x, "Mxy": M_xy, "My": M_y, "Qx": Q_x, "Qy": Q_y, "w": w}
 
-    outs = compute_outs(outs_pred["u"], x_faltten, y_faltten)
+    outs = compute_outs(outs_pred["u"], x_flatten, y_flatten)
 
     # plotting
-    griddata_points = paddle.concat([x_faltten, y_faltten], axis=-1).numpy()
+    griddata_points = paddle.concat([x_flatten, y_flatten], axis=-1).numpy()
     griddata_xi = (x_grad, y_grad)
     boundary = [0, cfg.LENGTH, 0, cfg.WIDTH]
     plotting(
@@ -410,12 +410,13 @@ def inference(cfg: DictConfig):
             start=0, stop=cfg.WIDTH, num=num_y, endpoint=True, dtype=np.float32
         ),
     )
-    x_faltten = x_grad.reshape(-1, 1)
-    y_faltten = y_grad.reshape(-1, 1)
+    x_flatten = x_grad.reshape(-1, 1)
+    y_flatten = y_grad.reshape(-1, 1)
 
-    output_dict = predictor.predict(
-        {"x": x_faltten, "y": y_faltten}, cfg.INFER.batch_size
-    )
+    with ppsci.misc.Timer("infer"):
+        output_dict = predictor.predict(
+            {"x": x_flatten, "y": y_flatten}, cfg.INFER.batch_size
+        )
 
     # mapping data to cfg.INFER.output_keys
     output_dict = {
@@ -424,7 +425,7 @@ def inference(cfg: DictConfig):
     }
 
     # plotting
-    griddata_points = np.concatenate([x_faltten, y_faltten], axis=-1)
+    griddata_points = np.concatenate([x_flatten, y_flatten], axis=-1)
     griddata_xi = (x_grad, y_grad)
     boundary = [0, cfg.LENGTH, 0, cfg.WIDTH]
     plotting(

ppsci/arch/chemprop_molecule_utils.py

@@ -1488,7 +1488,7 @@ def device(self, device: (paddle.CPUPlace, paddle.CUDAPlace, str)) -> None:
     @property
     def cuda(self) -> bool:
         """Whether to use CUDA (i.e., GPUs) or not."""
-        return not self.no_cuda and paddle.device.cuda.device_count() >= 1
+        return not self.no_cuda and paddle.device.device_count() >= 1
 
     @cuda.setter
     def cuda(self, cuda: bool) -> None:
@@ -1534,9 +1534,7 @@ def bond_features_size(self, bond_features_size: int) -> None:
         self._bond_features_size = bond_features_size
 
     def configure(self) -> None:
-        self.add_argument(
-            "--gpu", choices=list(range(paddle.device.cuda.device_count()))
-        )
+        self.add_argument("--gpu", choices=list(range(paddle.device.device_count())))
         self.add_argument(
             "--features_generator", choices=get_available_features_generators()
         )

ppsci/arch/regdgcnn.py

@@ -98,7 +98,7 @@ def get_graph_feature(x, k=20, idx=None):
         .contiguous()
     )
     del x, idx, idx_base
-    paddle.device.cuda.empty_cache()
+    paddle.device.empty_cache()
     return feature
 
 

ppsci/data/process/transform/preprocess.py

@@ -325,7 +325,9 @@ def __call__(
         self, *data: Tuple[Dict[str, np.ndarray], ...]
     ) -> Tuple[Dict[str, np.ndarray], ...]:
         data_dict, label_dict, weight_dict = data
-        data_dict_copy = {**data_dict}
-        label_dict_copy = {**label_dict}
-        weight_dict_copy = {**weight_dict} if weight_dict is not None else {}
+        data_dict_copy: Dict[str, np.ndarray] = {**data_dict}
+        label_dict_copy: Dict[str, np.ndarray] = {**label_dict}
+        weight_dict_copy: Dict[str, np.ndarray] = (
+            {**weight_dict} if weight_dict is not None else {}
+        )
         return self.transform_func(data_dict_copy, label_dict_copy, weight_dict_copy)

ppsci/experimental/math_module.py

@@ -24,6 +24,7 @@
 
 import numpy as np
 import paddle
+from scipy.special import gammaln
 from typing_extensions import Literal
 
 
@@ -149,7 +150,7 @@ def _compatible_meshgrid(*args: paddle.Tensor, **kwargs: paddle.Tensor):
             return paddle.meshgrid(*args, **kwargs)
 
     def _roots(N: int) -> np.ndarray:
-        return np.polynomial.legendre.leggauss(N)[0]
+        return np.polynomial.legendre.leggauss(N)[0].astype("float32")
 
     def _calculate_grid(
         N: int,
@@ -176,7 +177,7 @@ def _resize_roots(
         ):  # scale from [-1,1] to [a,b]
             a = integration_domain[0]
             b = integration_domain[1]
-            return ((b - a) / 2) * roots + ((a + b) / 2)
+            return (((b - a) / 2) * roots + ((a + b) / 2)).astype("float32")
 
         for dim in range(_dim):
             grid_1d.append(_resize_roots(integration_domains[dim], _roots(n_per_dim)))
@@ -220,7 +221,8 @@ def _evaluate_integrand(fn, points, weights=None, fn_args=None) -> paddle.Tensor
             ):  # if the the integrand is multi-dimensional, we need to reshape/repeat weights so they can be broadcast in the *=
                 integrand_shape = result.shape[1:]
                 weights = paddle.repeat_interleave(
-                    paddle.unsqueeze(weights, axis=1), np.prod(integrand_shape)
+                    paddle.unsqueeze(weights, axis=1),
+                    np.prod(integrand_shape, dtype="int64"),
                 ).reshape((weights.shape[0], *(integrand_shape)))
             result *= weights
 
@@ -393,7 +395,7 @@ def int_func(s):
         return _finite_derivative(func, s, dx=h) / (t - s) ** (alpha)
 
     result = (
-        1.0 / paddle.exp(paddle.lgamma(paddle.to_tensor(1.0 - alpha, dtype=dtype)))
+        1.0 / paddle.exp(paddle.to_tensor(gammaln(1.0 - alpha), dtype="float32"))
     ) * gaussian_integrate(
         int_func, dim=1, N=2**10 + 1, integration_domains=[[a, t]], dtype=dtype
     )

ppsci/solver/eval.py

@@ -247,7 +247,7 @@ def _eval_by_batch(
                 metric_dict = metric_func(output_dict, label_dict)
                 for var_name, metric_value in metric_dict.items():
                     if var_name not in metric_dict_group[metric_name]:
-                        metric_dict_group[metric_name] = []
+                        metric_dict_group[metric_name][var_name] = []
                     metric_dict_group[metric_name][var_name].append(
                         metric_value
                         if solver.world_size == 1

ppsci/solver/printer.py

@@ -81,10 +81,10 @@ def log_train_info(
     )
     if solver.benchmark_flag:
         max_mem_reserved_msg = (
-            f"max_mem_reserved: {device.cuda.max_memory_reserved() // (1 << 20)} MB"
+            f"max_mem_reserved: {device.max_memory_reserved() // (1 << 20)} MB"
         )
         max_mem_allocated_msg = (
-            f"max_mem_allocated: {device.cuda.max_memory_allocated() // (1 << 20)} MB"
+            f"max_mem_allocated: {device.max_memory_allocated() // (1 << 20)} MB"
         )
         log_str += f", {max_mem_reserved_msg}, {max_mem_allocated_msg}"
     logger.info(log_str)

ppsci/solver/solver.py

@@ -1225,11 +1225,7 @@ def register_callback_on_iter_end(
         Args:
             callback_fn : Callable[[Solver]]
                 A function that takes a Solver instance as an argument. This function
-                will be called at the end of every iteration.
-
-        Returns:
-        -------
-        None
+                will be called at the start of every iteration.
         """
         self.callbacks_on_iter_end.append(callback_fn)
 

ppsci/utils/reader.py

@@ -177,8 +177,8 @@ def load_vtk_file(
     Returns:
         Dict[str, np.ndarray]: Input coordinates dict, label coordinates dict
     """
-    input_dict = {var: [] for var in input_keys}
-    label_dict = {var: [] for var in label_keys}
+    input_dict: Dict[str, list] = {var: [] for var in input_keys}
+    label_dict: Dict[str, list] = {var: [] for var in label_keys}
     for index in time_index:
         file = filename_without_timeid + f"{index}.vtu"
         mesh = meshio.read(file)
@@ -227,7 +227,7 @@ def load_vtk_with_time_file(file: str) -> Dict[str, np.ndarray]:
 
 def load_dat_file(
     file_path: str,
-    keys: Tuple[str, ...] = None,
+    keys: Optional[Tuple[str, ...]] = None,
     alias_dict: Optional[Dict[str, str]] = None,
 ) -> Dict[str, np.ndarray]:
     """Load *.dat file and fetch data as given keys.