torch-pme


Nametorch-pme JSON
Version 0.3.0 PyPI version JSON
download
home_pageNone
SummaryParticle-mesh based calculations of long-range interactions in PyTorch
upload_time2025-02-21 17:09:07
maintainerNone
docs_urlNone
authortorch-pme developers
requires_python>=3.9
licenseBSD-3-Clause
keywords auto-differentiation computational materials science electrostatics fourier transform gpu acceleration high-performance computing machine learning molecular dynamics particle-mesh ewald pytorch scientific computing torchscript
VCS
bugtrack_url
requirements No requirements were recorded.
Travis-CI No Travis.
coveralls test coverage No coveralls.
            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 <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
   >>> 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
<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`_

            

Raw data

            {
    "_id": null,
    "home_page": null,
    "name": "torch-pme",
    "maintainer": null,
    "docs_url": null,
    "requires_python": ">=3.9",
    "maintainer_email": null,
    "keywords": "Auto-differentiation, Computational Materials Science, Electrostatics, Fourier Transform, GPU Acceleration, High-Performance Computing, Machine Learning, Molecular Dynamics, Particle-Mesh Ewald, PyTorch, Scientific Computing, TorchScript",
    "author": "torch-pme developers",
    "author_email": null,
    "download_url": "https://files.pythonhosted.org/packages/c6/34/a414288eccff5352e84271f2b09dfe6f9b18ff2da1b6f012dc35908d4a67/torch_pme-0.3.0.tar.gz",
    "platform": null,
    "description": "torch-pme\n=========\n\n.. image:: https://raw.githubusercontent.com/lab-cosmo/torch-pme/refs/heads/main/docs/logo/torch-pme.svg\n   :width: 200 px\n   :align: left\n\n|tests| |codecov| |docs|\n\n.. marker-introduction\n\n``torch-pme`` enables efficient and auto-differentiable computation of long-range\ninteractions in *PyTorch*. Auto-differentiation is supported for particle *positions*,\n\n*charges*/*dipoles*, and *cell* parameters, allowing not only the automatic computation\nof forces but also enabling general applications in machine learning tasks. For\n**monopoles** the library offers classes for Particle-Particle Particle-Mesh Ewald\n(``P3M``), Particle Mesh Ewald (``PME``), standard ``Ewald``, and non-periodic methods.\nThe library has the flexibility to calculate potentials beyond :math:`1/r`\nelectrostatics, including arbitrary order :math:`1/r^p` potentials. For **dipolar**\ninteraction we offer to calculate the :math:`1/r^3` potential using the standard\n``Ewald`` method.\n\nOptimized for both CPU and GPU devices, ``torch-pme`` is fully `TorchScriptable`_,\nallowing it to be converted into a format that runs independently of Python, such as in\nC++, making it ideal for high-performance production environments.\n\nWe also provide an experimental implementation for *JAX* in `jax-pme`_.\n\n.. _`TorchScriptable`: https://pytorch.org/docs/stable/jit.html\n.. _`jax-pme`: https://github.com/lab-cosmo/jax-pme\n\n.. marker-documentation\n\nDocumentation\n-------------\n\nFor details, tutorials, and examples, please have a look at our `documentation`_.\n\n.. _`documentation`: https://lab-cosmo.github.io/torch-pme\n\n.. marker-installation\n\nInstallation\n------------\n\nYou can install *torch-pme* using pip with\n\n.. code-block:: bash\n\n    pip install torch-pme\n\nor conda\n\n.. code-block:: bash\n\n    conda install -c conda-forge torch-pme\n\nand ``import torchpme`` to use it in your projects!\n\nWe also provide bindings to `metatensor <https://docs.metatensor.org>`_ which can\noptionally be installed together and used as ``torchpme.metatensor`` via\n\n.. code-block:: bash\n\n    pip install torch-pme[metatensor]\n\n.. marker-quickstart\n\nQuickstart\n----------\n\nHere is a simple example to get started with *torch-pme*:\n\n.. code-block:: python\n\n   >>> import torch\n   >>> import torchpme\n\n   >>> # Single charge in a cubic box\n   >>> positions = torch.zeros((1, 3), requires_grad=True)\n   >>> cell = 8 * torch.eye(3)\n   >>> charges = torch.tensor([[1.0]])\n\n   >>> # No neighbors for a single atom; use `vesin` for neighbors if needed\n   >>> neighbor_indices = torch.zeros((0, 2), dtype=torch.int64)\n   >>> neighbor_distances = torch.zeros((0,))\n\n   >>> # Tune P3M parameters\n   >>> smearing, p3m_parameters, _ = torchpme.tuning.tune_p3m(\n   ...    charges=charges,\n   ...    cell=cell,\n   ...    positions=positions,\n   ...    cutoff=5.0,\n   ...    neighbor_indices=neighbor_indices,\n   ...    neighbor_distances=neighbor_distances,\n   ... )\n\n   >>> # Initialize potential and calculator\n   >>> potential = torchpme.CoulombPotential(smearing)\n   >>> calculator = torchpme.P3MCalculator(potential, **p3m_parameters)\n\n   >>> # Compute (per-atom) potentials\n   >>> potentials = calculator.forward(\n   ...    charges=charges,\n   ...    cell=cell,\n   ...    positions=positions,\n   ...    neighbor_indices=neighbor_indices,\n   ...    neighbor_distances=neighbor_distances,\n   ... )\n\n   >>> # Calculate total energy and forces\n   >>> energy = torch.sum(charges * potentials)\n   >>> energy.backward()\n   >>> forces = -positions.grad\n\nFor more examples and details, please refer to the `documentation`_.\n\n.. marker-issues\n\nHaving problems or ideas?\n-------------------------\n\nHaving a problem with *torch-pme*? Please let us know by `submitting an issue\n<https://github.com/lab-cosmo/torch-pme/issues>`_.\n\nSubmit new features or bug fixes through a `pull request\n<https://github.com/lab-cosmo/torch-pme/pulls>`_.\n\n.. marker-cite\n\nReference\n---------\n\nIf you use *torch-pme* for your work, please read and cite our preprint available on\n`arXiv`_.\n\n.. code-block::\n\n   @article{loche_fast_2024,\n      title = {Fast and Flexible Range-Separated Models for Atomistic Machine Learning},\n      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},\n      year = {2024},\n      month = dec,\n      number = {arXiv:2412.03281},\n      eprint = {2412.03281},\n      primaryclass = {physics},\n      publisher = {arXiv},\n      doi = {10.48550/arXiv.2412.03281},\n      urldate = {2024-12-05},\n      archiveprefix = {arXiv}\n      }\n\n.. _`arXiv`: http://arxiv.org/abs/2412.03281\n\n.. marker-contributing\n\nContributors\n------------\n\nThanks goes to all people that make *torch-pme* possible:\n\n.. image:: https://contrib.rocks/image?repo=lab-cosmo/torch-pme\n   :target: https://github.com/lab-cosmo/torch-pme/graphs/contributors\n\n.. |tests| image:: https://github.com/lab-cosmo/torch-pme/workflows/Tests/badge.svg\n   :alt: Github Actions Tests Job Status\n   :target: https://github.com/lab-cosmo/torch-pme/actions?query=branch%3Amain\n\n.. |codecov| image:: https://codecov.io/gh/lab-cosmo/torch-pme/graph/badge.svg?token=srVKRy7r6m\n   :alt: Code coverage\n   :target: https://codecov.io/gh/lab-cosmo/torch-pme\n\n.. |docs| image:: https://img.shields.io/badge/documentation-latest-sucess\n   :alt: Documentation\n   :target: `documentation`_\n",
    "bugtrack_url": null,
    "license": "BSD-3-Clause",
    "summary": "Particle-mesh based calculations of long-range interactions in PyTorch",
    "version": "0.3.0",
    "project_urls": {
        "changelog": "https://lab-cosmo.github.io/torch-pme/latest/references/changelog.html",
        "documentation": "https://lab-cosmo.github.io/torch-pme",
        "homepage": "https://lab-cosmo.github.io/torch-pme",
        "issues": "https://github.com/lab-cosmo/torch-pme/issues",
        "repository": "https://github.com/lab-cosmo/torch-pme"
    },
    "split_keywords": [
        "auto-differentiation",
        " computational materials science",
        " electrostatics",
        " fourier transform",
        " gpu acceleration",
        " high-performance computing",
        " machine learning",
        " molecular dynamics",
        " particle-mesh ewald",
        " pytorch",
        " scientific computing",
        " torchscript"
    ],
    "urls": [
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "a1d09fa5d99d06c53be429dff2461d4230899b5ae78bb841cea84dbc8b5d96f4",
                "md5": "b7d257916432234197e264d499356774",
                "sha256": "fddcfa3900dab7ed3ed430b28301d1e31c3665885ed196495c94bec9a8a30d8f"
            },
            "downloads": -1,
            "filename": "torch_pme-0.3.0-py3-none-any.whl",
            "has_sig": false,
            "md5_digest": "b7d257916432234197e264d499356774",
            "packagetype": "bdist_wheel",
            "python_version": "py3",
            "requires_python": ">=3.9",
            "size": 64127,
            "upload_time": "2025-02-21T17:09:05",
            "upload_time_iso_8601": "2025-02-21T17:09:05.160120Z",
            "url": "https://files.pythonhosted.org/packages/a1/d0/9fa5d99d06c53be429dff2461d4230899b5ae78bb841cea84dbc8b5d96f4/torch_pme-0.3.0-py3-none-any.whl",
            "yanked": false,
            "yanked_reason": null
        },
        {
            "comment_text": null,
            "digests": {
                "blake2b_256": "c634a414288eccff5352e84271f2b09dfe6f9b18ff2da1b6f012dc35908d4a67",
                "md5": "5f7e5edd06ac85a4d3ab717b53080393",
                "sha256": "8ee9b387e6830c1a6d982873a49abcecdb7c3038a41d927a80784fb1fecb2385"
            },
            "downloads": -1,
            "filename": "torch_pme-0.3.0.tar.gz",
            "has_sig": false,
            "md5_digest": "5f7e5edd06ac85a4d3ab717b53080393",
            "packagetype": "sdist",
            "python_version": "source",
            "requires_python": ">=3.9",
            "size": 45063,
            "upload_time": "2025-02-21T17:09:07",
            "upload_time_iso_8601": "2025-02-21T17:09:07.492139Z",
            "url": "https://files.pythonhosted.org/packages/c6/34/a414288eccff5352e84271f2b09dfe6f9b18ff2da1b6f012dc35908d4a67/torch_pme-0.3.0.tar.gz",
            "yanked": false,
            "yanked_reason": null
        }
    ],
    "upload_time": "2025-02-21 17:09:07",
    "github": true,
    "gitlab": false,
    "bitbucket": false,
    "codeberg": false,
    "github_user": "lab-cosmo",
    "github_project": "torch-pme",
    "travis_ci": false,
    "coveralls": false,
    "github_actions": true,
    "tox": true,
    "lcname": "torch-pme"
}
        
Elapsed time: 0.43150s