deepmd-kit Documentation

1. Project Overview & Quickstart (deepmodeling/deepmd-kit) Copy Chapter

File: doc/data/data-conv.md Copy Chapter

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 Copy Chapter

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 Copy Chapter

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 Copy Chapter

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.

File: doc/data/dpdata.md Copy Chapter

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.

File: doc/data/index.rst Copy Chapter

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 :code:systems. Each :code:system is composed of a number of :code:frames. One may roughly view a :code:frame as a snapshot of an MD trajectory, but it does not necessarily come from an MD simulation. A :code: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 :code:frames in one :code:system share the same number of atoms with the same type.

.. toctree::
:maxdepth: 1

system
data-conv
dpdata

File: doc/data/system.md Copy Chapter

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.

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

The input frame properties contain the following property, the first axis of which is the number of 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:

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

Mixed type Copy Chapter

[!NOTE]
Only the DPA-1 and DPA-2 descriptors support this format.

In the standard data 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:

With these edited files, one can put together frames with the same Natoms, instead of the same formula (like H2O).

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.

:::{note}
Per-atom loss terms (force, atomic energy, dos, tensor) sum over a batch's real labels and divide by their total count, so a frame counts in proportion to its real atom number. This differs from averaging each frame's own per-label mean only when the frames of one batch pad a different number of virtual atoms; where the padding is uniform — including the case of no padding at all — the two are identical. Writing real_atom_types.npy with a per-frame varying number of -1 is therefore the one configuration whose loss values, and hence training trajectory, differ from a per-frame average. This applies to the PyTorch backend; the TensorFlow backend normalizes by the padded width regardless.
:::

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

File: doc/development/cicd.md Copy Chapter

CI/CD

CI Copy Chapter

Test CUDA Copy Section

Test CUDA action runs tests on a self-hosted runner with the NVIDIA card. It is not triggered by every PR. The developer who has the permission to manage the label can apply the label Test CUDA to a PR to trigger this action.

CD Copy Chapter

GitHub Actions is used to build pre-compiled packages for each commit. See the Easy install the latest development version section to learn how to install the latest development version.

File: doc/development/cmake.md Copy Chapter

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)

File: doc/development/coding-conventions.rst Copy Chapter

.. _coding conventions:

================== 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 pre-commit CI checks for python code style which will automatically fix the
PR.

Python

Rules Copy Chapter

The code must be compatible with the oldest supported version of python
which is 3.10.

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 <https://github.com/psf/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

  .. code-block:: python

  f"something {'this' if x else 'that'}"

• Use f-strings s = f"{x:.2f}" instead of old style formatting with "%f" % x.
  string format method "{x:.2f}".format() may be used sparsely where it is more
  convenient than f-strings.

Whitespace Copy Chapter

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.

.. _Style Guide for Python Code:
.. _PEP8: https://www.python.org/dev/peps/pep-0008/
.. _Docstring Conventions: https://www.python.org/dev/peps/pep-0257/
.. _Typing Conventions: https://www.python.org/dev/peps/pep-0484/
.. _Docutils project: http://docutils.sourceforge.net/docs/dev/policies.html
#python-coding-conventions
.. _trailing whitespaces: http://www.gnu.org/software/emacs/manual/html_node/
emacs/Useless-Whitespace.html

General advice Copy Chapter

• 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 Copy Chapter

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 <https://numpydoc.readthedocs.io/en/latest/format.html>_

C++

The customized Clang Format style is used for C++ code formatting. The style is defined in
.clang-format file in the root of the repository. The style is based on the Google C++
style with some modifications.

Run scripts to check the code

It's a good idea to install prek <https://github.com/j178/prek>_ on your repository:

.. code-block:: bash

$ uv tool install prek
$ prek install

The scripts will be run automatically before each commit and will fix the code style
issues automatically.

File: doc/development/create-a-model-pt.md Copy Chapter

Create a model in other backends {{ pytorch_icon }} {{ dpmodel_icon }}

[!NOTE]
Supported backends: PyTorch-TorchScript {{ pytorch_icon }}, DP {{ dpmodel_icon }}

In the following context, we use the PyTorch-TorchScript backend as the example, while it also applies to other backends listed above.

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, trainer, etc, you may want to read this section.

To incorporate your custom model you'll need to:

1. Register and implement new components (e.g. descriptor) in a Python file.
2. Register new arguments for user inputs.
3. Package new codes into a Python package.
4. Test new models.

