Your first model

This tutorial trains, continues and evaluates a compact MANDALA model on the two silicon calculations bundled with the repository: the ideal eight-atom diamond cell for training and a thermally perturbed cell for validation. It runs on a laptop CPU.

The configuration

# A compact MANDALA model for the bundled silicon calculations.
#
# Flat fields configure the default E3GNN (encoders, attention message
# passing, split per-pair readout heads). The data and training sections
# select the snapshots and the run settings.
cutoff_radius: 8.0
l_max: 2
hidden_irreps: 32x0e+8x0o+16x1o+8x1e+8x2e+4x2o+4x3o+4x3e+2x4e
n_radial: 16
radial_layers: [32]
num_layers_gnn: 2
node_update_message_agg: attention
node_update_attention_scalar_dim: 32
neck_depth: 1
head_e3mlp_layers: 1
matrix_targets: [hamiltonian, overlap, density]
loss_l1_fraction: 1.0
lr: 0.003
lr_scheduler_patience: 40
lr_scheduler_target: val/loss_total
revert_on_spike: false
checkpoint_monitor: val/loss_total
log_per_irrep_metrics: false
log_per_pair_loss_metrics: false
log_hamiltonian_irrep_contrib_metrics: false
log_hamiltonian_pair_contrib_metrics: false
data:
  representation: orbital_blocks
  splits:
    train: [data/small/Si]
    val: [data/small/Si_perturbed_scale1_078]
training:
  seed: 7
  dtype: float32
  max_epochs: 600

The flat fields select the default network (The MANDALA E3GNN) with two attention message-passing layers and features up to \(\ell=4\) — enough for the d–d couplings of the s2p2d1 basis (Sparse operators and their irreps). All three operators are predicted jointly and trained with an element-wise MAE (loss_l1_fraction: 1.0).

Check it

PYTHONPATH=src:. python -m scripts.train --config examples/first_model/silicon.yaml \
  --output runs/silicon --accelerator cpu --dry-run

The dry run parses both calculations, builds the 8 Å periodic graphs, builds the model, takes one optimization step and prints the parameter count, the loss and the split sizes — without writing anything.

Train

PYTHONPATH=src:. python -m scripts.train --config examples/first_model/silicon.yaml \
  --output runs/silicon --accelerator cpu

Training reads both calculations once (about two minutes for the 50 MB text dumps) and then takes one to two seconds per epoch on an eight-core CPU. The objective is the sum of the element-wise MAEs of H (Hartree), S and D:

Epoch

Training loss

Validation loss

1

0.0327

0.0250

100

0.00937

0.0126

200

0.00481

0.0114

400

0.00296

0.0111

600

0.00290

0.0111

With a single training structure the model fits it closely while the validation error levels off: generalization to new geometries comes from training on many structures, which is exactly what the paper models do.

The run directory contains run_manifest.json (resolved configuration, input hashes, environment), split_manifest.json, latest_checkpoint.pt and best_model.pt.

Continue

An interrupted run resumes exactly where it stopped — same optimizer state, learning-rate schedule, random streams and best-model ranking. Repeat the original command with --resume:

PYTHONPATH=src:. python -m scripts.train --config examples/first_model/silicon.yaml \
  --output runs/silicon --accelerator cpu --resume runs/silicon/latest_checkpoint.pt

Evaluate

PYTHONPATH=src:. python examples/quickstart/evaluate.py \
  --checkpoint runs/silicon/best_model.pt \
  --structure data/small/Si_perturbed_scale1_078 --plot bands.png

compared blocks: 696
hamiltonian  MAE   84.1019 meV  max 3717.2480 meV
overlap      MAE    0.0066   max    0.1580
density      MAE    0.0011   max    0.0392
band energy Tr(D_ref H) per atom: reference -16.4240 eV, predicted -15.3572 eV

The script compares every block that the model graph and the reference calculation share, evaluates the band energy with the reference density through the sparse trace \(\sum\operatorname{tr}(H_{ij}^{\bm L}D_{ji}^{-\bm L})\), and with --plot draws the bands of the predicted Hamiltonian, combined with the reference overlap, on top of the reference bands.

Going further