Merge branch 'devel' into devel

cherryWangY · web-flow · commit 7dd7f6aea325 · 2024-06-18T10:14:03.000+08:00
diff --git a/.github/workflows/build_wheel.yml b/.github/workflows/build_wheel.yml
@@ -75,7 +75,7 @@ jobs:
           # https://github.com/pypa/setuptools_scm/issues/480
           fetch-depth: 0
       - name: Install uv
-        run: curl -LsSf https://astral.sh/uv/0.2.10/install.sh | sh
+        run: curl -LsSf https://astral.sh/uv/install.sh | sh
         if: runner.os != 'Linux'
       - uses: docker/setup-qemu-action@v3
         name: Setup QEMU
@@ -161,7 +161,7 @@ jobs:
           images: ghcr.io/deepmodeling/deepmd-kit
 
       - name: Build and push Docker image
-        uses: docker/build-push-action@v5
+        uses: docker/build-push-action@v6
         with:
           context: source/install/docker
           push: ${{ github.repository_owner == 'deepmodeling' && github.event_name == 'push' && github.actor != 'dependabot[bot]' }}
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
@@ -29,7 +29,7 @@ repos:
         exclude: ^source/3rdparty
   - repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.4.8
+    rev: v0.4.9
     hooks:
       - id: ruff
         args: ["--fix"]
diff --git a/deepmd/pt/train/training.py b/deepmd/pt/train/training.py
@@ -31,7 +31,6 @@
     TensorLoss,
 )
 from deepmd.pt.model.model import (
-    DOSModel,
     get_model,
     get_zbl_model,
 )
@@ -601,15 +600,13 @@ def single_model_finetune(
                         _finetune_rule_single,
                         _sample_func,
                     ):
-                        # need fix for DOSModel
-                        if not isinstance(_model, DOSModel):
-                            _model = _model_change_out_bias(
-                                _model,
-                                _sample_func,
-                                _bias_adjust_mode="change-by-statistic"
-                                if not _finetune_rule_single.get_random_fitting()
-                                else "set-by-statistic",
-                            )
+                        _model = _model_change_out_bias(
+                            _model,
+                            _sample_func,
+                            _bias_adjust_mode="change-by-statistic"
+                            if not _finetune_rule_single.get_random_fitting()
+                            else "set-by-statistic",
+                        )
                         return _model
 
                     if not self.multi_task:
diff --git a/deepmd/pt/utils/stat.py b/deepmd/pt/utils/stat.py
@@ -399,6 +399,10 @@ def compute_output_stats_global(
     model_pred: Optional[Dict[str, np.ndarray]] = None,
 ):
     """This function only handle stat computation from reduced global labels."""
+    # return directly if model predict is empty for global
+    if model_pred == {}:
+        return {}, {}
+
     # get label dict from sample; for each key, only picking the system with global labels.
     outputs = {
         kk: [
diff --git a/deepmd/utils/econf_embd.py b/deepmd/utils/econf_embd.py
@@ -1,13 +1,8 @@
 # SPDX-License-Identifier: LGPL-3.0-or-later
 import numpy as np
-
-try:
-    import dpdata
-    from mendeleev import (
-        element,
-    )
-except ImportError:
-    pass
+from mendeleev import (
+    element,
+)
 
 ###
 # made by command
@@ -173,7 +168,7 @@
 maxl = maxn
 maxm = 2 * maxl + 1
 
-type_map = dpdata.periodic_table.ELEMENTS
+type_map = list(electronic_configuration_embedding.keys())
 ECONF_DIM = electronic_configuration_embedding[type_map[0]].shape[0]
 
 
diff --git a/doc/freeze/freeze.md b/doc/freeze/freeze.md
@@ -1,6 +1,7 @@
 # Freeze a model
 
-The trained neural network is extracted from a checkpoint and dumped into a protobuf(.pb) file. This process is called "freezing" a model. The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc). To freeze a model, typically one does
+The trained neural network is extracted from a checkpoint and dumped into a model file. This process is called "freezing" a model.
+To freeze a model, typically one does
 
 ::::{tab-set}
 
@@ -11,6 +12,7 @@ $ dp freeze -o model.pb
 ```
 
 in the folder where the model is trained. The output model is called `model.pb`.
+The idea and part of our code are from [Morgan](https://blog.metaflow.fr/tensorflow-how-to-freeze-a-model-and-serve-it-with-a-python-api-d4f3596b3adc).
 
 :::
 
diff --git a/doc/model/sel.md b/doc/model/sel.md
@@ -6,10 +6,26 @@ All descriptors require to set `sel`, which means the expected maximum number of
 
 To determine a proper `sel`, one can calculate the neighbor stat of the training data before training:
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
 ```sh
-dp neighbor-stat -s data -r 6.0 -t O H
+dp --tf neighbor-stat -s data -r 6.0 -t O H
 ```
 
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
+```sh
+dp --pt neighbor-stat -s data -r 6.0 -t O H
+```
+
+:::
+
+::::
+
 where `data` is the directory of data, `6.0` is the cutoff radius, and `O` and `H` is the type map. The program will give the `max_nbor_size`. For example, `max_nbor_size` of the water example is `[38, 72]`, meaning an atom may have 38 O neighbors and 72 H neighbors in the training data.
 
 The `sel` should be set to a higher value than that of the training data, considering there may be some extreme geometries during MD simulations. As a result, we set `sel` to `[46, 92]` in the water example.
diff --git a/doc/model/train-fitting-dos.md b/doc/model/train-fitting-dos.md
@@ -1,7 +1,7 @@
-# Fit electronic density of states (DOS) {{ tensorflow_icon }}
+# Fit electronic density of states (DOS) {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 
 :::{note}
-**Supported backends**: TensorFlow {{ tensorflow_icon }}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
 :::
 
 Here we present an API to DeepDOS model, which can be used to fit electronic density of state (DOS) (which is a vector).
@@ -82,10 +82,26 @@ To prepare the data, we recommend shifting the DOS data by the Fermi level.
 
 The training command is the same as `ener` mode, i.e.
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
+```bash
+dp --tf train input.json
+```
+
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
 ```bash
-dp train input.json
+dp --pt train input.json
 ```
 
+:::
+
+::::
+
 The detailed loss can be found in `lcurve.out`:
 
 ```
@@ -117,14 +133,33 @@ The detailed loss can be found in `lcurve.out`:
 
 In this earlier version, we can use `dp test` to infer the electronic density of state for given frames.
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
+```bash
+
+dp --tf freeze -o frozen_model.pb
+
+dp --tf test -m frozen_model.pb -s ../data/111/$k -d ${output_prefix} -a -n 100
+```
+
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
 ```bash
 
-$DP freeze -o frozen_model.pb
+dp --pt freeze -o frozen_model.pth
 
-$DP test -m frozen_model.pb -s ../data/111/$k -d ${output_prefix} -a -n 100
+dp --pt test -m frozen_model.pth -s ../data/111/$k -d ${output_prefix} -a -n 100
 ```
 
-if `dp test -d ${output_prefix} -a` is specified, the predicted DOS and atomic DOS for each frame is output in the working directory
+:::
+
+::::
+
+if `dp test -d ${output_prefix} -a` is specified, the predicted DOS and atomic DOS for each frame are output in the working directory
 
 ```
 ${output_prefix}.ados.out.0   ${output_prefix}.ados.out.1  ${output_prefix}.ados.out.2  ${output_prefix}.ados.out.3
diff --git a/doc/test/model-deviation.md b/doc/test/model-deviation.md
@@ -59,7 +59,7 @@ One can also use a subcommand to calculate the deviation of predicted forces or
 dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out
 ```
 
-where `-m` specifies graph files to be calculated, `-s` gives the data to be evaluated, `-o` the file to which model deviation results is dumped. Here is more information on this sub-command:
+where `-m` specifies model files to be calculated, `-s` gives the data to be evaluated, `-o` the file to which model deviation results are dumped. Here is more information on this sub-command:
 
 ```bash
 usage: dp model-devi [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}]
diff --git a/doc/third-party/gromacs.md b/doc/third-party/gromacs.md
@@ -105,7 +105,7 @@ Then, in your working directories, we have to write `input.json` file:
 
 Here is an explanation for these settings:
 
-- `graph_file` : The graph file (with suffix .pb) generated by `dp freeze` command
+- `graph_file` : The [model file](../backend.md) generated by `dp freeze` command
 - `type_file` : File to specify DP atom types (in space-separated format). Here, `type.raw` looks like
 
 ```
diff --git a/doc/third-party/lammps-command.md b/doc/third-party/lammps-command.md
@@ -70,6 +70,7 @@ pair_style deepmd models ... keyword value ...
 pair_style deepmd graph.pb
 pair_style deepmd graph.pb fparam 1.2
 pair_style deepmd graph_0.pb graph_1.pb graph_2.pb out_file md.out out_freq 10 atomic relative 1.0
+pair_style deepmd graph_0.pb graph_1.pth out_file md.out out_freq 100
 pair_coeff * * O H
 
 pair_style deepmd cp.pb fparam_from_compute TEMP
diff --git a/doc/train/training-advanced.md b/doc/train/training-advanced.md
@@ -170,9 +170,9 @@ One can set other environmental variables:
 | DP_AUTO_PARALLELIZATION | 0, 1          | 0             | Enable auto parallelization for CPU operators.                                                                      |
 | DP_JIT                  | 0, 1          | 0             | Enable JIT. Note that this option may either improve or decrease the performance. Requires TensorFlow supports JIT. |
 
-## Adjust `sel` of a frozen model
+## Adjust `sel` of a frozen model {{ tensorflow_icon }}
 
-One can use `--init-frz-model` features to adjust (increase or decrease) [`sel`](../model/sel.md) of a existing model. Firstly, one needs to adjust [`sel`](./train-input.rst) in `input.json`. For example, adjust from `[46, 92]` to `[23, 46]`.
+One can use `--init-frz-model` features to adjust (increase or decrease) [`sel`](../model/sel.md) of an existing model. Firstly, one needs to adjust [`sel`](./train-input.rst) in `input.json`. For example, adjust from `[46, 92]` to `[23, 46]`.
 
 ```json
 "model": {
diff --git a/doc/train/training.md b/doc/train/training.md
@@ -8,10 +8,26 @@ $ cd $deepmd_source_dir/examples/water/se_e2_a/
 
 After switching to that directory, the training can be invoked by
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
 ```bash
-$ dp train input.json
+$ dp --tf train input.json
 ```
 
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
+```bash
+$ dp --pt train input.json
+```
+
+:::
+
+::::
+
 where `input.json` is the name of the input script.
 
 By default, the verbosity level of the DeePMD-kit is `INFO`, one may see a lot of important information on the code and environment showing on the screen. Among them two pieces of information regarding data systems are worth special notice.
diff --git a/doc/troubleshooting/howtoset_num_nodes.md b/doc/troubleshooting/howtoset_num_nodes.md
@@ -72,13 +72,32 @@ There is no one general parallel configuration that works for all situations, so
 Here are some empirical examples.
 If you wish to use 3 cores of 2 CPUs on one node, you may set the environmental variables and run DeePMD-kit as follows:
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
+```bash
+export OMP_NUM_THREADS=3
+export DP_INTRA_OP_PARALLELISM_THREADS=3
+export DP_INTER_OP_PARALLELISM_THREADS=2
+dp --tf train input.json
+```
+
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
 ```bash
 export OMP_NUM_THREADS=3
 export DP_INTRA_OP_PARALLELISM_THREADS=3
 export DP_INTER_OP_PARALLELISM_THREADS=2
-dp train input.json
+dp --pt train input.json
 ```
 
+:::
+
+::::
+
 For a node with 128 cores, it is recommended to start with the following variables:
 
 ```bash
diff --git a/pyproject.toml b/pyproject.toml
@@ -47,6 +47,7 @@ dependencies = [
     'wcmatch',
     'packaging',
     'ml_dtypes',
+    'mendeleev',
 ]
 requires-python = ">=3.8"
 keywords = ["deepmd"]
@@ -78,7 +79,6 @@ test = [
     "pytest-sugar",
     "pytest-split",
     "dpgui",
-    "mendeleev",
 ]
 docs = [
     "sphinx>=3.1.1",
diff --git a/source/tests/common/test_econf_embd.py b/source/tests/common/test_econf_embd.py
@@ -6,15 +6,7 @@
     make_econf_embedding,
 )
 
-try:
-    import mendeleev  # noqa: F401
 
-    has_mendeleev = True
-except ImportError:
-    has_mendeleev = False
-
-
-@unittest.skipIf(not has_mendeleev, "does not have mendeleev installed, skip the UTs.")
 class TestEConfEmbd(unittest.TestCase):
     def test_fe(self):
         res = make_econf_embedding(["Fe"], flatten=False)["Fe"]
diff --git a/source/tests/pt/test_finetune.py b/source/tests/pt/test_finetune.py
