Table of Contents
ASC (Atomic Structure Classification, pronunced ask) is a high-performance framework designed to accelerate the development and benchmarking of Graph Neural Networks for atomistic simulations.
By decoupling the data engineering from the model design, this library allows researchers to ignore the data pipeline barrier and focus entirely on model innovation. An OVITO extension is also provided to allow researchers to quickly load trained models directly into OVITO to perform real-time inference and assess performances visually.
PyTorch Lightning - a PyTorch wrapper made to avoid boilerplate code and improve reproductibility, allowing fast iterations.
PyTorch Geometric - a flexible library build upon PyTorch to easily handle Graph Neural Networks.
- Python 3.11 or newer
We recommend using uv to install the library dependencies:
# Clone the repo
git clone https://github.com/akwarii/asc.git
cd asc
# If uv is not installed on your machine
pip install uv
# Install the virtual environment and dependencies
uv sync
# Add extra options if you want to download datasets and/or use hyperparameter optimization
uv sync --group api --group hpo
uv pip install pyg-lib -f https://data.pyg.org/whl/nightly/torch-${TORCH}+${CUDA}.html
# Activate the virtual environment
source .venv/bin/activatewhere
- ${TORCH} should be replaced by either 2.8.0, 2.9.0, or 2.10.0
- ${CUDA} should be replaced by either cpu, cu126, cu128, cu129, or cu130
Alternatively, you can use pip directly:
# Clone the repo
git clone https://github.com/akwarii/asc.git
cd asc
# (OPTIONAL) Create a conda environment
conda create -n asc python=3.11
conda activate asc
# Install the requirements
pip install -r requirements.txt
pip install pyg-lib -f https://data.pyg.org/whl/nightly/torch-${TORCH}+${CUDA}.htmlMaterials Project: If you intend to use the Materials Project dataset, get your API key here
Crystal Space Group: To use our preprocessed dataset build upon Materials Project, AFLOW and GNoME databases, you need a Kaggle account.
- Create an account on the Kaggle website
- Go to your User profile and click on
Create New Token - Move the downloaded
kaggle.jsonfile to the$HOME/.kagglefolder (if needed, create it withmkdir $HOME/.kaggle)
Once you got your credentials, create a .env file (using cp .env.example .env for example) and enter your credentials. Remember to not share this file with others.
To train the PaiNN model on the Si (silicon) dataset, you can run:
python main.py fit --config configs/main.yaml --trainer configs/trainer/default.yaml --model configs/model/painn.yaml --data configs/data/custom.yamlThe provided configuration file will train a model in the same way as what is presented in examples/silicon/silicon.ipynb.
main.py script can also be used to run validation/test steps as well as inference. The available subcommands are fit, validate, test, and predict. More details can be obtained by running
python main.py --helpThe CLI can give you more information about a specific argument by running:
python main.py fit --model.model.help src.models.PaiNN
python main.py fit --data.dataset.helpThe first will print the help message related to our implementation of the PaiNN model, while the other will display the arguments that can be passed to all our datasets.
Code and Framework:
- Config file for each dataset and model
- Document how to extend the code
- Node-level datamodule
- Improve documentation
- Switch between different CSG dataset version
- Make real tests
In order to install the dependencies required for building the documentation, run the following command:
uv sync --group docYou can then build the documentation by running:
make docsFrom there, you can open the docs/_build/html/index.html file in your browser to view the documentation.
Follow the generic coding conventions defined in PEP8.
Run pre-commit before submitting a PR by running make format.
If you use asc in your research, please consider citing the following work:
If you have any questions, please contact one of the contributors below:
If you found a bug or want to request a new feature, please create a new GitHub Issues