Devel update

.github/workflows/build_cc.yml

@@ -27,6 +27,10 @@ jobs:
         cache: 'pip'
     - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+         unzip libtorch.zip
     - run: |
          wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb \
          && sudo dpkg -i cuda-keyring_1.0-1_all.deb \
@@ -48,13 +52,17 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get install -y rocm-dev hipcub-dev
       if: matrix.variant == 'rocm'
-    - run: source/install/build_cc.sh
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: ${{ matrix.dp_variant }}
         DOWNLOAD_TENSORFLOW: "FALSE"
         CMAKE_GENERATOR: Ninja
       if: matrix.variant != 'clang'
-    - run: source/install/build_cc.sh
+    - run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: cpu
         DOWNLOAD_TENSORFLOW: "FALSE"

.github/workflows/build_wheel.yml

@@ -68,6 +68,17 @@ jobs:
       - uses: docker/setup-qemu-action@v3
         name: Setup QEMU
         if: matrix.platform_id == 'manylinux_aarch64' && matrix.os == 'ubuntu-latest'
+      # detect version in advance. See #3168
+      - uses: actions/setup-python@v5
+        name: Install Python
+        with:
+          python-version: '3.11'
+          cache: 'pip'
+        if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
+      - run: |
+          python -m pip install setuptools_scm
+          python -c "from setuptools_scm import get_version;print('SETUPTOOLS_SCM_PRETEND_VERSION_FOR_DEEPMD-KIT-CU11='+get_version())" >> $GITHUB_ENV
+        if: matrix.dp_pkg_name == 'deepmd-kit-cu11'
       - name: Build wheels
         uses: pypa/[email protected]
         env:

.github/workflows/codeql.yml

@@ -37,6 +37,8 @@ jobs:
          && sudo apt-get update \
          && sudo apt-get -y install cuda-cudart-dev-12-2 cuda-nvcc-12-2
         python -m pip install tensorflow
+        wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+        unzip libtorch.zip
       env:
         DEBIAN_FRONTEND: noninteractive
     # Initializes the CodeQL tools for scanning.
@@ -46,7 +48,9 @@ jobs:
         languages: ${{ matrix.language }}
         queries: security-extended,security-and-quality
     - name: "Run, Build Application using script"
-      run: source/install/build_cc.sh
+      run: |
+        export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+        source/install/build_cc.sh
       env:
         DP_VARIANT: cuda
         DOWNLOAD_TENSORFLOW: "FALSE"

.github/workflows/test_cc.yml

@@ -18,7 +18,13 @@ jobs:
         mpi: mpich
     - uses: lukka/get-cmake@latest
     - run: python -m pip install tensorflow
-    - run: source/install/test_cc_local.sh
+    - name: Download libtorch
+      run: |
+         wget https://download.pytorch.org/libtorch/cpu/libtorch-cxx11-abi-shared-with-deps-2.1.2%2Bcpu.zip -O libtorch.zip
+         unzip libtorch.zip
+    - run: |
+         export CMAKE_PREFIX_PATH=$GITHUB_WORKSPACE/libtorch
+         source/install/test_cc_local.sh
       env:
         OMP_NUM_THREADS: 1
         TF_INTRA_OP_PARALLELISM_THREADS: 1

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
       exclude: ^source/3rdparty
 -   repo: https://github.com/astral-sh/ruff-pre-commit
     # Ruff version.
-    rev: v0.1.13
+    rev: v0.1.14
     hooks:
     - id: ruff
       args: ["--fix"]

deepmd/infer/deep_pot.py

@@ -26,6 +26,7 @@
 from deepmd.utils.sess import (
     run_sess,
 )
