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.