Crystal Generation using Diffusion Models

Open In Colab

Crystal Generation using Diffusion Models#

In this tutorial, we’ll explore how to generate inorganic crystal structures using Chemeleon, a diffusion-based generative model for materials discovery.

Whereas autoregressive models like CrystaLLM generate structures one token at a time, diffusion models work by iteratively denoising a noisy latent representation to arrive at a valid crystal structure.

We will:

  • Understand the intuition behind diffusion-based generation

  • Use a pretrained Chemeleon model to generate new materials

  • Visualize and interpret the generated crystal structures

Reference: Chemeleon GitHub

Diffusion Models: A Brief Overview#

Diffusion models are a class of generative models inspired by physical processes of noise and denoising.

Forward process: Gradually add noise to a structure until it becomes pure noise.

Reverse process (learned): Train a neural network to remove noise step-by-step and reconstruct the original data.

Advantages:

  • High-quality, diverse generations

  • Better suited for continuous and spatially structured data (like crystal coordinates)

In Chemeleon, the diffusion model operates on a latent representation of crystal structures derived from atom types and fractional coordinates.

What is Chemeleon?#

Chemeleon is a graph-based diffusion model for crystal structure generation. It uses:

  • A crystal graph representation with node features (elements) and position embeddings

  • A denoising network trained on the Materials Project dataset

  • A score-based diffusion process in latent space to generate valid structures

Key components:

  • Graph Neural Networks (GNNs) for encoding

  • Denoising Score Matching for training

  • Support for conditional generation (e.g. given composition)

Paper: ChemRxiv

execute next cell only if you run on google colab

import locale
locale.getpreferredencoding = lambda: "UTF-8"

! pip install janus-core[all] data-tutorials chemeleon ase
get_ipython().kernel.do_shutdown(restart=True)

Load up the required functions#

from chemeleon import Chemeleon
from chemeleon.visualize import Visualizer
from ase.io import write
import os

Load the model#

We load up a pre-trained model. We start with the composition mode. This model only takes compositions as a prompt. Note that the download and loading takes some time (2 - 5 minutes).

%%time
composition_model = Chemeleon.load_composition_model()

Generate Structures#

Here we generate just one sample, to run quickly. But you can increase this later. Since we are using the composition only model, it can only take elements as a prompt.

# Set parameters
n_samples = 2
n_atoms = 8
prompt = "Li P S"

%%time
# Generate crystal structures
atoms_list = composition_model.sample(prompt, n_atoms, n_samples)

Output Interpretation#

The output of Chemeleon is a valid 3D crystal structure, represented as a pymatgen Structure object.

Each structure includes:

  • Lattice parameters

  • Atomic coordinates

  • Element types

You can export to CIF using:

structure.to(filename="generated_structure.cif")

# Visualise
visualizer = Visualizer(atoms_list)
visualizer.view(index=0)

from ase.io import write
write("generated_structures.extxyz",atoms_list)
write("generated_structure_0.cif",atoms_list[0])

Diffusion vs. Autoregressive Generation#

Feature

Diffusion (Chemeleon)

Autoregressive (CrystaLLM)

Generation style

Denoising from noise

Token-by-token sampling

Data representation

Continuous 3D + elements

Discrete sequence

Conditioning support

Via latent / graph inputs

Via token prompts

Output format

pymatgen Structure

CIF / token sequence

Typical advantages

Spatial realism, diversity

Interpretability, simplicity

Question for students:

  • Which approach seems more flexible for structure-property targeting?

  • Can you imagine combining both approaches in a hybrid model?

Try generation with more conditions#

We can use the general_text_model which allows us to use natural language to impose more conditions.

# Load default model checkpoint (general text types)
chemeleon = Chemeleon.load_general_text_model()

# Set parameters
n_samples = 3
n_atoms = 56
text_inputs = "A crystal structure of LiMn2O4 with cubic symmetry"

# Generate crystal structure
atoms_list = chemeleon.sample(text_inputs, n_atoms, n_samples)

# Visualize the generated crystal structure
visualizer = Visualizer(atoms_list)
visualizer.view(index=0)

View the generation trajectory#

We can visualise the diffusion Langevin dynamcis that is used to generate the final structure from the initial sampling. To achive this we return the trajectory from the genertation process.

In the visulaisation note how the composition, the positions and the lattice are updated at the same time.

n_samples = 1
n_atoms = 56
text_inputs = "A crystal structure of LiMn2O4 with cubic symmetry"

# Generate crystal structure with trajectory
trajectory = chemeleon.sample(text_inputs, n_atoms, n_samples, return_trajectory=True)

# Visualize the trajectory
idx = 0
traj_0 = [t[idx] for t in trajectory][::10] + [trajectory[-1][idx]]
visualizer = Visualizer(traj_0, resolution=15)
visualizer.view_trajectory(duration=0.1)

Exercise#

Use Chemeleon to generate some of the materials that you generated previously with CrystaLLM. Then export these structures and compare the formation energies of the two different generative models.

Contents