XtalPaint – A framework for crystal structure inpainting based on diffusion models

Welcome to the XtalPaint Documentation.

Overview

XtalPaint is a Python package that provides tools to perform crystal structure inpainting, i.e. adding atomic sites to a given host structure, using score-based diffusion models. Here, we provide retrained versions of the Mattergen architecture and the building blocks to set up the inpainting workflows. The initial application in our latest work: Score-based diffusion models for accurate crystal-structure inpainting and reconstruction of hydrogen positions, focuses on adding missing hydrogen sites to inorganic crystal structures, but the framework can be adapted to other inpainting tasks as well, i.e. general crystal structure prediction based on given host structures (see other interesting works in the field, e.g. by Zhong et al.).

Features

• Inpainting pipeline for crystal structures
• Integration with AiiDA workflow management
• Support for various relaxation methods
• Evaluation metrics for inpainting quality

Getting Started

Check out the examples on to run the inpainting pipeline:

• With AiiDA integration
• Without AiiDA integration

Installation

git clone https://github.com/psi-lms/XtalPaint.git
cd XtalPaint/

uv pip install .

This will install the default version. If you want to use it in combination with AiiDA, please also install the optional dependencies:

uv pip install .[aiida]

Model checkpoints for retrained versions of MatterGen

Model checkpoints for the retrained versions of MatterGen used in our work can be downloaded from Hugging Face. Currently, the repository contains the pos-only and TD-pos-only models discussed in the paper.

Acknowledgements

This project is developed to perform crystal structure inpainting, currently on top of Microsoft's MatterGen. Some parts of the codebase are adapted from MatterGen's implementation (as highlighted in the respective files).