+from deepmd_utils.infer.deep_pot import DeepPot as DeepPotBase
 
 if TYPE_CHECKING:
     from pathlib import (
@@ -35,7 +36,7 @@
 log = logging.getLogger(__name__)
 
 
-class DeepPot(DeepEval):
+class DeepPot(DeepEval, DeepPotBase):
     """Constructor.
 
     Parameters

deepmd_utils/infer/__init__.py

@@ -0,0 +1,6 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from .deep_pot import (
+    DeepPot,
+)
+
+__all__ = ["DeepPot"]

deepmd_utils/infer/backend.py

@@ -0,0 +1,33 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from enum import (
+    Enum,
+)
+
+
+class DPBackend(Enum):
+    """DeePMD-kit backend."""
+
+    TensorFlow = 1
+    PyTorch = 2
+    Paddle = 3
+    Unknown = 4
+
+
+def detect_backend(filename: str) -> DPBackend:
+    """Detect the backend of the given model file.
+
+    Parameters
+    ----------
+    filename : str
+        The model file name
+    """
+    if filename.endswith(".pb"):
+        return DPBackend.TensorFlow
+    elif filename.endswith(".pth") or filename.endswith(".pt"):
+        return DPBackend.PyTorch
+    elif filename.endswith(".pdmodel"):
+        return DPBackend.Paddle
+    return DPBackend.Unknown
+
+
+__all__ = ["DPBackend", "detect_backend"]

deepmd_utils/infer/deep_pot.py

@@ -0,0 +1,126 @@
+# SPDX-License-Identifier: LGPL-3.0-or-later
+from abc import (
+    ABC,
+    abstractmethod,
+)
+from typing import (
+    List,
+    Optional,
+    Tuple,
+    Union,
+)
+
+import numpy as np
+
+from deepmd_utils.utils.batch_size import (
+    AutoBatchSize,
+)
+
+from .backend import (
+    DPBackend,
+    detect_backend,
+)
+
+
+class DeepPot(ABC):
+    """Potential energy model.
+
+    Parameters
+    ----------
+    model_file : Path
+        The name of the frozen model file.
+    auto_batch_size : bool or int or AutoBatchSize, default: True
+        If True, automatic batch size will be used. If int, it will be used
+        as the initial batch size.
+    neighbor_list : ase.neighborlist.NewPrimitiveNeighborList, optional
+        The ASE neighbor list class to produce the neighbor list. If None, the
+        neighbor list will be built natively in the model.
+    """
+
+    @abstractmethod
+    def __init__(
+        self,
+        model_file,
+        *args,
+        auto_batch_size: Union[bool, int, AutoBatchSize] = True,
+        neighbor_list=None,
+        **kwargs,
+    ) -> None:
+        pass
+
+    def __new__(cls, model_file: str, *args, **kwargs):
+        if cls is DeepPot:
+            backend = detect_backend(model_file)
+            if backend == DPBackend.TensorFlow:
+                from deepmd.infer.deep_pot import DeepPot as DeepPotTF
+
+                return super().__new__(DeepPotTF)
+            elif backend == DPBackend.PyTorch:
+                from deepmd_pt.infer.deep_eval import DeepPot as DeepPotPT
+
+                return super().__new__(DeepPotPT)
+            else:
+                raise NotImplementedError("Unsupported backend: " + str(backend))
+        return super().__new__(cls)
+
+    @abstractmethod
+    def eval(
+        self,
+        coords: np.ndarray,
+        cells: np.ndarray,
+        atom_types: List[int],
+        atomic: bool = False,
+        fparam: Optional[np.ndarray] = None,
+        aparam: Optional[np.ndarray] = None,
+        efield: Optional[np.ndarray] = None,
+        mixed_type: bool = False,
+    ) -> Tuple[np.ndarray, ...]:
+        """Evaluate energy, force, and virial. If atomic is True,
+        also return atomic energy and atomic virial.
+
+        Parameters
+        ----------
+        coords : np.ndarray
+            The coordinates of the atoms, in shape (nframes, natoms, 3).
+        cells : np.ndarray
+            The cell vectors of the system, in shape (nframes, 9). If the system
+            is not periodic, set it to None.
+        atom_types : List[int]
+            The types of the atoms. If mixed_type is False, the shape is (natoms,);
+            otherwise, the shape is (nframes, natoms).
+        atomic : bool, optional
+            Whether to return atomic energy and atomic virial, by default False.
+        fparam : np.ndarray, optional
+            The frame parameters, by default None.
+        aparam : np.ndarray, optional
+            The atomic parameters, by default None.
+        efield : np.ndarray, optional
+            The electric field, by default None.
+        mixed_type : bool, optional
+            Whether the system contains mixed atom types, by default False.
+
+        Returns
+        -------
+        energy
+            The energy of the system, in shape (nframes,).
+        force
+            The force of the system, in shape (nframes, natoms, 3).
+        virial
+            The virial of the system, in shape (nframes, 9).
+        atomic_energy
+            The atomic energy of the system, in shape (nframes, natoms). Only returned
+            when atomic is True.
+        atomic_virial
+            The atomic virial of the system, in shape (nframes, natoms, 9). Only returned
+            when atomic is True.
+        """
+        # This method has been used by:
+        # documentation python.md
+        # dp model_devi: +fparam, +aparam, +mixed_type
+        # dp test: +atomic, +fparam, +aparam, +efield, +mixed_type
+        # finetune: +mixed_type
+        # dpdata
+        # ase
+
+
+__all__ = ["DeepPot"]

doc/_static/css/custom.css

@@ -7,8 +7,14 @@ pre{
 .wy-side-nav-search .wy-dropdown > a img.logo, .wy-side-nav-search > a img.logo {
     width: 275px;
 }
+img.platform-icon {
+	height: 2ex;
+}
 @media (prefers-color-scheme: dark) {
 	.wy-side-nav-search .wy-dropdown > a img.logo, .wy-side-nav-search > a img.logo {
 		content: url("../logo-dark.svg");
 	}
+	img.platform-icon {
+		filter: invert(1);
+	}
 }

doc/conf.py

@@ -213,7 +213,10 @@ def setup(app):
 exhale_projects_args = {
     "cc": {
         "containmentFolder": "./API_CC",
-        "exhaleDoxygenStdin": "INPUT = ../source/api_cc/include/",
+        "exhaleDoxygenStdin": """INPUT = ../source/api_cc/include/
+                                 PREDEFINED += BUILD_TENSORFLOW
+                                               BUILD_PYTORCH
+        """,
         "rootFileTitle": "C++ API",
         "rootFileName": "api_cc.rst",
     },
@@ -275,6 +278,11 @@ def setup(app):
 .. |PRECISION| replace:: {list_to_doc(PRECISION_DICT.keys())}
 """
 
+myst_substitutions = {
+    "tensorflow_icon": """![TensorFlow](/_static/tensorflow.svg){class=platform-icon}""",
+    "pytorch_icon": """![PyTorch](/_static/pytorch.svg){class=platform-icon}""",
+}
+
 # -- Options for HTML output -------------------------------------------------
 
 # The theme to use for HTML and HTML Help pages.  See the documentation for
@@ -298,6 +306,8 @@ def setup(app):
 myst_enable_extensions = [
     "dollarmath",
     "colon_fence",
+    "substitution",
+    "attrs_inline",
 ]
 myst_fence_as_directive = ("math",)
 # fix emoji issue in pdf

doc/freeze/compress.md

@@ -1,4 +1,8 @@
-# Compress a model
+# Compress a model {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 ## Theory
 

doc/model/dplr.md

@@ -1,4 +1,8 @@
-# Deep potential long-range (DPLR)
+# Deep potential long-range (DPLR) {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 Notice: **The interfaces of DPLR are not stable and subject to change**
 

doc/model/dprc.md

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

doc/model/linear.md

@@ -1,4 +1,8 @@
-## Linear model
+## Linear model {{ tensorflow_icon }}
+
+:::{note}
+**Supported backends**: TensorFlow {{ tensorflow_icon }}
+:::
 
 One can linearly combine existing models with arbitrary coefficients:
 

doc/model/pairtab.md

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

doc/model/train-energy-spin.md

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

doc/model/train-energy.md

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