torch-pme

torch-pme enables efficient, auto-differentiable computation of long-range interactions in PyTorch. Auto-differentiation is supported for particle positions, charges, and cell parameters, allowing not only the automatic computation of forces but also enabling general applications in machine learning tasks. The library offers classes for Particle-Particle Particle-Mesh Ewald (P3M), Particle Mesh Ewald (PME), standard Ewald, and non-periodic methods, with the flexibility to calculate potentials beyond 1/r electrostatics, including arbitrary order 1/r^p potentials.

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 get you started with torch-pme:

import torch
import torchpme

# Single charge in a cubic box
positions = torch.zeros((1, 3), requires_grad=True)
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 (cutoff optional, useful to set for ML with fixed cutoff)
smearing, p3m_parameters, _ = torchpme.utils.tune_p3m(
   sum_squared_charges=1,
   cell=cell,
   positions=positions,
   cutoff=5.0,
)

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

# 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

print("Energy:", energy.item())
print("Forces:", forces)

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 preprint available on arXiv.

@article{loche_fast_2024,
   title = {Fast and Flexible Range-Separated 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},
   year = {2024},
   month = dec,
   number = {arXiv:2412.03281},
   eprint = {2412.03281},
   primaryclass = {physics},
   publisher = {arXiv},
   doi = {10.48550/arXiv.2412.03281},
   urldate = {2024-12-05},
   archiveprefix = {arXiv}
   }

Contributors

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