Doc: Update PT tensor fitting

deepmd/utils/argcheck.py

@@ -1000,7 +1000,7 @@
     ]
 
 
-@fitting_args_plugin.register("polar", doc=doc_only_tf_supported)
+@fitting_args_plugin.register("polar")
 def fitting_polar():
     doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built."
     doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.'
@@ -1011,6 +1011,7 @@
     doc_fit_diag = "Fit the diagonal part of the rotational invariant polarizability matrix, which will be converted to normal polarizability matrix by contracting with the rotation matrix."
     doc_sel_type = "The atom types for which the atomic polarizability will be provided. If not set, all types will be selected."
     doc_seed = "Random seed for parameter initialization of the fitting net"
+    doc_exclude_types = "Atomic contributions of the excluded atom types are set zero."
 
     # YWolfeee: user can decide whether to use shift diag
     doc_shift_diag = "Whether to shift the diagonal of polar, which is beneficial to training. Default is true."
@@ -1044,7 +1045,14 @@
             [List[int], int, None],
             optional=True,
             alias=["pol_type"],
-            doc=doc_sel_type,
+            doc=doc_sel_type + doc_only_tf_supported,
+        ),
+        Argument(
+            "exclude_types",
+            [List[int], None],
+            optional=True,
+            default=[],
+            doc=doc_exclude_types + doc_only_pt_supported,
         ),
         Argument("seed", [int, None], optional=True, doc=doc_seed),
     ]
@@ -1054,14 +1062,15 @@
 #    return fitting_polar()
 
 
-@fitting_args_plugin.register("dipole", doc=doc_only_tf_supported)
+@fitting_args_plugin.register("dipole")
 def fitting_dipole():
     doc_neuron = "The number of neurons in each hidden layers of the fitting net. When two hidden layers are of the same size, a skip connection is built."
     doc_activation_function = f'The activation function in the fitting net. Supported activation functions are {list_to_doc(ACTIVATION_FN_DICT.keys())} Note that "gelu" denotes the custom operator version, and "gelu_tf" denotes the TF standard version. If you set "None" or "none" here, no activation function will be used.'
     doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
     doc_precision = f"The precision of the fitting net parameters, supported options are {list_to_doc(PRECISION_DICT.keys())} Default follows the interface precision."
     doc_sel_type = "The atom types for which the atomic dipole will be provided. If not set, all types will be selected."
     doc_seed = "Random seed for parameter initialization of the fitting net"
+    doc_exclude_types = "Atomic contributions of the excluded atom types are set zero."
     return [
         Argument(
             "neuron",
@@ -1085,7 +1094,14 @@
             [List[int], int, None],
             optional=True,
             alias=["dipole_type"],
-            doc=doc_sel_type,
+            doc=doc_sel_type + doc_only_tf_supported,
+        ),
+        Argument(
+            "exclude_types",
+            [List[int], None],
+            optional=True,
+            default=[],
+            doc=doc_exclude_types + doc_only_pt_supported,
         ),
         Argument("seed", [int, None], optional=True, doc=doc_seed),
     ]

doc/backend.md

