DeePMD-kit’s documentation

DeePMD-kit is a package written in Python/C++, designed to minimize the effort required to build deep learning-based models of interatomic potential energy and force field and to perform molecular dynamics (MD). This brings new hopes to addressing the accuracy-versus-efficiency dilemma in molecular simulations. Applications of DeePMD-kit span from finite molecules to extended systems and from metallic systems to chemically bonded systems.

Important

The project DeePMD-kit is licensed under GNU LGPLv3.0. If you use this code in any future publications, please cite this using Han Wang, Linfeng Zhang, Jiequn Han, and Weinan E. “DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics.” Computer Physics Communications 228 (2018): 178-184.

Getting Started

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

Easy install

There are 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 train models or run LAMMPS in parallel.

Note

Note: The off-line packages and conda packages require the GNU C Library 2.17 or above. The GPU version requires compatible NVIDIA driver to be installed in advance. It is possible to force conda to override detection when installation, but these requirements are still necessary during runtime.

Install off-line packages
Install with conda
Install with docker
Install Python interface with pip

Install off-line packages

Both CPU and GPU version offline packages are available in the Releases page.

Some packages are splited into two files due to size limit of GitHub. One may merge them into one after downloading:

cat deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh.0 deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh.1 > deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh

One may enable the environment using

conda activate /path/to/deepmd-kit

Install with conda

DeePMD-kit is available with conda. Install Anaconda or Miniconda first.

Official channel

One may create an environment that contains the CPU version of DeePMD-kit and LAMMPS:

conda create -n deepmd deepmd-kit=*=*cpu libdeepmd=*=*cpu lammps -c https://conda.deepmodeling.com -c defaults

Or one may want to create a GPU environment containing CUDA Toolkit:

conda create -n deepmd deepmd-kit=*=*gpu libdeepmd=*=*gpu lammps cudatoolkit=11.6 horovod -c https://conda.deepmodeling.com -c defaults

One could change the CUDA Toolkit version from 10.2 or 11.6.

One may specify the DeePMD-kit version such as 2.1.1 using

conda create -n deepmd deepmd-kit=2.1.1=*cpu libdeepmd=2.1.1=*cpu lammps horovod -c https://conda.deepmodeling.com -c defaults

One may enable the environment using

conda activate deepmd

conda-forge channel

DeePMD-kit is also available on the conda-forge channel:

conda create -n deepmd deepmd-kit lammps -c conda-forge

The supported platform includes Linux x86-64, macOS x86-64, and macOS arm64. Read conda-forge FAQ to learn how to install CUDA-enabled packages.

Install with docker

A docker for installing the DeePMD-kit is available here.

To pull the CPU version:

docker pull ghcr.io/deepmodeling/deepmd-kit:2.1.1_cpu

To pull the GPU version:

docker pull ghcr.io/deepmodeling/deepmd-kit:2.1.1_cuda11.6_gpu

To pull the ROCm version:

docker pull deepmodeling/dpmdkit-rocm:dp2.0.3-rocm4.5.2-tf2.6-lmp29Sep2021

Install Python interface with pip

If you have no existing TensorFlow installed, you can use pip to install the pre-built package of the Python interface with CUDA 11 supported:

pip install deepmd-kit[gpu,cu11]

cu11 is required only when CUDA Toolkit and cuDNN were not installed.

Or install the CPU version without CUDA supported:

pip install deepmd-kit[cpu]

The LAMMPS module and the i-Pi driver are only provided on Linux and macOS. To install LAMMPS and/or i-Pi, add lmp and/or ipi to extras:

pip install deepmd-kit[gpu,cu11,lmp,ipi]

MPICH is required for parallel running.

It is suggested to install the package into an isolated environment. The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64. A specific version of TensorFlow which is compatible with DeePMD-kit will be also installed.

Warning

If your platform is not supported, or want to build against the installed TensorFlow, or want to enable ROCM support, please build from source.

Prepare data with dpdata

One can use a convenient tool dpdata to convert data directly from the output of first principle packages to the DeePMD-kit format.

To install one can execute

pip install dpdata

An example of converting data VASP data in OUTCAR format to DeePMD-kit data can be found at

$deepmd_source_dir/examples/data_conv

Switch to that directory, then one can convert data by using the following python script

import dpdata

dsys = dpdata.LabeledSystem("OUTCAR")
dsys.to("deepmd/npy", "deepmd_data", set_size=dsys.get_nframes())

get_nframes() method gets the number of frames in the OUTCAR, and the argument set_size enforces that the set size is equal to the number of frames in the system, viz. only one set is created in the system.

The data in DeePMD-kit format is stored in the folder deepmd_data.

A list of all supported data format and more nice features of dpdata can be found on the official website.

Train a model

Several examples of training can be found in the examples directory:

$ cd $deepmd_source_dir/examples/water/se_e2_a/

After switching to that directory, the training can be invoked by

$ dp train input.json

where input.json is the name of the input script.

By default, the verbosity level of the DeePMD-kit is INFO, one may see a lot of important information on the code and environment showing on the screen. Among them two pieces of information regarding data systems are worth special notice.

DEEPMD INFO    ---Summary of DataSystem: training     -----------------------------------------------
DEEPMD INFO    found 3 system(s):
DEEPMD INFO                                        system  natoms  bch_sz   n_bch   prob  pbc
DEEPMD INFO                         ../data_water/data_0/     192       1      80  0.250    T
DEEPMD INFO                         ../data_water/data_1/     192       1     160  0.500    T
DEEPMD INFO                         ../data_water/data_2/     192       1      80  0.250    T
DEEPMD INFO    --------------------------------------------------------------------------------------
DEEPMD INFO    ---Summary of DataSystem: validation   -----------------------------------------------
DEEPMD INFO    found 1 system(s):
DEEPMD INFO                                        system  natoms  bch_sz   n_bch   prob  pbc
DEEPMD INFO                          ../data_water/data_3     192       1      80  1.000    T
DEEPMD INFO    --------------------------------------------------------------------------------------

The DeePMD-kit prints detailed information on the training and validation data sets. The data sets are defined by training_data and validation_data defined in the training section of the input script. The training data set is composed of three data systems, while the validation data set is composed by one data system. The number of atoms, batch size, the number of batches in the system and the probability of using the system are all shown on the screen. The last column presents if the periodic boundary condition is assumed for the system.

During the training, the error of the model is tested every disp_freq training steps with the batch used to train the model and with numb_btch batches from the validating data. The training error and validation error are printed correspondingly in the file disp_file (default is lcurve.out). The batch size can be set in the input script by the key batch_size in the corresponding sections for the training and validation data set. An example of the output

#  step      rmse_val    rmse_trn    rmse_e_val  rmse_e_trn    rmse_f_val  rmse_f_trn         lr
    3.33e+01    3.41e+01      1.03e+01    1.03e+01      8.39e-01    8.72e-01    1.0e-03
    2.57e+01    2.56e+01      1.87e+00    1.88e+00      8.03e-01    8.02e-01    1.0e-03
    2.45e+01    2.56e+01      2.26e-01    2.21e-01      7.73e-01    8.10e-01    1.0e-03
    1.62e+01    1.66e+01      5.01e-02    4.46e-02      5.11e-01    5.26e-01    1.0e-03
    1.36e+01    1.32e+01      1.07e-02    2.07e-03      4.29e-01    4.19e-01    1.0e-03
    1.07e+01    1.05e+01      2.45e-03    4.11e-03      3.38e-01    3.31e-01    1.0e-03

The file contains 8 columns, from left to right, which are the training step, the validation loss, training loss, root mean square (RMS) validation error of energy, RMS training error of energy, RMS validation error of force, RMS training error of force and the learning rate. The RMS error (RMSE) of the energy is normalized by the number of atoms in the system. One can visualize this file with a simple Python script:

import numpy as np
import matplotlib.pyplot as plt

data = np.genfromtxt("lcurve.out", names=True)
for name in data.dtype.names[1:-1]:
    plt.plot(data['step'], data[name], label=name)
plt.legend()
plt.xlabel('Step')
plt.ylabel('Loss')
plt.xscale('symlog')
plt.yscale('log')
plt.grid()
plt.show()

Checkpoints will be written to files with the prefix save_ckpt every save_freq training steps.

Warning

It is warned that the example water data (in folder examples/water/data) is of very limited amount, is provided only for testing purposes, and should not be used to train a production model.

Freeze a model

The trained neural network is extracted from a checkpoint and dumped into a protobuf(.pb) file. This process is called “freezing” a model. The idea and part of our code are from Morgan. To freeze a model, typically one does

$ dp freeze -o graph.pb

in the folder where the model is trained. The output model is called graph.pb.

In multi-task mode:

This process will in default output several models, each of which contains the common descriptor and one of the user-defined fitting nets in fitting_net_dict, let’s name it fitting_key, together frozen in graph_{fitting_key}.pb. Those frozen models are exactly the same as single-task output with fitting net fitting_key.
If you add --united-model option in this situation, the total multi-task model will be frozen into one unit graph.pb, which is mainly for multi-task initialization and can not be used directly for inference.

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

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

$ 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 prefix to files where details of energy, force and virial accuracy/accuracy per atom will be written
  -a, --atomic          Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)

Run MD with LAMMPS

Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows

pair_style     deepmd graph.pb
pair_coeff     * * O H

where graph.pb is the file name of the frozen model. pair_coeff maps atom names (O H) with LAMMPS atom types (integers from 1 to Ntypes, i.e. 1 2).

Installation

Easy install

There are 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 train models or run LAMMPS in parallel.

Note

Note: The off-line packages and conda packages require the GNU C Library 2.17 or above. The GPU version requires compatible NVIDIA driver to be installed in advance. It is possible to force conda to override detection when installation, but these requirements are still necessary during runtime.

Install off-line packages
Install with conda
Install with docker
Install Python interface with pip

Install off-line packages

Both CPU and GPU version offline packages are available in the Releases page.

Some packages are splited into two files due to size limit of GitHub. One may merge them into one after downloading:

cat deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh.0 deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh.1 > deepmd-kit-2.1.1-cuda11.6_gpu-Linux-x86_64.sh

One may enable the environment using

conda activate /path/to/deepmd-kit

Install with conda

DeePMD-kit is available with conda. Install Anaconda or Miniconda first.

Official channel

One may create an environment that contains the CPU version of DeePMD-kit and LAMMPS:

conda create -n deepmd deepmd-kit=*=*cpu libdeepmd=*=*cpu lammps -c https://conda.deepmodeling.com -c defaults

Or one may want to create a GPU environment containing CUDA Toolkit:

conda create -n deepmd deepmd-kit=*=*gpu libdeepmd=*=*gpu lammps cudatoolkit=11.6 horovod -c https://conda.deepmodeling.com -c defaults

One could change the CUDA Toolkit version from 10.2 or 11.6.

One may specify the DeePMD-kit version such as 2.1.1 using

conda create -n deepmd deepmd-kit=2.1.1=*cpu libdeepmd=2.1.1=*cpu lammps horovod -c https://conda.deepmodeling.com -c defaults

One may enable the environment using

conda activate deepmd

conda-forge channel

DeePMD-kit is also available on the conda-forge channel:

conda create -n deepmd deepmd-kit lammps -c conda-forge

The supported platform includes Linux x86-64, macOS x86-64, and macOS arm64. Read conda-forge FAQ to learn how to install CUDA-enabled packages.

Install with docker

A docker for installing the DeePMD-kit is available here.

To pull the CPU version:

docker pull ghcr.io/deepmodeling/deepmd-kit:2.1.1_cpu

To pull the GPU version:

docker pull ghcr.io/deepmodeling/deepmd-kit:2.1.1_cuda11.6_gpu

To pull the ROCm version:

docker pull deepmodeling/dpmdkit-rocm:dp2.0.3-rocm4.5.2-tf2.6-lmp29Sep2021

Install Python interface with pip

If you have no existing TensorFlow installed, you can use pip to install the pre-built package of the Python interface with CUDA 11 supported:

pip install deepmd-kit[gpu,cu11]

cu11 is required only when CUDA Toolkit and cuDNN were not installed.

Or install the CPU version without CUDA supported:

pip install deepmd-kit[cpu]

The LAMMPS module and the i-Pi driver are only provided on Linux and macOS. To install LAMMPS and/or i-Pi, add lmp and/or ipi to extras:

pip install deepmd-kit[gpu,cu11,lmp,ipi]

MPICH is required for parallel running.

It is suggested to install the package into an isolated environment. The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28 or above, macOS x86-64, and Windows x86-64. A specific version of TensorFlow which is compatible with DeePMD-kit will be also installed.

Warning

If your platform is not supported, or want to build against the installed TensorFlow, or want to enable ROCM support, please build from source.

Install from source code

Please follow our GitHub webpage to download the latest released version and development version.

Or get the DeePMD-kit source code by git clone

cd /some/workspace
git clone --recursive https://github.com/deepmodeling/deepmd-kit.git deepmd-kit

The --recursive option clones all submodules needed by DeePMD-kit.

For convenience, you may want to record the location of the source to a variable, saying deepmd_source_dir by

cd deepmd-kit
deepmd_source_dir=`pwd`

Install the python interface

Install Tensorflow’s python interface

First, check the python version on your machine

python --version

We follow the virtual environment approach to install TensorFlow’s Python interface. The full instruction can be found on the official TensorFlow website. TensorFlow 1.8 or later is supported. Now we assume that the Python interface will be installed to the virtual environment directory $tensorflow_venv

virtualenv -p python3 $tensorflow_venv
source $tensorflow_venv/bin/activate
pip install --upgrade pip
pip install --upgrade tensorflow

It is important that every time a new shell is started and one wants to use DeePMD-kit, the virtual environment should be activated by

source $tensorflow_venv/bin/activate

if one wants to skip out of the virtual environment, he/she can do

deactivate

If one has multiple python interpreters named something like python3.x, it can be specified by, for example

virtualenv -p python3.8 $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

pip install --upgrade tensorflow-cpu

To verify the installation, run

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.

One can also build the TensorFlow Python interface from source for custom hardware optimization, such as CUDA, ROCM, or OneDNN support.

Install the DeePMD-kit’s python interface

Check the compiler version on your machine

gcc --version

The compiler GCC 4.8 or later is supported in the DeePMD-kit. Note that TensorFlow may have specific requirements for the compiler version. It is recommended to use the same compiler version as TensorFlow, which can be printed by python -c "import tensorflow;print(tensorflow.version.COMPILER_VERSION)".

Execute

cd $deepmd_source_dir
pip install .

One may set the following environment variables before executing pip:

Environment variables	Allowed value	Default value	Usage
DP_VARIANT	`cpu`, `cuda`, `rocm`	`cpu`	Build CPU variant or GPU variant with CUDA or ROCM support.
CUDA_TOOLKIT_ROOT_DIR	Path	Detected automatically	The path to the CUDA toolkit directory. CUDA 7.0 or later is supported. NVCC is required.
ROCM_ROOT	Path	Detected automatically	The path to the ROCM toolkit directory.
TENSORFLOW_ROOT	Path	Detected automatically	The path to TensorFlow Python library. By default the installer only finds TensorFlow under user site-package directory (`site.getusersitepackages()`) or system site-package directory (`sysconfig.get_path("purelib")`) due to limitation of PEP-517. If not found, the latest TensorFlow (or the environment variable `TENSORFLOW_VERSION` if given) from PyPI will be built against.
DP_ENABLE_NATIVE_OPTIMIZATION	0, 1	0	Enable compilation optimization for the native machine’s CPU type. Do not enable it if generated code will run on different CPUs.

To test the installation, one should first jump out of the source directory

cd /some/other/workspace

then execute

dp -h

It will print the help information like

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 horovod and mpi4py

Horovod and mpi4py are used for parallel training. For better performance on GPU, please follow the tuning steps in Horovod on GPU.

# With GPU, prefer NCCL as a communicator.
HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 HOROVOD_GPU_OPERATIONS=NCCL HOROVOD_NCCL_HOME=/path/to/nccl pip install horovod mpi4py

If your work in a CPU environment, please prepare runtime as below:

# By default, MPI is used as communicator.
HOROVOD_WITHOUT_GLOO=1 HOROVOD_WITH_TENSORFLOW=1 pip install horovod mpi4py

To ensure Horovod has been built with proper framework support enabled, one can invoke the horovodrun --check-build command, e.g.,

$ horovodrun --check-build

Horovod v0.22.1:

Available Frameworks:
    [X] TensorFlow
    [X] PyTorch
    [ ] MXNet

Available Controllers:
    [X] MPI
    [X] Gloo

Available Tensor Operations:
    [X] NCCL
    [ ] DDL
    [ ] CCL
    [X] MPI
    [X] Gloo

Since version 2.0.1, Horovod and mpi4py with MPICH support are shipped with the installer.

If you don’t install Horovod, DeePMD-kit will fall back to serial mode.

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 Tensorflow’s C++ interface (optional)

Since TensorFlow 2.12, TensorFlow C++ library (libtensorflow_cc) is packaged inside the Python library. Thus, you can skip building TensorFlow C++ library manually. If that does not work for you, you can still build it manually.

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.8. Note that TensorFlow may have specific requirements for the compiler version.

First, the C++ interface of Tensorflow should be installed. It is noted that the version of Tensorflow should be consistent with the python interface. You may follow the instruction or run the script $deepmd_source_dir/source/install/build_tf.py to install the corresponding C++ interface.

Install DeePMD-kit’s C++ interface

Now go to the source code directory of DeePMD-kit and make a building place.

cd $deepmd_source_dir/source
mkdir build
cd build

I assume you have activated the TensorFlow Python environment and want to install DeePMD-kit into path $deepmd_root, then execute CMake

cmake -DUSE_TF_PYTHON_LIBS=TRUE -DCMAKE_INSTALL_PREFIX=$deepmd_root ..

If you specify -DUSE_TF_PYTHON_LIBS=FALSE, you need to give the location where TensorFlow’s C++ interface is installed to -DTENSORFLOW_ROOT=${tensorflow_root}.

One may add the following arguments to cmake:

CMake Aurgements	Allowed value	Default value	Usage
-DTENSORFLOW_ROOT=<value>	Path	-	The Path to TensorFlow’s C++ interface.
-DCMAKE_INSTALL_PREFIX=<value>	Path	-	The Path where DeePMD-kit will be installed.
-DUSE_CUDA_TOOLKIT=<value>	`TRUE` or `FALSE`	`FALSE`	If `TRUE`, Build GPU support with CUDA toolkit.
-DCUDA_TOOLKIT_ROOT_DIR=<value>	Path	Detected automatically	The path to the CUDA toolkit directory. CUDA 7.0 or later is supported. NVCC is required.
-DUSE_ROCM_TOOLKIT=<value>	`TRUE` or `FALSE`	`FALSE`	If `TRUE`, Build GPU support with ROCM toolkit.
-DCMAKE_HIP_COMPILER_ROCM_ROOT=<value>	Path	Detected automatically	The path to the ROCM toolkit directory.
-DLAMMPS_SOURCE_ROOT=<value>	Path	-	Only neccessary for LAMMPS plugin mode. The path to the LAMMPS source code. LAMMPS 8Apr2021 or later is supported. If not assigned, the plugin mode will not be enabled.
-DUSE_TF_PYTHON_LIBS=<value>	`TRUE` or `FALSE`	`FALSE`	If `TRUE`, Build C++ interface with TensorFlow’s Python libraries(TensorFlow’s Python Interface is required). And there’s no need for building TensorFlow’s C++ interface.
-DENABLE_NATIVE_OPTIMIZATION	`TRUE` or `FALSE`	`FALSE`	Enable compilation optimization for the native machine’s CPU type. Do not enable it if generated code will run on different CPUs.

If the CMake has been executed successfully, then run the following make commands to build the package:

make -j4
make install

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 have the following executable and libraries installed in $deepmd_root/bin and $deepmd_root/lib

$ ls $deepmd_root/bin
dp_ipi      dp_ipi_low
$ ls $deepmd_root/lib
libdeepmd_cc_low.so  libdeepmd_ipi_low.so  libdeepmd_lmp_low.so  libdeepmd_low.so          libdeepmd_op_cuda.so  libdeepmd_op.so
libdeepmd_cc.so      libdeepmd_ipi.so      libdeepmd_lmp.so      libdeepmd_op_cuda_low.so  libdeepmd_op_low.so   libdeepmd.so

Install LAMMPS

There are two ways to install LAMMPS: the built-in mode and the plugin mode. The built-in mode builds LAMMPS along with the DeePMD-kit and DeePMD-kit will be loaded automatically when running LAMMPS. The plugin mode builds LAMMPS and a plugin separately, so one needs to use plugin load command to load the DeePMD-kit’s LAMMPS plugin library.

Install LAMMPS’s DeePMD-kit module (built-in mode)

Before following this section, DeePMD-kit C++ interface should have be installed.

DeePMD-kit provides a module for running MD simulations with LAMMPS. Now make the DeePMD-kit module for LAMMPS.

cd $deepmd_source_dir/source/build
make lammps

DeePMD-kit will generate a module called USER-DEEPMD in the build directory. If you need the low-precision version, move env_low.sh to env.sh in the directory. Now download the LAMMPS code, and uncompress it.

cd /some/workspace
wget https://github.com/lammps/lammps/archive/stable_23Jun2022_update3.tar.gz
tar xf stable_23Jun2022_update3.tar.gz

The source code of LAMMPS is stored in the directory lammps-stable_23Jun2022_update3. Now go into the LAMMPS code and copy the DeePMD-kit module like this

cd lammps-stable_23Jun2022_update3/src/
cp -r $deepmd_source_dir/source/build/USER-DEEPMD .
make yes-kspace
make yes-extra-fix
make yes-user-deepmd

You can enable any other package you want. Now build LAMMPS

make mpi -j4

If everything works fine, you will end up with an executable lmp_mpi.

./lmp_mpi -h

The DeePMD-kit module can be removed from the LAMMPS source code by

make no-user-deepmd

Install LAMMPS (plugin mode)

Starting from 8Apr2021, LAMMPS also provides a plugin mode, allowing one to build LAMMPS and a plugin separately.

Now download the LAMMPS code (8Apr2021 or later), and uncompress it:

cd /some/workspace
wget https://github.com/lammps/lammps/archive/stable_23Jun2022_update3.tar.gz
tar xf stable_23Jun2022_update3.tar.gz

The source code of LAMMPS is stored in the directory lammps-stable_23Jun2022_update3. The directory of the source code should be specified as the CMAKE argument LAMMPS_SOURCE_ROOT during installation of the DeePMD-kit C++ interface. Now go into the LAMMPS directory and create a directory called build

mkdir -p lammps-stable_23Jun2022_update3/build/
cd lammps-stable_23Jun2022_update3/build/

Now build LAMMPS. Note that PLUGIN and KSPACE packages must be enabled, and BUILD_SHARED_LIBS must be set to yes. You can install any other package you want.

cmake -D PKG_PLUGIN=ON -D PKG_KSPACE=ON -D LAMMPS_INSTALL_RPATH=ON -D BUILD_SHARED_LIBS=yes -D CMAKE_INSTALL_PREFIX=${deepmd_root} -D CMAKE_INSTALL_LIBDIR=lib -D CMAKE_INSTALL_FULL_LIBDIR=${deepmd_root}/lib ../cmake
make -j4
make install

If everything works fine, you will end up with an executable ${deepmd_root}/bin/lmp.

${deepmd_root}/bin/lmp -h

Note

If ${tensorflow_root}, ${deepmd_root}, or the path to TensorFlow Python package if applicable is different from the prefix of LAMMPS, you need to append the library path to RUNPATH of liblammps.so. For example,

patchelf --set-rpath "${tensorflow_root}/lib" liblammps.so

Install 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, forces and virials). The server and client communicate via the Unix domain socket or the Internet socket. Full documentation for i-PI can be found here. The source code and a complete installation guide for i-PI can be found here. To use i-PI with already existing drivers, install and update using Pip:

pip install -U i-PI

Test with Pytest:

pip install pytest
pytest --pyargs ipi.tests

Install GROMACS with DeepMD

Before following this section, DeePMD-kit C++ interface should have be installed.

Patch source code of GROMACS

Download the source code of a supported GROMACS version (2020.2) from https://manual.gromacs.org/2020.2/download.html. Run the following command:

export PATH=$PATH:$deepmd_kit_root/bin
dp_gmx_patch -d $gromacs_root -v $version -p

where deepmd_kit_root is the directory where the latest version of DeePMD-kit is installed, and gromacs_root refers to the source code directory of GROMACS. And version represents the version of GROMACS, where only 2020.2 is supported now. If attempting to patch another version of GROMACS you will still need to set version to 2020.2 as this is the only supported version, we cannot guarantee that patching other versions of GROMACS will work.

Compile GROMACS with deepmd-kit

The C++ interface of Deepmd-kit 2.x and TensorFlow 2.x are required. And be aware that only DeePMD-kit with high precision is supported now since we cannot ensure single precision is enough for a GROMACS simulation. Here is a sample compile script:

#!/bin/bash
export CC=/usr/bin/gcc
export CXX=/usr/bin/g++
export CMAKE_PREFIX_PATH="/path/to/fftw-3.3.9" # fftw libraries
mkdir build
cd build

cmake3 .. -DCMAKE_CXX_STANDARD=14 \ # not required, but c++14 seems to be more compatible with higher version of tensorflow
          -DGMX_MPI=ON \
          -DGMX_GPU=CUDA \ # Gromacs on ROCm has not been fully developed yet
          -DCUDA_TOOLKIT_ROOT_DIR=/path/to/cuda \
          -DCMAKE_INSTALL_PREFIX=/path/to/gromacs-2020.2-deepmd
make -j
make install

Building conda packages

One may want to keep both convenience and personalization of the DeePMD-kit. To achieve this goal, one can consider building conda packages. We provide building scripts in deepmd-kit-recipes organization. These building tools are driven by conda-build and conda-smithy.

For example, if one wants to turn on MPIIO package in LAMMPS, go to lammps-feedstock repository and modify recipe/build.sh. -D PKG_MPIIO=OFF should be changed to -D PKG_MPIIO=ON. Then go to the main directory and execute

./build-locally.py

This requires that Docker has been installed. After the building, the packages will be generated in build_artifacts/linux-64 and build_artifacts/noarch, and then one can install then executing

conda create -n deepmd lammps -c file:///path/to/build_artifacts -c https://conda.deepmodeling.com -c nvidia

One may also upload packages to one’s Anaconda channel, so they can be installed on other machines:

