unify return type of metric module

Author: HydrogenSulfate

ppsci/metric/anomaly_coef.py

@@ -63,12 +63,11 @@ def __init__(
         unlog: bool = False,
         scale: float = 1e-5,
     ):
-        super().__init__()
+        super().__init__(keep_batch)
         self.num_lat = num_lat
         self.mean = (
             None if mean is None else paddle.to_tensor(mean, paddle.get_default_dtype())
         )
-        self.keep_batch = keep_batch
         self.variable_dict = variable_dict
         self.unlog = unlog
         self.scale = scale
@@ -110,14 +109,10 @@ def forward(self, output_dict, label_dict):
                     if self.keep_batch:
                         metric_dict[f"{key}.{variable_name}"] = rmse[:, idx]
                     else:
-                        metric_dict[f"{key}.{variable_name}"] = float(
-                            rmse[:, idx].mean()
-                        )
+                        metric_dict[f"{key}.{variable_name}"] = rmse[:, idx].mean()
             else:
                 if self.keep_batch:
-                    rmse = rmse.mean(axis=1)
-                    metric_dict[key] = rmse
+                    metric_dict[key] = rmse.mean(axis=1)
                 else:
-                    rmse = rmse.mean()
-                    metric_dict[key] = float(rmse)
+                    metric_dict[key] = rmse.mean()
         return metric_dict

ppsci/metric/base.py

@@ -18,5 +18,6 @@
 class Metric(nn.Layer):
     """Base class for metric."""
 
-    def __init__(self):
+    def __init__(self, keep_batch: str = False):
         super().__init__()
+        self.keep_batch = keep_batch

ppsci/metric/l2_rel.py

@@ -24,13 +24,18 @@ class L2Rel(base.Metric):
     metric = \dfrac{\Vert x-y \Vert_2}{\Vert y \Vert_2}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.L2Rel()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        if keep_batch is not False:
+            raise ValueError(f"keep_batch should be False, but got {keep_batch}.")
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
@@ -39,6 +44,6 @@ def forward(self, output_dict, label_dict):
             rel_l2 = paddle.norm(label_dict[key] - output_dict[key]) / paddle.norm(
                 label_dict[key]
             )
-            metric_dict[key] = float(rel_l2)
+            metric_dict[key] = rel_l2
 
         return metric_dict

ppsci/metric/mae.py

@@ -34,8 +34,7 @@ class MAE(base.Metric):
     """
 
     def __init__(self, keep_batch: bool = False):
-        super().__init__()
-        self.keep_batch = keep_batch
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
@@ -45,6 +44,6 @@ def forward(self, output_dict, label_dict):
             if self.keep_batch:
                 metric_dict[key] = mae.mean(axis=tuple(range(1, mae.ndim)))
             else:
-                metric_dict[key] = float(mae.mean())
+                metric_dict[key] = mae.mean()
 
         return metric_dict

ppsci/metric/mse.py

@@ -25,19 +25,25 @@ class MSE(base.Metric):
     metric = \dfrac{1}{N}\sum\limits_{i=1}^{N}{(x_i-y_i)^2}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.MSE()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
         metric_dict = {}
         for key in label_dict:
-            mse = F.mse_loss(output_dict[key], label_dict[key], "mean")
-            metric_dict[key] = float(mse)
+            mse = F.mse_loss(output_dict[key], label_dict[key], "none")
+            if self.keep_batch:
+                metric_dict[key] = mse.mean(axis=tuple(range(1, mse.ndim)))
+            else:
+                metric_dict[key] = mse.mean()
 
         return metric_dict

ppsci/metric/rmse.py

@@ -31,20 +31,25 @@ class RMSE(base.Metric):
     metric = \sqrt{\dfrac{1}{N}\sum\limits_{i=1}^{N}{(x_i-y_i)^2}}
     $$
 
+    Args:
+        keep_batch (bool, optional): Whether keep batch axis. Defaults to False.
+
     Examples:
         >>> import ppsci
         >>> metric = ppsci.metric.RMSE()
     """
 
