KO | EN

GitLite

Minimalist Open Source Repository Search Engine

gitlite — search

#dsh #dsh-plugin #developer-tools #ai-agent #claude-code #codex #typescript #ai-agents #deepseek-harness #mcp #cli #rust

MatterTune

★ 82 Open GitHub ↗

A unified platform for fine-tuning atomistic foundation models in chemistry and materials science

// repository documentation

Was this content helpful?

★ 0 (0 ratings)

# MatterTune: A Unified Platform for Atomistic Foundation Model Fine-Tuning

[![Documentation Status](https://github.com/Fung-Lab/MatterTune/actions/workflows/docs.yml/badge.svg)](https://fung-lab.github.io/MatterTune/)

**[📚 Documentation](https://fung-lab.github.io/MatterTune/) | [🔧 Installation Guide](https://fung-lab.github.io/MatterTune/installation.html)**

MatterTune is a flexible and powerful machine learning library designed specifically for fine-tuning state-of-the-art chemistry foundation models. It provides intuitive interfaces for computational chemists and materials scientists to fine-tune pre-trained models on their specific use cases.

## Features

- Foundation models support: 
  - [UMA](https://fung-lab.github.io/MatterTune/backbones/uma.html)
  - [MACE](https://fung-lab.github.io/MatterTune/backbones/mace.html)
  - [MatterSim](https://fung-lab.github.io/MatterTune/backbones/mattersim.html)
  - [NequIP/Allegro](https://fung-lab.github.io/MatterTune/backbones/nequip%26allegro.html)
  - [EquiformerV2](https://fung-lab.github.io/MatterTune/backbones/eqv2.html)
  - [JMP](https://fung-lab.github.io/MatterTune/backbones/jmp.html)
  - [ORB](https://fung-lab.github.io/MatterTune/backbones/orb.html)
  - and more to be added...
- Multiple property predictions: energy, forces, stress, and custom properties.
- Various supported dataset formats: XYZ, ASE databases, Materials Project, Matbench, and more.
- Comprehensive training features with automated data splitting and logging.

## Advanced Features

- Atomistic Foundation Model Acceleration via Message Passing Pruning and Graph Partitioning ([publication](), [docs]())

## Quick Start

```python
import mattertune as mt
import mattertune.configs as MC

# Phase 1: Fine-tuning the model
# -----------------------------

# Define the configuration for model, data, and training
config = mt.configs.MatterTunerConfig(
    # Configure the model: using JMP backbone with energy prediction
    model=MC.MACEBackboneConfig(
        pretrained_model = "mace-medium",  # Select pretrained model type
        properties = [
            MC.EnergyPropertyConfig(  # Configure energy prediction
                loss=MC.MAELossConfig(),  # Using MAE loss
            ),
            MC.ForcesPropertyConfig(
                loss=MC.MAELossConfig(), 
                conservative=True,
            ),
            MC.StressesPropertyConfig(
                loss=MC.MAELossConfig(),
                conservative=True
            )
        ],
        optimizer = MC.AdamWConfig(lr=8.0e-5)
    ),
    # Configure the data: loading from XYZ file with automatic train/val split
    data=MC.AutoSplitDataModuleConfig(
        dataset=MC.XYZDatasetConfig(src=Path("YOUR_XYZFILE_PATH")  # Path to your XYZ data),
        train_split=0.8,  # Use 80% of data for training
        batch_size=32  # Process 32 structures per batch
    ),
    # Configure the training process
    trainer=MC.TrainerConfig(
        max_epochs=10,  # Train for 10 epochs
        accelerator="gpu",  # Use GPU for training
        devices=[0]  # Use first GPU
        additional_trainer_kwargs={
            "inference_mode": False,
        }
    ),
)

# Create tuner and start training
tuner = mt.MatterTune(config)
model, trainer = tuner.tune()

# Save the fine-tuned model
trainer.save_checkpoint("mace_ft.ckpt")

# Phase 2: Using the fine-tuned model
# ----------------------------------

from ase.optimize import BFGS
from ase import Atoms

# Load the fine-tuned model
model = mt.backbones.MACEBackboneModule.load_from_checkpoint("finetuned_model.ckpt")

# Create an ASE calculator from the model
calculator = model.ase_calculator()

# Set up an atomic structure
atoms = Atoms('H2O',
             positions=[[0, 0, 0], [0, 0, 1], [0, 1, 0]],
             cell=[10, 10, 10],
             pbc=True)
atoms.calc = calculator

# Run geometry optimization
opt = BFGS(atoms)
opt.run(fmax=0.01)

# Get results
print("Final energy:", atoms.get_potential_energy())
print("Final forces:", atoms.get_forces())
```

## FAQ

We welcome anyone with questions or suggestions about the MatterTune project to open an issue in this repository’s Issues section. Before creating a new issue, please check the existing ones to see whether a solution has already been posted.

## License

MatterTune's core framework is licensed under the MIT License. Note that each supported model backbone is subject to its own licensing terms - see our [license information page of the documentation](https://fung-lab.github.io/MatterTune/license.html) for more details.

## Citation

```
@article{kong2025mattertune,
  title={Mattertune: An integrated, user-friendly platform for fine-tuning atomistic foundation models to accelerate materials simulation and discovery},
  author={Kong, Lingyu and Shoghi, Nima and Hu, Guoxiang and Li, Pan and Fung, Victor},
  journal={arXiv preprint arXiv:2504.10655},
  year={2025}
}
```

MatterTune

Explore Similar Repositories

Was this content helpful?

Recent Feedback

Download README