deepmodeling / CrystalFormer

A Foundation Model for Crystal Structure Generation and Prediction Generating Cs 2 ZnFe(CN) 6 Crystal ( mp-570545 ) Overview _CrystalFormer_ is a unified autoregressive transformer model for inorganic crystalline material generation that supports **both** de novo crystal generation (DNG) and crystal structure prediction (CSP) within a *single probabilistic framework*. It is specifically designed for space group-controlled generation of crystalline materials. The space group symmetry significantly simplifies the crystal space, which is crucial for data and compute efficient generative modeling of crystalline materials. The model can: • **De novo generation** $p(C|\varnothing)$: Generate plausible crystal structures from scratch, without any formula constraint. • **Crystal structure prediction** $p(C|f)$: Generate crystal structures conditioned on a given chemical formula $f$. We provide a pretrained checkpoint to support both functionalities. No architectural change is required — _CrystalFormer_ seamlessly switches behavior depending on whether a formula is supplied. Furthermore, the performance of both functionalities can be boosted via reinforcement fine-tuning. Contents • Contents • Model Card • Status • Get Started • Installation • CPU installation • CUDA (GPU) installation • Install required packages and command line tools • Available Weights • Crystal Structure Prediction • Sample • Relax • Energy Above Hull (Ehull) • End-to-End Pipeline • Model Context Protocol (MCP) Server • De Novo Generation • Sample • Evaluate • Advanced Usage • Reinforcement Fine-tuning • Writing custom reward functions • Pretrain • How to cite Model Card _CrystalFormer_ is an autoregressive transformer for the probability distribution of crystal structures: • **De novo generation**: $P(C|\varnothing) = P(g) P(W_1|...) P(A_1|...) P(X_1|...) ... P(L|...)$ • **Formula-conditioned prediction**: $P(C|f) = P(g|f) P(W_1|...) P(A_1|...) P(X_1|...) ... P(L|...)$ where the crystal structure $C$ is represented by the sequence $g-(W_{i}-A_{i}-X_{i})_{n}-L$: • $f$: chemical formula, e.g. Cu 12 Sb 4 S 13 • $g$: space group number 1-230 • $W$: Wyckoff letter ('a', 'b', ...,'A') • $A$: atom type ('H', 'He', ..., 'Og') in the chemical formula • $X$: fractional coordinates • $L$: lattice vector [a, b, c, alpha, beta, gamma] • $P(W_i| ...)$ and $P(A_i| ...)$ are categorical distributions. • $P(X_i| ...)$ is the mixture of von Mises distribution. • $P(L| ...)$ is the mixture of Gaussian distribution. We only consider symmetry inequivalent atoms in the crystal representation. The remaining atoms are restored based on the information of space group and Wyckoff letters. There is a natural alphabetical ordering for the Wyckoff letters, starting with 'a' for a position with the site-symmetry group of maximal order and ending with the highest letter for the general position. The sampling procedure starts from higher symmetry sites (with smaller multiplicities) and then goes on to lower symmetry ones (with larger multiplicities). Only for the cases where the Wyckoff letter cannot fully determine the structure, one needs to further consider fractional coordinates in the loss or sampling. Status Major milestones are summarized below. • v0.6: _CrystalFormer_ for unified de novo generation and crystal structure prediction. • v0.5: Initial release of _CrystalFormer-CSP_ for crystal structure prediction. • v0.4.2 : Add implementation of direct preference optimization. • v0.4.1 : Replace the absolute positional embedding with the Rotary Positional Embedding (RoPE). • v0.4 : Add reinforcement learning (proximal policy optimization). • v0.3 : Add conditional generation in the plug-and-play manner. • v0.2 : Add Markov chain Monte Carlo (MCMC) sampling for template-based structure generation. • v0.1 : Initial implementations of crystalline material generation conditioned on the space group. Get Started **Notebooks**: The quickest way to get started with _CrystalFormer_ is our notebooks in the Google Colab platform: • ColabCSP : Running _CrystalFormer-CSP_ Seamlessly on Google Colab • CrystalFormer-RL : Reinforcement fine-tuning for materials design Previous notebooks (only for reference but not actively maintained): • CrystalFormer Quickstart : GUI notebook demonstrating the conditional generation of crystalline materials with _CrystalFormer_ • CrystalFormer Application : Generating stable crystals with a given structure prototype. This workflow can be applied to tasks that are dominated by element substitution Installation Create a new environment and install the required packages, we recommend using python and conda to create the environment: Before installing the required packages, you need to install and first. CPU installation CUDA (GPU) installation If you intend to use CUDA (GPU) to speed up the training, it is important to install the appropriate version of and . It is recommended to check the jax docs for the installation guide. The basic installation command is given below: Install required packages and command line tools After installing and , you need to install the package: During installation, the command line tools in the cli directory will be automatically installed. Available Weights We release the weights of the model trained on the Alex20s dataset. More details are available in the model card. Crystal Structure Prediction Thinking fast and slow for crystal structure prediction Sample • : the optimizer to use, means no training, only sampling • : the path to the model weights • : the top-K number of space groups will be sampled uniformly. • : the number of samples to generate • : the chemical formula • : [Optional] the path to save the generated structures, if not provided, the structures will be saved in the folder. Instead of providing for top-K sampling, you may directly provide your favorite space group number • : the space group number [1-230] The sampled structure will be saved in the file…