Skip to content
BranchesTags

Repository files navigation

XxaCT-NN Reproduction

Reproduction of XxaCT-NN: Structure Agnostic Multimodal Learning for Materials Science (arXiv: 2507.01054).

XxaCT-NN predicts materials properties from elemental composition + simulated XRD patterns, without requiring crystal structure as input.

Architecture

Architecture

  • XRD Encoder: 17 tokens × 250 dims (4250-dim pattern reshaped), Transformer encoder
  • Composition Encoder: CrabNet — element identity + stoichiometric fraction → embeddings
  • Fusion Module: Transformer decoder (comp → K/V, XRD → Q)
  • Prediction Heads: regression (formation energy, band gap), classification (crystal system, space group)

Model Size vs. Paper

Component Paper This Implementation
XRD Encoder 3 layers, dim 512, 4 heads, FFN 1024 2 layers, dim 128, 2 heads, FFN 256
Fusion Module 12 layers, dim 768, 12 heads, FFN 2048 2 layers, dim 128, 2 heads, FFN 256
CrabNet 9.7M params same
Batch size 32 768 (8× H100) 64

Setup

Requires Python 3.12+ and uv.

# CPU (local dev)
uv sync --extra cpu --extra dev

# GPU (training)
uv sync --extra cuda --extra dev

# Experiment 1 (Materials Project data)
uv sync --extra cpu --extra dev --extra phase0

# Experiment 2 (Alexandria / HuggingFace datasets)
uv sync --extra cpu --extra dev --extra phase1

Activate the environment:

source .venv/bin/activate

Experiments

Experiment 1 — Train from Scratch on Materials Project

Train the model end-to-end on ~5 000 Materials Project structures, no pretraining.

Aspect Full Paper This Experiment
Dataset Alexandria 5M Materials Project ~5K
XRD Encoder 3 layers, dim 512, 4 heads 2 layers, dim 128, 2 heads
Fusion Module 12 layers, dim 768, 12 heads 2 layers, dim 128, 2 heads
Pretraining MXM + Contrastive None
Targets Ef, crystal system, band gap Ef only

Results

Test MAE  : 0.1654 eV/atom

1. Install dependencies

uv sync --extra cpu --extra dev --extra phase0   # CPU
# or
uv sync --extra cuda --extra dev --extra phase0  # GPU
source .venv/bin/activate

2. Run tests

pytest

All tests should pass before proceeding.

3. Prepare data

Requires a free Materials Project API key.

export MP_API_KEY=your_key_here
python scripts/prepare_data.py --n 5000 --data-dir data/

# On a multi-CPU machine, speed up XRD preprocessing with --workers:
python scripts/prepare_data.py --n 5000 --data-dir data/ --workers 32

This downloads ~5 000 structures and pre-computes XRD patterns to data/xrd/. XRD preprocessing runs in parallel across all available CPUs by default (--workers defaults to os.cpu_count()).

4. Train

python scripts/train_minimal.py \
    --data-dir data/ \
    --checkpoint-dir checkpoints/ \
    --epochs 50 \
    --batch-size 64 \
    --lr 5e-4 \
    --test-fraction 0.1 \
    --val-fraction 0.1 \
    --device cpu          # or cuda

Training prints train_MAE and val_MAE each epoch. The checkpoint is saved to checkpoints/model_phase0.pt and contains the model weights, all config values, and the full MAE history.

Key options:

Flag Default Description
--epochs 50 Number of training epochs
--batch-size 64 Samples per batch
--lr 5e-4 Peak learning rate (AdamW)
--test-fraction 0.1 Fraction held out for evaluation (never trained on)
--val-fraction 0.1 Fraction used for validation during training
--embed-dim 128 Hidden dimension for all transformer blocks
--device auto cpu or cuda

5. Evaluate on the held-out test set

python scripts/evaluate.py \
    --checkpoint checkpoints/model_phase0.pt \
    --data-dir data/ \
    --device cpu

The script reconstructs the exact test split from the checkpoint's saved config (same seed, val_fraction, test_fraction), then reports:

Test MAE  : 0.XXXX eV/atom
Baseline  : 0.XXXX eV/atom  (mean-prediction)
Beats baseline? YES

Loading a checkpoint in Python

import torch
from xxact_nn.models.xxact_nn import XxaCTNN

ckpt = torch.load("checkpoints/model_phase0.pt", map_location="cpu")
cfg = ckpt["config"]