Design a new component Copy Chapter

With DeePMD-kit v3, we have expanded support to include two additional backends alongside TensorFlow: the PyTorch-TorchScript backend and the framework-independent backend (dpmodel). The PyTorch-TorchScript backend adopts a highly modularized design to provide flexibility and extensibility. It ensures a consistent experience for both training and inference, aligning with the TensorFlow backend.

The framework-independent backend is implemented in pure NumPy, serving as a reference backend to ensure consistency in tests. Its design pattern closely parallels that of the PyTorch-TorchScript backend.

New descriptors Copy Section

When creating a new descriptor, it is essential to inherit from both the {py:class}deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor class and the {py:class}torch.nn.Module class. Abstract methods, including {py:class}deepmd.pt.model.descriptor.base_descriptor.BaseDescriptor.forward, must be implemented, while others remain optional. It is crucial to adhere to the original method arguments without any modifications. Once the implementation is complete, the next step involves registering the component with a designated key:

from deepmd.utils.data_system import DeepmdDataSystem
from deepmd.pt.model.descriptor.base_descriptor import (
    BaseDescriptor,
)


@BaseDescriptor.register("some_descrpt")
class SomeDescript(BaseDescriptor, torch.nn.Module):
    def __init__(self, arg1: bool, arg2: float) -> None:
        pass

    def get_rcut(self) -> float:
        pass

    def get_nnei(self) -> int:
        pass

    def get_ntypes(self) -> int:
        pass

    def get_dim_out(self) -> int:
        pass

    def get_dim_emb(self) -> int:
        pass

    def mixed_types(self) -> bool:
        pass

    def forward(
        self,
        coord_ext: torch.Tensor,
        atype_ext: torch.Tensor,
        nlist: torch.Tensor,
        mapping: Optional[torch.Tensor] = None,
    ):
        pass

    def serialize(self) -> dict:
        pass

    def deserialize(cls, data: dict) -> "SomeDescript":
        pass

    def update_sel(
        cls,
        train_data: DeepmdDataSystem,
        type_map: Optional[list[str]],
        local_jdata: dict,
    ):
        pass

The serialize and deserialize methods are important for cross-backend model conversion.

New fitting nets Copy Section

In many instances, there is no requirement to create a new fitting net. For fitting user-defined scalar properties, the {py:class}deepmd.pt.model.task.ener.InvarFitting class can be utilized. However, if there is a need for a new fitting net, one should inherit from both the {py:class}deepmd.pt.model.task.base_fitting.BaseFitting class and the {py:class}torch.nn.Module class. Alternatively, for a more straightforward approach, inheritance from the {py:class}deepmd.pt.model.task.fitting.GeneralFitting class is also an option.

from deepmd.pt.model.task.fitting import (
    GeneralFitting,
)
from deepmd.dpmodel import (
    FittingOutputDef,
    fitting_check_output,
)


@GeneralFitting.register("some_fitting")
@fitting_check_output
class SomeFittingNet(GeneralFitting):
    def __init__(self, arg1: bool, arg2: float) -> None:
        pass

    def forward(
        self,
        descriptor: torch.Tensor,
        atype: torch.Tensor,
        gr: Optional[torch.Tensor] = None,
        g2: Optional[torch.Tensor] = None,
        h2: Optional[torch.Tensor] = None,
        fparam: Optional[torch.Tensor] = None,
        aparam: Optional[torch.Tensor] = None,
    ):
        pass

    def output_def(self) -> FittingOutputDef:
        pass

New models Copy Section

The PyTorch-TorchScript backend's model architecture is meticulously structured with multiple layers of abstraction, ensuring a high degree of flexibility. Typically, the process commences with an atomic model responsible for atom-wise property calculations. This atomic model inherits from both the {py:class}deepmd.pt.model.atomic_model.base_atomic_model.BaseAtomicModel class and the {py:class}torch.nn.Module class.

Subsequently, the AtomicModel is encapsulated using the make_model(AtomicModel) function, which leverages the deepmd.pt.model.model.make_model.make_model function. The purpose of the make_model wrapper is to facilitate the translation between atomic property predictions and the extended property predictions and differentiation , e.g. the reduction of atomic energy contribution and the autodiff for calculating the forces and virial. The developers usually need to implement an AtomicModel not a Model.

from deepmd.pt.model.atomic_model.base_atomic_model import (
    BaseAtomicModel,
)