anaconda upload /path/to/build_artifacts/linux-64/*.tar.bz2 /path/to/build_artifacts/noarch/*.tar.bz2

Data

In this section, we will introduce how to convert the DFT-labeled data into the data format used by DeePMD-kit.

The DeePMD-kit organizes data in systems. Each system is composed of a number of frames. One may roughly view a frame as a snapshot of an MD trajectory, but it does not necessarily come from an MD simulation. A frame records the coordinates and types of atoms, cell vectors if the periodic boundary condition is assumed, energy, atomic forces and virials. It is noted that the frames in one system share the same number of atoms with the same type.

System

DeePMD-kit takes a system as the data structure. A snapshot of a system is called a frame. A system may contain multiple frames with the same atom types and numbers, i.e. the same formula (like H2O). To contains data with different formulas, one usually needs to divide data into multiple systems, which may sometimes result in sparse-frame systems. See a new system format to further combine different systems with the same atom numbers, when training with descriptor se_atten.

A system should contain system properties, input frame properties, and labeled frame properties. The system property contains the following property:

ID	Property	Raw file	Required/Optional	Shape	Description
type	Atom type indexes	type.raw	Required	Natoms	Integers that start with 0. If both the training parameter type_map is set and `type_map.raw` is provided, the system atom type should be mapped to `type_map.raw` in `type.raw` and will be mapped to the model atom type when training; otherwise, the system atom type will be always mapped to the model atom type (whether type_map is set or not)
type_map	Atom type names	type_map.raw	Optional	Ntypes	Atom names that map to atom type, which is unnecessary to be contained in the periodic table. Only works when the training parameter type_map is set
nopbc	Non-periodic system	nopbc	Optional	1	If True, this system is non-periodic; otherwise it’s periodic

The input frame properties contain the following property, the first axis of which is the number of frames:

ID	Property	Raw file	Unit	Required/Optional	Shape	Description
coord	Atomic coordinates	coord.raw	Å	Required	Nframes * Natoms * 3
box	Boxes	box.raw	Å	Required if periodic	Nframes * 3 * 3	in the order `XX XY XZ YX YY YZ ZX ZY ZZ`
fparam	Extra frame parameters	fparam.raw	Any	Optional	Nframes * Any
aparam	Extra atomic parameters	aparam.raw	Any	Optional	Nframes * aparam * Any
numb_copy	Each frame is copied by the `numb_copy` (int) times	prob.raw	1	Optional	Nframes	Integer; Default is 1 for all frames

The labeled frame properties are listed as follows, all of which will be used for training if and only if the loss function contains such property:

ID	Property	Raw file	Unit	Shape	Description
energy	Frame energies	energy.raw	eV	Nframes
force	Atomic forces	force.raw	eV/Å	Nframes * Natoms * 3
virial	Frame virial	virial.raw	eV	Nframes * 9	in the order `XX XY XZ YX YY YZ ZX ZY ZZ`
atom_ener	Atomic energies	atom_ener.raw	eV	Nframes * Natoms
atom_pref	Weights of atomic forces	atom_pref.raw	1	Nframes * Natoms
dipole	Frame dipole	dipole.raw	Any	Nframes * 3
atomic_dipole	Atomic dipole	atomic_dipole.raw	Any	Nframes * Natoms * 3
polarizability	Frame polarizability	polarizability.raw	Any	Nframes * 9	in the order `XX XY XZ YX YY YZ ZX ZY ZZ`
atomic_polarizability	Atomic polarizability	atomic_polarizability.raw	Any	Nframes * Natoms * 9	in the order `XX XY XZ YX YY YZ ZX ZY ZZ`

In general, we always use the following convention of units:

Property	Unit
Time	ps
Length	Å
Energy	eV
Force	eV/Å
Virial	eV
Pressure	Bar

Formats of a system

Two binary formats, NumPy and HDF5, are supported for training. The raw format is not directly supported, but a tool is provided to convert data from the raw format to the NumPy format.

NumPy format

In a system with the Numpy format, the system properties are stored as text files ending with .raw, such as type.raw and type_map.raw, under the system directory. If one needs to train a non-periodic system, an empty nopbc file should be put under the system directory. Both input and labeled frame properties are saved as the NumPy binary data (NPY) files ending with .npy in each of the set.* directories. Take an example, a system may contain the following files:

type.raw
type_map.raw
nopbc
set.000/coord.npy
set.000/energy.npy
set.000/force.npy
set.001/coord.npy
set.001/energy.npy
set.001/force.npy

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:

$ cat type.raw
0 1

Sometimes one needs to map the integer types to atom names. The mapping can be given by the file type_map.raw. For example

$ cat type_map.raw
O H

The type 0 is named by "O" and the type 1 is named by "H".

For training models with descriptor se_atten, a new system format is supported to put together the frame-sparse systems with the same atom number.

HDF5 format

A system with the HDF5 format has the same structure as the Numpy format, but in an HDF5 file, a system is organized as an HDF5 group. The file name of a Numpy file is the key in an HDF5 file, and the data is the value of the key. One needs to use # in a DP path to divide the path to the HDF5 file and the HDF5 path:

/path/to/data.hdf5#/H2O

Here, /path/to/data.hdf5 is the file path and /H2O is the HDF5 path. All HDF5 paths should start with /. There should be some data in the H2O group, such as /H2O/type.raw and /H2O/set.000/force.npy.

An HDF5 file with a large number of systems has better performance than multiple NumPy files in a large cluster.

Raw format and data conversion

A raw file is a plain text file with each information item written in one file and one frame written on one line. It’s not directly supported, but we provide a tool to convert them.

In the raw format, the property of one frame is provided per line, ending with .raw. Take an example, the default files that provide box, coordinate, force, energy and virial are box.raw, coord.raw, force.raw, energy.raw and virial.raw, respectively. Here is an example of force.raw:

$ 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. Other files are organized similarly. The number of lines of all raw files should be identical.

One can use the script $deepmd_source_dir/data/raw/raw_to_set.sh to convert the prepared raw files to the NumPy format. For example, if we have a raw file that contains 6000 frames,

$ 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 containing 2000 frames in the Numpy format.

Prepare data with dpdata

One can use a convenient tool dpdata to convert data directly from the output of first principle packages to the DeePMD-kit format.

To install one can execute

pip install dpdata

An example of converting data VASP data in OUTCAR format to DeePMD-kit data can be found at

$deepmd_source_dir/examples/data_conv

Switch to that directory, then one can convert data by using the following python script

import dpdata

dsys = dpdata.LabeledSystem("OUTCAR")
dsys.to("deepmd/npy", "deepmd_data", set_size=dsys.get_nframes())

get_nframes() method gets the number of frames in the OUTCAR, and the argument set_size enforces that the set size is equal to the number of frames in the system, viz. only one set is created in the system.

The data in DeePMD-kit format is stored in the folder deepmd_data.

A list of all supported data format and more nice features of dpdata can be found on the official website.

Model

Overall

A model has two parts, a descriptor that maps atomic configuration to a set of symmetry invariant features, and a fitting net that takes descriptor as input and predicts the atomic contribution to the target physical property. It’s defined in the model section of the input.json, for example,

    "model": {
        "type_map":	["O", "H"],
        "descriptor" :{
            "...": "..."
        },
        "fitting_net" : {
            "...": "..."
        }
    }

The two subsections, descriptor and fitting_net, define the descriptor and the fitting net, respectively.

The type_map is optional, which provides the element names (but not necessarily same as the actual name of the element) of the corresponding atom types. A water model, as in this example, has two kinds of atoms. The atom types are internally recorded as integers, e.g., 0 for oxygen and 1 for hydrogen here. A mapping from the atom type to their names is provided by type_map.

DeePMD-kit implements the following descriptors:

se_e2_a: DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.
se_e2_r: DeepPot-SE constructed from radial information of atomic configurations. The embedding takes the distance between atoms as input.
se_e3: DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes angles between two neighboring atoms as input.
se_a_mask: DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The input frames in one system can have a varied number of atoms. Input particles are padded with virtual particles of the same length.
loc_frame: Defines a local frame at each atom and compute the descriptor as local coordinates under this frame.
hybrid: Concate a list of descriptors to form a new descriptor.

The fitting of the following physical properties is supported

ener: Fit the energy of the system. The force (derivative with atom positions) and the virial (derivative with the box tensor) can also be trained.
dipole: The dipole moment.
polar: The polarizability.

Descriptor `"se_e2_a"`

The notation of se_e2_a is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The e2 stands for the embedding with two-atoms information. This descriptor was described in detail in the DeepPot-SE paper.

Note that it is sometimes called a “two-atom embedding descriptor” which means the input of the embedding net is atomic distances. The descriptor does encode multi-body information (both angular and radial information of neighboring atoms).

In this example, we will train a DeepPot-SE model for a water system. A complete training input script of this example can be found in the directory.

$deepmd_source_dir/examples/water/se_e2_a/input.json

With the training input script, data are also provided in the example directory. One may train the model with the DeePMD-kit from the directory.

The construction of the descriptor is given by section descriptor. An example of the descriptor is provided as follows

	"descriptor" :{
	    "type":		"se_e2_a",
	    "rcut_smth":	0.50,
	    "rcut":		6.00,
	    "sel":		[46, 92],
	    "neuron":		[25, 50, 100],
	    "type_one_side":	true,
	    "axis_neuron":	16,
	    "resnet_dt":	false,
	    "seed":		1
	}

The type of the descriptor is set to "se_e2_a".
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] denotes 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 layer from the input end to the output end, respectively. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a ResNet architecture is built between them.
If the option type_one_side is set to true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters.
The axis_neuron specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the DeepPot-SE paper
If the option resnet_dt is set to 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.

Descriptor `"se_e2_r"`

The notation of se_e2_r is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from the radial information of atomic configurations. The e2 stands for the embedding with two-atom information.

A complete training input script of this example can be found in the directory

$deepmd_source_dir/examples/water/se_e2_r/input.json

The training input script is very similar to that of se_e2_a. The only difference lies in the descriptor section

	"descriptor": {
	    "type":		"se_e2_r",
	    "sel":		[46, 92],
	    "rcut_smth":	0.50,
	    "rcut":		6.00,
	    "neuron":		[5, 10, 20],
	    "resnet_dt":	false,
	    "seed":		1,
	    "_comment": " that's all"
	},

The type of the descriptor is set by the key type.

Descriptor `"se_e3"`

The notation of se_e3 is short for the Deep Potential Smooth Edition (DeepPot-SE) constructed from all information (both angular and radial) of atomic configurations. The embedding takes angles between two neighboring atoms as input (denoted by e3).

A complete training input script of this example can be found in the directory

$deepmd_source_dir/examples/water/se_e3/input.json

The training input script is very similar to that of se_e2_a. The only difference lies in the descriptor <model/descriptor> section

	"descriptor": {
	    "type":		"se_e3",
	    "sel":		[40, 80],
	    "rcut_smth":	0.50,
	    "rcut":		6.00,
	    "neuron":		[2, 4, 8],
	    "resnet_dt":	false,
	    "seed":		1,
	    "_comment":		" that's all"
	},

The type of the descriptor is set by the key type.

Descriptor `"se_atten"`

DPA-1: Pretraining of Attention-based Deep Potential Model for Molecular Simulation

ALT

Here we propose DPA-1, a Deep Potential model with a novel attention mechanism, which is highly effective for representing the conformation and chemical spaces of atomic systems and learning the PES.

See this paper for more information. DPA-1 is implemented as a new descriptor "se_atten" for model training, which can be used after simply editing the input.json.

Installation

Follow the standard installation of Python interface in the DeePMD-kit. After that, you can smoothly use the DPA-1 model with the following instructions.

Introduction to new features of DPA-1

Next, we will list the detailed settings in input.json and the data format, especially for large systems with dozens of elements. An example of DPA-1 input can be found here.

Descriptor `"se_atten"`

The notation of se_atten is short for the smooth edition of Deep Potential with an attention mechanism. This descriptor was described in detail in the DPA-1 paper and the images above.

In this example, we will train a DPA-1 model for a water system. A complete training input script of this example can be found in the directory:

$deepmd_source_dir/examples/water/se_atten/input.json

With the training input script, data are also provided in the example directory. One may train the model with the DeePMD-kit from the directory.

An example of the DPA-1 descriptor is provided as follows

	"descriptor" :{
          "type":		"se_atten",
          "rcut_smth":	0.50,
          "rcut":		6.00,
          "sel":		120,
          "neuron":		[25, 50, 100],
          "axis_neuron":	16,
          "resnet_dt":	false,
          "attn":	128,
          "attn_layer":	2,
          "attn_mask":	false,
          "attn_dotr":	true,
          "seed":	1
	}

The type of the descriptor is set to "se_atten", which will use DPA-1 structures.
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 an int. Note that this number highly affects the efficiency of training, which we usually use less than 200. (We use 120 for training 56 elements in OC2M dataset)
The neuron specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layer from the input end to the output end, respectively. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a ResNet architecture is built between them.
The axis_neuron specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the DeepPot-SE paper
If the option resnet_dt is set to 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.
attn sets the length of a hidden vector during scale-dot attention computation.
attn_layer sets the number of layers in attention mechanism.
attn_mask determines whether to mask the diagonal in the attention weights and False is recommended.
attn_dotr determines whether to dot the relative coordinates on the attention weights as a gated scheme, True is recommended.

Fitting `"ener"`

DPA-1 only supports "ener" fitting type, and you can refer here for detailed information.

Type embedding

DPA-1 only supports models with type embeddings. And the default setting is as follows:

"type_embedding":{
            "neuron":           [8],
            "resnet_dt":        false,
            "seed":             1
        }

You can add these settings in input.json if you want to change the default ones, see here for detailed information.

Type map

For training large systems, especially those with dozens of elements, the type determines the element index of training data:

"type_map": [
   "Mg",
   "Al",
   "Cu"
  ]

which should include all the elements in the dataset you want to train on.

Data format

DPA-1 supports the standard data format, which is detailed in data-conv.md and system.md. Note that in this format, only those frames with the same fingerprint (i.e. the number of atoms of different elements) can be put together as a unified system. This may lead to sparse frame numbers in those rare systems.

An ideal way is to put systems with the same total number of atoms together, which is the way we trained DPA-1 on OC2M. This system format, which is called mixed_type, is proper to put frame-sparse systems together and is slightly different from the standard one. Take an example, a mixed_type may contain the following files:

type.raw
type_map.raw
set.*/box.npy
set.*/coord.npy
set.*/energy.npy
set.*/force.npy
set.*/real_atom_types.npy

This system contains Nframes frames with the same atom number Natoms, the total number of element types contained in all frames is Ntypes. Most files are the same as those in standard formats, here we only list the distinct ones:

ID	Property	File	Required/Optional	Shape	Description
/	Atom type indexes (place holder)	type.raw	Required	Natoms	All zeros to fake the type input
type_map	Atom type names	type_map.raw	Required	Ntypes	Atom names that map to atom type contained in all the frames, which is unnecessart to be contained in the periodic table
type	Atom type indexes of each frame	real_atom_types.npy	Required	Nframes * Natoms	Integers that describe atom types in each frame, corresponding to indexes in type_map. `-1` means virtual atoms.

With these edited files, one can put together frames with the same Natoms, instead of the same formula (like H2O). Note that this mixed_type format only supports se_atten descriptor.

To put frames with different Natoms into the same system, one can pad systems by adding virtual atoms whose type is -1. Virtual atoms do not contribute to any fitting property, so the atomic property of virtual atoms (e.g. forces) should be given zero.

The API to generate or transfer to mixed_type format is available on dpdata for a more convenient experience.

Training example

Here we upload the AlMgCu example shown in the paper, you can download it here: Baidu disk; Google disk.

Descriptor `"hybrid"`

This descriptor hybridizes multiple descriptors to form a new descriptor. For example, we have a list of descriptors denoted by $\mathcal D_1$, $\mathcal D_2$, …, $\mathcal D_N$, the hybrid descriptor this the concatenation of the list, i.e. $\mathcal D = (\mathcal D_1, \mathcal D_2, \cdots, \mathcal D_N)$.

To use the descriptor in DeePMD-kit, one firstly set the type to hybrid, then provide the definitions of the descriptors by the items in the list,

        "descriptor" :{
            "type": "hybrid",
            "list" : [
                {
		    "type" : "se_e2_a",
		    ...
                },
                {
		    "type" : "se_e2_r",
		    ...
                }
            ]
        },

A complete training input script of this example can be found in the directory

$deepmd_source_dir/examples/water/hybrid/input.json

Determine `sel`

All descriptors require to set sel, which means the expected maximum number of type-i neighbors of an atom. DeePMD-kit will allocate memory according to sel.

sel should not be too large or too small. If sel is too large, the computing will become much slower and cost more memory. If sel is not enough, the energy will be not conserved, making the accuracy of the model worse.

To determine a proper sel, one can calculate the neighbor stat of the training data before training:

dp neighbor-stat -s data -r 6.0 -t O H

where data is the directory of data, 6.0 is the cutoff radius, and O and H is the type map. The program will give the max_nbor_size. For example, max_nbor_size of the water example is [38, 72], meaning an atom may have 38 O neighbors and 72 H neighbors in the training data.

The sel should be set to a higher value than that of the training data, considering there may be some extreme geometries during MD simulations. As a result, we set sel to [46, 92] in the water example.

Fit energy

In this section, we will take $deepmd_source_dir/examples/water/se_e2_a/input.json as an example of the input file.

The fitting network

The construction of the fitting net is given by section fitting_net

	"fitting_net" : {
	    "neuron":		[240, 240, 240],
	    "resnet_dt":	true,
	    "seed":		1
	},

neuron specifies the size of the fitting net. If two neighboring layers are of the same size, then a ResNet architecture is built between them.
If the option resnet_dt is set to 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.

Loss

The loss function $L$ for training energy is given by

\[L = p_e L_e + p_f L_f + p_v L_v\]

where $L_e$, $L_f$, and $L_v$ denote the loss in energy, forces and virials, respectively. $p_e$, $p_f$, and $p_v$ give the prefactors of the energy, force and virial losses. The prefectors may not be a constant, rather it changes linearly with the learning rate. Taking the force prefactor for example, at training step $t$, it is given by

\[p_f(t) = p_f^0 \frac{ \alpha(t) }{ \alpha(0) } + p_f^\infty ( 1 - \frac{ \alpha(t) }{ \alpha(0) })\]

where $\alpha(t)$ denotes the learning rate at step $t$. $p_f^0$ and $p_f^\infty$ specifies the $p_f$ at the start of the training and the limit of $t \to \infty$ (set by start_pref_f and limit_pref_f, respectively), i.e.

pref_f(t) = start_pref_f * ( lr(t) / start_lr ) + limit_pref_f * ( 1 - lr(t) / start_lr )

The loss section in the input.json is

    "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
    }

The options start_pref_e, limit_pref_e, start_pref_f, limit_pref_f, start_pref_v and limit_pref_v determine the start and limit prefactors of energy, force and virial, respectively.

If one does not want to train with virial, then he/she may set the virial prefactors start_pref_v and limit_pref_v to 0.

Fit `tensor` like `Dipole` and `Polarizability`

Unlike energy, which is a scalar, one may want to fit some high dimensional physical quantity, like dipole (vector) and polarizability (matrix, shorted as polar). Deep Potential has provided different APIs to do this. In this example, we will show you how to train a model to fit a water system. A complete training input script of the examples can be found in

$deepmd_source_dir/examples/water_tensor/dipole/dipole_input.json
$deepmd_source_dir/examples/water_tensor/polar/polar_input.json

The training and validation data are also provided our examples. But note that the data provided along with the examples are of limited amount, and should not be used to train a production model.

Similar to the input.json used in ener mode, training JSON is also divided into model, learning_rate, loss and training. Most keywords remain the same as ener mode, and their meaning can be found here. To fit a tensor, one needs to modify model/fitting_net and loss.

The fitting Network

The fitting_net section tells DP which fitting net to use.

The JSON of dipole type should be provided like

	"fitting_net" : {
		"type": "dipole",
		"sel_type": [0],
		"neuron": [100,100,100],
		"resnet_dt": true,
		"seed": 1,
	},

The JSON of polar type should be provided like

	"fitting_net" : {
	   	"type": "polar",
		"sel_type": [0],
		"neuron": [100,100,100],
		"resnet_dt": true,
		"seed": 1,
	},

type specifies which type of fitting net should be used. It should be either dipole or polar. Note that global_polar mode in version 1.x is already deprecated and is merged into polar. To specify whether a system is global or atomic, please see here.
sel_type is a list specifying which type of atoms have the quantity you want to fit. For example, in the water system, sel_type is [0] since 0 represents atom O. If left unset, all types of atoms will be fitted.
The rest arguments have the same meaning as they do in ener mode.

Loss

DP supports a combinational training of the global system (only a global tensor label, i.e. dipole or polar, is provided in a frame) and atomic system (labels for each atom included in sel_type are provided). In a global system, each frame has just one tensor label. For example, when fitting polar, each frame will just provide a 1 x 9 vector which gives the elements of the polarizability tensor of that frame in order XX, XY, XZ, YX, YY, YZ, XZ, ZY, ZZ. By contrast, in an atomic system, each atom in sel_type has a tensor label. For example, when fitting a dipole, each frame will provide a #sel_atom x 3 matrices, where #sel_atom is the number of atoms whose type are in sel_type.

The loss section tells DP the weight of these two kinds of loss, i.e.

loss = pref * global_loss + pref_atomic * atomic_loss

The loss section should be provided like

	"loss" : {
		"type":		"tensor",
		"pref":		1.0,
		"pref_atomic":	1.0
	},

type should be written as tensor as a distinction from ener mode.
pref and pref_atomic respectively specify the weight of global loss and atomic loss. It can not be left unset. If set to 0, the corresponding label will NOT be included in the training process.

Training Data Preparation

In tensor mode, the identification of the label’s type (global or atomic) is derived from the file name. The global label should be named dipole.npy/raw or polarizability.npy/raw, while the atomic label should be named atomic_dipole.npy/raw or atomic_polarizability.npy/raw. If wrongly named, DP will report an error

ValueError: cannot reshape array of size xxx into shape (xx,xx). This error may occur when your label mismatch it's name, i.e. you might store global tensor in `atomic_tensor.npy` or atomic tensor in `tensor.npy`.

In this case, please check the file name of the label.

Train the Model

The training command is the same as ener mode, i.e.

dp train input.json

The detailed loss can be found in lcurve.out:

#  step    rmse_val   rmse_trn  rmse_lc_val rmse_lc_trn rmse_gl_val rmse_gl_trn  lr
   8.34e+00   8.26e+00   8.34e+00   8.26e+00    0.00e+00    0.00e+00   1.0e-02
   3.51e-02   8.55e-02   0.00e+00   8.55e-02    4.38e-03    0.00e+00   5.0e-03
   4.77e-02   5.61e-02   0.00e+00   5.61e-02    5.96e-03    0.00e+00   2.5e-03
   5.68e-02   1.47e-02   0.00e+00   0.00e+00    7.10e-03    1.84e-03   1.3e-03
   3.73e-02   3.48e-02   1.99e-02   0.00e+00    2.18e-03    4.35e-03   6.3e-04
   2.77e-02   5.82e-02   1.08e-02   5.82e-02    2.11e-03    0.00e+00   3.2e-04
   2.81e-02   5.43e-02   2.01e-02   0.00e+00    1.01e-03    6.79e-03   1.6e-04
   2.97e-02   3.28e-02   2.03e-02   0.00e+00    1.17e-03    4.10e-03   7.9e-05
   2.25e-02   6.19e-02   9.05e-03   0.00e+00    1.68e-03    7.74e-03   4.0e-05
   3.18e-02   5.54e-02   9.93e-03   5.54e-02    2.74e-03    0.00e+00   2.0e-05
   2.63e-02   5.02e-02   1.02e-02   5.02e-02    2.01e-03    0.00e+00   1.0e-05
   3.27e-02   5.89e-02   2.13e-02   5.89e-02    1.43e-03    0.00e+00   5.0e-06
   2.85e-02   2.42e-02   2.85e-02   0.00e+00    0.00e+00    3.02e-03   2.5e-06
   3.47e-02   5.71e-02   1.07e-02   5.71e-02    3.00e-03    0.00e+00   1.3e-06
   3.13e-02   5.76e-02   3.13e-02   5.76e-02    0.00e+00    0.00e+00   6.3e-07
   3.34e-02   1.11e-02   2.09e-02   0.00e+00    1.57e-03    1.39e-03   3.2e-07
   3.11e-02   5.64e-02   3.11e-02   5.64e-02    0.00e+00    0.00e+00   1.6e-07
   2.97e-02   5.05e-02   2.97e-02   5.05e-02    0.00e+00    0.00e+00   7.9e-08
   2.64e-02   7.70e-02   1.09e-02   0.00e+00    1.94e-03    9.62e-03   4.0e-08
   3.28e-02   2.56e-02   3.28e-02   0.00e+00    0.00e+00    3.20e-03   2.0e-08
   2.59e-02   5.71e-02   1.03e-02   5.71e-02    1.94e-03    0.00e+00   1.0e-08

One may notice that in each step, some of the local loss and global loss will be 0.0. This is because our training data and validation data consist of the global system and atomic system, i.e.

	--training_data
		>atomic_system
		>global_system
	--validation_data
		>atomic_system
		>global_system

During training, at each step when the lcurve.out is printed, the system used for evaluating the training (validation) error may be either with only global or only atomic labels, thus the corresponding atomic or global errors are missing and are printed as zeros.

Type embedding approach

We generate specific a type embedding vector for each atom type so that we can share one descriptor embedding net and one fitting net in total, which decline training complexity largely.

The training input script is similar to that of se_e2_a, but different by adding the type_embedding section.

Type embedding net

The model defines how the model is constructed, adding a section of type embedding net:

    "model": {
	"type_map":	["O", "H"],
	"type_embedding":{
			...
	},
	"descriptor" :{
            ...
	},
	"fitting_net" : {
            ...
	}
    }

The model will automatically apply the type embedding approach and generate type embedding vectors. If the type embedding vector is detected, the descriptor and fitting net would take it as a part of the input.

The construction of type embedding net is given by type_embedding. An example of type_embedding is provided as follows

	"type_embedding":{
	    "neuron":		[2, 4, 8],
	    "resnet_dt":	false,
	    "seed":		1
	}

The neuron specifies the size of the type embedding net. From left to right the members denote the sizes of each hidden layer from the input end to the output end, respectively. It takes a one-hot vector as input and output dimension equals to the last dimension of the neuron list. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a ResNet architecture is built between them.
If the option resnet_dt is set to 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.

A complete training input script of this example can be found in the directory.

$deepmd_source_dir/examples/water/se_e2_a_tebd/input.json

See here for further explanation of type embedding.

Note

You can’t apply the compression method while using the atom type embedding.

Descriptor `"se_a_mask"`

Descriptor se_a_mask is a concise implementation of the descriptor se_e2_a, but functions slightly differently. se_a_mask is specially designed for DP/MM simulations where the number of atoms in DP regions is dynamically changed in simulations.

Therefore, the descriptor se_a_mask is not supported for training with PBC systems for simplicity. Besides, to make the output shape of the descriptor matrix consistent, the input coordinates are padded with virtual particle coordinates to the maximum number of atoms (specified with sel in the descriptor setting) in the system. The real/virtual sign of the atoms is specified with the aparam.npy ( [ nframes * natoms ] ) file in the input systems set directory. The aparam.npy can also be seen as the mask of the atoms in the system, which is also the origin of the name se_a_mask.

In this example, we will train a DP Mask model for zinc protein interactions. The input systems are the collection of zinc and its coordinates residues. A sample input system that contains 2 frames is included in the directory.

$deepmd_source_dir/examples/zinc_protein/data_dp_mask

A complete training input script of this example can be found in the directory.

$deepmd_source_dir/examples/zinc_protein/zinc_se_a_mask.json

The construction of the descriptor is given by section descriptor. An example of the descriptor is provided as follows

	"descriptor" :{
	    "type":	"se_a_mask",
	    "sel":		[36, 16, 24, 64, 6, 1],
	    "neuron":		[25, 50, 100],
		"axis_neuron": 16,
	    "type_one_side":	false,
	    "resnet_dt":	false,
	    "seed":		1
	}

The type of the descriptor is set to "se_a_mask".
sel gives the maximum number of atoms in input coordinates. It is a list, the length of which is the same as the number of atom types in the system, and sel[i] denotes the maximum number of atoms with type i.
The neuron specifies the size of the embedding net. From left to right the members denote the sizes of each hidden layer from the input end to the output end, respectively. If the outer layer is twice the size of the inner layer, then the inner layer is copied and concatenated, then a ResNet architecture is built between them.
The axis_neuron specifies the size of the submatrix of the embedding matrix, the axis matrix as explained in the DeepPot-SE paper
If the option type_one_side is set to true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_\text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_\text{types}^2$ sets of embedding network parameters.
If the option resnet_dt is set to 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.

To make the aparam.npy used for descriptor se_a_mask, two variables in fitting_net section are needed.

	"fitting_net" :{
	    "neuron": [240, 240, 240],
      	"resnet_dt": true,
      	"seed": 1,
      	"numb_aparam": 1,
      	"use_aparam_as_mask": true
	}

neuron, resnet_dt and seed are the same as the fitting_net section for fitting energy.
numb_aparam gives the dimesion of the aparam.npy file. In this example, it is set to 1 and stores the real/virtual sign of the atoms. For real/virtual atoms, the corresponding sign in aparam.npy is set to 1/0.
use_aparam_as_mask is set to true to use the aparam.npy as the mask of the atoms in the descriptor se_a_mask.

Finally, to make a reasonable fitting task with se_a_mask descriptor for DP/MM simulations, the loss function with se_a_mask is designed to include the atomic forces difference in specific atoms of the input particles only. More details about the selection of the specific atoms can be found in paper [DP/MM](left to be filled). Thus, atom_pref.npy ( [ nframes * natoms ] ) is required as the indicator of the specific atoms in the input particles. And the loss section in the training input script should be set as follows.

"loss": {
    "type": "ener",
    "start_pref_e": 0.0,
    "limit_pref_e": 0.0,
    "start_pref_f": 0.0,
    "limit_pref_f": 0.0,
    "start_pref_pf": 1.0,
    "limit_pref_pf": 1.0,
    "_comment": " that's all"
  }

Deep potential long-range (DPLR)

Notice: The interfaces of DPLR are not stable and subject to change

The method of DPLR is described in this paper. One is recommended to read the paper before using the DPLR.

In the following, we take the DPLR model for example to introduce the training and LAMMPS simulation with the DPLR model. The DPLR model is trained in two steps.

Train a deep Wannier model for Wannier centroids

We use the deep Wannier model (DW) to represent the relative position of the Wannier centroid (WC) with the atom with which it is associated. One may consult the introduction of the dipole model for a detailed introduction. An example input wc.json and a small dataset data for tutorial purposes can be found in

$deepmd_source_dir/examples/water/dplr/train/

It is noted that the tutorial dataset is not enough for training a productive model. Two settings make the training input script different from an energy training input:

	"fitting_net": {
	    "type":		"dipole",
	    "dipole_type":	[0],
	    "neuron":		[128, 128, 128],
	    "seed":		1
	},

