Data preparation¶

This tutorial prepares the SPICE (MACE-OFF23) and MAD datasets for the examples in Training configurations. The scripts read energies and forces with ASE, store them under the reference keys used by FlashCart, and write extended XYZ files. They preserve the source units: eV for energies, Å for positions, and eV/Å for forces. No unit conversion is applied.

By default, FlashCart reads total reference energies from atoms.info["REF_energy"] and forces from atoms.arrays["REF_forces"]. Set energy_key and forces_key to use other names.

These examples use scripts from the cloned FlashCart repository. See Installation for installation instructions. Run the commands from the repository root, with wget and tar available for downloading and extracting the data.

SPICE (MACE-OFF23)¶

Download the training and test archives from the Cambridge Apollo repository and extract them into data/spice:

mkdir -p data/spice
(
    cd data/spice
    wget --user-agent="Mozilla/5.0" -O train.tar.gz \
        "https://www.repository.cam.ac.uk/bitstreams/b185b5ab-91cf-489a-9302-63bfac42824a/download"
    wget --user-agent="Mozilla/5.0" -O test.tar.gz \
        "https://www.repository.cam.ac.uk/bitstreams/cb8351dd-f09c-413f-921c-67a702a7f0c5/download"
    tar xzf train.tar.gz
    tar xzf test.tar.gz
)

The preparation script expects train_large_neut_no_bad_clean.xyz and test_large_neut_all.xyz directly in data/spice. The commands in parentheses run in that directory and leave the current shell in the repository root.

The script divides the supplied training set randomly into 95% training and 5% validation data. --seed controls the split. Different seeds produce separately sampled splits, which can overlap. The supplied test set is shared by all seeds.

Prepare the three splits used by the SPICE model configurations:

for seed in 0 1 2
do
    python examples/datasets/prepare_spice.py --dir data/spice --seed "$seed"
done

Each call writes train_<seed>.extxyz and val_<seed>.extxyz in data/spice. It also writes test.extxyz and groups the test structures by config_type into files under test_subsets. These test files are rewritten from the same supplied test set on each call.

SPICE preparation script

"""Prepare the SPICE (MACE-OFF23) training and test sets.

From the repository root, download and extract the two archives from Cambridge
Apollo (https://www.repository.cam.ac.uk/items/d50227cd-194f-4ba4-aeb7-2643a69f025f):

    mkdir -p data/spice
    (
        cd data/spice
        wget --user-agent="Mozilla/5.0" -O train.tar.gz \
            "https://www.repository.cam.ac.uk/bitstreams/b185b5ab-91cf-489a-9302-63bfac42824a/download"
        wget --user-agent="Mozilla/5.0" -O test.tar.gz \
            "https://www.repository.cam.ac.uk/bitstreams/cb8351dd-f09c-413f-921c-67a702a7f0c5/download"
        tar xzf train.tar.gz
        tar xzf test.tar.gz
    )

Then run from the repository root:

    python examples/datasets/prepare_spice.py --dir data/spice --seed 0

This writes train_<seed>.extxyz and val_<seed>.extxyz (a random 95/5 split of
the training set), test.extxyz, and one test_subsets/<config_type>.extxyz per
test subset, all under data/spice. Different seeds produce separately sampled
splits, which can overlap.
"""

import argparse
from collections import defaultdict
from pathlib import Path
from random import Random

from ase.io import read, write

VALID_FRACTION = 0.05


def to_reference(frames):
    """Move the calculator's energy and forces to the REF_energy / REF_forces keys FlashCart reads."""
    for atoms in frames:
        atoms.info["REF_energy"] = atoms.get_potential_energy()
        atoms.arrays["REF_forces"] = atoms.get_forces()
        atoms.calc = None
    return frames


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--dir", type=Path, default=Path("."), help="directory with the unpacked .xyz files")
    parser.add_argument("--seed", type=int, default=0, help="seed of the train/validation split")
    args = parser.parse_args()

    train = to_reference(read(args.dir / "train_large_neut_no_bad_clean.xyz", ":"))
    n_valid = round(len(train) * VALID_FRACTION)
    valid_idx = set(Random(args.seed).sample(range(len(train)), n_valid))
    write(args.dir / f"train_{args.seed}.extxyz", [a for i, a in enumerate(train) if i not in valid_idx])
    write(args.dir / f"val_{args.seed}.extxyz", [a for i, a in enumerate(train) if i in valid_idx])
    print(f"seed {args.seed}: {len(train) - n_valid} train / {n_valid} validation structures")

    test = to_reference(read(args.dir / "test_large_neut_all.xyz", ":"))
    write(args.dir / "test.extxyz", test)
    subsets = defaultdict(list)
    for atoms in test:
        subsets[atoms.info["config_type"]].append(atoms)
    (args.dir / "test_subsets").mkdir(exist_ok=True)
    for name, frames in sorted(subsets.items()):
        write(args.dir / "test_subsets" / f"{name.replace(' ', '_')}.extxyz", frames)
        print(f"test subset {name}: {len(frames)} structures")