class SomeAtomicModel(BaseAtomicModel, torch.nn.Module):
    def __init__(self, arg1: bool, arg2: float) -> None:
        pass

    def forward_atomic(self):
        pass

Floating-point precision Copy Section

When creating a new component, the floating-point precision should obey the Floating-point precision of the model section.
In implementation, the component should

• store parameters in the component precision, except those for output normalization;
• store output normalization parameters in {py:data}deepmd.pt.utils.env.GLOBAL_PT_FLOAT_PRECISION;
• before input normalization, cast the input tensor to the component precision; before output normalization, cast the output tensor to the {py:data}deepmd.pt.utils.env.GLOBAL_PT_FLOAT_PRECISION.

Register new arguments Copy Chapter

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 Copy Chapter

You may package new codes into a new Python package if you don't want to contribute it to the main DeePMD-kit repository.
A good example is DeePMD-GNN.
It's crucial to add your new component to project.entry-points."deepmd.pt" in pyproject.toml:

[project.entry-points."deepmd.pt"]
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.

Package customized C++ OPs Copy Section

You may need to use customized PyTorch C++ OPs in the new model.
Follow PyTorch documentation to create one library.

When using your customized C++ OPs in the Python interface, use {py:meth}torch.ops.load_library to load the OP library in the module defined in entry_points.

When using your customized C++ OPs in the C++ library, define the environment variable {envvar}DP_PLUGIN_PATH to load the OP library.

Unit tests Copy Chapter

Universal tests Copy Section

The source/tests/universal directory provides universal test suites for different models and backends.
The subdirectory cases defines fixtures for different test cases of models, atomic models, descriptors, and fitting networks.
The subdirectory dpmodel and pt are backend-specific test fixtures and suites.
Each test suite tests APIs and whether the serialized models give the same results as the original models.
Specially, the test suite for models also tests whether the model is permutation, translation, and rotation invariant, whether the model is differentiable and smooth near the cutoff radius, and whether the force is the negative gradient of the energy.

When adding a new model, add the fixture to the cases subdiretory and then apply the test fixture in the suite of different backends.
When implementing an existing model in a new backend, directly apply the existing test fixture to the test suite of that backend.

Consistent tests Copy Section

When transferring features from another backend to the PyTorch-TorchScript backend, it is essential to include a regression test in /source/tests/consistent to validate the consistency of the PyTorch-TorchScript backend with other backends. Presently, the regression tests cover self-consistency and cross-backend consistency between TensorFlow, PyTorch-TorchScript, and DP (NumPy) through the serialization/deserialization technique.

During the development of new components within the PyTorch-TorchScript backend, it is necessary to provide a DP (NumPy) implementation and incorporate corresponding regression tests. For PyTorch components, developers are also required to include a unit test using torch.jit.

File: doc/development/create-a-model-tf.md Copy Chapter

Create a model in TensorFlow {{ tensorflow_icon }}

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, trainer, etc, you may want to read this section.

To incorporate your custom model you'll need to:

1. Register and implement new components (e.g. descriptor) in a Python file. You may also want to register new TensorFlow OPs if necessary.
2. Register new arguments for user inputs.
3. Package new codes into a Python package.
4. Test new models.

Design a new component Copy Chapter

When creating a new component, take descriptor as the example, one should inherit from the {py:class}deepmd.tf.descriptor.descriptor.Descriptor class and override several methods. Abstract methods such as {py:class}deepmd.tf.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.tf.descriptor import Descriptor


@Descriptor.register("some_descrpt")
class SomeDescript(Descriptor):
    def __init__(self, arg1: bool, arg2: float) -> None:
        pass

Register new arguments Copy Chapter

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 Copy Chapter

You may package new codes into a new Python package if you don't want to contribute it to the main DeePMD-kit repository.
It's crucial to add your new component to project.entry-points."deepmd" in pyproject.toml:

[project.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.

Package customized C++ OPs Copy Section

You may need to use customized TensorFlow C++ OPs in the new model.
Follow TensorFlow documentation to create one library.

When using your customized C++ OPs in the Python interface, use {py:meth}tf.load_op_library to load the OP library in the module defined in entry_points.

When using your customized C++ OPs in the C++ library, define the environment variable {envvar}DP_PLUGIN_PATH to load the OP library.

File: doc/development/type-embedding.md Copy Chapter

Atom Type Embedding

Overview Copy Chapter

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 consists 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 the 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 reduces training complexity largely.

Preliminary Copy Chapter

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 Copy Chapter

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 Copy Chapter

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) Copy Section

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) Copy Section

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) Copy Section

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) Copy Section

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.

