Align docstring parameters with function signatures across PaddleScience

ppsci/arch/crystalgraphconvnet.py

@@ -123,10 +123,10 @@ def forward(self, input) -> paddle.Tensor:
 
         Args:
             input (list): List of input, which includes the following elements:
-                atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type.
-                nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors.
-                nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom.
-                crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx.
+                - atom_fea (paddle.Tensor): Shape (N, orig_atom_fea_len). Atom features from atom type.
+                - nbr_fea (paddle.Tensor): Shape (N, M, nbr_fea_len). Bond features of each atom's M neighbors.
+                - nbr_fea_idx (paddle.Tensor): Shape (N, M). Indices of M neighbors of each atom.
+                - crystal_atom_idx (list): List of paddle.Tensor of length N0. Mapping from the crystal idx to atom idx.
 
         Returns:
             paddle.Tensor: Shape (N,). Atom hidden features after convolution.

ppsci/arch/cuboid_transformer_decoder.py

@@ -56,10 +56,10 @@ def reset_parameters(self):
     def forward(self, x):
         """
         Args:
-            x : Shape (B, T, H, W, C)
+            x (paddle.Tensor): Shape (B, T, H, W, C).
 
         Returns:
-            out : the x + positional embeddings
+            out (paddle.Tensor): The input plus positional embeddings.
         """
 
         _, T, H, W, _ = x.shape
@@ -800,10 +800,10 @@ def forward(self, x):
         """
 
         Args:
-            x : (B, T, H, W, C) or (B, C, T, H, W)
+            x (paddle.Tensor): (B, T, H, W, C) or (B, C, T, H, W).
 
         Returns:
-            out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out)
+            out (paddle.Tensor): (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out).
         """
 
         if self.layout == "THWC":
@@ -1193,8 +1193,9 @@ def reset_parameters(self):
     def forward(self, x, mem_l, mem_global_vector_l=None):
         """
         Args:
-            x : Shape (B, T_top, H_top, W_top, C).
-            mem_l : A list of memory tensors.
+            x (paddle.Tensor): Shape (B, T_top, H_top, W_top, C).
+            mem_l (List[paddle.Tensor]): A list of memory tensors.
+            mem_global_vector_l (Optional[List[paddle.Tensor]]): Optional list of global memory vectors.
         """
 
         B, T_top, H_top, W_top, C = x.shape

ppsci/arch/cuboid_transformer_encoder.py

@@ -79,10 +79,10 @@ def forward(self, x):
         """
 
         Args:
-            x : (B, T, H, W, C)
+            x (paddle.Tensor): (B, T, H, W, C).
 
         Returns:
-            out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim)
+            out (paddle.Tensor): Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim).
         """
 
         B, T, H, W, C = x.shape
@@ -218,10 +218,10 @@ def reset_parameters(self):
     def forward(self, data):
         """
         Args:
-            x : Shape (B, seq_length, C_in)
+            data (paddle.Tensor): Shape (B, seq_length, C_in).
 
         Returns:
-            out : Shape (B, seq_length, C_out)
+            out (paddle.Tensor): Shape (B, seq_length, C_out).
         """
 
         residual = data
@@ -1480,7 +1480,8 @@ def get_mem_shapes(self):
     def forward(self, x, global_vectors=None):
         """
         Args:
-            x : Shape (B, T, H, W, C)
+            x (paddle.Tensor): Shape (B, T, H, W, C).
+            global_vectors (Optional[paddle.Tensor]): Optional global vector memory.
 
         Returns:
             out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For

ppsci/arch/extformer_moe_cuboid_decoder.py

@@ -65,10 +65,10 @@ def reset_parameters(self):
     def forward(self, x):
         """
         Args:
-            x : Shape (B, T, H, W, C)
+            x (paddle.Tensor): Shape (B, T, H, W, C).
 
         Returns:
-            out : the x + positional embeddings
+            out (paddle.Tensor): The input plus positional embeddings.
         """
 
         _, T, H, W, _ = x.shape
@@ -900,10 +900,10 @@ def forward(self, x):
         """
 
         Args:
-            x : (B, T, H, W, C) or (B, C, T, H, W)
+            x (paddle.Tensor): (B, T, H, W, C) or (B, C, T, H, W).
 
         Returns:
-            out : (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out)
+            out (paddle.Tensor): (B, T, H_new, W_out, C_out) or (B, C, T, H_out, W_out).
         """
 
         if self.layout == "THWC":
@@ -1304,8 +1304,9 @@ def reset_parameters(self):
     def forward(self, x, mem_l, mem_global_vector_l=None):
         """
         Args:
-            x : Shape (B, T_top, H_top, W_top, C).
-            mem_l : A list of memory tensors.
+            x (paddle.Tensor): Shape (B, T_top, H_top, W_top, C).
+            mem_l (List[paddle.Tensor]): A list of memory tensors.
+            mem_global_vector_l (Optional[List[paddle.Tensor]]): Optional list of global memory vectors.
         """
 
         B, T_top, H_top, W_top, C = x.shape

ppsci/arch/extformer_moe_cuboid_encoder.py

@@ -81,10 +81,10 @@ def forward(self, x):
         """
 
         Args:
-            x : (B, T, H, W, C)
+            x (paddle.Tensor): (B, T, H, W, C).
 
         Returns:
-            out : Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim)
+            out (paddle.Tensor): Shape (B, T // downsample[0], H // downsample[1], W // downsample[2], out_dim).
         """
 
         B, T, H, W, C = x.shape
@@ -242,10 +242,10 @@ def reset_parameters(self):
     def forward(self, data):
         """
         Args:
-            x : Shape (B, seq_length, C_in)
+            data (paddle.Tensor): Shape (B, seq_length, C_in).
 
         Returns:
-            out : Shape (B, seq_length, C_out)
+            out (paddle.Tensor): Shape (B, seq_length, C_out).
         """
 
         residual = data
@@ -1638,7 +1638,8 @@ def get_mem_shapes(self):
     def forward(self, x, global_vectors=None):
         """
         Args:
-            x : Shape (B, T, H, W, C)
+            x (paddle.Tensor): Shape (B, T, H, W, C).
+            global_vectors (Optional[paddle.Tensor]): Optional global vector memory.
 
         Returns:
             out (List[paddle.Tensor,..]): A list of tensors from the bottom layer to the top layer of the encoder. For

ppsci/arch/moflow_glow.py

@@ -303,7 +303,6 @@ def _squeeze(self, x):
 
         Args:
             x (paddle.Tensor): Input to squeeze or unsqueeze.
-            reverse (bool): Reverse the operation, i.e., unsqueeze.
 
         Returns:
             x (paddle.Tensor): Squeezed or unsqueezed tensor.

ppsci/data/dataset/cgcnn_dataset.py

@@ -120,7 +120,7 @@ def expand(self, distances):
         Apply Gaussian distance filter to a numpy distance array.
 
         Args:
-            distance (np.array): n-dimensional distance matrix of any shape.
+            distances (np.array): n-dimensional distance matrix of any shape.
 
         Returns:
             np.array: Expanded distance matrix with the last dimension of length len(self.filter).

ppsci/data/dataset/sevir_dataset.py

@@ -555,7 +555,7 @@ def _load_event_batch(self, event_idx, event_batch_size):
         """Loads a selected batch of events (not batch of sequences) into memory.
 
         Args:
-            idx (int): The index of the event in the batch.
+            event_idx (int): The index of the event in the batch.
             event_batch_size (int): Event_batch[i] = all_type_i_available_events[idx:idx + event_batch_size]
 
         Returns:
@@ -686,8 +686,10 @@ def downsample_data_dict(
 
         Args:
             data_dict (Dict[str, Union[np.array, paddle.Tensor]]): The dict of data.
+            data_types (Optional[Sequence[str]]): Data types to be downsampled. Defaults to all keys in `data_dict`.
             factors_dict (Optional[Dict[str, Sequence[int]]]):each element `factors` is
                 a Sequence of int, representing (t_factor, h_factor, w_factor).
+            layout (str): Layout string, such as \"NHWT\".
 
         Returns:
             downsampled_data_dict (Dict[str, paddle.Tensor]): Modify on a deep copy of

ppsci/data/dataset/synthemol_dataset.py

@@ -267,7 +267,9 @@ def make_mol(self, s: str, keep_h: bool, add_h: bool):
         Args:
             s: SMILES string.
             keep_h: Boolean whether to keep hydrogens in the input smiles. This does not add hydrogens, it only keeps them if they are specified.
-        return:
+            add_h: Whether to add hydrogens to the molecule.
+
+        Returns:
             RDKit molecule.
         """
         if keep_h:

ppsci/experimental/math_module.py

@@ -160,7 +160,7 @@ def _calculate_grid(
 
         Args:
             N (int): Number of points.
-            integration_domain (paddle.Tensor): Integration domain.
+            integration_domains (paddle.Tensor): Integration domains.
 
         Returns:
             Tuple[paddle.Tensor, paddle.Tensor, int]: Grid points, grid widths and
@@ -275,6 +275,7 @@ def _calculate_result(
             dim (int): Dimensionality.
             n_per_dim (int): Number of grid slices per dimension.
             hs (paddle.Tensor): Distances between grid slices for each dimension.
+            integration_domains (paddle.Tensor): Integration domains.
 
         Returns:
             paddle.Tensor: Quadrature result.

ppsci/geometry/geometry.py

@@ -139,10 +139,10 @@ def sample_interior(
 
         Args:
             n (int): Number of points.
-            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo".
-                pseudo: Pseudo random.
-                Halton: Halton sequence.
-                LHS: Latin Hypercube Sampling.
+            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". Options:
+                - "pseudo": Pseudo random.
+                - "Halton": Halton sequence.
+                - "LHS": Latin Hypercube Sampling.
             criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given
                 coords from different dimension and return a boolean array with shape [n,].
                 Defaults to None.
@@ -241,10 +241,10 @@ def sample_boundary(
 
         Args:
             n (int): Number of points.
-            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo".
-                pseudo: Pseudo random.
-                Halton: Halton sequence.
-                LHS: Latin Hypercube Sampling.
+            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". Options:
+                - "pseudo": Pseudo random.
+                - "Halton": Halton sequence.
+                - "LHS": Latin Hypercube Sampling.
             criteria (Optional[Callable[..., np.ndarray]]): Criteria function. Given
                 coords from different dimension and return a boolean array with shape [n,].
                 Defaults to None.
@@ -351,10 +351,10 @@ def random_points(
 
         Args:
             n (int): Number of points.
-            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo".
-                pseudo: Pseudo random.
-                Halton: Halton sequence.
-                LHS: Latin Hypercube Sampling.
+            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". Options:
+                - "pseudo": Pseudo random.
+                - "Halton": Halton sequence.
+                - "LHS": Latin Hypercube Sampling.
 
         Returns:
             np.ndarray: Random points in the geometry. The shape is [N, D].
@@ -400,10 +400,10 @@ def random_boundary_points(
 
         Args:
             n (int): Number of points.
-            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo".
-                pseudo: Pseudo random.
-                Halton: Halton sequence.
-                LHS: Latin Hypercube Sampling.
+            random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo". Options:
+                - "pseudo": Pseudo random.
+                - "Halton": Halton sequence.
+                - "LHS": Latin Hypercube Sampling.
 
         Returns:
             np.ndarray: Random points on the boundary. The shape is [N, D].

ppsci/geometry/mesh.py

@@ -1320,6 +1320,8 @@ def sample_in_triangle(v0, v1, v2, n, random="pseudo", criteria=None):
         v1 (np.ndarray): Coordinates of the second vertex of an triangle with shape of [3, ].
         v2 (np.ndarray): Coordinates of the third vertex of an triangle with shape of [3, ].
         n (int): Number of points to be sampled.
+        random (Literal["pseudo", "Halton", "LHS"]): Random method. Defaults to "pseudo".
+        criteria (Optional[Callable[..., np.ndarray]]): Criteria function to filter sampled points. Defaults to None.
 
     Returns:
         np.ndarray: Coordinates of sampled n points with shape of [n, 3].

ppsci/solver/solver.py

@@ -1070,14 +1070,15 @@ def no_sync_context_manager(
         ddp_model: paddle.DataParallel,
     ) -> contextlib.AbstractContextManager:
         """Smart no_sync context manager for given model.
-        NOTE: Only `paddle.DataParallel` object has `no_sync` interface.
-
-        Args:
-            enable (bool): Enable no_sync.
-
-        Returns:
-            contextlib.AbstractContextManager: Smart no_sync context manager.
-        """
+        NOTE: Only `paddle.DataParallel` object has `no_sync` interface.
+
+        Args:
+            enable (bool): Enable no_sync.
+            ddp_model (paddle.DataParallel): Model to apply `no_sync` on.
+
+        Returns:
+            contextlib.AbstractContextManager: Smart no_sync context manager.
+        """
         if enable:
             if isinstance(self.model, ppsci.arch.ModelList):
                 for model in self.model.model_list:

ppsci/utils/download.py

@@ -78,6 +78,8 @@ def get_path_from_url(url, root_dir, md5sum=None, check_exist=True, decompress=T
         root_dir (str): Root dir for downloading, it should be
                         WEIGHTS_HOME or DATASET_HOME
         md5sum (str): md5 sum of download package
+        check_exist (bool): Whether to skip download if file exists and matches md5. Defaults to True.
+        decompress (bool): Whether to decompress tar/zip files after download. Defaults to True.
 
     Returns:
         str: a local path to save downloaded models & weights & datasets.

ppsci/utils/reader.py

@@ -50,6 +50,7 @@ def load_csv_file(
         keys (Tuple[str, ...]): Required fetching keys.
         alias_dict (Optional[Dict[str, str]]): Alias for keys,
             i.e. {inner_key: outer_key}. Defaults to None.
+        delimiter (str, optional): Delimiter used to split columns. Defaults to ",".
         encoding (str, optional): Encoding code when open file. Defaults to "utf-8".
 
     Returns:

ppsci/visualize/__init__.py

@@ -62,8 +62,6 @@ def build_visualizer(cfg: DictConfig):
 
     Args:
         cfg (DictConfig): Visualizer(s) config list.
-        geom_dict (Dct[str, Geometry]): Geometry(ies) in dict.
-        equation_dict (Dct[str, Equation]): Equation(s) in dict.
 
     Returns:
         Dict[str, Visualizer]: Visualizer(s) in dict.