diff --git a/README.md b/README.md
index 0d61de70b6..cdec110a92 100644
--- a/README.md
+++ b/README.md
@@ -9,30 +9,7 @@
- [License and credits](#license-and-credits)
- [Deep Potential in a nutshell](#deep-potential-in-a-nutshell)
- [Download and install](#download-and-install)
- - [Easy installation methods](#easy-installation-methods)
- - [Offline packages](#offline-packages)
- - [With Docker](#with-docker)
- - [With conda](#with-conda)
- - [Install the python interaction](#install-the-python-interface)
- - [Install the Tensorflow's python interface](#install-the-tensorflows-python-interface)
- - [Install the DeePMD-kit's python interface](#install-the-deepmd-kits-python-interface)
- - [Install the C++ interface](#install-the-c-interface)
- - [Install the Tensorflow's C++ interface](#install-the-tensorflows-c-interface)
- - [Install the DeePMD-kit's C++ interface](#install-the-deepmd-kits-c-interface)
- - [Install LAMMPS's DeePMD-kit module](#install-lammpss-deepmd-kit-module)
- [Use DeePMD-kit](#use-deepmd-kit)
- - [Prepare data](#prepare-data)
- - [Train a model](#train-a-model)
- - [The DeePMD model](#the-deepmd-model)
- - [The DeepPot-SE model](#the-deeppot-se-model)
- - [Freeze a model](#freeze-a-model)
- - [Test a model](#test-a-model)
- - [Model inference](#model-inference)
- - [Run MD with Lammps](#run-md-with-lammps)
- - [Include deepmd in the pair style](#include-deepmd-in-the-pair-style)
- - [Long-range interaction](#long-range-interaction)
- - [Run path-integral MD with i-PI](#run-path-integral-md-with-i-pi)
- - [Use deep potential with ASE](#use-deep-potential-with-ase)
- [Troubleshooting](#troubleshooting)
# About DeePMD-kit
@@ -85,529 +62,27 @@ In addition to building up potential energy models, DeePMD-kit can also be used
# Download and install
-Please follow our [github](https://github.com/deepmodeling/deepmd-kit) webpage to see the latest released version and development version.
+Please follow our [github](https://github.com/deepmodeling/deepmd-kit) webpage to download the [latest released version](https://github.com/deepmodeling/deepmd-kit/tree/master) and [development version](https://github.com/deepmodeling/deepmd-kit/tree/devel).
-## Easy installation methods
-There various easy methods to install DeePMD-kit. Choose one that you prefer. If you want to build by yourself, jump to the next two sections.
+DeePMD-kit offers multiple installation methods. It is recommend using easily methods like [offline packages](doc/install.md#offline-packages), [conda](doc/install.md#with-conda) and [docker](doc/install.md#with-docker).
-After your easy installation, DeePMD-kit (`dp`) and LAMMPS (`lmp`) will be available to execute. You can try `dp -h` and `lmp -h` to see the help. `mpirun` is also available considering you may want to run LAMMPS in parallel.
+One may manually install DeePMD-kit by following the instuctions on [installing the python interface](doc/install.md#install-the-python-interface) and [installing the C++ interface](doc/install.md#install-the-c-interface). The C++ interface is necessary when using DeePMD-kit with LAMMPS and i-PI.
-### Offline packages
-Both CPU and GPU version offline packages are avaiable in [the Releases page](https://github.com/deepmodeling/deepmd-kit/releases).
-
-### With conda
-DeePMD-kit is avaiable with [conda](https://github.com/conda/conda). Install [Anaconda](https://www.anaconda.com/distribution/#download-section) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html) first.
-
-To install the CPU version:
-```bash
-conda install deepmd-kit=*=*cpu lammps-dp=*=*cpu -c deepmodeling
-```
-
-To install the GPU version containing [CUDA 10.1](https://docs.nvidia.com/deploy/cuda-compatibility/index.html#binary-compatibility__table-toolkit-driver):
-```bash
-conda install deepmd-kit=*=*gpu lammps-dp=*=*gpu -c deepmodeling
-```
-
-### With Docker
-A docker for installing the DeePMD-kit is available [here](https://github.com/orgs/deepmodeling/packages/container/package/deepmd-kit).
-
-To pull the CPU version:
-```bash
-docker pull ghcr.io/deepmodeling/deepmd-kit:1.2.2_cpu
-```
-
-To pull the GPU version:
-```bash
-docker pull ghcr.io/deepmodeling/deepmd-kit:1.2.2_cuda10.1_gpu
-```
-
-## Install the python interface
-### Install the Tensorflow's python interface
-First, check the python version on your machine
-```bash
-python --version
-```
-
-We follow the virtual environment approach to install the tensorflow's Python interface. The full instruction can be found on [the tensorflow's official website](https://www.tensorflow.org/install/pip). Now we assume that the Python interface will be installed to virtual environment directory `$tensorflow_venv`
-```bash
-virtualenv -p python3 $tensorflow_venv
-source $tensorflow_venv/bin/activate
-pip install --upgrade pip
-pip install --upgrade tensorflow==2.3.0
-```
-It is notice that everytime a new shell is started and one wants to use `DeePMD-kit`, the virtual environment should be activated by
-```bash
-source $tensorflow_venv/bin/activate
-```
-if one wants to skip out of the virtual environment, he/she can do
-```bash
-deactivate
-```
-If one has multiple python interpreters named like python3.x, it can be specified by, for example
-```bash
-virtualenv -p python3.7 $tensorflow_venv
-```
-If one does not need the GPU support of deepmd-kit and is concerned about package size, the CPU-only version of tensorflow should be installed by
-```bash
-pip install --upgrade tensorflow-cpu==2.3.0
-```
-To verify the installation, run
-```bash
-python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))"
-```
-One should remember to activate the virtual environment every time he/she uses deepmd-kit.
-
-### Install the DeePMD-kit's python interface
-
-Execute
-```bash
-pip install deepmd-kit
-```
-To test the installation, one may execute
-```bash
-dp -h
-```
-It will print the help information like
-```text
-usage: dp [-h] {train,freeze,test} ...
-
-DeePMD-kit: A deep learning package for many-body potential energy
-representation and molecular dynamics
-
-optional arguments:
- -h, --help show this help message and exit
-
-Valid subcommands:
- {train,freeze,test}
- train train a model
- freeze freeze the model
- test test the model
-```
-
-## Install the C++ interface
-
-If one does not need to use DeePMD-kit with Lammps or I-Pi, then the python interface installed in the previous section does everything and he/she can safely skip this section.
-
-### Install the Tensorflow's C++ interface
-
-Check the compiler version on your machine
-
-```
-gcc --version
-```
-
-The C++ interface of DeePMD-kit was tested with compiler gcc >= 4.8. It is noticed that the I-Pi support is only compiled with gcc >= 4.9.
-
-First the C++ interface of Tensorflow should be installed. It is noted that the version of Tensorflow should be in consistent with the python interface. We assume that you have followed our instruction and installed tensorflow python interface 1.14.0 with, then you may follow [the instruction for CPU](doc/install-tf.1.14.md) to install the corresponding C++ interface (CPU only). If one wants GPU supports, he/she should follow [the instruction for GPU](doc/install-tf.1.14-gpu.md) to install the C++ interface.
-
-### Install the DeePMD-kit's C++ interface
-
-Clone the DeePMD-kit source code
-```bash
-cd /some/workspace
-git clone --recursive https://github.com/deepmodeling/deepmd-kit.git deepmd-kit
-```
-
-For convenience, you may want to record the location of source to a variable, saying `deepmd_source_dir` by
-```bash
-cd deepmd-kit
-deepmd_source_dir=`pwd`
-```
-
-Now goto the source code directory of DeePMD-kit and make a build place.
-```bash
-cd $deepmd_source_dir/source
-mkdir build
-cd build
-```
-I assume you want to install DeePMD-kit into path `$deepmd_root`, then execute cmake
-```bash
-cmake -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
-```
-where the variable `tensorflow_root` stores the location where the tensorflow's C++ interface is installed. The DeePMD-kit will automatically detect if a CUDA tool-kit is available on your machine and build the GPU support accordingly. If you want to force the cmake to find CUDA tool-kit, you can speicify the key `USE_CUDA_TOOLKIT`,
-```bash
-cmake -DUSE_CUDA_TOOLKIT=true -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
-```
-and you may further asked to provide `CUDA_TOOLKIT_ROOT_DIR`. If the cmake has executed successfully, then
-```bash
-make
-make install
-```
-If everything works fine, you will have the following executable and libraries installed in `$deepmd_root/bin` and `$deepmd_root/lib`
-```bash
-$ ls $deepmd_root/bin
-dp_ipi
-$ ls $deepmd_root/lib
-libdeepmd_ipi.so libdeepmd_op.so libdeepmd.so
-```
-
-### Install LAMMPS's DeePMD-kit module
-DeePMD-kit provide module for running MD simulation with LAMMPS. Now make the DeePMD-kit module for LAMMPS.
-```bash
-cd $deepmd_source_dir/source/build
-make lammps
-```
-DeePMD-kit will generate a module called `USER-DEEPMD` in the `build` directory. Now download your favorite LAMMPS code, and uncompress it (I assume that you have downloaded the tar `lammps-stable.tar.gz`)
-```bash
-cd /some/workspace
-tar xf lammps-stable.tar.gz
-```
-The source code of LAMMPS is stored in directory, for example `lammps-31Mar17`. Now go into the LAMMPS code and copy the DeePMD-kit module like this
-```bash
-cd lammps-31Mar17/src/
-cp -r $deepmd_source_dir/source/build/USER-DEEPMD .
-```
-Now build LAMMPS
-```bash
-make yes-user-deepmd
-make mpi -j4
-```
-The option `-j4` means using 4 processes in parallel. You may want to use a different number according to your hardware.
-
-If everything works fine, you will end up with an executable `lmp_mpi`.
-
-The DeePMD-kit module can be removed from LAMMPS source code by
-```bash
-make no-user-deepmd
-```
# Use DeePMD-kit
-In this text, we will call the deep neural network that is used to represent the interatomic interactions (Deep Potential) the **model**. The typical procedure of using DeePMD-kit is
-
-1. Prepare data
-2. Train a model
-3. Freeze the model
-4. MD runs with the model (Native MD code or LAMMPS)
-
-## Prepare data
-One needs to provide the following information to train a model: the atom type, the simulation box, the atom coordinate, the atom force, system energy and virial. A snapshot of a system that contains these information is called a **frame**. We use the following convention of units:
-
-Property| Unit
---- | :---:
-Time | ps
-Length | Å
-Energy | eV
-Force | eV/Å
-Pressure| Bar
-
-The frames of the system are stored in two formats. A raw file is a plain text file with each information item written in one file and one frame written on one line. The default files that provide box, coordinate, force, energy and virial are `box.raw`, `coord.raw`, `force.raw`, `energy.raw` and `virial.raw`, respectively. *We recommend you use these file names*. Here is an example of force.raw:
-```bash
-$ cat force.raw
--0.724 2.039 -0.951 0.841 -0.464 0.363
- 6.737 1.554 -5.587 -2.803 0.062 2.222
--1.968 -0.163 1.020 -0.225 -0.789 0.343
-```
-This `force.raw` contains 3 frames with each frame having the forces of 2 atoms, thus it has 3 lines and 6 columns. Each line provides all the 3 force components of 2 atoms in 1 frame. The first three numbers are the 3 force components of the first atom, while the second three numbers are the 3 force components of the second atom. The coordinate file `coord.raw` is organized similarly. In `box.raw`, the 9 components of the box vectors should be provided on each line. In `virial.raw`, the 9 components of the virial tensor should be provided on each line. The number of lines of all raw files should be identical.
-
-We assume that the atom types do not change in all frames. It is provided by `type.raw`, which has one line with the types of atoms written one by one. The atom types should be integers. For example the `type.raw` of a system that has 2 atoms with 0 and 1:
-```bash
-$ cat type.raw
-0 1
-```
-
-The second format is the data sets of `numpy` binary data that are directly used by the training program. User can use the script `$deepmd_source_dir/data/raw/raw_to_set.sh` to convert the prepared raw files to data sets. For example, if we have a raw file that contains 6000 frames,
-```bash
-$ ls
-box.raw coord.raw energy.raw force.raw type.raw virial.raw
-$ $deepmd_source_dir/data/raw/raw_to_set.sh 2000
-nframe is 6000
-nline per set is 2000
-will make 3 sets
-making set 0 ...
-making set 1 ...
-making set 2 ...
-$ ls
-box.raw coord.raw energy.raw force.raw set.000 set.001 set.002 type.raw virial.raw
-```
-It generates three sets `set.000`, `set.001` and `set.002`, with each set contains 2000 frames. The last set (`set.002`) is used as testing set, while the rest sets (`set.000` and `set.001`) are used as training sets. One do not need to take care of the binary data files in each of the `set.*` directories. The path containing `set.*` and `type.raw` is called a *system*.
-
-## Train a model
-
-### Write the input script
-
-The method of training is explained in our [DeePMD][2] and [DeepPot-SE][3] papers. With the source code we provide a small training dataset taken from 400 frames generated by NVT ab-initio water MD trajectory with 300 frames for training and 100 for testing. [An example training parameter file](./examples/water/train/water_se_a.json) is provided. One can try with the training by
-```bash
-$ cd $deepmd_source_dir/examples/water/train/
-$ dp train water_se_a.json
-```
-where `water_se_a.json` is the `json` format parameter file that controls the training. It is also possible to use `yaml` format file with the same keys as json (see `water_se_a.yaml` example). You can use script `json2yaml.py` in `data/json/` dir to convert your json files to yaml. The components of the `water.json` contains four parts, `model`, `learning_rate`, `loss` and `training`.
-
-The `model` section specify how the deep potential model is built. An example of the smooth-edition is provided as follows
-```json
- "model": {
- "type_map": ["O", "H"],
- "descriptor" :{
- "type": "se_a",
- "rcut_smth": 5.80,
- "rcut": 6.00,
- "sel": [46, 92],
- "neuron": [25, 50, 100],
- "axis_neuron": 16,
- "resnet_dt": false,
- "seed": 1,
- "_comment": " that's all"
- },
- "fitting_net" : {
- "neuron": [240, 240, 240],
- "resnet_dt": true,
- "seed": 1,
- "_comment": " that's all"
- },
- "_comment": " that's all"
- }
-```
-The **`type_map`** is optional, which provide the element names (but not restricted to) for corresponding atom types.
-The construction of the descriptor is given by option **`descriptor`**. The **`type`** of the descriptor is set to `"se_a"`, which means smooth-edition, angular infomation. The **`rcut`** is the cut-off radius for neighbor searching, and the **`rcut_smth`** gives where the smoothing starts. **`sel`** gives the maximum possible number of neighbors in the cut-off radius. It is a list, the length of which is the same as the number of atom types in the system, and `sel[i]` denote the maximum possible number of neighbors with type `i`. The **`neuron`** specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layers from input end to the output end, respectively. The **`axis_neuron`** specifies the size of submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper][3]. If the outer layer is of twice size as the inner layer, then the inner layer is copied and concatenated, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is build between them. If the option **`resnet_dt`** is set `true`, then a timestep is used in the ResNet. **`seed`** gives the random seed that is used to generate random numbers when initializing the model parameters.
+The typical procedure of using DeePMD-kit includes 5 steps
-The construction of the fitting net is give by **`fitting_net`**. The key **`neuron`** specifies the size of the fitting net. If two neighboring layers are of the same size, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is build between them. If the option **`resnet_dt`** is set `true`, then a timestep is used in the ResNet. **`seed`** gives the random seed that is used to generate random numbers when initializing the model parameters.
+1. [Prepare data](doc/use-deepmd-kit.md#prepare-data)
+2. [Train a model](doc/use-deepmd-kit.md#train-a-model)
+3. [Freeze the model](doc/use-deepmd-kit.md#freeze-a-model)
+4. [Test the model](doc/use-deepmd-kit.md#test-a-model)
+5. [Inference the model in python](doc/use-deepmd-kit.md#model-inference) or using the model in other molecular simulation packages like [LAMMPS](doc/use-deepmd-kit.md#run-md-with-lammps), [i-PI](doc/use-deepmd-kit.md#run-path-integral-md-with-i-pi) or [ASE](doc/use-deepmd-kit.md#use-deep-potential-with-ase).
-An example of the `learning_rate` is given as follows
-```json
- "learning_rate" :{
- "type": "exp",
- "start_lr": 0.005,
- "decay_steps": 5000,
- "decay_rate": 0.95,
- "_comment": "that's all"
- }
-```
-The option **`start_lr`**, **`decay_rate`** and **`decay_steps`** specify how the learning rate changes. For example, the `t`th batch will be trained with learning rate:
-```math
-lr(t) = start_lr * decay_rate ^ ( t / decay_steps )
-```
-
-An example of the `loss` is
-```json
- "loss" : {
- "start_pref_e": 0.02,
- "limit_pref_e": 1,
- "start_pref_f": 1000,
- "limit_pref_f": 1,
- "start_pref_v": 0,
- "limit_pref_v": 0,
- "_comment": " that's all"
- }
-```
-The options **`start_pref_e`**, **`limit_pref_e`**, **`start_pref_f`**, **`limit_pref_f`**, **`start_pref_v`** and **`limit_pref_v`** determine how the prefactors of energy error, force error and virial error changes in the loss function (see the appendix of the [DeePMD paper][2] for details). Taking the prefactor of force error for example, the prefactor at batch `t` is
-```math
-w_f(t) = start_pref_f * ( lr(t) / start_lr ) + limit_pref_f * ( 1 - lr(t) / start_lr )
-```
-Since we do not have virial data, the virial prefactors `start_pref_v` and `limit_pref_v` are set to 0.
-
-An example of `training` is
-```json
- "training" : {
- "systems": ["../data1/", "../data2/"],
- "set_prefix": "set",
- "stop_batch": 1000000,
- "_comment": " batch_size can be supplied with, e.g. 1, or auto (string) or [10, 20]",
- "batch_size": 1,
-
- "seed": 1,
-
- "_comment": " display and restart",
- "_comment": " frequencies counted in batch",
- "disp_file": "lcurve.out",
- "disp_freq": 100,
- "_comment": " numb_test can be supplied with, e.g. 1, or XX% (string) or [10, 20]",
- "numb_test": 10,
- "save_freq": 1000,
- "save_ckpt": "model.ckpt",
- "load_ckpt": "model.ckpt",
- "disp_training":true,
- "time_training":true,
- "profiling": false,
- "profiling_file":"timeline.json",
- "_comment": "that's all"
- }
-```
-The option **`systems`** provide location of the systems (path to `set.*` and `type.raw`). It is a vector, thus DeePMD-kit allows you to provide multiple systems. DeePMD-kit will train the model with the systems in the vector one by one in a cyclic manner. **It is warned that the example water data (in folder `examples/data/water`) is of very limited amount, is provided only for testing purpose, and should not be used to train a productive model.**
-
-The option **`batch_size`** specifies the number of frames in each batch. It can be set to `"auto"` to enable a automatic batch size or it can be input as a list setting batch size individually for each system.
-The option **`stop_batch`** specifies the total number of batches will be used in the training.
-
-The option **`numb_test`** specifies the number of tests that will be used for each system. If it is an integer each system will be tested with the same number of tests. It can be set to percentage `"XX%"` to use XX% of frames of each system for its testing or it can be input as a list setting numer of tests individually for each system (the order should correspond to ordering of the systems key in json).
-
-### Training
-
-The training can be invoked by
-```bash
-$ dp train water_se_a.json
-```
-
-During the training, the error of the model is tested every **`disp_freq`** batches with **`numb_test`** frames from the last set in the **`systems`** directory on the fly, and the results are output to **`disp_file`**. A typical `disp_file` looks like
-```bash
-# batch l2_tst l2_trn l2_e_tst l2_e_trn l2_f_tst l2_f_trn lr
- 0 2.67e+01 2.57e+01 2.21e-01 2.22e-01 8.44e-01 8.12e-01 1.0e-03
- 100 6.14e+00 5.40e+00 3.01e-01 2.99e-01 1.93e-01 1.70e-01 1.0e-03
- 200 5.02e+00 4.49e+00 1.53e-01 1.53e-01 1.58e-01 1.42e-01 1.0e-03
- 300 4.36e+00 3.71e+00 7.32e-02 7.27e-02 1.38e-01 1.17e-01 1.0e-03
- 400 4.04e+00 3.29e+00 3.16e-02 3.22e-02 1.28e-01 1.04e-01 1.0e-03
-```
-The first column displays the number of batches. The second and third columns display the loss function evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The fourth and fifth columns display the RMS energy error (normalized by number of atoms) evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The sixth and seventh columns display the RMS force error (component-wise) evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The last column displays the current learning rate.
+A quick-start on using DeePMD-kit can be found [here](doc/use-deepmd-kit.md).
-Checkpoints will be written to files with prefix **`save_ckpt`** every **`save_freq`** batches. If **`restart`** is set to `true`, then the training will start from the checkpoint named **`load_ckpt`**, rather than from scratch.
-
-Several command line options can be passed to `dp train`, which can be checked with
-```bash
-$ dp train --help
-```
-An explanation will be provided
-```
-positional arguments:
- INPUT the input json database
-
-optional arguments:
- -h, --help show this help message and exit
- --init-model INIT_MODEL
- Initialize a model by the provided checkpoint
- --restart RESTART Restart the training from the provided checkpoint
-```
-The keys `intra_op_parallelism_threads` and `inter_op_parallelism_threads` are Tensorflow configurations for multithreading, which are explained [here](https://www.tensorflow.org/performance/performance_guide#optimizing_for_cpu). Skipping `-t` and `OMP_NUM_THREADS` leads to the default setting of these keys in the Tensorflow.
+A full [document](doc/train-input.rst) on options in the training input script is available.
-**`--init-model model.ckpt`**, for example, initializes the model training with an existing model that is stored in the checkpoint `model.ckpt`, the network architectures should match.
-
-**`--restart model.ckpt`**, continues the training from the checkpoint `model.ckpt`.
-
-On some resources limited machines, one may want to control the number of threads used by DeePMD-kit. This is achieved by three environmental variables: `OMP_NUM_THREADS`, `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`. `OMP_NUM_THREADS` controls the multithreading of DeePMD-kit implemented operations. `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS` controls `intra_op_parallelism_threads` and `inter_op_parallelism_threads`, which are Tensorflow configurations for multithreading. An explanation is found [here](https://stackoverflow.com/questions/41233635/meaning-of-inter-op-parallelism-threads-and-intra-op-parallelism-threads).
-
-For example 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:
-```bash
-export OMP_NUM_THREADS=6
-export TF_INTRA_OP_PARALLELISM_THREADS=3
-export TF_INTER_OP_PARALLELISM_THREADS=2
-dp train input.json
-```
-
-## Freeze a model
-
-The trained neural network is extracted from a checkpoint and dumped into a database. 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
-```bash
-$ dp freeze -o graph.pb
-```
-in the folder where the model is trained. The output database is called `graph.pb`.
-
-
-## Test a model
-
-The frozen model can be used in many ways. The most straightforward test can be performed using `dp test`. A typical usage of `dp test` is
-```bash
-dp test -m graph.pb -s /path/to/system -n 30
-```
-where `-m` gives the tested model, `-s` the path to the tested system and `-n` the number of tested frames. Several other command line options can be passed to `dp test`, which can be checked with
-```bash
-$ dp test --help
-```
-An explanation will be provided
-```
-usage: dp test [-h] [-m MODEL] [-s SYSTEM] [-S SET_PREFIX] [-n NUMB_TEST]
- [-r RAND_SEED] [--shuffle-test] [-d DETAIL_FILE]
-
-optional arguments:
- -h, --help show this help message and exit
- -m MODEL, --model MODEL
- Frozen model file to import
- -s SYSTEM, --system SYSTEM
- The system dir
- -S SET_PREFIX, --set-prefix SET_PREFIX
- The set prefix
- -n NUMB_TEST, --numb-test NUMB_TEST
- The number of data for test
- -r RAND_SEED, --rand-seed RAND_SEED
- The random seed
- --shuffle-test Shuffle test data
- -d DETAIL_FILE, --detail-file DETAIL_FILE
- The file containing details of energy force and virial
- accuracy
-```
-
-## Model inference
-One may use the python interface of DeePMD-kit for model inference, an example is given as follows
-```python
-import deepmd.DeepPot as DP
-import numpy as np
-dp = DP('graph.pb')
-coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
-cell = np.diag(10 * np.ones(3)).reshape([1, -1])
-atype = [1,0,1]
-e, f, v = dp.eval(coord, cell, atype)
-```
-where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively.
-
-
-## Run MD with LAMMPS
-### Include deepmd in the pair style
-Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows
-```bash
-pair_style deepmd graph.pb
-pair_coeff
-```
-where `graph.pb` is the file name of the frozen model. The `pair_coeff` should be left blank. It should be noted that LAMMPS counts atom types starting from 1, therefore, all LAMMPS atom type will be firstly subtracted by 1, and then passed into the DeePMD-kit engine to compute the interactions. [A detailed documentation of this pair style is available.](doc/lammps-pair-style-deepmd.md).
-
-### Long-range interaction
-The reciprocal space part of the long-range interaction can be calculated by LAMMPS command `kspace_style`. To use it with DeePMD-kit, one writes
-```bash
-pair_style deepmd graph.pb
-pair_coeff
-kspace_style pppm 1.0e-5
-kspace_modify gewald 0.45
-```
-Please notice that the DeePMD does nothing to the direct space part of the electrostatic interaction, because this part is assumed to be fitted in the DeePMD model (the direct space cut-off is thus the cut-off of the DeePMD model). The splitting parameter `gewald` is modified by the `kspace_modify` command.
-
-## Run path-integral MD with i-PI
-The i-PI works in a client-server model. The i-PI provides the server for integrating the replica positions of atoms, while the DeePMD-kit provides a client named `dp_ipi` that computes the interactions (including energy, force and virial). The server and client communicates via the Unix domain socket or the Internet socket. The client can be started by
-```bash
-$ dp_ipi water.json
-```
-It is noted that multiple instances of the client is allow for computing, in parallel, the interactions of multiple replica of the path-integral MD.
-
-`water.json` is the parameter file for the client `dp_ipi`, and [an example](./examples/ipi/water.json) is provided:
-```json
-{
- "verbose": false,
- "use_unix": true,
- "port": 31415,
- "host": "localhost",
- "graph_file": "graph.pb",
- "coord_file": "conf.xyz",
- "atom_type" : {
- "OW": 0,
- "HW1": 1,
- "HW2": 1
- }
-}
-```
-The option **`use_unix`** is set to `true` to activate the Unix domain socket, otherwise, the Internet socket is used.
-
-The option **`graph_file`** provides the file name of the frozen model.
-
-The `dp_ipi` gets the atom names from an [XYZ file](https://en.wikipedia.org/wiki/XYZ_file_format) provided by **`coord_file`** (meanwhile ignores all coordinates in it), and translates the names to atom types by rules provided by **`atom_type`**.
-
-## Use deep potential with ASE
-
-Deep potential can be set up as a calculator with ASE to obtain potential energies and forces.
-```python
-from ase import Atoms
-from deepmd.calculator import DP
-
-water = Atoms('H2O',
- positions=[(0.7601, 1.9270, 1),
- (1.9575, 1, 1),
- (1., 1., 1.)],
- cell=[100, 100, 100],
- calculator=DP(model="frozen_model.pb"))
-print(water.get_potential_energy())
-print(water.get_forces())
-```
-
-Optimization is also available:
-```python
-from ase.optimize import BFGS
-dyn = BFGS(water)
-dyn.run(fmax=1e-6)
-print(water.get_positions())
-```
# Troubleshooting
In consequence of various differences of computers or systems, problems may occur. Some common circumstances are listed as follows.
diff --git a/doc/install.md b/doc/install.md
new file mode 100644
index 0000000000..cf5ddbddbd
--- /dev/null
+++ b/doc/install.md
@@ -0,0 +1,200 @@
+- [Download and install](#download-and-install)
+ - [Easy installation methods](#easy-installation-methods)
+ - [Offline packages](#offline-packages)
+ - [With Docker](#with-docker)
+ - [With conda](#with-conda)
+ - [Install the python interaction](#install-the-python-interface)
+ - [Install the Tensorflow's python interface](#install-the-tensorflows-python-interface)
+ - [Install the DeePMD-kit's python interface](#install-the-deepmd-kits-python-interface)
+ - [Install the C++ interface](#install-the-c-interface)
+ - [Install the Tensorflow's C++ interface](#install-the-tensorflows-c-interface)
+ - [Install the DeePMD-kit's C++ interface](#install-the-deepmd-kits-c-interface)
+ - [Install LAMMPS's DeePMD-kit module](#install-lammpss-deepmd-kit-module)
+
+# Download and install
+
+Please follow our [github](https://github.com/deepmodeling/deepmd-kit) webpage to download the [latest released version](https://github.com/deepmodeling/deepmd-kit/tree/master) and [development version](https://github.com/deepmodeling/deepmd-kit/tree/devel).
+
+## Easy installation methods
+There various easy methods to install DeePMD-kit. Choose one that you prefer. If you want to build by yourself, jump to the next two sections.
+
+After your easy installation, DeePMD-kit (`dp`) and LAMMPS (`lmp`) will be available to execute. You can try `dp -h` and `lmp -h` to see the help. `mpirun` is also available considering you may want to run LAMMPS in parallel.
+
+### Offline packages
+Both CPU and GPU version offline packages are avaiable in [the Releases page](https://github.com/deepmodeling/deepmd-kit/releases).
+
+### With conda
+DeePMD-kit is avaiable with [conda](https://github.com/conda/conda). Install [Anaconda](https://www.anaconda.com/distribution/#download-section) or [Miniconda](https://docs.conda.io/en/latest/miniconda.html) first.
+
+To install the CPU version:
+```bash
+conda install deepmd-kit=*=*cpu lammps-dp=*=*cpu -c deepmodeling
+```
+
+To install the GPU version containing [CUDA 10.1](https://docs.nvidia.com/deploy/cuda-compatibility/index.html#binary-compatibility__table-toolkit-driver):
+```bash
+conda install deepmd-kit=*=*gpu lammps-dp=*=*gpu -c deepmodeling
+```
+
+### With Docker
+A docker for installing the DeePMD-kit is available [here](https://github.com/orgs/deepmodeling/packages/container/package/deepmd-kit).
+
+To pull the CPU version:
+```bash
+docker pull ghcr.io/deepmodeling/deepmd-kit:1.2.2_cpu
+```
+
+To pull the GPU version:
+```bash
+docker pull ghcr.io/deepmodeling/deepmd-kit:1.2.2_cuda10.1_gpu
+```
+
+## Install the python interface
+### Install the Tensorflow's python interface
+First, check the python version on your machine
+```bash
+python --version
+```
+
+We follow the virtual environment approach to install the tensorflow's Python interface. The full instruction can be found on [the tensorflow's official website](https://www.tensorflow.org/install/pip). Now we assume that the Python interface will be installed to virtual environment directory `$tensorflow_venv`
+```bash
+virtualenv -p python3 $tensorflow_venv
+source $tensorflow_venv/bin/activate
+pip install --upgrade pip
+pip install --upgrade tensorflow==2.3.0
+```
+It is notice that everytime a new shell is started and one wants to use `DeePMD-kit`, the virtual environment should be activated by
+```bash
+source $tensorflow_venv/bin/activate
+```
+if one wants to skip out of the virtual environment, he/she can do
+```bash
+deactivate
+```
+If one has multiple python interpreters named like python3.x, it can be specified by, for example
+```bash
+virtualenv -p python3.7 $tensorflow_venv
+```
+If one does not need the GPU support of deepmd-kit and is concerned about package size, the CPU-only version of tensorflow should be installed by
+```bash
+pip install --upgrade tensorflow-cpu==2.3.0
+```
+To verify the installation, run
+```bash
+python -c "import tensorflow as tf;print(tf.reduce_sum(tf.random.normal([1000, 1000])))"
+```
+One should remember to activate the virtual environment every time he/she uses deepmd-kit.
+
+### Install the DeePMD-kit's python interface
+
+Execute
+```bash
+pip install deepmd-kit
+```
+To test the installation, one may execute
+```bash
+dp -h
+```
+It will print the help information like
+```text
+usage: dp [-h] {train,freeze,test} ...
+
+DeePMD-kit: A deep learning package for many-body potential energy
+representation and molecular dynamics
+
+optional arguments:
+ -h, --help show this help message and exit
+
+Valid subcommands:
+ {train,freeze,test}
+ train train a model
+ freeze freeze the model
+ test test the model
+```
+
+## Install the C++ interface
+
+If one does not need to use DeePMD-kit with Lammps or I-Pi, then the python interface installed in the previous section does everything and he/she can safely skip this section.
+
+### Install the Tensorflow's C++ interface
+
+Check the compiler version on your machine
+
+```
+gcc --version
+```
+
+The C++ interface of DeePMD-kit was tested with compiler gcc >= 4.8. It is noticed that the I-Pi support is only compiled with gcc >= 4.9.
+
+First the C++ interface of Tensorflow should be installed. It is noted that the version of Tensorflow should be in consistent with the python interface. We assume that you have followed our instruction and installed tensorflow python interface 1.14.0 with, then you may follow [the instruction for CPU](doc/install-tf.1.14.md) to install the corresponding C++ interface (CPU only). If one wants GPU supports, he/she should follow [the instruction for GPU](doc/install-tf.1.14-gpu.md) to install the C++ interface.
+
+### Install the DeePMD-kit's C++ interface
+
+Clone the DeePMD-kit source code
+```bash
+cd /some/workspace
+git clone --recursive https://github.com/deepmodeling/deepmd-kit.git deepmd-kit
+```
+
+For convenience, you may want to record the location of source to a variable, saying `deepmd_source_dir` by
+```bash
+cd deepmd-kit
+deepmd_source_dir=`pwd`
+```
+
+Now goto the source code directory of DeePMD-kit and make a build place.
+```bash
+cd $deepmd_source_dir/source
+mkdir build
+cd build
+```
+I assume you want to install DeePMD-kit into path `$deepmd_root`, then execute cmake
+```bash
+cmake -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
+```
+where the variable `tensorflow_root` stores the location where the tensorflow's C++ interface is installed. The DeePMD-kit will automatically detect if a CUDA tool-kit is available on your machine and build the GPU support accordingly. If you want to force the cmake to find CUDA tool-kit, you can speicify the key `USE_CUDA_TOOLKIT`,
+```bash
+cmake -DUSE_CUDA_TOOLKIT=true -DTENSORFLOW_ROOT=$tensorflow_root -DCMAKE_INSTALL_PREFIX=$deepmd_root ..
+```
+and you may further asked to provide `CUDA_TOOLKIT_ROOT_DIR`. If the cmake has executed successfully, then
+```bash
+make
+make install
+```
+If everything works fine, you will have the following executable and libraries installed in `$deepmd_root/bin` and `$deepmd_root/lib`
+```bash
+$ ls $deepmd_root/bin
+dp_ipi
+$ ls $deepmd_root/lib
+libdeepmd_ipi.so libdeepmd_op.so libdeepmd.so
+```
+
+### Install LAMMPS's DeePMD-kit module
+DeePMD-kit provide module for running MD simulation with LAMMPS. Now make the DeePMD-kit module for LAMMPS.
+```bash
+cd $deepmd_source_dir/source/build
+make lammps
+```
+DeePMD-kit will generate a module called `USER-DEEPMD` in the `build` directory. Now download your favorite LAMMPS code, and uncompress it (I assume that you have downloaded the tar `lammps-stable.tar.gz`)
+```bash
+cd /some/workspace
+tar xf lammps-stable.tar.gz
+```
+The source code of LAMMPS is stored in directory, for example `lammps-31Mar17`. Now go into the LAMMPS code and copy the DeePMD-kit module like this
+```bash
+cd lammps-31Mar17/src/
+cp -r $deepmd_source_dir/source/build/USER-DEEPMD .
+```
+Now build LAMMPS
+```bash
+make yes-user-deepmd
+make mpi -j4
+```
+The option `-j4` means using 4 processes in parallel. You may want to use a different number according to your hardware.
+
+If everything works fine, you will end up with an executable `lmp_mpi`.
+
+The DeePMD-kit module can be removed from LAMMPS source code by
+```bash
+make no-user-deepmd
+```
diff --git a/doc/train-input.rst b/doc/train-input.rst
new file mode 100644
index 0000000000..e29dd5799b
--- /dev/null
+++ b/doc/train-input.rst
@@ -0,0 +1,1023 @@
+.. raw:: html
+
+
+model:
+ | type: ``dict``
+ | argument path: ``model``
+
+ .. raw:: html
+
+
+ type_map:
+ | type: ``list``, optional
+ | argument path: ``model/type_map``
+
+ A list of strings. Give the name to each type of atoms.
+
+ .. raw:: html
+
+
+ data_stat_nbatch:
+ | type: ``int``, optional, default: ``10``
+ | argument path: ``model/data_stat_nbatch``
+
+ The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.
+
+ .. raw:: html
+
+
+ descriptor:
+ | type: ``dict``
+ | argument path: ``model/descriptor``
+
+ The descriptor of atomic environment.
+
+
+ Depending on the value of *type*, different sub args are accepted.
+
+ .. raw:: html
+
+
+ type:
+ | type: ``str`` (flag key)
+ | argument path: ``model/descriptor/type``
+
+ The type of the descritpor. Valid types are `loc_frame`, `se_a`, `se_r` and `se_ar`.
+
+ - `loc_frame`: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.
+
+ - `se_a`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.
+
+ - `se_r`: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.
+
+ - `se_ar`: A hybrid of `se_a` and `se_r`. Typically `se_a` has a smaller cut-off while the `se_r` has a larger cut-off.
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``loc_frame``:
+
+ .. raw:: html
+
+
+ sel_a:
+ | type: ``list``
+ | argument path: ``model/descriptor[loc_frame]/sel_a``
+
+ A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_a[i]` gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.
+
+ .. raw:: html
+
+
+ sel_r:
+ | type: ``list``
+ | argument path: ``model/descriptor[loc_frame]/sel_r``
+
+ A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_r[i]` gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.
+
+ .. raw:: html
+
+
+ rcut:
+ | type: ``float``, optional, default: ``6.0``
+ | argument path: ``model/descriptor[loc_frame]/rcut``
+
+ The cut-off radius. The default value is 6.0
+
+ .. raw:: html
+
+
+ axis_rule:
+ | type: ``list``
+ | argument path: ``model/descriptor[loc_frame]/axis_rule``
+
+ A list of integers. The length should be 6 times of the number of types.
+
+ - axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.
+
+ - axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom.
+
+ - axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance.
+
+ - axis_rule[i*6+3]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.
+
+ - axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom.
+
+ - axis_rule[i*6+5]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``se_a``:
+
+ .. raw:: html
+
+
+ sel:
+ | type: ``list``
+ | argument path: ``model/descriptor[se_a]/sel``
+
+ A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.
+
+ .. raw:: html
+
+
+ rcut:
+ | type: ``float``, optional, default: ``6.0``
+ | argument path: ``model/descriptor[se_a]/rcut``
+
+ The cut-off radius.
+
+ .. raw:: html
+
+
+ rcut_smth:
+ | type: ``float``, optional, default: ``0.5``
+ | argument path: ``model/descriptor[se_a]/rcut_smth``
+
+ Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[10, 20, 40]``
+ | argument path: ``model/descriptor[se_a]/neuron``
+
+ Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.
+
+ .. raw:: html
+
+
+ axis_neuron:
+ | type: ``int``, optional, default: ``4``
+ | argument path: ``model/descriptor[se_a]/axis_neuron``
+
+ Size of the submatrix of G (embedding matrix).
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/descriptor[se_a]/activation_function``
+
+ The activation function in the embedding net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_a]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ type_one_side:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_a]/type_one_side``
+
+ Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/descriptor[se_a]/precision``
+
+ The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ trainable:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/descriptor[se_a]/trainable``
+
+ If the parameters in the embedding net is trainable
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/descriptor[se_a]/seed``
+
+ Random seed for parameter initialization
+
+ .. raw:: html
+
+
+ exclude_types:
+ | type: ``list``, optional, default: ``[]``
+ | argument path: ``model/descriptor[se_a]/exclude_types``
+
+ The Excluded types
+
+ .. raw:: html
+
+
+ set_davg_zero:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_a]/set_davg_zero``
+
+ Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``se_r``:
+
+ .. raw:: html
+
+
+ sel:
+ | type: ``list``
+ | argument path: ``model/descriptor[se_r]/sel``
+
+ A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.
+
+ .. raw:: html
+
+
+ rcut:
+ | type: ``float``, optional, default: ``6.0``
+ | argument path: ``model/descriptor[se_r]/rcut``
+
+ The cut-off radius.
+
+ .. raw:: html
+
+
+ rcut_smth:
+ | type: ``float``, optional, default: ``0.5``
+ | argument path: ``model/descriptor[se_r]/rcut_smth``
+
+ Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[10, 20, 40]``
+ | argument path: ``model/descriptor[se_r]/neuron``
+
+ Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/descriptor[se_r]/activation_function``
+
+ The activation function in the embedding net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_r]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ type_one_side:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_r]/type_one_side``
+
+ Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/descriptor[se_r]/precision``
+
+ The precision of the embedding net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ trainable:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/descriptor[se_r]/trainable``
+
+ If the parameters in the embedding net is trainable
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/descriptor[se_r]/seed``
+
+ Random seed for parameter initialization
+
+ .. raw:: html
+
+
+ exclude_types:
+ | type: ``list``, optional, default: ``[]``
+ | argument path: ``model/descriptor[se_r]/exclude_types``
+
+ The Excluded types
+
+ .. raw:: html
+
+
+ set_davg_zero:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``model/descriptor[se_r]/set_davg_zero``
+
+ Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``se_ar``:
+
+ .. raw:: html
+
+
+ a:
+ | type: ``dict``
+ | argument path: ``model/descriptor[se_ar]/a``
+
+ The parameters of descriptor `se_a <#model/descriptor[se_a]>`__
+
+ .. raw:: html
+
+
+ r:
+ | type: ``dict``
+ | argument path: ``model/descriptor[se_ar]/r``
+
+ The parameters of descriptor `se_r <#model/descriptor[se_r]>`__
+
+ .. raw:: html
+
+
+ fitting_net:
+ | type: ``dict``
+ | argument path: ``model/fitting_net``
+
+ The fitting of physical properties.
+
+
+ Depending on the value of *type*, different sub args are accepted.
+
+ .. raw:: html
+
+
+ type:
+ | type: ``str`` (flag key), default: ``ener``
+ | argument path: ``model/fitting_net/type``
+
+ The type of the fitting. Valid types are `ener`, `dipole`, `polar` and `global_polar`.
+
+ - `ener`: Fit an energy model (potential energy surface).
+
+ - `dipole`: Fit an atomic dipole model. Atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file has number of frames lines and 3 times of number of selected atoms columns.
+
+ - `polar`: Fit an atomic polarizability model. Atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file has number of frames lines and 9 times of number of selected atoms columns.
+
+ - `global_polar`: Fit a polarizability model. Polarizability labels should be provided by `polarizability.npy` in each data system. The file has number of frames lines and 9 columns.
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``ener``:
+
+ .. raw:: html
+
+
+ numb_fparam:
+ | type: ``int``, optional, default: ``0``
+ | argument path: ``model/fitting_net[ener]/numb_fparam``
+
+ The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams.
+
+ .. raw:: html
+
+
+ numb_aparam:
+ | type: ``int``, optional, default: ``0``
+ | argument path: ``model/fitting_net[ener]/numb_aparam``
+
+ The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams.
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[120, 120, 120]``
+ | argument path: ``model/fitting_net[ener]/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.
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/fitting_net[ener]/activation_function``
+
+ The activation function in the fitting net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/fitting_net[ener]/precision``
+
+ The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[ener]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ trainable:
+ | type: ``bool`` | ``list``, optional, default: ``True``
+ | argument path: ``model/fitting_net[ener]/trainable``
+
+ Whether the parameters in the fitting net are trainable. This option can be
+
+ - bool: True if all parameters of the fitting net are trainable, False otherwise.
+
+ - list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of tihs list should be equal to len(`neuron`)+1.
+
+ .. raw:: html
+
+
+ rcond:
+ | type: ``float``, optional, default: ``0.001``
+ | argument path: ``model/fitting_net[ener]/rcond``
+
+ The condition number used to determine the inital energy shift for each type of atoms.
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/fitting_net[ener]/seed``
+
+ Random seed for parameter initialization of the fitting net
+
+ .. raw:: html
+
+
+ atom_ener:
+ | type: ``list``, optional, default: ``[]``
+ | argument path: ``model/fitting_net[ener]/atom_ener``
+
+ Specify the atomic energy in vacuum for each type
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``dipole``:
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[120, 120, 120]``
+ | argument path: ``model/fitting_net[dipole]/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.
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/fitting_net[dipole]/activation_function``
+
+ The activation function in the fitting net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[dipole]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/fitting_net[dipole]/precision``
+
+ The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ sel_type:
+ | type: ``int`` | ``NoneType`` | ``list``, optional
+ | argument path: ``model/fitting_net[dipole]/sel_type``
+
+ The atom types for which the atomic dipole will be provided. If not set, all types will be selected.
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/fitting_net[dipole]/seed``
+
+ Random seed for parameter initialization of the fitting net
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``polar``:
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[120, 120, 120]``
+ | argument path: ``model/fitting_net[polar]/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.
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/fitting_net[polar]/activation_function``
+
+ The activation function in the fitting net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[polar]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/fitting_net[polar]/precision``
+
+ The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ fit_diag:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[polar]/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.
+
+ .. raw:: html
+
+
+ scale:
+ | type: ``float`` | ``list``, optional, default: ``1.0``
+ | argument path: ``model/fitting_net[polar]/scale``
+
+ The output of the fitting net (polarizability matrix) will be scaled by ``scale``
+
+ .. raw:: html
+
+
+ diag_shift:
+ | type: ``float`` | ``list``, optional, default: ``0.0``
+ | argument path: ``model/fitting_net[polar]/diag_shift``
+
+ The diagonal part of the polarizability matrix will be shifted by ``diag_shift``. The shift operation is carried out after ``scale``.
+
+ .. raw:: html
+
+
+ sel_type:
+ | type: ``int`` | ``NoneType`` | ``list``, optional
+ | argument path: ``model/fitting_net[polar]/sel_type``
+
+ The atom types for which the atomic polarizability will be provided. If not set, all types will be selected.
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/fitting_net[polar]/seed``
+
+ Random seed for parameter initialization of the fitting net
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``global_polar``:
+
+ .. raw:: html
+
+
+ neuron:
+ | type: ``list``, optional, default: ``[120, 120, 120]``
+ | argument path: ``model/fitting_net[global_polar]/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.
+
+ .. raw:: html
+
+
+ activation_function:
+ | type: ``str``, optional, default: ``tanh``
+ | argument path: ``model/fitting_net[global_polar]/activation_function``
+
+ The activation function in the fitting net. Supported activation functions are "relu", "relu6", "softplus", "sigmoid", "tanh", "gelu".
+
+ .. raw:: html
+
+
+ resnet_dt:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[global_polar]/resnet_dt``
+
+ Whether to use a "Timestep" in the skip connection
+
+ .. raw:: html
+
+
+ precision:
+ | type: ``str``, optional, default: ``float64``
+ | argument path: ``model/fitting_net[global_polar]/precision``
+
+ The precision of the fitting net parameters, supported options are "float64", "float32", "float16".
+
+ .. raw:: html
+
+
+ fit_diag:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``model/fitting_net[global_polar]/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.
+
+ .. raw:: html
+
+
+ scale:
+ | type: ``float`` | ``list``, optional, default: ``1.0``
+ | argument path: ``model/fitting_net[global_polar]/scale``
+
+ The output of the fitting net (polarizability matrix) will be scaled by ``scale``
+
+ .. raw:: html
+
+
+ diag_shift:
+ | type: ``float`` | ``list``, optional, default: ``0.0``
+ | argument path: ``model/fitting_net[global_polar]/diag_shift``
+
+ The diagonal part of the polarizability matrix will be shifted by ``diag_shift``. The shift operation is carried out after ``scale``.
+
+ .. raw:: html
+
+
+ sel_type:
+ | type: ``int`` | ``NoneType`` | ``list``, optional
+ | argument path: ``model/fitting_net[global_polar]/sel_type``
+
+ The atom types for which the atomic polarizability will be provided. If not set, all types will be selected.
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``model/fitting_net[global_polar]/seed``
+
+ Random seed for parameter initialization of the fitting net
+
+
+.. raw:: html
+
+
+loss:
+ | type: ``dict``
+ | argument path: ``loss``
+
+ The definition of loss function. The type of the loss depends on the type of the fitting. For fitting type `ener`, the prefactors before energy, force, virial and atomic energy losses may be provided. For fitting type `dipole`, `polar` and `global_polar`, the loss may be an empty `dict` or unset.
+
+
+ Depending on the value of *type*, different sub args are accepted.
+
+ .. raw:: html
+
+
+ type:
+ | type: ``str`` (flag key), default: ``ener``
+ | argument path: ``loss/type``
+
+ The type of the loss. For fitting type `ener`, the loss type should be set to `ener` or left unset. For tensorial fitting types `dipole`, `polar` and `global_polar`, the type should be left unset.
+ \.
+
+
+ .. raw:: html
+
+
+ When *type* is set to ``ener``:
+
+ .. raw:: html
+
+
+ start_pref_e:
+ | type: ``float`` | ``int``, optional, default: ``0.02``
+ | argument path: ``loss[ener]/start_pref_e``
+
+ The prefactor of energy loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the energy label should be provided by file energy.npy in each data system. If both start_pref_energy and limit_pref_energy are set to 0, then the energy will be ignored.
+
+ .. raw:: html
+
+
+ limit_pref_e:
+ | type: ``float`` | ``int``, optional, default: ``1.0``
+ | argument path: ``loss[ener]/limit_pref_e``
+
+ The prefactor of energy loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.
+
+ .. raw:: html
+
+
+ start_pref_f:
+ | type: ``float`` | ``int``, optional, default: ``1000``
+ | argument path: ``loss[ener]/start_pref_f``
+
+ The prefactor of force loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the force label should be provided by file force.npy in each data system. If both start_pref_force and limit_pref_force are set to 0, then the force will be ignored.
+
+ .. raw:: html
+
+
+ limit_pref_f:
+ | type: ``float`` | ``int``, optional, default: ``1.0``
+ | argument path: ``loss[ener]/limit_pref_f``
+
+ The prefactor of force loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.
+
+ .. raw:: html
+
+
+ start_pref_v:
+ | type: ``float`` | ``int``, optional, default: ``0.0``
+ | argument path: ``loss[ener]/start_pref_v``
+
+ The prefactor of virial loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the virial label should be provided by file virial.npy in each data system. If both start_pref_virial and limit_pref_virial are set to 0, then the virial will be ignored.
+
+ .. raw:: html
+
+
+ limit_pref_v:
+ | type: ``float`` | ``int``, optional, default: ``0.0``
+ | argument path: ``loss[ener]/limit_pref_v``
+
+ The prefactor of virial loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.
+
+ .. raw:: html
+
+
+ start_pref_ae:
+ | type: ``float`` | ``int``, optional, default: ``0.0``
+ | argument path: ``loss[ener]/start_pref_ae``
+
+ The prefactor of virial loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the virial label should be provided by file virial.npy in each data system. If both start_pref_virial and limit_pref_virial are set to 0, then the virial will be ignored.
+
+ .. raw:: html
+
+
+ limit_pref_ae:
+ | type: ``float`` | ``int``, optional, default: ``0.0``
+ | argument path: ``loss[ener]/limit_pref_ae``
+
+ The prefactor of virial loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.
+
+ .. raw:: html
+
+
+ relative_f:
+ | type: ``float`` | ``NoneType``, optional
+ | argument path: ``loss[ener]/relative_f``
+
+ If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label.
+
+
+.. raw:: html
+
+
+learning_rate:
+ | type: ``dict``
+ | argument path: ``learning_rate``
+
+ The learning rate options
+
+ .. raw:: html
+
+
+ start_lr:
+ | type: ``float``, optional, default: ``0.001``
+ | argument path: ``learning_rate/start_lr``
+
+ The learning rate the start of the training.
+
+ .. raw:: html
+
+
+ stop_lr:
+ | type: ``float``, optional, default: ``1e-08``
+ | argument path: ``learning_rate/stop_lr``
+
+ The desired learning rate at the end of the training.
+
+ .. raw:: html
+
+
+ decay_steps:
+ | type: ``int``, optional, default: ``5000``
+ | argument path: ``learning_rate/decay_steps``
+
+ The learning rate is decaying every this number of training steps.
+
+
+.. raw:: html
+
+
+training:
+ | type: ``dict``
+ | argument path: ``training``
+
+ The training options
+
+ .. raw:: html
+
+
+ systems:
+ | type: ``list`` | ``str``
+ | argument path: ``training/systems``
+
+ The data systems. This key can be provided with a listthat specifies the systems, or be provided with a string by which the prefix of all systems are given and the list of the systems is automatically generated.
+
+ .. raw:: html
+
+
+ set_prefix:
+ | type: ``str``, optional, default: ``set``
+ | argument path: ``training/set_prefix``
+
+ The prefix of the sets in the systems.
+
+ .. raw:: html
+
+
+ stop_batch:
+ | type: ``int``
+ | argument path: ``training/stop_batch``
+
+ Number of training batch. Each training uses one batch of data.
+
+ .. raw:: html
+
+
+ batch_size:
+ | type: ``int`` | ``list`` | ``str``, optional, default: ``auto``
+ | argument path: ``training/batch_size``
+
+ This key can be
+
+ - list: the length of which is the same as the `systems`. The batch size of each system is given by the elements of the list.
+
+ - int: all `systems` uses the same batch size.
+
+ - string "auto": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than 32.
+
+ - string "auto:N": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than N.
+
+ .. raw:: html
+
+
+ seed:
+ | type: ``int`` | ``NoneType``, optional
+ | argument path: ``training/seed``
+
+ The random seed for training.
+
+ .. raw:: html
+
+
+ disp_file:
+ | type: ``str``, optional, default: ``lcueve.out``
+ | argument path: ``training/disp_file``
+
+ The file for printing learning curve.
+
+ .. raw:: html
+
+
+ disp_freq:
+ | type: ``int``, optional, default: ``1000``
+ | argument path: ``training/disp_freq``
+
+ The frequency of printing learning curve.
+
+ .. raw:: html
+
+
+ numb_test:
+ | type: ``int``, optional, default: ``1``
+ | argument path: ``training/numb_test``
+
+ Number of frames used for the test during training.
+
+ .. raw:: html
+
+
+ save_freq:
+ | type: ``int``, optional, default: ``1000``
+ | argument path: ``training/save_freq``
+
+ The frequency of saving check point.
+
+ .. raw:: html
+
+
+ save_ckpt:
+ | type: ``str``, optional, default: ``model.ckpt``
+ | argument path: ``training/save_ckpt``
+
+ The file name of saving check point.
+
+ .. raw:: html
+
+
+ disp_training:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``training/disp_training``
+
+ Displaying verbose information during training.
+
+ .. raw:: html
+
+
+ time_training:
+ | type: ``bool``, optional, default: ``True``
+ | argument path: ``training/time_training``
+
+ Timing durining training.
+
+ .. raw:: html
+
+
+ profiling:
+ | type: ``bool``, optional, default: ``False``
+ | argument path: ``training/profiling``
+
+ Profiling during training.
+
+ .. raw:: html
+
+
+ profiling_file:
+ | type: ``str``, optional, default: ``timeline.json``
+ | argument path: ``training/profiling_file``
+
+ Output file for profiling.
+
diff --git a/doc/use-deepmd-kit.md b/doc/use-deepmd-kit.md
new file mode 100644
index 0000000000..cbf7a91cbc
--- /dev/null
+++ b/doc/use-deepmd-kit.md
@@ -0,0 +1,352 @@
+- [Use DeePMD-kit](#use-deepmd-kit)
+ - [Prepare data](#prepare-data)
+ - [Train a model](#train-a-model)
+ - [The DeePMD model](#the-deepmd-model)
+ - [The DeepPot-SE model](#the-deeppot-se-model)
+ - [Freeze a model](#freeze-a-model)
+ - [Test a model](#test-a-model)
+ - [Model inference](#model-inference)
+ - [Run MD with Lammps](#run-md-with-lammps)
+ - [Include deepmd in the pair style](#include-deepmd-in-the-pair-style)
+ - [Long-range interaction](#long-range-interaction)
+ - [Run path-integral MD with i-PI](#run-path-integral-md-with-i-pi)
+ - [Use deep potential with ASE](#use-deep-potential-with-ase)
+
+# Use DeePMD-kit
+In this text, we will call the deep neural network that is used to represent the interatomic interactions (Deep Potential) the **model**. The typical procedure of using DeePMD-kit is
+
+1. Prepare data
+2. Train a model
+3. Freeze the model
+4. Test the model
+5. Inference with the model
+
+## Prepare data
+One needs to provide the following information to train a model: the atom type, the simulation box, the atom coordinate, the atom force, system energy and virial. A snapshot of a system that contains these information is called a **frame**. We use the following convention of units:
+
+Property| Unit
+--- | :---:
+Time | ps
+Length | Å
+Energy | eV
+Force | eV/Å
+Pressure| Bar
+
+The frames of the system are stored in two formats. A raw file is a plain text file with each information item written in one file and one frame written on one line. The default files that provide box, coordinate, force, energy and virial are `box.raw`, `coord.raw`, `force.raw`, `energy.raw` and `virial.raw`, respectively. *We recommend you use these file names*. Here is an example of force.raw:
+```bash
+$ cat force.raw
+-0.724 2.039 -0.951 0.841 -0.464 0.363
+ 6.737 1.554 -5.587 -2.803 0.062 2.222
+-1.968 -0.163 1.020 -0.225 -0.789 0.343
+```
+This `force.raw` contains 3 frames with each frame having the forces of 2 atoms, thus it has 3 lines and 6 columns. Each line provides all the 3 force components of 2 atoms in 1 frame. The first three numbers are the 3 force components of the first atom, while the second three numbers are the 3 force components of the second atom. The coordinate file `coord.raw` is organized similarly. In `box.raw`, the 9 components of the box vectors should be provided on each line. In `virial.raw`, the 9 components of the virial tensor should be provided on each line. The number of lines of all raw files should be identical.
+
+We assume that the atom types do not change in all frames. It is provided by `type.raw`, which has one line with the types of atoms written one by one. The atom types should be integers. For example the `type.raw` of a system that has 2 atoms with 0 and 1:
+```bash
+$ cat type.raw
+0 1
+```
+
+The second format is the data sets of `numpy` binary data that are directly used by the training program. User can use the script `$deepmd_source_dir/data/raw/raw_to_set.sh` to convert the prepared raw files to data sets. For example, if we have a raw file that contains 6000 frames,
+```bash
+$ ls
+box.raw coord.raw energy.raw force.raw type.raw virial.raw
+$ $deepmd_source_dir/data/raw/raw_to_set.sh 2000
+nframe is 6000
+nline per set is 2000
+will make 3 sets
+making set 0 ...
+making set 1 ...
+making set 2 ...
+$ ls
+box.raw coord.raw energy.raw force.raw set.000 set.001 set.002 type.raw virial.raw
+```
+It generates three sets `set.000`, `set.001` and `set.002`, with each set contains 2000 frames. The last set (`set.002`) is used as testing set, while the rest sets (`set.000` and `set.001`) are used as training sets. One do not need to take care of the binary data files in each of the `set.*` directories. The path containing `set.*` and `type.raw` is called a *system*.
+
+## Train a model
+
+### Write the input script
+
+The method of training is explained in our [DeePMD][2] and [DeepPot-SE][3] papers. With the source code we provide a small training dataset taken from 400 frames generated by NVT ab-initio water MD trajectory with 300 frames for training and 100 for testing. [An example training parameter file](./examples/water/train/water_se_a.json) is provided. One can try with the training by
+```bash
+$ cd $deepmd_source_dir/examples/water/train/
+$ dp train water_se_a.json
+```
+where `water_se_a.json` is the `json` format parameter file that controls the training. It is also possible to use `yaml` format file with the same keys as json (see `water_se_a.yaml` example). You can use script `json2yaml.py` in `data/json/` dir to convert your json files to yaml. The components of the `water.json` contains four parts, `model`, `learning_rate`, `loss` and `training`.
+
+The `model` section specify how the deep potential model is built. An example of the smooth-edition is provided as follows
+```json
+ "model": {
+ "type_map": ["O", "H"],
+ "descriptor" :{
+ "type": "se_a",
+ "rcut_smth": 5.80,
+ "rcut": 6.00,
+ "sel": [46, 92],
+ "neuron": [25, 50, 100],
+ "axis_neuron": 16,
+ "resnet_dt": false,
+ "seed": 1,
+ "_comment": " that's all"
+ },
+ "fitting_net" : {
+ "neuron": [240, 240, 240],
+ "resnet_dt": true,
+ "seed": 1,
+ "_comment": " that's all"
+ },
+ "_comment": " that's all"
+ }
+```
+The **`type_map`** is optional, which provide the element names (but not restricted to) for corresponding atom types.
+
+The construction of the descriptor is given by option **`descriptor`**. The **`type`** of the descriptor is set to `"se_a"`, which means smooth-edition, angular infomation. The **`rcut`** is the cut-off radius for neighbor searching, and the **`rcut_smth`** gives where the smoothing starts. **`sel`** gives the maximum possible number of neighbors in the cut-off radius. It is a list, the length of which is the same as the number of atom types in the system, and `sel[i]` denote the maximum possible number of neighbors with type `i`. The **`neuron`** specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layers from input end to the output end, respectively. The **`axis_neuron`** specifies the size of submatrix of the embedding matrix, the axis matrix as explained in the [DeepPot-SE paper][3]. If the outer layer is of twice size as the inner layer, then the inner layer is copied and concatenated, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is build between them. If the option **`resnet_dt`** is set `true`, then a timestep is used in the ResNet. **`seed`** gives the random seed that is used to generate random numbers when initializing the model parameters.
+
+The construction of the fitting net is give by **`fitting_net`**. The key **`neuron`** specifies the size of the fitting net. If two neighboring layers are of the same size, then a [ResNet architecture](https://arxiv.org/abs/1512.03385) is build between them. If the option **`resnet_dt`** is set `true`, then a timestep is used in the ResNet. **`seed`** gives the random seed that is used to generate random numbers when initializing the model parameters.
+
+An example of the `learning_rate` is given as follows
+```json
+ "learning_rate" :{
+ "type": "exp",
+ "start_lr": 0.005,
+ "decay_steps": 5000,
+ "decay_rate": 0.95,
+ "_comment": "that's all"
+ }
+```
+The option **`start_lr`**, **`decay_rate`** and **`decay_steps`** specify how the learning rate changes. For example, the `t`th batch will be trained with learning rate:
+```math
+lr(t) = start_lr * decay_rate ^ ( t / decay_steps )
+```
+
+An example of the `loss` is
+```json
+ "loss" : {
+ "start_pref_e": 0.02,
+ "limit_pref_e": 1,
+ "start_pref_f": 1000,
+ "limit_pref_f": 1,
+ "start_pref_v": 0,
+ "limit_pref_v": 0,
+ "_comment": " that's all"
+ }
+```
+The options **`start_pref_e`**, **`limit_pref_e`**, **`start_pref_f`**, **`limit_pref_f`**, **`start_pref_v`** and **`limit_pref_v`** determine how the prefactors of energy error, force error and virial error changes in the loss function (see the appendix of the [DeePMD paper][2] for details). Taking the prefactor of force error for example, the prefactor at batch `t` is
+```math
+w_f(t) = start_pref_f * ( lr(t) / start_lr ) + limit_pref_f * ( 1 - lr(t) / start_lr )
+```
+Since we do not have virial data, the virial prefactors `start_pref_v` and `limit_pref_v` are set to 0.
+
+An example of `training` is
+```json
+ "training" : {
+ "systems": ["../data1/", "../data2/"],
+ "set_prefix": "set",
+ "stop_batch": 1000000,
+ "_comment": " batch_size can be supplied with, e.g. 1, or auto (string) or [10, 20]",
+ "batch_size": 1,
+
+ "seed": 1,
+
+ "_comment": " display and restart",
+ "_comment": " frequencies counted in batch",
+ "disp_file": "lcurve.out",
+ "disp_freq": 100,
+ "_comment": " numb_test can be supplied with, e.g. 1, or XX% (string) or [10, 20]",
+ "numb_test": 10,
+ "save_freq": 1000,
+ "save_ckpt": "model.ckpt",
+ "load_ckpt": "model.ckpt",
+ "disp_training":true,
+ "time_training":true,
+ "profiling": false,
+ "profiling_file":"timeline.json",
+ "_comment": "that's all"
+ }
+```
+The option **`systems`** provide location of the systems (path to `set.*` and `type.raw`). It is a vector, thus DeePMD-kit allows you to provide multiple systems. DeePMD-kit will train the model with the systems in the vector one by one in a cyclic manner. **It is warned that the example water data (in folder `examples/data/water`) is of very limited amount, is provided only for testing purpose, and should not be used to train a productive model.**
+
+The option **`batch_size`** specifies the number of frames in each batch. It can be set to `"auto"` to enable a automatic batch size or it can be input as a list setting batch size individually for each system.
+The option **`stop_batch`** specifies the total number of batches will be used in the training.
+
+The option **`numb_test`** specifies the number of tests that will be used for each system. If it is an integer each system will be tested with the same number of tests. It can be set to percentage `"XX%"` to use XX% of frames of each system for its testing or it can be input as a list setting numer of tests individually for each system (the order should correspond to ordering of the systems key in json).
+
+### Training
+
+The training can be invoked by
+```bash
+$ dp train water_se_a.json
+```
+
+During the training, the error of the model is tested every **`disp_freq`** batches with **`numb_test`** frames from the last set in the **`systems`** directory on the fly, and the results are output to **`disp_file`**. A typical `disp_file` looks like
+```bash
+# batch l2_tst l2_trn l2_e_tst l2_e_trn l2_f_tst l2_f_trn lr
+ 0 2.67e+01 2.57e+01 2.21e-01 2.22e-01 8.44e-01 8.12e-01 1.0e-03
+ 100 6.14e+00 5.40e+00 3.01e-01 2.99e-01 1.93e-01 1.70e-01 1.0e-03
+ 200 5.02e+00 4.49e+00 1.53e-01 1.53e-01 1.58e-01 1.42e-01 1.0e-03
+ 300 4.36e+00 3.71e+00 7.32e-02 7.27e-02 1.38e-01 1.17e-01 1.0e-03
+ 400 4.04e+00 3.29e+00 3.16e-02 3.22e-02 1.28e-01 1.04e-01 1.0e-03
+```
+The first column displays the number of batches. The second and third columns display the loss function evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The fourth and fifth columns display the RMS energy error (normalized by number of atoms) evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The sixth and seventh columns display the RMS force error (component-wise) evaluated by `numb_test` frames randomly chosen from the test set and that evaluated by the current training batch, respectively. The last column displays the current learning rate.
+
+Checkpoints will be written to files with prefix **`save_ckpt`** every **`save_freq`** batches. If **`restart`** is set to `true`, then the training will start from the checkpoint named **`load_ckpt`**, rather than from scratch.
+
+Several command line options can be passed to `dp train`, which can be checked with
+```bash
+$ dp train --help
+```
+An explanation will be provided
+```
+positional arguments:
+ INPUT the input json database
+
+optional arguments:
+ -h, --help show this help message and exit
+ --init-model INIT_MODEL
+ Initialize a model by the provided checkpoint
+ --restart RESTART Restart the training from the provided checkpoint
+```
+The keys `intra_op_parallelism_threads` and `inter_op_parallelism_threads` are Tensorflow configurations for multithreading, which are explained [here](https://www.tensorflow.org/performance/performance_guide#optimizing_for_cpu). Skipping `-t` and `OMP_NUM_THREADS` leads to the default setting of these keys in the Tensorflow.
+
+**`--init-model model.ckpt`**, for example, initializes the model training with an existing model that is stored in the checkpoint `model.ckpt`, the network architectures should match.
+
+**`--restart model.ckpt`**, continues the training from the checkpoint `model.ckpt`.
+
+On some resources limited machines, one may want to control the number of threads used by DeePMD-kit. This is achieved by three environmental variables: `OMP_NUM_THREADS`, `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS`. `OMP_NUM_THREADS` controls the multithreading of DeePMD-kit implemented operations. `TF_INTRA_OP_PARALLELISM_THREADS` and `TF_INTER_OP_PARALLELISM_THREADS` controls `intra_op_parallelism_threads` and `inter_op_parallelism_threads`, which are Tensorflow configurations for multithreading. An explanation is found [here](https://stackoverflow.com/questions/41233635/meaning-of-inter-op-parallelism-threads-and-intra-op-parallelism-threads).
+
+For example 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:
+```bash
+export OMP_NUM_THREADS=6
+export TF_INTRA_OP_PARALLELISM_THREADS=3
+export TF_INTER_OP_PARALLELISM_THREADS=2
+dp train input.json
+```
+
+## Freeze a model
+
+The trained neural network is extracted from a checkpoint and dumped into a database. 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
+```bash
+$ dp freeze -o graph.pb
+```
+in the folder where the model is trained. The output database is called `graph.pb`.
+
+
+## Test a model
+
+The frozen model can be used in many ways. The most straightforward test can be performed using `dp test`. A typical usage of `dp test` is
+```bash
+dp test -m graph.pb -s /path/to/system -n 30
+```
+where `-m` gives the tested model, `-s` the path to the tested system and `-n` the number of tested frames. Several other command line options can be passed to `dp test`, which can be checked with
+```bash
+$ dp test --help
+```
+An explanation will be provided
+```
+usage: dp test [-h] [-m MODEL] [-s SYSTEM] [-S SET_PREFIX] [-n NUMB_TEST]
+ [-r RAND_SEED] [--shuffle-test] [-d DETAIL_FILE]
+
+optional arguments:
+ -h, --help show this help message and exit
+ -m MODEL, --model MODEL
+ Frozen model file to import
+ -s SYSTEM, --system SYSTEM
+ The system dir
+ -S SET_PREFIX, --set-prefix SET_PREFIX
+ The set prefix
+ -n NUMB_TEST, --numb-test NUMB_TEST
+ The number of data for test
+ -r RAND_SEED, --rand-seed RAND_SEED
+ The random seed
+ --shuffle-test Shuffle test data
+ -d DETAIL_FILE, --detail-file DETAIL_FILE
+ The file containing details of energy force and virial
+ accuracy
+```
+
+## Model inference
+One may use the python interface of DeePMD-kit for model inference, an example is given as follows
+```python
+import deepmd.DeepPot as DP
+import numpy as np
+dp = DP('graph.pb')
+coord = np.array([[1,0,0], [0,0,1.5], [1,0,3]]).reshape([1, -1])
+cell = np.diag(10 * np.ones(3)).reshape([1, -1])
+atype = [1,0,1]
+e, f, v = dp.eval(coord, cell, atype)
+```
+where `e`, `f` and `v` are predicted energy, force and virial of the system, respectively.
+
+
+## Run MD with LAMMPS
+### Include deepmd in the pair style
+Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows
+```bash
+pair_style deepmd graph.pb
+pair_coeff
+```
+where `graph.pb` is the file name of the frozen model. The `pair_coeff` should be left blank. It should be noted that LAMMPS counts atom types starting from 1, therefore, all LAMMPS atom type will be firstly subtracted by 1, and then passed into the DeePMD-kit engine to compute the interactions. [A detailed documentation of this pair style is available.](doc/lammps-pair-style-deepmd.md).
+
+### Long-range interaction
+The reciprocal space part of the long-range interaction can be calculated by LAMMPS command `kspace_style`. To use it with DeePMD-kit, one writes
+```bash
+pair_style deepmd graph.pb
+pair_coeff
+kspace_style pppm 1.0e-5
+kspace_modify gewald 0.45
+```
+Please notice that the DeePMD does nothing to the direct space part of the electrostatic interaction, because this part is assumed to be fitted in the DeePMD model (the direct space cut-off is thus the cut-off of the DeePMD model). The splitting parameter `gewald` is modified by the `kspace_modify` command.
+
+## Run path-integral MD with i-PI
+The i-PI works in a client-server model. The i-PI provides the server for integrating the replica positions of atoms, while the DeePMD-kit provides a client named `dp_ipi` that computes the interactions (including energy, force and virial). The server and client communicates via the Unix domain socket or the Internet socket. The client can be started by
+```bash
+$ dp_ipi water.json
+```
+It is noted that multiple instances of the client is allow for computing, in parallel, the interactions of multiple replica of the path-integral MD.
+
+`water.json` is the parameter file for the client `dp_ipi`, and [an example](./examples/ipi/water.json) is provided:
+```json
+{
+ "verbose": false,
+ "use_unix": true,
+ "port": 31415,
+ "host": "localhost",
+ "graph_file": "graph.pb",
+ "coord_file": "conf.xyz",
+ "atom_type" : {
+ "OW": 0,
+ "HW1": 1,
+ "HW2": 1
+ }
+}
+```
+The option **`use_unix`** is set to `true` to activate the Unix domain socket, otherwise, the Internet socket is used.
+
+The option **`graph_file`** provides the file name of the frozen model.
+
+The `dp_ipi` gets the atom names from an [XYZ file](https://en.wikipedia.org/wiki/XYZ_file_format) provided by **`coord_file`** (meanwhile ignores all coordinates in it), and translates the names to atom types by rules provided by **`atom_type`**.
+
+## Use deep potential with ASE
+
+Deep potential can be set up as a calculator with ASE to obtain potential energies and forces.
+```python
+from ase import Atoms
+from deepmd.calculator import DP
+
+water = Atoms('H2O',
+ positions=[(0.7601, 1.9270, 1),
+ (1.9575, 1, 1),
+ (1., 1., 1.)],
+ cell=[100, 100, 100],
+ calculator=DP(model="frozen_model.pb"))
+print(water.get_potential_energy())
+print(water.get_forces())
+```
+
+Optimization is also available:
+```python
+from ase.optimize import BFGS
+dyn = BFGS(water)
+dyn.run(fmax=1e-6)
+print(water.get_positions())
+```
diff --git a/setup.py b/setup.py
index db86971589..e03d2dcf7f 100644
--- a/setup.py
+++ b/setup.py
@@ -17,7 +17,7 @@
with open(readme_file) as f:
readme = f.read()
-install_requires=['numpy', 'scipy', 'pyyaml']
+install_requires=['numpy', 'scipy', 'pyyaml', 'dargs']
setup_requires=['setuptools_scm', 'scikit-build']
tf_version = os.environ.get('TENSORFLOW_VERSION', '2.3')
diff --git a/source/train/CMakeLists.txt b/source/train/CMakeLists.txt
index 1875d2097c..34363a6826 100644
--- a/source/train/CMakeLists.txt
+++ b/source/train/CMakeLists.txt
@@ -2,7 +2,7 @@
configure_file("RunOptions.py.in" "${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py" @ONLY)
-file(GLOB LIB_PY main.py common.py env.py compat.py calculator.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py EwaldRecp.py DataModifier.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py transform.py)
+file(GLOB LIB_PY main.py common.py env.py compat.py calculator.py Network.py Deep*.py Data.py DataSystem.py Model*.py Descrpt*.py Fitting.py Loss.py LearningRate.py Trainer.py TabInter.py EwaldRecp.py DataModifier.py ${CMAKE_CURRENT_BINARY_DIR}/RunOptions.py transform.py argcheck.py doc.py)
file(GLOB CLS_PY Local.py Slurm.py)
diff --git a/source/train/argcheck.py b/source/train/argcheck.py
new file mode 100644
index 0000000000..89034180ba
--- /dev/null
+++ b/source/train/argcheck.py
@@ -0,0 +1,378 @@
+from dargs import Argument, Variant
+from deepmd.common import activation_fn_dict
+
+def list_to_doc (xx):
+ items = []
+ for ii in xx:
+ if len(items) == 0:
+ items.append(f'"{ii}"')
+ else:
+ items.append(f', "{ii}"')
+ items.append('.')
+ return ''.join(items)
+
+
+def supported_precision() :
+ return list_to_doc(['float64', 'float32', 'float16'])
+
+
+def make_link(content, ref_key) :
+ return f'`{content} <#{ref_key}>`__'
+
+
+def descrpt_local_frame_args ():
+ doc_sel_a = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_a[i]` gives the selected number of type-i neighbors. The full relative coordinates of the neighbors are used by the descriptor.'
+ doc_sel_r = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel_r[i]` gives the selected number of type-i neighbors. Only relative distance of the neighbors are used by the descriptor. sel_a[i] + sel_r[i] is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.'
+ doc_rcut = 'The cut-off radius. The default value is 6.0'
+ doc_axis_rule = 'A list of integers. The length should be 6 times of the number of types. \n\n\
+- axis_rule[i*6+0]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\
+- axis_rule[i*6+1]: type of the atom defining the first axis of type-i atom.\n\n\
+- axis_rule[i*6+2]: index of the axis atom defining the first axis. Note that the neighbors with the same class and type are sorted according to their relative distance.\n\n\
+- axis_rule[i*6+3]: class of the atom defining the first axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.\n\n\
+- axis_rule[i*6+4]: type of the atom defining the second axis of type-i atom.\n\n\
+- axis_rule[i*6+5]: class of the atom defining the second axis of type-i atom. 0 for neighbors with full coordinates and 1 for neighbors only with relative distance.'
+
+ return [
+ Argument("sel_a", list, optional = False, doc = doc_sel_a),
+ Argument("sel_r", list, optional = False, doc = doc_sel_r),
+ Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut),
+ Argument("axis_rule", list, optional = False, doc = doc_axis_rule)
+ ]
+
+
+def descrpt_se_a_args():
+ doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.'
+ doc_rcut = 'The cut-off radius.'
+ doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`'
+ doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.'
+ doc_axis_neuron = 'Size of the submatrix of G (embedding matrix).'
+ doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(activation_fn_dict.keys())}'
+ doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
+ doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets'
+ doc_precision = f'The precision of the embedding net parameters, supported options are {supported_precision()}'
+ doc_trainable = 'If the parameters in the embedding net is trainable'
+ doc_seed = 'Random seed for parameter initialization'
+ doc_exclude_types = 'The Excluded types'
+ doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used'
+
+ return [
+ Argument("sel", list, optional = False, doc = doc_sel),
+ Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut),
+ Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth),
+ Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron),
+ Argument("axis_neuron", int, optional = True, default = 4, doc = doc_axis_neuron),
+ Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function),
+ Argument("resnet_dt", bool, optional = True, default = False, doc = doc_resnet_dt),
+ Argument("type_one_side", bool, optional = True, default = False, doc = doc_type_one_side),
+ Argument("precision", str, optional = True, default = "float64", doc = doc_precision),
+ Argument("trainable", bool, optional = True, default = True, doc = doc_trainable),
+ Argument("seed", [int,None], optional = True, doc = doc_seed),
+ Argument("exclude_types", list, optional = True, default = [], doc = doc_exclude_types),
+ Argument("set_davg_zero", bool, optional = True, default = False, doc = doc_set_davg_zero)
+ ]
+
+
+def descrpt_se_r_args():
+ doc_sel = 'A list of integers. The length of the list should be the same as the number of atom types in the system. `sel[i]` gives the selected number of type-i neighbors. `sel[i]` is recommended to be larger than the maximally possible number of type-i neighbors in the cut-off radius.'
+ doc_rcut = 'The cut-off radius.'
+ doc_rcut_smth = 'Where to start smoothing. For example the 1/r term is smoothed from `rcut` to `rcut_smth`'
+ doc_neuron = 'Number of neurons in each hidden layers of the embedding net. When two layers are of the same size or one layer is twice as large as the previous layer, a skip connection is built.'
+ doc_activation_function = f'The activation function in the embedding net. Supported activation functions are {list_to_doc(activation_fn_dict.keys())}'
+ doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
+ doc_type_one_side = 'Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets'
+ doc_precision = f'The precision of the embedding net parameters, supported options are {supported_precision()}'
+ doc_trainable = 'If the parameters in the embedding net is trainable'
+ doc_seed = 'Random seed for parameter initialization'
+ doc_exclude_types = 'The Excluded types'
+ doc_set_davg_zero = 'Set the normalization average to zero. This option should be set when `atom_ener` in the energy fitting is used'
+
+ return [
+ Argument("sel", list, optional = False, doc = doc_sel),
+ Argument("rcut", float, optional = True, default = 6.0, doc = doc_rcut),
+ Argument("rcut_smth", float, optional = True, default = 0.5, doc = doc_rcut_smth),
+ Argument("neuron", list, optional = True, default = [10,20,40], doc = doc_neuron),
+ Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function),
+ Argument("resnet_dt", bool, optional = True, default = False, doc = doc_resnet_dt),
+ Argument("type_one_side", bool, optional = True, default = False, doc = doc_type_one_side),
+ Argument("precision", str, optional = True, default = "float64", doc = doc_precision),
+ Argument("trainable", bool, optional = True, default = True, doc = doc_trainable),
+ Argument("seed", [int,None], optional = True, doc = doc_seed),
+ Argument("exclude_types", list, optional = True, default = [], doc = doc_exclude_types),
+ Argument("set_davg_zero", bool, optional = True, default = False, doc = doc_set_davg_zero)
+ ]
+
+
+def descrpt_se_ar_args():
+ link = make_link('se_a', 'model/descriptor[se_a]')
+ doc_a = f'The parameters of descriptor {link}'
+ link = make_link('se_r', 'model/descriptor[se_r]')
+ doc_r = f'The parameters of descriptor {link}'
+
+ return [
+ Argument("a", dict, optional = False, doc = doc_a),
+ Argument("r", dict, optional = False, doc = doc_r),
+ ]
+
+
+def descrpt_variant_type_args():
+ doc_descrpt_type = 'The type of the descritpor. Valid types are `loc_frame`, `se_a`, `se_r` and `se_ar`. \n\n\
+- `loc_frame`: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.\n\n\
+- `se_a`: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.\n\n\
+- `se_r`: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.\n\n\
+- `se_ar`: A hybrid of `se_a` and `se_r`. Typically `se_a` has a smaller cut-off while the `se_r` has a larger cut-off.'
+
+ return Variant("type", [
+ Argument("loc_frame", dict, descrpt_local_frame_args()),
+ Argument("se_a", dict, descrpt_se_a_args()),
+ Argument("se_r", dict, descrpt_se_r_args()),
+ Argument("se_ar", dict, descrpt_se_ar_args())
+ ], doc = doc_descrpt_type)
+
+
+def fitting_ener():
+ doc_numb_fparam = 'The dimension of the frame parameter. If set to >0, file `fparam.npy` should be included to provided the input fparams.'
+ doc_numb_aparam = 'The dimension of the atomic parameter. If set to >0, file `aparam.npy` should be included to provided the input aparams.'
+ 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())}'
+ doc_precision = f'The precision of the fitting net parameters, supported options are {supported_precision()}'
+ doc_resnet_dt = 'Whether to use a "Timestep" in the skip connection'
+ doc_trainable = 'Whether the parameters in the fitting net are trainable. This option can be\n\n\
+- bool: True if all parameters of the fitting net are trainable, False otherwise.\n\n\
+- list of bool: Specifies if each layer is trainable. Since the fitting net is composed by hidden layers followed by a output layer, the length of tihs list should be equal to len(`neuron`)+1.'
+ doc_rcond = 'The condition number used to determine the inital energy shift for each type of atoms.'
+ doc_seed = 'Random seed for parameter initialization of the fitting net'
+ doc_atom_ener = 'Specify the atomic energy in vacuum for each type'
+
+ return [
+ Argument("numb_fparam", int, optional = True, default = 0, doc = doc_numb_fparam),
+ Argument("numb_aparam", int, optional = True, default = 0, doc = doc_numb_aparam),
+ Argument("neuron", list, optional = True, default = [120,120,120], doc = doc_neuron),
+ Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function),
+ Argument("precision", str, optional = True, default = 'float64', doc = doc_precision),
+ Argument("resnet_dt", bool, optional = True, default = True, doc = doc_resnet_dt),
+ Argument("trainable", [list,bool], optional = True, default = True, doc = doc_trainable),
+ Argument("rcond", float, optional = True, default = 1e-3, doc = doc_rcond),
+ Argument("seed", [int,None], optional = True, doc = doc_seed),
+ Argument("atom_ener", list, optional = True, default = [], doc = doc_atom_ener)
+ ]
+
+
+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())}'
+ 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 {supported_precision()}'
+ doc_scale = 'The output of the fitting net (polarizability matrix) will be scaled by ``scale``'
+ doc_diag_shift = 'The diagonal part of the polarizability matrix will be shifted by ``diag_shift``. The shift operation is carried out after ``scale``.'
+ 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'
+
+ return [
+ Argument("neuron", list, optional = True, default = [120,120,120], doc = doc_neuron),
+ Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function),
+ Argument("resnet_dt", bool, optional = True, default = True, doc = doc_resnet_dt),
+ Argument("precision", str, optional = True, default = 'float64', doc = doc_precision),
+ Argument("fit_diag", bool, optional = True, default = True, doc = doc_fit_diag),
+ Argument("scale", [list,float], optional = True, default = 1.0, doc = doc_scale),
+ Argument("diag_shift", [list,float], optional = True, default = 0.0, doc = doc_diag_shift),
+ Argument("sel_type", [list,int,None], optional = True, doc = doc_sel_type),
+ Argument("seed", [int,None], optional = True, doc = doc_seed)
+ ]
+
+
+def fitting_global_polar():
+ return fitting_polar()
+
+
+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())}'
+ 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 {supported_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'
+ return [
+ Argument("neuron", list, optional = True, default = [120,120,120], doc = doc_neuron),
+ Argument("activation_function", str, optional = True, default = 'tanh', doc = doc_activation_function),
+ Argument("resnet_dt", bool, optional = True, default = True, doc = doc_resnet_dt),
+ Argument("precision", str, optional = True, default = 'float64', doc = doc_precision),
+ Argument("sel_type", [list,int,None], optional = True, doc = doc_sel_type),
+ Argument("seed", [int,None], optional = True, doc = doc_seed)
+ ]
+
+
+def fitting_variant_type_args():
+ doc_descrpt_type = 'The type of the fitting. Valid types are `ener`, `dipole`, `polar` and `global_polar`. \n\n\
+- `ener`: Fit an energy model (potential energy surface).\n\n\
+- `dipole`: Fit an atomic dipole model. Atomic dipole labels for all the selected atoms (see `sel_type`) should be provided by `dipole.npy` in each data system. The file has number of frames lines and 3 times of number of selected atoms columns.\n\n\
+- `polar`: Fit an atomic polarizability model. Atomic polarizability labels for all the selected atoms (see `sel_type`) should be provided by `polarizability.npy` in each data system. The file has number of frames lines and 9 times of number of selected atoms columns.\n\n\
+- `global_polar`: Fit a polarizability model. Polarizability labels should be provided by `polarizability.npy` in each data system. The file has number of frames lines and 9 columns.'
+
+ return Variant("type", [Argument("ener", dict, fitting_ener()),
+ Argument("dipole", dict, fitting_dipole()),
+ Argument("polar", dict, fitting_polar()),
+ Argument("global_polar", dict, fitting_global_polar())],
+ optional = True,
+ default_tag = 'ener',
+ doc = doc_descrpt_type)
+
+
+def model_args ():
+ doc_type_map = 'A list of strings. Give the name to each type of atoms.'
+ doc_data_stat_nbatch = 'The model determines the normalization from the statistics of the data. This key specifies the number of `frames` in each `system` used for statistics.'
+ doc_descrpt = 'The descriptor of atomic environment.'
+ doc_fitting = 'The fitting of physical properties.'
+ ca = Argument("model", dict,
+ [Argument("type_map", list, optional = True, doc = doc_type_map),
+ Argument("data_stat_nbatch", int, optional = True, default = 10, doc = doc_data_stat_nbatch),
+ Argument("descriptor", dict, [], [descrpt_variant_type_args()], doc = doc_descrpt),
+ Argument("fitting_net", dict, [], [fitting_variant_type_args()], doc = doc_fitting)
+ ])
+ # print(ca.gen_doc())
+ return ca
+
+
+def learning_rate_args():
+ doc_start_lr = 'The learning rate the start of the training.'
+ doc_stop_lr = 'The desired learning rate at the end of the training.'
+ doc_decay_steps = 'The learning rate is decaying every this number of training steps.'
+
+ args = [
+ Argument("start_lr", float, optional = True, default = 1e-3, doc = doc_start_lr),
+ Argument("stop_lr", float, optional = True, default = 1e-8, doc = doc_stop_lr),
+ Argument("decay_steps", int, optional = True, default = 5000, doc = doc_decay_steps)
+ ]
+
+ doc_lr = "The learning rate options"
+ return Argument("learning_rate", dict, args, [], doc = doc_lr)
+
+
+def start_pref(item):
+ return f'The prefactor of {item} loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the {item} label should be provided by file {item}.npy in each data system. If both start_pref_{item} and limit_pref_{item} are set to 0, then the {item} will be ignored.'
+
+def limit_pref(item):
+ return f'The prefactor of {item} loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.'
+
+def loss_ener():
+ doc_start_pref_e = start_pref('energy')
+ doc_limit_pref_e = limit_pref('energy')
+ doc_start_pref_f = start_pref('force')
+ doc_limit_pref_f = limit_pref('force')
+ doc_start_pref_v = start_pref('virial')
+ doc_limit_pref_v = limit_pref('virial')
+ doc_start_pref_ae = start_pref('atom_ener')
+ doc_start_pref_ae = limit_pref('atom_ener')
+ doc_relative_f = 'If provided, relative force error will be used in the loss. The difference of force will be normalized by the magnitude of the force in the label with a shift given by `relative_f`, i.e. DF_i / ( || F || + relative_f ) with DF denoting the difference between prediction and label and || F || denoting the L2 norm of the label.'
+ return [
+ Argument("start_pref_e", [float,int], optional = True, default = 0.02, doc = doc_start_pref_e),
+ Argument("limit_pref_e", [float,int], optional = True, default = 1.00, doc = doc_limit_pref_e),
+ Argument("start_pref_f", [float,int], optional = True, default = 1000, doc = doc_start_pref_f),
+ Argument("limit_pref_f", [float,int], optional = True, default = 1.00, doc = doc_limit_pref_f),
+ Argument("start_pref_v", [float,int], optional = True, default = 0.00, doc = doc_start_pref_v),
+ Argument("limit_pref_v", [float,int], optional = True, default = 0.00, doc = doc_limit_pref_v),
+ Argument("start_pref_ae", [float,int], optional = True, default = 0.00, doc = doc_start_pref_v),
+ Argument("limit_pref_ae", [float,int], optional = True, default = 0.00, doc = doc_limit_pref_v),
+ Argument("relative_f", [float,None], optional = True, doc = doc_relative_f)
+ ]
+
+
+def loss_variant_type_args():
+ doc_loss = 'The type of the loss. For fitting type `ener`, the loss type should be set to `ener` or left unset. For tensorial fitting types `dipole`, `polar` and `global_polar`, the type should be left unset.\n\.'
+
+ return Variant("type",
+ [Argument("ener", dict, loss_ener())],
+ optional = True,
+ default_tag = 'ener',
+ doc = doc_loss)
+
+def loss_args():
+ doc_loss = 'The definition of loss function. The type of the loss depends on the type of the fitting. For fitting type `ener`, the prefactors before energy, force, virial and atomic energy losses may be provided. For fitting type `dipole`, `polar` and `global_polar`, the loss may be an empty `dict` or unset.'
+ ca = Argument('loss', dict, [],
+ [loss_variant_type_args()],
+ doc = doc_loss)
+ return ca
+
+def training_args():
+ doc_systems = 'The data systems. This key can be provided with a listthat specifies the systems, or be provided with a string by which the prefix of all systems are given and the list of the systems is automatically generated.'
+ doc_set_prefix = 'The prefix of the sets in the systems.'
+ doc_stop_batch = 'Number of training batch. Each training uses one batch of data.'
+ doc_batch_size = 'This key can be \n\n\
+- list: the length of which is the same as the `systems`. The batch size of each system is given by the elements of the list.\n\n\
+- int: all `systems` uses the same batch size.\n\n\
+- string "auto": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than 32.\n\n\
+- string "auto:N": automatically determines the batch size os that the batch_size times the number of atoms in the system is no less than N.'
+ doc_seed = 'The random seed for training.'
+ doc_disp_file = 'The file for printing learning curve.'
+ doc_disp_freq = 'The frequency of printing learning curve.'
+ doc_numb_test = 'Number of frames used for the test during training.'
+ doc_save_freq = 'The frequency of saving check point.'
+ doc_save_ckpt = 'The file name of saving check point.'
+ doc_disp_training = 'Displaying verbose information during training.'
+ doc_time_training = 'Timing durining training.'
+ doc_profiling = 'Profiling during training.'
+ doc_profiling_file = 'Output file for profiling.'
+
+ args = [
+ Argument("systems", [list,str], optional = False, doc = doc_systems),
+ Argument("set_prefix", str, optional = True, default = 'set', doc = doc_set_prefix),
+ Argument("stop_batch", int, optional = False, doc = doc_stop_batch),
+ Argument("batch_size", [list,int,str], optional = True, default = 'auto', doc = doc_batch_size),
+ Argument("seed", [int,None], optional = True, doc = doc_seed),
+ Argument("disp_file", str, optional = True, default = 'lcueve.out', doc = doc_disp_file),
+ Argument("disp_freq", int, optional = True, default = 1000, doc = doc_disp_freq),
+ Argument("numb_test", int, optional = True, default = 1, doc = doc_numb_test),
+ Argument("save_freq", int, optional = True, default = 1000, doc = doc_save_freq),
+ Argument("save_ckpt", str, optional = True, default = 'model.ckpt', doc = doc_save_ckpt),
+ Argument("disp_training", bool, optional = True, default = True, doc = doc_disp_training),
+ Argument("time_training", bool, optional = True, default = True, doc = doc_time_training),
+ Argument("profiling", bool, optional = True, default = False, doc = doc_profiling),
+ Argument("profiling_file", str, optional = True, default = 'timeline.json', doc = doc_profiling_file)
+ ]
+
+ doc_training = 'The training options'
+ return Argument("training", dict, args, [], doc = doc_training)
+
+
+def make_index(keys):
+ ret = []
+ for ii in keys:
+ ret.append(make_link(ii, ii))
+ return ', '.join(ret)
+
+
+def gen_doc(**kwargs):
+ ma = model_args()
+ lra = learning_rate_args()
+ la = loss_args()
+ ta = training_args()
+ ptr = []
+ ptr.append(ma.gen_doc(**kwargs))
+ ptr.append(la.gen_doc(**kwargs))
+ ptr.append(lra.gen_doc(**kwargs))
+ ptr.append(ta.gen_doc(**kwargs))
+
+ key_words = []
+ for ii in "\n\n".join(ptr).split('\n'):
+ if 'argument path' in ii:
+ key_words.append(ii.split(':')[1].replace('`','').strip())
+ #ptr.insert(0, make_index(key_words))
+
+ return "\n\n".join(ptr)
+
+def normalize(data):
+ ma = model_args()
+ lra = learning_rate_args()
+ la = loss_args()
+ ta = training_args()
+
+ base = Argument("base", dict, [ma, lra, la, ta])
+ data = base.normalize_value(data, trim_pattern = "_*")
+ base.check_value(data)
+
+ return data
+
+
+if __name__ == '__main__':
+ gen_doc()
+
diff --git a/source/train/common.py b/source/train/common.py
index c420cab0d3..c14dcca6ee 100644
--- a/source/train/common.py
+++ b/source/train/common.py
@@ -111,6 +111,8 @@ def add (self,
def _add_single(self, key, data) :
vtype = type(data)
+ if data is None:
+ return data
if not(vtype in self.arg_dict[key]['types']) :
# ! altered by Marián Rynik
# try the type convertion to one of the types
diff --git a/source/train/doc.py b/source/train/doc.py
new file mode 100644
index 0000000000..599f64e247
--- /dev/null
+++ b/source/train/doc.py
@@ -0,0 +1,5 @@
+from deepmd.argcheck import gen_doc
+
+def doc_train_input(args):
+ doc_str = gen_doc(make_anchor=True)
+ print(doc_str)
diff --git a/source/train/main.py b/source/train/main.py
index 56100ec54e..8cb039bc0e 100644
--- a/source/train/main.py
+++ b/source/train/main.py
@@ -5,6 +5,7 @@
from .config import config
from .test import test
from .transform import transform
+from .doc import doc_train_input
def main () :
parser = argparse.ArgumentParser(
@@ -32,6 +33,9 @@ def main () :
parser_train.add_argument('--restart', type = str,
help=
'Restart the training from the provided checkpoint.')
+ parser_train.add_argument('-o','--output', type = str, default = 'out.json',
+ help=
+ 'The output file of the parameters used in training.')
parser_frz = subparsers.add_parser('freeze', help='freeze the model')
parser_frz.add_argument("-d", "--folder", type=str, default = ".",
@@ -57,6 +61,9 @@ def main () :
parser_tst.add_argument("-d", "--detail-file", type=str,
help="The file containing details of energy force and virial accuracy")
+ parser_train = subparsers.add_parser('doc-train-input',
+ help='print the documentation (in rst format) of input training parameters.')
+
args = parser.parse_args()
if args.command is None :
@@ -72,5 +79,7 @@ def main () :
test(args)
elif args.command == 'transform' :
transform(args)
+ elif args.command == 'doc-train-input' :
+ doc_train_input(args)
else :
raise RuntimeError('unknown command ' + args.command)
diff --git a/source/train/train.py b/source/train/train.py
index 3e7ba2955b..e7978361b2 100755
--- a/source/train/train.py
+++ b/source/train/train.py
@@ -3,6 +3,7 @@
import os
import sys
import time
+import json
import numpy as np
from deepmd.env import tf
from deepmd.compat import convert_input_v0_v1
@@ -11,6 +12,7 @@
from deepmd.Trainer import NNPTrainer
from deepmd.common import data_requirement, expand_sys_str, j_loader
from deepmd.DataModifier import DipoleChargeModifier
+from deepmd.argcheck import normalize
def create_done_queue(cluster_spec, task_index):
with tf.device("/job:ps/task:%d" % (task_index)):
@@ -54,6 +56,11 @@ def train (args) :
jdata = convert_input_v0_v1(jdata,
warning = True,
dump = 'input_v1_compat.json')
+
+ jdata = normalize(jdata)
+ with open(args.output, 'w') as fp:
+ json.dump(jdata, fp, indent=4)
+
# run options
with_distrib = False
if 'with_distrib' in jdata: