Metadata-Version: 2.4
Name: lammps-jax
Version: 0.2.0
Summary: CUDA-first LAMMPS/Kokkos plugin support for exported JAX models
License-Expression: MIT
Project-URL: Repository, https://github.com/abhijeetgangan/lammps-jax
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
License-File: NOTICE
Requires-Dist: jax>=0.8
Requires-Dist: jaxlib!=0.10.2,>=0.8
Provides-Extra: cpu
Requires-Dist: lammps-jax-pjrt-cpu>=0.11.2; (sys_platform == "linux" and platform_machine == "x86_64") and extra == "cpu"
Provides-Extra: dist
Requires-Dist: ase>=3.22; extra == "dist"
Requires-Dist: e3nn-jax>=0.20; extra == "dist"
Requires-Dist: jax!=0.11.2; extra == "dist"
Requires-Dist: jax-md>=0.2.27; extra == "dist"
Requires-Dist: nequix>=0.4; extra == "dist"
Provides-Extra: test
Requires-Dist: pytest>=7; extra == "test"
Requires-Dist: numpy>=1.24; extra == "test"
Dynamic: license-file

# LAMMPS-JAX

[![DOI](https://zenodo.org/badge/1278620240.svg)](https://doi.org/10.5281/zenodo.20838693)

[LAMMPS](https://www.lammps.org) plugin for running JAX forcefields with
`jax.export` and a PJRT runtime.

## Setup

The Python package is on PyPI; the LAMMPS plugin and the examples build from
this repository.

```bash
pip install lammps-jax 'jax[cuda12]'  # GPU
pip install 'lammps-jax[cpu]'         # CPU, with the XLA CPU PJRT plugin
```

From a checkout:

```bash
git clone https://github.com/abhijeetgangan/lammps-jax && cd lammps-jax
uv venv .venv --python python3
source .venv/bin/activate
uv pip install -e '.[test]'
uv pip install -U 'jax[cuda12]'
```

## Build

The pair style needs the KOKKOS precision layer from LAMMPS 10 Sep 2025 or
newer; stable releases do not carry it yet, and older trees stop at a
compile-time error. Build the plugin against the same LAMMPS source tree as
the `lmp` binary that loads it, and rebuild both together after pulling.

```bash
cmake -S cpp -B build-plugin-gpu-pjrt \
  -D CMAKE_CXX_COMPILER="$KOKKOS_NVCC_WRAPPER" \
  -D CMAKE_BUILD_TYPE=Release \
  -D CMAKE_CXX_FLAGS="-fno-lto" \
  -D CMAKE_SHARED_LINKER_FLAGS="-fno-lto" \
  -D CMAKE_INTERPROCEDURAL_OPTIMIZATION=OFF \
  -D LAMMPS_HEADER_DIR="$LAMMPS_SRC/src" \
  -D JAXLIB_INCLUDE_DIR="$JAXLIB_INCLUDE_DIR" \
  -D KOKKOS_CONFIG_INCLUDE_DIR="$LAMMPS_BUILD/lib/kokkos"
cmake --build build-plugin-gpu-pjrt -j4
```

## Run

```bash
python examples/export_model.py lj examples/lj.lammps-jax.json \
  --max-atoms 1024 --edges-per-atom 96
```

The bundle loads as a pair style:

```
pair_style jax/kk ${pjrt}
pair_coeff * * examples/lj.lammps-jax.json
```

```bash
LAMMPS_PLUGIN_PATH=build-plugin-gpu-pjrt \
  "$LAMMPS_INSTALL/bin/lmp" -k on g 1 -pk kokkos newton off neigh full -sf kk \
  -var pjrt /absolute/path/xla_cuda_plugin.so -in examples/in.lj_jax
```

## CPU

Kokkos OpenMP builds run bundles exported with `platform="cpu"` on the XLA CPU
PJRT plugin from the `cpu` extra (Linux x86-64):

```bash
uv pip install -e '.[cpu]'
python -c 'import lammps_jax_pjrt_cpu as p; print(p.PLUGIN_PATH)'
```

Build LAMMPS with `-C cmake/presets/kokkos-openmp.cmake -D PKG_PLUGIN=on
-D BUILD_MPI=on -D Kokkos_ENABLE_IMPL_VIEW_LEGACY=on` and the plugin as above
against that build, without `CMAKE_CXX_COMPILER`. One MPI rank per core runs
fastest:

```bash
python examples/export_model.py eam cuzr.lammps-jax.json \
  --setfl examples/potentials/CuZr.eam.alloy.gz \
  --max-atoms 8192 --edges-per-atom 128 --precision float64 --platform cpu
LAMMPS_PLUGIN_PATH=build-plugin-cpu-pjrt \
  mpirun -np 4 "$LAMMPS_CPU_INSTALL/bin/lmp" -k on t 1 -sf kk -pk kokkos newton on neigh half \
  -var pjrt /absolute/path/pjrt_c_api_cpu_plugin.so -in examples/in.eam_cuzr
```

cuEquivariance and OpenEquivariance kernels are CUDA only.

## Test

```bash
JAX_PLATFORMS=cpu python -m pytest
```

Integration tests, which require a GPU and LAMMPS:

```bash
LAMMPS_BIN=$LAMMPS_INSTALL/bin/lmp \
PJRT_PLUGIN=/absolute/path/xla_cuda_plugin.so \
LAMMPS_PLUGIN_PATH=$PWD/build-plugin-gpu-pjrt \
  python -m pytest tests/test_lammps.py -v
```

For a host Kokkos LAMMPS and the CPU plugin, add `LAMMPS_JAX_TEST_PLATFORM=cpu`.

## Design notes

Models are exported with `jax.export` into a JSON bundle holding the
serialized program and the settings the pair style enforces at load: atom and
edge capacities, cutoff, precision, unit style, and the distribution scheme.
Exported programs have static shapes, so positions and the edge list are
padded to capacity and an edge mask marks the live entries. The pair style
distance-filters pairs to the model cutoff and repacks the edge list every
step, so edge capacity sizes to the cutoff rather than cutoff plus skin.

The pair style executes the bundle through the PJRT C API, loading the same
plugin library jax uses on the GPU, so the build never needs XLA. Kokkos kernels
pack the LAMMPS neighbor list into sender, receiver, and mask arrays on
device, the program runs on the LAMMPS CUDA stream, and forces come back as
device buffers, so array data never passes through the host.

Multi-rank runs pick one of two schemes at export. A ghost export widens the
ghost shell to n_hops cutoffs and stays communication-free inside the
program; a comm export keeps the one-cutoff shell and exchanges per-layer
features through LAMMPS forward and reverse communication, called from inside
the program as an FFI callback. Exchanges are linear JAX primitives whose
transposes swap forward and reverse, so models may call either direction and
every differentiation mode traces through them. Communicating half-edge
exports pack each pair on one rank, matching the native half neighbor list,
and reverse-communicate ghost density partials into their owners. Exchanged
rows move in place on device under the Kokkos brick and tiled comm styles and
stage through pinned host memory otherwise.

cuEquivariance and OpenEquivariance kernels stay in the exported program as
custom call targets, resolved at run time from LAMMPS_JAX_FFI_HANDLERS;
contrib/ffi-replay handles libraries whose compiled kernels only exist inside
the exporting process. NequIP checkpoints export through the uv-runnable
scripts in examples/torch, which trace the torch blocks via torchax into
the serialized program; the bundles load and run like any native export.

## Citation

If you use `lammps-jax` or any algorithms implemented here, please cite
the archived Zenodo release: https://doi.org/10.5281/zenodo.20838694

## Acknowledgements

Thanks to @mitkotak for API discussions and @wcwitt for guidance on
distributed inference. See
[lammps/lammps#4691](https://github.com/lammps/lammps/pull/4691) for the
original discussion on running JAX models in LAMMPS, and
[openmm-jax](https://github.com/atomicarchitects/openmm-jax) for the OpenMM
counterpart.

## License

MIT. See `LICENSE`.