The type of fitting is set to dipole. The dipole is associated with type 0 atoms (oxygens), by the setting "dipole_type": [0]. What we trained is the displacement of the WC from the corresponding oxygen atom. It shares the same training input as the atomic dipole because both are 3-dimensional vectors defined on atoms. The loss section is provided as follows

    "loss": {
	"type":		"tensor",
	"pref":		0.0,
	"pref_atomic":	1.0
    },

so that the atomic dipole is trained as labels. Note that the NumPy compressed file atomic_dipole.npy should be provided in each dataset.

The training and freezing can be started from the example directory by

dp train dw.json && dp freeze -o dw.pb

Train the DPLR model

The training of the DPLR model is very similar to the standard short-range DP models. An example input script can be found in the example directory. The following section is introduced to compute the long-range energy contribution of the DPLR model, and modify the short-range DP model by this part.

        "modifier": {
            "type":             "dipole_charge",
            "model_name":       "dw.pb",
            "model_charge_map": [-8],
            "sys_charge_map":   [6, 1],
            "ewald_h":          1.00,
            "ewald_beta":       0.40
        },

The model_name specifies which DW model is used to predict the position of WCs. model_charge_map gives the amount of charge assigned to WCs. sys_charge_map provides the nuclear charge of oxygen (type 0) and hydrogen (type 1) atoms. ewald_beta (unit $\text{Å}^{-1}$) gives the spread parameter controls the spread of Gaussian charges, and ewald_h (unit Å) assigns the grid size of Fourier transformation. The DPLR model can be trained and frozen by (from the example directory)

dp train ener.json && dp freeze -o ener.pb

Molecular dynamics simulation with DPLR

In MD simulations, the long-range part of the DPLR is calculated by the LAMMPS kspace support. Then the long-range interaction is back-propagated to atoms by DeePMD-kit. This setup is commonly used in classical molecular dynamics simulations as the “virtual site”. Unfortunately, LAMMPS does not natively support virtual sites, so we have to hack the LAMMPS code, which makes the input configuration and script a little wired.

An example of an input configuration file and script can be found in

$deepmd_source_dir/examples/water/dplr/lmp/

We use atom_style full for DPLR simulations. the coordinates of the WCs are explicitly written in the configuration file. Moreover, a virtual bond is established between the oxygens and the WCs to indicate they are associated together. The configuration file containing 128 H2O molecules is thus written as

512 atoms
3 atom types
128 bonds
1 bond types

0 16.421037674 xlo xhi
0 16.421037674 ylo yhi
0 16.421037674 zlo zhi
0 0 0 xy xz yz

Masses

1 16
2 2
3 16

Atoms

       1        1 1  6 8.4960699081e+00 7.5073699951e+00 9.6371297836e+00
       2        2 1  6 4.0597701073e+00 6.8156299591e+00 1.2051420212e+01
...

     385        1 3 -8 8.4960699081e+00 7.5073699951e+00 9.6371297836e+00
     386        2 3 -8 4.0597701073e+00 6.8156299591e+00 1.2051420212e+01
...

Bonds

1 1 1 385
2 1 2 386
...

The oxygens and hydrogens are assigned with atom types 1 and 2 (corresponding to training atom types 0 and 1), respectively. The WCs are assigned with atom type 3. We want to simulate heavy water so the mass of hydrogens is set to 2.

An example input script is provided in

$deepmd_source_dir/examples/water/dplr/lmp/in.lammps

Here are some explanations

# groups of real and virtual atoms
group           real_atom type 1 2
group           virtual_atom type 3

# bond between real and its corresponding virtual site should be given
# to setup a map between real and virtual atoms. However, no real
# bonded interaction is applied, thus bond_sytle "zero" is used.
pair_style      deepmd ener.pb
pair_coeff      * *
bond_style      zero
bond_coeff      *
special_bonds   lj/coul 1 1 1 angle no

Type 1 and 2 (O and H) are real_atoms, while type 3 (WCs) are virtual_atoms. The model file ener.pb stores both the DW and DPLR models, so the position of WCs and the energy can be inferred from it. A virtual bond type is specified by bond_style zero. The special_bonds command switches off the exclusion of intramolecular interactions.

# kspace_style "pppm/dplr" should be used. in addition the
# gewald(1/distance) should be set the same as that used in
# training. Currently only ik differentiation is supported.
kspace_style	pppm/dplr 1e-5
kspace_modify	gewald ${BETA} diff ik mesh ${KMESH} ${KMESH} ${KMESH}

The long-range part is calculated by the kspace support of LAMMPS. The kspace_style pppm/dplr is required. The spread parameter set by variable BETA should be set the same as that used in training. The KMESH should be set dense enough so the long-range calculation is converged.

# "fix dplr" set the position of the virtual atom, and spread the
# electrostatic interaction asserting on the virtual atom to the real
# atoms. "type_associate" associates the real atom type its
# corresponding virtual atom type. "bond_type" gives the type of the
# bond between the real and virtual atoms.
fix		0 all dplr model ener.pb type_associate 1 3 bond_type 1
fix_modify	0 virial yes

The fix command dplr calculates the position of WCs by the DW model and back-propagates the long-range interaction on virtual atoms to real toms. At this time, the training parameter type_map will be mapped to LAMMPS atom types.

# compute the temperature of real atoms, excluding virtual atom contribution
compute		real_temp real_atom temp
compute		real_press all pressure real_temp
fix		1 real_atom nvt temp ${TEMP} ${TEMP} ${TAU_T}
fix_modify	1 temp real_temp

The temperature of the system should be computed from the real atoms. The kinetic contribution in the pressure tensor is also computed from the real atoms. The thermostat is applied to only real atoms. The computed temperature and pressure of real atoms can be accessed by, e.g.

fix             thermo_print all print ${THERMO_FREQ} "$(step) $(pe) $(ke) $(etotal) $(enthalpy) $(c_real_temp) $(c_real_press) $(vol) $(c_real_press[1]) $(c_real_press[2]) $(c_real_press[3])" append thermo.out screen no title "# step pe ke etotal enthalpy temp press vol pxx pyy pzz"

The LAMMPS simulation can be started from the example directory by

lmp -i in.lammps

If LAMMPS complains that no model file ener.pb exists, it can be copied from the training example directory.

The MD simulation lasts for only 20 steps. If one runs a longer simulation, it will blow up, because the model is trained with a very limited dataset for very short training steps, thus is of poor quality.

Another restriction that should be noted is that the energies printed at the zero steps are not correct. This is because at the zero steps the position of the WC has not been updated with the DW model. The energies printed in later steps are correct.

Deep Potential - Range Correction (DPRc)

Deep Potential - Range Correction (DPRc) is designed to combine with QM/MM method, and corrects energies from a low-level QM/MM method to a high-level QM/MM method:

\[ E=E_\text{QM}(\mathbf R; \mathbf P) + E_\text{QM/MM}(\mathbf R; \mathbf P) + E_\text{MM}(\mathbf R) + E_\text{DPRc}(\mathbf R) \]

See the JCTC paper for details.

Training data

Instead the normal ab initio data, one needs to provide the correction from a low-level QM/MM method to a high-level QM/MM method:

\[ E = E_\text{high-level QM/MM} - E_\text{low-level QM/MM} \]

Two levels of data use the same MM method, so $E_\text{MM}$ is eliminated.

Training the DPRc model

In a DPRc model, QM atoms and MM atoms have different atom types. Assuming we have 4 QM atom types (C, H, O, P) and 2 MM atom types (HW, OW):

"type_map": ["C", "H", "HW", "O", "OW", "P"]

As described in the paper, the DPRc model only corrects $E_\text{QM}$ and $E_\text{QM/MM}$ within the cutoff, so we use a hybrid descriptor to describe them separatedly:

"descriptor" :{
    "type":             "hybrid",
    "list" : [
        {
            "type":     "se_e2_a",
            "sel":              [6, 11, 0, 6, 0, 1],
            "rcut_smth":        1.00,
            "rcut":             9.00,
            "neuron":           [12, 25, 50],
            "exclude_types":    [[2, 2], [2, 4], [4, 4], [0, 2], [0, 4], [1, 2], [1, 4], [3, 2], [3, 4], [5, 2], [5, 4]],
            "axis_neuron":      12,
            "set_davg_zero":    true,
            "_comment": " QM/QM interaction"
        },
        {
            "type":     "se_e2_a",
            "sel":              [6, 11, 100, 6, 50, 1],
            "rcut_smth":        0.50,
            "rcut":             6.00,
            "neuron":           [12, 25, 50],
            "exclude_types":    [[0, 0], [0, 1], [0, 3], [0, 5], [1, 1], [1, 3], [1, 5], [3, 3], [3, 5], [5, 5], [2, 2], [2, 4], [4, 4]],
            "axis_neuron":      12,
            "set_davg_zero":    true,
            "_comment": " QM/MM interaction"
        }
    ]
}

exclude_types can be generated by the following Python script:

from itertools import combinations_with_replacement, product
qm = (0, 1, 3, 5)
mm = (2, 4)
print("QM/QM:", list(map(list, list(combinations_with_replacement(mm, 2)) + list(product(qm, mm)))))
print("QM/MM:", list(map(list, list(combinations_with_replacement(qm, 2)) + list(combinations_with_replacement(mm, 2)))))

Also, DPRc assumes MM atom energies (atom_ener) are zero:

"fitting_net": {
   "neuron": [240, 240, 240],
   "resnet_dt": true,
   "atom_ener": [null, null, 0.0, null, 0.0, null]
}

Note that atom_ener only works when descriptor/set_davg_zero is true.

Run MD simulations

The DPRc model has the best practices with the AMBER QM/MM module. An example is given by GitLab RutgersLBSR/AmberDPRc. In theory, DPRc is able to be used with any QM/MM package, as long as the DeePMD-kit package accepts QM atoms and MM atoms within the cutoff range and returns energies and forces.

Training

Train a model

Several examples of training can be found in the examples directory:

$ cd $deepmd_source_dir/examples/water/se_e2_a/

After switching to that directory, the training can be invoked by

$ dp train input.json

where input.json is the name of the input script.

By default, the verbosity level of the DeePMD-kit is INFO, one may see a lot of important information on the code and environment showing on the screen. Among them two pieces of information regarding data systems are worth special notice.

DEEPMD INFO    ---Summary of DataSystem: training     -----------------------------------------------
DEEPMD INFO    found 3 system(s):
DEEPMD INFO                                        system  natoms  bch_sz   n_bch   prob  pbc
DEEPMD INFO                         ../data_water/data_0/     192       1      80  0.250    T
DEEPMD INFO                         ../data_water/data_1/     192       1     160  0.500    T
DEEPMD INFO                         ../data_water/data_2/     192       1      80  0.250    T
DEEPMD INFO    --------------------------------------------------------------------------------------
DEEPMD INFO    ---Summary of DataSystem: validation   -----------------------------------------------
DEEPMD INFO    found 1 system(s):
DEEPMD INFO                                        system  natoms  bch_sz   n_bch   prob  pbc
DEEPMD INFO                          ../data_water/data_3     192       1      80  1.000    T
DEEPMD INFO    --------------------------------------------------------------------------------------

The DeePMD-kit prints detailed information on the training and validation data sets. The data sets are defined by training_data and validation_data defined in the training section of the input script. The training data set is composed of three data systems, while the validation data set is composed by one data system. The number of atoms, batch size, the number of batches in the system and the probability of using the system are all shown on the screen. The last column presents if the periodic boundary condition is assumed for the system.

During the training, the error of the model is tested every disp_freq training steps with the batch used to train the model and with numb_btch batches from the validating data. The training error and validation error are printed correspondingly in the file disp_file (default is lcurve.out). The batch size can be set in the input script by the key batch_size in the corresponding sections for the training and validation data set. An example of the output

#  step      rmse_val    rmse_trn    rmse_e_val  rmse_e_trn    rmse_f_val  rmse_f_trn         lr
    3.33e+01    3.41e+01      1.03e+01    1.03e+01      8.39e-01    8.72e-01    1.0e-03
    2.57e+01    2.56e+01      1.87e+00    1.88e+00      8.03e-01    8.02e-01    1.0e-03
    2.45e+01    2.56e+01      2.26e-01    2.21e-01      7.73e-01    8.10e-01    1.0e-03
    1.62e+01    1.66e+01      5.01e-02    4.46e-02      5.11e-01    5.26e-01    1.0e-03
    1.36e+01    1.32e+01      1.07e-02    2.07e-03      4.29e-01    4.19e-01    1.0e-03
    1.07e+01    1.05e+01      2.45e-03    4.11e-03      3.38e-01    3.31e-01    1.0e-03

The file contains 8 columns, from left to right, which are the training step, the validation loss, training loss, root mean square (RMS) validation error of energy, RMS training error of energy, RMS validation error of force, RMS training error of force and the learning rate. The RMS error (RMSE) of the energy is normalized by the number of atoms in the system. One can visualize this file with a simple Python script:

import numpy as np
import matplotlib.pyplot as plt

data = np.genfromtxt("lcurve.out", names=True)
for name in data.dtype.names[1:-1]:
    plt.plot(data['step'], data[name], label=name)
plt.legend()
plt.xlabel('Step')
plt.ylabel('Loss')
plt.xscale('symlog')
plt.yscale('log')
plt.grid()
plt.show()

Checkpoints will be written to files with the prefix save_ckpt every save_freq training steps.

Warning

It is warned that the example water data (in folder examples/water/data) is of very limited amount, is provided only for testing purposes, and should not be used to train a production model.

Advanced options

In this section, we will take $deepmd_source_dir/examples/water/se_e2_a/input.json as an example of the input file.

Learning rate

The learning_rate section in input.json is given as follows

    "learning_rate" :{
	"type":		"exp",
	"start_lr":	0.001,
	"stop_lr":	3.51e-8,
	"decay_steps":	5000,
	"_comment":	"that's all"
    }

start_lr gives the learning rate at the beginning of the training.
stop_lr gives the learning rate at the end of the training. It should be small enough to ensure that the network parameters satisfactorily converge.
During the training, the learning rate decays exponentially from start_lr to stop_lr following the formula:

\[ \alpha(t) = \alpha_0 \lambda ^ { t / \tau } \]

where $t$ is the training step, $\alpha$ is the learning rate, $\alpha_0$ is the starting learning rate (set by start_lr), $\lambda$ is the decay rate, and $\tau$ is the decay steps, i.e.

```
lr(t) = start_lr * decay_rate ^ ( t / decay_steps )
```

Training parameters

Other training parameters are given in the training section.

    "training": {
 	"training_data": {
	    "systems":		["../data_water/data_0/", "../data_water/data_1/", "../data_water/data_2/"],
	    "batch_size":	"auto"
	},
	"validation_data":{
	    "systems":		["../data_water/data_3"],
	    "batch_size":	1,
	    "numb_btch":	3
	},
	"mixed_precision": {
	    "output_prec":      "float32",
	    "compute_prec":     "float16"
	},

	"numb_steps":	1000000,
	"seed":		1,
	"disp_file":	"lcurve.out",
	"disp_freq":	100,
	"save_freq":	1000
    }

The sections training_data and validation_data give the training dataset and validation dataset, respectively. Taking the training dataset for example, the keys are explained below:

systems provide paths of the training data systems. DeePMD-kit allows you to provide multiple systems with different numbers of atoms. This key can be a list or a str.
- list: systems gives the training data systems.
- str: systems should be a valid path. DeePMD-kit will recursively search all data systems in this path.
At each training step, DeePMD-kit randomly picks batch_size frame(s) from one of the systems. The probability of using a system is by default in proportion to the number of batches in the system. More options are available for automatically determining the probability of using systems. One can set the key auto_prob to
- "prob_uniform" all systems are used with the same probability.
- "prob_sys_size" the probability of using a system is proportional to its size (number of frames).
- "prob_sys_size; sidx_0:eidx_0:w_0; sidx_1:eidx_1:w_1;..." the list of systems is divided into blocks. Block i has systems ranging from sidx_i to eidx_i. The probability of using a system from block i is proportional to w_i. Within one block, the probability of using a system is proportional to its size.
An example of using "auto_prob" is given below. The probability of using systems[2] is 0.4, and the sum of the probabilities of using systems[0] and systems[1] is 0.6. If the number of frames in systems[1] is twice of system[0], then the probability of using system[1] is 0.4 and that of system[0] is 0.2.

 	"training_data": {
	    "systems":		["../data_water/data_0/", "../data_water/data_1/", "../data_water/data_2/"],
	    "auto_prob":	"prob_sys_size; 0:2:0.6; 2:3:0.4",
	    "batch_size":	"auto"
	}

The probability of using systems can also be specified explicitly with key sys_probs which is a list having the length of the number of systems. For example

 	"training_data": {
	    "systems":		["../data_water/data_0/", "../data_water/data_1/", "../data_water/data_2/"],
	    "sys_probs":	[0.5, 0.3, 0.2],
	    "batch_size":	"auto:32"
	}

The key batch_size specifies the number of frames used to train or validate the model in a training step. It can be set to
- 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 use the same batch size.
- "auto": the same as "auto:32", see "auto:N"
- "auto:N": automatically determines the batch size so that the batch_size times the number of atoms in the system is no less than N.
The key numb_batch in validate_data gives the number of batches of model validation. Note that the batches may not be from the same system

The section mixed_precision specifies the mixed precision settings, which will enable the mixed precision training workflow for DeePMD-kit. The keys are explained below:

output_prec precision used in the output tensors, only float32 is supported currently.
compute_prec precision used in the computing tensors, only float16 is supported currently. Note there are several limitations about mixed precision training:
Only se_e2_a type descriptor is supported by the mixed precision training workflow.
The precision of the embedding net and the fitting net are forced to be set to float32.

Other keys in the training section are explained below:

numb_steps The number of training steps.
seed The random seed for getting frames from the training data set.
disp_file The file for printing learning curve.
disp_freq The frequency of printing learning curve. Set in the unit of training steps
save_freq The frequency of saving checkpoint.

Options and environment variables

Several command line options can be passed to dp train, which can be checked with

$ 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

  --init-frz-model INIT_FRZ_MODEL
                        Initialize the training from the frozen model.
  --skip-neighbor-stat  Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled. (default: False)

--init-model model.ckpt, 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.

--init-frz-model frozen_model.pb, initializes the training with an existing model that is stored in frozen_model.pb.

--skip-neighbor-stat will skip calculating neighbor statistics if one is concerned about performance. Some features will be disabled.

To maximize the performance, one should follow FAQ: How to control the parallelism of a job to control the number of threads.

One can set other environmental variables:

Environment variables	Allowed value	Default value	Usage
DP_INTERFACE_PREC	`high`, `low`	`high`	Control high (double) or low (float) precision of training.
DP_AUTO_PARALLELIZATION	0, 1	0	Enable auto parallelization for CPU operators.
DP_JIT	0, 1	0	Enable JIT. Note that this option may either improve or decrease the performance. Requires TensorFlow supports JIT.

Adjust `sel` of a frozen model

One can use --init-frz-model features to adjust (increase or decrease) sel of a existing model. Firstly, one needs to adjust sel in input.json. For example, adjust from [46, 92] to [23, 46].

"model": {
	"descriptor": {
		"sel": [23, 46]
	}
}

To obtain the new model at once, numb_steps should be set to zero:

"training": {
	"numb_steps": 0
}

Then, one can initialize the training from the frozen model and freeze the new model at once:

dp train input.json --init-frz-model frozen_model.pb
dp freeze -o frozen_model_adjusted_sel.pb

Two models should give the same result when the input satisfies both constraints.

Note: At this time, this feature is only supported by se_e2_a descriptor with set_davg_true enabled, or hybrid composed of the above descriptors.

Training Parameters

Note

One can load, modify, and export the input file by using our effective web-based tool DP-GUI. All training parameters below can be set in DP-GUI. By clicking “SAVE JSON”, one can download the input file for furthur training.

model:

type: dict

argument path: model

type_map:

type: list, optional

argument path: model/type_map

A list of strings. Give the name to each type of atoms. It is noted that the number of atom type of training system must be less than 128 in a GPU environment. If not given, type.raw in each system should use the same type indexes, and type_map.raw will take no effect.

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.

data_stat_protect:

type: float, optional, default: 0.01

argument path: model/data_stat_protect

Protect parameter for atomic energy regression.

data_bias_nsample:

type: int, optional, default: 10

argument path: model/data_bias_nsample

The number of training samples in a system to compute and change the energy bias.

use_srtab:

type: str, optional

argument path: model/use_srtab

The table for the short-range pairwise interaction added on top of DP. The table is a text data file with (N_t + 1) * N_t / 2 + 1 columes. The first colume is the distance between atoms. The second to the last columes are energies for pairs of certain types. For example we have two atom types, 0 and 1. The columes from 2nd to 4th are for 0-0, 0-1 and 1-1 correspondingly.

smin_alpha:

type: float, optional

argument path: model/smin_alpha

The short-range tabulated interaction will be swithed according to the distance of the nearest neighbor. This distance is calculated by softmin. This parameter is the decaying parameter in the softmin. It is only required when use_srtab is provided.

sw_rmin:

type: float, optional

argument path: model/sw_rmin

The lower boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

sw_rmax:

type: float, optional

argument path: model/sw_rmax

The upper boundary of the interpolation between short-range tabulated interaction and DP. It is only required when use_srtab is provided.

type_embedding:

type: dict, optional

argument path: model/type_embedding

The type embedding.

neuron:

type: list, optional, default: [8]

argument path: model/type_embedding/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.

activation_function:

type: str, optional, default: tanh

argument path: model/type_embedding/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/type_embedding/resnet_dt

Whether to use a “Timestep” in the skip connection

precision:

type: str, optional, default: default

argument path: model/type_embedding/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/type_embedding/trainable

If the parameters in the embedding net are trainable

seed:

type: NoneType | int, optional, default: None

argument path: model/type_embedding/seed

Random seed for parameter initialization

descriptor:

type: dict

argument path: model/descriptor

The descriptor of atomic environment.

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key)

argument path: model/descriptor/type

possible choices: loc_frame, se_e2_a, se_e3, se_a_tpe, se_e2_r, hybrid, se_atten, se_a_mask

The type of the descritpor. See explanation below.

loc_frame: Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.

se_e2_a: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor.

se_e2_r: Used by the smooth edition of Deep Potential. Only the distance between atoms is used to construct the descriptor.

se_e3: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Three-body embedding will be used by this descriptor.

se_a_tpe: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Type embedding will be used by this descriptor.

se_atten: Used by the smooth edition of Deep Potential. The full relative coordinates are used to construct the descriptor. Attention mechanism will be used by this descriptor.

se_a_mask: Used by the smooth edition of Deep Potential. It can accept a variable number of atoms in a frame (Non-PBC system). aparam are required as an indicator matrix for the real/virtual sign of input atoms.

hybrid: Concatenate of a list of descriptors as a new descriptor.

When type is set to loc_frame:

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.

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.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[loc_frame]/rcut

The cut-off radius. The default value is 6.0

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 second 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]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.

When type is set to se_e2_a (or its alias se_a):

sel:

type: str | list, optional, default: auto

argument path: model/descriptor[se_e2_a]/sel

This parameter set the number of selected neighbors for each type of atom. It can be:

List[int]. 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. It is noted that the total sel value must be less than 4096 in a GPU environment.

str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[se_e2_a]/rcut

The cut-off radius.

rcut_smth:

type: float, optional, default: 0.5

argument path: model/descriptor[se_e2_a]/rcut_smth

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_e2_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.

axis_neuron:

type: int, optional, default: 4, alias: n_axis_neuron

argument path: model/descriptor[se_e2_a]/axis_neuron

Size of the submatrix of G (embedding matrix).

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_e2_a]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_a]/resnet_dt

Whether to use a “Timestep” in the skip connection

type_one_side:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_a]/type_one_side

If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_text{types}^2$ sets of embedding network parameters.

precision:

type: str, optional, default: default

argument path: model/descriptor[se_e2_a]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_e2_a]/trainable

If the parameters in the embedding net is trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_e2_a]/seed

Random seed for parameter initialization

exclude_types:

type: list, optional, default: []

argument path: model/descriptor[se_e2_a]/exclude_types

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_a]/set_davg_zero

Set the normalization average to zero. This option should be set when atom_ener in the energy fitting is used

When type is set to se_e3 (or its aliases se_at, se_a_3be, se_t):

sel:

type: str | list, optional, default: auto

argument path: model/descriptor[se_e3]/sel

This parameter set the number of selected neighbors for each type of atom. It can be:

List[int]. 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. It is noted that the total sel value must be less than 4096 in a GPU environment.

str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[se_e3]/rcut

The cut-off radius.

rcut_smth:

type: float, optional, default: 0.5

argument path: model/descriptor[se_e3]/rcut_smth

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_e3]/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.

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_e3]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_e3]/resnet_dt

Whether to use a “Timestep” in the skip connection

precision:

type: str, optional, default: default

argument path: model/descriptor[se_e3]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_e3]/trainable

If the parameters in the embedding net are trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_e3]/seed

Random seed for parameter initialization

set_davg_zero:

type: bool, optional, default: False

argument path: model/descriptor[se_e3]/set_davg_zero

Set the normalization average to zero. This option should be set when atom_ener in the energy fitting is used

When type is set to se_a_tpe (or its alias se_a_ebd):

sel:

type: str | list, optional, default: auto

argument path: model/descriptor[se_a_tpe]/sel

This parameter set the number of selected neighbors for each type of atom. It can be:

List[int]. 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. It is noted that the total sel value must be less than 4096 in a GPU environment.

str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[se_a_tpe]/rcut

The cut-off radius.

rcut_smth:

type: float, optional, default: 0.5

argument path: model/descriptor[se_a_tpe]/rcut_smth

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_a_tpe]/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.

axis_neuron:

type: int, optional, default: 4, alias: n_axis_neuron

argument path: model/descriptor[se_a_tpe]/axis_neuron

Size of the submatrix of G (embedding matrix).

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_a_tpe]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_a_tpe]/resnet_dt

Whether to use a “Timestep” in the skip connection

type_one_side:

type: bool, optional, default: False

argument path: model/descriptor[se_a_tpe]/type_one_side

If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_text{types}^2$ sets of embedding network parameters.

precision:

type: str, optional, default: default

argument path: model/descriptor[se_a_tpe]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_a_tpe]/trainable

If the parameters in the embedding net is trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_a_tpe]/seed

Random seed for parameter initialization

exclude_types:

type: list, optional, default: []

argument path: model/descriptor[se_a_tpe]/exclude_types

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero:

type: bool, optional, default: False

argument path: model/descriptor[se_a_tpe]/set_davg_zero

Set the normalization average to zero. This option should be set when atom_ener in the energy fitting is used

type_nchanl:

type: int, optional, default: 4

argument path: model/descriptor[se_a_tpe]/type_nchanl

number of channels for type embedding

type_nlayer:

type: int, optional, default: 2

argument path: model/descriptor[se_a_tpe]/type_nlayer

number of hidden layers of type embedding net

numb_aparam:

type: int, optional, default: 0

argument path: model/descriptor[se_a_tpe]/numb_aparam

dimension of atomic parameter. if set to a value > 0, the atomic parameters are embedded.

When type is set to se_e2_r (or its alias se_r):

sel:

type: str | list, optional, default: auto

argument path: model/descriptor[se_e2_r]/sel

This parameter set the number of selected neighbors for each type of atom. It can be:

List[int]. 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. It is noted that the total sel value must be less than 4096 in a GPU environment.

str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[se_e2_r]/rcut

The cut-off radius.

rcut_smth:

type: float, optional, default: 0.5

argument path: model/descriptor[se_e2_r]/rcut_smth

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_e2_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.

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_e2_r]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_r]/resnet_dt

Whether to use a “Timestep” in the skip connection

type_one_side:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_r]/type_one_side

If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_text{types}^2$ sets of embedding network parameters.

precision:

type: str, optional, default: default

argument path: model/descriptor[se_e2_r]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_e2_r]/trainable

If the parameters in the embedding net are trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_e2_r]/seed

Random seed for parameter initialization

exclude_types:

type: list, optional, default: []

argument path: model/descriptor[se_e2_r]/exclude_types

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero:

type: bool, optional, default: False

argument path: model/descriptor[se_e2_r]/set_davg_zero

Set the normalization average to zero. This option should be set when atom_ener in the energy fitting is used

When type is set to hybrid:

list:

type: list

argument path: model/descriptor[hybrid]/list

A list of descriptor definitions

When type is set to se_atten:

sel:

type: str | list | int, optional, default: auto

argument path: model/descriptor[se_atten]/sel

This parameter set the number of selected neighbors. Note that this parameter is a little different from that in other descriptors. Instead of separating each type of atoms, only the summation matters. And this number is highly related with the efficiency, thus one should not make it too large. Usually 200 or less is enough, far away from the GPU limitation 4096. It can be:

int. The maximum number of neighbor atoms to be considered. We recommend it to be less than 200.

List[int]. 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. Only the summation of sel[i] matters, and it is recommended to be less than 200. - str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

rcut:

type: float, optional, default: 6.0

argument path: model/descriptor[se_atten]/rcut

The cut-off radius.

rcut_smth:

type: float, optional, default: 0.5

argument path: model/descriptor[se_atten]/rcut_smth

Where to start smoothing. For example the 1/r term is smoothed from rcut to rcut_smth

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_atten]/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.

axis_neuron:

type: int, optional, default: 4, alias: n_axis_neuron

argument path: model/descriptor[se_atten]/axis_neuron

Size of the submatrix of G (embedding matrix).

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_atten]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_atten]/resnet_dt

Whether to use a “Timestep” in the skip connection

type_one_side:

type: bool, optional, default: False

argument path: model/descriptor[se_atten]/type_one_side

If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_text{types}^2$ sets of embedding network parameters.

precision:

type: str, optional, default: default

argument path: model/descriptor[se_atten]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_atten]/trainable

If the parameters in the embedding net is trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_atten]/seed

Random seed for parameter initialization

exclude_types:

type: list, optional, default: []

argument path: model/descriptor[se_atten]/exclude_types

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

set_davg_zero:

type: bool, optional, default: True

argument path: model/descriptor[se_atten]/set_davg_zero

Set the normalization average to zero. This option should be set when se_atten descriptor or atom_ener in the energy fitting is used

attn:

type: int, optional, default: 128

argument path: model/descriptor[se_atten]/attn

The length of hidden vectors in attention layers

attn_layer:

type: int, optional, default: 2

argument path: model/descriptor[se_atten]/attn_layer

The number of attention layers

attn_dotr:

type: bool, optional, default: True

argument path: model/descriptor[se_atten]/attn_dotr

Whether to do dot product with the normalized relative coordinates

attn_mask:

type: bool, optional, default: False

argument path: model/descriptor[se_atten]/attn_mask

Whether to do mask on the diagonal in the attention matrix

When type is set to se_a_mask:

sel:

type: str | list, optional, default: auto

argument path: model/descriptor[se_a_mask]/sel

This parameter sets the number of selected neighbors for each type of atom. It can be:

List[int]. 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. It is noted that the total sel value must be less than 4096 in a GPU environment.

str. Can be “auto:factor” or “auto”. “factor” is a float number larger than 1. This option will automatically determine the sel. In detail it counts the maximal number of neighbors with in the cutoff radius for each type of neighbor, then multiply the maximum by the “factor”. Finally the number is wraped up to 4 divisible. The option “auto” is equivalent to “auto:1.1”.

neuron:

type: list, optional, default: [10, 20, 40]

argument path: model/descriptor[se_a_mask]/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.

axis_neuron:

type: int, optional, default: 4, alias: n_axis_neuron

argument path: model/descriptor[se_a_mask]/axis_neuron

Size of the submatrix of G (embedding matrix).

activation_function:

type: str, optional, default: tanh

argument path: model/descriptor[se_a_mask]/activation_function

The activation function in the embedding net. Supported activation functions are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: False

argument path: model/descriptor[se_a_mask]/resnet_dt

Whether to use a “Timestep” in the skip connection

type_one_side:

type: bool, optional, default: False

argument path: model/descriptor[se_a_mask]/type_one_side

If true, the embedding network parameters vary by types of neighbor atoms only, so there will be $N_text{types}$ sets of embedding network parameters. Otherwise, the embedding network parameters vary by types of centric atoms and types of neighbor atoms, so there will be $N_text{types}^2$ sets of embedding network parameters.

exclude_types:

type: list, optional, default: []

argument path: model/descriptor[se_a_mask]/exclude_types

The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

precision:

type: str, optional, default: default

argument path: model/descriptor[se_a_mask]/precision

The precision of the embedding net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

trainable:

type: bool, optional, default: True

argument path: model/descriptor[se_a_mask]/trainable

If the parameters in the embedding net is trainable

seed:

type: NoneType | int, optional

argument path: model/descriptor[se_a_mask]/seed

Random seed for parameter initialization

fitting_net:

type: dict, optional

argument path: model/fitting_net

The fitting of physical properties.

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key), default: ener

argument path: model/fitting_net/type

possible choices: ener, dipole, polar

The type of the fitting. See explanation below.

ener: Fit an energy model (potential energy surface).

dipole: Fit an atomic dipole model. Global dipole labels or atomic dipole labels for all the selected atoms (see sel_type) should be provided by dipole.npy in each data system. The file either has number of frames lines and 3 times of number of selected atoms columns, or has number of frames lines and 3 columns. See loss parameter.

polar: Fit an atomic polarizability model. Global polarizazbility labels or atomic polarizability labels for all the selected atoms (see sel_type) should be provided by polarizability.npy in each data system. The file eith has number of frames lines and 9 times of number of selected atoms columns, or has number of frames lines and 9 columns. See loss parameter.

When type is set to ener:

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.

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.

neuron:

type: list, optional, default: [120, 120, 120], alias: n_neuron

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.

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”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

precision:

type: str, optional, default: default

argument path: model/fitting_net[ener]/precision

The precision of the fitting net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

resnet_dt:

type: bool, optional, default: True

argument path: model/fitting_net[ener]/resnet_dt

Whether to use a “Timestep” in the skip connection

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.

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.

seed:

type: NoneType | int, optional

argument path: model/fitting_net[ener]/seed

Random seed for parameter initialization of the fitting net

atom_ener:

type: list, optional, default: []

argument path: model/fitting_net[ener]/atom_ener

Specify the atomic energy in vacuum for each type

layer_name:

type: list, optional

argument path: model/fitting_net[ener]/layer_name

The name of the each layer. The length of this list should be equal to n_neuron + 1. If two layers, either in the same fitting or different fittings, have the same name, they will share the same neural network parameters. The shape of these layers should be the same. If null is given for a layer, parameters will not be shared.

use_aparam_as_mask:

type: bool, optional, default: False

argument path: model/fitting_net[ener]/use_aparam_as_mask

Whether to use the aparam as a mask in input.If True, the aparam will not be used in fitting net for embedding.When descrpt is se_a_mask, the aparam will be used as a mask to indicate the input atom is real/virtual. And use_aparam_as_mask should be set to True.

When type is set to dipole:

neuron:

type: list, optional, default: [120, 120, 120], alias: n_neuron

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.

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”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: True

argument path: model/fitting_net[dipole]/resnet_dt

Whether to use a “Timestep” in the skip connection

precision:

type: str, optional, default: default

argument path: model/fitting_net[dipole]/precision

The precision of the fitting net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

sel_type:

type: int | NoneType | list, optional, alias: dipole_type

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.

seed:

type: NoneType | int, optional

argument path: model/fitting_net[dipole]/seed

Random seed for parameter initialization of the fitting net

When type is set to polar:

neuron:

type: list, optional, default: [120, 120, 120], alias: n_neuron

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.

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”, “gelu_tf”, “None”, “none”. Note that “gelu” denotes the custom operator version, and “gelu_tf” denotes the TF standard version. If you set “None” or “none” here, no activation function will be used.

resnet_dt:

type: bool, optional, default: True

argument path: model/fitting_net[polar]/resnet_dt

Whether to use a “Timestep” in the skip connection

precision:

type: str, optional, default: default

argument path: model/fitting_net[polar]/precision

The precision of the fitting net parameters, supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”. Default follows the interface precision.

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.

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

shift_diag:

type: bool, optional, default: True

argument path: model/fitting_net[polar]/shift_diag

Whether to shift the diagonal of polar, which is beneficial to training. Default is true.

sel_type:

type: int | NoneType | list, optional, alias: pol_type

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.

seed:

type: NoneType | int, optional

argument path: model/fitting_net[polar]/seed

Random seed for parameter initialization of the fitting net

fitting_net_dict:

type: dict, optional

argument path: model/fitting_net_dict

The dictionary of multiple fitting nets in multi-task mode. Each fitting_net_dict[fitting_key] is the single definition of fitting of physical properties with user-defined name fitting_key.

modifier:

type: dict, optional

argument path: model/modifier

The modifier of model output.

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key)

argument path: model/modifier/type

possible choices: dipole_charge

The type of modifier. See explanation below.

-dipole_charge: Use WFCC to model the electronic structure of the system. Correct the long-range interaction

When type is set to dipole_charge:

model_name:

type: str

argument path: model/modifier[dipole_charge]/model_name

The name of the frozen dipole model file.

model_charge_map:

type: list

argument path: model/modifier[dipole_charge]/model_charge_map

The charge of the WFCC. The list length should be the same as the sel_type.

sys_charge_map:

type: list

argument path: model/modifier[dipole_charge]/sys_charge_map

The charge of real atoms. The list length should be the same as the type_map

ewald_beta:

type: float, optional, default: 0.4

argument path: model/modifier[dipole_charge]/ewald_beta

The splitting parameter of Ewald sum. Unit is A^-1

ewald_h:

type: float, optional, default: 1.0

argument path: model/modifier[dipole_charge]/ewald_h

The grid spacing of the FFT grid. Unit is A

compress:

type: dict, optional

argument path: model/compress

Model compression configurations

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key), default: se_e2_a

argument path: model/compress/type

possible choices: se_e2_a

The type of model compression, which should be consistent with the descriptor type.

When type is set to se_e2_a (or its alias se_a):

model_file:

type: str

argument path: model/compress[se_e2_a]/model_file

The input model file, which will be compressed by the DeePMD-kit.

table_config:

type: list

argument path: model/compress[se_e2_a]/table_config

The arguments of model compression, including extrapolate(scale of model extrapolation), stride(uniform stride of tabulation’s first and second table), and frequency(frequency of tabulation overflow check).

min_nbor_dist:

type: float

argument path: model/compress[se_e2_a]/min_nbor_dist

The nearest distance between neighbor atoms saved in the frozen model.

learning_rate:

type: dict, optional

argument path: learning_rate

The definitio of learning rate

scale_by_worker:

type: str, optional, default: linear

argument path: learning_rate/scale_by_worker

When parallel training or batch size scaled, how to alter learning rate. Valid values are linear`(default), `sqrt or none.

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key), default: exp

argument path: learning_rate/type

possible choices: exp

The type of the learning rate.

When type is set to exp:

start_lr:

type: float, optional, default: 0.001

argument path: learning_rate[exp]/start_lr

The learning rate the start of the training.

stop_lr:

type: float, optional, default: 1e-08

argument path: learning_rate[exp]/stop_lr

The desired learning rate at the end of the training.

decay_steps:

type: int, optional, default: 5000

argument path: learning_rate[exp]/decay_steps

The learning rate is decaying every this number of training steps.

learning_rate_dict:: type: dict, optional

argument path: learning_rate_dict

The dictionary of definitions of learning rates in multi-task mode. Each learning_rate_dict[fitting_key], with user-defined name fitting_key in model/fitting_net_dict, is the single definition of learning rate.

loss:

type: dict, optional

argument path: loss

The definition of loss function. The loss type should be set to tensor, ener or left unset.

Depending on the value of type, different sub args are accepted.

type:

type: str (flag key), default: ener

argument path: loss/type

possible choices: ener, tensor

The type of the loss. When the fitting type is ener, the loss type should be set to ener or left unset. When the fitting type is dipole or polar, the loss type should be set to tensor.

When type is set to ener:

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.

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.

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.

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.

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.

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.

start_pref_ae:

type: float | int, optional, default: 0.0

argument path: loss[ener]/start_pref_ae

The prefactor of atom_ener loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the atom_ener label should be provided by file atom_ener.npy in each data system. If both start_pref_atom_ener and limit_pref_atom_ener are set to 0, then the atom_ener will be ignored.

limit_pref_ae:

type: float | int, optional, default: 0.0

argument path: loss[ener]/limit_pref_ae

The prefactor of atom_ener loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.

start_pref_pf:

type: float | int, optional, default: 0.0

argument path: loss[ener]/start_pref_pf

The prefactor of atom_pref loss at the start of the training. Should be larger than or equal to 0. If set to none-zero value, the atom_pref label should be provided by file atom_pref.npy in each data system. If both start_pref_atom_pref and limit_pref_atom_pref are set to 0, then the atom_pref will be ignored.

limit_pref_pf:

type: float | int, optional, default: 0.0

argument path: loss[ener]/limit_pref_pf

The prefactor of atom_pref loss at the limit of the training, Should be larger than or equal to 0. i.e. the training step goes to infinity.

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.

enable_atom_ener_coeff:

type: bool, optional, default: False

argument path: loss[ener]/enable_atom_ener_coeff

If true, the energy will be computed as sum_i c_i E_i. c_i should be provided by file atom_ener_coeff.npy in each data system, otherwise it’s 1.

When type is set to tensor:

pref:

type: float | int

argument path: loss[tensor]/pref

