jax-pme: Particle-mesh based calculations of long-range interactions in JAX

This is an experimental version of the torch-pme package, written using jax instead of torch. It currently offers a subset of its features. Note that the API is not yet finalised -- we appreciate any feedback! There will be breaking changes without announcement.

To learn more about Ewald summation and its particle-mesh variants, which this package implements, please have a look at our preprint:

Title: Fast and flexible range-separated models for atomistic machine learning
Authors: Philip Loche, Kevin K. Huguenin-Dumittan, Melika Honarmand, Qianjun Xu, Egor Rumiantsev, Wei Bin How, Marcel F. Langer, Michele Ceriotti
Preprint: arXiv:2412.03281 (2024)

You should also check the torch-pme documentation!

Installation

This package requires jax. Please make sure to install the appropriate version for your setup.

Once this has been done, you can install the latest release from PyPI:

pip install jax-pme

For development, clone this repository and run pip install -e . inside the folder. The dev extra contains development dependencies, and can be installed by adding .[dev] to the command.

Usage

The interface of this package is not yet fully designed. Please file an issue or get in touch via marcel.langer@epfl.ch if you have a particular use-case in mind and would like to chat about how to best support it.

Currently, the high-level API of this package is designed to compute (a) potentials, (b) total energy, (c) forces, and (d) stress for pairwise potentials of the form 1/r**p, with p=1 defining the important case of Coulomb interactions, i.e., electrostatics. The API is designed for standalone computation of these quantities, and is not particularly optimised for integration into machine learning model architectures yet. For example, it expects to be given positions and the cell as inputs, as opposed to graph edges, and it expects a half neighborlist by default (full_neighbor_list=False).

Design

This package has to respect the constraints of jax, and is therefore designed differently than its torch counterpart. The relevant issue here is that jax is not designed to manage stateful classes, and it requires array shapes to be known ahead of time. Preprocessing steps (k-grid shape computation, charge handling) use plain numpy to avoid unnecessary JAX device allocation. For example, the shape of the reciprocal-space grid (the $k$-grid), which depends both on convergence parameters like lr_wavelength (the cutoff in reciprocal space, i.e., the minimum wavelength) and the periodic cell of the system, has to be known ahead of time. This is different to the torch version, where we can just compute it in forward. We also need to be careful to ensure that all operations relevant to differentiation happen inside the scope of the calculation function, since jax traces the computation, and doesn't track the arrays across compuations like torch.

jax-pme is designed accordingly: The actual compute functions are pure functions that accept the relevant arguments for differentiation (positions, charges, cell) as well as information like the shape of the k-grid. As a consequence, they can be traced and transformed by jax, for instance with jax.grad and jax.jit. Provided inputs are padded appropriately, even vmap and scan can be used.

For convenience, we also provide helper functions that transform more conventional inputs, together with the relevant convergence parameters, into the inputs required by jax-pme. We currently expect the user to take care of padding inputs to common shapes, but we support masking out padded inputs to enable this.

API

The high-level API is provided by Calculator classes, which are simple namedtuples of functions:

Calculator = namedtuple(
    "Calculator",
    ("prepare", "potentials", "energy", "energy_forces", "energy_forces_stress"),
)

They are instantiated just like any other class, by calling

from jaxpme import Ewald, PME, P3M

calculator = Ewald(
    exponent=1,  # p in 1/r^p, integer 1-6; 1 corresponds to electrostatics
    exclusion_radius=None,  # if this is not None, purely long-range potentials are computed (see preprint)
    prefactor=1.0,  # default to Gauss units. jaxpme.prefactors.eV_A for standard ase units
    custom_potential=None,  # mostly for testing -- you can define custom potential functions
    full_neighbor_list=False,  # set True if your neighborlist includes both i->j and j->i
)

calculator = PME(
    exponent=1,  # p in 1/r^p, integer 1-6; 1 corresponds to electrostatics
    exclusion_radius=None,  # if this is not None, purely long-range potentials are computed (see preprint)
    prefactor=1.0,  # default to Gauss units. jaxpme.prefactors.eV_A for standard ase units
    interpolation_nodes=4,  # currently only 4 is supported
    custom_potential=None,  # mostly for testing -- you can define custom potential functions
    full_neighbor_list=False,  # set True if your neighborlist includes both i->j and j->i
)

calculator = P3M(
    exponent=1,  # p in 1/r^p, integer 1-6; 1 corresponds to electrostatics
    exclusion_radius=None,  # if this is not None, purely long-range potentials are computed (see preprint)
    prefactor=1.0,  # default to Gauss units. jaxpme.prefactors.eV_A for standard ase units
    interpolation_nodes=4,  # B-spline interpolation, supports 1-5
    custom_potential=None,  # mostly for testing -- you can define custom potential functions
    full_neighbor_list=False,  # set True if your neighborlist includes both i->j and j->i
)

# -> calculator.prepare, .energy, etc ... can be called

The functions exposed by Calculator consist of a prepare function that arranges all the inputs required for calculations of some input structure, including determining the shape of the reciprocal-space grid, and a bundle of functions that then execute different calculations.

prepare expects the arguments atoms (ase.Atoms instance), charges, cutoff (for the real-space neighborlist), mesh_spacing (PME/P3M) or lr_wavelength (Ewald) (defining the resolution/cutoff in reciprocal space), smearing (range separation parameter, related to cutoff). The parameters can be tuned with torch-pme or set heuristically (see torch-pme docs). It returns a tuple of inputs charges, *graph, k_grid, smearing, where *graph collects cell, positions, neighbor indices i and j, and cell_shifts. k_grid is a dummy array that defines the shape of the reciprocal-space grid via its shape, its values are not used. prepare is not jax.jit-able as it returns variable-shape output. (Ewald additionally appends pbc to the tuple, used to route non-periodic structures; see below.)

