Merge branch 'devel' into 3831

deepmodeling · Jun 5, 2024 · c982d0e · c982d0e
2 parents b39fcce + eb474d4
commit c982d0e
Show file tree

Hide file tree

Showing 23 changed files with 544 additions and 18 deletions.
diff --git a/.github/workflows/test_cuda.yml b/.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 \

diff --git a/deepmd/dpmodel/descriptor/dpa1.py b/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]

diff --git a/deepmd/dpmodel/descriptor/repformers.py b/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

diff --git a/deepmd/dpmodel/descriptor/se_e2_a.py b/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")

diff --git a/deepmd/dpmodel/descriptor/se_r.py b/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")

diff --git a/deepmd/dpmodel/descriptor/se_t.py b/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

diff --git a/deepmd/entrypoints/doc.py b/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
diff --git a/deepmd/main.py b/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",
     )

diff --git a/deepmd/pt/model/descriptor/repformers.py b/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,

diff --git a/deepmd/pt/model/descriptor/se_atten.py b/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(

diff --git a/deepmd/tf/descriptor/se_atten.py b/deepmd/tf/descriptor/se_atten.py
@@ -718,6 +718,12 @@ def _pass_filter(
                 tf.shape(inputs_i)[0],
                 self.nei_type_vec,  # extra input for atten
             )
+            #  (nframes * nloc * nnei, 1)
+            nei_exclude_mask = tf.slice(
+                tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]),
+                [0, 0],
+                [-1, 1],
+            )
             if self.smooth:
                 inputs_i = tf.where(
                     tf.cast(mask, tf.bool),
@@ -727,15 +733,18 @@ def _pass_filter(
                         tf.reshape(self.avg_looked_up, [-1, 1]), [1, self.ndescrpt]
                     ),
                 )
+                #  (nframes, nloc, nnei)
                 self.recovered_switch *= tf.reshape(
-                    tf.slice(
-                        tf.reshape(tf.cast(mask, self.filter_precision), [-1, 4]),
-                        [0, 0],
-                        [-1, 1],
-                    ),
+                    nei_exclude_mask,
                     [-1, natoms[0], self.sel_all_a[0]],
                 )
             else:
+                #  (nframes * nloc, 1,  nnei)
+                self.nmask *= tf.reshape(
+                    nei_exclude_mask,
+                    [-1, 1, self.sel_all_a[0]],
+                )
+                self.negative_mask = -(2 << 32) * (1.0 - self.nmask)
                 inputs_i *= mask
         if nvnmd_cfg.enable and nvnmd_cfg.quantize_descriptor:
             inputs_i = descrpt2r4(inputs_i, atype)

diff --git a/deepmd/utils/argcheck.py b/deepmd/utils/argcheck.py
@@ -14,7 +14,13 @@
     Variant,
     dargs,
 )
+from dargs.json_schema import (
+    generate_json_schema,
+)
 
+from deepmd import (
+    __version__,
+)
 from deepmd.common import (
     VALID_ACTIVATION,
     VALID_PRECISION,
@@ -2527,6 +2533,18 @@ def gen_args(**kwargs) -> List[Argument]:
     ]
 
 
+def gen_json_schema() -> str:
+    """Generate JSON schema.
+
+    Returns
+    -------
+    str
+        JSON schema.
+    """
+    arg = Argument("DeePMD-kit", dict, gen_args(), doc=f"DeePMD-kit {__version__}")
+    return json.dumps(generate_json_schema(arg))
+
+
 def normalize(data):
     base = Argument("base", dict, gen_args())
     data = base.normalize_value(data, trim_pattern="_*")

diff --git a/doc/train/train-input.rst b/doc/train/train-input.rst
@@ -1,8 +1,42 @@
 Training Parameters
 ======================================
 .. note::
-   One can load, modify, and export the input file by using our effective web-based tool `DP-GUI <https://deepmodeling.com/dpgui/input/deepmd-kit-2.0>`_ online or hosted using the :ref:`command line interface <cli>` :code:`dp gui`. All training parameters below can be set in DP-GUI. By clicking "SAVE JSON", one can download the input file for furthur training.
+   One can load, modify, and export the input file by using our effective web-based tool `DP-GUI <https://deepmodeling.com/dpgui/input/deepmd-kit-2.0>`_ online or hosted using the :ref:`command line interface <cli>` :code:`dp gui`. All training parameters below can be set in DP-GUI. By clicking "SAVE JSON", one can download the input file for further training.
+
+.. note::
+   One can benefit from IntelliSense and validation when
+   :ref:`writing JSON files using Visual Studio Code <json_vscode>`.
+   See :ref:`here <json_vscode>` to learn how to configure.
 
 .. dargs::
-   :module: deepmd.tf.utils.argcheck
+   :module: deepmd.utils.argcheck
    :func: gen_args
+
+.. _json_vscode:
+
+Writing JSON files using Visual Studio Code
+-------------------------------------------
+
+When writing JSON files using `Visual Studio Code <https://code.visualstudio.com/>`_, one can benefit from IntelliSense and
+validation by adding a `JSON schema <https://json-schema.org/>`_.
+To do so, in a VS Code workspace, one can generate a JSON schema file for the input file by running the following command:
+
+.. code-block:: bash
+
+   dp doc-train-input --out-type json_schema > deepmd.json
+
+Then one can `map the schema <https://code.visualstudio.com/docs/languages/json#_mapping-to-a-schema-in-the-workspace>`_
+by updating the workspace settings in the `.vscode/settings.json` file as follows:
+
+.. code-block:: json
+
+   {
+      "json.schemas": [
+         {
+               "fileMatch": [
+                  "/**/*.json"
+               ],
+               "url": "./deepmd.json"
+         }
+      ]
+   }
diff --git a/pyproject.toml b/pyproject.toml
@@ -39,7 +39,7 @@ dependencies = [
     'numpy',
     'scipy',
     'pyyaml',
-    'dargs >= 0.4.1',
+    'dargs >= 0.4.6',
     'typing_extensions; python_version < "3.8"',
     'importlib_metadata>=1.4; python_version < "3.8"',
     'h5py',

diff --git a/source/tests/common/test_doc_train_input.py b/source/tests/common/test_doc_train_input.py
@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+import io
+import json
+import unittest
+from contextlib import (
+    redirect_stdout,
+)
+
+from deepmd.entrypoints.doc import (
+    doc_train_input,
+)
+
+
+class TestDocTrainInput(unittest.TestCase):
+    def test_rst(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="rst")
+        self.assertNotEqual(f.getvalue(), "")
+
+    def test_json(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="json")
+        # validate json
+        json.loads(f.getvalue())
+
+    def test_json_schema(self):
+        f = io.StringIO()
+        with redirect_stdout(f):
+            doc_train_input(out_type="json_schema")
+        # validate json
+        json.loads(f.getvalue())