ArrayFF

The Force Field Library Based on Multiple Array Backends like NumPy, PyTorch and JAX.

Kernels are written once with autoray, so the backend is chosen by the arrays you pass in. Structural I/O uses ase.Atoms throughout.

The only required array backends are NumPy and autograd (for NumPy autodiff); PyTorch and JAX are optional extras.

Install

uv pip install -e .
uv pip install -e ".[torch]"     # optional PyTorch backend + nn wrappers
uv pip install -e ".[jax]"       # optional JAX backend
uv pip install -e ".[builders]"  # vesin / matscipy / pymatgen neighbours

Quick start

from ase.build import bulk
from arrayff.potentials import EMT, EMTCalculator

atoms = bulk("Cu", "fcc", a=3.6, cubic=True).repeat((2, 2, 2))

calc = EMT()                       # numpy backend + autograd for forces
energy, forces = calc.energy_and_forces(atoms, method="analytic")

atoms.calc = EMTCalculator()       # ASE calculator interface
energy = atoms.get_potential_energy()

Backends are selected at construction:

calc = EMT(backend="torch")        # or "jax"
calc = EMT(backend="jax", jit=True)

LennardJones() without arguments uses the OpenKIM universal shifted parameters (118 elements, cutoff 4*sigma per pair, Lorentz-Berthelot mixing: arithmetic r0/cutoff and geometric epsilon).

Each potential exposes analytic and autodiff forces/stress:

from arrayff.potentials import LennardJones

calc = LennardJones(cutoff=8.0, sigma=3.405, epsilon=0.01032)
calc.energy_and_forces(atoms, method="analytic")
calc.energy_and_forces(atoms, method="autodiff")
calc.stress(atoms, method="analytic")

jit=True compiles the array-level kernel with autoray.autojit (JAX and PyTorch are the useful backends; NumPy's compiler is best-effort).

Charge equilibration and composite potentials

from arrayff.electrostatics import QEq

qeq = QEq(chi={8: 8.5, 1: 4.5}, J={8: 10.0, 1: 12.0}, shielding=1.0)
charges = qeq.charges(atoms)       # sum == total charge

# additive composite: energy/forces/stress add, charges come from QEq
composite = LennardJones(cutoff=8.0, sigma=3.405, epsilon=0.01032) + qeq
energy, forces = composite.energy_and_forces(atoms)

QEq supports charge_response="frozen" (Hellmann-Feynman, default) and "relaxed" (differentiate through the charge solve).

Training (optional)

from arrayff.potentials import LJTorch

model = LJTorch(species=[18])           # epsilon/sigma trainable
loss = (model(atoms) - target_energy) ** 2
loss.backward()

Tests and benchmarks

pytest -q
pytest -s src/tests-benchmark/bench_potentials.py
pytest -s src/tests-benchmark/bench_derivatives.py

The benchmark tables compare analytic vs autodiff forces/stress and jit=False/True across backends.

Status