The prefactor of the weight of global loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to global label, i.e. ‘polarizability.npy` or dipole.npy, whose shape should be #frames x [9 or 3]. If it’s larger than 0.0, this npy should be included.

pref_atomic:

type: float | int

argument path: loss[tensor]/pref_atomic

The prefactor of the weight of atomic loss. It should be larger than or equal to 0. If controls the weight of loss corresponding to atomic label, i.e. atomic_polarizability.npy or atomic_dipole.npy, whose shape should be #frames x ([9 or 3] x #selected atoms). If it’s larger than 0.0, this npy should be included. Both pref and pref_atomic should be provided, and either can be set to 0.0.

loss_dict:: type: dict, optional

argument path: loss_dict

The dictionary of definitions of multiple loss functions in multi-task mode. Each loss_dict[fitting_key], with user-defined name fitting_key in model/fitting_net_dict, is the single definition of loss function, whose type should be set to tensor, ener or left unset.

training:

type: dict

argument path: training

The training options.

training_data:

type: dict, optional

argument path: training/training_data

Configurations of training data.

systems:

type: str | list

argument path: training/training_data/systems

The data systems for training. This key can be provided with a list that 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.

set_prefix:

type: str, optional, default: set

argument path: training/training_data/set_prefix

The prefix of the sets in the systems.

batch_size:

type: int | str | list, optional, default: auto

argument path: training/training_data/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 use the same batch size.

string “auto”: automatically determines the batch size so 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 so that the batch_size times the number of atoms in the system is no less than N.

string “mixed:N”: the batch data will be sampled from all systems and merged into a mixed system with the batch size N. Only support the se_atten descriptor.

auto_prob:

type: str, optional, default: prob_sys_size, alias: auto_prob_style

argument path: training/training_data/auto_prob

Determine the probability of systems automatically. The method is assigned by this key and can be

“prob_uniform” : the probability all the systems are equal, namely 1.0/self.get_nsystems()

“prob_sys_size” : the probability of a system is proportional to the number of batches in the system

“prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;…” : the list of systems is devided into blocks. A block is specified by stt_idx:end_idx:weight, where stt_idx is the starting index of the system, end_idx is then ending (not including) index of the system, the probabilities of the systems in this block sums up to weight, and the relatively probabilities within this block is proportional to the number of batches in the system.

sys_probs:

type: NoneType | list, optional, default: None, alias: sys_weights

argument path: training/training_data/sys_probs

A list of float if specified. Should be of the same length as systems, specifying the probability of each system.

validation_data:

type: dict | NoneType, optional, default: None

argument path: training/validation_data

Configurations of validation data. Similar to that of training data, except that a numb_btch argument may be configured

systems:

type: str | list

argument path: training/validation_data/systems

The data systems for validation. This key can be provided with a list that 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.

set_prefix:

type: str, optional, default: set

argument path: training/validation_data/set_prefix

The prefix of the sets in the systems.

batch_size:

type: int | str | list, optional, default: auto

argument path: training/validation_data/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 use the same batch size.

string “auto”: automatically determines the batch size so 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 so that the batch_size times the number of atoms in the system is no less than N.

auto_prob:

type: str, optional, default: prob_sys_size, alias: auto_prob_style

argument path: training/validation_data/auto_prob

Determine the probability of systems automatically. The method is assigned by this key and can be

“prob_uniform” : the probability all the systems are equal, namely 1.0/self.get_nsystems()

“prob_sys_size” : the probability of a system is proportional to the number of batches in the system

“prob_sys_size;stt_idx:end_idx:weight;stt_idx:end_idx:weight;…” : the list of systems is devided into blocks. A block is specified by stt_idx:end_idx:weight, where stt_idx is the starting index of the system, end_idx is then ending (not including) index of the system, the probabilities of the systems in this block sums up to weight, and the relatively probabilities within this block is proportional to the number of batches in the system.

sys_probs:

type: NoneType | list, optional, default: None, alias: sys_weights

argument path: training/validation_data/sys_probs

A list of float if specified. Should be of the same length as systems, specifying the probability of each system.

numb_btch:

type: int, optional, default: 1, alias: numb_batch

argument path: training/validation_data/numb_btch

An integer that specifies the number of batches to be sampled for each validation period.

mixed_precision:

type: dict, optional

argument path: training/mixed_precision

Configurations of mixed precision.

output_prec:

type: str, optional, default: float32

argument path: training/mixed_precision/output_prec

The precision for mixed precision params. ” “The trainable variables precision during the mixed precision training process, ” “supported options are float32 only currently.

compute_prec:

type: str

argument path: training/mixed_precision/compute_prec

The precision for mixed precision compute. ” “The compute precision during the mixed precision training process, “” “supported options are float16 and bfloat16 currently.

numb_steps:

type: int, alias: stop_batch

argument path: training/numb_steps

Number of training batch. Each training uses one batch of data.

seed:

type: NoneType | int, optional

argument path: training/seed

The random seed for getting frames from the training data set.

disp_file:

type: str, optional, default: lcurve.out

argument path: training/disp_file

The file for printing learning curve.

disp_freq:

type: int, optional, default: 1000

argument path: training/disp_freq

The frequency of printing learning curve.

save_freq:

type: int, optional, default: 1000

argument path: training/save_freq

The frequency of saving check point.

save_ckpt:

type: str, optional, default: model.ckpt

argument path: training/save_ckpt

The file name of saving check point.

disp_training:

type: bool, optional, default: True

argument path: training/disp_training

Displaying verbose information during training.

time_training:

type: bool, optional, default: True

argument path: training/time_training

Timing durining training.

profiling:

type: bool, optional, default: False

argument path: training/profiling

Profiling during training.

profiling_file:

type: str, optional, default: timeline.json

argument path: training/profiling_file

Output file for profiling.

enable_profiler:

type: bool, optional, default: False

argument path: training/enable_profiler

Enable TensorFlow Profiler (available in TensorFlow 2.3) to analyze performance. The log will be saved to tensorboard_log_dir.

tensorboard:

type: bool, optional, default: False

argument path: training/tensorboard

Enable tensorboard

tensorboard_log_dir:

type: str, optional, default: log

argument path: training/tensorboard_log_dir

The log directory of tensorboard outputs

tensorboard_freq:

type: int, optional, default: 1

argument path: training/tensorboard_freq

The frequency of writing tensorboard events.

data_dict:

type: dict, optional

argument path: training/data_dict

The dictionary of multi DataSystems in multi-task mode. Each data_dict[fitting_key], with user-defined name fitting_key in model/fitting_net_dict, contains training data and optional validation data definitions.

fitting_weight:

type: dict, optional

argument path: training/fitting_weight

Each fitting_weight[fitting_key], with user-defined name fitting_key in model/fitting_net_dict, is the training weight of fitting net fitting_key. Fitting nets with higher weights will be selected with higher probabilities to be trained in one step. Weights will be normalized and minus ones will be ignored. If not set, each fitting net will be equally selected when training.

nvnmd:

type: dict, optional

argument path: nvnmd

The nvnmd options.

net_size:

type: int

argument path: nvnmd/net_size

configuration the number of nodes of fitting_net, just can be set as 128

map_file:

type: str

argument path: nvnmd/map_file

A file containing the mapping tables to replace the calculation of embedding nets

config_file:

type: str

argument path: nvnmd/config_file

A file containing the parameters about how to implement the model in certain hardware

weight_file:

type: str

argument path: nvnmd/weight_file

a *.npy file containing the weights of the model

enable:

type: bool

argument path: nvnmd/enable

enable the nvnmd training

restore_descriptor:

type: bool

argument path: nvnmd/restore_descriptor

enable to restore the parameter of embedding_net from weight.npy

restore_fitting_net:

type: bool

argument path: nvnmd/restore_fitting_net

enable to restore the parameter of fitting_net from weight.npy

quantize_descriptor:

type: bool

argument path: nvnmd/quantize_descriptor

enable the quantizatioin of descriptor

quantize_fitting_net:

type: bool

argument path: nvnmd/quantize_fitting_net

enable the quantizatioin of fitting_net

Parallel training

Currently, parallel training is enabled in a synchronized way with help of Horovod. Depending on the number of training processes (according to MPI context) and the number of GPU cards available, DeePMD-kit will decide whether to launch the training in parallel (distributed) mode or in serial mode. Therefore, no additional options are specified in your JSON/YAML input file.

Tuning learning rate

Horovod works in the data-parallel mode, resulting in a larger global batch size. For example, the real batch size is 8 when batch_size is set to 2 in the input file and you launch 4 workers. Thus, learning_rate is automatically scaled by the number of workers for better convergence. Technical details of such heuristic rule are discussed at Accurate, Large Minibatch SGD: Training ImageNet in 1 Hour.

The number of decay steps required to achieve the same accuracy can decrease by the number of cards (e.g., 1/2 of steps in the above case), but needs to be scaled manually in the input file.

In some cases, it won’t work well when scaling the learning rate by worker count in a linear way. Then you can try sqrt or none by setting argument scale_by_worker like below.

    "learning_rate" :{
        "scale_by_worker": "none",
        "type": "exp"
    }

Scaling test

Testing examples/water/se_e2_a on an 8-GPU host, linear acceleration can be observed with the increasing number of cards.

Num of GPU cards	Seconds every 100 samples	Samples per second	Speed up
1	1.4515	68.89	1.00
2	1.5962	62.65*2	1.82
4	1.7635	56.71*4	3.29
8	1.7267	57.91*8	6.72

How to use

Training workers can be launched with horovodrun. The following command launches 4 processes on the same host:

CUDA_VISIBLE_DEVICES=4,5,6,7 horovodrun -np 4 \
    dp train --mpi-log=workers input.json

Need to mention, the environment variable CUDA_VISIBLE_DEVICES must be set to control parallelism on the occupied host where one process is bound to one GPU card.

To maximize the performance, one should follow FAQ: How to control the parallelism of a job to control the number of threads.

When using MPI with Horovod, horovodrun is a simple wrapper around mpirun. In the case where fine-grained control over options is passed to mpirun, mpirun can be invoked directly, and it will be detected automatically by Horovod, e.g.,

CUDA_VISIBLE_DEVICES=4,5,6,7 mpirun -l -launcher=fork -hosts=localhost -np 4 \
    dp train --mpi-log=workers input.json

this is sometimes necessary for an HPC environment.

Whether distributed workers are initiated can be observed in the “Summary of the training” section in the log (world size > 1, and distributed).

[0] DEEPMD INFO    ---Summary of the training---------------------------------------
[0] DEEPMD INFO    distributed
[0] DEEPMD INFO    world size:           4
[0] DEEPMD INFO    my rank:              0
[0] DEEPMD INFO    node list:            ['exp-13-57']
[0] DEEPMD INFO    running on:           exp-13-57
[0] DEEPMD INFO    computing device:     gpu:0
[0] DEEPMD INFO    CUDA_VISIBLE_DEVICES: 0,1,2,3
[0] DEEPMD INFO    Count of visible GPU: 4
[0] DEEPMD INFO    num_intra_threads:    0
[0] DEEPMD INFO    num_inter_threads:    0
[0] DEEPMD INFO    -----------------------------------------------------------------

Logging

What’s more, 2 command-line arguments are defined to control the logging behavior when performing parallel training with MPI.

optional arguments:
  -l LOG_PATH, --log-path LOG_PATH
                        set log file to log messages to disk, if not
                        specified, the logs will only be output to console
                        (default: None)
  -m {master,collect,workers}, --mpi-log {master,collect,workers}
                        Set the manner of logging when running with MPI.
                        'master' logs only on main process, 'collect'
                        broadcasts logs from workers to master and 'workers'
                        means each process will output its own log (default:
                        master)

Multi-task training

Perform the multi-task training

Training on multiple data sets (each data set contains several data systems) can be performed in multi-task mode, with one common descriptor and multiple specific fitting nets for each data set. One can simply switch the following parameters in training input script to perform multi-task mode:

fitting_net –> fitting_net_dict, each key of which can be one individual fitting net.
training_data, validation_data –> data_dict, each key of which can be one individual data set contains several data systems for corresponding fitting net, the keys must be consistent with those in fitting_net_dict.
loss –> loss_dict, each key of which can be one individual loss setting for corresponding fitting net, the keys must be consistent with those in fitting_net_dict, if not set, the corresponding fitting net will use the default loss.
(Optional) fitting_weight, each key of which can be a non-negative integer or float, deciding the chosen probability for corresponding fitting net in training, if not set or invalid, the corresponding fitting net will not be used.

The training procedure will automatically choose single-task or multi-task mode, based on the above parameters. Note that parameters of single-task mode and multi-task mode can not be mixed.

An example input for training energy and dipole in water system can be found here: multi-task input on water.

The supported descriptors for multi-task mode are listed:

The supported fitting nets for multi-task mode are listed:

The output of dp freeze command in multi-task mode can be seen in freeze command.

Initialization from pretrained multi-task model

For advance training in multi-task mode, one can first train the descriptor on several upstream datasets and then transfer it on new downstream ones with newly added fitting nets. At the second step, you can also inherit some fitting nets trained on upstream datasets, by merely adding fitting net keys in fitting_net_dict and optional fitting net weights in fitting_weight.

Take multi-task input on water again for example. You can first train a multi-task model using input script with the following model part:

    "model": {
        "type_map": ["O", "H"],
        "descriptor": {
            "type":     "se_e2_a",
            "sel":      [46, 92],
            "rcut_smth":    0.5,
            "rcut":     6.0,
            "neuron":       [25, 50, 100],
        },
        "fitting_net_dict": {
            "water_dipole": {
                "type":         "dipole",
                "neuron":       [100, 100, 100],
            },
            "water_ener": {
                "neuron":       [240, 240, 240],
                "resnet_dt":    true,
            }
        },
    }

After training, you can freeze this multi-task model into one unit graph:

$ dp freeze -o graph.pb --united-model

Then if you want to transfer the trained descriptor and some fitting nets (take water_ener for example) to newly added datasets with new fitting net water_ener_2, you can modify the model part of the new input script in a more simplified way:

    "model": {
        "type_map": ["O", "H"],
        "descriptor": {},
        "fitting_net_dict": {
            "water_ener": {},
            "water_ener_2": {
                "neuron":       [240, 240, 240],
                "resnet_dt":    true,
            }
        },
    }

It will autocomplete the configurations according to the frozen graph.

Note that for newly added fitting net keys, other parts in the input script, including data_dict and loss_dict (optionally fitting_weight), should be set explicitly. While for old fitting net keys, it will inherit the old configurations if not set.

Finally, you can perform the modified multi-task training from the frozen model with command:

$ dp train input.json --init_frz_model graph.pb

TensorBoard Usage

TensorBoard provides the visualization and tooling needed for machine learning experimentation. Full instructions for TensorBoard can be found here.

Highlighted features

DeePMD-kit can now use most of the interesting features enabled by TensorBoard!

Tracking and visualizing metrics, such as l2_loss, l2_energy_loss and l2_force_loss
Visualizing the model graph (ops and layers)
Viewing histograms of weights, biases, or other tensors as they change over time.
Viewing summaries of trainable variables

How to use Tensorboard with DeePMD-kit

Before running TensorBoard, make sure you have generated summary data in a log directory by modifying the input script, setting tensorboard to true in the training subsection will enable the TensorBoard data analysis. eg. water_se_a.json.

    "training" : {
	"systems":	["../data/"],
	"set_prefix":	"set",
	"stop_batch":	1000000,
	"batch_size":	1,

	"seed":		1,

	"_comment": " display and restart",
	"_comment": " frequencies counted in batch",
	"disp_file":	"lcurve.out",
	"disp_freq":	100,
	"numb_test":	10,
	"save_freq":	1000,
	"save_ckpt":	"model.ckpt",

	"disp_training":true,
	"time_training":true,
	"tensorboard":	true,
	"tensorboard_log_dir":"log",
	"tensorboard_freq": 1000,
	"profiling":	false,
	"profiling_file":"timeline.json",
	"_comment":	"that's all"
    }

Once you have event files, run TensorBoard and provide the log directory. This should print that TensorBoard has started. Next, connect to http://tensorboard_server_ip:6006.

TensorBoard requires a logdir to read logs from. For info on configuring TensorBoard, run TensorBoard –help. One can easily change the log name with “tensorboard_log_dir” and the sampling frequency with “tensorboard_freq”.

tensorboard --logdir path/to/logs

Examples

Tracking and visualizing loss metrics(red:train, blue:test)

l2 loss

l2 energy loss

l2 force loss

Visualizing DeePMD-kit model graph

DeePMD-kit graph

Viewing histograms of weights, biases, or other tensors as they change over time

DeePMD-kit histograms

DeePMD-kit distribution

Viewing summaries of trainable variables

DeePMD-kit scalar

Attention

Allowing the tensorboard analysis will takes extra execution time.(eg, 15% increasing @Nvidia GTX 1080Ti double precision with default water sample)

TensorBoard can be used in Google Chrome or Firefox. Other browsers might work, but there may be bugs or performance issues.

Known limitations of using GPUs

If you use DeePMD-kit in a GPU environment, the acceptable value range of some variables is additionally restricted compared to the CPU environment due to the software’s GPU implementations:

The number of atom types of a given system must be less than 128.
The maximum distance between an atom and its neighbors must be less than 128. It can be controlled by setting the rcut value of training parameters.
Theoretically, the maximum number of atoms that a single GPU can accept is about 10,000,000. However, this value is limited by the GPU memory size currently, usually within 1000,000 atoms even in the model compression mode.
The total sel value of training parameters(in model/descriptor section) must be less than 4096.
The size of the last layer of the embedding net must be less than 1024 during the model compression process.

Finetune the pretrained model

Pretraining-and-finetuning is a widely used approach in other fields such as Computer Vision (CV) or Natural Language Processing (NLP) to vastly reduce the training cost, while it’s not trivial in potential models. Compositions and configurations of data samples or even computational parameters in upstream software (such as VASP) may be different between the pretrained and target datasets, leading to energy shifts or other diversities of training data.

Recently the emerging of methods such as DPA-1 has brought us to a new stage where we can perform similar pretraining-finetuning approaches. DPA-1 can hopefully learn the common knowledge in the pretrained dataset (especially the force information) and thus reduce the computational cost in downstream training tasks. If you have a pretrained model pretrained.pb (here we support models using se_atten descriptor and ener fitting net) on a large dataset (for example, OC2M in DPA-1 paper), a finetuning strategy can be performed by simply running:

$ dp train input.json --finetune pretrained.pb

The command above will change the energy bias in the last layer of the fitting net in pretrained.pb, according to the training dataset in input.json.

Warning

Note that the elements in the training dataset must be contained in the pretrained dataset.

The finetune procedure will inherit the model structures in pretrained.pb, and thus it will ignore the model parameters in input.json, such as descriptor, fitting_net, type_embedding and type_map. However, you can still set the trainable parameters in each part of input.json to control the training procedure.

To obtain a more simplified script, for example, you can change the model part in input.json to perform finetuning:

    "model": {
        "type_map":     ["O", "H"],
        "type_embedding": {"trainable":  true},
        "descriptor" :  {},
        "fitting_net" : {}
    }

Freeze and Compress

Freeze a model

The trained neural network is extracted from a checkpoint and dumped into a protobuf(.pb) file. This process is called “freezing” a model. The idea and part of our code are from Morgan. To freeze a model, typically one does

$ dp freeze -o graph.pb

in the folder where the model is trained. The output model is called graph.pb.

In multi-task mode:

This process will in default output several models, each of which contains the common descriptor and one of the user-defined fitting nets in fitting_net_dict, let’s name it fitting_key, together frozen in graph_{fitting_key}.pb. Those frozen models are exactly the same as single-task output with fitting net fitting_key.
If you add --united-model option in this situation, the total multi-task model will be frozen into one unit graph.pb, which is mainly for multi-task initialization and can not be used directly for inference.

Compress a model

Once the frozen model is obtained from DeePMD-kit, we can get the neural network structure and its parameters (weights, biases, etc.) from the trained model, and compress it in the following way:

dp compress -i graph.pb -o graph-compress.pb

where -i gives the original frozen model, -o gives the compressed model. Several other command line options can be passed to dp compress, which can be checked with

$ dp compress --help

An explanation will be provided

usage: dp compress [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
                   [-m {master,collect,workers}] [-i INPUT] [-o OUTPUT]
                   [-s STEP] [-e EXTRAPOLATE] [-f FREQUENCY]
                   [-c CHECKPOINT_FOLDER]

optional arguments:
  -h, --help            show this help message and exit
  -v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}, --log-level {DEBUG,3,INFO,2,WARNING,1,ERROR,0}
                        set verbosity level by string or number, 0=ERROR,
                        1=WARNING, 2=INFO and 3=DEBUG (default: INFO)
  -l LOG_PATH, --log-path LOG_PATH
                        set log file to log messages to disk, if not
                        specified, the logs will only be output to console
                        (default: None)
  -m {master,collect,workers}, --mpi-log {master,collect,workers}
                        Set the manner of logging when running with MPI.
                        'master' logs only on main process, 'collect'
                        broadcasts logs from workers to master and 'workers'
                        means each process will output its own log (default:
                        master)
  -i INPUT, --input INPUT
                        The original frozen model, which will be compressed by
                        the code (default: frozen_model.pb)
  -o OUTPUT, --output OUTPUT
                        The compressed model (default:
                        frozen_model_compressed.pb)
  -s STEP, --step STEP  Model compression uses fifth-order polynomials to
                        interpolate the embedding-net. It introduces two
                        tables with different step size to store the
                        parameters of the polynomials. The first table covers
                        the range of the training data, while the second table
                        is an extrapolation of the training data. The domain
                        of each table is uniformly divided by a given step
                        size. And the step(parameter) denotes the step size of
                        the first table and the second table will use 10 *
                        step as it's step size to save the memory. Usually the
                        value ranges from 0.1 to 0.001. Smaller step means
                        higher accuracy and bigger model size (default: 0.01)
  -e EXTRAPOLATE, --extrapolate EXTRAPOLATE
                        The domain range of the first table is automatically
                        detected by the code: [d_low, d_up]. While the second
                        table ranges from the first table's upper
                        boundary(d_up) to the extrapolate(parameter) * d_up:
                        [d_up, extrapolate * d_up] (default: 5)
  -f FREQUENCY, --frequency FREQUENCY
                        The frequency of tabulation overflow check(Whether the
                        input environment matrix overflow the first or second
                        table range). By default do not check the overflow
                        (default: -1)
  -c CHECKPOINT_FOLDER, --checkpoint-folder CHECKPOINT_FOLDER
                        path to checkpoint folder (default: .)
  -t TRAINING_SCRIPT, --training-script TRAINING_SCRIPT
                        The training script of the input frozen model
                        (default: None)

Parameter explanation

Model compression, which includes tabulating the embedding net. The table is composed of fifth-order polynomial coefficients and is assembled from two sub-tables. For model descriptor with se_e2_a type, the first sub-table takes the stride(parameter) as its uniform stride, while the second sub-table takes 10 * stride as its uniform stride; For model descriptor with se_e3 type, the first sub-table takes 10 * stride as it’s uniform stride, while the second sub-table takes 100 * stride as it’s uniform stride. The range of the first table is automatically detected by DeePMD-kit, while the second table ranges from the first table’s upper boundary(upper) to the extrapolate(parameter) * upper. Finally, we added a check frequency parameter. It indicates how often the program checks for overflow(if the input environment matrix overflows the first or second table range) during the MD inference.

Justification of model compression

Model compression, with little loss of accuracy, can greatly speed up MD inference time. According to different simulation systems and training parameters, the speedup can reach more than 10 times at both CPU and GPU devices. At the same time, model compression can greatly change memory usage, reducing as much as 20 times under the same hardware conditions.

Acceptable original model version

The model compression interface requires the version of DeePMD-kit used in the original model generation should be 2.0.0-alpha.0 or above. If one has a frozen 1.2 or 1.3 model, one can upgrade it through the dp convert-from interface. (eg: dp convert-from 1.2/1.3 -i old_frozen_model.pb -o new_frozen_model.pb)

Acceptable descriptor type

Descriptors with se_e2_a, se_e3, and se_e2_r types are supported by the model compression feature. Hybrid mixed with the above descriptors is also supported.

Available activation functions for descriptor:

tanh
gelu
relu
relu6
softplus
sigmoid

Test

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

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

$ 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 prefix to files where details of energy, force and virial accuracy/accuracy per atom will be written
  -a, --atomic          Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)

Calculate Model Deviation

One can also use a subcommand to calculate the deviation of predicted forces or virials for a bunch of models in the following way:

dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out

where -m specifies graph files to be calculated, -s gives the data to be evaluated, -o the file to which model deviation results is dumped. Here is more information on this sub-command:

usage: dp model-devi [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}]
                     [-l LOG_PATH] [-m MODELS [MODELS ...]] [-s SYSTEM]
                     [-S SET_PREFIX] [-o OUTPUT] [-f FREQUENCY] [-i ITEMS]

optional arguments:
  -h, --help            show this help message and exit
  -v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}, --log-level {DEBUG,3,INFO,2,WARNING,1,ERROR,0}
                        set verbosity level by string or number, 0=ERROR,
                        1=WARNING, 2=INFO and 3=DEBUG (default: INFO)
  -l LOG_PATH, --log-path LOG_PATH
                        set log file to log messages to disk, if not
                        specified, the logs will only be output to console
                        (default: None)
  -m MODELS [MODELS ...], --models MODELS [MODELS ...]
                        Frozen models file to import (default:
                        ['graph.000.pb', 'graph.001.pb', 'graph.002.pb',
                        'graph.003.pb'])
  -s SYSTEM, --system SYSTEM
                        The system directory, not support recursive detection.
                        (default: .)
  -S SET_PREFIX, --set-prefix SET_PREFIX
                        The set prefix (default: set)
  -o OUTPUT, --output OUTPUT
                        The output file for results of model deviation
                        (default: model_devi.out)
  -f FREQUENCY, --frequency FREQUENCY
                        The trajectory frequency of the system (default: 1)

For more details concerning the definition of model deviation and its application, please refer to Yuzhi Zhang, Haidi Wang, Weijie Chen, Jinzhe Zeng, Linfeng Zhang, Han Wang, and Weinan E, DP-GEN: A concurrent learning platform for the generation of reliable deep learning based potential energy models, Computer Physics Communications, 2020, 253, 107206.

Inference

Note that the model for inference is required to be compatible with the DeePMD-kit package. See Model compatibility for details.

Python interface

One may use the python interface of DeePMD-kit for model inference, an example is given as follows

from deepmd.infer import DeepPot
import numpy as np

dp = DeepPot("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.

Furthermore, one can use the python interface to calculate model deviation.

from deepmd.infer import calc_model_devi
from deepmd.infer import DeepPot as DP
import numpy as np

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]
graphs = [DP("graph.000.pb"), DP("graph.001.pb")]
model_devi = calc_model_devi(coord, cell, atype, graphs)

Note that if the model inference or model deviation is performed cyclically, one should avoid calling the same model multiple times. Otherwise, tensorFlow will never release the memory and this may lead to an out-of-memory (OOM) error.

C/C++ interface

C++ interface

The C++ interface of DeePMD-kit is also available for the model interface, which is considered faster than the Python interface. An example infer_water.cpp is given below:

#include "deepmd/DeepPot.h"

int main(){
  deepmd::DeepPot dp ("graph.pb");
  std::vector<double > coord = {1., 0., 0., 0., 0., 1.5, 1. ,0. ,3.};
  std::vector<double > cell = {10., 0., 0., 0., 10., 0., 0., 0., 10.};
  std::vector<int > atype = {1, 0, 1};
  double e;
  std::vector<double > f, v;
  dp.compute (e, f, v, coord, atype, cell);
}

where e, f and v are predicted energy, force and virial of the system, respectively. See deepmd::DeepPot for details.

You can compile infer_water.cpp using gcc:

gcc infer_water.cpp -L $deepmd_root/lib -L $tensorflow_root/lib -I $deepmd_root/include -Wl,--no-as-needed -ldeepmd_cc -lstdc++ -ltensorflow_cc -Wl,-rpath=$deepmd_root/lib -Wl,-rpath=$tensorflow_root/lib -o infer_water

and then run the program:

./infer_water

C interface

Although C is harder to write, the C library will not be affected by different versions of C++ compilers.

An example infer_water.c is given below:

#include <stdio.h>
#include <stdlib.h>
#include "deepmd/c_api.h"

int main(){
  const char* model = "graph.pb";
  double coord[] = {1., 0., 0., 0., 0., 1.5, 1. ,0. ,3.};
  double cell[] = {10., 0., 0., 0., 10., 0., 0., 0., 10.};
  int atype[] = {1, 0, 1};
  // init C pointers with given memory
  double* e = malloc(sizeof(*e));
  double* f = malloc(sizeof(*f) * 9); // natoms * 3
  double* v = malloc(sizeof(*v) * 9);
  double* ae = malloc(sizeof(*ae) * 9); // natoms
  double* av = malloc(sizeof(*av) * 27); // natoms * 9
  // DP model
  DP_DeepPot* dp = DP_NewDeepPot(model);
  DP_DeepPotCompute (dp, 3, coord, atype, cell, e, f, v, ae, av);
  // print results
  printf("energy: %f\n", *e);
  for (int ii = 0; ii < 9; ++ii)
    printf("force[%d]: %f\n", ii, f[ii]);
  for (int ii = 0; ii < 9; ++ii)
    printf("force[%d]: %f\n", ii, v[ii]);
  // free memory
  free(e);
  free(f);
  free(v);
  free(ae);
  free(av);
  free(dp);
}

where e, f and v are predicted energy, force and virial of the system, respectively. ae and av are atomic energy and atomic virials, respectively. See DP_DeepPotCompute() for details.

You can compile infer_water.c using gcc:

gcc infer_water.c -L $deepmd_root/lib -L $tensorflow_root/lib -I $deepmd_root/include -Wl,--no-as-needed -ldeepmd_c -Wl,-rpath=$deepmd_root/lib -Wl,-rpath=$tensorflow_root/lib -o infer_water

and then run the program:

./infer_water

Header-only C++ library interface (recommended)

The header-only C++ library is built based on the C library. Thus, it has the same ABI compatibility as the C library but provides a powerful C++ interface. To use it, include deepmd/deepmd.hpp.

#include "deepmd/deepmd.hpp"

int main(){
  deepmd::hpp::DeepPot dp ("graph.pb");
  std::vector<double > coord = {1., 0., 0., 0., 0., 1.5, 1. ,0. ,3.};
  std::vector<double > cell = {10., 0., 0., 0., 10., 0., 0., 0., 10.};
  std::vector<int > atype = {1, 0, 1};
  double e;
  std::vector<double > f, v;
  dp.compute (e, f, v, coord, atype, cell);
}

Note that the feature of the header-only C++ library is still limited compared to the original C++ library. See deepmd::hpp::DeepPot for details.

You can compile infer_water_hpp.cpp using gcc:

gcc infer_water_hpp.hpp -L $deepmd_root/lib -L $tensorflow_root/lib -I $deepmd_root/include -Wl,--no-as-needed -ldeepmd_c -Wl,-rpath=$deepmd_root/lib -Wl,-rpath=$tensorflow_root/lib -o infer_water_hpp

and then run the program:

./infer_water_hpp

In some cases, one may want to pass the custom neighbor list instead of the native neighbor list. The above code can be revised as follows:

  // neighbor list
  std::vector<std::vector<int >> nlist_vec = {
    {1, 2},
    {0, 2},
    {0, 1}
    };
  std::vector<int> ilist(3), numneigh(3);
  std::vector<int*> firstneigh(3);
  InputNlist nlist(3, &ilist[0], &numneigh[0], &firstneigh[0]);
  convert_nlist(nlist, nlist_vec);
  dp.compute (e, f, v, coord, atype, cell, 0, nlist, 0);

Here, nlist_vec means the neighbors of atom 0 are atom 1 and atom 2, the neighbors of atom 1 are atom 0 and atom 2, and the neighbors of atom 2 are atom 0 and atom 1.

Command line interface

DeePMD-kit: A deep learning package for many-body potential energy representation and molecular dynamics

usage: dp [-h] [--version]
          {config,transfer,train,freeze,test,compress,doc-train-input,model-devi,convert-from,neighbor-stat,train-nvnmd}
          ...

Named Arguments

--version: show program’s version number and exit

Valid subcommands

command: Possible choices: config, transfer, train, freeze, test, compress, doc-train-input, model-devi, convert-from, neighbor-stat, train-nvnmd

Sub-commands

config

fast configuration of parameter file for smooth model

dp config [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
          [-o OUTPUT]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-o, --output

the output json file

Default: “input.json”

transfer

pass parameters to another model

dp transfer [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
            [-r RAW_MODEL] [-O OLD_MODEL] [-o OUTPUT]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-r, --raw-model

the model receiving parameters

Default: “raw_frozen_model.pb”

-O, --old-model

the model providing parameters

Default: “old_frozen_model.pb”

-o, --output

the model after passing parameters

Default: “frozen_model.pb”

train

train a model

dp train [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
         [-m {master,collect,workers}]
         [-i INIT_MODEL | -r RESTART | -f INIT_FRZ_MODEL | -t FINETUNE]
         [-o OUTPUT] [--skip-neighbor-stat]
         INPUT

Positional Arguments

INPUT: the input parameter file in json or yaml format

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-m, --mpi-log

Possible choices: master, collect, workers

Set the manner of logging when running with MPI. ‘master’ logs only on main process, ‘collect’ broadcasts logs from workers to master and ‘workers’ means each process will output its own log

Default: “master”

-i, --init-model

Initialize the model by the provided checkpoint.

-r, --restart

Restart the training from the provided checkpoint.

-f, --init-frz-model

Initialize the training from the frozen model.

-t, --finetune

Finetune the frozen pretrained model.

-o, --output

The output file of the parameters used in training.

Default: “out.json”

--skip-neighbor-stat

Skip calculating neighbor statistics. Sel checking, automatic sel, and model compression will be disabled.

Default: False

examples:: dp train input.json dp train input.json –restart model.ckpt dp train input.json –init-model model.ckpt

freeze

freeze the model

dp freeze [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
          [-c CHECKPOINT_FOLDER] [-o OUTPUT] [-n NODE_NAMES] [-w NVNMD_WEIGHT]
          [--united-model]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-c, --checkpoint-folder

path to checkpoint folder

Default: “.”

-o, --output

name of graph, will output to the checkpoint folder

Default: “frozen_model.pb”

-n, --node-names

the frozen nodes, if not set, determined from the model type

-w, --nvnmd-weight

the name of weight file (.npy), if set, save the model’s weight into the file

--united-model

When in multi-task mode, freeze all nodes into one united model

Default: False

examples:: dp freeze dp freeze -o graph.pb

test

test the model

dp test [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH] [-m MODEL]
        [-s SYSTEM | -f DATAFILE] [-S SET_PREFIX] [-n NUMB_TEST]
        [-r RAND_SEED] [--shuffle-test] [-d DETAIL_FILE] [-a]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-m, --model

Frozen model file to import

Default: “frozen_model.pb”

-s, --system

The system dir. Recursively detect systems in this directory

Default: “.”

-f, --datafile

The path to file of test list.

-S, --set-prefix

The set prefix

Default: “set”

-n, --numb-test

The number of data for test

Default: 100

-r, --rand-seed

The random seed

--shuffle-test

Shuffle test data

Default: False

-d, --detail-file

The prefix to files where details of energy, force and virial accuracy/accuracy per atom will be written

-a, --atomic

Test the accuracy of atomic label, i.e. energy / tensor (dipole, polar)

Default: False

examples:: dp test -m graph.pb -s /path/to/system -n 30

compress

compress a model

dp compress [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
            [-m {master,collect,workers}] [-i INPUT] [-o OUTPUT] [-s STEP]
            [-e EXTRAPOLATE] [-f FREQUENCY] [-c CHECKPOINT_FOLDER]
            [-t TRAINING_SCRIPT]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-m, --mpi-log

Possible choices: master, collect, workers

Set the manner of logging when running with MPI. ‘master’ logs only on main process, ‘collect’ broadcasts logs from workers to master and ‘workers’ means each process will output its own log

Default: “master”

-i, --input

The original frozen model, which will be compressed by the code

Default: “frozen_model.pb”

-o, --output

The compressed model

Default: “frozen_model_compressed.pb”

-s, --step

Model compression uses fifth-order polynomials to interpolate the embedding-net. It introduces two tables with different step size to store the parameters of the polynomials. The first table covers the range of the training data, while the second table is an extrapolation of the training data. The domain of each table is uniformly divided by a given step size. And the step(parameter) denotes the step size of the first table and the second table will use 10 * step as it’s step size to save the memory. Usually the value ranges from 0.1 to 0.001. Smaller step means higher accuracy and bigger model size

Default: 0.01

-e, --extrapolate

The domain range of the first table is automatically detected by the code: [d_low, d_up]. While the second table ranges from the first table’s upper boundary(d_up) to the extrapolate(parameter) * d_up: [d_up, extrapolate * d_up]

Default: 5

-f, --frequency

The frequency of tabulation overflow check(Whether the input environment matrix overflow the first or second table range). By default do not check the overflow

Default: -1

-c, --checkpoint-folder

path to checkpoint folder

Default: “model-compression”

-t, --training-script

The training script of the input frozen model

examples:: dp compress dp compress -i graph.pb -o compressed.pb

doc-train-input

print the documentation (in rst format) of input training parameters.

dp doc-train-input [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
                   [--out-type OUT_TYPE]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

--out-type

The output type

Default: “rst”

model-devi

calculate model deviation

dp model-devi [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
              [-m MODELS [MODELS ...]] [-s SYSTEM] [-S SET_PREFIX] [-o OUTPUT]
              [-f FREQUENCY]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-m, --models

Frozen models file to import

Default: [‘graph.000.pb’, ‘graph.001.pb’, ‘graph.002.pb’, ‘graph.003.pb’]

-s, --system

The system directory. Recursively detect systems in this directory.

Default: “.”

-S, --set-prefix

The set prefix

Default: “set”

-o, --output

The output file for results of model deviation

Default: “model_devi.out”

-f, --frequency

The trajectory frequency of the system

Default: 1

examples:: dp model-devi -m graph.000.pb graph.001.pb graph.002.pb graph.003.pb -s ./data -o model_devi.out

convert-from

convert lower model version to supported version

dp convert-from [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
                [-i INPUT_MODEL] [-o OUTPUT_MODEL]
                [{auto,0.12,1.0,1.1,1.2,1.3,2.0,pbtxt}]

Positional Arguments

FROM

Possible choices: auto, 0.12, 1.0, 1.1, 1.2, 1.3, 2.0, pbtxt

The original model compatibility

Default: “auto”

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-i, --input-model

the input model

Default: “frozen_model.pb”

-o, --output-model

the output model

Default: “convert_out.pb”

examples:: dp convert-from -i graph.pb -o graph_new.pb dp convert-from auto -i graph.pb -o graph_new.pb dp convert-from 1.0 -i graph.pb -o graph_new.pb

neighbor-stat

Calculate neighbor statistics

dp neighbor-stat [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
                 [-s SYSTEM] -r RCUT -t TYPE_MAP [TYPE_MAP ...] [--one-type]

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-s, --system

The system dir. Recursively detect systems in this directory

Default: “.”

-r, --rcut

cutoff radius

-t, --type-map

type map

--one-type

treat all types as a single type. Used with se_atten descriptor.

Default: False

examples:: dp neighbor-stat -s data -r 6.0 -t O H

train-nvnmd

train nvnmd model

dp train-nvnmd [-h] [-v {DEBUG,3,INFO,2,WARNING,1,ERROR,0}] [-l LOG_PATH]
               [-r RESTART] [-s {s1,s2}]
               INPUT

Positional Arguments

INPUT: the input parameter file in json format

Named Arguments

-v, --log-level

Possible choices: DEBUG, 3, INFO, 2, WARNING, 1, ERROR, 0

set verbosity level by string or number, 0=ERROR, 1=WARNING, 2=INFO and 3=DEBUG

Default: “INFO”

-l, --log-path

set log file to log messages to disk, if not specified, the logs will only be output to console

-r, --restart

Restart the training from the provided checkpoint.

-s, --step

Possible choices: s1, s2

steps to train model of NVNMD: s1 (train CNN), s2 (train QNN)

Default: “s1”

Integrate with third-party packages

Note that the model for inference is required to be compatible with the DeePMD-kit package. See Model compatibility for details.

Use deep potential with ASE

Deep potential can be set up as a calculator with ASE to obtain potential energies and forces.

from ase import Atoms
from deepmd.calculator import DP

water = Atoms(
    "H2O",
    positions=[(0.7601, 1.9270, 1), (1.9575, 1, 1), (1.0, 1.0, 1.0)],
    cell=[100, 100, 100],
    calculator=DP(model="frozen_model.pb"),
)
print(water.get_potential_energy())
print(water.get_forces())

Optimization is also available:

from ase.optimize import BFGS

dyn = BFGS(water)
dyn.run(fmax=1e-6)
print(water.get_positions())

Run MD with LAMMPS

Running an MD simulation with LAMMPS is simpler. In the LAMMPS input file, one needs to specify the pair style as follows

pair_style     deepmd graph.pb
pair_coeff     * * O H

where graph.pb is the file name of the frozen model. pair_coeff maps atom names (O H) with LAMMPS atom types (integers from 1 to Ntypes, i.e. 1 2).

LAMMPS commands

Enable DeePMD-kit plugin (plugin mode)

If you are using the plugin mode, enable DeePMD-kit package in LAMMPS with plugin command:

plugin load libdeepmd_lmp.so

After LAMMPS version patch_24Mar2022, another way to load plugins is to set the environmental variable LAMMPS_PLUGIN_PATH:

LAMMPS_PLUGIN_PATH=$deepmd_root/lib/deepmd_lmp

where $deepmd_root is the directory to install C++ interface.

The built-in mode doesn’t need this step.

pair_style `deepmd`

The DeePMD-kit package provides the pair_style deepmd

pair_style deepmd models ... keyword value ...

deepmd = style of this pair_style
models = frozen model(s) to compute the interaction. If multiple models are provided, then only the first model serves to provide energy and force prediction for each timestep of molecular dynamics, and the model deviation will be computed among all models every out_freq timesteps.
keyword = out_file or out_freq or fparam or fparam_from_compute or atomic or relative or relative_v or aparam or ttm

    out_file value = filename
        filename = The file name for the model deviation output. Default is model_devi.out
    out_freq value = freq
        freq = Frequency for the model deviation output. Default is 100.
    fparam value = parameters
        parameters = one or more frame parameters required for model evaluation.
    fparam_from_compute value = id
        id = compute id used to update the frame parameter.
    atomic = no value is required.
        If this keyword is set, the model deviation of each atom will be output.
    relative value = level
        level = The level parameter for computing the relative model deviation of the force
    relative_v value = level
        level = The level parameter for computing the relative model deviation of the virial
    aparam value = parameters
        parameters = one or more atomic parameters of each atom required for model evaluation
    ttm value = id
        id = fix ID of fix ttm

Examples

pair_style deepmd graph.pb
pair_style deepmd graph.pb fparam 1.2
pair_style deepmd graph_0.pb graph_1.pb graph_2.pb out_file md.out out_freq 10 atomic relative 1.0
pair_coeff * * O H

pair_style deepmd cp.pb fparam_from_compute TEMP
compute    TEMP all temp

Description

Evaluate the interaction of the system by using Deep Potential or Deep Potential Smooth Edition. It is noticed that deep potential is not a “pairwise” interaction, but a multi-body interaction.

This pair style takes the deep potential defined in a model file that usually has the .pb extension. The model can be trained and frozen by package DeePMD-kit.

The model deviation evalulates the consistency of the force predictions from multiple models. By default, only the maximal, minimal and average model deviations are output. If the key atomic is set, then the model deviation of force prediction of each atom will be output.

By default, the model deviation is output in absolute value. If the keyword relative is set, then the relative model deviation of the force will be output, including values output by the keyword atomic. The relative model deviation of the force on atom $i$ is defined by

\[E_{f_i}=\frac{\left|D_{f_i}\right|}{\left|f_i\right|+l}\]

where $D_{f_i}$ is the absolute model deviation of the force on atom $i$, $f_i$ is the norm of the force and $l$ is provided as the parameter of the keyword relative. If the keyword relative_v is set, then the relative model deviation of the virial will be output instead of the absolute value, with the same definition of that of the force:

\[E_{v_i}=\frac{\left|D_{v_i}\right|}{\left|v_i\right|+l}\]

If the keyword fparam is set, the given frame parameter(s) will be fed to the model. If the keyword fparam_from_compute is set, the global parameter(s) from compute command (e.g., temperature from compute temp command) will be fed to the model as the frame parameter(s). If the keyword aparam is set, the given atomic parameter(s) will be fed to the model, where each atom is assumed to have the same atomic parameter(s). If the keyword ttm is set, electronic temperatures from fix ttm command will be fed to the model as the atomic parameters.

Only a single pair_coeff command is used with the deepmd style which specifies atom names. These are mapped to LAMMPS atom types (integers from 1 to Ntypes) by specifying Ntypes additional arguments after * * in the pair_coeff command. If atom names are not set in the pair_coeff command, the training parameter type_map will be used by default. If the training parameter type_map is not set, atom names in the pair_coeff command cannot be set. In this case, atom type indexes in type.raw (integers from 0 to Ntypes-1) will map to LAMMPS atom types.

Restrictions

The deepmd pair style is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the DeePMD-kit website for more information.

Compute tensorial properties

The DeePMD-kit package provides the compute deeptensor/atom for computing atomic tensorial properties.

compute ID group-ID deeptensor/atom model_file

ID: user-assigned name of the computation
group-ID: ID of the group of atoms to compute
deeptensor/atom: the style of this compute
model_file: the name of the binary model file.

At this time, the training parameter type_map will be mapped to LAMMPS atom types.

Examples

compute         dipole all deeptensor/atom dipole.pb

The result of the compute can be dumped to trajectory file by

dump            1 all custom 100 water.dump id type c_dipole[1] c_dipole[2] c_dipole[3]

Restrictions

The deeptensor/atom compute is provided in the USER-DEEPMD package, which is compiled from the DeePMD-kit, visit the DeePMD-kit website for more information.

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

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.

Use of the centroid/stress/atom to get the full 3x3 “atomic-virial”

The DeePMD-kit allows also the computation of per-atom stress tensor defined as:

\[dvatom=\sum_{m}( \mathbf{r}_n- \mathbf{r}_m) \frac{de_m}{d\mathbf{r}_n}\]

Where $\mathbf{r}_n$ is the atomic position of nth atom, $\mathbf{v}_n$ velocity of the atom and $\frac{de_m}{d\mathbf{r}_n}$ the derivative of the atomic energy.

In LAMMPS one can get the per-atom stress using the command centroid/stress/atom:

compute ID group-ID centroid/stress/atom NULL virial

see LAMMPS doc page for more details on the meaning of the keywords.

Examples

In order of computing the 9-component per-atom stress

compute stress all centroid/stress/atom NULL virial

Thus c_stress is an array with 9 components in the order xx,yy,zz,xy,xz,yz,yx,zx,zy.

If you use this feature please cite D. Tisi, L. Zhang, R. Bertossa, H. Wang, R. Car, S. Baroni - arXiv preprint arXiv:2108.10850, 2021

Computation of heat flux

Using a per-atom stress tensor one can, for example, compute the heat flux defined as:

\[\mathbf J = \sum_n e_n \mathbf v_n + \sum_{n,m} ( \mathbf r_m- \mathbf r_n) \frac{de_m}{d\mathbf r_n} \mathbf v_n\]

to compute the heat flux with LAMMPS:

compute ke_ID all ke/atom
compute pe_ID all pe/atom
compute stress_ID group-ID centroid/stress/atom NULL virial
compute flux_ID all heat/flux ke_ID pe_ID stress_ID

Examples

compute ke all ke/atom
compute pe all pe/atom
compute stress all centroid/stress/atom NULL virial
compute flux all heat/flux ke pe stress

c_flux is a global vector of length 6. The first three components are the $x$, $y$ and $z$ components of the full heat flux vector. The others are the components of the so-called convective portion, see LAMMPS doc page for more detailes.

If you use these features please cite D. Tisi, L. Zhang, R. Bertossa, H. Wang, R. Car, S. Baroni - arXiv preprint arXiv:2108.10850, 2021

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 (or dp_ipi_low for low precision) that computes the interactions (including energy, forces and virials). The server and client communicate via the Unix domain socket or the Internet socket. Installation instructions for i-PI can be found here. The client can be started by

i-pi input.xml &
dp_ipi water.json

It is noted that multiple instances of the client allow for computing, in parallel, the interactions of multiple replicas of the path-integral MD.

water.json is the parameter file for the client dp_ipi, and an example is provided:

{
    "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 port should be the same as that in input.xml:

<port>31415</port>

The option graph_file provides the file name of the frozen model.

The dp_ipi gets the atom names from an XYZ file provided by coord_file (meanwhile ignores all coordinates in it) and translates the names to atom types by rules provided by atom_type.

Running MD with GROMACS

DP/MM Simulation

This part gives a simple tutorial on how to run a DP/MM simulation for methane in water, which means using DP for methane and TIP3P for water. All relevant files can be found in examples/methane.

Topology Preparation

Similar to QM/MM simulation, the internal interactions (including bond, angle, dihedrals, LJ, Columb) of the region described by a neural network potential (NNP) have to be turned off. In GROMACS, bonded interactions can be turned off by modifying [ bonds ], [ angles ], [ dihedrals ] and [ pairs ] sections. And LJ and Columb interactions must be turned off by [ exclusions ] section.

For example, if one wants to simulate ethane in water, using DeepPotential for methane and TIP3P for water, the topology of methane should be like the following (as presented in examples/methane/methane.itp):

[ atomtypes ]
;name btype  mass  charge ptype    sigma  epsilon
  c3    c3   0.0     0.0     A 0.339771 0.451035
  hc    hc   0.0     0.0     A 0.260018 0.087027

[ moleculetype ]
;name            nrexcl
 methane          3

[ atoms ]
; nr type  resnr residue atom  cgnr  charge   mass
  1   c3      1     MOL   C1     1 -0.1068 12.010
  2   hc      1     MOL   H1     2  0.0267  1.008
  3   hc      1     MOL   H2     3  0.0267  1.008
  4   hc      1     MOL   H3     4  0.0267  1.008
  5   hc      1     MOL   H4     5  0.0267  1.008

[ bonds ]
; i  j  func  b0  kb
 1  2     5
 1  3     5
 1  4     5
 1  5     5

[ exclusions ]
; ai  aj1  aj2  aj3  aj4
  1    2    3    4    5
  2    1    3    4    5
  3    1    2    4    5
  4    1    2    3    5
  5    1    2    3    4

For comparison, the original topology file generated by acpype will be:

; methane_GMX.itp created by acpype (v: 2021-02-05T22:15:50CET) on Wed Sep  8 01:21:53 2021

[ atomtypes ]
;name   bond_type     mass     charge   ptype   sigma         epsilon       Amb
 c3       c3          0.00000  0.00000   A     3.39771e-01   4.51035e-01 ; 1.91  0.1078
 hc       hc          0.00000  0.00000   A     2.60018e-01   8.70272e-02 ; 1.46  0.0208

[ moleculetype ]
;name            nrexcl
 methane          3

[ atoms ]
;   nr  type  resi  res  atom  cgnr     charge      mass       ; qtot   bond_type
     1   c3     1   MOL    C1    1    -0.106800     12.01000 ; qtot -0.107
     2   hc     1   MOL    H1    2     0.026700      1.00800 ; qtot -0.080
     3   hc     1   MOL    H2    3     0.026700      1.00800 ; qtot -0.053
     4   hc     1   MOL    H3    4     0.026700      1.00800 ; qtot -0.027
     5   hc     1   MOL    H4    5     0.026700      1.00800 ; qtot 0.000

[ bonds ]
;   ai     aj funct   r             k
     1      2   1    1.0970e-01    3.1455e+05 ;     C1 - H1
     1      3   1    1.0970e-01    3.1455e+05 ;     C1 - H2
     1      4   1    1.0970e-01    3.1455e+05 ;     C1 - H3
     1      5   1    1.0970e-01    3.1455e+05 ;     C1 - H4

[ angles ]
;   ai     aj     ak    funct   theta         cth
     2      1      3      1    1.0758e+02    3.2635e+02 ;     H1 - C1     - H2
     2      1      4      1    1.0758e+02    3.2635e+02 ;     H1 - C1     - H3
     2      1      5      1    1.0758e+02    3.2635e+02 ;     H1 - C1     - H4
     3      1      4      1    1.0758e+02    3.2635e+02 ;     H2 - C1     - H3
     3      1      5      1    1.0758e+02    3.2635e+02 ;     H2 - C1     - H4
     4      1      5      1    1.0758e+02    3.2635e+02 ;     H3 - C1     - H4

DeepMD Settings

Before running simulations, we need to tell GROMACS to use DeepPotential by setting the environment variable GMX_DEEPMD_INPUT_JSON:

export GMX_DEEPMD_INPUT_JSON=input.json

Then, in your working directories, we have to write input.json file:

{
    "graph_file": "/path/to/graph.pb",
    "type_file": "type.raw",
    "index_file": "index.raw",
    "lambda": 1.0,
    "pbc": false
}

Here is an explanation for these settings:

graph_file : The graph file (with suffix .pb) generated by dp freeze command
type_file : File to specify DP atom types (in space-separated format). Here, type.raw looks like

1 0 0 0 0

index_file : File containing indices of DP atoms (in space-separated format), which should be consistent with the indices’ order in .gro file but starting from zero. Here, index.raw looks like

0 1 2 3 4

lambda: Optional, default 1.0. Used in alchemical calculations.
pbc: Optional, default true. If true, the GROMACS periodic condition is passed to DeepMD.

Run Simulation

Finally, you can run GROMACS using gmx mdrun as usual.

All-atom DP Simulation

This part gives an example of how to simulate all atoms described by a DeepPotential with Gromacs, taking water as an example. Instead of using [ exclusions ] to turn off the non-bonded energies, we can simply do this by setting LJ parameters (i.e. epsilon and sigma) and partial charges to 0, as shown in examples/water/gmx/water.top:

[ atomtypes ]
; name      at.num  mass     charge ptype  sigma      epsilon
HW           1       1.008   0.0000  A   0.00000e+00  0.00000e+00
OW           8      16.00    0.0000  A   0.00000e+00  0.00000e+00

As mentioned in the above section, input.json and relevant files (index.raw, type.raw) should also be created. Then, we can start the simulation under the NVT ensemble and plot the radial distribution function (RDF) by gmx rdf command. We can see that the RDF given by Gromacs+DP matches perfectly with Lammps+DP, which further provides an evidence on the validity of our simulation. rdf

However, we still recommend you run an all-atom DP simulation using LAMMPS since it is more stable and efficient.

Interfaces out of DeePMD-kit

The codes of the following interfaces are not a part of the DeePMD-kit package and maintained by other repositories. We list these interfaces here for user convenience.

dpdata

dpdata provides the predict method for System class:

import dpdata
dsys = dpdata.LabeledSystem('OUTCAR')
dp_sys = dsys.predict("frozen_model_compressed.pb")

By inferring with the DP model frozen_model_compressed.pb, dpdata will generate a new labeled system dp_sys with inferred energies, forces, and virials.

OpenMM plugin for DeePMD-kit

An OpenMM plugin is provided from JingHuangLab/openmm_deepmd_plugin, written by the Huang Lab at Westlake University.

AMBER interface to DeePMD-kit

An AMBER interface to DeePMD-kit is written by the [York Lab from Rutgers University. It is open-source at GitLab RutgersLBSR/AmberDPRc. Details can be found in this paper.

DP-GEN

DP-GEN provides a workflow to generate accurate DP models by calling DeePMD-kit’s command line interface (CLI) in the local or remote server. Details can be found in this paper.

MLatom

Mlatom provides an interface to the DeePMD-kit within MLatom’s workflow by calling DeePMD-kit’s CLI. Details can be found in this paper.

ABACUS

ABACUS can run molecular dynamics with a DP model. User is required to build ABACUS with DeePMD-kit.

Use NVNMD

Introduction

NVNMD stands for non-von Neumann molecular dynamics.

This is the training code we used to generate the results in our paper entitled “Accurate and Efficient Molecular Dynamics based on Machine Learning and non von Neumann Architecture”, which has been accepted by npj Computational Materials (DOI: 10.1038/s41524-022-00773-z).

Any user can follow two consecutive steps to run molecular dynamics (MD) on the proposed NVNMD computer, which has been released online: (i) to train a machine learning (ML) model that can decently reproduce the potential energy surface (PES); and (ii) to deploy the trained ML model on the proposed NVNMD computer, then run MD there to obtain the atomistic trajectories.

Training

Our training procedure consists of not only continuous neural network (CNN) training but also quantized neural network (QNN) training which uses the results of CNN as inputs. It is performed on CPU or GPU by using the training codes we open-sourced online.

To train an ML model that can decently reproduce the PES, a training and testing data set should be prepared first. This can be done by using either the state-of-the-art active learning tools or the outdated (i.e., less efficient) brute-force density functional theory (DFT)-based ab-initio molecular dynamics (AIMD) sampling.

If you just want to simply test the training function, you can use the example in the $deepmd_source_dir/examples/nvnmd directory. If you want to fully experience training and running MD functions, you can download the complete example from the website.

Then, copy the data set to the working directory

mkdir -p $workspace
cd $workspace
mkdir -p data
cp -r $dataset data

where $dataset is the path to the data set and $workspace is the path to the working directory.

Input script

Create and go to the training directory.

mkdir train
cd train

Then copy the input script train_cnn.json and train_qnn.json to the directory train

cp -r $deepmd_source_dir/examples/nvnmd/train/train_cnn.json train_cnn.json
cp -r $deepmd_source_dir/examples/nvnmd/train/train_qnn.json train_qnn.json

The structure of the input script is as follows

{
    "nvnmd" : {},
    "learning_rate" : {},
    "loss" : {},
    "training": {}
}

nvnmd

The “nvnmd” section is defined as

{
    "net_size":128,
    "sel":[60, 60],
    "rcut":6.0,
    "rcut_smth":0.5
}

where items are defined as:

Item	Mean	Optional Value
net_size	the size of nueral network	128
sel	the number of neighbors	integer list of lengths 1 to 4 are acceptable
rcut	the cutoff radial	(0, 8.0]
rcut_smth	the smooth cutoff parameter	(0, 8.0]

learning_rate

The “learning_rate” section is defined as

{
    "type":"exp",
    "start_lr": 1e-3,
    "stop_lr": 3e-8,
    "decay_steps": 5000
}

where items are defined as:

Item	Mean	Optional Value
type	learning rate variant type	exp
start_lr	the learning rate at the beginning of the training	a positive real number
stop_lr	the desired learning rate at the end of the training	a positive real number
decay_stops	the learning rate is decaying every {decay_stops} training steps	a positive integer

loss

The “loss” section is defined as

{
    "start_pref_e": 0.02,
    "limit_pref_e": 2,
    "start_pref_f": 1000,
    "limit_pref_f": 1,
    "start_pref_v": 0,
    "limit_pref_v": 0
}

where items are defined as:

Item	Mean	Optional Value
start_pref_e	the loss factor of energy at the beginning of the training	zero or positive real number
limit_pref_e	the loss factor of energy at the end of the training	zero or positive real number
start_pref_f	the loss factor of force at the beginning of the training	zero or positive real number
limit_pref_f	the loss factor of force at the end of the training	zero or positive real number
start_pref_v	the loss factor of virial at the beginning of the training	zero or positive real number
limit_pref_v	the loss factor of virial at the end of the training	zero or positive real number

training

The “training” section is defined as

{
  "seed": 1,
    "stop_batch": 1000000,
    "numb_test": 1,
    "disp_file": "lcurve.out",
    "disp_freq": 1000,
    "save_ckpt": "model.ckpt",
    "save_freq": 10000,
    "training_data":{
      "systems":["system1_path", "system2_path", "..."],
      "set_prefix": "set",
      "batch_size": ["batch_size_of_system1", "batch_size_of_system2", "..."]
    }
}

where items are defined as:

Item	Mean	Optional Value
seed	the randome seed	a integer
stop_batch	the total training steps	a positive integer
numb_test	the accuracy is test by using {numb_test} sample	a positive integer
disp_file	the log file where the training message display	a string
disp_freq	display frequency	a positive integer
save_ckpt	check point file	a string
save_freq	save frequency	a positive integer
systems	a list of data directory which contains the dataset	string list
set_prefix	the prefix of dataset	a string
batch_size	a list of batch size of corresponding dataset	a integer list

Training

Training can be invoked by

# step1: train CNN
dp train-nvnmd train_cnn.json -s s1
# step2: train QNN
dp train-nvnmd train_qnn.json -s s2

After the training process, you will get two folders: nvnmd_cnn and nvnmd_qnn. The nvnmd_cnn contains the model after continuous neural network (CNN) training. The nvnmd_qnn contains the model after quantized neural network (QNN) training. The binary file nvnmd_qnn/model.pb is the model file that is used to perform NVNMD in the server [http://nvnmd.picp.vip].

You can also restart the CNN training from the checkpoint (nvnmd_cnn/model.ckpt) by

dp train-nvnmd train_cnn.json -r nvnmd_cnn/model.ckpt -s s1

Testing

The frozen model can be used in many ways. The most straightforward testing can be invoked by

mkdir test
dp test -m ./nvnmd_qnn/frozen_model.pb -s path/to/system -d ./test/detail -n 99999 -l test/output.log

where the frozen model file to import is given via the -m command line flag, the path to the testing data set is given via the -s command line flag, and the file containing details of energy, forces and virials accuracy is given via the -d command line flag, the amount of data for testing is given via the -n command line flag.

Running MD

After CNN and QNN training, you can upload the ML model to our online NVNMD system and run MD there.

Account application

The server website of NVNMD is available at http://nvnmd.picp.vip. You can visit the URL and enter the login interface (Figure.1).

ALT

Figure.1 The login interface

To obtain an account, please send your application to the email (jie_liu@hnu.edu.cn, liujie@uw.edu). The username and password will be sent to you by email.

Adding task

After successfully obtaining the account, enter the username and password in the login interface, and click “Login” to enter the homepage (Figure.2).

ALT

Figure.2 The homepage

The homepage displays the remaining calculation time and all calculation records not deleted. Click Add a new task to enter the interface for adding a new task (Figure.3).

ALT

Figure.3 The interface for adding a new task

Task name: name of the task
Upload mode: two modes of uploading results to online data storage, including Manual upload and Automatic upload. Results need to be uploaded manually to online data storage with Manual upload mode and will be uploaded automatically with Automatic upload mode.
Input script: input file of the MD simulation.

In the input script, one needs to specify the pair style as follows

pair_style nvnmd model.pb
pair_coeff * *

Model file: the ML model named model.pb obtained by QNN training.
Data files: data files containing the information required for running an MD simulation (e.g., coord.lmp containing initial atom coordinates).

Next, you can click Submit to submit the task and then automatically return to the homepage (Figure.4).

ALT

Figure.4 The homepage with a new record

Then, click Refresh to view the latest status of all calculation tasks.

Cancelling calculation

For the task whose calculation status is Pending and Running, you can click the corresponding Cancel on the homepage to stop the calculation (Figure.5).

ALT

Figure.5 The homepage with a canceled task

Downloading results

For the task whose calculation status is Completed, Failed and Cancelled, you can click the corresponding Package or Separate files in the Download results bar on the homepage to download results.

Click Package to download a zipped package of all files including input files and output results (Figure.6).

ALT

Figure.6 The interface for downloading a zipped package

Click Separate files to download the required separate files (Figure.7).

ALT

Figure.7 The interface for downloading separate files

If Manual upload mode is selected or the file has expired, click Upload on the download interface to upload manually.

Deleting record

For the task no longer needed, you can click the corresponding Delete on the homepage to delete the record.

Records cannot be retrieved after deletion.

Clearing records

Click Clear calculation records on the homepage to clear all records.

Records cannot be retrieved after clearing.

FAQs

As a consequence of differences in computers or systems, problems may occur. Some common circumstances are listed as follows. In addition, some frequently asked questions are listed as follows. If other unexpected problems occur, you’re welcome to contact us for help.

How to tune Fitting/embedding-net size ?

Here are some test forms on fitting-net size tuning or embedding-net size tuning performed on several different systems.

Al2O3

Fitting net size tuning form on Al2O3: (embedding-net size: [25,50,100])

Fitting-net size	Energy L2err(eV)	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[240,240,240]	1.742252e-02	7.259383e-05	4.014115e-02
[80,80,80]	1.799349e-02	7.497287e-05	4.042977e-02
[40,40,40]	1.799036e-02	7.495984e-05	4.068806e-02
[20,20,20]	1.834032e-02	7.641801e-05	4.094784e-02
[10,10,10]	1.913058e-02	7.971073e-05	4.154775e-02
[5,5,5]	1.932914e-02	8.053808e-05	4.188052e-02
[4,4,4]	1.944832e-02	8.103467e-05	4.217826e-02
[3,3,3]	2.068631e-02	8.619296e-05	4.300497e-02
[2,2,2]	2.267962e-02	9.449840e-05	4.413609e-02
[1,1,1]	2.813596e-02	1.172332e-04	4.781115e-02
[]	3.135002e-02	1.306251e-04	5.373120e-02

[] means no hidden layer, but there is still a linear output layer. This situation is equal to the linear regression.

Embedding net size tuning form on Al2O3: (Fitting-net size: [240,240,240])

Embedding-net size	Energy L2err(eV)	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[25,50,100]	1.742252e-02	7.259383e-05	4.014115e-02
[10,20,40]	2.909990e-02	1.212496e-04	4.734667e-02
[5,10,20]	3.357767e-02	1.399070e-04	5.706385e-02
[4,8,16]	6.060367e-02	2.525153e-04	7.333304e-02
[3,6,12]	5.656043e-02	2.356685e-04	7.793539e-02
[2,4,8]	5.277023e-02	2.198759e-04	7.459995e-02
[1,2,4]	1.302282e-01	5.426174e-04	9.672238e-02

Cu

Fitting net size tuning form on Cu: (embedding-net size: [25,50,100])

Fitting-net size	Energy L2err(eV)	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[240,240,240]	4.135548e-02	1.615449e-04	8.940946e-02
[20,20,20]	4.323858e-02	1.689007e-04	8.955762e-02
[10,10,10]	4.399364e-02	1.718502e-04	8.962891e-02
[5,5,5]	4.468404e-02	1.745470e-04	8.970111e-02
[4,4,4]	4.463580e-02	1.743586e-04	8.972011e-02
[3,3,3]	4.493758e-02	1.755374e-04	8.971303e-02
[2,2,2]	4.500736e-02	1.758100e-04	8.973878e-02
[1,1,1]	4.542073e-02	1.774247e-04	8.964761e-02
[]	4.545168e-02	1.775456e-04	8.983201e-02

Embedding net size tuning form on Cu: (Fitting-net size: [240,240,240])

Embedding-net size	Energy L2err(eV)	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[25,50,100]	4.135548e-02	1.615449e-04	8.940946e-02
[20,40,80]	4.203562e-02	1.642016e-04	8.925881e-02
[15,30,60]	4.146672e-02	1.619794e-04	8.936911e-02
[10,20,40]	4.263060e-02	1.665258e-04	8.955818e-02
[5,10,20]	4.994913e-02	1.951138e-04	9.007786e-02
[4,8,16]	1.022157e-01	3.992802e-04	9.532119e-02
[3,6,12]	1.362098e-01	5.320695e-04	1.073860e-01
[2,4,8]	7.061800e-02	2.758515e-04	9.126418e-02
[1,2,4] && seed = 1	9.843161e-02	3.844985e-04	9.348505e-02
[1,2,4] && seed = 2	9.404335e-02	3.673568e-04	9.304089e-02
[1,2,4] && seed = 3	1.508016e-01	5.890688e-04	1.382356e-01
[1,2,4] && seed = 4	9.686949e-02	3.783965e-04	9.294820e-02

Water

Fitting net size tuning form on water: (embedding-net size: [25,50,100])

Fitting-net size	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[240,240,240]	9.1589E-04	5.1540E-02
[200,200,200]	9.3221E-04	5.2366E-02
[160,160,160]	9.4274E-04	5.3403E-02
[120,120,120]	9.5407E-04	5.3093E-02
[80,80,80]	9.4605E-04	5.3402E-02
[40,40,40]	9.8533E-04	5.5790E-02
[20,20,20]	1.0057E-03	5.8232E-02
[10,10,10]	1.0466E-03	6.2279E-02
[5,5,5]	1.1154E-03	6.7994E-02
[4,4,4]	1.1289E-03	6.9613E-02
[3,3,3]	1.2368E-03	7.9786E-02
[2,2,2]	1.3558E-03	9.7042E-02
[1,1,1]	1.4633E-03	1.1265E-01
[]	1.5193E-03	1.2136E-01

Embedding net size tuning form on water: (Fitting-net size: [240,240,240])

Embedding-net size	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[25,50,100]	9.1589E-04	5.1540E-02
[20,40,80]	9.5080E-04	5.3593E-02
[15,30,60]	9.7996E-04	5.6338E-02
[10,20,40]	1.0353E-03	6.2776E-02
[5,10,20]	1.1254E-03	7.3195E-02
[4,8,16]	1.2495E-03	8.0371E-02
[3,6,12]	1.3604E-03	9.9883E-02
[2,4,8]	1.4358E-03	9.7389E-02
[1,2,4]	2.1765E-03	1.7276E-01

Mg-Al

Fitting net size tuning form on Mg-Al: (embedding-net size: [25,50,100])

Fitting-net size	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[240,240,240]	3.9606e-03	1.6289e-02
[200,200,200]	3.9449e-03	1.6471e-02
[160,160,160]	4.0947e-03	1.6413e-02
[120,120,120]	3.9234e-03	1.6283e-02
[80,80,80]	3.9758e-03	1.6506e-02
[40,40,40]	3.9142e-03	1.6348e-02
[20,20,20]	4.1302e-03	1.7006e-02
[10,10,10]	4.3433e-03	1.7524e-02
[5,5,5]	5.3154e-03	1.9716e-02
[4,4,4]	5.4210e-03	1.9710e-02
[2,2,2]	6.2667e-03	2.2568e-02
[1,1,1]	7.3676e-03	2.6375e-02
[]	7.3999e-03	2.6097e-02

Embedding net size tuning form on Mg-Al: (Fitting-net size: [240,240,240])

Embedding-net size	Energy L2err/Natoms(eV)	Force L2err(eV/Angstrom)
[25,50,100]	3.9606e-03	1.6289e-02
[20,40,80]	4.0292e-03	1.6555e-02
[15,30,60]	4.1743e-03	1.7026e-02
[10,20,40]	4.8138e-03	1.8516e-02
[5,10,20]	5.6052e-03	2.0709e-02
[4,8,16]	6.1335e-03	2.1450e-02
[3,6,12]	6.6469e-03	2.3003e-02
[2,4,8]	6.8222e-03	2.6318e-02
[1,2,4]	1.0678e-02	3.9559e-02

How to control the parallelism of a job?

DeePMD-kit has three levels of parallelism. To get the best performance, one should control the number of threads used by DeePMD-kit. One should make sure the product of the parallel numbers is less than or equal to the number of cores available.

MPI (optional)

Parallelism for MPI is optional and used for multiple nodes, multiple GPU cards, or sometimes multiple CPU cores.

To enable MPI support for training, one should install horovod in advance. Note that the parallelism mode is data parallelism, so it is not expected to see the training time per batch decreases.

MPI support for inference is not directly supported by DeePMD-kit, but indirectly supported by the third-party software. For example, LAMMPS enables running simulations in parallel using the MPI parallel communication standard with distributed data. That software has to build against MPI.

Set the number of processes with:

mpirun -np $num_nodes dp

Note that mpirun here should be the same as the MPI used to build software. For example, one can use mpirun -h and lmp -h to see if mpirun and LAMMPS has the same MPI version.

Sometimes, $num_nodes and the nodes information can be directly given by the HPC scheduler system, if the MPI used here is the same as the MPI used to build the scheduler system. Otherwise, one have to manually assign these information.

Parallelism between independent operators

For CPU devices, TensorFlow use multiple streams to run independent operators (OP).

export TF_INTER_OP_PARALLELISM_THREADS=3

However, for GPU devices, TensorFlow uses only one compute stream and multiple copy streams. Note that some of DeePMD-kit OPs do not have GPU support, so it is still encouraged to set environmental variables even if one has a GPU.

Parallelism within an individual operators

For CPU devices, TF_INTRA_OP_PARALLELISM_THREADS controls parallelism within TensorFlow native OPs when TensorFlow is built against Eigen.

export TF_INTRA_OP_PARALLELISM_THREADS=2

OMP_NUM_THREADS is threads for OpenMP parallelism. It controls parallelism within TensorFlow native OPs when TensorFlow is built by Intel OneDNN and DeePMD-kit custom CPU OPs. It may also control parallelsim for NumPy when NumPy is built against OpenMP, so one who uses GPUs for training should also care this environmental variable.

export OMP_NUM_THREADS=2

There are several other environmental variables for OpenMP, such as KMP_BLOCKTIME. See Intel documentation for detailed information.

Tune the performance

There is no one general parallel configuration that works for all situations, so you are encouraged to tune parallel configurations yourself after empirical testing.

Here are some empirical examples. If you wish to use 3 cores of 2 CPUs on one node, you may set the environmental variables and run DeePMD-kit as follows:

export OMP_NUM_THREADS=3
export TF_INTRA_OP_PARALLELISM_THREADS=3
export TF_INTER_OP_PARALLELISM_THREADS=2
dp train input.json

For a node with 128 cores, it is recommended to start with the following variables:

export OMP_NUM_THREADS=16
export TF_INTRA_OP_PARALLELISM_THREADS=16
export TF_INTER_OP_PARALLELISM_THREADS=8

Again, in general, one should make sure the product of the parallel numbers is less than or equal to the number of cores available. In the above case, $16 \times 8 = 128$, so threads will not compete with each other.

Do we need to set rcut < half boxsize?

When seeking the neighbors of atom i under periodic boundary conditions, DeePMD-kit considers all j atoms within cutoff rcut from atom i in all mirror cells.

So, there is no limitation on the setting of rcut.

PS: The reason why some software requires rcut < half box size is that they only consider the nearest mirrors from the center cell. DeePMD-kit is different from them.

How to set sel?

sel is short for “selected number of atoms in rcut”.

sel_a[i] is 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 number of the selected number of type i neighbors within rcut. To ensure that the results are strictly accurate, sel_a[i] should be larger than the largest number of type i neighbors in the rcut.

However, the computation overhead increases with sel_a[i], therefore, sel_a[i] should be as small as possible.

The setting of sel_a[i] should balance the above two considerations.

Installation

Inadequate versions of gcc/g++

Sometimes you may use a gcc/g++ of version < 4.8. In this way, you can still compile all the parts of TensorFlow and most of the parts of DeePMD-kit, but i-Pi and GROMACS plugins will be disabled automatically. Or if you have a gcc/g++ of version > 4.8, say, 7.2.0, you may choose to use it by doing

export CC=/path/to/gcc-7.2.0/bin/gcc
export CXX=/path/to/gcc-7.2.0/bin/g++

Build files left in DeePMD-kit

When you try to build a second time when installing DeePMD-kit, files produced before may contribute to failure. Thus, you may clear them by

cd build
rm -r *

and redo the cmake process.

The temperature undulates violently during the early stages of MD

This is probably because your structure is too far from the equilibrium configuration.

To make sure the potential model is truly accurate, we recommend checking model deviation.

MD: cannot run LAMMPS after installing a new version of DeePMD-kit

This typically happens when you install a new version of DeePMD-kit and copy directly the generated USER-DEEPMD to a LAMMPS source code folder and re-install LAMMPS.

To solve this problem, it suffices to first remove USER-DEEPMD from the LAMMPS source code by

make no-user-deepmd

and then install the new USER-DEEPMD.

If this does not solve your problem, try to decompress the LAMMPS source tarball and install LAMMPS from scratch again, which typically should be very fast.

Model compatibility

When the version of DeePMD-kit used to train the model is different from the that of DeePMD-kit running MDs, one has the problem of model compatibility.

DeePMD-kit guarantees that the codes with the same major and minor revisions are compatible. That is to say, v0.12.5 is compatible with v0.12.0, but is not compatible with v0.11.0 or v1.0.0.

One can execute dp convert-from to convert an old model to a new one.

Model version	v0.12	v1.0	v1.1	v1.2	v1.3	v2.0	v2.1
Compatibility	😊	😊	😊	😊	😊	😄	😄

Legend:

😄: The model is compatible with the DeePMD-kit package.
😊: The model is incompatible with the DeePMD-kit package, but one can execute dp convert-from to convert an old model to v2.1.
😢: The model is incompatible with the DeePMD-kit package, and there is no way to convert models.

Why does a model have low precision?

Many phenomena are caused by model accuracy. For example, during simulations, temperatures explode, structures fall apart, and atoms are lost. One can test the model to confirm whether the model has the enough accuracy.

There are many reasons for a low-quality model. Some common reasons are listed below.

Data

Data units and signs

The unit of training data should follow what is listed in data section. Usually, the package to calculate the training data has different units from those of the DeePMD-kit. It is noted that some software label the energy gradient as forces, instead of the negative energy gradient. It is neccessary to check them carefully to avoid inconsistent data.

SCF coverage and data accuracy

The accuracy of models will not exceed the accuracy of training data, so the training data should reach enough accuracy. Here is a checklist for the accuracy of data:

SCF should converge to a suitable threshold for all points in the training data.
The convergence of the energy, force and virial with respect to the energy cutoff and k-spacing sample is checked.
Sometimes, QM software may generate unstable outliers, which should be removed.
The data should be extracted with enough digits and stored with the proper precision. Large energies may have low precision when they are stored as the single-precision floating-point format (FP32).

Enough data

If the model performs good on the training data, but has bad accuracy on another data, this means some data space is not covered by the training data. It can be validated by evaluting the model deviation with multiple models. If the model deviation of these data is high for some data, try to collect more data using DP-GEN.

Values of data

One should be aware that the errors of some data is also affected by the absolute values of this data. Stable structures tend to be more precise than unstable structures because unstable structures may have larger forces. Also, errors will be introduced in the Projector augmented wave (PAW) DFT calculations when the atoms are very close due to the overlap of pseudo-potentials. It is expected to see that data with large forces has larger errors and it is better to compare different models only with the same data.

Model

Enough `sel`

The sel of the descriptors must be enough for both training and test data. Otherwise, the model will be unreliable and give wrong results.

Cutoff radius

The model cannot fit the long-term interaction out of the cutoff radius. This is a designed approximation for performance, but one has to choose proper cutoff radius for the system.

Neural network size

The size of neural networks will affect the accuracy, but if one follows the parameters in the examples, this effect is insignificant. See FAQ: How to tune Fitting/embedding-net size for details.

Neural network precision

In some cases, one may want to use the FP32 precision to make the model faster. For some applications, FP32 is enough and thus is recommended, but one should still be aware that the precision of FP32 is not as high as that of FP64.

Training

Training steps

Generally speaking, the longer the number of training steps, the better the model. A balance between model accuracy and training time can be achieved. If one finds that model accuracy decreases with training time, there may be a problem with the data. See the data section for details.

Learning rate

Both too large and too small learning rate may affect the training. It is recommended to start with a large learning rate and end with a small learning rate. The learning rate from the examples is a good choice to start.

Find DeePMD-kit C/C++ library from CMake

After DeePMD-kit C/C++ library is installed, one can find DeePMD-kit from CMake:

find_package(DeePMD REQUIRED)

Note that you may need to add ${deepmd_root} to the cached CMake variable CMAKE_PREFIX_PATH.

To link against the C interface library, using

target_link_libraries(some_library PRIVATE DeePMD::deepmd_c)

To link against the C++ interface library, using

target_link_libraries(some_library PRIVATE DeePMD::deepmd_cc)

Coding Conventions

Preface

The aim of these coding standards is to help create a codebase with a defined and consistent coding style that every contributor can get easily familiar with. This will in enhance code readability as there will be no different coding styles from different contributors and everything will be documented. Also, PR diffs will be smaller because of the unified coding style. Finally, static typing will help in hunting down potential bugs before the code is even run.

Contributed code will not be refused merely because it does not strictly adhere to these conditions; as long as it’s internally consistent, clean, and correct, it probably will be accepted. But don’t be surprised if the “offending” code gets fiddled with overtime to conform to these conventions.

There are also GitHub actions CI checks for python code style which will annotate the PR diff for you to see the areas where your code is lacking compared to the set standard.

Rules

The code must be compatible with the oldest supported version of python which is 3.7

The project follows the generic coding conventions as specified in the Style Guide for Python Code, Docstring Conventions and Typing Conventions PEPs, clarified and extended as follows:

Do not use “*” imports such as from module import *. Instead, list imports explicitly.
Use 4 spaces per indentation level. No tabs.
No one-liner compound statements (i.e., no if x: return: use two lines).
Maximum line length is 88 characters as recommended by black which is less strict than Docstring Conventions suggests.
Use “StudlyCaps” for class names.
Use “lowercase” or “lowercase_with_underscores” for function, method, variable names and module names. For short names, joined lowercase may be used (e.g. “tagname”). Choose what is most readable.
No single-character variable names, except indices in loops that encompass a very small number of lines (for i in range(5): ...).
Avoid lambda expressions. Use named functions instead.
Avoid functional constructs (filter, map, etc.). Use list comprehensions instead.
Use "double quotes" for string literals, and """triple double quotes""" for docstring’s. Single quotes are OK for something like
```
f"something {'this' if x else 'that'}"
```
Use f-strings s = f"{x:.2f}" instead of old style formating with "%f" % x. string format method "{x:.2f}".format() may be used sparsely where it is more convenient than f-strings.

Whitespace

Python is not C/C++ so whitespace should be used sparingly to maintain code readability

Read the Whitespace in Expressions and Statements section of PEP8.
Avoid trailing whitespaces.
Do not use excessive whitespace in your expressions and statements.
You should have blank spaces after commas, colons, and semi-colons if it isn’t trailing next to the end of a bracket, brace, or parentheses.
With any operators you should use space on both sides of the operator.
Colons for slicing are considered a binary operator, and should not have any spaces between them.
You should have parentheses with no space, directly next to the function when calling functions function().
When indexing or slicing the brackets should be directly next to the collection with no space collection["index"].
Whitespace used to line up variable values is not recommended.
Make sure you are consistent with the formats you choose when optional choices are available.

General advice

Get rid of as many break and continue statements as possible.
Write short functions. All functions should fit within a standard screen.
Use descriptive variable names.

Writing documentation in the code

Here is an example of how to write good docstrings:

https://github.com/numpy/numpy/blob/master/doc/example.py

The NumPy docstring documentation can be found here

It is a good practice to run pydocstyle check on your code or use a text editor that does it automatically):

$ pydocstyle filename.py

Run pycodestyle on your code

It’s a good idea to run pycodestyle on your code (or use a text editor that does it automatically):

$ pycodestyle filename.py

Run mypy on your code

It’s a good idea to run mypy on your code (or use a text editor that does it automatically):

$ mypy filename.py

Run pydocstyle on your code

It’s a good idea to run pycodestyle on your code (or use a text editor that does it automatically):

$ pycodestyle filename.py --max-line-length=88

Run black on your code

Another method of enforcing PEP8 is using a tool such as black. These tools tend to be very effective at cleaning up code but should be used carefully and code should be retested after cleaning it. Try:

$ black --help

Create a model

If you’d like to create a new model that isn’t covered by the existing DeePMD-kit library, but reuse DeePMD-kit’s other efficient modules such as data processing, trainner, etc, you may want to read this section.

To incorporate your custom model you’ll need to:

Register and implement new components (e.g. descriptor) in a Python file. You may also want to register new TensorFlow OPs if necessary.
Register new arguments for user inputs.
Package new codes into a Python package.
Test new models.

Design a new component

When creating a new component, take descriptor as the example, you should inherit deepmd.descriptor.descriptor.Descriptor class and override several methods. Abstract methods such as deepmd.descriptor.descriptor.Descriptor.build must be implemented and others are not. You should keep arguments of these methods unchanged.

After implementation, you need to register the component with a key:

from deepmd.descriptor import Descriptor

@Descriptor.register("some_descrpt")
class SomeDescript(Descriptor):
    def __init__(self, arg1: bool, arg2: float) -> None:
        pass

Register new arguments

To let someone uses your new component in their input file, you need to create a new method that returns some Argument of your new component, and then register new arguments. For example, the code below

from typing import List

from dargs import Argument
from deepmd.utils.argcheck import descrpt_args_plugin

@descrpt_args_plugin.register("some_descrpt")
def descrpt_some_args() -> List[Argument]:
    return [
        Argument("arg1", bool, optional=False, doc="balabala"),
        Argument("arg2", float, optional=True, default=6.0, doc="haha"),
    ]

allows one to use your new descriptor as below:

"descriptor" :{
    "type": "some_descrpt",
    "arg1": true,
    "arg2": 6.0
}

The arguments here should be consistent with the class arguments of your new component.

Package new codes

You may use setuptools to package new codes into a new Python package. It’s crucial to add your new component to entry_points['deepmd'] in setup.py:

    entry_points={
        'deepmd': [
            'some_descrpt=deepmd_some_descrtpt:SomeDescript',
        ],
    },

where deepmd_some_descrtpt is the module of your codes. It is equivalent to from deepmd_some_descrtpt import SomeDescript.

If you place SomeDescript and descrpt_some_args into different modules, you are also expected to add descrpt_some_args to entry_points.

After you install your new package, you can now use dp train to run your new model.

Atom Type Embedding

Overview

Here is an overview of the DeePMD-kit algorithm. Given a specific centric atom, we can obtain the matrix describing its local environment, named $\mathcal R$. It is consist of the distance between the centric atom and its neighbors, as well as a direction vector. We can embed each distance into a vector of $M_1$ dimension by an embedding net, so the environment matrix $\mathcal R$ can be embedded into matrix $\mathcal G$. We can thus extract a descriptor vector (of $M_1 \times M_2$ dim) of the centric atom from the $\mathcal G$ by some matrix multiplication, and put the descriptor into fitting net to get predicted energy $E$. The vanilla version of DeePMD-kit builds embedding net and fitting net relying on the atom type, resulting in $O(N)$ memory usage. After applying atom type embedding, in DeePMD-kit v2.0, we can share one embedding net and one fitting net in total, which decline training complexity largely.

Preliminary

In the following chart, you can find the meaning of symbols used to clarify the atom-type embedding algorithm.

$i$: Type of centric atom

$j$: Type of neighbor atom

$s_{ij}$: Distance between centric atom and neighbor atom

$\mathcal G_{ij}(\cdot)$: Origin embedding net, take $s_{ij}$ as input and output embedding vector of $M_1$ dim

$\mathcal G(\cdot)$: Shared embedding net

$\text{Multi}(\cdot)$: Matrix multiplication and flattening, output the descriptor vector of $M_1\times M_2$ dim

$F_i(\cdot)$: Origin fitting net, take the descriptor vector as input and output energy

$F(\cdot)$: Shared fitting net

$A(\cdot)$: Atom type embedding net, input is atom type, the output is type embedding vector of dim nchanl

So, we can formulate the training process as follows. Vanilla DeePMD-kit algorithm:

\[E = F_i( \text{Multi}( \mathcal G_{ij}( s_{ij} ) ) )\]

DeePMD-kit applying atom type embedding:

\[E = F( [ \text{Multi}( \mathcal G( [s_{ij}, A(i), A(j)] ) ), A(j)] )\]

or

\[E = F( [ \text{Multi}( \mathcal G( [s_{ij}, A(j)] ) ), A(j)] )\]

The difference between the two variants above is whether using the information of centric atom when generating the descriptor. Users can choose by modifying the type_one_side hyper-parameter in the input JSON file.

How to use

A detailed introduction can be found at se_e2_a_tebd. Looking for a fast start-up, you can simply add a type_embedding section in the input JSON file as displayed in the following, and the algorithm will adopt the atom type embedding algorithm automatically. An example of type_embedding is like

    "type_embedding":{
       "neuron":    [2, 4, 8],
       "resnet_dt": false,
       "seed":      1
    }

Code Modification

Atom-type embedding can be applied to varied embedding net and fitting net, as a result, we build a class TypeEmbedNet to support this free combination. In the following, we will go through the execution process of the code to explain our code modification.

trainer (train/trainer.py)

In trainer.py, it will parse the parameter from the input JSON file. If a type_embedding section is detected, it will build a TypeEmbedNet, which will be later input in the model. model will be built in the function _build_network.

model (model/ener.py)

When building the operation graph of the model in model.build. If a TypeEmbedNet is detected, it will build the operation graph of type embed net, embedding net and fitting net by order. The building process of type embed net can be found in TypeEmbedNet.build, which output the type embedding vector of each atom type (of [$\text{ntypes} \times \text{nchanl}$] dimensions). We then save the type embedding vector into input_dict, so that they can be fetched later in embedding net and fitting net.

embedding net (descriptor/se*.py)

In embedding net, we shall take local environment $\mathcal R$ as input and output matrix $\mathcal G$. Functions called in this process by the order is

build -> _pass_filter -> _filter -> _filter_lower

_pass_filter: It will first detect whether an atom type embedding exists, if so, it will apply atom type embedding algorithm and doesn’t divide the input by type.

_filter: It will call _filter_lower function to obtain the result of matrix multiplication ($\mathcal G^T\cdot \mathcal R$), do further multiplication involved in $\text{Multi}(\cdot)$, and finally output the result of descriptor vector of $M_1 \times M_2$ dim.

_filter_lower: The main function handling input modification. If type embedding exists, it will call _concat_type_embedding function to concat the first column of input $\mathcal R$ (the column of $s_{ij}$) with the atom type embedding information. It will decide whether to use the atom type embedding vector of the centric atom according to the value of type_one_side (if set True, then we only use the vector of the neighbor atom). The modified input will be put into the fitting net to get $\mathcal G$ for further matrix multiplication stage.

fitting net (fit/ener.py)

In fitting net, it takes the descriptor vector as input, whose dimension is [natoms, $M_1\times M_2$]. Because we need to involve information on the centric atom in this step, we need to generate a matrix named atype_embed (of dim [natoms, nchanl]), in which each row is the type embedding vector of the specific centric atom. The input is sorted by type of centric atom, we also know the number of a particular atom type (stored in natoms[2+i]), thus we get the type vector of the centric atom. In the build phase of the fitting net, it will check whether type embedding exists in input_dict and fetch them. After that, call embed_atom_type function to look up the embedding vector for the type vector of the centric atom to obtain atype_embed, and concat input with it ([input, atype_embed]). The modified input goes through fitting net` to get predicted energy.

