Add LAMMPS pair_matgl + Kokkos plugin (Phases 1-3)

[shyuep]

Summary

Adds a LAMMPS plugin that runs MatGL TensorNet (PyG) potentials at production scale, mirroring the MACE-Kokkos pattern. Three phases:

• Phase 1 — Python TorchScript export. New matgl.ext.lammps.LAMMPSMatGLModel wraps a Potential so it takes Cartesian positions, edge_index, integer image shifts, cell, atomic numbers, and a ghost mask, and returns total local energy / per-atom energies / forces / virials. New mgl create-lammps-model CLI subcommand. 4 parity tests (eager vs Potential.forward, ghost-mask additivity, virial vs stress·volume, TorchScript save/load round-trip).
• Phase 2 — pair_matgl CPU pair style under lammps/src/ML-MATGL/. Requests a full neighbor list with ghost atoms, builds (positions / edge_index / unit_shifts / cell / atomic_numbers / local_or_ghost) tensors, calls the scripted forward, accumulates forces + 3×3 virial back into LAMMPS arrays.
• Phase 3 — Kokkos GPU/host variant pair_matgl/kk under lammps/src/KOKKOS/. Two parallel_for / parallel_scan kernels build edges on-device; Kokkos Views are wrapped as torch tensors via torch::from_blob (zero-copy on matching device); device picked from lmp -k on g 1. Single-GPU only — multi-rank Kokkos with libtorch is unreliable (MACE issues #1294, Bump boto3 from 1.34.160 to 1.34.161 #322).

CI

A new lammps-build GitHub Actions workflow runs on every push that touches lammps/, the export wrapper, or the workflow itself. It uses the materialyzeai/lammps:latest Docker image, installs the build toolchain, downloads CXX11-ABI libtorch (cached), clones LAMMPS source (cached), exports a tiny in-tree TensorNet via LAMMPSMatGLModel, builds LAMMPS with the pair_matgl source dropped into src/, runs the in.matgl_si deck, and diffs PotEng against the Python reference within 1e-3 eV.

The Kokkos variant is not exercised in CI — GitHub-hosted runners have no GPU. Hardware-accelerated CI on a self-hosted CUDA runner is the next step.

Upstream PyG layer tweaks

A few files in src/matgl/layers/_*_pyg.py and src/matgl/utils/cutoff.py got minor TorchScript-friendliness changes (explicit type annotations, returning a tuple instead of a Dict[str, Tensor], locally-defined pi). Pure additions, no runtime behavior changed; the full PyG test suite (178 tests) still passes.

Test plan

• uv run pytest tests/ext/test_lammps_export.py — 4 tests, all green.
• uv run pytest tests/ — 178 passed, 31 backend-skipped.
• uv run mypy -p matgl — only the 2 pre-existing _pymatgen_pyg.py errors remain.
• uv run ruff check src tests lammps — clean.
• lammps-build GitHub Actions workflow — green (run 25344207159, 5m52s end-to-end including LAMMPS source build).
• Single-GPU pair_matgl/kk parity check — needs CUDA hardware; deferred to follow-up.
• End-to-end run against a real foundation potential (materialyze/TensorNet-MatPES-r2SCAN) — recommended before announcing the feature.

🤖 Generated with Claude Code

[shyuep]

@kenko911 @drakeyu I seem to have gotten a Kokkos plugin for TensorNet working. The partiy tests suggest the LAMMPS and pure python output are the same. I suggest @drakeyu try out the new Kokkos plugin first and make sure it is working on a real system (e.g., MD).

[shyuep]

Just a note that the failing tests are due to the new uploaded TensorNet pre-trained models. It has nothing to do iwth the LAMMPS stuff.

[codecov]

Codecov Report

❌ Patch coverage is 83.21168% with 46 lines in your changes missing coverage. Please review.
✅ Project coverage is 94.72%. Comparing base (fcf19dc) to head (bc225e0).

Files with missing lines	Patch %	Lines
src/matgl/cli.py	26.66%	22 Missing ⚠️
src/matgl/ext/_lammps.py	92.17%	14 Missing ⚠️
src/matgl/graph/_compute_pyg.py	86.11%	5 Missing ⚠️
src/matgl/ext/lammps.py	0.00%	3 Missing ⚠️
src/matgl/layers/_graph_convolution_pyg.py	66.66%	2 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #770      +/-   ##
==========================================
- Coverage   95.16%   94.72%   -0.44%     
==========================================
  Files          81       83       +2     
  Lines        7388     7647     +259     
==========================================
+ Hits         7031     7244     +213     
- Misses        357      403      +46     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[shyuep]

Automated PR review generated by Claude (Opus 4.7) on behalf of @shyuep. Based purely on the diff and current GitHub Actions state — no code was executed locally.

This is a substantial and well-scoped piece of work — three logically separate phases (Python TorchScript export, CPU pair_matgl, Kokkos pair_matgl/kk) all delivered together with a real CI gate. Mirroring the MACE-Kokkos pattern is the right design choice.

CI status: all functional checks green — pair_matgl (CPU) build + in.matgl_si parity diff against the Python reference is the most important signal here, and that's passing in 5m52s. lint, test_pyg, test_dgl, coverage, and pre-commit.ci are all green. The only red marks are codecov/patch and codecov/project, expected for a large addition where ~all of lammps/ (C++) and the Kokkos kernels are not covered by Python tests. Not a blocker.

Strengths:

• The 4 parity tests in tests/ext/test_lammps_export.py (eager vs Potential.forward, ghost-mask additivity, virial vs stress · volume, TorchScript round-trip) hit exactly the right invariants. These are the things that silently break when a wrapper drifts.
• Honest scoping in the description re: single-GPU Kokkos limitation, with concrete upstream MACE issue references (#1294, #322). Good engineering hygiene.
• materialyzeai/lammps:latest Docker image + cached libtorch / LAMMPS source keeps CI runtime sane.
• 1e-3 eV PotEng tolerance against the Python reference is a tight, meaningful gate.

Things to verify before announcing the feature (already in your test plan but flagging for the record):

1. Single-GPU pair_matgl/kk parity check on real CUDA hardware — the Kokkos variant is unexercised in CI. Consider standing up a self-hosted CUDA runner or at minimum a documented manual smoke procedure committed under lammps/tests/ so it's reproducible.
2. End-to-end run with a real foundation potential (materialyze/TensorNet-MatPES-r2SCAN) before announcing — the in-tree tiny TensorNet exercises the plumbing but doesn't catch numerical-stability quirks at production scale.
3. Multi-rank guard: since multi-rank Kokkos+libtorch is unreliable, a clear error->all(...) if comm->nprocs > 1 in pair_matgl/kk setup would protect users from surprising hangs. (Worth checking whether this is already in pair_matgl_kokkos.cpp.)

Minor / non-blocking:

• src/matgl/ext/_lammps.py is 498 lines — fine for now, but if the export-wrapper logic grows it might be worth splitting the LAMMPSMatGLModel class from the helper functions (graph construction, ghost-mask handling).
• Pre-existing _pymatgen_pyg.py mypy errors are noted as out-of-scope; consider opening a follow-up issue to track those so they don't get forgotten.
• The PyG layer tweaks in _basis.py, _core.py, _embedding*.py, _graph_convolution_pyg.py, cutoff.py for TorchScript-friendliness are pure additions per the description and the 178-test PyG suite stays green — good. Worth one targeted test that explicitly TorchScripts the resulting Potential end-to-end (not just the wrapper) so future PyG-layer changes can't silently break scriptability.

Overall: ready to merge once the items in your own test-plan checklist (single-GPU Kokkos parity, real foundation-potential end-to-end run) are checked off. The Phase 1+2 surface is in good shape.

[shyuep]

🤖 Automated PR review (generated by Claude on behalf of @shyuep)

• ✅ Phase 1 (TorchScript export wrapper + parity tests covering eager-vs-Potential, ghost-mask additivity, virial-vs-stress·V, and round-trip) is the right shape — those are the assertions that actually catch regressions in a scripted model.
• ✅ Phase 2 CPU pair_matgl integrates cleanly with LAMMPS's full-with-ghost neighbor list. The new lammps-build workflow that diffs PotEng against the Python reference within 1e-3 eV is a strong CI gate.
• ⚠️ pre-commit.ci - pr is failing on the head commit; please clean up before merging.
• ⚠️ Test plan still has two unchecked items: single-GPU pair_matgl/kk parity and an end-to-end run against a real foundation potential (materialyze/TensorNet-MatPES-r2SCAN). At minimum the real-potential run should gate the merge — Kokkos can ship behind a clear "untested on hardware" doc note since no GPU CI is available.
• 💡 23 files / 2717 additions is a lot. The PyG layer tweaks in src/matgl/layers/_*_pyg.py and src/matgl/utils/cutoff.py (TorchScript-friendliness changes) are claimed to be pure additions — would be worth calling out explicitly in the changelog so downstream PyG users know what to expect.
• 💡 Codecov: patch coverage 82.4%, with src/matgl/cli.py at 26.7% (22 missing). The new mgl create-lammps-model CLI subcommand should at minimum get a smoke test (export + reload) so the CLI surface doesn't regress silently.
• 💡 Single-GPU only is a real constraint; consider adding a runtime check in pair_matgl/kk that errors with a clear message when launched with multi-rank Kokkos rather than failing opaquely inside libtorch.

---

Generated by Claude Code

[shyuep]

Automated PR review generated by Claude (Opus 4.7) on behalf of @shyuep. Follow-up note — no code was executed locally.

Since the prior automated review at 6b4edfb, the only new commits on this branch are four Merge branch 'main' into lammps merges (5df9481, 3faa2c8, c32dcac, af0f6ed) — no new functional changes. The substantive review still stands. Two CI signals worth flagging from the latest run:

1. New lint failure (real, from a ruff bump): ruff 0.15.12 now flags RET504 Unnecessary assignment to 'atomic' before 'return' statement in _lammps.py:255. One-line fix — drop the intermediate atomic variable and return the expression directly. The pre-commit.ci failure is the same rule via the autofix path.
2. test_pyg (ubuntu-latest, 3.13, ops) is not a real failure — the runner hit remote: Internal Server Error (HTTP 500) on git fetch ... pull/770/merge three times in a row before the job aborted. Pure GitHub infra flake; re-run after the lint fix lands.

codecov/patch and codecov/project remain red for the same reason called out previously (C++/Kokkos kernels uncovered by Python tests) — non-blocking.

Otherwise unchanged: ready to merge once the items in your own test-plan checklist (single-GPU Kokkos parity, real foundation-potential end-to-end run) are checked off.

[shyuep]

🤖 Automated PR review (generated by Claude on behalf of @shyuep) — re-review after the latest commits.

Two new functional commits since the prior review (af0f6ed):

• ✅ 580255d — script-safe SphericalBessel bond expansion. A real fix: the pretrained TensorNet-MatPES-r2SCAN checkpoint uses rbf_type='SphericalBessel', whose _call_smooth_sbf is @torch.jit.ignore, so torch.jit.script(...).save(...) previously errored. The new _SmoothSBFExpansion adapter wraps the pure-tensor spherical_bessel_smooth, and a tensornet_sb parity fixture now exercises that path. This directly clears one of the open items from the prior review (real-foundation-potential export), and the explicit NotImplementedError for the non-smooth case is the right call.
• ✅ 6bd13af — narrows bond_expansion typing for --strict mypy, consistent with the existing cast pattern in the file.
• ✅ CI: all functional jobs green — lint, the pair_matgl (CPU) parity diff, and the PyG/DGL test suites. Only codecov/patch + codecov/project are red, expected for the C++/Kokkos additions that Python tests don't cover — non-blocking.
• ⚠️ Still open from the prior review: single-GPU pair_matgl/kk parity on real CUDA hardware remains unexercised in CI. That's the last meaningful gate before announcing the feature — Phase 1+2 look merge-ready.

---

Generated by Claude Code — automated review, no code executed locally.

---

Generated by Claude Code

[shyuep]

🤖 Automated PR review (generated by Claude on behalf of @shyuep) — re-review after the latest commits (580255d, 6bd13af).

Reviewed the delta since the prior review from the GitHub diff + CI only — no code executed locally.

• ✅ Adds _SmoothSBFExpansion, a pure-tensor TorchScript stand-in for BondExpansion(rbf_type='SphericalBessel', smooth=True). _TensorNetKernel swaps it in for SphericalBessel models and raises a clear NotImplementedError on the non-smooth path. This is the right fix — the default foundation potential (TensorNet-MatPES-r2SCAN) uses the SphericalBessel basis, so without it the flagship model can't be exported to LAMMPS.
• ✅ _basis.py: spherical_bessel_smooth now inlines pi/sqrt(2) as locals so it survives torch.jit.script; @torch.jit.ignore / @torch.jit.unused annotations added where needed. Import and call signature in _lammps.py line up.
• ✅ Test fixture parametrized to cover tensornet_sb (SphericalBessel-smooth) alongside Gaussian and M3GNet — good coverage of the new path.
• 💡 Minor: from typing import cast is imported inside _TensorNetKernel.__init__; consider hoisting it to the module-level typing import.
• CI: all functional test jobs pass. The 3 red checks are non-blocking — codecov/patch + codecov/project (coverage thresholds) and pre-commit.ci (ruff PLW2901/SIM105 on examples/*.ipynb notebooks not touched by this PR — pre-existing lint debt worth a separate cleanup on main).

LGTM.