ml-property-predictor

Files (8)

• SKILL.md
• README.md · examples/mace_bulk_modulus/
• run_mace.py · examples/mace_bulk_modulus/
• README.md · examples/matgl_bulk_modulus/
• fetch_bulk_modulus.py · examples/matgl_bulk_modulus/
• run_matgl.py · examples/matgl_bulk_modulus/
• train_mace_property.py · scripts/
• train_matgl_property.py · scripts/

SKILL.md (first lines)

# MLIP Property Predictor Training

## Goal

To leverage pre-trained GNN representations from MLIPs to train an independent readout head for any custom scalar target property (e.g., bulk modulus, bandgap, formation energy, or spin states) directly from crystal or molecular structures.

## Overview

This skill allows you to leverage pre-trained GNN representations from MLIPs to train an independent readout head for any custom scalar target property, such as bulk modulus, bandgap, formation energy, or spin states.

To keep the core MLIP wrappers clean, property prediction in AtomisticSkills is handled by standalone training scripts located in the `${CLAUDE_SKILL_DIR}/scripts/` directory.

## Workflow

1. **Prepare Data**: Build a `.json` or `.xyz` dataset containing structures and the corresponding scalar property labels. JSON datasets should be lists of dicts containing a `structure` key (Pymatgen format) and your target property key.
2. **Determine Property Type**: Determine if the property is `"intensive"` (e.g. Bandgap, Bulk Modulus) or `"extensive"` (e.g. Total Energy).
   - *MatGL logic*: For intensive targets, node features undergo a global graph readout (like `Set2Set`) before passing through an MLP. For extensive targets, the MLP outputs atomic properties which are then sum-pooled.
   - *MACE logic*: MACE natively supports extensive targets by predicting site-wise scalar outputs and sum-pooling them. When intensive properties are targeted, MACE still sum-pools site-wise outputs, forcing the model to internally learn the intensive invariant.
3. **Execute Script**: Run the MACE or MatGL property prediction script in the `mlip` environment.

---

## Example 1: Training a MACE Property Predictor

MACE property training is handled by `scripts/train_mace_property.py`. It dynamically patches the `mace.cli.run_train` module to freeze the backbone (if requested) and inject a custom intensive/extensive property readout.

```bash

# Run the standalone MACE property training script
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/train_mace_property.py \
    --data_path .agents/test/mp_bulk_modulus.json \
    --model_name MACE-OMAT-0-small \
    --target_property bulk_modulus \
    --property_type intensive \
    --epochs 30 \
    --batch_size 16 \
    --lr 0.001 \
    --output_dir custom_mace_results/
```

- Adds `--freeze_backbone` automatically by default to preserve the MACE representation and avoid catastrophic forgetting.
- The custom weights will be saved to `custom_mace_results/`.

---

## Example 2: Training a MatGL (CHGNet) Property Predictor

MatGL property training is handled by `scripts/train_matgl_property.py`. It loads a pretrained CHGNet model (`CHGNet-PES-MatPES-PBE-1M-2026.9` by default), replaces the data collater to safely handle graph caching, and trains the property explicitly.

```bash

# Run the standalone MatGL property training script
${CLAUDE_SKILL_DIR}/../../venv/run mlip python ${CLAUDE_SKILL_DIR}/scripts/train_matgl_property.py \
    --data_path .agents/test/mp_bulk_modulus.json \
    --model_name CHGNet-PES-MatPES-PBE-1M-2026.9 \
    --target_property bulk_modulus \
    --property_type intensive \
    --epochs 30 \
    --batch_size 32 \
    --lr 0.001 \
    --freeze_backbone \
    --output_dir custom_matgl_results/
```

- `--freeze_backbone` is optional (defaults to `False` for MatGL). If passed, only the final `readout` MLP layers of the CHGNet model will be fine-tuned.
- Intensive extensive targets are handled seamlessly without breaking the pretrained model architecture.
- Model checkpoints are securely saved in `custom_matgl_results/matgl_model/`.
- The example dataset `.agents/test/mp_bulk_modulus.json` (500 MP K_VRH values) is produced by
  [`examples/matgl_bulk_modulus/fetch_bulk_modulus.py`](examples/matgl_bulk_modulus/fetch_bulk_modulus.py);
  see [the MatGL example](examples/matgl_bulk_modulus/README.md) for results against a mean-value baseline.

## Constraints

- **Environments**: The MACE and MatGL predictors both require `mlip`. **Each code block MUST specify the environment.**
- **Data Format**: The dataset must be `.json` or XYZ formatted with the raw structures or ASE atoms.
- **Subprocess Dependency**: The `train_mace_property.py` script spawns an underlying `mace.cli.run_train` subprocess to maintain compatibility with MACE's native optimizers.
- **Pre-trained Architecture**: For MatGL, changing the intensive/extensive nature of a pre-trained model changes the head dimensions. Extensive model predictions are mathematically scaled down by the number of atoms dynamically at training time if an intensive property is targeted.

## References

- Batatia et al., "MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields", *NeurIPS*, 2022. [DOI](https://doi.org/10.48550/arXiv.2206.07697)
- Deng et al., "CHGNet as a pretrained universal neural network potential for charge-informed atomistic modelling", *Nat. Mach. Intell.* 5, 1031 (2023). [DOI](https://doi.org/10.1038/s42256-023-00716-3)

---

**Author:** Bowen Deng
**Contact:** [GitHub @learningmatter-mit](https://github.com/learningmatter-mit)

From learningmatter-mit/AtomisticSkills, MIT.

More from learningmatter-mit/AtomisticSkills