File: doc/freeze/compress.md Copy Chapter

Compress a model {{ tensorflow_icon }} {{ pytorch_icon }} {{ jax_icon }} {{ dpmodel_icon }}

[!NOTE]
Backends covered below: TensorFlow and TensorFlow 2
{{ tensorflow_icon }}, PyTorch-TorchScript and PyTorch-Exportable
{{ pytorch_icon }}, JAX {{ jax_icon }}, and DP {{ dpmodel_icon }}.

Theory Copy Chapter

The compression of the DP model uses three techniques, tabulated inference, operator merging, and precise neighbor indexing, to improve the performance of model training and inference when the model parameters are properly trained.

For better performance, the NN inference can be replaced by tabulated function evaluations if the input of the NN is of dimension one.
The idea is to approximate the output of the NN by a piece-wise polynomial fitting.
The input domain (a compact domain in $\mathbb R$) is divided into $L_c$ equally spaced intervals, in which we apply a fifth-order polynomial $g^l_m(x)$ approximation of the $m$-th output component of the NN function:

g^l_m(x) = a^l_m x^5 + b^l_m x^4 + c^l_m x^3 + d^l_m x^2 + e^l_m x + f^l_m,\quad
    x \in [x_l, x_{l+1}),

where $l=1,2,\dots,L_c$ is the index of the intervals, $x_1, \dots, x_{L_c}, x_{L_c+1}$ are the endpoints of the intervals, and $a^l_m$, $b^l_m$, $c^l_m$, $d^l_m$, $e^l_m$, and $f^l_m$ are the fitting parameters.
The fitting parameters can be computed by the equations below:

a^l_m = \frac{1}{2\Delta x_l^5}[12h_{m,l}-6(y'_{m,l+1}+y'_{m,l})\Delta x_l + (y''_{m,l+1}-y''_{m,l})\Delta x_l^2],

b^l_m = \frac{1}{2\Delta x_l^4}[-30h_{m,l} +(14y'_{m,l+1}+16y'_{m,l})\Delta x_l + (-2y''_{m,l+1}+3y''_{m,l})\Delta x_l^2],

c^l_m = \frac{1}{2\Delta x_l^3}[20h_{m,l}-(8y'_{m,l+1}+12y'_{m,l})\Delta x_l + (y''_{m,l+1}-3y''_{m,l})\Delta x_l^2],

d^l_m = \frac{1}{2}y''_{m,l},

e^l_m = y_{m,l}',

f^l_m = y_{m,l},

where $\Delta x_l=x_{l+1}-x_l$ denotes the size of the interval. $h_{m,l}=y_{m,l+1}-y_{m,l}$. $y_{m,l} = y_m(x_l)$, $y'_{m,l} = y'm(x_l)$ and $y''{m,l} = y''_m(x_l)$ are the value, the first-order derivative, and the second-order derivative of the $m$-th component of the target NN function at the interval point $x_l$, respectively.
The first and second-order derivatives are easily calculated by the back-propagation of the NN functions.

In the standard DP model inference, taking the two-body embedding descriptor as an example, the matrix product $(\mathcal G^i)^T \mathcal R$ requires the transfer of the tensor $\mathcal G^i$ between the register and the host/device memories, which usually becomes the bottle-neck of the computation due to the relatively small memory bandwidth of the GPUs.
The compressed DP model merges the matrix multiplication $(\mathcal G^i)^T \mathcal R$ with the tabulated inference step.
More specifically, once one column of the $(\mathcal G^i)^T$ is evaluated, it is immediately multiplied with one row of the environment matrix in the register, and the outer product is deposited to the result of $(\mathcal G^i)^T \mathcal R$.
By the operator merging technique, the allocation of $\mathcal G^i$ and the memory movement between register and host/device memories is avoided.
The operator merging of the three-body embedding can be derived analogously.

The first dimension, $N_c$, of the environment ($\mathcal R^i$) and embedding ($\mathcal G^i$) matrices is the expected maximum number of neighbors.
If the number of neighbors of an atom is smaller than $N_c$, the corresponding positions of the matrices are pad with zeros.
In practice, if the real number of neighbors is significantly smaller than $N_c$, a notable operation is spent on the multiplication of padding zeros.
In the compressed DP model, the number of neighbors is precisely indexed at the tabulated inference stage, further saving computational costs.[^1]

