Merge branch '3831' of github.com:Chengqian-Zhang/deepmd-kit into 3831

Author: Chengqian-Zhang

.github/workflows/test_cuda.yml

@@ -34,12 +34,12 @@ jobs:
       uses: mpi4py/setup-mpi@v1
       with:
         mpi: mpich
+    - name: Install wget and unzip
+      run: apt-get update && apt-get install -y wget unzip
     - uses: lukka/get-cmake@latest
       with:
         useLocalCache: true
         useCloudCache: false
-    - name: Install wget and unzip
-      run: apt-get update && apt-get install -y wget unzip
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \

deepmd/dpmodel/descriptor/dpa1.py

@@ -804,7 +804,10 @@ def call(
         nf, nloc, nnei, _ = dmatrix.shape
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
         # nfnl x nnei
+        exclude_mask = exclude_mask.reshape(nf * nloc, nnei)
+        # nfnl x nnei
         nlist = nlist.reshape(nf * nloc, nnei)
+        nlist = np.where(exclude_mask, nlist, -1)
         # nfnl x nnei x 4
         dmatrix = dmatrix.reshape(nf * nloc, nnei, 4)
         # nfnl x nnei x 1
@@ -824,8 +827,6 @@ def call(
             nf * nloc, nnei, self.tebd_dim
         )
         ng = self.neuron[-1]
-        # nfnl x nnei
-        exclude_mask = exclude_mask.reshape(nf * nloc, nnei)
         # nfnl x nnei x 4
         rr = dmatrix.reshape(nf * nloc, nnei, 4)
         rr = rr * exclude_mask[:, :, None]

deepmd/dpmodel/descriptor/repformers.py

@@ -338,7 +338,7 @@ def call(
         mapping: Optional[np.ndarray] = None,
     ):
         exclude_mask = self.emask.build_type_exclude_mask(nlist, atype_ext)
-        nlist = nlist * exclude_mask
+        nlist = np.where(exclude_mask, nlist, -1)
         # nf x nloc x nnei x 4
         dmatrix, diff, sw = self.env_mat.call(
             coord_ext, atype_ext, nlist, self.mean, self.stddev

deepmd/dpmodel/descriptor/se_e2_a.py

@@ -117,6 +117,9 @@ class DescrptSeA(NativeOP, BaseDescriptor):
             The precision of the embedding net parameters. Supported options are |PRECISION|
     spin
             The deepspin object.
+    ntypes : int
+            Number of element types.
+            Not used in this descriptor, only to be compat with input.
 
     Limitations
     -----------
@@ -150,9 +153,11 @@ def __init__(
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         spin: Optional[Any] = None,
+        ntypes: Optional[int] = None,  # to be compat with input
         # consistent with argcheck, not used though
         seed: Optional[int] = None,
     ) -> None:
+        del ntypes
         ## seed, uniform_seed, not included.
         if spin is not None:
             raise NotImplementedError("spin is not implemented")

deepmd/dpmodel/descriptor/se_r.py

@@ -75,6 +75,9 @@ class DescrptSeR(NativeOP, BaseDescriptor):
             The precision of the embedding net parameters. Supported options are |PRECISION|
     spin
             The deepspin object.
+    ntypes : int
+            Number of element types.
+            Not used in this descriptor, only to be compat with input.
 
     Limitations
     -----------
@@ -107,9 +110,11 @@ def __init__(
         activation_function: str = "tanh",
         precision: str = DEFAULT_PRECISION,
         spin: Optional[Any] = None,
+        ntypes: Optional[int] = None,  # to be compat with input
         # consistent with argcheck, not used though
         seed: Optional[int] = None,
     ) -> None:
+        del ntypes
         ## seed, uniform_seed, not included.
         if not type_one_side:
             raise NotImplementedError("type_one_side == False not implemented")

deepmd/dpmodel/descriptor/se_t.py

@@ -78,6 +78,9 @@ class DescrptSeT(NativeOP, BaseDescriptor):
             If the weights of embedding net are trainable.
     seed : int, Optional
             Random seed for initializing the network parameters.
+    ntypes : int
+            Number of element types.
+            Not used in this descriptor, only to be compat with input.
     """
 
     def __init__(
@@ -94,7 +97,9 @@ def __init__(
         precision: str = DEFAULT_PRECISION,
         trainable: bool = True,
         seed: Optional[int] = None,
+        ntypes: Optional[int] = None,  # to be compat with input
     ) -> None:
+        del ntypes
         self.rcut = rcut
         self.rcut_smth = rcut_smth
         self.sel = sel

deepmd/entrypoints/doc.py

@@ -4,6 +4,7 @@
 from deepmd.utils.argcheck import (
     gen_doc,
     gen_json,
+    gen_json_schema,
 )
 
 __all__ = ["doc_train_input"]
@@ -15,6 +16,8 @@ def doc_train_input(*, out_type: str = "rst", **kwargs):
         doc_str = gen_doc(make_anchor=True)
     elif out_type == "json":
         doc_str = gen_json()
+    elif out_type == "json_schema":
+        doc_str = gen_json_schema()
     else:
         raise RuntimeError(f"Unsupported out type {out_type}")
     print(doc_str)  # noqa: T201

deepmd/main.py

@@ -500,7 +500,7 @@ def main_parser() -> argparse.ArgumentParser:
     parsers_doc.add_argument(
         "--out-type",
         default="rst",
-        choices=["rst", "json"],
+        choices=["rst", "json", "json_schema"],
         type=str,
         help="The output type",
     )

deepmd/pt/model/descriptor/repformers.py

@@ -410,7 +410,7 @@ def forward(
         atype = extended_atype[:, :nloc]
         # nb x nloc x nnei
         exclude_mask = self.emask(nlist, extended_atype)
-        nlist = nlist * exclude_mask
+        nlist = torch.where(exclude_mask != 0, nlist, -1)
         # nb x nloc x nnei x 4, nb x nloc x nnei x 3, nb x nloc x nnei x 1
         dmatrix, diff, sw = prod_env_mat(
             extended_coord,

deepmd/pt/model/descriptor/se_atten.py

@@ -478,11 +478,12 @@ def forward(
             self.rcut_smth,
             protection=self.env_protection,
         )
+        # nb x nloc x nnei
+        exclude_mask = self.emask(nlist, extended_atype)
+        nlist = torch.where(exclude_mask != 0, nlist, -1)
         nlist_mask = nlist != -1
         nlist = torch.where(nlist == -1, 0, nlist)
         sw = torch.squeeze(sw, -1)
-        # beyond the cutoff sw should be 0.0
-        sw = sw.masked_fill(~nlist_mask, 0.0)
         # nf x nloc x nt -> nf x nloc x nnei x nt
         atype_tebd = extended_atype_embd[:, :nloc, :]
         atype_tebd_nnei = atype_tebd.unsqueeze(2).expand(-1, -1, self.nnei, -1)
@@ -495,8 +496,10 @@ def forward(
         atype_tebd_nlist = torch.gather(atype_tebd_ext, dim=1, index=index)
         # nb x nloc x nnei x nt
         atype_tebd_nlist = atype_tebd_nlist.view(nb, nloc, nnei, nt)
+        # beyond the cutoff sw should be 0.0
+        sw = sw.masked_fill(~nlist_mask, 0.0)
         # (nb x nloc) x nnei
-        exclude_mask = self.emask(nlist, extended_atype).view(nb * nloc, nnei)
+        exclude_mask = exclude_mask.view(nb * nloc, nnei)
         if self.old_impl:
             assert self.filter_layers_old is not None
             dmatrix = dmatrix.view(