model = XxaCTNN.from_config(
    embed_dim=cfg["embed_dim"],
    num_heads=cfg["num_heads"],
    num_layers=cfg["num_layers"],
    ffn_dim=cfg["ffn_dim"],
    dropout=cfg["dropout"],
)
model.load_state_dict(ckpt["model_state_dict"])
model.eval()

# Training history
print(ckpt["results"]["val_mae_history"])

Experiment 2 — Pretrain on Alexandria + Fine-tune on Materials Project

Self-supervised pretraining (MXM + Contrastive) on a 100K Alexandria subset, then fine-tune on the Materials Project dataset.

Results

Test MAE  : 0.0832 eV/atom
Improvement over Experiment 1: 0.1654 → 0.0832 eV/atom (−50%)

Step A — Alexandria Data Pipeline

1. Install dependencies
uv sync --extra cpu --extra dev --extra phase1   # CPU
# or
uv sync --extra cuda --extra dev --extra phase1  # GPU
source .venv/bin/activate
2. Run tests
pytest tests/data/test_alexandria.py
3. Prepare Alexandria cache

Pre-compute the XRD cache (downloads data automatically from the Alexandria server):

# Full 100K subset (GPU server, parallelized across CPUs)
python scripts/prepare_alexandria.py \
    --url https://alexandria.icams.rub.de/data/pbe/2024.12.15/ \
    --cache-dir data/alexandria_cache \
    --max-samples 100000 \
    --workers 16

# Resume an interrupted run — already-cached entries are skipped automatically
python scripts/prepare_alexandria.py \
    --url https://alexandria.icams.rub.de/data/pbe/2024.12.15/ \
    --cache-dir data/alexandria_cache \
    --max-samples 100000

# Use a manually downloaded local file instead
python scripts/prepare_alexandria.py \
    --source data/alexandria_000.json.bz2 \
    --cache-dir data/alexandria_cache \
    --max-samples 100000

# Local smoke test — no download, no Pymatgen, finishes in seconds
python scripts/prepare_alexandria.py \
    --synthetic 20 \
    --cache-dir /tmp/alex_test

All flags:

Flag Default Description
--url URL — Base URL to stream chunks from (mutually exclusive with --source/--synthetic)
--source PATH — Path to a local .json.bz2 chunk file (mutually exclusive with --url/--synthetic)
--synthetic N — Generate N fake entries locally, no download needed (mutually exclusive with --url/--source)
--cache-dir required Output directory for .npy files and metadata.csv
--max-samples 100000 Number of entries to process (real data only)
--seed 42 Shuffle seed for deterministic subset selection
--workers 4 Parallel XRD workers, real data only (ProcessPoolExecutor)

The script is resumable: already-cached .npy files are skipped on re-run. It produces:

  • data/alexandria_cache/{entry_id}.npy — XRD tensor (17, 250) per compound
  • data/alexandria_cache/metadata.csv — entry_id, formula, ef, crystal_system
4. Use the Dataset
from xxact_nn.data.alexandria import AlexandriaDataset
from torch.utils.data import DataLoader

train_ds = AlexandriaDataset("data/alexandria_cache", split="train")
val_ds   = AlexandriaDataset("data/alexandria_cache", split="val")

loader = DataLoader(train_ds, batch_size=256, shuffle=True)
batch = next(iter(loader))
# batch["xrd"]           → (B, 17, 250) float32
# batch["ef"]            → list of B floats  (formation energy, eV/atom)
# batch["crystal_system"] → list of B ints   (0=triclinic … 6=cubic)
# batch["formula"]       → list of B strings

Step B+C — Pretraining (MXM + Contrastive)

Run tests first:

.venv/bin/pytest tests/pretraining/ -q

Then run pretraining on the Alexandria cache:

python scripts/pretrain.py \
    --cache-dir data/alexandria_cache/ \
    --loss combined \
    --epochs 20 \
    --batch-size 256 \
    --lr 1e-4 \
    --device cuda          # or cpu

Saves the checkpoint to checkpoints/pretrained.pt.

Resuming pretraining

To continue training from an existing checkpoint (optimizer and scheduler state are preserved):

python scripts/pretrain.py \
    --resume checkpoints/pretrained.pt \
    --cache-dir data/alexandria_cache/ \
    --loss combined \
    --epochs 20 \
    --batch-size 256 \
    --lr 1e-4 \
    --device cuda

The epoch counter continues from where it left off. The updated checkpoint is saved back to checkpoints/pretrained.pt, overwriting the previous one — copy it first if you want to keep it.

