deepmodeling · wanghan-iapcm · Jan 24, 2024 · Jan 23, 2024
diff --git a/doc/_static/css/custom.css b/doc/_static/css/custom.css
@@ -7,8 +7,14 @@ pre{
 .wy-side-nav-search .wy-dropdown > a img.logo, .wy-side-nav-search > a img.logo {
     width: 275px;
 }
+img.platform-icon {
+	height: 2ex;
+}
 @media (prefers-color-scheme: dark) {
 	.wy-side-nav-search .wy-dropdown > a img.logo, .wy-side-nav-search > a img.logo {
 		content: url("../logo-dark.svg");
 	}
+	img.platform-icon {
+		filter: invert(1);
+	}
 }
diff --git a/doc/_static/pytorch.svg b/doc/_static/pytorch.svg
diff --git a/doc/_static/tensorflow.svg b/doc/_static/tensorflow.svg
diff --git a/doc/conf.py b/doc/conf.py
@@ -275,6 +275,11 @@ def setup(app):
 .. |PRECISION| replace:: {list_to_doc(PRECISION_DICT.keys())}
 """
 
+myst_substitutions = {
+    "tensorflow_icon": """![TensorFlow](/_static/tensorflow.svg){class=platform-icon}""",
+    "pytorch_icon": """![PyTorch](/_static/pytorch.svg){class=platform-icon}""",
+}
+
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
@@ -298,6 +303,8 @@ def setup(app):
 myst_enable_extensions = [
     "dollarmath",
     "colon_fence",
+    "substitution",
+    "attrs_inline",
 ]
 myst_fence_as_directive = ("math",)
 # fix emoji issue in pdf

diff --git a/doc/freeze/compress.md b/doc/freeze/compress.md
@@ -1,4 +1,8 @@
-# Compress a model
+# Compress a model {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 ## Theory
 

diff --git a/doc/model/dplr.md b/doc/model/dplr.md
@@ -1,4 +1,8 @@
-# Deep potential long-range (DPLR)
+# Deep potential long-range (DPLR) {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Notice: **The interfaces of DPLR are not stable and subject to change**
 

diff --git a/doc/model/dprc.md b/doc/model/dprc.md
@@ -1,4 +1,8 @@
-# Deep Potential - Range Correction (DPRc)
+# Deep Potential - Range Correction (DPRc) {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Deep Potential - Range Correction (DPRc) is designed to combine with QM/MM method, and corrects energies from a low-level QM/MM method to a high-level QM/MM method:
 

diff --git a/doc/model/linear.md b/doc/model/linear.md
@@ -1,4 +1,8 @@
-## Linear model
+## Linear model {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 One can linearly combine existing models with arbitrary coefficients:
 

diff --git a/doc/model/pairtab.md b/doc/model/pairtab.md
@@ -1,4 +1,8 @@
-# Interpolation or combination with a pairwise potential
+# Interpolation or combination with a pairwise potential {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 ## Theory
 In applications like the radiation damage simulation, the interatomic distance may become too close, so that the DFT calculations fail.

diff --git a/doc/model/train-energy-spin.md b/doc/model/train-energy-spin.md
@@ -1,4 +1,8 @@
-# Fit spin energy
+# Fit spin energy {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 In this section, we will take `$deepmd_source_dir/examples/NiO/se_e2_a/input.json` as an example of the input file.
 

diff --git a/doc/model/train-energy.md b/doc/model/train-energy.md
@@ -1,4 +1,8 @@
-# Fit energy
+# Fit energy {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 In this section, we will take `$deepmd_source_dir/examples/water/se_e2_a/input.json` as an example of the input file.
 

diff --git a/doc/model/train-fitting-dos.md b/doc/model/train-fitting-dos.md
@@ -1,4 +1,8 @@
-# Fit electronic density of states (DOS)
+# Fit electronic density of states (DOS) {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Here we present an API to DeepDOS model, which can be used to fit electronic density of state (DOS) (which is a vector).
 

diff --git a/doc/model/train-fitting-tensor.md b/doc/model/train-fitting-tensor.md
@@ -1,4 +1,8 @@
-# Fit `tensor` like `Dipole` and `Polarizability`
+# Fit `tensor` like `Dipole` and `Polarizability` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_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
 

diff --git a/doc/model/train-hybrid.md b/doc/model/train-hybrid.md
@@ -1,4 +1,8 @@
-# Descriptor `"hybrid"`
+# Descriptor `"hybrid"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_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)$.
 

diff --git a/doc/model/train-se-a-mask.md b/doc/model/train-se-a-mask.md
@@ -1,4 +1,8 @@
-# Descriptor `"se_a_mask"`
+# Descriptor `"se_a_mask"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 
 Descriptor `se_a_mask` is a concise implementation of the descriptor `se_e2_a`,

diff --git a/doc/model/train-se-atten.md b/doc/model/train-se-atten.md
@@ -1,4 +1,8 @@
-# Descriptor `"se_atten"`
+# Descriptor `"se_atten"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 ## DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation
 

diff --git a/doc/model/train-se-e2-a-tebd.md b/doc/model/train-se-e2-a-tebd.md
@@ -1,4 +1,8 @@
-# Type embedding approach
+# Type embedding approach {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 We generate specific a type embedding vector for each atom type so that we can share one descriptor embedding net and one fitting net in total, which decline training complexity largely.
 

diff --git a/doc/model/train-se-e2-a.md b/doc/model/train-se-e2-a.md
@@ -1,4 +1,8 @@
-# Descriptor `"se_e2_a"`
+# Descriptor `"se_e2_a"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_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).
 

diff --git a/doc/model/train-se-e2-r.md b/doc/model/train-se-e2-r.md
@@ -1,4 +1,8 @@
-# Descriptor `"se_e2_r"`
+# Descriptor `"se_e2_r"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_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.
 

diff --git a/doc/model/train-se-e3.md b/doc/model/train-se-e3.md
@@ -1,4 +1,8 @@
-# Descriptor `"se_e3"`
+# Descriptor `"se_e3"` {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 The notation of `se_e3` is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The embedding takes bond angles between a central atom and its two neighboring atoms as input (denoted by `e3`).
 

diff --git a/doc/nvnmd/nvnmd.md b/doc/nvnmd/nvnmd.md
@@ -1,4 +1,8 @@
-# Introduction
+# Introduction {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 NVNMD stands for non-von Neumann molecular dynamics.
 

diff --git a/doc/train/finetuning.md b/doc/train/finetuning.md
@@ -1,4 +1,8 @@
-# Finetune the pretrained model
+# Finetune the pretrained model {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Pretraining-and-finetuning is a widely used approach in other fields such as Computer Vision (CV) or Natural Language Processing (NLP)
 to vastly reduce the training cost, while it's not trivial in potential models.

diff --git a/doc/train/gpu-limitations.md b/doc/train/gpu-limitations.md
@@ -1,4 +1,5 @@
-# Known limitations of using GPUs
+# Known limitations of using GPUs {{ tensorflow_icon }}
+
 If you use DeePMD-kit in a GPU environment, the acceptable value range of some variables is additionally restricted compared to the CPU environment due to the software's GPU implementations:
 1. The number of atom types of a given system must be less than 128.
 2. The maximum distance between an atom and its neighbors must be less than 128. It can be controlled by setting the rcut value of training parameters.

diff --git a/doc/train/multi-task-training.md b/doc/train/multi-task-training.md
@@ -1,4 +1,8 @@
-# Multi-task training
+# Multi-task training {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 ## Theory
 

diff --git a/doc/train/parallel-training.md b/doc/train/parallel-training.md
@@ -1,4 +1,8 @@
-# Parallel training
+# Parallel training {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Currently, parallel training is enabled in a synchronized way with help of [Horovod](https://github.com/horovod/horovod).
 Depending on the number of training processes (according to MPI context) and the number of GPU cards available, DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode. Therefore, no additional options are specified in your JSON/YAML input file.

diff --git a/doc/train/tensorboard.md b/doc/train/tensorboard.md
@@ -1,4 +1,8 @@
-# TensorBoard Usage
+# TensorBoard Usage {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 TensorBoard provides the visualization and tooling needed for machine learning
 experimentation. Full instructions for TensorBoard can be found