MLIP fine tunings

This page contains a tutorial on how to perform a self-consistent biased fine-tuning with MACE for a reference potential energy surface with Caracal. In this example, a MACE MLIP will be fine-tuned on the fly for the intermolecular proton transfer in malonaldehyde. It is a benchmark example, where not an ab initio method is learned, but the analytical PES for the system by Mizukami et al., which is included into Caracal.

All relevant input files for the treated example system can also be found on a Zenodo archive.

Before starting the calculations, ensure that you have installed the ASE and MACE Python modules (see, e.g., here).

In general, it is advisable to make a subfolder for each planned fine-tuning cycle (e.g., cycle1, cycle2 etc.), and within these folders, generate another subfolder for the steps A to C (optionally D) described in the following:

A. Perform biased MACE MD calculations

First, the umbrella structure generation phase needs to be sampled by the selected MACE foundation model or the fine-tuned MLIP from the last cycle to obtain a representative set of structures along the reaction path. For this, a usual calc_rate.x keyfile (calc_rate.key) can be used:

pes mace

mace {
  mlip_file mace-mp-0b3-medium.model
  coord_file ts.xyz
  device gpu
}

ts_struc ts.xyz

rpmd_beads 4
deltat 0.2
tdump 10
nvt {
  temp 500
  thermostat andersen
  andersen_step 80
}

mecha {
  type rearrange
  reactant1  1 2 3 4 5 6 7 8 9
  bond_form 3-7
  bond_break 5-7
  dist_inf 10
  reactants_struc reactant.xyz
  n_paths 1
}

umbrella {
  bias 0.015
  bonds -0.10  1.30
  dist 0.02
  gen_steps 10000
  equi_steps 2000
  sample_steps 5000
  sample_trajs 20   
}

pmf {
  xi_range -0.08  1.28
  bins 5000
  method integration
  minloc pmf_min
}

 recross {
   equi_steps 4000
   child_total 50
   child_interval 1000
   child_perpoint 10
   child_steps 1000
   manual_pos 1.0
   recross_nocheck
   recross_mpi
 }

print_gen 625
print_cross 50
print_train mace

This input file assumes a foundation model calculation in the first fine-tuning cycle (MP-0-b3 in the example). Replace mlip_file mace-mp-0b3-medium.model with the fine-tuned MACE MLIP if it is the second or later cycle.

With the keyword print_gen 625 it is ordered that only the umbrella structure equilibration along the reaction path with one RPMD bead is done and each 625th structure is printed out. This counter starts again for each single umbrella trajectory. In this case, 100000 MD steps are done per umbrella structure generation trajectory (gen_steps 10000) and thus 16 frames are written per trajectory, which equals to 1136 structures in total (71 umbrella windows between -0.10 and 1.30). The keyword print_train mace orders the print to be done in the MACE training set format, which is a usual xyz trajectory with more information at the right side (energies and forces). The shown input file does the MD samplings at 500 K. If a MLIP with the ability to describe the system at various different temperatures shall be optimized, the calculation should be redone at a couple of different temperatures.

For unimolecular reactions like the H exchange in malonaldehyde, it is important to also include information about recrossing trajectories to the training set. Therefore, the keyword print_cross has been added as well. Now, the recrossing calculation is done after the umbrella structure generation, skipping the actual umbrella samplings. The location of the TS dividing surface, usually determined by the umbrella samplings, must therefore be set manually (manual_pos 1.0), usually 1.0, i.e., the potential energy TS, should be good enough, since we are not interested in the exact transmission coefficient but only in sufficient samplings on both sides of the TS region. The number of written structures is determined with the print_cross keyword, here is each 50th frame of the recrossing trajectories written out, which corresponds to 1000 structures (20 frames per trajectory, 50 trajectories in total). For the 50 trajectories, I set child_total 50 and child_perpoint 10 such that 10 trajectories are started from another parent trajectory frame with 1000 MD steps distance to each other, respectively.

Then, execute the calculation on a GPU cluster:

 calc_rate.x calc_rate.key

The selected MD frames of the structure generation phase are then written to traj_gen.xyz, those of the recrossing phase are written to traj_cross.xyz.

B. Recalculate energies and forces with reference method

For an efficient MACE fine-tuning, the energies of the isolated atoms should be given. For the analytical reference potential, this is not possible, I therefore just took the energies of the foundation model, calculated with the following python file:

 from ase.io import read
 from mace.calculators import MACECalculator

 model_file = "mace-mp-0b3-medium.model"

 calc = MACECalculator(model_path=model_file)

 atoms_list = read("atoms.xyz", index=":")  # read all frames

 for atom in atoms_list:
     energy = calc.get_potential_energy(atom)
     print(f"{energy:.6f} eV")

together with the structure file atoms.xyz:

 1
 H atom  
 H 0 0 0
 1
 C atom
 C 0 0 0
 1
 O atom
 O 0 0 0