@@ -23,18 +23,18 @@ DeePMD-kit does not use the TensorFlow v2 API but uses the TensorFlow v1 API (`t
 [PyTorch](https://pytorch.org/) 2.0 or above is required.
 While `.pth` and `.pt` are the same in the PyTorch package, they have different meanings in the DeePMD-kit to distinguish the model and the checkpoint.
 
-### DPModel {{ dpmodel_icon }}
+### DP {{ dpmodel_icon }}
 
 :::{note}
 This backend is only for development and should not take into production.
 :::
 
 - Model filename extension: `.dp`
 
-DPModel is a reference backend for development, which uses pure [NumPy](https://numpy.org/) to implement models without using any heavy deep-learning frameworks.
+DP is a reference backend for development, which uses pure [NumPy](https://numpy.org/) to implement models without using any heavy deep-learning frameworks.
 Due to the limitation of NumPy, it doesn't support gradient calculation and thus cannot be used for training.
 As a reference backend, it is not aimed at the best performance, but only the correct results.
-The DPModel backend uses [HDF5](https://docs.h5py.org/) to store model serialization data, which is backend-independent.
+The DP backend uses [HDF5](https://docs.h5py.org/) to store model serialization data, which is backend-independent.
 Only Python inference interface can load this format.
 
 ## Switch the backend

doc/conf.py

@@ -186,7 +186,7 @@ def setup(app):
 myst_substitutions = {
     "tensorflow_icon": """![TensorFlow](/_static/tensorflow.svg){class=platform-icon}""",
     "pytorch_icon": """![PyTorch](/_static/pytorch.svg){class=platform-icon}""",
-    "dpmodel_icon": """![DPModel](/_static/logo_icon.svg){class=platform-icon}""",
+    "dpmodel_icon": """![DP](/_static/logo_icon.svg){class=platform-icon}""",
 }
 
 # -- Options for HTML output -------------------------------------------------

doc/model/train-fitting-tensor.md

@@ -1,16 +1,33 @@
-# Fit `tensor` like `Dipole` and `Polarizability` {{ tensorflow_icon }}
+# Fit `tensor` like `Dipole` and `Polarizability` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 
 :::{note}
-**Supported backends**: TensorFlow {{ tensorflow_icon }}
+**Supported backends**: TensorFlow {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 :::
 
 Unlike `energy`, which is a scalar, one may want to fit some high dimensional physical quantity, like `dipole` (vector) and `polarizability` (matrix, shorted as `polar`). Deep Potential has provided different APIs to do this. In this example, we will show you how to train a model to fit a water system. A complete training input script of the examples can be found in
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
 ```bash
 $deepmd_source_dir/examples/water_tensor/dipole/dipole_input.json
 $deepmd_source_dir/examples/water_tensor/polar/polar_input.json
 ```
 
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
+```bash
+$deepmd_source_dir/examples/water_tensor/dipole/dipole_input_torch.json
+$deepmd_source_dir/examples/water_tensor/polar/polar_input_torch.json
+```
+
+:::
+
+::::
+
 The training and validation data are also provided our examples. But note that **the data provided along with the examples are of limited amount, and should not be used to train a production model.**
 
 Similar to the `input.json` used in `ener` mode, training JSON is also divided into {ref}`model <model>`, {ref}`learning_rate <learning_rate>`, {ref}`loss <loss>` and {ref}`training <training>`. Most keywords remain the same as `ener` mode, and their meaning can be found [here](train-se-e2-a.md). To fit a tensor, one needs to modify {ref}`model/fitting_net <model/fitting_net>` and {ref}`loss <loss>`.
@@ -53,6 +70,10 @@ The tensorial models can be used to calculate IR spectrum and Raman spectrum.[^1
 
 The {ref}`fitting_net <model/fitting_net>` section tells DP which fitting net to use.
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
 The JSON of `dipole` type should be provided like
 
 ```json
@@ -81,9 +102,48 @@ The JSON of `polar` type should be provided like
 -   `sel_type` is a list specifying which type of atoms have the quantity you want to fit. For example, in the water system, `sel_type` is `[0]` since `0` represents atom `O`. If left unset, all types of atoms will be fitted.
 -   The rest arguments have the same meaning as they do in `ener` mode.
 
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
+The JSON of `dipole` type should be provided like
+```json
+	"atom_exclude_types": [
+      1
+    ],
+	"fitting_net" : {
+		"type": "dipole",
+		"neuron": [100,100,100],
+		"resnet_dt": true,
+		"seed": 1,
+	},
+```
+
+The JSON of `polar` type should be provided like
+
+```json
+	"atom_exclude_types": [
+      1
+    ],
+	"fitting_net" : {
+	   	"type": "polar",
+		"neuron": [100,100,100],
+		"resnet_dt": true,
+		"seed": 1,
+	},
+```
+-   `type` specifies which type of fitting net should be used. It should be either `dipole` or `polar`. Note that `global_polar` mode in version 1.x is already **deprecated** and is merged into `polar`. To specify whether a system is global or atomic, please see [here](train-se-e2-a.md).
+-   `atom_exclude_types` is a list specifying the which type of atoms have the quantity you want to set to zero. For example, in the water system, `atom_exclude_types` is `[1]` since `1` represents atom `H`.
+-   The rest arguments have the same meaning as they do in `ener` mode.
+:::
+
+::::
+
+
+
 ## Loss
 
-DP supports a combinational training of the global system (only a global `tensor` label, i.e. dipole or polar, is provided in a frame) and atomic system (labels for **each** atom included in `sel_type` are provided). In a global system, each frame has just **one** `tensor` label. For example, when fitting `polar`, each frame will just provide a `1 x 9` vector which gives the elements of the polarizability tensor of that frame in order XX, XY, XZ, YX, YY, YZ, XZ, ZY, ZZ. By contrast, in an atomic system, each atom in `sel_type` has a `tensor` label. For example, when fitting a dipole, each frame will provide a `#sel_atom x 3` matrices, where `#sel_atom` is the number of atoms whose type are in `sel_type`.
+DP supports a combinational training of the global system (only a global `tensor` label, i.e. dipole or polar, is provided in a frame) and atomic system (labels for **each** atom included in `sel_type`/ not included in `atom_exclude_types` are provided). In a global system, each frame has just **one** `tensor` label. For example, when fitting `polar`, each frame will just provide a `1 x 9` vector which gives the elements of the polarizability tensor of that frame in order XX, XY, XZ, YX, YY, YZ, XZ, ZY, ZZ. By contrast, in an atomic system, each atom in `sel_type` has a `tensor` label. For example, when fitting a dipole, each frame will provide a `#sel_atom x 3` matrices, where `#sel_atom` is the number of atoms whose type are in `sel_type`.
 
 The {ref}`loss <loss>` section tells DP the weight of these two kinds of loss, i.e.
 
@@ -118,9 +178,24 @@ In this case, please check the file name of the label.
 
 The training command is the same as `ener` mode, i.e.
 
+::::{tab-set}
+
+:::{tab-item} TensorFlow {{ tensorflow_icon }}
+
 ```bash
 dp train input.json
 ```
+:::
+
+:::{tab-item} PyTorch {{ pytorch_icon }}
+
+```bash
+dp --pt train input.json
+```
+:::
+
+::::
+
 
 The detailed loss can be found in `lcurve.out`:
 

doc/model/train-hybrid.md

@@ -1,7 +1,7 @@
 # Descriptor `"hybrid"`  {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 
 :::{note}
-**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DPModel {{ dpmodel_icon }}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
 :::
 
 This descriptor hybridizes multiple descriptors to form a new descriptor. For example, we have a list of descriptors denoted by $\mathcal D_1$, $\mathcal D_2$, ..., $\mathcal D_N$, the hybrid descriptor this the concatenation of the list, i.e. $\mathcal D = (\mathcal D_1, \mathcal D_2, \cdots, \mathcal D_N)$.

doc/model/train-se-e2-a.md

@@ -1,7 +1,7 @@
 # Descriptor `"se_e2_a"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 
 :::{note}
-**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DPModel {{ dpmodel_icon }}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
 :::
 
 The notation of `se_e2_a` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The `e2` stands for the embedding with two-atoms information. This descriptor was described in detail in [the DeepPot-SE paper](https://arxiv.org/abs/1805.09003).

doc/model/train-se-e2-r.md

@@ -1,7 +1,7 @@
 # Descriptor `"se_e2_r"` {{ tensorflow_icon }} {{ pytorch_icon }} {{ dpmodel_icon }}
 
 :::{note}
-**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DPModel {{ dpmodel_icon }}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}, DP {{ dpmodel_icon }}
 :::
 
 The notation of `se_e2_r` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from the radial information of atomic configurations. The `e2` stands for the embedding with two-atom information.