Optimized Crystallographic Graph Generation for Material Science

TL;DR

This paper introduces pyMatGraph, a GPU-optimized framework for real-time generation of graph representations of crystalline materials, facilitating efficient training of graph neural networks for material discovery.

Contribution

It presents a novel GPU-based tool for generating cutoff and k-nearest-neighbors graphs of periodic structures during neural network training.

Findings

Enables real-time graph updates during training.

Supports efficient processing of periodic structures.

Open-source implementation available at GitHub.

Abstract

Graph neural networks are widely used in machine learning applied to chemistry, and in particular for material science discovery. For crystalline materials, however, generating graph-based representation from geometrical information for neural networks is not a trivial task. The periodicity of crystalline needs efficient implementations to be processed in real-time under a massively parallel environment. With the aim of training graph-based generative models of new material discovery, we propose an efficient tool to generate cutoff graphs and k-nearest-neighbours graphs of periodic structures within GPU optimization. We provide pyMatGraph a Pytorch-compatible framework to generate graphs in real-time during the training of neural network architecture. Our tool can update a graph of a structure, making generative models able to update the geometry and process the updated graph during the forward propagation on the GPU side. Our code is publicly available at https://github.com/aklipf/mat-graph.