Differentiable Transition Temperature Correction (DiffTTC)

Code supporting the paper Refining machine learning potentials through thermodynamic theory of phase transitions.

Overview

This repository contains the following files:

examples/titanium/
| data/ -- Refererence data (properties)
| ...
| model/ -- Trained and refined model
| ...
| notebooks/ -- Evaluation of training and simulations
| ...
| scripts/ -- Training scripts
| eval_utils.py -- Evaluate free energy computations
| export.py -- Export model to LAMMPS
| predict.py -- Predict U, F, etc. on DFT samples
| train_all_parallel.py -- Refine model on phase diagram/
| train_utils.py -- Model definition, optimizer, etc.
| train.py -- Pretrain model with Force Matching
| visualize.py -- Visualize results from multiple simulations
| simulations/
| lammps/ -- Simulation input scripts
| ...
| mace-mp/ -- MACE-MP-0b3 foundational model files
| ..
| scripts/ -- Run multiple bundled lammps simulations
| ...
| external/ -- Dependencies to be installed
| chemtrain/ -- Modified version of chemtrain
| ...
| chemutils/ -- MACE model and dataset loader
| ...
| lammps/ -- Used LAMMPS version (subrepo)
| ...
| juwels/ -- Extra scripts to run on HPC cluster
| ...

Quickstart

First, install the external dependencies in the following order:

pip install -e "./external/chemtrain[all]"
pip install -e "./external/chemutils"
pip install "jax[cuda12]==0.4.37" --force-reinstall

To pretrain a model with the given hyperparameter on the CUDA device 0, run the following command:

cd examples/titanium/scripts
python train.py "0"

To refine the provided pretrained model via DiffTTC on the prepared setup, run the following command:

cd examples/titanium/scripts
python train_all_parallel.py "0"

To install the LAMMPS plugin, please refer to the documentation.

Full Cycle - Pretraining and Refining with DiffTTC

First, train an MLP using the command given in Quickstart. This command will create a model folder titanium__train* in the parent directory examples/titanium/scripts/output. The model can be copied into the models folder examples/titanium/models. Then, the LAMMPS simulations can be run by specifying the model in the run_*.sh scripts located in the folder examples/titanium/simulations by executing them outside their parent directory. For example, to run the coexistence simulations, enter the command:

cd examples/titanium/simulations
bash scripts/run_coexistence.sh

Some simulations have to be run in a specific order, including postprocessing scripts.

For the HCP-BCC transition:

1. run_prepare.sh and compute hcp_constants.txtand bcc_constants.txt.
2. run_ti_bcc.sh and run_ti_hcp.sh
3. run_rs_bcc-hcp.sh
4. run_create_bcc_box.sh and run_create_hcp_box.sh and convert to pdb

For the BCC-Liquid transition:

1. run_coexistence.sh and compute temps.txt
2. run_prepare_liquid.sh and compute bcc_constants.txt and liquid_constants.txt
3. run_rs_bcc-liquid.sh
4. run_create_bcc_box.sh and run_create_liquid_box.sh and convert to pdb

The thermal expansion, MSD and VACF have to be computed last:

1. run_thermal_expansion.sh
2. run_compute_msd.sh and run_compute_vacf.sh

The prepared free energy targets and simulation boxes have to be defined in the training script train_all_parallel.py. Then, the script can be run via the instructions given in Quickstart.

Citation

If you use DiffTTC, please cite:

@article{fuchsDiffTTC2026,
title = {Refining Machine Learning Potentials through Thermodynamic Theory of Phase Transitions},
author = {Fuchs, Paul and Zavadlav, Julija},
date = {2026-06-22},
journaltitle = {npj Computational Materials},
shortjournal = {npj Comput Mater},
issn = {2057-3960},
doi = {10.1038/s41524-026-02195-7},
url = {https://doi.org/10.1038/s41524-026-02195-7},
}