Repository for training, testing and developing machine learned force fields using the So3krates transformer (check
out the NeurIPS paper).
Assuming you have already set up an virtual environment with python version >= 3.8. In order to ensure compatibility
with CUDA jax/jaxlib have to be installed manually. Therefore before you install MLFF run one of the following
commands (depending on your CUDA version)
pip install --upgrade pip
# CUDA 12 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
# CUDA 11 installation
# Note: wheels only available on linux.
pip install --upgrade "jax[cuda11_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
for details check the official JAX
repository.
Since the optax package has a tensorflow dependency for checkpointing and does not install it itself, we also have to
install tensorflow, although the models in mlff are written in JAX. It can be installed running
pip install tensorflow
Apple ARM M1/M2
In case you have an Apple ARM M1/M2 chip, the designated tensorflow package called tensorflow-maxos has to be
installed. More details on doing so can be found e.g. here
or here (the metal plugin is not needed).
Unfortunately, we didn't found an M1 for testing, so we are happy about reported issues.
Next clone the mlff repository by running
git clone https://github.com/thorben-frank/mlff.git
Now do
cd mlff
pip install -e .
which completes the installation and installs remaining dependencies.
If you do not have a weights and bias account already you can create on here. After installing
mlff run
wandb login
and log in with your account.
Following we will give a quick start how to train, evaluate and run an MD simulation with the
So3krates model.
Train your fist So3krates model by running
train_so3krates --data_file data.xyz --n_train 1000 --n_valid 100 --wandb_init project=so3krates,name=first_run
The --data_file can be any format digestible by the ase.io.read method. In case minimal image convention
should be applied, add --mic to the command. The model parameters will be saved per default to module/. Another
directory can be specified using --ckpt_dir $CKPT_DIR, which will safe the model parameters to $CKPT_DIR/.
More details on training can be found in the detailed training section below.
After training, change into the model directory, e.g. and run the evaluate command
cd module
evaluate
As before, when your data is not in eV and Angstrom add the --units keyword. The reported metrics are then in eV and
Angstrom (e.g. --units energy='kcal/mol',force='kcal/(mol*Ang)' if the energy in your data is in kcal/mol).
You can use the mdx package which is the mlff internal MD package, fully relying on jax and thus fully
optimized for XLA compilation on GPU. Before you can use it make sure you install the glp
package by cloning the glp repository and install it
git clone git@github.com:sirmarcel/glp.git
cd glp
pip install .
First, lets create a relaxed structure, using the LBFGS optimizer
run_relaxation --qn_max_steps 1000 --qn_tol 0.0001 --use_mdx
which will save the relaxed geometry to relaxed_structure.h5. Next, convert the .h5 file to an
xyz file, by running
trajectory_to_xyz --trajectory relaxed_structure.h5 --output relaxed_structure.xyz
We now run an MD with the relaxed structure as start geometry
run_md --start_geometry relaxed_structure.xyz --thermostat velocity_verlet --temperature_init 600 --time_step 0.5 --total_time 1 --use_mdx
Temperature is in Kelvin, time step in femto seconds and total time in nano seconds. It will save a trajectory.h5
file to the current working directory.
After the MD is finished you can either work with the trajectory.h5 using e.g. a jupyter notebook and h5py.
Alternatively, you can run
trajectory_to_xyz --trajectory trajectory.h5 --output trajectory.xyz
which will create an xyz file. The resulting xyz file can be used as input to the
MDAnalysis python package, which provides a broad range of functions
to analyse the MD simulations. The central Universe object can be creates easily as
import MDAnalysis as mda
# Load MD simulation results from xyz
u = mda.Universe('trajectory.xyz')In the quickstart section we went through a few basic steps, allowing to train, validate a so3krates model as well
as running an MD simulation. If you want to learn more about each of the steps, check the following sections.
Lets start start from the training command already shown in the quickstart section
train_so3krates --ckpt_dir first_module --data_file atoms.xyz --n_train 1000 --n_valid 100
for which all data in atoms.xyz is loaded an split into 1000 data points for training,
100 data points for validation and the remaining data points n_test = n_tot - 1000 - 100 is hold back for testing
the potential after training. The validation data points are used to determine the best performing model during training,
for which the parameters are saved to first_module/checkpoint_loss_XXX where XXX denotes the training step for which the best performing
model was found. We will show later, how to load the checkpoint such that one use the trained potential directly in
Python.
mlff can deal with any input file that can be read by the ase.io.read method.
Further --data_file admits to pass *.npz files. *.npz files allow to store numpy.ndarray under different keys.
Thus, mlff needs to "know" under which key to find e.g. the positions, the forces and so on .. Per default, mlff
assumes the following relations between property and key
{
atomic_position: R, # shape: (n_data, n, 3)
atomic_type: z, # shape: (n_data, n) or (n)
energy: E, # shape: (n_data, 1)
force: F # shape: (n_data, n, 3)
}
If you have an *.npz file which uses a different convention, you can specify the keys customizing the property keys
via
train_so3krates --ckpt_dir second_module --data_file file.npz --n_train 1000 --n_valid 100 --prop_keys atomic_position=pos,atomic_type=numbers
The above examples would assume that the properties energy and force are still found under the keys E and F,
respectively but position and atomic_type are found under pos and numbers.
Per default, mlff assumes the ASE default units which are eV for energy and Angstrom for coordinates. Some data
sets, however, differ from these convention, e.g. the MD17 or the MD22 data set. You can download the corresponding
*.npz files here (Note that the *.xyz files provided there are not formatted in a
way that allows reading them via ase.io.read method). For both data sets the energy is is in kcal/mol such that
the forces are in kcal/(mol*Ang). You can either pre-process the data yourself by applying the proper conversion
factors and pass the data directly into the train_so3krates command. Alternatively, you can set them manually by
train_so3krates --ckpt_dir dha_module --train_file dha.npz --n_train 1000 --n_valid 500 --units energy='kcal/mol',force='kcal/(mol*Ang)'
Note the character strings for the units, which are necessary in the course of internal processing. This will internally
rescale the energy and the forces to eV and eV/Ang.
In [1] So3krates was used to calculate EOS and heat flux in solids, such that it must be capable of
handling periodic boundary conditions. If you want to apply the minimal image convention, you can specify this by
adding the corresponding flag to the training command
train_so3krates --ckpt_dir pbc_module --train_file file_with_pbc.xyz --n_train 100 --n_valid 100 --mic
Internally, mlff uses the ase.neighborlist.primitive_neighborlist
for computing the atomic neighborhoods.
The energy scale of the data, tend to increase rapidly with the number of atoms in the system. However, relevant scale
of energy changes is usually on the scale of a few eV, such that one typically shifts the energy by the mean of the
energies in the --n_train training data points. This is done by default when running the train_so3krates command.
However, it might be desired to atom type specific shifts, which is possible via
train_so3krates --energy_shift_module --atoms.xyz --n_train 1000 --n_valid 200 --shifts 1=-500.30, 6=-6000.25
and would shift the energy by -500.30 for each hydrogen in the structure and by -6000.25 for each carbon.
One can further vary different model hyperparameters: --L sets the number of message passing steps, --F sets the
feature dimension, --degrees sets the degrees for spherical harmonic coordinates and --r_cut sets the cutoff for the
atomic neighborhoods. E.g. a model with 2 message passing layers, feature dimension 64, degree 1 and 2 and cutoff of 4 Angstrom
can be trained by running
train_so3krates --new_hypers_module --atoms.xyz --n_train 1000 --n_valid 200 --L 2 --F 64 --degrees 1 2 --r_cut 4
TODO
To keep track of the performed experiments, you can organize your trainings using weights and bias. You can specify the project and name of the current trainings run via
train_so3krates --wandb_module --atoms.xyz --n_train 1000 --n_valid 200 --wandb_init project=so3krates,name=deep_dive_run
The arguments passed to --wandb_init are passed as is to the wandb.init for which you can find all possible
arguments here.
In order to validate the models performance, the errors on the unseen test set are often a good starting point. Go to
model you want to evaluate by going to the model directory (the directory in which the checkpoint_loss_XXX file lies)
and use the evaluate command
cd module
evaluate
This command will evaluate the model on the test data points and print the mean absolute error (MAE), the root mean squared
error (RMSE) and the R2 spearson correlation. It will further save two files in the current directory which are called
metrics.json and evaluate_predictions.npz. The former contains the metrics which have also been printed and the latter
contains the per data point predictions of the model. Thus, the metrics can be calculated from the evaluate_predictions.npz
by using the following code snippet
import numpy as np
eval_data = np.load('evaluate_predictions.npz', allow_pickle=True)
predicted_energies = eval_data['predictions'].item()['E']
target_energies = eval_data['targets'].item()['E']
mae = np.abs(predicted_energies - target_energies).mean()
print(f'MAE: {mae:.4f} (eV)')While it is convenient to train and validate a model using the CLI, playing around with the potential is often much
easier in Python in particular if it is of interest to couple it with other methods and packages.
Below you find a code snippet how to load a trained So3krates model and use it in Python.
import jax.numpy as jnp
from mlff.mdx import MLFFPotential
ckpt_dir = 'path/to/ckpt_dir/'
dtype = jnp.float64
pot = MLFFPotential.create_from_ckpt_dir(ckpt_dir=ckpt_dir, dtype=dtype)The resulting potential is a Potential in the sense of the
the glp package. Thus, you can directly use your trained so3krates model to
perform energy, force, stress and heat flux calculations. glp also offers a binding to vibes, which
further allows you to do Phonon calculations and much more. To calculate thermal conductivities you can use the gkx
package.
TODO
The test suite can be run with pytest as:
pytest tests/
If you use parts of the code please cite the corresponding paper
@article{frank2022so3krates,
title={So3krates: Equivariant attention for interactions on arbitrary length-scales in molecular systems},
author={Frank, Thorben and Unke, Oliver and M{\"u}ller, Klaus-Robert},
journal={Advances in Neural Information Processing Systems},
volume={35},
pages={29400--29413},
year={2022}
}
- [1] M.F. Langer, J.T. Frank, and F. Knoop. Stress and heat flux via automatic differentiation, arxiv 2305.01401 (2023)