refine type hint and docstring

Author: HydrogenSulfate

ppsci/constraint/boundary_constraint.py

@@ -47,7 +47,7 @@ class BoundaryConstraint(base.Constraint):
             Defaults to None.
         evenly (bool, optional): Whether to use evenly distribution sampling.
             Defaults to False.
-        weight_dict (Optional[Dict[str, Callable]]): Define the weight of each
+        weight_dict (Optional[Dict[str, Union[float, Callable]]]): Define the weight of each
             constraint variable. Defaults to None.
         name (str, optional): Name of constraint object. Defaults to "BC".
 
@@ -78,7 +78,7 @@ def __init__(
         random: Literal["pseudo", "LHS"] = "pseudo",
         criteria: Optional[Callable] = None,
         evenly: bool = False,
-        weight_dict: Optional[Dict[str, Callable]] = None,
+        weight_dict: Optional[Dict[str, Union[float, Callable]]] = None,
         name: str = "BC",
     ):
         self.output_expr = output_expr

ppsci/data/dataset/array_dataset.py

@@ -49,6 +49,8 @@ def __init__(
         super().__init__()
         self.input = input
         self.label = label
+        self.input_keys = tuple(input.keys())
+        self.label_keys = tuple(label.keys())
         self.weight = weight
         self.transforms = transforms
         self._len = len(next(iter(input.values())))

ppsci/data/dataset/csv_dataset.py

@@ -178,7 +178,7 @@ def __init__(
         label_keys: Tuple[str, ...],
         alias_dict: Optional[Dict[str, str]] = None,
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
-        timestamps: Optional[Tuple[Union[int, float], ...]] = None,
+        timestamps: Optional[Tuple[float, ...]] = None,
         transforms: Optional[vision.Compose] = None,
     ):
         super().__init__()

ppsci/data/dataset/mat_dataset.py

@@ -178,7 +178,7 @@ def __init__(
         label_keys: Tuple[str, ...] = (),
         alias_dict: Optional[Dict[str, str]] = None,
         weight_dict: Optional[Dict[str, Union[Callable, float]]] = None,
-        timestamps: Optional[Tuple[Union[int, float], ...]] = None,
+        timestamps: Optional[Tuple[float, ...]] = None,
         transforms: Optional[vision.Compose] = None,
     ):
         super().__init__()

ppsci/data/process/transform/preprocess.py

@@ -31,7 +31,7 @@ class Translate:
         >>> translate = ppsci.data.transform.Translate({"x": 1.0, "y": -1.0})
     """
 
-    def __init__(self, offset: Dict[str, Union[int, float]]):
+    def __init__(self, offset: Dict[str, float]):
         self.offset = offset
 
     def __call__(self, data_dict):
@@ -53,7 +53,7 @@ class Scale:
         >>> translate = ppsci.data.transform.Scale({"x": 1.5, "y": 2.0})
     """
 
-    def __init__(self, scale: Dict[str, Union[int, float]]):
+    def __init__(self, scale: Dict[str, float]):
         self.scale = scale
 
     def __call__(self, data_dict):