feat(lammps): support Cartesian native-spin magnetic forces

.github/workflows/test_cc.yml

@@ -94,6 +94,7 @@ jobs:
           TF_INTRA_OP_PARALLELISM_THREADS: 1
           TF_INTER_OP_PARALLELISM_THREADS: 1
           LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
+          DPA4SPIN_LMP_REQUIRED: ${{ matrix.enable_pytorch && '1' || '0' }}
           ENABLE_TENSORFLOW: ${{ matrix.enable_tensorflow && '1' || '0' }}
           ENABLE_PYTORCH: ${{ matrix.enable_pytorch && '1' || '0' }}
           ENABLE_JAX: ${{ matrix.enable_tensorflow && '1' || '0' }}

.github/workflows/test_cuda.yml

@@ -92,7 +92,7 @@ jobs:
         with:
           python-version: "3.11"
           # cache: 'pip'
-      - name: Install wget and unzip
+      - name: Install build prerequisites
         run: sudo apt-get update && sudo apt-get install -y wget unzip
       - uses: lukka/get-cmake@latest
         with:
@@ -115,14 +115,42 @@ jobs:
       - run: |
           export PYTORCH_ROOT=$(python -c 'import torch;print(torch.__path__[0])')
           export TENSORFLOW_ROOT=$(python -c 'import importlib.util,pathlib;print(pathlib.Path(importlib.util.find_spec("tensorflow").origin).parent)')
-          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit
+          # Use the LAMMPS wheel's MPI distribution for every build and launch.
+          # Ubuntu 24.04 MPICH can start independent singleton worlds:
+          # https://github.com/mpi4py/mpi4py/issues/547
+          source/install/uv_with_retry.sh pip install --system mpich
+          export MPI_HOME=$(python -c 'import sys; print(sys.prefix)')
+          export MPICC=$MPI_HOME/bin/mpicc
+          export LD_LIBRARY_PATH=$MPI_HOME/lib:${LD_LIBRARY_PATH:-}
+          echo "MPI_HOME=$MPI_HOME" >> "$GITHUB_ENV"
+          source/install/uv_with_retry.sh pip install --system -v -e .[gpu,test,lmp,cu12,torch,jax] mpi4py --reinstall-package deepmd-kit --reinstall-package mpi4py --no-binary mpi4py
           # See https://github.com/jax-ml/jax/issues/29042
           source/install/uv_with_retry.sh pip install --system -U 'nvidia-cublas-cu12>=12.9.0.13'
         env:
           DP_VARIANT: cuda
           DP_ENABLE_NATIVE_OPTIMIZATION: 1
           DP_ENABLE_PYTORCH: 1
       - run: dp --version
+      - name: Verify the two-rank MPI runtime
+        run: |
+          export LD_LIBRARY_PATH=$MPI_HOME/lib:${LD_LIBRARY_PATH:-}
+          timeout 60 "$MPI_HOME/bin/mpiexec" -n 2 python -c '
+          import os
+          from pathlib import Path
+          from mpi4py import MPI
+          comm = MPI.COMM_WORLD
+          assert comm.size == 2, MPI.Get_library_version()
+          assert comm.allgather(comm.rank) == [0, 1]
+          libraries = {
+              Path(line.split()[-1]).resolve()
+              for line in Path("/proc/self/maps").read_text().splitlines()
+              if Path(line.split()[-1]).name.startswith("libmpi")
+          }
+          expected = (Path(os.environ["MPI_HOME"]) / "lib").resolve()
+          assert libraries and all(p.parent == expected for p in libraries), libraries
+          if comm.rank == 0:
+              print(MPI.Get_library_version(), libraries)
+          '
       - name: Test the default CUDA dtype contract
         run: |
           python -c "import torch; import deepmd.pt.cxx_op; from deepmd.pt_expt.kernels.cuda.dpa4.so2_conv_train import op_available, float64_available; assert torch.cuda.is_available(); assert op_available(); assert float64_available() is False"
@@ -133,18 +161,28 @@ jobs:
       - name: Convert models
         run: source/tests/infer/convert-models.sh
       - run: |
