Inpainting crystal structure generations with score-based denoising

Searching for the optimal atomic position of additive atoms in a given host structure is crucial in designing materials with intercalation chemistry for energy storage applications. In this study, we present an application of the SE(3)-equivariant diffusion model for such conditional crystal structure predictions using inpainting methods. The model, built upon the e3nn framework, was pre-trained on the Materials Project structure database via denoising score matching. By solving the reverse stochastic differential equation using the predictor-corrector approach, the model is capable of de novo crystal generation as well as conditional generation -- finding atomic sites of additive atoms within a host structure. We benchmarked the model performance on the WBM dataset and showcased examples of ion intercalation in different MnO$_2$ polymorphs. This efficient, probabilistic site-finding tool offers the potential for accelerating the materials discovery.