Key options:

Flag Default Description
--cache-dir required Alexandria XRD cache dir (output of prepare_alexandria.py)
--resume None Path to an existing checkpoint to resume from
--checkpoint-dir checkpoints/ Where to write pretrained.pt
--loss combined mxm / contrastive / combined (recommended)
--epochs 20 Additional epochs to train
--batch-size 256 Samples per batch
--lr 1e-4 Learning rate (AdamW)
--mask-ratio 0.05 Fraction of XRD tokens masked per sample (MXM loss); paper uses 0.05
--embed-dim 128 Hidden dimension
--num-workers 4 DataLoader worker processes
--max-samples None Limit dataset size (for smoke tests)
--device auto cpu or cuda

Smoke test (CPU, 1 epoch):

python scripts/pretrain.py \
    --cache-dir data/alexandria_cache/ \
    --max-samples 64 \
    --epochs 1 \
    --device cpu

Step D — Fine-tuning on Materials Project

Run tests first:

.venv/bin/pytest tests/scripts/test_finetune.py -q

Fine-tune the pretrained checkpoint on the Materials Project dataset:

python scripts/finetune.py \
    --pretrained-checkpoint checkpoints/pretrained.pt \
    --data-dir data/ \
    --checkpoint-dir checkpoints/ \
    --epochs 50 \
    --batch-size 64 \
    --lr 5e-4 \
    --device cuda          # or cpu

The script:

  • Loads architecture config from the pretrained checkpoint
  • Initialises model weights from the checkpoint (strict=False — the regression head uses random init if absent)
  • Fine-tunes on the Materials Project Ef regression task
  • Reports train / val / test MAE and whether it beats the Experiment 1 baseline
  • Saves the fine-tuned checkpoint to checkpoints/model_phase1.pt

Key options:

Flag Default Description
--pretrained-checkpoint required Path to pretrained .pt checkpoint
--data-dir data/ Directory containing mp_raw.json and xrd/
--checkpoint-dir checkpoints/ Where to write model_phase1.pt
--epochs 50 Fine-tuning epochs
--batch-size 64 Samples per batch
--lr 5e-4 Peak learning rate (AdamW)
--val-fraction 0.1 Fraction used for validation
--test-fraction 0.1 Fraction held out for evaluation
--max-samples None Limit dataset size (for smoke tests)
--device auto cpu or cuda

Smoke test (CPU, 1 epoch):

python scripts/finetune.py \
    --pretrained-checkpoint checkpoints/pretrained.pt \
    --data-dir data/ \
    --max-samples 64 \
    --epochs 1 \
    --device cpu

Results Summary

Experiment Test MAE (eV/atom)
1 — Train from scratch (MP 5K) 0.1654
2 — Pretrain Alexandria + Fine-tune 0.0832

Project Structure

xxact_nn/
├── data/
│   ├── alexandria.py        # AlexandriaDataset + load_alexandria_subset (Experiment 2)
│   ├── collate.py           # MPCollator: formula strings → element/fraction tensors
│   ├── dataset.py           # MPDataset (train/val/test split) + make_dataloaders
│   ├── mp_loader.py         # Materials Project API download
│   └── xrd_utils.py         # XRD simulation, smearing, normalization
├── models/
│   ├── composition_encoder.py
│   ├── xrd_encoder.py
│   ├── fusion.py
│   ├── heads.py
│   └── xxact_nn.py          # Full model assembly
├── training/
│   ├── config.py            # TrainingConfig dataclass
│   ├── scheduler.py         # Linear warmup + cosine decay
│   └── trainer.py           # Training loop (AdamW, AMP-aware, MAE logging)
└── scripts/
    ├── prepare_data.py         # Download MP subset + pre-compute XRD (Experiment 1)
    ├── prepare_alexandria.py   # Download Alexandria subset + pre-compute XRD (Experiment 2)
    ├── train_minimal.py        # Experiment 1 training entry point
    ├── pretrain.py             # Experiment 2 pretraining on Alexandria (MXM + contrastive)
    ├── finetune.py             # Experiment 2 fine-tuning on Materials Project
    └── evaluate.py             # Evaluate on held-out test split

License

MIT

About

Unofficial implementation of XxaCT-NN model (arXiv:2507.01054)

Topics

machine-learningmaterial-sciencexrd

Resources

Readme
MIT license
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages

Generated from aerubanov/runpod-exp-setup