# -*- coding: utf-8 -*-
from setuptools import setup

packages = \
['faenet']

package_data = \
{'': ['*']}

install_requires = \
['mendeleev>=0.12', 'torch>=1.12']

setup_kwargs = {
    'name': 'faenet',
    'version': '0.1.3',
    'description': "PyTorch implementation for FAENet from 'FAENet: Frame Averaging Equivariant GNN for Materials Modeling'",
    'long_description': '<p align="center">\n<strong><a href="https://github.com/vict0rsch/faenet" target="_blank">💻&nbsp;&nbsp;Code</a></strong>\n<strong>&nbsp;&nbsp;•&nbsp;&nbsp;</strong>\n<strong><a href="https://faenet.readthedocs.io/" target="_blank">Docs&nbsp;&nbsp;📑</a></strong>\n</p>\n\n<p align="center">\n    <a>\n\t    <img src=\'https://img.shields.io/badge/python-3.8%2B-blue\' alt=\'Python\' />\n\t</a>\n\t<a href=\'https://faenet.readthedocs.io/en/latest/?badge=latest\'>\n    \t<img src=\'https://readthedocs.org/projects/faenet/badge/?version=latest\' alt=\'Documentation Status\' />\n\t</a>\n    <a href="https://github.com/psf/black">\n\t    <img src=\'https://img.shields.io/badge/code%20style-black-black\' />\n\t</a>\n<a href="https://pytorch.org">\n<img src="https://img.shields.io/badge/PyTorch-%23EE4C2C.svg?logo=PyTorch&logoColor=white"/>\n</a>\n</p>\n<br/>\n\n# FAENet: Frame Averaging Equivariant GNN for Materials modeling\n\n\nThis repository contains an implementation of the paper [*FAENet: Frame Averaging Equivariant GNN for Materials modeling*](https://arxiv.org/pdf/2305.05577.pdf), accepted at ICML 2023. More precisely, you will find:\n\n* `FrameAveraging`: the transform that projects your pytorch-geometric data into a canonical space of all euclidean transformations, as defined in the paper.\n* `FAENet`: a GNN architecture for material modeling.  \n* `model_forward`: a high-level forward function that computes appropriate equivariant model predictions for the Frame Averaging method, i.e. handling the different frames and mapping to equivariant predictions.\n\nAlso: https://github.com/vict0rsch/faenet\n\n## Installation\n\n```\npip install faenet\n```\n\n⚠️ The above installation requires `Python >= 3.8`, [`torch > 1.11`](https://pytorch.org/get-started/locally/), [`torch_geometric > 2.1`](https://pytorch-geometric.readthedocs.io/en/latest/notes/installation.html#) to the best of our knowledge. Both `mendeleev` and `pandas` package are also required to derive physics-aware atom embeddings in FAENet.\n\n## Getting started\n\n### Frame Averaging Transform\n\n`FrameAveraging` is a Transform method applicable to pytorch-geometric `Data` object, which shall be used in the `get_item()` function of your `Dataset` class. This method derives a new canonical position for the atomic graph, identical for all euclidean symmetries, and stores it under the data attribute `fa_pos`. You can choose among several options for the frame averaging, ranging from *Full FA* to *Stochastic FA* (in 2D or 3D) including traditional data augmentation *DA* with rotated samples. See the full [doc](https://faenet.readthedocs.io/en/latest/autoapi/faenet/transforms/index.html#faenet.transforms.FrameAveraging) for more details. Note that, although this transform is specific to pytorch-geometric data objects, it can be easily extended to new settings since the core functions `frame_averaging_2D()` and `frame_averaging_3D()` generalise to other data format. \n\n```python\nimport torch\nfrom faenet.transforms import FrameAveraging\n\nframe_averaging = "3D"  # symmetry preservation method used: {"3D", "2D", "DA", ""}:\nfa_method = "stochastic"  # the frame averaging method: {"det", "all", "se3-stochastic", "se3-det", "se3-all", ""}:\ntransform = FrameAveraging(frame_averaging, fa_method)\ntransform(data)  # transform the PyG graph data\n```\n\n### Model forward for Frame Averaging\n\n`model_forward()` aggregates the predictions of a chosen ML model (e.g FAENet) when Frame Averaging is applied, as stipulated by the Equation (1) of the paper. INded, applying the model on canonical positions (`fa_pos`) directly would not yield equivariant predictions. This method must be applied at training and inference time to compute all model predictions. It requires `batch` to have pos, batch and frame averaging attributes (see [docu](https://faenet.readthedocs.io/en/latest/autoapi/faenet/fa_forward/index.html)). \n\n```python\nfrom faenet.fa_forward import model_forward\n\npreds = model_forward(\n    batch=batch,   # batch from, dataloader\n    model=model,  # FAENet(**kwargs)\n    frame_averaging="3D", # ["2D", "3D", "DA", ""]\n    mode="train",  # for training \n    crystal_task=True,  # for crystals, with pbc conditions\n)\n```\n\n### FAENet GNN \n\nImplementation of the FAENet GNN model, compatible with any dataset or transform. In short, FAENet is a very simple, scalable and expressive model. Since does not explicitly preserve data symmetries, it has the ability to process directly and unrestrictedly atom relative positions, which is very efficient and powerful. Although it was specifically designed to be applied with Frame Averaging above, to preserve symmetries without any design restrictions, note that it can also be applied without. When applied with Frame Averaging, we need to use the `model_forward()` function above to compute model predictions, `model(data)` is not enough. Note that the training procedure is not given here, you should refer to the original [github repository](https://github.com/RolnickLab/ocp). Check the [documentation](https://faenet.readthedocs.io/en/latest/autoapi/faenet/model/index.html) to see all input parameters. \n\nNote that the model assumes input data (e.g.`batch` below) to have certain attributes, like atomic_numbers, batch, pos or edge_index. If your data does not have these attributes, you can apply custom pre-processing functions, taking `pbc_preprocess` or `base_preprocess` in [utils.py](https://faenet.readthedocs.io/en/latest/autoapi/faenet/utils/index.html) as inspiration. You simply need to pass them as argument to FAENet (`preprocess`).\n\n```python\nfrom faenet.model import FAENet\n\npreds = FAENet(**kwargs)\nmodel(batch)\n```\n\n![FAENet architecture](https://github.com/vict0rsch/faenet/blob/main/examples/data/faenet-archi.png)\n\n### Eval \n\nThe `eval_model_symmetries()` function helps you evaluate the equivariant, invariant and other properties of a model, as we did in the paper. \n\nNote: you can predict any atom-level or graph-level property, although the code explicitly refers to energy and forces.\n\n### Tests\n\nThe `/tests` folder contains several useful unit-tests. Feel free to have a look at them to explore how the model can be used. For more advanced examples, please refer to the full [repository](https://github.com/RolnickLab/ocp) used in our ICML paper to make predictions on OC20 IS2RE, S2EF, QM9 and QM7-X dataset. \n\nThis requires [`poetry`](https://python-poetry.org/docs/). Make sure to have `torch` and `torch_geometric` installed in your environment before you can run the tests. Unfortunately because of CUDA/torch compatibilities, neither `torch` nor `torch_geometric` are part of the explicit dependencies and must be installed independently.\n\n```bash\ngit clone git@github.com:vict0rsch/faenet.git\npoetry install --with dev\npytest --cov=faenet --cov-report term-missing\n```\n\nTesting on Macs you may encounter a [Library Not Loaded Error](https://github.com/pyg-team/pytorch_geometric/issues/6530)\n\n### Contact\n\nAuthors: Alexandre Duval (alexandre.duval@mila.quebec) and Victor Schmidt (schmidtv@mila.quebec). We welcome your questions and feedback via email or GitHub Issues.\n\n',
    'author': 'Victor Schmidt',
    'author_email': 'vsch@pm.me',
    'maintainer': 'None',
    'maintainer_email': 'None',
    'url': 'None',
    'packages': packages,
    'package_data': package_data,
    'install_requires': install_requires,
    'python_requires': '>=3.8',
}


setup(**setup_kwargs)