-          export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
+          export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$LD_LIBRARY_PATH
           source/install/test_cc_local.sh
         env:
           OMP_NUM_THREADS: 1
           TF_INTRA_OP_PARALLELISM_THREADS: 1
           TF_INTER_OP_PARALLELISM_THREADS: 1
           CMAKE_GENERATOR: Ninja
           DP_VARIANT: cuda
-          DP_USE_MPICH2: 1
           ENABLE_PADDLE: FALSE
+      - name: Test native-spin magnetic forces with CUDA Kokkos LAMMPS
+        run: |
+          export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:${LD_LIBRARY_PATH:-}
+          source/install/test_lammps_kokkos.sh
+        env:
+          CMAKE_GENERATOR: Ninja
+          DP_VARIANT: cuda
+          CUDA_VISIBLE_DEVICES: 0
+          OMP_NUM_THREADS: 1
+          TF_INTRA_OP_PARALLELISM_THREADS: 1
+          TF_INTER_OP_PARALLELISM_THREADS: 1
       - run: |
-          export LD_LIBRARY_PATH=$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH
+          export LD_LIBRARY_PATH=$MPI_HOME/lib:$CUDA_PATH/lib64:/usr/lib/x86_64-linux-gnu/:$GITHUB_WORKSPACE/dp_test/lib:$LD_LIBRARY_PATH
           export PATH=$GITHUB_WORKSPACE/dp_test/bin:$PATH
           python -m pytest -s source/lmp/tests || (cat log.lammps && exit 1)
           python -m pytest source/ipi/tests
@@ -153,6 +191,7 @@ jobs:
           TF_INTRA_OP_PARALLELISM_THREADS: 1
           TF_INTER_OP_PARALLELISM_THREADS: 1
           LAMMPS_PLUGIN_PATH: ${{ github.workspace }}/dp_test/lib/deepmd_lmp
+          DPA4SPIN_LMP_REQUIRED: 1
           CUDA_VISIBLE_DEVICES: 0
           ENABLE_PADDLE: 0
   pass:

doc/model/dpa4c.md

@@ -604,6 +604,20 @@ already performs, and the magnetic force is reduced back onto owning atoms
 alongside the conservative force, so domain decomposition needs no additional
 exchange.
 
+By default the pair styles store the precession magnetic force expected by the
+LAMMPS SPIN package (the model energy gradient scaled by the moment magnitude
+over `hbar`). A Newtonian Cartesian integrator instead needs the raw energy
+gradient `-dE/dS`. Select it with `magnetic_force cartesian` (requires
+`units metal`); the default is `magnetic_force precession`:
+
+```lammps
+pair_style        dpa4spin/kk compressed_model.pt2 magnetic_force cartesian
+```
+
+The selected convention is exposed as the scalar integer
+`Pair::extract("fm_cartesian_force", dim)` (zero for precession, one for
+Cartesian) so a consuming integrator can reject an incompatible configuration.
+
 `min_style spin` reads the magnetic force from the pair style and relaxes the
 moment directions. Spin dynamics through `fix nve/spin` requires a LAMMPS build
 whose fix recognizes this pair style, because the stock fix accumulates the

examples/spin/dpa4c/lmp/README.md

@@ -67,3 +67,12 @@ they stay fixed through that stage.
 Spin dynamics through `fix nve/spin` requires a LAMMPS build whose fix
 recognizes this pair style, because the stock fix accumulates the magnetic
 force only from pair styles matching its own name pattern.
+
+The default magnetic-force convention is precession (`|m| / hbar` scaling of
+the model energy gradient). A separately maintained Newtonian Cartesian
+integrator can request `-dE/dS` with only the energy-per-model-spin unit
+conversion via `magnetic_force cartesian` on the `pair_style` line
+(`units metal`) and can confirm the choice through
+`Pair::extract("fm_cartesian_force", dim)`. On `dpa4spin/kk`, pair that keyword
+with the compressed compact-canonical archive (not the graph-lower freeze
+output); host `dpa4spin` may use either archive.

examples/spin/dpa4c/lmp/in.lammps

@@ -11,6 +11,15 @@
 # graph and the moment in device memory:
 #
 #   lmp -k on g 1 -sf kk -in in.lammps
