Metadata-Version: 2.4
Name: jax-pme
Version: 0.1.0
Summary: (Particle Mesh) Ewald summation in JAX
Author-email: Marcel Langer <dev@marcel.science>
License-Expression: MIT
Project-URL: Homepage, https://github.com/lab-cosmo/jax-pme
Project-URL: Repository, https://github.com/lab-cosmo/jax-pme
Project-URL: Issues, https://github.com/lab-cosmo/jax-pme/issues
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Requires-Python: >=3.10
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: jax>=0.4.30
Requires-Dist: vesin>=0.4.0
Requires-Dist: ase>=3.22.1
Requires-Dist: numpy
Provides-Extra: dev
Requires-Dist: pytest; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Dynamic: license-file

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

This is an experimental version of the [`torch-pme`](https://github.com/lab-cosmo/torch-pme) package, written using [`jax`](https://github.com/jax-ml/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](https://arxiv.org/abs/2412.03281):

```
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](https://lab-cosmo.github.io/torch-pme/latest)!

## 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:

```bash
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 `namedtuple`s of functions:

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

They are instantiated just like any other class, by calling

```python
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:

```python
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](https://github.com/lab-cosmo/.github/blob/main/Maintainers.md) by [@E-Rum](https://github.com/E-Rum), and [@sirmarcel](https://github.com/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.
