https://github.com/lab-cosmo/torch-pme
Particle-mesh based calculations of long-range interactions in PyTorch
https://github.com/lab-cosmo/torch-pme
electrostatics fast-fourier-transform machine-learning material-science molecular-dynamics particle-mesh-ewald pytorch
Last synced: 12 days ago
JSON representation
Particle-mesh based calculations of long-range interactions in PyTorch
- Host: GitHub
- URL: https://github.com/lab-cosmo/torch-pme
- Owner: lab-cosmo
- License: bsd-3-clause
- Created: 2023-08-28T21:02:26.000Z (over 1 year ago)
- Default Branch: main
- Last Pushed: 2025-05-05T13:58:34.000Z (14 days ago)
- Last Synced: 2025-05-07T17:12:57.708Z (12 days ago)
- Topics: electrostatics, fast-fourier-transform, machine-learning, material-science, molecular-dynamics, particle-mesh-ewald, pytorch
- Language: Python
- Homepage: https://lab-cosmo.github.io/torch-pme/
- Size: 8.07 MB
- Stars: 47
- Watchers: 17
- Forks: 5
- Open Issues: 18
-
Metadata Files:
- Readme: README.rst
- Contributing: CONTRIBUTING.rst
- License: LICENSE
- Code of conduct: CODE_OF_CONDUCT.md
- Citation: CITATION.cff
Awesome Lists containing this project
README
torch-pme
=========
.. image:: https://raw.githubusercontent.com/lab-cosmo/torch-pme/refs/heads/main/docs/logo/torch-pme.svg
:width: 200 px
:align: left
|tests| |codecov| |docs|
.. marker-introduction
``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 :math:`1/r`
electrostatics, including arbitrary order :math:`1/r^p` potentials. For **dipolar**
interaction we offer to calculate the :math:`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`_.
.. _`TorchScriptable`: https://pytorch.org/docs/stable/jit.html
.. _`jax-pme`: https://github.com/lab-cosmo/jax-pme
.. marker-documentation
Documentation
-------------
For details, tutorials, and examples, please have a look at our `documentation`_.
.. _`documentation`: https://lab-cosmo.github.io/torch-pme
.. marker-installation
Installation
------------
You can install *torch-pme* using pip with
.. code-block:: bash
pip install torch-pme
or conda
.. code-block:: bash
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
.. code-block:: bash
pip install torch-pme[metatensor]
.. marker-quickstart
Quickstart
----------
Here is a simple example to get started with *torch-pme*:
.. code-block:: python
>>> 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`_.
.. marker-issues
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
`_.
.. marker-cite
Reference
---------
If you use *torch-pme* for your work, please read and cite our preprint available on
`arXiv`_.
.. code-block::
@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}
}
.. _`arXiv`: http://arxiv.org/abs/2412.03281
.. marker-contributing
Contributors
------------
Thanks goes to all people that make *torch-pme* possible:
.. image:: https://contrib.rocks/image?repo=lab-cosmo/torch-pme
:target: https://github.com/lab-cosmo/torch-pme/graphs/contributors
.. |tests| image:: https://github.com/lab-cosmo/torch-pme/workflows/Tests/badge.svg
:alt: Github Actions Tests Job Status
:target: https://github.com/lab-cosmo/torch-pme/actions?query=branch%3Amain
.. |codecov| image:: https://codecov.io/gh/lab-cosmo/torch-pme/graph/badge.svg?token=srVKRy7r6m
:alt: Code coverage
:target: https://codecov.io/gh/lab-cosmo/torch-pme
.. |docs| image:: https://img.shields.io/badge/documentation-latest-sucess
:alt: Documentation
:target: `documentation`_