XQueryer Lightweight

Overview

This project is a lightweight improved version of XQueryer, the intelligent crystal structure identifier for powder X-ray diffraction. It keeps the 100315-class identification task and changes the data path from pre-imported XRD arrays to dynamic simulation from crystal structures stored in MP500.db. All patterns are aligned to 3500 points over 10-90 degrees.

Data

Download MP500.db from the GitHub release of this project and place it at data/MP500.db. It is an ASE database with 100315 structures. ASE row ids start at 1; model labels are entry_id - 1, so labels are 0..100314.

The split is XRD-level. Train, validation, and test all cover the full set of 100315 structures, but each split uses a different simulation seed range so the same structure appears with different simulated parameter combinations.

Dynamic XRD Simulation

For each structure, the dataset samples physical and instrumental parameters and calls Pysimxrd.generator.parser(database, entry_id, ...). By default, train simulates 20 XRD patterns per structure, validation simulates 1, and test simulates 2.

grainsize: 10-100
prefect_orientation: -0.3 to 0.3 for both values
thermo_vibration: 0-2
zero_shift: 0-2
dis_detector2sample: 300-600
half_height_slit_detector: 3-8
half_height_sample: 1-4
xrd: reciprocal
background_order: 6
background_ratio: 0.02-0.2
mixture_noise_ratio: 0.01-0.1

Model

The lightweight model keeps the same conceptual framework as XQueryer: FFT processing, CNN feature extraction, Cross-Attention, and a final classification head. The main change is that the very large element-to-sequence expansion is replaced by a small set of element-conditioned queries.

1. FFT branch

Each normalized XRD pattern is transformed with real FFT. Several low-pass frequency views are reconstructed with inverse FFT and concatenated with the original signal. This keeps the original XQueryer idea of frequency-aware peak denoising and improves robustness to background, broadening, and high-frequency noise.

2. CNN peak encoder

The four-channel signal is passed through residual depthwise-separable 1D convolution blocks. Large kernels capture peak width and local shifts, while strided blocks progressively reduce sequence length. This gives a lighter local feature extractor than standard wide convolutions while preserving peak-shape information.

3. Residual Cross-Attention

CNN features are pooled into a fixed number of XRD tokens. A 92-dimensional element-presence vector is encoded into several structure queries. These queries attend to the XRD tokens through multi-head Cross-Attention, so chemical composition guides which diffraction tokens matter most. This replaces the original expensive element -> 768 x 3500 mapping with compact learned queries. Each attention layer uses residual cross-attention, query self-attention, and an FFN refinement block, giving the queries multiple chances to match peak tokens while remaining much smaller than the original expansion.

4. Classification head

The attended queries are fused with CNN global mean/max features and FFT global statistics. The fused representation is passed through an MLP neck and a normalized cosine classifier over 100315 MP entries. This is different from a plain linear head: normalized features encourage better angular separation between many similar structures. The default model uses base_channels=64, attn_dim=192, num_heads=6, num_tokens=96, num_queries=4, and num_attn_layers=2, giving about 28.0M trainable parameters. Most parameters are in the unavoidable 100315-class final layer.

Differences from original XQueryer

The main retained ideas are FFT enhancement, CNN peak features, element-conditioned Cross-Attention, and direct MP-entry classification. The main changes are dynamic Pysimxrd training data, 10-90 degree alignment, compact element queries instead of huge element sequence expansion, residual attention refinement, CNN/FFT feature fusion, and a cosine retrieval-style classifier.

Training

pip install torch ase scipy tqdm Pysimxrd

torchrun --nproc_per_node=4 trainer.py \
  --db_path data/MP500.db \
  --epochs 100 \
  --batch_size 64 \
  --num_workers 8 \
  --simulations_per_entry 20 \
  --output_dir outputs/lightweight

The trainer saves checkpoint_0000.pth, latest.pth, per-epoch checkpoints, and best.pth. Checkpoints include EMA weights and inference uses EMA by default when available.

Architecture knobs for cluster experiments are --base_channels, --attn_dim, --num_heads, --num_tokens, and --num_queries. Increasing them may improve accuracy at higher memory cost; reducing them gives a smaller model for ablation.

Training improvements include warmup-cosine learning-rate scheduling, AdamW, gradient clipping, label smoothing, top-1/top-5/top-10 reporting, optional weak XRD mixup, and EMA validation/checkpointing.

Inference

Experimental files must be CSV files with angle and intensity columns. The script interpolates each pattern to 3500 points over 10 to 90 degrees.

python inference.py \
  --checkpoint outputs/lightweight/checkpoints/best.pth \
  --inputs "exp_data/*.csv" \
  --topk 5

Citation

@article{cao2025xqueryer,
  title={XQueryer: an intelligent crystal structure identifier for powder X-ray diffraction},
  author={Cao, Bin and Zheng, Zinan and Liu, Yang and Zhang, Longhan and Wong, Lawrence WY and Weng, Lu-Tao and Li, Jia and Li, Haoxiang and Zhang, Tong-Yi},
  journal={National Science Review},
  volume={12},
  number={12},
  pages={nwaf421},
  year={2025},
  publisher={Oxford University Press}
}