Note

You can’t apply the compression method while using atom-type embedding.

Python API

deepmd package

Root of the deepmd package, exposes all public classes and submodules.

class deepmd.DeepEval(model_file: Path, load_prefix: str = 'load', default_tf_graph: bool = False, auto_batch_size: bool | int | AutoBatchSize = False)[source]

Bases: object

Common methods for DeepPot, DeepWFC, DeepPolar, …

Parameters:

model_filePath: The name of the frozen model file.
load_prefix: str: The prefix in the load computational graph
default_tf_graphbool: If uses the default tf graph, otherwise build a new tf graph for evaluation
auto_batch_sizebool or int or AutomaticBatchSize, default: False: If True, automatic batch size will be used. If int, it will be used as the initial batch size.

Attributes:

model_type: Get type of model.
model_version: Get version of model.
sess: Get TF session.

Methods

`eval_typeebd`()	Evaluate output of type embedding network by using this model.
`make_natoms_vec`(atom_types[, mixed_type])	Make the natom vector used by deepmd-kit.
`reverse_map`(vec, imap)	Reverse mapping of a vector according to the index map.
`sort_input`(coord, atom_type[, sel_atoms, ...])	Sort atoms in the system according their types.

eval_typeebd() → ndarray[source]

Evaluate output of type embedding network by using this model.