+#
+# By default atom->fm holds the precession magnetic force. For a Newtonian
+# Cartesian integrator that needs -dE/dS, append `magnetic_force cartesian`
+# (units metal) and read Pair::extract("fm_cartesian_force", dim). Host
+# `dpa4spin` accepts the graph-lower archive from freeze; `dpa4spin/kk`
+# requires the compact canonical archive from `dp --pt-expt compress`
+# (see README.md):
+#
+#   pair_style        dpa4spin/kk compressed_model.pt2 magnetic_force cartesian
 
 units             metal
 dimension         3

source/install/test_cc_local.sh

@@ -134,6 +134,11 @@ else:
 	# exports NO with-comm artifact (single-rank only by construction).
 	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4_spin_zbl.py &
 	PID14=$!
+	# Minimal native-spin DPA4C for LAMMPS pair_style dpa4spin (no
+	# fparam/aparam/chg_spin). CPU-safe; used by the magnetic-force
+	# convention behavioral tests.
+	env ${_GEN_ENV} python ${INFER_SCRIPT_PATH}/gen_dpa4c_spin_lmp.py &
+	PID14B=$!
 	# Native-spin DPA4C: the archives above all ship a with-comm artifact
 	# and are therefore served by DeepSpinPTExpt, which leaves
 	# NativeSpinPTExpt -- the class the dispatch selects for a native-spin
@@ -149,6 +154,7 @@ else:
 	wait $PID12
 	wait $PID13
 	wait $PID14
+	wait $PID14B
 	if [ "${DP_VARIANT}" = "cuda" ]; then
 		wait $PID15
 	fi

source/install/test_lammps_kokkos.sh