Non-periodic structures (pbc=[False, False, False]) are supported by Ewald: there is no periodic image to sum, so prepare builds the all-pairs list (half by default, both directions with full_neighbor_list=True) and potentials evaluates a bare $1/r^p$ sum (no real-space cutoff, no range separation). Serial PME/P3M do not yet support non-PBC inputs (their reciprocal block divides by the zero cell volume) — use Ewald or a batched calculator instead.

The following calculation functions are implemented:

• potentials: Accepts the above inputs and returns the potential values at each position.
• energy: Computes the total energy obtained by multplying the charge at each position with the potential at each positions and summing up.
• energy_forces: Energy as above, and its derivative with respect to positions.
• energy_forces_stress: The above, with additionally the stress.

All these other functions can be jit-ed and support function transformations like vmap and grad. They optionally accept boolean mask arrays atom_mask and pair_mask to exclude irrelevant inputs from the output, typically introduced by padding. We currently do not support padding the $k$-grid, you should simply use the biggest grid consistently. Make sure that padding indices do not connect non-padded edges. Additionally, Ewald's potentials_fn accepts optional distances (to skip recomputing pairwise distances) and pbc kwargs; PME and P3M accept pbc.

The low-level API is not yet ready for public consumption. We split the calculation task into sub-problems: solvers.py defines the actual implementations of the Ewald and PME method, potentials.py defines the actual potential functions, and the other files implement various helper functionality.

Recommendations

We find in benchmarks that for moderately-sized systems up to a few thousand atoms, the asymptotically less efficient Ewald method works best. For large systems, PME and P3M are preferable, as they scale $O(N \log N)$. Note that PME is not smooth in the forces -- be careful when using it for dynamics. P3M uses B-spline interpolation (with interpolation_nodes 1-5) and an influence function correction, giving smoother forces and better accuracy. Use P3M for molecular dynamics.

It is highly recommended to tune convergence parameters for your specific system. torchpme.utils.tune_ewald (and its pme version) exists for this purpose. Paramters can be used directly in jax-pme. You should typically tune the parameters for the largest system in a given dataset.

Batched computation

Warning: The batched API is highly experimental and subject to breaking changes. Use at your own risk.

For computing energies/forces across multiple structures (e.g. for training), batched implementations are available:

from jaxpme.batched_mixed import Ewald  # rectangular max-padding, supports cutoff-only API
from jaxpme.batched_tiled import Ewald  # per-system sum-padding + tile dispatch (see below)
from jaxpme.batched_flat import Ewald  # alternative flat padding strategy

All three accept lists of ase.Atoms in prepare and handle padding/masking internally. Currently, only batched Ewald is implemented. 2D PBC (slab geometries) is supported for arbitrary triclinic cells; large vacuum gaps are automatically shrunk to keep the k-grid efficient. The shrunk cell travels separately from the raw Batch.cell and the calculators compose the two at entry — see jaxpme.utils.compose_cell for the mechanism and gradient policy (2D-slab stress no longer includes the shrink's artifact gradient). Consequence: for 2D PBC the stress components touching the non-periodic direction carry only the per-atom term and should be treated as meaningless; the in-plane block is correct.

batched_tiled is a second Ewald backend designed for heterogeneous batches: atoms are sum-padded per system (each system padded to ⌈N_b/BM⌉·BM atoms, concatenated into one flat array) rather than max-padded to the batch's largest system. The reciprocal sum runs through a pure-JAX tile-dispatched kernel over fixed-size (BM × BK) work tiles — vmap + segment_sum for pass 1's structure factors, vmap + reshape-sum for pass 2's per-atom potential. Trade-offs vs batched_mixed:

• num_k is required on prepare and fixes the reciprocal grid (per-cell K target via lr_wavelength_for_num_k; the K axis stays rectangular so all systems share K_pad). The real-space cutoff follows it: omit it and it is derived as lr_wavelength · 8 (with smearing = lr_wavelength · 2), matching batched_mixed, so real and reciprocal space stay balanced. An explicit cutoff overrides only the real-space radius (smearing still tracks num_k).
• Tile sizes (BM, BK) are fixed at prepare time (defaults BM=32, BK=128). They drive both the per-system atom padding and the kernel tile dimensions.
• Lower memory and faster on heterogeneous batches where system sizes vary by a lot (small molecules + larger crystals/MOFs in the same batch).
• Host-side batching is exposed for external pipelines: batched_tiled.batching.sample_shapes gives the per-sample size accounting get_batch itself uses (for batch-size planners), and get_batch(samples=[], dtype=...) builds a pure-padding batch at explicit sizes; int_dtype= sets the neighbor-list index dtype (default int64).

Development

The package uses ruff for linting and formatting and pytest for testing. Please run ruff format . && ruff check --fix . before every commit or set up a commit hook to do it. Tests can be run in the tests/ folder with pytest. Be aware that the test suite can take a few minutes to run.

To release a new version, push a v* tag. This triggers a workflow that builds the package and publishes it to PyPI; the version is taken from the tag via setuptools_scm.

This project is maintained by @E-Rum, and @sirmarcel, who will reply to issues and pull requests opened on this repository as soon as possible. You can mention them directly if you did not receive an answer after a couple of days.