CrystAF — Crystal AnyFlow

Few-step, all-atom 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), then post-trains that flow map with a PoseBusters-ranked NFT objective. One 16-LoRA adapter serves NFE 8 / 16 / 32 / 50 — you change only the evaluation time grid, never the weights.

Code, environment setup, and every eval script: https://github.com/HaCTang/Crystal-NFT

Results — no stereochemistry correction

Plain model quality. 200 CSD validation families × 20 samples, summary.paper_bootstrap, L1 EMD PDD, interval flow-map sampler, no inference-time correctors and no cell calibration.

	NFE	rho	PB % ↑	clash % ↓	Vol.Err ↓	EMD PDD ↓
Clari-M backbone (Heun)	16	—	57.05	10.65	1.85	10.36
+ Clari PB-NFT rank800 = distillation teacher (Heun)	16	—	75.38	9.96	2.00	10.29
CrystAF distilled (cont3), uniform grid	16	—	77.65	14.71	2.07	10.59
CrystAF distilled, report grid	8	0.30	83.67	30.73	2.46	11.54
CrystAF distilled	16	0.75	85.49	13.85	2.09	10.92
CrystAF distilled	50	1	89.44	10.88	1.86	10.47
+ Stage-2 MeanFlowNFT	16	0.75	85.71	11.29	1.85	10.45
+ Stage-2 MeanFlowNFT	32	1	93.21	7.86	1.61	10.07
+ Stage-2 MeanFlowNFT	50	1	93.84	7.37	1.70	10.22

Reference, 1000 families × 20 × 50 Heun steps (wider protocol, don't subtract row-wise): Clari-M 88.43 PB / 8.57 clash, Clari-L 86.88 / 6.92 (published Clari-L: 85.89 / 7.69 / 1.50 / 9.28).

Reproducibility caveat. A second seed of the identical recipe reproduces PB ≈ 93 (93.08 vs 93.21) but not clash < 8 or Vol < 1.7 — it bottoms out at 8.60 / 1.88. The 7.86 / 1.61 above is about 1.1 SE better than the other seed's best, i.e. the favourable tail. The defensible summary is PB ≈ 93, clash ≈ 8.6–8.9, Vol.Err ≈ 1.7–2.0. Both seeds are published here so you can check this yourself.

Results — with stereochemistry correction

The correctors are training-free and isometric, applied at sampling time: CRYSTAF_MIRROR_FIX=body CRYSTAF_STEREO_REFLECT=1 CRYSTAF_MMFF=1 CRYSTAF_RELAX_CLASH=1, plus CRYSTAF_VOL_SCALE, a lattice-only cell calibration that moves no atom.

	NFE	PB % ↑	clash % ↓	Vol.Err ↓	EMD PDD ↓	stereo % ↑
baseline, no correction	16	85.49	13.85	2.09	10.92	49.91
cont3 + correctors	16	92.16	2.27	2.06	11.03	95.24
cont3 + correctors	50	92.63	2.40	1.78	10.57	95.59
Stage-2 + correctors	32	93.73	2.58	1.64	10.42	95.26
Stage-2 + correctors + VOL_SCALE=0.9850	32	93.58	2.12	1.51	10.47	95.20
CRYSTAF_PCFM=rs (max chirality, no relaxation)	16	70.13	14.14	2.06	10.79	99.97

Stereochemistry is not learned: the backbone's atom/bond features are identical for the two enantiomers, so the base model sits at chance (49.9%) on genuine R/S centres, and training-time conditioning is a measured negative result (chirality stayed at 49.21 while PB fell 11.3 points). The 95% comes entirely from the sampling-time correctors.

Vol.Err 1.51 is a dispersion floor — two different calibration factors (0.9850, 0.9925) both land on it, so the residual is spread, not bias. EMD PDD is the one column still short of Clari-L (10.47 vs 9.28); UMA relaxation was measured to do nothing for PDD (12.322 → 12.329), so that column needs a better base model.

Files

File	Use
crystaf-nft-mfpure-epoch12.pt	The report checkpoint (Stage-2 MeanFlowNFT). Reproduces the bold rows above.
crystaf-cont3-step2000.pt	Distilled, pre-post-training baseline.
crystaf-nft-seed2-epoch12.pt	Second seed of the same recipe — for the reproducibility caveat.
teacher-rank800-pbnft-epoch1.pt	rank800 PB-NFT teacher (LoRA merged). Only needed to re-run distillation.

Each file holds net_state_dict + ema_state_dict + meta; evaluate with the EMA weights. The Clari-M backbone is not redistributed here — fetch it from the-matter-lab/clari; you need it to build the DiT.

from huggingface_hub import hf_hub_download

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

Evaluation also needs CSD-derived tensors, which are CCDC-licensed and cannot be redistributed — build them yourself with scripts/build_clari_csd.sh (see doc/env.md).

Four things that silently produce wrong numbers

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 16, 1 at 32/50. NFE 8 at rho=0.75 scores 48 instead of 83.67.
3. Report summary.paper_bootstrap.pb_score_pct, not raw mean_pb_score — invalid crystals naively score 1 and inflate it. Check mean_pb_valid ≈ 0.695.
4. Two stereo keys. mean_stereo_agreement_pct includes molecules with no R/S centre (79 here); mean_stereo_agreement_defined_pct counts only defined centres (95). Every stereo number above is the _defined one. Reading the other key produces an apparent 16-point regression that does not exist.

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, and 60-family grids cannot rank configurations at all.

Caveat on the clash and volume columns

With the correctors on, clash and Vol.Err each have a corrector aimed directly at them, so they are no longer independent checks of packing quality — the clash relaxation optimises exactly what clash_rate measures. The bare-model table above is the one to read for packing quality. EMD PDD stays independent in both tables.

License

CC-BY-NC-4.0, inherited from the Clari model weights these derive 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.