@@ -0,0 +1,108 @@
+#!/bin/bash
+set -euxo pipefail
+
+# Run after test_cc_local.sh: reuse its C API/backend installation and the
+# matching LAMMPS checkout instead of rebuilding DeePMD or using the pip wheel.
+SCRIPT_PATH=$(dirname "$(realpath -s "$0")")
+REPO_ROOT=$(realpath "${SCRIPT_PATH}/../..")
+DEEPMD_C_ROOT=${DEEPMD_C_ROOT:-${REPO_ROOT}/dp_test}
+LAMMPS_SOURCE_ROOT=${LAMMPS_SOURCE_ROOT:-${REPO_ROOT}/source/build_tests/_deps/lammps_download-src}
+BUILD_ROOT=${REPO_ROOT}/source/build_lammps_kokkos
+LAMMPS_BUILD_DIR=${BUILD_ROOT}/lammps
+PLUGIN_BUILD_DIR=${BUILD_ROOT}/plugin
+PLUGIN_DIR=${BUILD_ROOT}/autoload
+NVCC_WRAPPER=${LAMMPS_SOURCE_ROOT}/lib/kokkos/bin/nvcc_wrapper
+# CUDA compilation can use substantial memory per process.
+BUILD_JOBS=${CMAKE_BUILD_PARALLEL_LEVEL:-4}
+
+if [[ ${DP_VARIANT:-} != cuda ]]; then
+	echo "This regression requires DP_VARIANT=cuda." >&2
+	exit 1
+fi
+test -f "${DEEPMD_C_ROOT}/lib/libdeepmd_c.so"
+test -f "${LAMMPS_SOURCE_ROOT}/cmake/CMakeLists.txt"
+test -x "${NVCC_WRAPPER}"
+command -v nvcc
+python -c 'import torch; assert torch.cuda.is_available(), "CUDA device is required"'
+
+# Use one MPI prefix for headers, libraries, Python bindings and launcher.
+# The CUDA job installs the MPICH wheel before building mpi4py and DeePMD.
+MPI_ROOT=${MPI_HOME:-$(python -c 'import sys; print(sys.prefix)')}
+MPI_C_COMPILER=$(command -v "${MPI_ROOT}/bin/mpicc")
+MPI_CXX_COMPILER=$(command -v "${MPI_ROOT}/bin/mpicxx")
+export MPI4PY_MPIABI=mpich
+export LD_LIBRARY_PATH=${MPI_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+export DPA4SPIN_KOKKOS_MPIEXEC
+DPA4SPIN_KOKKOS_MPIEXEC=$(command -v "${MPI_ROOT}/bin/mpiexec")
+MPI_ARGS=(
+	"-DMPI_C_COMPILER=${MPI_C_COMPILER}"
+	"-DMPI_CXX_COMPILER=${MPI_CXX_COMPILER}"
+)
+
+# Kokkos detects the GPU architecture on the runner. Both builds use its
+# nvcc_wrapper and the plugin imports the exact Kokkos build configuration.
+cmake -S "${LAMMPS_SOURCE_ROOT}/cmake" -B "${LAMMPS_BUILD_DIR}" \
+	-D CMAKE_BUILD_TYPE=Release \
+	-D CMAKE_CXX_COMPILER="${NVCC_WRAPPER}" \
+	-D CMAKE_CXX_STANDARD=17 \
+	-D BUILD_SHARED_LIBS=ON \
+	-D BUILD_MPI=ON \
+	-D BUILD_OMP=OFF \
+	-D PKG_KOKKOS=ON \
+	-D Kokkos_ENABLE_CUDA=ON \
+	-D Kokkos_ENABLE_SERIAL=ON \
+	-D PKG_SPIN=ON \
+	-D PKG_PLUGIN=ON \
+	-D PKG_KSPACE=ON \
+	-D PKG_EXTRA-FIX=ON \
+	-D FFT=KISS \
+	-D FFT_KOKKOS=KISS \
+	"${MPI_ARGS[@]}"
+cmake --build "${LAMMPS_BUILD_DIR}" --target lammps --parallel "${BUILD_JOBS}"
+
+cmake -S "${REPO_ROOT}/source" -B "${PLUGIN_BUILD_DIR}" \
+	-D CMAKE_BUILD_TYPE=Release \
+	-D CMAKE_CXX_COMPILER="${NVCC_WRAPPER}" \
+	-D CMAKE_CXX_STANDARD=17 \
+	-D CMAKE_DISABLE_PRECOMPILE_HEADERS=ON \
+	-D BUILD_TESTING=OFF \
+	-D DEEPMD_C_ROOT="${DEEPMD_C_ROOT}" \
+	-D DP_USING_C_API=ON \
+	-D LAMMPS_SOURCE_ROOT="${LAMMPS_SOURCE_ROOT}" \
+	-D DEEPMD_LAMMPS_KOKKOS=ON \
+	-D Kokkos_DIR="${LAMMPS_BUILD_DIR}/lib/kokkos" \
+	"${MPI_ARGS[@]}"
+cmake --build "${PLUGIN_BUILD_DIR}" --target deepmd_lmp --parallel "${BUILD_JOBS}"
+mkdir -p "${PLUGIN_DIR}"
+ln -sfn "${PLUGIN_BUILD_DIR}/lmp/plugin/libdeepmd_lmp.so" "${PLUGIN_DIR}/dpplugin.so"
+
+export LD_LIBRARY_PATH=${LAMMPS_BUILD_DIR}:${DEEPMD_C_ROOT}/lib${LD_LIBRARY_PATH:+:${LD_LIBRARY_PATH}}
+python "${REPO_ROOT}/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py"
+
+# The wheel's Python package loads its own liblammps by absolute path. Merely
+# setting LD_LIBRARY_PATH would therefore still test the non-Kokkos wheel.
+export PYTHONPATH=${LAMMPS_SOURCE_ROOT}/python${PYTHONPATH:+:${PYTHONPATH}}
+export LAMMPS_PLUGIN_PATH=${PLUGIN_DIR}
+python - "${LAMMPS_SOURCE_ROOT}" "${LAMMPS_BUILD_DIR}" <<'PY'
+import sys
+from pathlib import Path
+
+import lammps
+
+source, build = map(Path, sys.argv[1:])
+assert Path(lammps.__file__).resolve() == (
+    source / "python/lammps/__init__.py"
+).resolve(), lammps.__file__
+with lammps.lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp:
+    mapped_libraries = {
+        Path(line.split()[-1]).resolve()
+        for line in Path("/proc/self/maps").read_text().splitlines()
+        if "/liblammps" in line
+    }
+    assert mapped_libraries == {(build / "liblammps.so").resolve()}, mapped_libraries
+    assert lmp.has_style("pair", "dpa4spin/kk"), "Kokkos plugin was not loaded"
+    print(f"Verified matching LAMMPS Python and shared library: {mapped_libraries}")
+PY
+
+DPA4SPIN_KOKKOS_REQUIRED=1 python -m pytest -v -s \
+	"${REPO_ROOT}/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py"