fix bug in test. add doc for forward_common_atomic

deepmodeling · Mar 15, 2024 · 8d05a4e · 8d05a4e
1 parent a15875d
commit 8d05a4e
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 2 deletions.
diff --git a/deepmd/dpmodel/atomic_model/base_atomic_model.py b/deepmd/dpmodel/atomic_model/base_atomic_model.py
@@ -79,6 +79,37 @@ def forward_common_atomic(
         fparam: Optional[np.ndarray] = None,
         aparam: Optional[np.ndarray] = None,
     ) -> Dict[str, np.ndarray]:
+        """Common interface for atomic inference.
+
+        This method accept extended coordinates, extended atom typs, neighbor list,
+        and predict the atomic contribution of the fit property.
+
+        Parameters
+        ----------
+        extended_coord
+            extended coodinates, shape: nf x (nall x 3)
+        extended_atype
+            extended atom typs, shape: nf x nall
+            for a type < 0 indicating the atomic is virtual.
+        nlist
+            neighbor list, shape: nf x nloc x nsel
+        mapping
+            extended to local index mapping, shape: nf x nall
+        fparam
+            frame parameters, shape: nf x dim_fparam
+        aparam
+            atomic parameter, shape: nf x nloc x dim_aparam
+
+        Returns
+        -------
+        ret_dict
+            dict of output atomic properties.
+            should implement the definition of `fitting_output_def`.
+            ret_dit["mask"] of shape nf x nloc will be provided.
+            ret_dit["mask"][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real.
+            ret_dit["mask"][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual.
+
+        """
         _, nloc, _ = nlist.shape
         atype = extended_atype[:, :nloc]
         if self.pair_excl is not None:

diff --git a/deepmd/pt/model/atomic_model/base_atomic_model.py b/deepmd/pt/model/atomic_model/base_atomic_model.py
@@ -108,7 +108,33 @@ def forward_common_atomic(
     ) -> Dict[str, torch.Tensor]:
         """Common interface for atomic inference.
 
-        This method accept
+        This method accept extended coordinates, extended atom typs, neighbor list,
+        and predict the atomic contribution of the fit property.
+
+        Parameters
+        ----------
+        extended_coord
+            extended coodinates, shape: nf x (nall x 3)
+        extended_atype
+            extended atom typs, shape: nf x nall
+            for a type < 0 indicating the atomic is virtual.
+        nlist
+            neighbor list, shape: nf x nloc x nsel
+        mapping
+            extended to local index mapping, shape: nf x nall
+        fparam
+            frame parameters, shape: nf x dim_fparam
+        aparam
+            atomic parameter, shape: nf x nloc x dim_aparam
+
+        Returns
+        -------
+        ret_dict
+            dict of output atomic properties.
+            should implement the definition of `fitting_output_def`.
+            ret_dit["mask"] of shape nf x nloc will be provided.
+            ret_dit["mask"][ff,ii] == 1 indicating the ii-th atom of the ff-th frame is real.
+            ret_dit["mask"][ff,ii] == 0 indicating the ii-th atom of the ff-th frame is virtual.
 
         """
         _, nloc, _ = nlist.shape

diff --git a/source/tests/common/dpmodel/test_nlist.py b/source/tests/common/dpmodel/test_nlist.py
@@ -144,7 +144,7 @@ def setUp(self):
         self.nsel = [10, 10]
         self.ref_nlist = np.array(
             [
-                [-1] * sum(self.sel),
+                [-1] * sum(self.nsel),
                 [1, 1, 1, 1, 1, 1, -1, -1, -1, -1, 2, 2, 2, 2, -1, -1, -1, -1, -1, -1],
                 [1, 1, 1, 1, -1, -1, -1, -1, -1, -1, 2, 2, 2, 2, 2, 2, -1, -1, -1, -1],
             ]