v-baturin/mattergen-mp40

MatterGen MP-40 full checkpoint

This is an unconditional MatterGen crystal-generation model trained on an MP-40 dataset containing periodic structures with 1–40 atoms per unit cell. The selected checkpoint is epoch 894 (loss_val=0.57).

The checkpoint is complete, not inference-only. It contains model weights, optimizer and learning-rate-scheduler states, Lightning callback state, epoch, global step, and the resolved training configuration. It can therefore be used for inference or to resume training.

Files

File	Purpose
last.ckpt	Full PyTorch Lightning checkpoint, including optimizer state
mp40_num_atoms_distribution.json	Empirical MP-40 training-set prior for sampling 1–40 atoms
sample_mp40.py	Self-contained inference helper that registers the prior and generates structures
mp_40_data_module.yaml	Portable Hydra data-module configuration for training resume

Checkpoint SHA-256:

72a140e838404540a590453e74b8b621329dbf8e5c6c873b69926df173e0fc40

Only load a PyTorch checkpoint obtained from a trusted source.

Requirements

Install the MatterGen code before using this bundle. The training codebase was the v-baturin/mattergenbis_vb scout-matter fork at Git revision d670944. Use that compatible fork/revision (or a demonstrably compatible newer version), create its environment, and install it as described by that repository.

This model was trained with no conditioning properties (properties: []). It supports unconditional generation. Do not pass chemical-system, band-gap, or other property conditioning unless a separately trained conditional checkpoint is used.

Download

Replace <HF_REPO_ID> with the actual Hugging Face model ID, for example username/mattergen-mp40:

hf download <HF_REPO_ID> --local-dir mattergen-mp40

For a private repository, first run hf auth login with a token that has read access.

Verify the checkpoint after downloading:

sha256sum mattergen-mp40/last.ckpt

Recommended inference method

Run the included helper from an environment in which the compatible MatterGen package is installed:

python mattergen-mp40/sample_mp40.py \
  --output results/mp40_samples \
  --batch-size 4 \
  --num-batches 1 \
  --gpu-memory-gb 8

Use --force-gpu 0 to select GPU 0. Reduce --batch-size if GPU memory is insufficient. batch-size * num-batches is the total number of generated structures.

The helper performs all required portability steps:

1. Loads mp40_num_atoms_distribution.json and verifies that probabilities are non-negative, sum to one, and cover every atom count from 1 through 40.
2. Registers it as NUM_ATOMS_DISTRIBUTIONS["MP_40"] at runtime. No source-code modification is required for this path.
3. Recovers config.yaml from the full checkpoint and replaces the original machine-specific GemNet scale-file path with the installed MatterGen path.
4. Loads last.ckpt, samples with num_atoms_distribution="MP_40", and disables trajectory recording to reduce disk and memory use.

Outputs:

results/mp40_samples/generated_crystals.extxyz
results/mp40_samples/generated_crystals_cif.zip

The generated structures are unrelaxed. Validate interatomic distances, composition, and charge plausibility, then relax and evaluate them with an appropriate interatomic potential or DFT workflow before scientific use.

Why the JSON distribution is required

MatterGen first samples the requested number of atoms and then generates a structure conditional on that count. The atom-count prior is not a learned parameter and is therefore not part of state_dict. Stock MatterGen commonly defaults to ALEX_MP_20, which only samples 1–20 atoms. Using that default with this checkpoint would leave the trained 21–40 atom range unused and would not reproduce the MP-40 training distribution.

Optional permanent MatterGen integration

The included sample_mp40.py is the recommended approach because it works without changing the installed package. To register MP-40 permanently, copy the JSON beside MatterGen's distribution module:

cp mattergen-mp40/mp40_num_atoms_distribution.json \
  /path/to/mattergen/mattergen/common/data/mp40_num_atoms_distribution.json

Then add these imports near the top of mattergen/common/data/num_atoms_distribution.py:

import json
from pathlib import Path

After the existing NUM_ATOMS_DISTRIBUTIONS dictionary is defined, add:

_mp40_path = Path(__file__).with_name("mp40_num_atoms_distribution.json")
with _mp40_path.open(encoding="utf-8") as _handle:
    _mp40_payload = json.load(_handle)

NUM_ATOMS_DISTRIBUTIONS["MP_40"] = {
    int(num_atoms): float(probability)
    for num_atoms, probability in _mp40_payload["probabilities"].items()
}

The current mattergen-generate wrapper also hardcodes the generator's default atom prior. To expose it through the CLI, add this argument to main() in mattergen/scripts/generate.py:

num_atoms_distribution: str = "ALEX_MP_20",

and pass it when constructing CrystalGenerator:

num_atoms_distribution=num_atoms_distribution,

After extracting the embedded configuration (the helper does this automatically), direct CLI generation can then use:

mattergen-generate results/mp40_samples \
  --model_path=mattergen-mp40 \
  --checkpoint_epoch=last \
  --num_atoms_distribution=MP_40 \
  --batch_size=4 \
  --num_batches=1 \
  --record_trajectories=False

Resume training

Resuming requires the original MP-40 dataset cache; the dataset itself is not included in this model bundle. It must have this layout:

datasets/cache/mp_40/
├── train/
├── val/
└── test/

Copy the supplied data-module configuration into the compatible MatterGen checkout:

cp mattergen-mp40/mp_40_data_module.yaml \
  /path/to/mattergen/mattergen/conf/data_module/mp_40.yaml

From that MatterGen checkout, resume with an absolute checkpoint path:

mattergen-train \
  data_module=mp_40 \
  checkpoint_path=/absolute/path/to/mattergen-mp40/last.ckpt \
  auto_resume=false

The saved run targeted 900 epochs and the checkpoint is from epoch 894. To continue beyond epoch 900, override the maximum, for example:

mattergen-train \
  data_module=mp_40 \
  data_module.max_epochs=1000 \
  checkpoint_path=/absolute/path/to/mattergen-mp40/last.ckpt \
  auto_resume=false

Adjust dataset location, devices, worker counts, and output directory for the new machine. PyTorch Lightning will restore the model, optimizer, scheduler, epoch, and global-step state from last.ckpt.