if __name__ == "__main__":
    main()

MAD¶

Download the supplied training, validation, and test splits from Materials Cloud, and the benchmark structures from the PET-MAD repository:

mkdir -p data/mad
for split in train val test
do
    wget -O "data/mad/mad-$split.xyz" \
        "https://archive.materialscloud.org/records/c4ene-0mv14/files/mad-$split.xyz?download=1"
done
wget -O data/mad/mad-bench.xyz \
    "https://huggingface.co/lab-cosmo/pet-mad/resolve/main/benchmarks/mad-test-mad-settings.xyz"

The preparation script expects mad-train.xyz, mad-val.xyz, mad-test.xyz, and mad-bench.xyz in data/mad. It preserves the supplied splits and groups the MAD test set by subset and the PET-MAD benchmark by dataset. Structures without the relevant key are assigned to an unlabelled group.

Run the preparation script:

python examples/datasets/prepare_mad.py --dir data/mad

This writes mad-train.extxyz, mad-val.extxyz, mad-test.extxyz, and mad-bench.extxyz in data/mad, together with the grouped files in data/mad/subsets. The reference energies already have the isolated-atom baseline subtracted. Keep that reference when training with the supplied MAD configuration.

MAD preparation script

"""Prepare the MAD dataset and the PET-MAD benchmark.

From the repository root, download the official splits from Materials Cloud
(https://archive.materialscloud.org/records/c4ene-0mv14) and the 1562-structure
benchmark shipped with PET-MAD (Table 1 of arXiv:2503.14118):

    mkdir -p data/mad
    for split in train val test
    do
        wget -O "data/mad/mad-$split.xyz" \
            "https://archive.materialscloud.org/records/c4ene-0mv14/files/mad-$split.xyz?download=1"
    done
    wget -O data/mad/mad-bench.xyz \
        "https://huggingface.co/lab-cosmo/pet-mad/resolve/main/benchmarks/mad-test-mad-settings.xyz"

Then run from the repository root:

    python examples/datasets/prepare_mad.py --dir data/mad

This writes mad-{train,val,test,bench}.extxyz and one
subsets/<file>-<group>.extxyz per MAD test subset and per benchmark source
dataset, all under data/mad. The reference energies already have the
isolated-atom baseline subtracted. The training configuration starts from
zero prior offsets and fits the remaining per-element energy shifts.
"""

import argparse
from collections import defaultdict
from pathlib import Path

from ase.io import read, write

GROUP_KEYS = {"mad-test": "subset", "mad-bench": "dataset"}


def to_reference(frames):
    """Move the calculator's energy and forces to the REF_energy / REF_forces keys FlashCart reads."""
    for atoms in frames:
        atoms.info["REF_energy"] = atoms.get_potential_energy()
        atoms.arrays["REF_forces"] = atoms.get_forces()
        atoms.calc = None
    return frames


def main():
    parser = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    parser.add_argument("--dir", type=Path, default=Path("."), help="directory with the downloaded .xyz files")
    args = parser.parse_args()

    (args.dir / "subsets").mkdir(exist_ok=True)
    for name in ("mad-train", "mad-val", "mad-test", "mad-bench"):
        frames = to_reference(read(args.dir / f"{name}.xyz", ":"))
        write(args.dir / f"{name}.extxyz", frames)
        print(f"{name}: {len(frames)} structures")
        if name not in GROUP_KEYS:
            continue
        groups = defaultdict(list)
        for atoms in frames:
            groups[str(atoms.info.get(GROUP_KEYS[name], "unlabelled")).replace("/", "_")].append(atoms)
        for group, members in sorted(groups.items()):
            write(args.dir / "subsets" / f"{name}-{group}.extxyz", members)
            print(f"  {group}: {len(members)} structures")


if __name__ == "__main__":
    main()