breaking: enable PyTorch backend for PyPI LAMMPS

backend/read_env.py

@@ -83,7 +83,7 @@ def get_argument_from_env() -> tuple[str, list, list, dict, str, str]:
         cmake_args.append(f"-DLAMMPS_VERSION={dp_lammps_version}")
     if dp_ipi == "1":
         cmake_args.append("-DENABLE_IPI:BOOL=TRUE")
-        extra_scripts["dp_ipi"] = "deepmd.tf.entrypoints.ipi:dp_ipi"
+        extra_scripts["dp_ipi"] = "deepmd.entrypoints.ipi:dp_ipi"
 
     if os.environ.get("DP_ENABLE_TENSORFLOW", "1") == "1":
         tf_install_dir, _ = find_tensorflow()

deepmd/tf/entrypoints/ipi.py → deepmd/entrypoints/ipi.py

@@ -8,7 +8,7 @@
     NoReturn,
 )
 
-from deepmd.tf.lmp import (
+from deepmd.lmp import (
     get_op_dir,
 )
 

deepmd/tf/lmp.py → deepmd/lmp.py

@@ -13,12 +13,15 @@
     Optional,
 )
 
+import torch  # noqa: TID253
 from packaging.version import (
     Version,
 )
 
-from deepmd.tf.env import (
+from deepmd.env import (
     SHARED_LIB_DIR,
+)
+from deepmd.tf.env import (  # noqa: TID253
     TF_VERSION,
     tf,
 )
@@ -75,6 +78,7 @@ def get_library_path(module: str, filename: str) -> list[str]:
     raise RuntimeError("Unsupported platform")
 
 tf_dir = tf.sysconfig.get_lib()
+pt_dir = os.path.join(torch.__path__[0], "lib")
 op_dir = str(SHARED_LIB_DIR)
 
 
@@ -106,6 +110,7 @@ def get_library_path(module: str, filename: str) -> list[str]:
         os.environ.get(lib_env),
         tf_dir,
         os.path.join(tf_dir, "python"),
+        pt_dir,
         op_dir,
     ]
 )

doc/install/easy-install.md

@@ -114,14 +114,6 @@ pip install deepmd-kit[cpu]
 
 :::::
 
-[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are only provided on Linux and macOS for the TensorFlow backend. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` to extras:
-
-```bash
-pip install deepmd-kit[gpu,cu12,lmp,ipi]
-```
-
-MPICH is required for parallel running.
-
 ::::::
 
 ::::::{tab-item} PyTorch {{ pytorch_icon }}
@@ -228,3 +220,11 @@ The supported platform includes Linux x86-64 and aarch64 with GNU C Library 2.28
 :::{Warning}
 If your platform is not supported, or you want to build against the installed backends, or you want to enable ROCM support, please [build from source](install-from-source.md).
 :::
+
+[The LAMMPS module](../third-party/lammps-command.md) and [the i-PI driver](../third-party/ipi.md) are provided on Linux and macOS for the TensorFlow, PyTorch, and JAX backend. It requires both TensorFlow and PyTorch. To install LAMMPS and/or i-PI, add `lmp` and/or `ipi` to extras:
+
+```bash
+pip install deepmd-kit[gpu,cu12,lmp,ipi]
+```
+
+MPICH is required for parallel running.

pyproject.toml

@@ -58,7 +58,7 @@ requires-python = ">=3.9"
 keywords = ["deepmd"]
 
 [project.entry-points."lammps.plugins"]
-deepmd = "deepmd.tf.lmp:get_op_dir"
+deepmd = "deepmd.lmp:get_op_dir"
 
 [project.entry-points."dpgui"]
 "DeePMD-kit" = "deepmd.utils.argcheck:gen_args"
@@ -227,7 +227,7 @@ test-command = [
     """python -c "import deepmd.tf;import deepmd.pt;import deepmd.pd" """,
     "dp -h",
     "dp_ipi",
-    "pytest {project}/source/tests/tf/test_lammps.py"
+    "pytest {project}/source/tests/common/test_lammps.py"
 ]
 test-extras = ["cpu", "test", "lmp", "ipi", "torch", "paddle"]
 build = ["cp311-*"]

source/CMakeLists.txt

@@ -382,6 +382,10 @@ if(ENABLE_PYTORCH AND NOT DEEPMD_C_ROOT)
   get_filename_component(PyTorch_LIBRARY_PATH ${_TORCH_LOCATION} DIRECTORY)
   list(APPEND BACKEND_LIBRARY_PATH ${PyTorch_LIBRARY_PATH})
   list(APPEND BACKEND_INCLUDE_DIRS ${TORCH_INCLUDE_DIRS})
+  if(USE_PT_PYTHON_LIBS OR BUILD_PY_IF)
+    # when libtorch.so is in a wheel
+    link_directories(${PyTorch_LIBRARY_PATH}/../../torch.libs)
+  endif()
 endif()
 if(ENABLE_JAX
    AND BUILD_CPP_IF

source/api_cc/CMakeLists.txt

@@ -16,10 +16,7 @@ if(ENABLE_TENSORFLOW)
                                            TensorFlow::tensorflow_framework)
   target_compile_definitions(${libname} PRIVATE BUILD_TENSORFLOW)
 endif()
-if(ENABLE_PYTORCH
-   AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}"
-       # LAMMPS and i-PI in the Python package are not ready - needs more work
-   AND NOT BUILD_PY_IF)
+if(ENABLE_PYTORCH AND "${OP_CXX_ABI_PT}" EQUAL "${OP_CXX_ABI}")
   target_link_libraries(${libname} PRIVATE "${TORCH_LIBRARIES}")
   target_compile_definitions(${libname} PRIVATE BUILD_PYTORCH)
 endif()

source/tests/tf/test_lammps.py → source/tests/common/test_lammps.py

@@ -3,7 +3,7 @@
 import subprocess
 import unittest
 
-from deepmd.tf.utils.convert import (
+from deepmd.tf.utils.convert import (  # noqa: TID253
     convert_pbtxt_to_pb,
 )
 
@@ -27,6 +27,10 @@ def setUpClass(cls) -> None:
             str(cls.work_dir / "deeppot.pbtxt"), str(cls.work_dir / "deep_pot.pb")
         )
 
+    @unittest.skipIf(
+        os.environ.get("CUDA_VERSION", "").startswith("11"),
+        "CUDA 11.x wheel uses PyTorch 2.3 which is not ABI compatible with TensorFlow",
+    )
     def test_lmp(self) -> None:
         in_file = (self.work_dir / "in.test").absolute()
         subprocess.check_call(["lmp", "-in", str(in_file)], cwd=str(self.work_dir))

source/tests/infer/in.test

@@ -11,7 +11,7 @@ read_data	../../../examples/water/lmp/water.lmp
 mass 		1 16
 mass		2 2
 
-pair_style	deepmd deep_pot.pb
+pair_style	deepmd deep_pot.pb deeppot_sea.pth
 pair_coeff  * *
 
 velocity        all create 330.0 23456789