Metadata-Version: 2.2
Name: torch-pme
Version: 0.2.0
Summary: Particle-mesh based calculations of long-range interactions in PyTorch
Author: torch-pme developers
License: BSD-3-Clause
Project-URL: homepage, https://lab-cosmo.github.io/torch-pme
Project-URL: documentation, https://lab-cosmo.github.io/torch-pme
Project-URL: repository, https://github.com/lab-cosmo/torch-pme
Project-URL: issues, https://github.com/lab-cosmo/torch-pme/issues
Project-URL: changelog, https://lab-cosmo.github.io/torch-pme/latest/references/changelog.html
Keywords: PyTorch,Auto-differentiation,Particle-Mesh Ewald,Electrostatics,Computational Materials Science,Machine Learning,Molecular Dynamics,GPU Acceleration,High-Performance Computing,Fourier Transform,TorchScript,Scientific Computing
Classifier: Development Status :: 4 - Beta
Classifier: Intended Audience :: Science/Research
Classifier: License :: OSI Approved :: BSD License
Classifier: Natural Language :: English
Classifier: Operating System :: POSIX
Classifier: Operating System :: MacOS :: MacOS X
Classifier: Operating System :: Microsoft :: Windows
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Classifier: Topic :: Scientific/Engineering
Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
Classifier: Topic :: Scientific/Engineering :: Chemistry
Classifier: Topic :: Scientific/Engineering :: Physics
Classifier: Topic :: Software Development :: Libraries
Classifier: Topic :: Software Development :: Libraries :: Python Modules
Requires-Python: >=3.9
Description-Content-Type: text/x-rst
License-File: LICENSE
Requires-Dist: torch>=2.3
Provides-Extra: examples
Requires-Dist: ase>=3.22.1; extra == "examples"
Requires-Dist: chemiscope; extra == "examples"
Requires-Dist: matplotlib; extra == "examples"
Requires-Dist: metatensor-operations; extra == "examples"
Requires-Dist: scipy; extra == "examples"
Requires-Dist: vesin>=0.3.0; extra == "examples"
Requires-Dist: vesin-torch>=0.3.0; extra == "examples"
Provides-Extra: metatensor
Requires-Dist: metatensor-torch<0.7,>=0.6; extra == "metatensor"

torch-pme
=========

.. image:: https://raw.githubusercontent.com/lab-cosmo/torch-pme/refs/heads/main/docs/src/logo/torch-pme.svg
   :width: 200 px
   :align: left

|tests| |codecov| |docs|

.. marker-introduction

``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 :math:`1/r` electrostatics, including arbitrary order :math:`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`_.

.. _`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 <https://docs.metatensor.org>`_ 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 (cutoff optional, useful to set for ML with fixed cutoff)
   smearing, p3m_parameters, _ = torchpme.tuning.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`_.

.. marker-issues

Having problems or ideas?
-------------------------

Having a problem with *torch-pme*? Please let us know by `submitting an issue
<https://github.com/lab-cosmo/torch-pme/issues>`_.

Submit new features or bug fixes through a `pull request
<https://github.com/lab-cosmo/torch-pme/pulls>`_.

.. 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`_