For ab initio calculations, you can just calculate the atomic energies with the reference method. The energies then need to be placed of the header of the actual training set file, like:

 1
 Properties=species:S:1:pos:R:3:REF_forces:R:3 config_type=IsolatedAtom REF_energy=      -1.117340  pbc="F F F"
 H 0.0 0.0 0.0 0.0 0.0 0.0

Next, the training set itself needs to be calculated with the reference method. For this, combine all traj_gen.xyz and traj_cross.xyz files frames into a file named trajectory.xyz and write the following egrad.key file:

 pes ana_malon

 xyzstart trajectory.xyz

 job egrad

Just execute

 explore.x egrad.key

and the energies and forces of the structures are written to files energies.dat and gradients.dat, respectively.

To convert the results back into the MACE training set format, compile the Fortran program egrad2mace.f90 from the scripts folder of the Caracal repository and execute it:

 egrad2mace

Then, a file xyz_grads.xyz is written, which has the format of a MACE training set file.

If you do ab initio calculations, write your own script to translate the results into the MACE training set format. If you do VASP calculations, you can use the vasp2trainset program of our utils4VASP repository. Ensure that energies are given in eV and forces in eV/Angstrom!

Now, cat this file together with the atomic energies to the overall trainset.xyz file.

C. MACE fine-tuning

Copy the trainset.xyz file into a new directory and execute the following command on a GPU node, ideally with a local Python environment that supports cueq for GPU speedup (look into the MACE Wiki for details):

 source /home/hpc/b146dc/b146dc10/work_vault/python_envs/mace_cueq/bin/activate

 python3.12 /home/hpc/b146dc/b146dc10/work_vault/python_envs/mace_cueq/lib/python3.12/site-packages/mace/cli/run_train.py \
       --name="mace_fine_tuning" \
       --train_file="trainset.xyz" \
       --foundation_model="mace-mp-0b3-medium.model" \
       --valid_fraction=0.2 \
       --energy_weight=1 \
       --forces_weight=10 \
       --eval_interval=1 \
       --error_table='PerAtomMAE' \
       --energy_key="REF_energy" \
       --forces_key="REF_forces" \
       --interaction_first="RealAgnosticResidualInteractionBlock" \
       --interaction="RealAgnosticResidualInteractionBlock" \
       --scaling='rms_forces_scaling' \
       --lr=0.001 \
       --enable_cueq=True \
       --E0s="isolated" \
       --weight_decay=1e-8 \
       --ema \
       --ema_decay=0.995 \
       --scheduler_patience=5 \
       --batch_size=10 \
       --valid_batch_size=20 \
       --max_num_epochs=100 \
       --patience=20 \
       --amsgrad \
       --device="cuda" \
       --seed=1 \
       --clip_grad=100 \
       --keep_checkpoints \
       --restart_latest \
       --save_cpu > logfile

If the optimization appears to be unstable (spiking loss function), lower --lr and increase --batch_size as well as --valid_batch_size. If there are memory issues, on the other hand, lower --batch_size and --valid_batch_size. Further, the total number of optimization steps might be increased or decreased (--max_num_epochs).

If this is the second or later fine-tuning cycle, copy the MLIP file from the last cycle into the folder and replace --foundation_model="mace-mp-0b3-medium.model" by the name of the MLIP from the last cycle!

After the fine-tuning has been finished, copy the mace_fine_tuning_run-1.model (or whatever the name of your calculation was) from the `checkpoints´ folder.

If you want to verify the quality of the MLIP in this cycle, continue with step D, else, go back to step A for the next cycle.

D. Investigation of MLIP quality (optional)

In order to investigate how well the fine-tuned MLIP of this cycle is able to reproduce the energies and forces of the true PES, another set of umbrella structure generation calculations should be done to generate an independent verification set.

For this, take the shown keyfile from part A and redo the calculations for the investigated temperatures with the fine-tuned MACE MLIP from part C. You might want to change the print_gen 100 and equi_steps keywords for a higher printout frequency, e.g., to print_gen 100 and equi_steps 4000`, or take less structures if the ab initio method is very expensive.

Then, execute again on a GPU:

calc_rate.x caracal.key

Now, cat all traj_gen.xyz files together and copy them into a new folder. Also cat the traj_gen_xi.dat files together and copy them into the same folder (for caracal_quality, see below).

There, do the same as in B for the recalculation with the reference method, in this case, just take the egrad.key file:

 pes ana_malon

 xyzstart traj_gen.xyz

 job egrad

and start the calculation with

 explore.x egrad.key

The files energies.dat and gradients.dat are then written. After the calculations have been done, compile the program caracal_quality.f90 from the scripts folder of the Caracal repository and execute it:

 caracal_quality > quality.log

In the quality.log file, the MAE of energies and forces is given. Further, two plots were done and written to files: 2d_gradient_histo.png with the deviations of force directions and plot_xi_ener.png with the energy error of each calculations, along the reaction path (for this, the traj_gen_xi.dat file is needed). If the style of these plots shall be changed, vary the respective plot_2d_grad_histo.gnu and plot_xi_ener.gnu files.

If you train against ab initio just write a helper script that translates the ab initio results into the files energies.dat and ´gradients.datto use thecaracal_quality` program there as well.