Instructions Copy Chapter

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:

::::{tab-set}

:::{tab-item} TensorFlow {{ tensorflow_icon }}

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

:::

:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}

dp --tf2 compress -i model.ckpt.tf2 -o model-compress.savedmodeltf

TensorFlow 2 compression reads a .tf2 training checkpoint directory or a
checkpoint prefix and writes a compressed .savedmodeltf model. See the
descriptor documentation for model-specific SavedModel export requirements.
:::

:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }}

dp --pt compress -i model.pth -o model-compress.pth

:::

:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}

dp --pt-expt compress -i dpa1-graph.pt2 -o dpa1-graph-compress.pt2

This command produces an executable compressed artifact only for descriptors
that support the PyTorch-Exportable compression path. See the
DPA-1 model-compression requirements
for the graph-lowered .pt2 route used in this example.
:::

:::{tab-item} JAX {{ jax_icon }}

dp --jax compress -i frozen_model.jax -o compressed_model.jax

JAX compression accepts .jax and .hlo inputs. Use .jax for the general,
lossless compressed serialization path. Descriptor pages document whether a
compressed model can also be exported to StableHLO .hlo.
:::

:::{tab-item} DP {{ dpmodel_icon }}

dp --dp compress -i model.dp -o model-compress.dp

DP compression accepts native .dp, .yaml, and .yml models.
:::

::::

where -i gives the original frozen model, -o gives the compressed model.
The compression entrypoints resolve the minimum neighbor distance and tabulate
supported descriptor embedding networks. If the model does not contain a
minimum neighbor distance, pass the training script with -t or
--training-script so it can be computed from the training data.

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

dp compress --help

An explanation will be provided

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

Not any descriptor supports model compression.
See the documentation of a specific descriptor to see whether it supports model compression.

Available activation functions for descriptor:

• tanh
• gelu
• relu
• relu6
• softplus
• sigmoid

Requirements of installation {{ pytorch_icon }} Copy Chapter

When compressing models in the PyTorch-TorchScript backend, the customized OP library for the Python interface must be installed when freezing the model.

The customized OP library for the Python interface is installed by default when building DeePMD-kit from source; see the installation guide for details.

[^1]: This section is built upon Jinzhe Zeng, Duo Zhang, Denghui Lu, Pinghui Mo, Zeyu Li, Yixiao Chen, Marián Rynik, Li'ang Huang, Ziyao Li, Shaochen Shi, Yingze Wang, Haotian Ye, Ping Tuo, Jiabin Yang, Ye Ding, Yifan Li, Davide Tisi, Qiyu Zeng, Han Bao, Yu Xia, Jiameng Huang, Koki Muraoka, Yibo Wang, Junhan Chang, Fengbo Yuan, Sigbjørn Løland Bore, Chun Cai, Yinnian Lin, Bo Wang, Jiayan Xu, Jia-Xin Zhu, Chenxing Luo, Yuzhi Zhang, Rhys E. A. Goodall, Wenshuo Liang, Anurag Kumar Singh, Sikai Yao, Jingchao Zhang, Renata Wentzcovitch, Jiequn Han, Jie Liu, Weile Jia, Darrin M. York, Weinan E, Roberto Car, Linfeng Zhang, Han Wang, J. Chem. Phys. 159, 054801 (2023) licensed under a Creative Commons Attribution (CC BY) license.

File: doc/freeze/freeze.md Copy Chapter

Freeze a model

The trained neural network is extracted from a checkpoint and dumped into a model file. This process is called "freezing" a model.
To freeze a model, typically one does

::::{tab-set}

:::{tab-item} TensorFlow {{ tensorflow_icon }}

dp freeze -o model.pb

in the folder where the model is trained. The output model is called model.pb.
The idea and part of our code are from Morgan.
:::

:::{tab-item} TensorFlow 2 {{ tensorflow_icon }}

dp --tf2 freeze -c model.ckpt -o model.savedmodeltf

When -c names a checkpoint prefix, the backend also checks the corresponding
path with .tf2 appended, so the example reads model.ckpt.tf2 and writes the
TensorFlow SavedModel to model.savedmodeltf. If -c is omitted, it defaults
to the current directory. For a multi-task checkpoint, select a branch with
--head CHOSEN_BRANCH. SavedModel export requires graph-traceable model code;
descriptor-specific export requirements are documented on the corresponding
model pages.
:::

