Welcome to mlcg’s documentation!

This is the base code to create models such as the ones described in [TransCGSchnet] and [CGSchnet].

Docs badge License pipeline status

This repository collects a set of tools to apply machine learning techniques to coarse grain atomic systems.

Installation

Requires Python 3.12. Clone the repo:

git clone git@github.com:ClementiGroup/mlcg.git
cd mlcg

We recommend uv for installation.

uv sync creates a .venv in the repo root and installs all dependencies from the correct pre-built wheel pages. If you manage your own virtual environment instead, install reproducibly from one of the provided per-platform lock files with uv pip. Pick the row matching your hardware:

Hardware

uv sync

Self-managed environment

NVIDIA CUDA 13.0

uv sync --extra cu130

uv pip install -r pylock.cu130.toml

NVIDIA CUDA 12.8

uv sync --extra cu128

uv pip install -r pylock.cu128.toml

AMD ROCm 7.2

uv sync --extra rocm72

uv pip install -r pylock.rocm72.toml

CPU only

uv sync --extra cpu

uv pip install -r pylock.cpu.toml

In a self-managed environment you can also skip the lock file and install the package directly with the extra matching your hardware:

uv pip install ".[cu130]"

For developers

Add --group dev to install additional development dependencies (black, pytest, coverage):

uv sync --extra cu130 --group dev

or, for a self-managed environment:

uv pip install ".[cu130]" --group dev

With pip

Installing with pip instead of uv is also possible, but requires installing dependencies in a specific order to avoid version conflicts between optional model backends — see `Installation with pip`_ below.

Examples

Please take a look into the examples folder of the repository to see how to use this code to train a model over an existing dataset.

CLI

The models defined in this library can be conveniently trained using the pytorch-lightning CLI utilities.

Contents

Indices and tables