Returns:

np.ndarray: The output of type embedding network. The shape is [ntypes, o_size], where ntypes is the number of types, and o_size is the number of nodes in the output layer.

Raises:

KeyError: If the model does not enable type embedding.

See also

deepmd.utils.type_embed.TypeEmbedNet: The type embedding network.

Examples

Get the output of type embedding network of graph.pb:

>>> from deepmd.infer import DeepPotential
>>> dp = DeepPotential('graph.pb')
>>> dp.eval_typeebd()

load_prefix: str

make_natoms_vec(atom_types: ndarray, mixed_type: bool = False) → ndarray[source]

Make the natom vector used by deepmd-kit.

Parameters:

atom_types: The type of atoms
mixed_type: Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

Returns:

natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

property model_type: str

Get type of model.

:type:str

property model_version: str

Get version of model.

Returns:

str: version of model

static reverse_map(vec: ndarray, imap: List[int]) → ndarray[source]

Reverse mapping of a vector according to the index map.

Parameters:

vec: Input vector. Be of shape [nframes, natoms, -1]
imap: Index map. Be of shape [natoms]

Returns:

vec_out: Reverse mapped vector.

property sess: Session: Get TF session.

static sort_input(coord: ndarray, atom_type: ndarray, sel_atoms: List[int] | None = None, mixed_type: bool = False)[source]

Sort atoms in the system according their types.

Parameters:

coord: The coordinates of atoms. Should be of shape [nframes, natoms, 3]
atom_type: The type of atoms Should be of shape [natoms]
sel_atoms: The selected atoms by type
mixed_type: Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.

Returns:

coord_out: The coordinates after sorting
atom_type_out: The atom types after sorting
idx_map: The index mapping from the input to the output. For example coord_out = coord[:,idx_map,:]
sel_atom_type: Only output if sel_atoms is not None The sorted selected atom types
sel_idx_map: Only output if sel_atoms is not None The index mapping from the selected atoms to sorted selected atoms.

Factory function that will inialize appropriate potential read from model_file.

Parameters:

model_filestr: The name of the frozen model file.
load_prefixstr: The prefix in the load computational graph
default_tf_graphbool: If uses the default tf graph, otherwise build a new tf graph for evaluation

Returns:

Union[DeepDipole, DeepGlobalPolar, DeepPolar, DeepPot, DeepWFC]: one of the available potentials

Raises:

RuntimeError: if model file does not correspond to any implementd potential

class deepmd.DipoleChargeModifier(model_name: str, model_charge_map: List[float], sys_charge_map: List[float], ewald_h: float = 1, ewald_beta: float = 1)[source]

Bases: DeepDipole

Parameters:

model_name: The model file for the DeepDipole model
model_charge_map: Gives the amount of charge for the wfcc
sys_charge_map: Gives the amount of charge for the real atoms
ewald_h: Grid spacing of the reciprocal part of Ewald sum. Unit: A
ewald_beta: Splitting parameter of the Ewald sum. Unit: A^{-1}

Attributes:

model_type: Get type of model.
model_version: Get version of model.
sess: Get TF session.

Methods

`build_fv_graph`()	Build the computational graph for the force and virial inference.
`eval`(coord, box, atype[, eval_fv])	Evaluate the modification.
`eval_full`(coords, cells, atom_types[, ...])	Evaluate the model with interface similar to the energy model.
`eval_typeebd`()	Evaluate output of type embedding network by using this model.
`get_dim_aparam`()	Unsupported in this model.
`get_dim_fparam`()	Unsupported in this model.
`get_ntypes`()	Get the number of atom types of this model.
`get_rcut`()	Get the cut-off radius of this model.
`get_sel_type`()	Get the selected atom types of this model.
`get_type_map`()	Get the type map (element name of the atom types) of this model.
`make_natoms_vec`(atom_types[, mixed_type])	Make the natom vector used by deepmd-kit.
`modify_data`(data)	Modify data.
`reverse_map`(vec, imap)	Reverse mapping of a vector according to the index map.
`sort_input`(coord, atom_type[, sel_atoms, ...])	Sort atoms in the system according their types.

build_fv_graph() → Tensor[source]: Build the computational graph for the force and virial inference.

eval(coord: ndarray, box: ndarray, atype: ndarray, eval_fv: bool = True) → Tuple[ndarray, ndarray, ndarray][source]

Evaluate the modification.

Parameters:

coord: The coordinates of atoms
box: The simulation region. PBC is assumed
atype: The atom types
eval_fv: Evaluate force and virial

Returns:

tot_e: The energy modification
tot_f: The force modification
tot_v: The virial modification

load_prefix: str

modify_data(data: dict) → None[source]

Modify data.

Parameters:

data: Internal data of DeepmdData. Be a dict, has the following keys - coord coordinates - box simulation box - type atom types - find_energy tells if data has energy - find_force tells if data has force - find_virial tells if data has virial - energy energy - force force - virial virial

Subpackages

deepmd.cluster package

Module that reads node resources, auto detects if running local or on SLURM.

deepmd.cluster.get_resource() → Tuple[str, List[str], List[int] | None][source]

Get local or slurm resources: nodename, nodelist, and gpus.

Returns:

Tuple[str, List[str], Optional[List[int]]]: nodename, nodelist, and gpus

Submodules

deepmd.cluster.local module

Get local GPU resources.

deepmd.cluster.local.get_gpus()[source]

Get available IDs of GPU cards at local. These IDs are valid when used as the TensorFlow device ID.

Returns:

Optional[List[int]]: List of available GPU IDs. Otherwise, None.

deepmd.cluster.local.get_resource() → Tuple[str, List[str], List[int] | None][source]

Get local resources: nodename, nodelist, and gpus.

Returns:

Tuple[str, List[str], Optional[List[int]]]: nodename, nodelist, and gpus

deepmd.cluster.slurm module

MOdule to get resources on SLURM cluster.

References

https://github.com/deepsense-ai/tensorflow_on_slurm ####

deepmd.cluster.slurm.get_resource() → Tuple[str, List[str], List[int] | None][source]

Get SLURM resources: nodename, nodelist, and gpus.

Returns:

Tuple[str, List[str], Optional[List[int]]]: nodename, nodelist, and gpus

Raises:

RuntimeError: if number of nodes could not be retrieved
ValueError: list of nodes is not of the same length sa number of nodes
ValueError: if current nodename is not found in node list

deepmd.descriptor package

class deepmd.descriptor.Descriptor(*args, **kwargs)[source]

Bases: PluginVariant

The abstract class for descriptors. All specific descriptors should be based on this class.

The descriptor $\mathcal{D}$ describes the environment of an atom, which should be a function of coordinates and types of its neighbour atoms.

Notes

Only methods and attributes defined in this class are generally public, that can be called by other classes.

Examples

>>> descript = Descriptor(type="se_e2_a", rcut=6., rcut_smth=0.5, sel=[50])
>>> type(descript)
<class 'deepmd.descriptor.se_a.DescrptSeA'>

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_tf.Tensor: The coordinate of atoms
atype_tf.Tensor: The type of atoms
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of frames
meshtf.Tensor: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dictdict[str, Any]: Dictionary for additional inputs
reusebool, optional: The weights in the networks should be reused when get the variable.
suffixstr, optional: Name suffix to identify this descriptor

Returns:

descriptor: tf.Tensor: The output descriptor

Notes

This method must be implemented, as it’s called by other classes.

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) → Tensor[source]

Build the type exclude mask for the descriptor.

Parameters:

exclude_typesList[Tuple[int, int]]: The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
ntypesint: The number of types.
selList[int]: The list of the number of selected neighbors for each type.
ndescrptint: The number of descriptors for each atom.
atypetf.Tensor: The type of atoms, with the size of shape0.
shape0tf.Tensor: The shape of the first dimension of the inputs, which is equal to nsamples * natoms.

Returns:

tf.Tensor: The type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.

Notes

To exclude the interaction between two types, the derivative of energy with respect to distances (or angles) between two atoms should be zero[R08579741114c-1]_, i.e.

\[\forall i \in \text{type 1}, j \in \text{type 2}, \frac{\partial{E}}{\partial{r_{ij}}} = 0\]

When embedding networks between every two types are built, we can just remove that network. But when type_one_side is enabled, a network may be built for multiple pairs of types. In this case, we need to build a mask to exclude the interaction between two types.

The mask assumes the descriptors are sorted by neighbro type with the fixed number of given sel and each neighbor has the same number of descriptors (for example 4).

References

[1]

Jinzhe Zeng, Timothy J. Giese, ̧Sölen Ekesan, Darrin M. York, Development of Range-Corrected Deep Learning Potentials for Fast, Accurate Quantum Mechanical/molecular Mechanical Simulations of Chemical Reactions in Solution, J. Chem. Theory Comput., 2021, 17 (11), 6993-7009.

abstract compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[ndarray]]) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coordlist[np.ndarray]: The coordinates. Can be generated by deepmd.model.model_stat.make_stat_input()
data_boxlist[np.ndarray]: The box. Can be generated by deepmd.model.model_stat.make_stat_input()
data_atypelist[np.ndarray]: The atom types. Can be generated by deepmd.model.model_stat.make_stat_input()
natoms_veclist[np.ndarray]: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.model_stat.make_stat_input()
meshlist[np.ndarray]: The mesh for neighbor searching. Can be generated by deepmd.model.model_stat.make_stat_input()
input_dictdict[str, list[np.ndarray]]: Dictionary for additional input

Notes

This method must be implemented, as it’s called by other classes.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph definition of the model
table_extrapolatefloat, default: 5.: The scale of model extrapolation
table_stride_1float, default: 0.01: The uniform stride of the first table
table_stride_2float, default: 0.1: The uniform stride of the second table
check_frequencyint, default: -1: The overflow check frequency
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported compression.

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]

Reveive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

Notes

This method is called by others when the descriptor supported compression.

abstract get_dim_out() → int[source]

Returns the output dimension of this descriptor.

Returns:

int: the output dimension of this descriptor

Notes

This method must be implemented, as it’s called by other classes.

get_dim_rot_mat_1() → int[source]

Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

Returns:

int: the first dimension of the rotation matrix

