diff --git a/.github/workflows/test_cc.yml b/.github/workflows/test_cc.yml index 13e3e51a32..361b709704 100644 --- a/.github/workflows/test_cc.yml +++ b/.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' }} diff --git a/.github/workflows/test_cuda.yml b/.github/workflows/test_cuda.yml index 01f8dcfa76..58284a2e4c 100644 --- a/.github/workflows/test_cuda.yml +++ b/.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,7 +115,15 @@ 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: @@ -123,6 +131,26 @@ jobs: 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,7 +161,7 @@ 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 @@ -141,10 +169,20 @@ jobs: 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: diff --git a/doc/model/dpa4c.md b/doc/model/dpa4c.md index 7e79d8098e..4564c2f460 100644 --- a/doc/model/dpa4c.md +++ b/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 diff --git a/examples/spin/dpa4c/lmp/README.md b/examples/spin/dpa4c/lmp/README.md index aa4153f4f6..d13be52950 100644 --- a/examples/spin/dpa4c/lmp/README.md +++ b/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. diff --git a/examples/spin/dpa4c/lmp/in.lammps b/examples/spin/dpa4c/lmp/in.lammps index 050ce53903..d55fde1cdf 100644 --- a/examples/spin/dpa4c/lmp/in.lammps +++ b/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 diff --git a/source/install/test_cc_local.sh b/source/install/test_cc_local.sh index 17b4f2f593..8fc4fcc4d9 100755 --- a/source/install/test_cc_local.sh +++ b/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 diff --git a/source/install/test_lammps_kokkos.sh b/source/install/test_lammps_kokkos.sh new file mode 100755 index 0000000000..b5f2506c16 --- /dev/null +++ b/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" diff --git a/source/lmp/pair_dpa4spin.cpp b/source/lmp/pair_dpa4spin.cpp index 0a8d6e1ba8..01ba78d4d2 100644 --- a/source/lmp/pair_dpa4spin.cpp +++ b/source/lmp/pair_dpa4spin.cpp @@ -107,9 +107,7 @@ static const char cite_user_deepmd_package[] = "}\n\n"; namespace { -// Reduced Planck constant in eV.ps. The model reports the magnetic force as -// the energy gradient with respect to the magnetic moment, while LAMMPS stores -// the precession force; the two differ by the factor hbar / |m|. +// Reduced Planck constant in eV.ps, used by the default precession convention. constexpr double kHBar = 6.5821191e-04; // Positions of the LAMMPS global virial components (xx, yy, zz, xy, xz, yz) // within the nine-component tensor the model reports. @@ -120,7 +118,11 @@ constexpr int kCentroidVirialMap[9] = {0, 4, 8, 3, 6, 7, 1, 2, 5}; } // namespace PairDPA4Spin::PairDPA4Spin(LAMMPS* lmp) - : Pair(lmp), scale(nullptr), cutoff(0.0), commdata_(nullptr) { + : Pair(lmp), + scale(nullptr), + cutoff(0.0), + fm_cartesian_force(0), + commdata_(nullptr) { if (lmp->citeme) { lmp->citeme->add(cite_user_deepmd_package); } @@ -237,46 +239,81 @@ void PairDPA4Spin::settings(int narg, char** arg) { } cutoff = deep_spin.cutoff() * dist_unit_cvt_factor; - // How many values name a charge state is a property of the artifact, so the - // keyword can only be read once the model is loaded. + // The artifact determines how many values follow charge_spin. + // magnetic_force independently selects the convention stored in atom->fm. const int dim_chg_spin = deep_spin.dim_chg_spin(); std::vector charge_spin; + bool magnetic_force_seen = false; + fm_cartesian_force = 0; int iarg = 1; while (iarg < narg) { - if (std::string(arg[iarg]) != "charge_spin") { + const std::string key(arg[iarg]); + if (key == "charge_spin") { + // One charge state holds for the whole run, so a second occurrence names + // no state the style could serve. Values already read mark the keyword + // as seen: a successful read always contributes at least one value. + if (!charge_spin.empty()) { + error->all(FLERR, + "Illegal pair_style command: keyword charge_spin names the " + "single charge state of the whole run, so it may appear " + "only once."); + } + if (dim_chg_spin == 0) { + error->all(FLERR, + "Illegal pair_style command: the artifact served by pair " + "style " + + style + + " carries no charge/spin condition, so keyword " + "charge_spin names nothing it can serve."); + } + if (iarg + dim_chg_spin >= narg) { + error->all(FLERR, + "Illegal pair_style command: keyword charge_spin names a " + "charge state with " + + std::to_string(dim_chg_spin) + " value(s)."); + } + for (int ii = 0; ii < dim_chg_spin; ++ii) { + charge_spin.push_back( + utils::numeric(FLERR, arg[iarg + 1 + ii], false, lmp)); + } + iarg += 1 + dim_chg_spin; + } else if (key == "magnetic_force") { + if (magnetic_force_seen) { + error->all(FLERR, + "Illegal pair_style command: keyword magnetic_force may " + "appear only once."); + } + if (iarg + 1 >= narg) { + error->all(FLERR, + "Illegal pair_style command: keyword magnetic_force takes " + "one of precession or cartesian."); + } + const std::string value(arg[iarg + 1]); + if (value == "precession") { + fm_cartesian_force = 0; + } else if (value == "cartesian") { + fm_cartesian_force = 1; + } else { + error->all(FLERR, + "Illegal pair_style command: keyword magnetic_force takes " + "one of precession or cartesian, not '" + + value + "'."); + } + magnetic_force_seen = true; + iarg += 2; + } else { error->all(FLERR, "Illegal pair_style command: pair style " + style + " takes the artifact path followed by the optional " - "keyword charge_spin, not '" + - std::string(arg[iarg]) + "'."); - } - // One charge state holds for the whole run, so a second occurrence names - // no state the style could serve. Values already read mark the keyword - // as seen: a successful read always contributes at least one value. - if (!charge_spin.empty()) { - error->all(FLERR, - "Illegal pair_style command: keyword charge_spin names the " - "single charge state of the whole run, so it may appear " - "only once."); - } - if (dim_chg_spin == 0) { - error->all(FLERR, - "Illegal pair_style command: the artifact served by pair " - "style " + - style + - " carries no charge/spin condition, so keyword " - "charge_spin names nothing it can serve."); - } - if (iarg + dim_chg_spin >= narg) { - error->all(FLERR, - "Illegal pair_style command: keyword charge_spin names a " - "charge state with " + - std::to_string(dim_chg_spin) + " value(s)."); + "keywords charge_spin and magnetic_force, not '" + + key + "'."); } - for (int ii = 0; ii < dim_chg_spin; ++ii) { - charge_spin.push_back( - utils::numeric(FLERR, arg[iarg + 1 + ii], false, lmp)); - } - iarg += 1 + dim_chg_spin; + } + + // Cartesian moments use the model's spin units and energies in eV. + if (fm_cartesian_force && strcmp(update->unit_style, "metal") != 0) { + error->all(FLERR, + "pair style " + style + + " keyword magnetic_force cartesian requires units metal."); } // A charge/spin condition named on the pair_style line holds for the whole @@ -305,6 +342,8 @@ void PairDPA4Spin::settings(int narg, char** arg) { } utils::logmesg(lmp, " using charge_spin: {}\n", values); } + utils::logmesg(lmp, " magnetic_force: {}\n", + fm_cartesian_force ? "cartesian" : "precession"); } /* ---------------------------------------------------------------------- @@ -534,18 +573,19 @@ void PairDPA4Spin::compute(int eflag, int vflag) { } // === Step 3. Accumulate the model outputs === - // The model reports a force and a magnetic force for every atom the neighbor - // list covers, ghosts included; the spin atom style folds the ghost rows - // onto their owners in its reverse communication. Dividing the magnetic - // force by hbar / |m| turns the energy gradient with respect to the moment - // into the precession force LAMMPS stores. Multiplication by |m| preserves - // that relation while making a zero moment well-defined. + // Ghost contributions are folded onto their owners by the spin atom style. + // Precession scales -dE/dS by |m|/hbar; Cartesian keeps the energy-per-spin + // force, including at zero moment. const double force_scale = scale[1][1] * force_unit_cvt_factor; - const double magnetic_force_scale = force_scale / kHBar; + const double magnetic_force_scale = fm_cartesian_force + ? scale[1][1] * ener_unit_cvt_factor + : force_scale / kHBar; for (int ii = 0; ii < nall; ++ii) { + const double moment_factor = fm_cartesian_force ? 1.0 : sp[ii][3]; for (int dd = 0; dd < 3; ++dd) { f[ii][dd] += dforce[3 * ii + dd] * force_scale; - fm[ii][dd] += dforce_mag[3 * ii + dd] * sp[ii][3] * magnetic_force_scale; + fm[ii][dd] += + dforce_mag[3 * ii + dd] * moment_factor * magnetic_force_scale; } } @@ -579,5 +619,9 @@ void* PairDPA4Spin::extract(const char* str, int& dim) { dim = 2; return (void*)scale; } + if (strcmp(str, "fm_cartesian_force") == 0) { + dim = 0; + return (void*)&fm_cartesian_force; + } return nullptr; } diff --git a/source/lmp/pair_dpa4spin.h b/source/lmp/pair_dpa4spin.h index 7d0aba8977..1f389b2142 100644 --- a/source/lmp/pair_dpa4spin.h +++ b/source/lmp/pair_dpa4spin.h @@ -66,8 +66,10 @@ class PairDPA4Spin : public Pair { PairDPA4Spin(class LAMMPS*); ~PairDPA4Spin() override; - // Load the artifact named by ``pair_style dpa4spin [charge_spin - // ... ]`` and fix the charge state the keyword names. + // Load the artifact named by ``pair_style dpa4spin + // [charge_spin ... ] [magnetic_force precession|cartesian]`` and + // fix the optional charge state and magnetic-force convention the keywords + // name. ``magnetic_force`` defaults to precession. void settings(int, char**) override; // Resolve the LAMMPS atom types onto the element list of the model. void coeff(int, char**) override; @@ -98,6 +100,10 @@ class PairDPA4Spin : public Pair { std::vector type_idx_map; // LAMMPS unit system relative to the eV / Angstrom system of the model. double ener_unit_cvt_factor, dist_unit_cvt_factor, force_unit_cvt_factor; + // Convention in atom->fm: 0 = precession (|m|/hbar scaling), 1 = Cartesian + // (-dE/dS, including at zero moment). An int for the scalar Pair::extract + // API. + int fm_cartesian_force; private: CommBrickDPA4Spin* commdata_; diff --git a/source/lmp/pair_dpa4spin_kokkos.cpp b/source/lmp/pair_dpa4spin_kokkos.cpp index f479b27559..53b8109cf3 100644 --- a/source/lmp/pair_dpa4spin_kokkos.cpp +++ b/source/lmp/pair_dpa4spin_kokkos.cpp @@ -18,9 +18,7 @@ using namespace LAMMPS_NS; namespace { -// Reduced Planck constant in eV.ps. The model reports the magnetic force as -// the energy gradient with respect to the magnetic moment, while LAMMPS stores -// the precession force; the two differ by the moment magnitude over hbar. +// Reduced Planck constant in eV.ps, used by the default precession convention. constexpr double kHBar = 6.5821191e-04; } // namespace @@ -333,7 +331,11 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) { // communication stages. auto model2loc = compact_graph.d_model2loc; const double fscale = scale[1][1] * force_unit_cvt_factor; - const double fmscale = scale[1][1] * force_unit_cvt_factor / kHBar; + // Precession: energy gradient times |m| / hbar. Cartesian: energy gradient + // with only the energy-per-model-spin unit conversion (no |m|, no hbar). + const int cartesian = fm_cartesian_force; + const double fmscale = + cartesian ? scale[1][1] * ener_unit_cvt_factor : fscale / kHBar; reverse_virial = false; reverse_used_host = false; // The KOKKOS package runs 'newton off', so the integrator's force_clear only @@ -364,15 +366,14 @@ void PairDPA4SpinKokkos::compute(int eflag, int vflag) { "dpa4spin/kk:scatter_f", Kokkos::RangePolicy(0, nnode_m), KOKKOS_LAMBDA(const int m) { const int i = model2loc(m); - // The precession force carried by ``fm`` is the energy gradient with - // respect to the moment, rescaled by the moment magnitude over hbar. - const double moment = fmscale * sp(i, 3); + // Precession multiplies by the moment magnitude; Cartesian does not. + const double moment_factor = cartesian ? 1.0 : sp(i, 3); f(i, 0) += fscale * out_force(3 * m + 0); f(i, 1) += fscale * out_force(3 * m + 1); f(i, 2) += fscale * out_force(3 * m + 2); - fm(i, 0) += moment * out_force_mag(3 * m + 0); - fm(i, 1) += moment * out_force_mag(3 * m + 1); - fm(i, 2) += moment * out_force_mag(3 * m + 2); + fm(i, 0) += fmscale * moment_factor * out_force_mag(3 * m + 0); + fm(i, 1) += fmscale * moment_factor * out_force_mag(3 * m + 1); + fm(i, 2) += fmscale * moment_factor * out_force_mag(3 * m + 2); }); atomKK->modified(execution_space, F_MASK | FM_MASK); if (multi_rank) { diff --git a/source/lmp/plugin/deepmdplugin.cpp b/source/lmp/plugin/deepmdplugin.cpp index 9bc0b685d2..46922f4ff1 100644 --- a/source/lmp/plugin/deepmdplugin.cpp +++ b/source/lmp/plugin/deepmdplugin.cpp @@ -13,6 +13,7 @@ #endif #include "pair_deepmd.h" #include "pair_deepspin.h" +#include "pair_dpa4spin.h" #include "version.h" #if LAMMPS_VERSION_NUMBER >= 20220328 #include "pppm_dplr.h" @@ -22,6 +23,7 @@ using namespace LAMMPS_NS; static Pair* pairdeepmd(LAMMPS* lmp) { return new PairDeepMD(lmp); } static Pair* pairdeepspin(LAMMPS* lmp) { return new PairDeepSpin(lmp); } +static Pair* pairdpa4spin(LAMMPS* lmp) { return new PairDPA4Spin(lmp); } #ifdef LMP_KOKKOS // Runtime plugins do not consume the PairStyle declarations used by LAMMPS' @@ -109,6 +111,15 @@ extern "C" void lammpsplugin_init(void* lmp, void* handle, void* regfunc) { plugin.handle = handle; (*register_plugin)(&plugin, lmp); + plugin.version = LAMMPS_VERSION; + plugin.style = "pair"; + plugin.name = "dpa4spin"; + plugin.info = "dpa4spin pair style " STR_GIT_SUMM; + plugin.author = "Tiancheng Li"; + plugin.creator.v1 = (lammpsplugin_factory1*)&pairdpa4spin; + plugin.handle = handle; + (*register_plugin)(&plugin, lmp); + plugin.style = "compute"; plugin.name = "deeptensor/atom"; plugin.info = "compute deeptensor/atom " STR_GIT_SUMM; diff --git a/source/lmp/tests/run_mpi_dpa4spin_conventions.py b/source/lmp/tests/run_mpi_dpa4spin_conventions.py new file mode 100644 index 0000000000..7412ffc0c6 --- /dev/null +++ b/source/lmp/tests/run_mpi_dpa4spin_conventions.py @@ -0,0 +1,178 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Run native-spin conventions through a real host or CUDA Kokkos plugin. + +Each MPI process uses the same GPU, so two-rank coverage needs only one device. +Host-staged communication deliberately does not require CUDA-aware MPI. +""" + +from __future__ import ( + annotations, +) + +import argparse +import ctypes +import json +from pathlib import ( + Path, +) + +import numpy as np +from lammps import ( + lammps, +) +from mpi4py import ( + MPI, +) + + +def _pair_pointer(lmp, name, pointer_type): + extract = lmp.lib.lammps_extract_pair + previous = extract.restype + extract.restype = pointer_type + try: + pointer = extract(lmp.lmp, name.encode()) + finally: + extract.restype = previous + assert pointer, f"Pair::extract({name!r}) returned null" + return pointer + + +def _snapshot(lmp): + # gather() returns atom-ID order, including contributions reverse-communicated + # from ghost nodes. All ranks must participate in these collective calls. + def gather(name, count): + return np.array(lmp.gather(name, 1, count)).reshape(-1, count).tolist() + + return { + "ids": list(lmp.gather_atoms("id", 0, 1)), + "force": gather("f", 3), + "magnetic_force": gather("c_fmprop", 3), + "spin": gather("c_spprop", 4), + "energy": lmp.get_thermo("pe"), + } + + +def _cmdargs(kokkos): + cmdargs = ["-log", "none", "-screen", "none"] + if kokkos: + cmdargs += [ + "-k", + "on", + "g", + "1", + "-pk", + "kokkos", + "neigh", + "full", + "newton", + "off", + "comm", + "host", + "gpu/aware", + "off", + ] + return cmdargs + + +def _run(lmp, args, selection, overlay=False): + comm = MPI.COMM_WORLD + assert lmp.get_mpi_comm().size == args.ranks + pair = "dpa4spin/kk" if args.kokkos else "dpa4spin" + atom = "spin/kk" if args.kokkos else "spin" + integrator = "verlet/kk" if args.kokkos else "verlet" + assert lmp.has_style("pair", pair), f"Missing plugin pair style {pair}" + pair_args = "" if selection == "default" else f" magnetic_force {selection}" + substyle = f"{pair} {args.model.resolve()}{pair_args}" + if overlay: + style = "hybrid/overlay/kk" if args.kokkos else "hybrid/overlay" + pair_command = f"pair_style {style} {substyle} {substyle}" + coefficients = [f"pair_coeff * * {pair} {i} Ni O" for i in (1, 2)] + else: + style = pair + pair_command = f"pair_style {substyle}" + coefficients = ["pair_coeff * * Ni O"] + for command in ( + "units metal", + f"processors {comm.size} 1 1", + f"atom_style {atom}", + "atom_modify map yes", + f"read_data {args.data.resolve()}", + # No global -sf kk: select the matching integrator explicitly too. + f"run_style {integrator}", + "mass 1 58.6934", + "mass 2 15.9994", + pair_command, + *coefficients, + "neighbor 1.0 bin", + "compute fmprop all property/atom fmx fmy fmz", + "compute spprop all property/atom spx spy spz sp", + "run 0", + ): + lmp.command(command) + assert lmp.extract_pair_dimension("fm_cartesian_force") == 0 + flag = _pair_pointer(lmp, "fm_cartesian_force", ctypes.POINTER(ctypes.c_int))[0] + result = { + "convention": flag, + "pair_style": lmp.extract_global("pair_style"), + "run_style_active": lmp.eval(f"is_active(run_style,{integrator})"), + "atom_style": lmp.extract_global("atom_style"), + "kokkos_active": lmp.extract_setting("kokkos_active"), + "kokkos_ngpus": lmp.extract_setting("kokkos_ngpus"), + "newton_pair": lmp.extract_setting("newton_pair"), + "local_counts": comm.allgather(lmp.extract_setting("nlocal")), + "ghost_counts": comm.allgather(lmp.extract_setting("nghost")), + "snapshots": [_snapshot(lmp)], + } + assert result["pair_style"] == style + assert result["atom_style"] == atom + # Repeated computes expose stale ghost magnetic-force accumulation. + for _ in range(2): + lmp.command("run 0") + result["snapshots"].append(_snapshot(lmp)) + if not overlay: + scale = _pair_pointer( + lmp, "scale", ctypes.POINTER(ctypes.POINTER(ctypes.c_double)) + ) + scale[1][1] = 2.75 + lmp.command("run 0") + result["scaled"] = _snapshot(lmp) + return result + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("data", type=Path) + parser.add_argument("model", type=Path) + parser.add_argument("output", type=Path) + parser.add_argument("--kokkos", action="store_true") + parser.add_argument("--ranks", type=int, default=1) + args = parser.parse_args() + assert MPI.COMM_WORLD.size == args.ranks, ( + f"Expected {args.ranks} MPI ranks, got {MPI.COMM_WORLD.size}: " + f"{MPI.Get_library_version()}" + ) + results = {} + library = None + try: + for selection, overlay in ( + ("default", False), + ("precession", False), + ("cartesian", False), + ("precession", True), + ("cartesian", True), + ): + with lammps(cmdargs=_cmdargs(args.kokkos), comm=MPI.COMM_WORLD) as lmp: + library = lmp.lib + key = f"overlay_{selection}" if overlay else selection + results[key] = _run(lmp, args, selection, overlay) + if MPI.COMM_WORLD.rank == 0: + args.output.write_text(json.dumps(results)) + finally: + # close() destroys each instance but keeps Kokkos initialized for later + # instances. Finalize once, even on error, before mpi4py finalizes MPI. + if args.kokkos and library is not None: + library.lammps_kokkos_finalize() + + +if __name__ == "__main__": + main() diff --git a/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py b/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py new file mode 100644 index 0000000000..7141e3a044 --- /dev/null +++ b/source/lmp/tests/test_dpa4spin_kokkos_magnetic_force.py @@ -0,0 +1,229 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Actual CUDA /kk convention and MPI regression, separate from host FD tests. + +Run source/install/test_lammps_kokkos.sh after test_cc_local.sh on a CUDA host. +That job supplies a compressed canonical fixture and a matching Kokkos LAMMPS +library/plugin, and sets DPA4SPIN_KOKKOS_REQUIRED=1: no missing prerequisite may +silently skip it. This checks adapter equivalence, not GPU energy differences. +""" + +from __future__ import ( + annotations, +) + +import importlib +import json +import os +import shutil +import signal +import subprocess +import sys +import zipfile +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from write_lmp_data import ( + write_lmp_data_spin, +) + +HBAR_METAL = 6.5821191e-04 +_MODEL = ( + Path(__file__).resolve().parents[2] + / "tests/infer/deeppot_dpa4c_spin_lmp_kokkos.pt2" +) +_BOX = np.array([0.0, 14.0, 0.0, 14.0, 0.0, 14.0, 0.0, 0.0, 0.0]) +# Both ranks own atoms and interact across the x=7 domain boundary. The box +# exceeds twice the cutoff, as required by the single-rank folded /kk graph. +_COORD = np.array( + [ + [5.5, 4.0, 4.0], + [7.7, 4.4, 4.1], + [5.8, 4.8, 4.0], + [4.9, 4.2, 4.6], + [8.1, 5.0, 4.3], + [7.9, 3.7, 4.7], + ] +) +_SPIN = np.array( + [ + [0.0, 0.0, 0.0], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] +) +_TYPE = np.array([1, 1, 1, 2, 2, 2], dtype=np.int32) + + +def _unavailable(reason): + if os.environ.get("DPA4SPIN_KOKKOS_REQUIRED") == "1": + pytest.fail(reason) + pytest.skip(reason) + + +@pytest.fixture(scope="module") +def kokkos_model(): + # The ordinary LAMMPS suite also runs on CUDA, but uses a non-Kokkos wheel. + # Only the dedicated matching library/plugin job opts into this regression. + if os.environ.get("DPA4SPIN_KOKKOS_REQUIRED") != "1": + pytest.skip("Enable the dedicated CUDA Kokkos plugin regression explicitly") + try: + torch = importlib.import_module("torch") + except ImportError as exc: + _unavailable(f"CUDA PyTorch unavailable: {exc}") + raise + if not torch.cuda.is_available(): + _unavailable("CUDA device required for the actual dpa4spin/kk adapter") + for module in ("lammps", "mpi4py"): + if importlib.util.find_spec(module) is None: + _unavailable(f"Required Python package missing: {module}") + if shutil.which(os.environ.get("DPA4SPIN_KOKKOS_MPIEXEC", "mpirun")) is None: + _unavailable("mpirun required for the two-rank /kk regression") + model = Path(os.environ.get("DPA4SPIN_KOKKOS_MODEL", _MODEL)) + if not model.is_file(): + _unavailable(f"Generate the compressed canonical fixture first: {model}") + with zipfile.ZipFile(model) as archive: + meta = json.loads(archive.read("model/extra/metadata.json")) + assert meta["lower_input_kind"] == "dpa4c_canonical" + assert meta["canonical_index_dtype"] == "uint32" + assert meta["graph_edge_dtype"] == "float32" + assert meta["spin_scheme"] == "native" + assert meta["type_map"] == ["Ni", "O"] + assert meta["use_spin"] == [True, False] + for key in ("dim_fparam", "dim_aparam", "dim_chg_spin"): + assert meta[key] == 0 + assert not meta.get("has_comm_artifact", False) + return model + + +def _run(tmp_path, data, model, *, ranks, kokkos): + output = tmp_path / f"{'kokkos' if kokkos else 'host'}-{ranks}.json" + runner = Path(__file__).with_name("run_mpi_dpa4spin_conventions.py") + command = [ + os.environ.get("DPA4SPIN_KOKKOS_MPIEXEC", "mpirun"), + "-n", + str(ranks), + sys.executable, + str(runner), + str(data), + str(model), + str(output), + "--ranks", + str(ranks), + ] + if kokkos: + command.append("--kokkos") + # As in dpa4_spin_harness, bound the whole MPI process group: killing only + # mpirun can leave workers holding the pipes open after a timeout. + with subprocess.Popen( + command, + stdout=subprocess.PIPE, + stderr=subprocess.PIPE, + text=True, + start_new_session=True, + ) as process: + try: + stdout, stderr = process.communicate(timeout=600) + except subprocess.TimeoutExpired: + try: + os.killpg(process.pid, signal.SIGKILL) + except ProcessLookupError: + # The process group already exited between timeout and cleanup. + pass + stdout, stderr = process.communicate() + pytest.fail( + f"MPI runner timed out (process group killed):\n{stdout}{stderr}" + ) + assert process.returncode == 0, stdout + stderr + return json.loads(output.read_text()) + + +@pytest.fixture(scope="module", params=[0.0, 1e-8, 0.2], ids=["zero", "tiny", "finite"]) +def samples(kokkos_model, tmp_path_factory, request): + directory = tmp_path_factory.mktemp(f"dpa4spin-kk-{request.param}") + data = directory / "nio.lmp" + spin = _SPIN.copy() + spin[0] = request.param * np.array([2.0, -1.0, 2.0]) / 3.0 + write_lmp_data_spin(_BOX, _COORD, spin, _TYPE, data) + return ( + _run(directory, data, kokkos_model, ranks=1, kokkos=False), + _run(directory, data, kokkos_model, ranks=1, kokkos=True), + _run(directory, data, kokkos_model, ranks=2, kokkos=True), + ) + + +def _assert_snapshot(actual, expected, scale=1.0): + assert actual["ids"] == expected["ids"] == list(range(1, 7)) + np.testing.assert_allclose(actual["spin"], expected["spin"], atol=0, rtol=0) + # The canonical kernels use float32; edge/reduction ordering differs between + # host, folded single-rank and extended multi-rank graphs. + for name, atol in (("force", 2e-7), ("magnetic_force", 5e-9)): + np.testing.assert_allclose( + actual[name], scale * np.array(expected[name]), rtol=5e-4, atol=atol + ) + assert actual["energy"] == pytest.approx( + scale * expected["energy"], rel=5e-6, abs=2e-7 + ) + + +def test_kokkos_conventions_and_mpi(samples): + host, single, parallel = samples + reference = host["cartesian"]["snapshots"][0] + assert np.linalg.norm(np.array(reference["magnetic_force"])[0]) > 1e-7 + for ranks, result in ((1, single), (2, parallel)): + cartesian = result["cartesian"]["snapshots"][0] + magnetic = np.array(cartesian["magnetic_force"]) + # In particular, forcing the GPU kernel to precession at S=0 fails here. + assert np.linalg.norm(magnetic[0]) > 1e-7 + np.testing.assert_array_equal(magnetic[3:], 0.0) + _assert_snapshot(cartesian, reference) + for selection in ("default", "precession", "cartesian"): + sample = result[selection] + assert sample["convention"] == int(selection == "cartesian") + assert sample["pair_style"] == "dpa4spin/kk" + assert sample["atom_style"] == "spin/kk" + assert sample["run_style_active"] == 1 + assert sample["kokkos_active"] == sample["kokkos_ngpus"] == 1 + assert sample["newton_pair"] == 0 + assert len(sample["local_counts"]) == ranks + assert sum(sample["local_counts"]) == 6 + assert min(sample["local_counts"]) > 0 + if ranks == 2: + assert min(sample["ghost_counts"]) > 0 + actual = sample["snapshots"][0] + _assert_snapshot(actual, host[selection]["snapshots"][0]) + _assert_snapshot(actual, single[selection]["snapshots"][0]) + if selection != "cartesian": + moments = np.array(actual["spin"])[:, 3] + finite = moments > 0 + precession = np.array(actual["magnetic_force"]) + # Normalize BEFORE comparing so the tiny-moment case cannot + # pass with zero force hidden by a fixed absolute tolerance. + np.testing.assert_allclose( + precession[finite] * (HBAR_METAL / moments[finite, None]), + magnetic[finite], + rtol=5e-4, + atol=5e-9, + ) + np.testing.assert_array_equal(precession[~finite], 0.0) + for repeated in sample["snapshots"][1:]: + _assert_snapshot(repeated, actual) + _assert_snapshot(sample["scaled"], actual, scale=2.75) + + +def test_kokkos_overlay_accumulates(samples): + host, single, parallel = samples + for result in (single, parallel): + for selection in ("precession", "cartesian"): + overlay = result[f"overlay_{selection}"] + assert overlay["pair_style"] == "hybrid/overlay/kk" + assert overlay["run_style_active"] == 1 + assert overlay["convention"] == int(selection == "cartesian") + for snapshot in overlay["snapshots"]: + _assert_snapshot(snapshot, host[selection]["snapshots"][0], scale=2.0) + _assert_snapshot(snapshot, result[selection]["snapshots"][0], scale=2.0) diff --git a/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py b/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py new file mode 100644 index 0000000000..7855803905 --- /dev/null +++ b/source/lmp/tests/test_dpa4spin_magnetic_force_convention.py @@ -0,0 +1,245 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Exercise magnetic-force conventions through the registered LAMMPS adapter. + +Generate the CPU graph fixture with ``source/tests/infer/gen_dpa4c_spin_lmp.py``. +PyTorch-enabled LAMMPS CI sets DPA4SPIN_LMP_REQUIRED=1 so missing fixtures, +Python bindings, or plugin registration fail instead of hiding regressions. +These tests exercise the host adapter; the GPU-only Kokkos adapter requires a +separate compressed canonical archive and is not covered by this CPU fixture. +""" + +from __future__ import ( + annotations, +) + +import ctypes +import importlib +import os +from contextlib import ( + contextmanager, +) +from pathlib import ( + Path, +) + +import numpy as np +import pytest +from write_lmp_data import ( + write_lmp_data_spin, +) + +HBAR_METAL = 6.5821191e-04 +_MODEL = ( + Path(__file__).resolve().parents[2] + / "tests" + / "infer" + / "deeppot_dpa4c_spin_lmp.pt2" +) +_BOX = np.array([0.0, 6.0, 0.0, 6.0, 0.0, 6.0, 0.0, 0.0, 0.0]) +_COORD = np.array( + [ + [1.0, 1.0, 1.0], + [3.2, 1.4, 1.1], + [1.3, 1.8, 1.0], + [0.4, 1.2, 1.6], + [3.6, 2.0, 1.3], + [3.4, 0.7, 1.7], + ] +) +_SPIN = np.array( + [ + [0.0, 0.0, 0.0], + [-0.07, 0.09, 0.03], + [0.02, -0.06, 0.08], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + [0.0, 0.0, 0.0], + ] +) +_TYPE = np.array([1, 1, 1, 2, 2, 2], dtype=np.int32) + + +def _unavailable(reason: str) -> None: + if os.environ.get("DPA4SPIN_LMP_REQUIRED") == "1": + pytest.fail(reason) + pytest.skip(reason) + + +@pytest.fixture(scope="module") +def adapter(): + if os.environ.get("ENABLE_PYTORCH", "1") != "1": + pytest.skip("PyTorch backend is disabled") + try: + module = importlib.import_module("lammps") + except ImportError as exc: + _unavailable(f"LAMMPS Python bindings unavailable: {exc}") + raise # Keep the import failure path explicit for static analysis. + model = Path(os.environ.get("DPA4SPIN_LMP_MODEL", _MODEL)) + if not model.is_file(): + _unavailable(f"Generate the native-spin LAMMPS fixture first: {model}") + with module.lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp: + if not lmp.has_style("pair", "dpa4spin"): + _unavailable("LAMMPS plugin did not register pair style dpa4spin") + return module.lammps, model + + +@pytest.fixture +def data_file(tmp_path: Path) -> Path: + path = tmp_path / "nio.lmp" + write_lmp_data_spin(_BOX, _COORD, _SPIN, _TYPE, path) + return path + + +@contextmanager +def _system(adapter, data_file: Path, pair_args: str = "", units: str = "metal"): + lammps_cls, model = adapter + with lammps_cls(cmdargs=["-log", "none", "-screen", "none"]) as lmp: + for command in ( + f"units {units}", + "atom_style spin", + "atom_modify map yes", + f"read_data {data_file.resolve()}", + "mass 1 58.6934", + "mass 2 15.9994", + f"pair_style dpa4spin {model.resolve()} {pair_args}", + "pair_coeff * * Ni O", + "neighbor 2.0 bin", + "compute fmprop all property/atom fmx fmy fmz", + "compute spprop all property/atom spx spy spz sp", + "run 0", + ): + lmp.command(command) + yield lmp + + +def _pair_pointer(lmp, name: str, pointer_type): + # extract_pair() assumes double values, whereas the convention flag is int. + extract = lmp.lib.lammps_extract_pair + previous = extract.restype + extract.restype = pointer_type + try: + pointer = extract(lmp.lmp, name.encode()) + finally: + extract.restype = previous + assert pointer, f"Pair::extract({name!r}) returned null" + return pointer + + +def _convention(lmp) -> int: + assert lmp.extract_pair_dimension("fm_cartesian_force") == 0 + return _pair_pointer(lmp, "fm_cartesian_force", ctypes.POINTER(ctypes.c_int))[0] + + +def _gather(lmp, name: str, columns: int) -> np.ndarray: + return np.array(lmp.gather(name, 1, columns)).reshape(-1, columns) + + +@pytest.mark.parametrize("magnitude", [0.0, 1e-8, 0.2], ids=["zero", "tiny", "finite"]) +def test_cartesian_matches_energy_gradient(adapter, data_file: Path, magnitude: float): + """Check every Cartesian component, including the derivative at S=0.""" + spin = _SPIN.copy() + spin[0] = magnitude * np.array([2.0, -1.0, 2.0]) / 3.0 + write_lmp_data_spin(_BOX, _COORD, spin, _TYPE, data_file) + with _system(adapter, data_file, "magnetic_force cartesian") as lmp: + assert _convention(lmp) == 1 + analytic = _gather(lmp, "c_fmprop", 3)[0] + assert np.all(np.abs(analytic) > 1e-9), ( + "Every tested spin component must be live", + analytic, + ) + initial = _gather(lmp, "c_spprop", 4)[0] + moment = initial[:3] * initial[3] + eps = 1e-5 + finite_difference = np.empty(3) + for component in range(3): + energies = [] + for sign in (1, -1): + displaced = moment.copy() + displaced[component] += sign * eps + norm = np.linalg.norm(displaced) + direction = displaced / norm + # Both displaced norms are positive even when the centre is S=0. + # LAMMPS's set command rejects an exactly zero magnitude. + lmp.command( + f"set atom 1 spin/atom {norm:.17g} " + + " ".join(f"{value:.17g}" for value in direction) + ) + lmp.command("run 0") + energies.append(lmp.get_thermo("pe")) + finite_difference[component] = -(energies[0] - energies[1]) / (2 * eps) + # Check components separately: a norm-only comparison could hide a + # sign error in the small transverse component of the zero-spin atom. + np.testing.assert_allclose(finite_difference, analytic, rtol=5e-3, atol=1e-11) + + +def test_precession_default_and_explicit_match_cartesian(adapter, data_file: Path): + with _system(adapter, data_file, "magnetic_force cartesian") as lmp: + cartesian = _gather(lmp, "c_fmprop", 3) + moments = _gather(lmp, "c_spprop", 4)[:, 3] + force = _gather(lmp, "f", 3) + energy = lmp.get_thermo("pe") + assert np.linalg.norm(cartesian[0]) > 1e-8 + np.testing.assert_allclose(cartesian[3:], 0.0, atol=1e-14) + for selection in ("", "magnetic_force precession"): + with _system(adapter, data_file, selection) as lmp: + assert _convention(lmp) == 0 + actual = _gather(lmp, "c_fmprop", 3) + np.testing.assert_allclose( + actual, + cartesian * (moments / HBAR_METAL)[:, None], + rtol=1e-10, + atol=1e-14, + ) + np.testing.assert_array_equal(actual[0], 0.0) + np.testing.assert_allclose( + _gather(lmp, "f", 3), force, rtol=1e-12, atol=1e-14 + ) + assert lmp.get_thermo("pe") == pytest.approx(energy, rel=1e-12) + + +@pytest.mark.parametrize("selection", ["precession", "cartesian"]) +def test_pair_scale_and_repeated_compute(adapter, data_file: Path, selection: str): + """Exercise the actual extract/scale and force accumulation paths.""" + with _system(adapter, data_file, f"magnetic_force {selection}") as lmp: + magnetic_force = _gather(lmp, "c_fmprop", 3) + force = _gather(lmp, "f", 3) + energy = lmp.get_thermo("pe") + scale = _pair_pointer( + lmp, "scale", ctypes.POINTER(ctypes.POINTER(ctypes.c_double)) + ) + scale[1][1] = 2.75 + for _ in range(2): + lmp.command("run 0") + np.testing.assert_allclose( + _gather(lmp, "c_fmprop", 3), + 2.75 * magnetic_force, + rtol=1e-10, + atol=1e-14, + ) + np.testing.assert_allclose( + _gather(lmp, "f", 3), 2.75 * force, rtol=1e-10, atol=1e-14 + ) + assert lmp.get_thermo("pe") == pytest.approx(2.75 * energy, rel=1e-12) + + +@pytest.mark.parametrize( + "pair_args, units, message", + [ + ("magnetic_force", "metal", "takes one of precession or cartesian"), + ("magnetic_force unknown", "metal", "takes one of precession or cartesian"), + ( + "magnetic_force cartesian magnetic_force precession", + "metal", + "magnetic_force may appear only once", + ), + ("magnetic_force cartesian", "real", "requires units metal"), + ], +) +def test_invalid_magnetic_force_settings( + adapter, data_file: Path, pair_args, units, message +): + with ( + pytest.raises(Exception, match=message), + _system(adapter, data_file, pair_args, units), + ): + pytest.fail("Invalid magnetic_force setting was accepted") diff --git a/source/lmp/tests/test_lammps_plugin.py b/source/lmp/tests/test_lammps_plugin.py new file mode 100644 index 0000000000..d7073fe1d4 --- /dev/null +++ b/source/lmp/tests/test_lammps_plugin.py @@ -0,0 +1,22 @@ +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Check style registration without requiring a model or inference backend. + +CI loads the runtime plugin through LAMMPS_PLUGIN_PATH. A standalone plugin +can instead be loaded explicitly with DEEPMD_TEST_PLUGIN. +""" + +import os + +from lammps import ( + lammps, +) + + +def test_dpa4spin_host_style_registered() -> None: + """The unsuffixed host style must also be available without Kokkos.""" + with lammps(cmdargs=["-log", "none", "-screen", "none"]) as lmp: + if plugin := os.environ.get("DEEPMD_TEST_PLUGIN"): + lmp.command(f'plugin load "{plugin}"') + assert lmp.has_style("pair", "dpa4spin"), ( + "LAMMPS plugin did not register pair style dpa4spin" + ) diff --git a/source/tests/infer/gen_dpa4c_spin_lmp.py b/source/tests/infer/gen_dpa4c_spin_lmp.py new file mode 100644 index 0000000000..bdd13c0599 --- /dev/null +++ b/source/tests/infer/gen_dpa4c_spin_lmp.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate a minimal native-spin DPA4C graph ``.pt2`` for LAMMPS ``dpa4spin``. + +Unlike ``gen_dpa4c_spin.py``, this archive carries no fparam, aparam, or +charge/spin condition so ``pair_style dpa4spin`` can load it without extra +keywords. Used by ``source/lmp/tests/test_dpa4spin_magnetic_force_convention.py``. +""" + +from __future__ import ( + annotations, +) + +import copy +import json +import os +import sys +import zipfile + +import numpy as np + +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..")) + +from dpa4_fixtures import ( + jitter_zero_arrays, +) +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, +) + +CFG = { + "type_map": ["Ni", "O"], + "descriptor": { + "type": "dpa4c", + "rcut": 4.0, + "channels": 8, + "lmax": 2, + "n_radial": 4, + "precision": "float64", + "seed": 7, + }, + "fitting_net": {"neuron": [8, 8], "precision": "float64", "seed": 7}, + "spin": {"use_spin": [True, False], "scheme": "native"}, +} + +_JITTER_SEED = 20260817 +OUT = os.path.join(os.path.dirname(__file__), "deeppot_dpa4c_spin_lmp.pt2") + + +def main() -> None: + from torch._inductor import config as inductor_config + + from deepmd.dpmodel.model.model import ( + get_model, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + ensure_inductor_compiler() + load_custom_ops() + model = get_model(copy.deepcopy(CFG)) + md = jitter_zero_arrays(model.serialize(), np.random.default_rng(_JITTER_SEED)) + data = { + "model": md, + "model_def_script": CFG, + "backend": "dpmodel", + "software": "deepmd-kit", + "version": "3.0.0", + } + print(f"Exporting {OUT}") # noqa: T201 + # PyTorch 2.11's single-threaded CPU SIMD gather codegen can use an + # index temporary before declaring it. This tiny fixture checks adapter + # magnetic-force conventions, so scalar codegen avoids that compiler bug + # without changing production exports or CUDA compilation. + with inductor_config.patch({"cpp.simdlen": 1}): + deserialize_to_file(OUT, data, do_atomic_virial=True, lower_kind="graph") + with zipfile.ZipFile(OUT) as zf: + meta = json.loads(zf.read("model/extra/metadata.json").decode()) + print( # noqa: T201 + json.dumps( + { + k: meta.get(k) + for k in ( + "spin_scheme", + "has_comm_artifact", + "dim_fparam", + "dim_aparam", + "dim_chg_spin", + "lower_input_kind", + ) + }, + indent=2, + ) + ) + + +if __name__ == "__main__": + main() diff --git a/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py b/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py new file mode 100644 index 0000000000..76c0a6c0d0 --- /dev/null +++ b/source/tests/infer/gen_dpa4c_spin_lmp_kokkos.py @@ -0,0 +1,131 @@ +#!/usr/bin/env python3 +# SPDX-License-Identifier: LGPL-3.0-or-later +"""Generate the compressed CUDA fixture for LAMMPS ``dpa4spin/kk``. + +The Kokkos adapter requires the compact canonical graph ABI, which is only +available for a compressed float32 DPA4C model with a fused fitting network. +Keep this archive separate from the float64 CPU graph fixture. Both carry +native spin and no frame, atomic, or charge/spin conditioning parameters. + +Run with a CUDA-enabled DeePMD build and a visible GPU. A CPU-only invocation +fails rather than silently exporting an archive the Kokkos adapter cannot use. +""" + +from __future__ import ( + annotations, +) + +import copy +import json +import os +import zipfile + +import gen_dpa4c_spin_lmp as host_fixture +import numpy as np +from gen_common import ( + ensure_inductor_compiler, + load_custom_ops, +) + +CFG = copy.deepcopy(host_fixture.CFG) +CFG["descriptor"]["precision"] = "float32" +CFG["fitting_net"]["precision"] = "float32" +# The fused fitting kernel does not support per-layer residual timesteps. +CFG["fitting_net"]["resnet_dt"] = False +OUT = os.path.join(os.path.dirname(__file__), "deeppot_dpa4c_spin_lmp_kokkos.pt2") + + +def _build_model_data() -> dict: + """Build reproducible weights and compression tables on the CPU.""" + from deepmd.dpmodel.model.model import ( + get_model, + ) + from deepmd.pt_expt.kernels.dpa4c.graph_compress import ( + mega_eligible, + ) + from deepmd.pt_expt.kernels.graph_fitting import ( + fitting_eligible, + ) + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + + model_dict = host_fixture.jitter_zero_arrays( + get_model(copy.deepcopy(CFG)).serialize(), + np.random.default_rng(host_fixture._JITTER_SEED), + ) + model = BaseModel.deserialize(model_dict).to("cpu").eval() + descriptor = model.get_descriptor() + assert mega_eligible(descriptor), "the fixture needs a compiled DPA4C profile" + assert fitting_eligible(model.get_fitting_net()), "the fixture needs fused fitting" + # Jitter before compression so the tabulated snapshot retains the live + # native-spin branch, including its derivative at exactly zero spin. + descriptor.enable_compression(min_nbor_dist=0.5) + assert descriptor.compress + return { + "model": model.serialize(), + "model_def_script": copy.deepcopy(CFG), + "backend": "pt_expt", + "software": "deepmd-kit", + "version": "3.0.0", + } + + +def _check_metadata(path: str) -> None: + """Reject an archive that would exercise a different deployment route.""" + with zipfile.ZipFile(path) as zf: + metadata = json.loads(zf.read("model/extra/metadata.json")) + expected = { + "type_map": ["Ni", "O"], + "spin_scheme": "native", + "use_spin": [True, False], + "has_comm_artifact": False, + "dim_fparam": 0, + "dim_aparam": 0, + "dim_chg_spin": 0, + "lower_input_kind": "dpa4c_canonical", + "canonical_index_dtype": "uint32", + "graph_edge_dtype": "float32", + } + for key, value in expected.items(): + assert metadata.get(key) == value, (key, metadata.get(key), value) + for key in ("energy", "force", "force_mag", "virial", "atom_virial"): + assert key in metadata["output_keys"], key + print(json.dumps({key: metadata[key] for key in expected}, indent=2)) # noqa: T201 + + +def main() -> None: + from deepmd.pt_expt.kernels.dpa4c.canonical import ( + canonical_model_eligible, + op_available, + ) + from deepmd.pt_expt.model.model import ( + BaseModel, + ) + from deepmd.pt_expt.utils.env import ( + DEVICE, + ) + from deepmd.pt_expt.utils.serialization import ( + deserialize_to_file, + ) + + if DEVICE.type != "cuda": + raise RuntimeError( + "The dpa4spin/kk fixture requires a CUDA-enabled DeePMD build " + "and a visible GPU (without DEVICE=cpu)." + ) + ensure_inductor_compiler() + load_custom_ops() + if not op_available(): + raise RuntimeError( + "The compressed canonical DPA4C CUDA operators are required." + ) + data = _build_model_data() + assert canonical_model_eligible(BaseModel.deserialize(data["model"])) + print(f"Exporting {OUT}") # noqa: T201 + deserialize_to_file(OUT, data, do_atomic_virial=True, lower_kind="dpa4c_canonical") + _check_metadata(OUT) + + +if __name__ == "__main__": + main()