Metadata-Version: 2.1
Name: graph-pes
Version: 0.0.0
Summary: Potential Energy Surfaces on Graphs
Author-email: John Gardner <gardner.john97@gmail.com>
License: MIT License
        
        Copyright (c) 2023 John Gardner
        
        Permission is hereby granted, free of charge, to any person obtaining a copy
        of this software and associated documentation files (the "Software"), to deal
        in the Software without restriction, including without limitation the rights
        to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
        copies of the Software, and to permit persons to whom the Software is
        furnished to do so, subject to the following conditions:
        
        The above copyright notice and this permission notice shall be included in all
        copies or substantial portions of the Software.
        
        THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
        IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
        FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
        AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
        LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
        OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
        SOFTWARE.
        
Project-URL: Homepage, https://github.com/jla-gardner/graph-pes
Classifier: License :: OSI Approved :: MIT License
Classifier: Programming Language :: Python
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.8
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: torch-geometric
Requires-Dist: pytorch-lightning
Requires-Dist: ase
Requires-Dist: numpy
Requires-Dist: jaxtyping
Provides-Extra: test
Requires-Dist: pytest; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Provides-Extra: dev
Requires-Dist: notebook; extra == "dev"
Requires-Dist: ruff; extra == "dev"
Requires-Dist: bumpver; extra == "dev"
Provides-Extra: doc
Requires-Dist: sphinx; extra == "doc"
Requires-Dist: furo; extra == "doc"
Requires-Dist: nbsphinx; extra == "doc"
Requires-Dist: sphinx-autobuild; extra == "doc"
Requires-Dist: sphinxext-opengraph; extra == "doc"
Requires-Dist: sphinx-copybutton; extra == "doc"
Provides-Extra: publish
Requires-Dist: build; extra == "publish"
Requires-Dist: twine; extra == "publish"

# `graph-pes` - Potential Energy Surfaces on Graphs

`graph-pes` is a Python framework for training Potential Energy Surface (PES) 
models that operate on graph representations of atomic structures.
Under-the-hood, this relies on Pytorch/Geometric for efficient tensor operations.

Batteries are included:
- **easy data manipulations** : see docs relating to `AtomicGraph`, intuitive batching, easy conversion from ase etc.
- **easy construction of PES models** : implement `predict_local_energies` , easy to save, load and share
- **useful primitives** : PerSpeciesParameter
- **easy training** : forces, energies, well conditioned losses etc.

## Installation

```bash
pip install graph-pes
```

## Minimal example

```python
from graph_pes import GraphPESModel
from graph_pes.data import convert_to_graphs, AtomicGraph
from graph_pes.training import train_model
from ase.io import read
import torch

# 1. define a model
class LennardJones(GraphPES):
    def __init__(self):
        super().__init__()
        self.sigma = nn.Parameter(torch.tensor(1.0))
        self.epsilon = nn.Parameter(torch.tensor(1.0))

    def predict_local_energies(self, graph: AtomicGraph) -> torch.Tensor:
        central_atoms, neighbours = graph.neighbour_index
        distances = graph.neighbour_distances
        
        # calculate pairwise interactions
        x = (self.sigma / distances)**6
        pairwise_interaction = 4 * self.epsilon * (x**2 - x)

        # sum over neighbours to get per-atom contributions
        return torch.scatter_add(
            torch.zeros(graph.n_atoms, device=graph.device),
            central_atoms,
            pairwise_interaction
        )

# 2. load some structures
structures = read('structures.xyz', index=':10')

# 3. convert to graphs (e.g. using a radius cutoff)
graphs = convert_to_graphs(structures, cutoff=5.0)

# 4. train the model
model = LennardJones()
train_model(model, graphs, max_epochs=100)
```