get_feed_dict(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor) → Dict[str, Tensor][source]

Generate the feed_dict for current descriptor.

Parameters:

coord_tf.Tensor: The coordinate of atoms
atype_tf.Tensor: The type of atoms
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
boxtf.Tensor: The box. Can be generated by deepmd.model.make_stat_input
meshtf.Tensor: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.

Returns:

feed_dictdict[str, tf.Tensor]: The output feed_dict of current descriptor

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns:

nlisttf.Tensor: Neighbor list
rijtf.Tensor: The relative distance between the neighbor and the center atom.
sel_alist[int]: The number of neighbors with full information
sel_rlist[int]: The number of neighbors with only radial information

abstract get_ntypes() → int[source]

Returns the number of atom types.

Returns:

int: the number of atom types

Notes

This method must be implemented, as it’s called by other classes.

abstract get_rcut() → float[source]

Returns the cut-off radius.

Returns:

float: the cut-off radius

Notes

This method must be implemented, as it’s called by other classes.

get_tensor_names(suffix: str = '') → Tuple[str][source]

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

Tuple[str]: Names of tensors

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported initialization from the given variables.

pass_tensors_from_frz_model(*tensors: Tensor) → None[source]

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor: passed tensors

Notes

The number of parameters in the method must be equal to the numbers of returns in get_tensor_names().

abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_enertf.Tensor: The atomic energy
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

forcetf.Tensor: The force on atoms
virialtf.Tensor: The total virial
atom_virialtf.Tensor: The atomic virial

static register(key: str) → Descriptor[source]

Register a descriptor plugin.

Parameters:

keystr: the key of a descriptor

Returns:

Descriptor: the registered descriptor

Examples

>>> @Descriptor.register("some_descrpt")
    class SomeDescript(Descriptor):
        pass

class deepmd.descriptor.DescrptHybrid(*args, **kwargs)[source]

Bases: Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters:

listlist: Build a descriptor from the concatenation of the list of descriptors.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.
`get_nlist_i`(ii)	Get the neighbor information of the ii-th descriptor.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_def of the model
table_extrapolatefloat, default: 5.: The scale of model extrapolation
table_stride_1float, default: 0.01: The uniform stride of the first table
table_stride_2float, default: 0.1: The uniform stride of the second table
check_frequencyint, default: -1: The overflow check frequency
suffixstr, optional: The suffix of the scope

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]

Reveive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_nlist_i(ii: int) → Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the ii-th descriptor.

Parameters:

iiint: The index of the descriptor

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_tensor_names(suffix: str = '') → Tuple[str][source]

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

Tuple[str]: Names of tensors

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr: The sum of radial statisitcs.
suma: The sum of relative coord statisitcs.
sumn: The sum of neighbor numbers.
sumr2: The sum of square of radial statisitcs.
suma2: The sum of square of relative coord statisitcs.

pass_tensors_from_frz_model(*tensors: Tensor) → None[source]

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor: passed tensors

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

class deepmd.descriptor.DescrptLocFrame(*args, **kwargs)[source]

Bases: Descriptor

Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.

Parameters:

rcut: The cut-off radius
sel_alist[str]: 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.
sel_rlist[str]: 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.
axis_rule: list[int]: 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 second 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]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns:
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_rot_mat() → Tensor[source]: Get rotational matrix.

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

class deepmd.descriptor.DescrptSeA(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor $\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}$ is given by [1]

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where $\mathcal{R}^i \in \mathbb{R}^{N \times 4}$ is the coordinate matrix, and each row of $\mathcal{R}^i$ can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where $\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})$ is the relative coordinate and $r_{ji}=\lVert \mathbf{R}_{ji} \lVert$ is its norm. The switching function $s(r)$ is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix $\mathcal{G}^i \in \mathbb{R}^{N \times M_1}$ consists of outputs of a embedding network $\mathcal{N}$ of $s(r_{ji})$:

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

$\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}$ takes first $M_2$ columns of $\mathcal{G}^i$. The equation of embedding network $\mathcal{N}$ can be found at deepmd.utils.network.embedding_net().

Parameters:

rcut: The cut-off radius $r_c$
rcut_smth: From where the environment matrix should be smoothed $r_s$
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net $\mathcal{N}$
axis_neuron: Number of the axis neuron $M_2$ (number of columns of the sub-matrix of the embedding matrix)
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
precision: The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
multi_task: If the model has multi fitting nets to train.

References

[1] (1,2)

Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018. End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451.

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_dist: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_def of the model
table_extrapolate: The scale of model extrapolation
table_stride_1: The uniform stride of the first table
table_stride_2: The uniform stride of the second table
check_frequency: The overflow check frequency
suffixstr, optional: The suffix of the scope

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]

Reveive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_dim_rot_mat_1() → int[source]: Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_rot_mat() → Tensor[source]: Get rotational matrix.

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr: The sum of radial statisitcs.
suma: The sum of relative coord statisitcs.
sumn: The sum of neighbor numbers.
sumr2: The sum of square of radial statisitcs.
suma2: The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

class deepmd.descriptor.DescrptSeAEbd(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE descriptor with type embedding approach.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
axis_neuron: Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
type_nchanl: Number of channels for type representation
type_nlayer: Number of hidden layers for the type embedding net (skip connected).
numb_aparam: Number of atomic parameters. If >0 it will be embedded with atom types.
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are {0}
precision: The precision of the embedding net parameters. Supported options are {1}
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

class deepmd.descriptor.DescrptSeAEf(*args, **kwargs)[source]

Bases: Descriptor

Smooth edition descriptor with Ef.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
axis_neuron: Number of the axis neuron (number of columns of the sub-matrix of the embedding matrix)
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
precision: The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs. Should have ‘efield’.
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_dim_rot_mat_1() → int[source]: Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_rot_mat() → Tensor[source]: Get rotational matrix.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

class deepmd.descriptor.DescrptSeAEfLower(*args, **kwargs)[source]

Bases: DescrptSeA

Helper class for implementing DescrptSeAEf.

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_, atype_, natoms, box_, mesh, input_dict, suffix='', reuse=None)[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict)[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

class deepmd.descriptor.DescrptSeAMask(*args, **kwargs)[source]

Bases: DescrptSeA

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations. The embedding takes the distance between atoms as input.

The descriptor $\mathcal{D}^i \in \mathcal{R}^{M_1 \times M_2}$ is given by [1]_

\[\mathcal{D}^i = (\mathcal{G}^i)^T \mathcal{R}^i (\mathcal{R}^i)^T \mathcal{G}^i_<\]

where $\mathcal{R}^i \in \mathbb{R}^{N \times 4}$ is the coordinate matrix, and each row of $\mathcal{R}^i$ can be constructed as follows

\[(\mathcal{R}^i)_j = [ \begin{array}{c} s(r_{ji}) & \frac{s(r_{ji})x_{ji}}{r_{ji}} & \frac{s(r_{ji})y_{ji}}{r_{ji}} & \frac{s(r_{ji})z_{ji}}{r_{ji}} \end{array} ]\]

where $\mathbf{R}_{ji}=\mathbf{R}_j-\mathbf{R}_i = (x_{ji}, y_{ji}, z_{ji})$ is the relative coordinate and $r_{ji}=\lVert \mathbf{R}_{ji} \lVert$ is its norm. The switching function $s(r)$ is defined as:

\[\begin{split}s(r)= \begin{cases} \frac{1}{r}, & r<r_s \\ \frac{1}{r} \{ {(\frac{r - r_s}{ r_c - r_s})}^3 (-6 {(\frac{r - r_s}{ r_c - r_s})}^2 +15 \frac{r - r_s}{ r_c - r_s} -10) +1 \}, & r_s \leq r<r_c \\ 0, & r \geq r_c \end{cases}\end{split}\]

Each row of the embedding matrix $\mathcal{G}^i \in \mathbb{R}^{N \times M_1}$ consists of outputs of a embedding network $\mathcal{N}$ of $s(r_{ji})$:

\[(\mathcal{G}^i)_j = \mathcal{N}(s(r_{ji}))\]

$\mathcal{G}^i_< \in \mathbb{R}^{N \times M_2}$ takes first $M_2$ columns of $\mathcal{G}^i$. The equation of embedding network $\mathcal{N}$ can be found at deepmd.utils.network.embedding_net(). Specially for descriptor se_a_mask is a concise implementation of se_a. The difference is that se_a_mask only considered a non-pbc system. And accept a mask matrix to indicate the atom i in frame j is a real atom or not. (1 means real atom, 0 means ghost atom) Thus se_a_mask can accept a variable number of atoms in a frame.

Parameters:

sellist[str]: sel[i] specifies the maxmum number of type i atoms in the neighbor list.
neuronlist[int]: Number of neurons in each hidden layers of the embedding net $\mathcal{N}$
axis_neuron: Number of the axis neuron $M_2$ (number of columns of the sub-matrix of the embedding matrix)
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
activation_function: The activation function in the embedding net. Supported options are {0}
precision: The precision of the embedding net parameters. Supported options are {1}
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
References
———-
.. [1] Linfeng Zhang, Jiequn Han, Han Wang, Wissam A. Saidi, Roberto Car, and E. Weinan. 2018.: End-to-end symmetry preserving inter-atomic potential energy model for finite and extended systems. In Proceedings of the 32nd International Conference on Neural Information Processing Systems (NIPS’18). Curran Associates Inc., Red Hook, NY, USA, 4441–4451.

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cutoff radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

get_rcut() → float[source]: Returns the cutoff radius.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: None for se_a_mask op
atom_virial: None for se_a_mask op

class deepmd.descriptor.DescrptSeAtten(*args, **kwargs)[source]

Bases: DescrptSeA

Smooth version descriptor with attention.

Parameters:

rcut: The cut-off radius $r_c$
rcut_smth: From where the environment matrix should be smoothed $r_s$
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net $\mathcal{N}$
axis_neuron: Number of the axis neuron $M_2$ (number of columns of the sub-matrix of the embedding matrix)
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
precision: The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed
attn: The length of hidden vector during scale-dot attention computation.
attn_layer: The number of layers in attention mechanism.
attn_dotr: Whether to dot the relative coordinates on the attention weights as a gated scheme.
attn_mask: Whether to mask the diagonal in the attention weights.
multi_task: If the model has multi fitting nets to train.

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the attention descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor, nei_type_vec: Tensor) → Tensor[source]

Build the type exclude mask for the attention descriptor.

Parameters:

exclude_typesList[Tuple[int, int]]: The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
ntypesint: The number of types.
selList[int]: The list of the number of selected neighbors for each type.
ndescrptint: The number of descriptors for each atom.
atypetf.Tensor: The type of atoms, with the size of shape0.
shape0tf.Tensor: The shape of the first dimension of the inputs, which is equal to nsamples * natoms.
nei_type_vectf.Tensor: The type of neighbors, with the size of (shape0, nnei).

Returns:

tf.Tensor: The type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.

See also

deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask

Notes

This method has the similiar way to build the type exclude mask as deepmd.descriptor.descriptor.Descriptor.build_type_exclude_mask(). The mathmatical expression has been explained in that method. The difference is that the attention descriptor has provided the type of the neighbors (idx_j) that is not in order, so we use it from an extra input.

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict, mixed_type: bool = False, real_natoms_vec: list | None = None) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. If mixed_type is True, this para is blank. See real_natoms_vec.
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input
mixed_type: Whether to perform the mixed_type mode. If True, the input data has the mixed_type format (see doc/model/train_se_atten.md), in which frames in a system may have different natoms_vec(s), with the same nloc.
real_natoms_vec: If mixed_type is True, it takes in the real natoms_vec for each frame.

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

class deepmd.descriptor.DescrptSeR(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from radial information of atomic configurations.

The embedding takes the distance between atoms as input.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
type_one_side: Try to build N_types embedding nets. Otherwise, building N_types^2 embedding nets
exclude_typesList[List[int]]: The excluded pairs of types which have no interaction with each other. For example, [[0, 1]] means no interaction between type 0 and type 1.
activation_function: The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
precision: The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord, data_box, data_atype, natoms_vec, mesh, input_dict)[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_dist: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_def of the model
table_extrapolate: The scale of model extrapolation
table_stride_1: The uniform stride of the first table
table_stride_2: The uniform stride of the second table
check_frequency: The overflow check frequency
suffixstr, optional: The suffix of the scope

get_dim_out()[source]: Returns the output dimension of this descriptor.

get_nlist()[source]

Returns neighbor information.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes()[source]: Returns the number of atom types.

get_rcut()[source]: Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr: The sum of radial statisitcs.
sumn: The sum of neighbor numbers.
sumr2: The sum of square of radial statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

class deepmd.descriptor.DescrptSeT(*args, **kwargs)[source]

Bases: DescrptSe

DeepPot-SE constructed from all information (both angular and radial) of atomic configurations.

The embedding takes angles between two neighboring atoms as input.

Parameters:

rcut: The cut-off radius
rcut_smth: From where the environment matrix should be smoothed
sellist[str]: sel[i] specifies the maxmum number of type i atoms in the cut-off radius
neuronlist[int]: Number of neurons in each hidden layers of the embedding net
resnet_dt: Time-step dt in the resnet construction: y = x + dt * phi (Wx + b)
trainable: If the weights of embedding net are trainable.
seed: Random seed for initializing the network parameters.
set_davg_zero: Set the shift of embedding net input to zero.
activation_function: The activation function in the embedding net. Supported options are “relu”, “relu6”, “softplus”, “sigmoid”, “tanh”, “gelu”, “gelu_tf”, “None”, “none”.
precision: The precision of the embedding net parameters. Supported options are “default”, “float16”, “float32”, “float64”, “bfloat16”.
uniform_seed: Only for the purpose of backward compatibility, retrieves the old behavior of using the random seed

Attributes:

precision: Precision of filter network.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_dist: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_def of the model
table_extrapolate: The scale of model extrapolation
table_stride_1: The uniform stride of the first table
table_stride_2: The uniform stride of the second table
check_frequency: The overflow check frequency
suffixstr, optional: The suffix of the scope

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr: The sum of radial statisitcs.
suma: The sum of relative coord statisitcs.
sumn: The sum of neighbor numbers.
sumr2: The sum of square of radial statisitcs.
suma2: The sum of square of relative coord statisitcs.

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

Submodules

deepmd.descriptor.descriptor module

class deepmd.descriptor.descriptor.Descriptor(*args, **kwargs)[source]

Bases: PluginVariant

The abstract class for descriptors. All specific descriptors should be based on this class.

The descriptor $\mathcal{D}$ describes the environment of an atom, which should be a function of coordinates and types of its neighbour atoms.

Notes

Only methods and attributes defined in this class are generally public, that can be called by other classes.

Examples

>>> descript = Descriptor(type="se_e2_a", rcut=6., rcut_smth=0.5, sel=[50])
>>> type(descript)
<class 'deepmd.descriptor.se_a.DescrptSeA'>

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

abstract build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: Dict[str, Any], reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_tf.Tensor: The coordinate of atoms
atype_tf.Tensor: The type of atoms
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of frames
meshtf.Tensor: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dictdict[str, Any]: Dictionary for additional inputs
reusebool, optional: The weights in the networks should be reused when get the variable.
suffixstr, optional: Name suffix to identify this descriptor

Returns:

descriptor: tf.Tensor: The output descriptor

Notes

This method must be implemented, as it’s called by other classes.

build_type_exclude_mask(exclude_types: List[Tuple[int, int]], ntypes: int, sel: List[int], ndescrpt: int, atype: Tensor, shape0: Tensor) → Tensor[source]

Build the type exclude mask for the descriptor.

Parameters:

exclude_typesList[Tuple[int, int]]: The list of excluded types, e.g. [(0, 1), (1, 0)] means the interaction between type 0 and type 1 is excluded.
ntypesint: The number of types.
selList[int]: The list of the number of selected neighbors for each type.
ndescrptint: The number of descriptors for each atom.
atypetf.Tensor: The type of atoms, with the size of shape0.
shape0tf.Tensor: The shape of the first dimension of the inputs, which is equal to nsamples * natoms.

Returns:

tf.Tensor: The type exclude mask, with the shape of (shape0, ndescrpt), and the precision of GLOBAL_TF_FLOAT_PRECISION. The mask has the value of 1 if the interaction between two types is not excluded, and 0 otherwise.

Notes

To exclude the interaction between two types, the derivative of energy with respect to distances (or angles) between two atoms should be zero[Rafc1ae60e195-1]_, i.e.

\[\forall i \in \text{type 1}, j \in \text{type 2}, \frac{\partial{E}}{\partial{r_{ij}}} = 0\]

When embedding networks between every two types are built, we can just remove that network. But when type_one_side is enabled, a network may be built for multiple pairs of types. In this case, we need to build a mask to exclude the interaction between two types.

The mask assumes the descriptors are sorted by neighbro type with the fixed number of given sel and each neighbor has the same number of descriptors (for example 4).

References

[1]

Jinzhe Zeng, Timothy J. Giese, ̧Sölen Ekesan, Darrin M. York, Development of Range-Corrected Deep Learning Potentials for Fast, Accurate Quantum Mechanical/molecular Mechanical Simulations of Chemical Reactions in Solution, J. Chem. Theory Comput., 2021, 17 (11), 6993-7009.

abstract compute_input_stats(data_coord: List[ndarray], data_box: List[ndarray], data_atype: List[ndarray], natoms_vec: List[ndarray], mesh: List[ndarray], input_dict: Dict[str, List[ndarray]]) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coordlist[np.ndarray]: The coordinates. Can be generated by deepmd.model.model_stat.make_stat_input()
data_boxlist[np.ndarray]: The box. Can be generated by deepmd.model.model_stat.make_stat_input()
data_atypelist[np.ndarray]: The atom types. Can be generated by deepmd.model.model_stat.make_stat_input()
natoms_veclist[np.ndarray]: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.model_stat.make_stat_input()
meshlist[np.ndarray]: The mesh for neighbor searching. Can be generated by deepmd.model.model_stat.make_stat_input()
input_dictdict[str, list[np.ndarray]]: Dictionary for additional input

Notes

This method must be implemented, as it’s called by other classes.

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph definition of the model
table_extrapolatefloat, default: 5.: The scale of model extrapolation
table_stride_1float, default: 0.01: The uniform stride of the first table
table_stride_2float, default: 0.1: The uniform stride of the second table
check_frequencyint, default: -1: The overflow check frequency
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported compression.

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]

Reveive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

Notes

This method is called by others when the descriptor supported compression.

abstract get_dim_out() → int[source]

Returns the output dimension of this descriptor.

Returns:

int: the output dimension of this descriptor

Notes

This method must be implemented, as it’s called by other classes.

get_dim_rot_mat_1() → int[source]

Returns the first dimension of the rotation matrix. The rotation is of shape dim_1 x 3.

Returns:

int: the first dimension of the rotation matrix

get_feed_dict(coord_: Tensor, atype_: Tensor, natoms: Tensor, box: Tensor, mesh: Tensor) → Dict[str, Tensor][source]

Generate the feed_dict for current descriptor.

Parameters:

coord_tf.Tensor: The coordinate of atoms
atype_tf.Tensor: The type of atoms
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
boxtf.Tensor: The box. Can be generated by deepmd.model.make_stat_input
meshtf.Tensor: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.

Returns:

feed_dictdict[str, tf.Tensor]: The output feed_dict of current descriptor

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns neighbor information.

Returns:

nlisttf.Tensor: Neighbor list
rijtf.Tensor: The relative distance between the neighbor and the center atom.
sel_alist[int]: The number of neighbors with full information
sel_rlist[int]: The number of neighbors with only radial information

abstract get_ntypes() → int[source]

Returns the number of atom types.

Returns:

int: the number of atom types

Notes

This method must be implemented, as it’s called by other classes.

abstract get_rcut() → float[source]

Returns the cut-off radius.

Returns:

float: the cut-off radius

Notes

This method must be implemented, as it’s called by other classes.

get_tensor_names(suffix: str = '') → Tuple[str][source]

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

Tuple[str]: Names of tensors

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

Notes

This method is called by others when the descriptor supported initialization from the given variables.

pass_tensors_from_frz_model(*tensors: Tensor) → None[source]

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor: passed tensors

Notes

The number of parameters in the method must be equal to the numbers of returns in get_tensor_names().

abstract prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_enertf.Tensor: The atomic energy
natomstf.Tensor: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

forcetf.Tensor: The force on atoms
virialtf.Tensor: The total virial
atom_virialtf.Tensor: The atomic virial

static register(key: str) → Descriptor[source]

Register a descriptor plugin.

Parameters:

keystr: the key of a descriptor

Returns:

Descriptor: the registered descriptor

Examples

>>> @Descriptor.register("some_descrpt")
    class SomeDescript(Descriptor):
        pass

deepmd.descriptor.hybrid module

class deepmd.descriptor.hybrid.DescrptHybrid(*args, **kwargs)[source]

Bases: Descriptor

Concate a list of descriptors to form a new descriptor.

Parameters:

listlist: Build a descriptor from the concatenation of the list of descriptors.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.
`get_nlist_i`(ii)	Get the neighbor information of the ii-th descriptor.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`merge_input_stats`(stat_dict)	Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

enable_compression(min_nbor_dist: float, graph: Graph, graph_def: GraphDef, table_extrapolate: float = 5.0, table_stride_1: float = 0.01, table_stride_2: float = 0.1, check_frequency: int = -1, suffix: str = '') → None[source]

Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.

Parameters:

min_nbor_distfloat: The nearest distance between atoms
graphtf.Graph: The graph of the model
graph_deftf.GraphDef: The graph_def of the model
table_extrapolatefloat, default: 5.: The scale of model extrapolation
table_stride_1float, default: 0.01: The uniform stride of the first table
table_stride_2float, default: 0.1: The uniform stride of the second table
check_frequencyint, default: -1: The overflow check frequency
suffixstr, optional: The suffix of the scope

enable_mixed_precision(mixed_prec: dict | None = None) → None[source]

Reveive the mixed precision setting.

Parameters:

mixed_prec: The mixed precision setting used in the embedding net

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the descriptor, returns the nlist of the descriptor with the largest cut-off radius.

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_nlist_i(ii: int) → Tuple[Tensor, Tensor, List[int], List[int]][source]

Get the neighbor information of the ii-th descriptor.

Parameters:

iiint: The index of the descriptor

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_tensor_names(suffix: str = '') → Tuple[str][source]

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

Tuple[str]: Names of tensors

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

merge_input_stats(stat_dict)[source]

Merge the statisitcs computed from compute_input_stats to obtain the self.davg and self.dstd.

Parameters:

stat_dict

The dict of statisitcs computed from compute_input_stats, including:

sumr: The sum of radial statisitcs.
suma: The sum of relative coord statisitcs.
sumn: The sum of neighbor numbers.
sumr2: The sum of square of radial statisitcs.
suma2: The sum of square of relative coord statisitcs.

pass_tensors_from_frz_model(*tensors: Tensor) → None[source]

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

*tensorstf.Tensor: passed tensors

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

deepmd.descriptor.loc_frame module

class deepmd.descriptor.loc_frame.DescrptLocFrame(*args, **kwargs)[source]

Bases: Descriptor

Defines a local frame at each atom, and the compute the descriptor as local coordinates under this frame.

Parameters:

rcut: The cut-off radius
sel_alist[str]: 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.
sel_rlist[str]: 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.
axis_rule: list[int]: 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 second 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]: index of the axis atom defining the second axis. Note that the neighbors with the same class and type are sorted according to their relative distance.

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns:
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_rot_mat`()	Get rotational matrix.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(*tensors)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

build(coord_: Tensor, atype_: Tensor, natoms: Tensor, box_: Tensor, mesh: Tensor, input_dict: dict, reuse: bool | None = None, suffix: str = '') → Tensor[source]

Build the computational graph for the descriptor.

Parameters:

coord_: The coordinate of atoms
atype_: The type of atoms
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms
box_tf.Tensor: The box of the system
mesh: For historical reasons, only the length of the Tensor matters. if size of mesh == 6, pbc is assumed. if size of mesh == 0, no-pbc is assumed.
input_dict: Dictionary for additional inputs
reuse: The weights in the networks should be reused when get the variable.
suffix: Name suffix to identify this descriptor

Returns:

descriptor: The output descriptor

compute_input_stats(data_coord: list, data_box: list, data_atype: list, natoms_vec: list, mesh: list, input_dict: dict) → None[source]

Compute the statisitcs (avg and std) of the training data. The input will be normalized by the statistics.

Parameters:

data_coord: The coordinates. Can be generated by deepmd.model.make_stat_input
data_box: The box. Can be generated by deepmd.model.make_stat_input
data_atype: The atom types. Can be generated by deepmd.model.make_stat_input
natoms_vec: The vector for the number of atoms of the system and different types of atoms. Can be generated by deepmd.model.make_stat_input
mesh: The mesh for neighbor searching. Can be generated by deepmd.model.make_stat_input
input_dict: Dictionary for additional input

get_dim_out() → int[source]: Returns the output dimension of this descriptor.

get_nlist() → Tuple[Tensor, Tensor, List[int], List[int]][source]

Returns:

nlist: Neighbor list
rij: The relative distance between the neighbor and the center atom.
sel_a: The number of neighbors with full information
sel_r: The number of neighbors with only radial information

get_ntypes() → int[source]: Returns the number of atom types.

get_rcut() → float[source]: Returns the cut-off radius.

get_rot_mat() → Tensor[source]: Get rotational matrix.

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

prod_force_virial(atom_ener: Tensor, natoms: Tensor) → Tuple[Tensor, Tensor, Tensor][source]

Compute force and virial.

Parameters:

atom_ener: The atomic energy
natoms: The number of atoms. This tensor has the length of Ntypes + 2 natoms[0]: number of local atoms natoms[1]: total number of atoms held by this processor natoms[i]: 2 <= i < Ntypes+2, number of type i atoms

Returns:

force: The force on atoms
virial: The total virial
atom_virial: The atomic virial

deepmd.descriptor.se module

class deepmd.descriptor.se.DescrptSe(*args, **kwargs)[source]

Bases: Descriptor

A base class for smooth version of descriptors.

Notes

All of these descriptors have an environmental matrix and an embedding network (deepmd.utils.network.embedding_net()), so they can share some similiar methods without defining them twice.

Attributes:

embedding_net_variablesdict: initial embedding network variables
descrpt_reshapetf.Tensor: the reshaped descriptor
descrpt_derivtf.Tensor: the descriptor derivative
rijtf.Tensor: distances between two atoms
nlisttf.Tensor: the neighbor list

Methods

`build`(coord_, atype_, natoms, box_, mesh, ...)	Build the computational graph for the descriptor.
`build_type_exclude_mask`(exclude_types, ...)	Build the type exclude mask for the descriptor.
`compute_input_stats`(data_coord, data_box, ...)	Compute the statisitcs (avg and std) of the training data.
`enable_compression`(min_nbor_dist, graph, ...)	Reveive the statisitcs (distance, max_nbor_size and env_mat_range) of the training data.
`enable_mixed_precision`([mixed_prec])	Reveive the mixed precision setting.
`get_dim_out`()	Returns the output dimension of this descriptor.
`get_dim_rot_mat_1`()	Returns the first dimension of the rotation matrix.
`get_feed_dict`(coord_, atype_, natoms, box, mesh)	Generate the feed_dict for current descriptor.
`get_nlist`()	Returns neighbor information.
`get_ntypes`()	Returns the number of atom types.
`get_rcut`()	Returns the cut-off radius.
`get_tensor_names`([suffix])	Get names of tensors.
`init_variables`(graph, graph_def[, suffix])	Init the embedding net variables with the given dict.
`pass_tensors_from_frz_model`(descrpt_reshape, ...)	Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.
`prod_force_virial`(atom_ener, natoms)	Compute force and virial.
`register`(key)	Register a descriptor plugin.

get_tensor_names(suffix: str = '') → Tuple[str][source]

Get names of tensors.

Parameters:

suffixstr: The suffix of the scope

Returns:

Tuple[str]: Names of tensors

init_variables(graph: Graph, graph_def: GraphDef, suffix: str = '') → None[source]

Init the embedding net variables with the given dict.

Parameters:

graphtf.Graph: The input frozen model graph
graph_deftf.GraphDef: The input frozen model graph_def
suffixstr, optional: The suffix of the scope

pass_tensors_from_frz_model(descrpt_reshape: Tensor, descrpt_deriv: Tensor, rij: Tensor, nlist: Tensor)[source]

Pass the descrpt_reshape tensor as well as descrpt_deriv tensor from the frz graph_def.

Parameters:

descrpt_reshape: The passed descrpt_reshape tensor
descrpt_deriv: The passed descrpt_deriv tensor
rij: The passed rij tensor
nlist: The passed nlist tensor

property precision: DType: Precision of filter network.

deepmd.descriptor.se_a module

class deepmd.descriptor.se_a.DescrptSeA(*args, **kwargs)[source]