CrystAF — Crystal AnyFlow

Few-step, all-atom, stereochemistry-aware molecular crystal structure generation.

CrystAF distills a 50-step Clari crystal generator into a dual-time flow map U(z, r, t) that jumps z += (t - r) · U(z, r, t). One 16-LoRA adapter serves NFE 8 / 16 / 50 — you change only the evaluation time grid, never the weights. Two training-free sampling-time correctors then fix stereochemistry and steric clashes.

Results

200 CSD validation families × 20 samples, paper_bootstrap metrics, with the stereochemistry + clash recipe enabled:

NFE	rho	PB ↑	Clash ↓	Vol.Err ↓	EMD PDD ↓	stereo(defined) ↑
8	0.30	83.30	2.11	2.46	11.71	84.4
16	0.75	85.28	1.22	2.06	10.80	86.0
50	1.0	88.47	0.77	1.87	10.40	86.3

Without the correctors (plain CrystAF): clash 30.73 / 13.85 / 10.88 and stereo ≈ 50% (chance) at NFE 8 / 16 / 50 respectively. Reference Clari-M at 50 Heun steps: PB 88.43, clash 8.57.

stereo(defined) counts only genuine R/S centres, so 50% is chance.

Files

File	Use
crystaf-cont3-step2000.pt	The report checkpoint. Reproduces every number above. Contains net_state_dict + ema_state_dict (evaluate with the EMA weights).
teacher-rank800-pbnft-epoch1.pt	rank800 PB-NFT teacher (LoRA already merged). Only needed to re-run distillation.

The Clari-M backbone is not redistributed here — fetch it from the-matter-lab/clari.

Usage

Code, environment setup and evaluation scripts: https://github.com/HaCTang/Crystal-NFT

from huggingface_hub import hf_hub_download

student = hf_hub_download("Haocheng1/CrystAF", "crystaf-cont3-step2000.pt")
backbone = hf_hub_download("the-matter-lab/clari", "clari-med.ckpt")

Two things that will silently give wrong numbers if you get them wrong:

1. Sample with the interval flow map, not Clari's Heun sampler — MEANFLOW_SAMPLER_MODE=interval. Heun drops the second time argument.
2. Use the reported rho per NFE (t_i = (i/N)^rho): 0.30 at NFE 8, 0.75 at NFE 16, 1.0 at NFE 50. NFE 8 at rho=0.75 scores 48% instead of 84%.

Enable the stereochemistry + clash correctors with CRYSTAF_MIRROR_FIX=body CRYSTAF_RELAX_CLASH=1. Both are isometric transforms, so pb_score and volume_error are bit-identical with and without them (verified paired, 0/160 crystals changed).

Report PB as summary.paper_bootstrap.pb_score_pct, not raw mean_pb_score — invalid crystals naively score 1 and inflate it. Judge changes on all four Table 1 columns, not PB alone. Repeat runs of an identical config vary by roughly ±0.5 PB / ±0.6 clash / ±0.3 PDD.

Caveat on the clash number

The clash corrector directly optimises the quantity clash_rate measures, driving it to ~1%. That is legitimate sampling-time constraint enforcement, but it means clash_rate is no longer an independent check of packing quality for these configurations. EMD PDD stays independent, and it is unchanged within run noise. Report both together.

License

CC-BY-NC-4.0, inherited from the Clari model weights these are derived from (Clari's code is MIT; its weights are CC-BY-NC-4.0). Non-commercial use only; please credit the-matter-lab/clari upstream.