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 to get 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
>>> 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)

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