-    def __init__(self):
-        super().__init__()
+    def __init__(self, keep_batch: bool = False):
+        if keep_batch is not False:
+            raise ValueError(f"keep_batch should be False, but got {keep_batch}.")
+        super().__init__(keep_batch)
 
     @paddle.no_grad()
     def forward(self, output_dict, label_dict):
         metric_dict = {}
         for key in label_dict:
             rmse = F.mse_loss(output_dict[key], label_dict[key], "mean") ** 0.5
-            metric_dict[key] = float(rmse)
+            metric_dict[key] = rmse
 
         return metric_dict
 
@@ -88,18 +93,16 @@ def __init__(
         unlog: bool = False,
         scale: float = 1e-5,
     ):
-        super().__init__()
+        super().__init__(keep_batch)
         self.num_lat = num_lat
         self.std = (
             None
             if std is None
             else paddle.to_tensor(std, paddle.get_default_dtype()).reshape((1, -1))
         )
-        self.keep_batch = keep_batch
         self.variable_dict = variable_dict
         self.unlog = unlog
         self.scale = scale
-
         self.weight = self.get_latitude_weight(num_lat)
 
     def get_latitude_weight(self, num_lat: int = 720):
@@ -127,18 +130,10 @@ def forward(self, output_dict, label_dict):
                 rmse = rmse * self.std
             if self.variable_dict is not None:
                 for variable_name, idx in self.variable_dict.items():
-                    if self.keep_batch:
-                        metric_dict[f"{key}.{variable_name}"] = rmse[:, idx]
-                    else:
-                        metric_dict[f"{key}.{variable_name}"] = float(
-                            rmse[:, idx].mean()
-                        )
+                    metric_dict[f"{key}.{variable_name}"] = (
+                        rmse[:, idx] if self.keep_batch else rmse[:, idx].mean()
+                    )
             else:
-                if self.keep_batch:
-                    rmse = rmse.mean(axis=1)
-                    metric_dict[key] = rmse
-                else:
-                    rmse = rmse.mean()
-                    metric_dict[key] = float(rmse)
+                metric_dict[key] = rmse.mean(axis=1) if self.keep_batch else rmse.mean()
 
         return metric_dict

ppsci/solver/eval.py

@@ -134,7 +134,9 @@ def _eval_by_dataset(solver, epoch_id: int, log_freq: int) -> float:
                     solver.eval_output_info[metric_str] = misc.AverageMeter(
                         metric_str, ".5f"
                     )
-                solver.eval_output_info[metric_str].update(metric_value, num_samples)
+                solver.eval_output_info[metric_str].update(
+                    float(metric_value), num_samples
+                )
 
         if target_metric is None:
             tmp = metric

ppsci/solver/train.py

@@ -169,6 +169,8 @@ def closure():
             return total_loss
 
         solver.optimizer.step(closure)
+
+        # update learning rate by step
         if solver.lr_scheduler is not None and not solver.lr_scheduler.by_epoch:
             solver.lr_scheduler.step()
 

ppsci/utils/misc.py

@@ -14,6 +14,9 @@
 
 import collections
 import random
+from typing import Dict
+from typing import List
+from typing import Tuple
 
 import numpy as np
 import paddle
@@ -94,7 +97,16 @@ def __str__(self):
         return "".join([str((k, v)) for k, v in self.items()])
 
 
-def convert_to_dict(array, keys):
+def convert_to_dict(array: np.ndarray, keys: Tuple[str, ...]) -> Dict[str, np.ndarray]:
+    """Split given array into single channel array at axis -1 in order of given keys.
+
+    Args:
+        array (np.ndarray): Array to be splited.
+        keys (Tuple[str, ...]):Keys used in split.
+
+    Returns:
+        Dict[str, np.ndarray]: Splited dict.
+    """
     if array.shape[-1] != len(keys):
         raise ValueError(
             f"dim of array({array.shape[-1]}) must equal to " f"len(keys)({len(keys)})"
@@ -104,7 +116,9 @@ def convert_to_dict(array, keys):
     return {key: split_array[i] for i, key in enumerate(keys)}
 
 
-def all_gather(tensor, concat=True, axis=0):
+def all_gather(
+    tensor: paddle.Tensor, concat: bool = True, axis: int = 0
+) -> List[paddle.Tensor]:
     """Gather tensor from all devices, concatenate them along given axis if specified.
 
     Args:
@@ -122,29 +136,78 @@ def all_gather(tensor, concat=True, axis=0):
     return result
 
 
-def convert_to_array(dict, keys):
+def convert_to_array(dict: Dict[str, np.ndarray], keys: Tuple[str, ...]) -> np.ndarray:
+    """Concatenate arrays in axis -1 in order of given keys.
+
+    Args:
+        dict (Dict[str, np.ndarray]): Dict contains arrays.
+        keys (Tuple[str, ...]): Concatenate keys used in concatenation.
+
+    Returns:
+        np.ndarray: Concatenated array.
+    """
     return np.concatenate([dict[key] for key in keys], axis=-1)
 
 
-def concat_dict_list(dict_list):
+def concat_dict_list(
+    dict_list: Tuple[Dict[str, np.ndarray], ...]
+) -> Dict[str, np.ndarray]:
+    """concatenate arrays in tuple of dicts at axis 0.
+
+    Args:
+        dict_list (Tuple[Dict[str, np.ndarray], ...]): Tuple of dicts.
+
+    Returns:
+        Dict[str, np.ndarray]: A dict with concatenated arrays for each key.
+    """
     ret = {}
     for key in dict_list[0].keys():
         ret[key] = np.concatenate([_dict[key] for _dict in dict_list], axis=0)
     return ret
 
 
-def stack_dict_list(dict_list):
+def stack_dict_list(
+    dict_list: Tuple[Dict[str, np.ndarray], ...]
+) -> Dict[str, np.ndarray]:
+    """Stack arrays in tuple of dicts at axis 0.
+
+    Args:
+        dict_list (Tuple[Dict[str, np.ndarray], ...]): Tuple of dicts.
+
+    Returns:
+        Dict[str, np.ndarray]: A dict with stacked arrays for each key.
+    """
     ret = {}
     for key in dict_list[0].keys():
         ret[key] = np.stack([_dict[key] for _dict in dict_list], axis=0)
     return ret
 
 
-def typename(object):
+def typename(object: object) -> str:
+    """Return type name of given object.
+
+    Args:
+        object (object): Python object which is instantiated from a class.
+
+    Returns:
+        str: Class name of given object.
+    """
     return object.__class__.__name__
 
 
-def combine_array_with_time(x, t):
+def combine_array_with_time(x: np.ndarray, t: Tuple[int, ...]) -> np.ndarray:
+    """Combine given data x with time sequence t.
+    Given x with shape (N, D) and t with shape (T, ),
+    this function will repeat t_i for N times and will concat it with data x for each t_i in t,
+    finally return the stacked result, whic is of shape (NxT, D+1).
+
+    Args:
+        x (np.ndarray): Points data with shape (N, D).
+        t (Tuple[int, ...]): Time sequence with shape (T, ).
+
+    Returns:
+        np.ndarray: Combined data with shape of (NxT, D+1).
+    """
     nx = len(x)
     tx = []
     for ti in t:
@@ -158,6 +221,11 @@ def combine_array_with_time(x, t):
 
 
 def set_random_seed(seed: int):
+    """Set numpy, random, paddle random_seed to given seed.
+
+    Args:
+        seed (int): Random seed.
+    """
     paddle.seed(seed)
     np.random.seed(seed)
     random.seed(seed)

ppsci/visualize/visualizer.py

@@ -19,7 +19,6 @@
 from typing import Tuple
 
 import numpy as np
-import paddle
 
 from ppsci.visualize import base
 from ppsci.visualize import plot