Merge branch 'lammps'

Author: shyuep

.codespellrc

@@ -1,2 +1,2 @@
 [codespell]
-ignore-words-list = COO,Mater,ket
+ignore-words-list = COO,Mater,ket,nd,te

.github/workflows/lammps-build.yml

@@ -0,0 +1,212 @@
+# MatGL → LAMMPS pair_style
+
+`pair_matgl` is a LAMMPS pair style that loads a TorchScript-compiled
+**MatGL TensorNet** PES (PyG backend, no-Warp, extensive head) and uses
+LibTorch to evaluate energies, forces, and the virial tensor on every
+timestep.
+
+This directory ships:
+
+- `src/ML-MATGL/pair_matgl.{cpp,h}` — the CPU/serial pair style.
+- `src/KOKKOS/pair_matgl_kokkos.{cpp,h}` — the Kokkos GPU/host variant
+  (`pair_style matgl/kk`).
+- `cmake/ML-MATGL.cmake` and `cmake/ML-MATGL-KOKKOS.cmake` — drop-in
+  CMake snippets.
+- `tests/in.matgl_si` — sample input deck for a single-point parity check.
+
+The Python side (one repo up) ships `mgl create-lammps-model`, which
+produces the `.pt` artifact these pair styles consume.
+
+> **Status — Phases 2 + 3 of the MatGL LAMMPS-Kokkos plugin.** CPU CI
+> exists (`.github/workflows/lammps-build.yml`); GPU runs require a
+> CUDA-capable runner (not in CI yet).
+
+## Building
+
+### 1. Export a LAMMPS-loadable model
+
+```bash
+# From your matgl checkout:
+uv run mgl create-lammps-model \
+    -m materialyze/TensorNet-MatPES-r2SCAN \
+    -o tensornet_matpes_r2scan.pt \
+    --dtype float32
+```
+
+The CLI prints `r_max`, `n_species`, the dtype, and the species list — all
+of which you'll need for `pair_coeff`.
+
+### 2. Build LAMMPS with the package
+
+Drop the package into a stock LAMMPS source tree and configure:
+
+```bash
+# 1) Copy or symlink the source files.
+ln -s /path/to/matgl/lammps/src/ML-MATGL <lammps>/src/ML-MATGL
+
+# 2) Tell LAMMPS' CMake about the package.
+echo 'include(/path/to/matgl/lammps/cmake/ML-MATGL.cmake)' \
+    >> <lammps>/cmake/CMakeLists.txt
+
+# 3) Configure + build. Match libtorch's CXX11 ABI to LAMMPS'.
+cmake -B build -S <lammps>/cmake \
+    -D PKG_ML-MATGL=ON \
+    -D CMAKE_PREFIX_PATH=/path/to/libtorch \
+    -D CMAKE_BUILD_TYPE=Release \
+    -D BUILD_MPI=ON
+cmake --build build -j 8
+```
+
+### 2b. Build the Kokkos GPU variant
+
+To get the `matgl/kk` pair style, also enable Kokkos and append the
+matching snippet to LAMMPS' CMake. CUDA example for an Ampere card
+(A100/A30):
+
+```bash
+echo 'include(/path/to/matgl/lammps/cmake/ML-MATGL-KOKKOS.cmake)' \
+    >> <lammps>/cmake/CMakeLists.txt
+
+cmake -B build -S <lammps>/cmake \
+    -D PKG_ML-MATGL=ON \
+    -D PKG_KOKKOS=ON \
+    -D Kokkos_ENABLE_CUDA=ON \
+    -D Kokkos_ARCH_AMPERE80=ON \
+    -D CMAKE_PREFIX_PATH=/path/to/libtorch \
+    -D CMAKE_CXX_COMPILER=<lammps>/lib/kokkos/bin/nvcc_wrapper \
+    -D CMAKE_BUILD_TYPE=Release
+cmake --build build -j 8
+```
+
+Run with:
+
+```bash
+mpirun -n 1 build/lmp -k on g 1 -sf kk -in in.matgl_si
+```
+
+`-sf kk` makes LAMMPS prefer Kokkos pair styles, so `pair_style matgl`
+in your input deck dispatches to `matgl/kk` automatically. If you'd
+rather force it explicitly, write `pair_style matgl/kk` instead.
+
+**Single-GPU only.** Multi-rank Kokkos with libtorch is unreliable
+(MACE issues #1294 and #322); the package emits a CMake message making
+this explicit.
+
+Tested with:
+
+- LibTorch 2.2.x – 2.5.x (CXX11 ABI, CPU build).
+- LAMMPS develop branch (Aug 2024 or newer for the `add_request` /
+  `REQ_GHOST` neighbor-list API).
+- C++17, MPI optional.
+
+## LAMMPS input syntax
+
+```lammps
+units           metal
+atom_style      atomic
+atom_modify     map yes        # required: pair_matgl needs the atom map
+newton          on             # required: ghost contributions
+
+pair_style      matgl
+pair_coeff      * * tensornet_matpes_r2scan.pt Si C O
+```
+
+`pair_coeff` arguments after the `.pt` path are **species symbols** in
+LAMMPS atom-type order: type 1 = first symbol, type 2 = second, …
+
+The cutoff (`r_max`) is read from the model — you don't pass it.
+
+### Optional pair_style flags
+
+```lammps
+pair_style matgl no_domain_decomposition
+```
+
+Reserved for future single-rank optimisations (mirrors the MACE flag).
+Currently a no-op.
+
+## Limitations
+
+- **No per-atom energies / virials.** `eflag_atom`, `vflag_atom`, and
+  `compute … pe/atom` will error. The model returns a single
+  `total_energy_local` scalar plus a 3×3 virial tensor; per-atom
+  decompositions would require a different export.
+- **`atom_style atomic` only** for now. Charged systems aren't supported
+  (the model has no charge head).
+- **TorchScript artifacts are dtype-specific.** Re-run
+  `mgl create-lammps-model --dtype float64` to get a double-precision
+  model; mixing dtypes between LAMMPS and the model will error at load
+  time.
+- **Multi-rank**: works for CPU MPI, but each rank loads the model
+  independently (memory adds up). The `data_mean` buffer baked into the
+  TorchScript is added once per rank — keep `data_mean = 0` (the default
+  for trained MatGL PES models). Non-zero `data_mean` will over-count
+  proportionally to the number of ranks.
+- **No restart support.** The model lives on disk; `restart` files don't
+  capture the path. Re-issue `pair_style` / `pair_coeff` after a restart.
+- **TensorNet only.** M3GNet, CHGNet, MEGNet, SO3Net, QET are DGL-only
+  in the matgl repo and would need PyG ports first.
+
+## Continuous integration
+
+`.github/workflows/lammps-build.yml` builds the **CPU** pair style on
+every push that touches the `lammps/` tree, the Python wrapper, or the
+workflow itself. The job runs inside the `lammps/lammps-build:ubuntu_latest`
+public Docker image, downloads a CXX11-ABI libtorch, clones LAMMPS at a
+pinned tag, builds with `PKG_ML-MATGL=ON`, exports a tiny in-tree model
+through `LAMMPSMatGLModel`, runs the `in.matgl_si` deck, and diffs the
+LAMMPS energy against the Python reference.
+
+The Kokkos variant is **not** exercised in CI today — GitHub-hosted
+runners have no GPU. Hardware-accelerated CI is on the Phase-3 follow-up
+list and likely lives on a self-hosted CUDA runner.
+
+## Verifying a build
+
+```bash
+cd lammps/tests
+<lammps>/build/lmp -in in.matgl_si
+```
+
+The test deck prints energy, forces, and stress on a small Si supercell.
+Compare against the Python reference:
+
+```bash
+uv run python tests/python_reference.py    # in this directory
+```
+
+Energies should match within `1e-5 eV`, forces within `1e-4 eV/Å`, and
+stresses (when nonzero) within `1e-3 GPa`.
+
+## Implementation notes
+
+- The pair style requests a **full neighbor list with ghost atoms**
+  (`REQ_FULL | REQ_GHOST`). The model expects edge indices that span both
+  owned and ghost atoms.
+- Every edge is folded back onto the *local* row of the atom it represents
+  (via `atom->map(atom->tag[j])`) rather than pointing at the ghost row
+  directly. TensorNet's message-passing layers need one consistent row
+  per physical atom — a ghost row never propagates outgoing messages back
+  to the atom it duplicates. The periodic image is recovered explicitly as
+  an integer `unit_shifts` (the ghost/local position difference,
+  transformed through the box's inverse deformation matrix and rounded to
+  the nearest integer), rather than relying on LAMMPS' already-imaged
+  ghost positions with `unit_shifts = 0`. Single-rank only: ghost atoms
+  can be owned by a different MPI rank, so there is no local row to fold
+  onto in a multi-rank run.
+- Forces are accumulated for **all** atoms (owned + ghost). LAMMPS' usual
+  `comm->reverse_comm` step then sums ghost contributions back to the
+  rank that owns each atom. This requires `newton on`.
+- Virials are written into the global `virial[6]` array directly as
+  `virial -= va` (the model returns `virials = dE/dstrain = -W`, while
+  LAMMPS' convention is `W = sum_i r_i ⊗ f_i`). We set
+  `no_virial_fdotr_compute = 1` in the constructor so LAMMPS doesn't
+  recompute the virial from forces.
+
+## Reference
+
+Plan and design notes:
+[`develop-a-kokkos-plugin-eventual-hare.md`](https://github.com/materialyzeai/matgl/tree/lammps).
+
+The Python wrapper is documented inline at
+`src/matgl/ext/_lammps.py` in the matgl repo.

lammps/README.md

@@ -0,0 +1,43 @@
+# ML-MATGL Kokkos variant — drop-in CMake snippet.
+#
+# Layered on top of ML-MATGL.cmake: include() this *after* the base snippet,
+# OR set PKG_ML-MATGL=ON and PKG_KOKKOS=ON together.
+#
+# Usage (from a stock LAMMPS source tree):
+#   cmake -B build \
+#       -D PKG_ML-MATGL=ON -D PKG_KOKKOS=ON \
+#       -D Kokkos_ENABLE_CUDA=ON \
+#       -D Kokkos_ARCH_AMPERE80=ON \
+#       -D CMAKE_PREFIX_PATH=/path/to/libtorch \
+#       -D CMAKE_CXX_COMPILER=$LAMMPS/lib/kokkos/bin/nvcc_wrapper \
+#       <other flags>
+#
+# The `pair_matgl/kk` style is registered via the standard LAMMPS Kokkos
+# pair-style macro so users invoke it with `pair_style matgl/kk` or by
+# launching LAMMPS with `-sf kk -k on g 1`.
+
+if(NOT PKG_ML-MATGL OR NOT PKG_KOKKOS)
+    return()
+endif()
+
+if(NOT DEFINED ML_MATGL_KOKKOS_DIR)
+    get_filename_component(ML_MATGL_KOKKOS_DIR
+        "${CMAKE_CURRENT_LIST_DIR}/../src/KOKKOS" ABSOLUTE)
+endif()
+
+if(NOT EXISTS "${ML_MATGL_KOKKOS_DIR}/pair_matgl_kokkos.cpp")
+    message(FATAL_ERROR
+        "ML-MATGL-KOKKOS source not found at ${ML_MATGL_KOKKOS_DIR}. "
+        "Set -DML_MATGL_KOKKOS_DIR=<path/to/lammps/src/KOKKOS>.")
+endif()
+
+file(GLOB ML_MATGL_KOKKOS_SOURCES "${ML_MATGL_KOKKOS_DIR}/*.cpp")
+
+target_sources(lammps PRIVATE ${ML_MATGL_KOKKOS_SOURCES})
+target_include_directories(lammps PRIVATE ${ML_MATGL_KOKKOS_DIR})
+
+# Single-GPU only: warn loudly. MACE upstream issues #1294 and #322 cover
+# the multi-rank-with-libtorch breakage we inherit.
+message(STATUS
+    "ML-MATGL-KOKKOS: enabled. Single-GPU runs only — multi-rank Kokkos with "
+    "libtorch is unreliable (see MACE issues #1294, #322).")

lammps/cmake/ML-MATGL-KOKKOS.cmake

@@ -0,0 +1,67 @@
+# ML-MATGL package — drop-in CMake snippet for a stock LAMMPS source tree.
+#
+# Usage (from a stock LAMMPS source tree):
+#   1. Copy or symlink lammps/src/ML-MATGL  →  <lammps>/src/ML-MATGL
+#   2. Append to <lammps>/cmake/CMakeLists.txt (anywhere after the `set(STANDARD_PACKAGES …)`
+#      block):
+#        include(/path/to/matgl/lammps/cmake/ML-MATGL.cmake)
+#   3. Configure with:
+#        cmake -B build \
+#            -D PKG_ML-MATGL=ON \
+#            -D CMAKE_PREFIX_PATH=/path/to/libtorch \
+#            -D CMAKE_BUILD_TYPE=Release \
+#            <other flags>
+#
+# CMake variables consumed:
+#   PKG_ML-MATGL           - turn the package on/off (default OFF).
+#   CMAKE_PREFIX_PATH      - must point at a libtorch install (CXX11 ABI build).
+#   ML_MATGL_DIR           - override path to lammps/src/ML-MATGL (defaults to
+#                            ${CMAKE_CURRENT_LIST_DIR}/../src/ML-MATGL).
+
+option(PKG_ML-MATGL "Build the matgl pair_style backed by libtorch" OFF)
+
+if(NOT PKG_ML-MATGL)
+    return()
+endif()
+
+# Locate the source directory.
+if(NOT DEFINED ML_MATGL_DIR)
+    get_filename_component(ML_MATGL_DIR
+        "${CMAKE_CURRENT_LIST_DIR}/../src/ML-MATGL" ABSOLUTE)
+endif()
+
+if(NOT EXISTS "${ML_MATGL_DIR}/pair_matgl.cpp")
+    message(FATAL_ERROR
+        "ML-MATGL source not found at ${ML_MATGL_DIR}. "
+        "Set -DML_MATGL_DIR=<path/to/lammps/src/ML-MATGL>.")
+endif()
+
+# Pull in libtorch.
+find_package(Torch REQUIRED)
+if(NOT TORCH_LIBRARIES)
+    message(FATAL_ERROR
+        "find_package(Torch) succeeded but TORCH_LIBRARIES is empty. "
+        "Did you set CMAKE_PREFIX_PATH to a libtorch install?")
+endif()
+
+# Compose the source list.
+file(GLOB ML_MATGL_SOURCES "${ML_MATGL_DIR}/*.cpp")
+
+# Hook into the LAMMPS build. This file is included from
+# <lammps>/cmake/CMakeLists.txt; the `lammps` target already exists by then.
+target_sources(lammps PRIVATE ${ML_MATGL_SOURCES})
+target_include_directories(lammps PRIVATE ${ML_MATGL_DIR})
+target_compile_features(lammps PRIVATE cxx_std_17)
+target_link_libraries(lammps PRIVATE ${TORCH_LIBRARIES})
+
+# Make sure libtorch's headers come ahead of any system Eigen/torch shims.
+target_include_directories(lammps PRIVATE ${TORCH_INCLUDE_DIRS})
+
+# LibTorch ships with -D_GLIBCXX_USE_CXX11_ABI=…; propagate it so consumers
+# (e.g. KOKKOS in Phase 3) see the same ABI.
+if(DEFINED TORCH_CXX_FLAGS)
+    set_property(TARGET lammps APPEND_STRING PROPERTY COMPILE_FLAGS " ${TORCH_CXX_FLAGS}")
+endif()
+
+message(STATUS "ML-MATGL: enabled, sources from ${ML_MATGL_DIR}")
+message(STATUS "ML-MATGL: linking against TORCH_LIBRARIES=${TORCH_LIBRARIES}")