Welcome to the torch-pme Documentation!

Overview

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.

Caution

To compute real-space short-range interactions, torch-pme requires a neighbor list, specifically the neighbor indices and their distances. We do not provide functionality to compute these neighbor lists, as there are already highly efficient libraries, such as vesin, that specialize in this task and also support auto-differentiation.

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}
   }