:::{tab-item} PyTorch-TorchScript {{ pytorch_icon }}

dp --pt freeze -o model.pth

in the folder where the model is trained. The output model is called model.pth.

In multi-task mode, you need to choose one available heads (e.g. CHOSEN_BRANCH) by --head
to specify which model branch you want to freeze:

dp --pt freeze -o model_branch1.pth --head CHOSEN_BRANCH

The output model is called model_branch1.pth, which is the specifically frozen model with the CHOSEN_BRANCH head.
:::

:::{tab-item} PyTorch-Exportable {{ pytorch_icon }}

dp --pt-expt freeze -c model.ckpt.pt -o model

The backend writes .pte for the dense neighbor-list lower form and .pt2 for
the graph lower form. A suffixless output lets DeePMD-kit select the matching
extension. --lower-kind graph requires a graph-eligible model. Conversely, a
graph-capable DPA model may override a requested nlist lower with the graph
form and emit a warning. In multi-task mode, select a model branch with
--head CHOSEN_BRANCH.
:::

:::{tab-item} Paddle {{ paddle_icon }}

dp --pd freeze -o model

in the folder where the model is trained. The output model is called model.json and model.pdiparams.

In multi-task mode, you need to choose one available heads (e.g. CHOSEN_BRANCH) by --head
to specify which model branch you want to freeze:

dp --pd freeze -o model_branch1 --head CHOSEN_BRANCH

The output model is called model_branch1.json, which is the specifically frozen model with the CHOSEN_BRANCH head.
:::

:::{tab-item} JAX {{ jax_icon }}

dp --jax freeze -c model.ckpt.jax -o model.hlo

The JAX backend can write a StableHLO .hlo model, a lossless .jax
serialization, or a JAX2TF .savedmodel model. The .hlo and .savedmodel
formats work with the normal dp test/DeepPot route; .jax is intended for
checkpoint round-tripping and JAX-MD and is not a DeepEval model format. The
.savedmodel format requires TensorFlow and is the JAX format that supports the
C++ inference interface.
:::

::::

Freeze a JAX model with Hessian output {{ jax_icon }} Copy Chapter

Use --hessian to add coordinate-Hessian output to a frozen JAX energy model:

dp --jax freeze -c model.ckpt.jax -o model-hessian.hlo --hessian

The option applies to JAX .hlo, .jax, and .savedmodel outputs. A model
whose serialized definition already enables Hessian mode retains that mode even
when --hessian is omitted.

2. Official Technical Reference & Guides (deepmodeling/docs) Copy Chapter

File: README.md Copy Chapter

docs

The root of docs.deepmodeling.com.

File: index.rst Copy Chapter

========================== DeepModeling documentations

.. toctree::
:maxdepth: 1

DeePMD-kit
DP-GEN
dpdata
DPDispatcher
dargs
ABACUS https://abacus.deepmodeling.com/
dpgen2
DeepFlame https://deepflame.deepmodeling.com/
DPTI
ReacNetGenerator

.. toctree::
:caption: FAQ

faq/conda

• :ref:genindex
• :ref:modindex
• :ref:search

.. _feedback:
.. _affiliated packages:

File: requirements.txt Copy Chapter

sphinx_rtd_theme
deepmodeling_sphinx
myst-parser

File: .readthedocs.yaml Copy Chapter

.readthedocs.yaml

Read the Docs configuration file

See https://docs.readthedocs.io/en/stable/config-file/v2.html for details

Required

version: 2

Set the version of Python and other tools you might need

build:
os: ubuntu-22.04
tools:
python: "3.10"

Build documentation in the docs/ directory with Sphinx

sphinx:
configuration: conf.py

If using Sphinx, optionally build your docs in additional formats such as PDF

formats: all

Optionally declare the Python requirements required to build your docs

python:
install:

• requirements: requirements.txt

File: faq/conda.md Copy Chapter

Conda/pip environments

:::{note}
The following instructions are for the Linux operating system.
Most DeepModeling projects fully support the Linux OS but lack support for other OS.
If you use the Windows OS, you can use Windows Subsystem for Linux (WSL).
:::

How to setup a conda/pip base environment? Copy Chapter

:::{note}
The following instructions are for Miniforge but should also work for the offline packages provided by the DeepModeling projects.
:::

An environment is a directory that contains a specific collection of packages that you have installed.
Different environments are isolated from each other.
A base environment is required before setting up multiple environments.

First of all, download and execute conda-forge installer:

wget https://github.com/conda-forge/miniforge/releases/latest/download/Miniforge3-Linux-x86_64.sh -O miniforge.sh
bash miniforge.sh -b -p $HOME/miniforge3

This will install the conda base environment into $HOME/miniforge3. You can install the environment into a different path.

After installation, you can activate the base environment using

source $HOME/miniforge3/bin/activate

In order to initialize the environment whenever starting the bash shell, run

conda init --all

Now you can use conda install or pip install command to install conda or pip packages.

What is the difference between conda and pip? Copy Chapter

The largest difference is that pip can only install Python packages,
but conda can install any packages including C, C++, and Fortran packages without Python involved.
We suggest reading the Anaconda blog for more information.

How can I check what packages have been installed? Copy Chapter

Run conda list and pip list, which show all conda/pip packages that have been installed in the current environment.

How to create a new conda/pip environment? Copy Chapter

To create a conda/pip environment named my_env, run

conda create -n my_env python=3.11

Now you can activate the environment using

conda activate my_env

If you want to activate the environment in a shell script, use source instead:

source activate my_env

Why do I need multiple environments instead of one? Copy Chapter

Due to multiple reasons, dependencies of packages often conflict with each other.
For example, Python 3.11 and 3.12 cannot be installed in the same environment;
CUDA 11 and CUDA 12 packages cannot be installed in the same environment;
Open MPI and MPICH cannot be installed in the same environment;
packages built against zlib 1.2 and 1.3 cannot be installed in the same environment.
While conda-forge has tried its best to reduce package conflicts, it may still happen when installing multiple packages.
Installing different packages into different environments will reduce the risk of meeting the conflict packages.

How to use conda to install CUDA packages without CUDA installed? Copy Chapter

Usually, conda detects the CUDA driver on the system to install the compatible CUDA packages.
You don't need to worry about installing the CUDA Toolkit in advance.
However, sometimes, you may want to install CUDA packages into a machine that doesn't install CUDA drivers (for example, the login node of an HPC cluster).
In this case, you must manually set the environment variable CONDA_OVERRIDE_CUDA to tell conda which packages are compatible.

For example, you would like to use the CUDA packages on another machine with installed CUDA drivers. Run nvidia-smi on that machine,

+---------------------------------------------------------------------------------------+
| NVIDIA-SMI 545.29.06              Driver Version: 545.29.06    CUDA Version: 12.3     |
|-----------------------------------------+----------------------+----------------------+

The output shows that the CUDA driver version is 12.3.
Then, on the machine for performing installation, you may set the environment variable export CONDA_OVERRIDE_CUDA=12.3 before running conda install.

How to install conda/pip packages without the Internet access? Copy Chapter

To install conda/pip packages without Internet access, you need to have a machine that can connect to the Internet, and use conda or pip to install all packages in an environment.
Then, install conda-pack:

conda install conda-pack

Next, running conda pack to pack an environment, saying my_env:

conda pack -n my_env -o my_env.tar.gz

Transfer my_env.tar.gz to the machine that doesn't have the Internet access.
On that machine, you can unpack the environment:

mkdir -p my_env
tar -xzf my_env.tar.gz -C my_env
source my_env/bin/activate
conda-unpack

How to speed up conda/pip in China? Copy Chapter

When you are in China, we suggest adding Tsinghua University tuna mirrors to your conda/pip configs to speed up downloading conda/pip packages.

In order to set up Anaconda mirrors, append the following lines to the $HOME/.condarc file (if it doesn't exist, you can create one):

default_channels:
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
custom_channels:
  conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud

In order to set up PyPI mirrors, run

pip config set global.index-url https://pypi.tuna.tsinghua.edu.cn/simple

File: .github/workflows/mirror_gitee.yml Copy Chapter

name: Mirror to Gitee Repository

on: [ push, delete, create ]

Ensures that only one mirror task will run at a time.

concurrency:
group: git-mirror

jobs:
git-mirror:
uses: deepmodeling/workflows/.github/workflows/mirror_gitee.yml@main
secrets:
SYNC_GITEE_PRIVATE_KEY: ${{ secrets.SYNC_GITEE_PRIVATE_KEY }}