torch-pme

torch-pme enables efficient and auto-differentiable computation of long-range interactions in PyTorch. Auto-differentiation is supported for particle positions,

charges/dipoles, and cell parameters, allowing not only the automatic computation of forces but also enabling general applications in machine learning tasks. For monopoles the library offers classes for Particle-Particle Particle-Mesh Ewald (P3M), Particle Mesh Ewald (PME), standard Ewald, and non-periodic methods. The library has the flexibility to calculate potentials beyond 1/r electrostatics, including arbitrary order 1/r^p potentials. For dipolar interaction we offer to calculate the 1/r^3 potential using the standard Ewald method.

Optimized for both CPU and GPU devices, torch-pme is fully TorchScriptable, allowing it to be converted into a format that runs independently of Python, such as in C++, making it ideal for high-performance production environments.

We also provide an experimental implementation for JAX in jax-pme.

Documentation

For details, tutorials, and examples, please have a look at our documentation.

Installation

You can install torch-pme using pip with

pip install torch-pme

or conda

conda install -c conda-forge torch-pme

and import torchpme to use it in your projects!

We also provide bindings to metatensor which can optionally be installed together and used as torchpme.metatensor via

pip install torch-pme[metatensor]

Quickstart

Here is a simple example to get started with torch-pme:

>>> import torch
>>> import torchpme

>>> # Single charge in a cubic box
>>> positions = torch.zeros((1, 3))
>>> cell = 8 * torch.eye(3)
>>> charges = torch.tensor([[1.0]])

>>> # No neighbors for a single atom; use `vesin` for neighbors if needed
>>> neighbor_indices = torch.zeros((0, 2), dtype=torch.int64)
>>> neighbor_distances = torch.zeros((0,))

>>> # Tune P3M parameters
>>> smearing, p3m_parameters, _ = torchpme.tuning.tune_p3m(
...    charges=charges,
...    cell=cell,
...    positions=positions,
...    cutoff=5.0,
...    neighbor_indices=neighbor_indices,
...    neighbor_distances=neighbor_distances,
... )

>>> # Initialize potential and calculator
>>> potential = torchpme.CoulombPotential(smearing)
>>> calculator = torchpme.P3MCalculator(potential, **p3m_parameters)

>>> # Start recording operations done to ``positions``
>>> _ = positions.requires_grad_()

>>> # Compute (per-atom) potentials
>>> potentials = calculator.forward(
...    charges=charges,
...    cell=cell,
...    positions=positions,
...    neighbor_indices=neighbor_indices,
...    neighbor_distances=neighbor_distances,
... )

>>> # Calculate total energy and forces
>>> energy = torch.sum(charges * potentials)
>>> energy.backward()
>>> forces = -positions.grad

For more examples and details, please refer to the documentation.

Having problems or ideas?

Having a problem with torch-pme? Please let us know by submitting an issue.

Submit new features or bug fixes through a pull request.

Reference

If you use torch-pme for your work, please read and cite our publication available on JCP.

@article{10.1063/5.0251713,
   title = {Fast and flexible long-range models for atomistic machine learning},
   author = {Loche, Philip and Huguenin-Dumittan, Kevin K. and Honarmand, Melika and Xu, Qianjun and Rumiantsev, Egor and How, Wei Bin and Langer, Marcel F. and Ceriotti, Michele},
   journal = {The Journal of Chemical Physics},
   volume = {162},
   number = {14},
   pages = {142501},
   year = {2025},
   month = {04},
   issn = {0021-9606},
   doi = {10.1063/5.0251713},
   url = {https://doi.org/10.1063/5.0251713},
}

Contributors

Thanks goes to all people that make torch-pme possible: