MACE in Practice II

MACE in Practice II#

In this tutorial, you will learn how to improve MLIP models by using iterative training and active learning. We illustrate these training workflows on MACE, but they are broadly applicable to all MLIPs. We will also showcase the state-of-the-art foundational models - the latest development in the field of MLIPs. These models are trained on massive training sets of inorganic and organic databases and show a great deal of out-of-the-box MD stability in an extensive variety of applications. We will discuss fine-tunning which is an actively-researched technique to tweak these foundational models to new systems (out of training) and/or new levels of reference theory.

Learning Objectives for today:#

Iterative Training: improving stability and accuracy
Active learning: committee models
Foundational models and fine-tuning

1. Iterative Training#

1.1 MD with a smaller MACE model#

The model we trained in our previous tutorial was already stable in MD and quite accurate with little training. This is both because MACE models are smooth and very accurate but also because the task of simulating a single molecule for a few picoseconds is not all that difficult. In general, in real research applications, achieving MD stability and accuracy is not always straightforward from the get-go. Models can be improved through iterative training and active learning which expands the training data to fix errors on the model’s potential energy surface.

To illustrate these concepts in practice, let’s first train a less accurate MACE by reducing a lot the model size and amount of training data: we will make a model with just 20 training configurations. Note that a larger model would probably work out of the box for this example.

from ase.io import read, write

device = 'cuda'
# device = 'cpu'

db = read('data/solvent_xtb.xyz', ':')
write('data/solvent_xtb_train_23.xyz', db[:23]) #first 20 configs plus the 3 E0s will be used for training (parameter optimization)
write('data/solvent_xtb_valid_20.xyz', db[23:43]) #next 20 configs will be used for validation (independent, stop condition)

Let’s make an input config.

%%writefile config/config-03.yml

model: "MACE"
num_channels: 32
max_L: 0
r_max: 4.0
max_ell: 2
name: "mace02_com1"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
train_file: "data/solvent_xtb_train_23.xyz"
valid_file: "data/solvent_xtb_valid_20.xyz"
test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
batch_size: 10
max_num_epochs: 200
swa: True
seed: 123

dev = f'device: {device}'
%store dev >>"config/config-03.yml"

import warnings
warnings.filterwarnings("ignore")
from mace.cli.run_train import main as mace_run_train_main
import sys
import logging

def train_mace(config_file_path):
    logging.getLogger().handlers.clear()
    sys.argv = ["program", "--config", config_file_path]
    mace_run_train_main()

train_mace("config/config-03.yml")

#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture
import glob
import os
for file in glob.glob("MACE_models/*.pt"):
    os.remove(file)

Notice, we are getting substantially larger errors than before. Now, let’s run some dynamics:

from ase.io import read, write
from ase import units
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution
from aseMolec import extAtoms as ea

import os
import time
import numpy as np
import pylab as pl
from IPython import display
np.random.seed(701) #just making sure the MD failure is reproducible

def simpleMD(init_conf, temp, calc, fname, s, T):
    init_conf.set_calculator(calc)

    #initialize the temperature

    MaxwellBoltzmannDistribution(init_conf, temperature_K=300) #initialize temperature at 300
    Stationary(init_conf)
    ZeroRotation(init_conf)

    dyn = Langevin(init_conf, 1.0*units.fs, temperature_K=temp, friction=0.1) #drive system to desired temperature

    %matplotlib inline

    time_fs = []
    temperature = []
    energies = []

    #remove previously stored trajectory with the same name
    os.system('rm -rfv '+fname)

    fig, ax = pl.subplots(2, 1, figsize=(6,6), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})

    def write_frame():
            dyn.atoms.info['energy_mace'] = dyn.atoms.get_potential_energy()
            dyn.atoms.arrays['force_mace'] = dyn.atoms.calc.get_forces()
            dyn.atoms.write(fname, append=True)
            time_fs.append(dyn.get_time()/units.fs)
            temperature.append(dyn.atoms.get_temperature())
            energies.append(dyn.atoms.get_potential_energy()/len(dyn.atoms))

            ax[0].plot(np.array(time_fs), np.array(energies), color="b")
            ax[0].set_ylabel('E (eV/atom)')

            # plot the temperature of the system as subplots
            ax[1].plot(np.array(time_fs), temperature, color="r")
            ax[1].set_ylabel('T (K)')
            ax[1].set_xlabel('Time (fs)')

            display.clear_output(wait=True)
            display.display(pl.gcf())
            time.sleep(0.01)

    dyn.attach(write_frame, interval=s)
    t0 = time.time()
    dyn.run(T)
    t1 = time.time()
    print("MD finished in {0:.2f} minutes!".format((t1-t0)/60))

#let us start with a single molecule
init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()

#we can use MACE as a calculator in ASE!
from mace.calculators import MACECalculator
mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_run-123_stagetwo.model'], device=device)

simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md.xyz', s=10, T=2000)

from ase.io import read, write
import nglview as nv

traj = read('moldyn/mace02_md.xyz', ':')
view = nv.show_asetraj(traj)
view._set_size(w='100%', h='500px')
view

If you go to the end of the trajectory, you should find that the bond angles are actually very strange - it looks unphsyical.

Task for you

Sometimes, when the model is very bad you may find the simulation exploding. Try for yourself, reduce the size of the model (e.g. number of channels) and see if a small model will generated exploding simulations.

1.2 Identify the problem and expand training#

Something doesn’t look right with the hydrogen atoms. Let’s re-evaluate the first 100 steps from the trajectory on the reference XTB potential energy surface and then plot it against MACE energy.

from ase.io import read, write
from tqdm import tqdm
from tblite.ase import TBLite
xtb_calc = TBLite(method="GFN2-xTB",verbosity=0,max_iterations=1000)

#compute true reference XTB values
print("Evaluating MACE configurations with XTB")
traj = read('moldyn/mace02_md.xyz', ':')
for at in tqdm(traj[:100]):
    at.calc =  TBLite(method="GFN2-xTB",verbosity=0)
    at.info['energy_xtb'] = at.get_potential_energy()
    at.arrays['forces_xtb'] = at.get_forces()
    at.calc = None

write('data/mace02_md_100_xtb.xyz', traj[:100]) #save full result

import numpy as np
from aseMolec import extAtoms as ea
from matplotlib import pyplot as plt

traj = read('data/mace02_md_100_xtb.xyz', ':')
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB')
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace', peratom=True), label='MACE')
plt.legend()
plt.xlabel('Time (fs)')
plt.ylabel('Total Energy per Atom (eV)')

Indeed, in the second half of the trajectory, the XTB energy diverges because MACE finds unphysical configurations. With older potentials, at this point the MD would explode, but because MACE is much smoother the simulation keeps going, albeit generating the wrong dynamics. Let’s take twelve of these unphysical configs, add them back to the training set and refit a new model. This is called iterative training:

db = read('data/solvent_xtb_train_23.xyz', ':')
db += traj[40:100:5] #add failed configs to the training set
write('data/solvent_xtb_train_35_gen1.xyz', db)

A good exercise is to try to pick the configs that have the largest error on forces and energies. Try coding it yourself.

1.3 Train a new MACE model and run MD again#

%%writefile config/config-04.yml

model: "MACE"
num_channels: 32
max_L: 0
r_max: 4.0
max_ell: 2
name: "mace02_com1_gen1"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
train_file: "data/solvent_xtb_train_35_gen1.xyz"
valid_file: "data/solvent_xtb_valid_20.xyz"
test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
batch_size: 10
max_num_epochs: 200
swa: True
seed: 123

dev = f'device: {device}'
%store dev >>"config/config-04.yml"

train_mace("config/config-04.yml")

#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture
import glob
import os
for file in glob.glob("MACE_models/*.pt"):
    os.remove(file)

For some reason the energy error on the training set is now huge - can you work out why this is?

What does this imply about how we do iterative training?

from mace.calculators import MACECalculator

mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_gen1_run-123_stagetwo.model'], device=device)
init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()
simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md_gen1.xyz', s=10, T=2000)

from weas_widget import WeasWidget

traj = read('moldyn/mace02_md_gen1.xyz', ':')
w = WeasWidget()
w.from_ase(traj)
w.avr.model_style = 1
w.avr.show_hydrogen_bonds = True
w

Great! The dynamics is already looking better, however it can be difficult to tell if it is really correct or not.

Task for you

Compute the true XTB enery and check if they agree with the MACE energy predicted along the trajectory
Compute radial distribution functions (RDFs) and compared to a ground truth RDFs
If the MACE trajectory does not agree with the ground truth, continue the iterative process and gradually improve the potential
For every new potential, repeat the dynamics and check against grount truth observables.

This is an arduous process, because we need to carefully investigate the trajectories and decide which configs to add back to training. If the ground truth is expensive to compute (DFT, QM) this becomes a very slow process. We could instead automate this protocol by approximating the errors (e.g. via a committee of models) on the fly and select configs which are not well predicted: this is called active learning.

2. Active Learning with MACE#

2.1 Preparing a committee of models#

We can compute errors by evaluating the reference energy and forces (in our case XTB) and computing the difference to MACE predictions. In real research applications, this can be very expensive to evaluate depending on the reference level of theory. Alternatively, we can estimate errors based on a committee of models. Let’s train a committee of MACE models by adding some randomness to the optimization process. We can achieve this by changing the --seed. We already have a model, we will fit two more, on the same data:

%%writefile config/config-05.yml

model: "MACE"
num_channels: 32
max_L: 0
r_max: 4.0
name: "mace02_com2"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
train_file: "data/solvent_xtb_train_23.xyz"
valid_file: "data/solvent_xtb_valid_20.xyz"
test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
batch_size: 10
max_num_epochs: 200
swa: True
seed: 345

dev = f'device: {device}'
%store dev >>"config/config-05.yml"

%%writefile config/config-06.yml

model: "MACE"
num_channels: 32
max_L: 0
r_max: 4.0
name: "mace02_com3"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
train_file: "data/solvent_xtb_train_23.xyz"
valid_file: "data/solvent_xtb_valid_20.xyz"
test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
device: cuda
batch_size: 10
max_num_epochs: 200
swa: True
seed: 567

dev = f'device: {device}'
%store dev >>"config/config-06.yml"

train_mace("config/config-05.yml")
train_mace("config/config-06.yml")

#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture
import glob
import os
for file in glob.glob("MACE_models/*.pt"):
    os.remove(file)

Perfect, we have two new models. Let’s start by testing the commitee on the first 100 frames of the first trajectory we generated. The MACECalculator can conveniently take a list of calculators as input and will compute separate energies from each calculator.

import numpy as np
from ase.io import read
from aseMolec import extAtoms as ea
from matplotlib import pyplot as plt
from mace.calculators import MACECalculator
from tqdm import tqdm
import warnings
warnings.filterwarnings("ignore")

model_paths = ['MACE_models/mace02_com1_run-123_stagetwo.model',
               'MACE_models/mace02_com2_run-345_stagetwo.model',
               'MACE_models/mace02_com3_run-567_stagetwo.model',]
mace_calcs = MACECalculator(model_paths=model_paths, device=device)

traj = read('data/mace02_md_100_xtb.xyz', ':')
for at in tqdm(traj):
    at.calc = mace_calcs
    engs = at.get_potential_energies()
    at.info['energy_mace_1'] = at.info.pop('energy_mace') #rename value obtained with first member of the committee
    at.info['energy_mace_2'] = engs[1]
    at.info['energy_mace_3'] = engs[2]
    at.info['variance'] = np.std(engs)
    at.info['average_mace_energy'] = np.average(engs)
    at.info['true_error'] = np.abs(at.info['average_mace_energy'] - at.info['energy_xtb'])

#Let's check the energies of the MACE committee vs XTB energy
plt.figure(figsize=(12,6))
plt.subplot(1,2,1)

plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB');
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_1', peratom=True), label='MACE_1');
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_2', peratom=True), label='MACE_2');
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_3', peratom=True), label='MACE_3');
plt.legend()
plt.xlabel('Time (fs)');
plt.ylabel('Energy per Atom (eV)');

plt.subplot(1,2,2)
plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'variance', peratom=True), label='committee variance', color='black');
plt.xlabel('Time (fs)');
plt.ylabel('Committee Energy Variance');

Notice how the variance (disagreement between models) increases around the same config where the true error with respect to XTB diverges. This is good news because it indicates the variance is a good proxy for true error.

Task for you

Make a correlation plot between true error and estimated error. How well does the estimation perform?
Try a commitee of models with both different seeds and trained on different selections of data. Does it perform better?

Now we can run dynamics with a commitee of models and monitor the variance in the energy prediction. Because XTB is cheap enough we can also compare that variance with the true error. Do they correlate?

2.2 Running MD with the MACE committee#

from aseMolec import extAtoms as ea
from ase import units
from ase.md.langevin import Langevin
from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution
from ase.io import read, write

import random
import numpy as np
import time
import pylab as pl
from IPython import display

from tblite.ase import TBLite
from mace.calculators import MACECalculator

model_paths = ['MACE_models/mace02_com1_run-123_stagetwo.model',
               'MACE_models/mace02_com2_run-345_stagetwo.model',
               'MACE_models/mace02_com3_run-567_stagetwo.model']

xtb_calc = TBLite(method="GFN2-xTB",verbosity=0)
mace_calc = MACECalculator(model_paths=model_paths, device=device)

init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()
init_conf.calc = mace_calc

#initialize the temperature
np.random.seed(701)
MaxwellBoltzmannDistribution(init_conf, temperature_K=300)
Stationary(init_conf)
ZeroRotation(init_conf)

dyn = Langevin(init_conf, 1*units.fs, temperature_K=1200, friction=0.1)

%matplotlib inline

time_fs = []
temperature = []
energies_1 = []
energies_2 = []
energies_3 = []
variances = []
xtb_energies = []
true_errors = []

!rm -rfv moldyn/mace02_md_committee.xyz
fig, ax = pl.subplots(4, 1, figsize=(8,8), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})


def write_frame():
        at = dyn.atoms.copy()
        at.calc = xtb_calc
        xtb_energy = at.get_potential_energy()

        dyn.atoms.write('moldyn/mace02_md_committee.xyz', append=True, write_results=False)
        time_fs.append(dyn.get_time()/units.fs)
        temperature.append(dyn.atoms.get_temperature())
        energies_1.append(dyn.atoms.calc.results["energies"][0]/len(dyn.atoms))
        energies_2.append(dyn.atoms.calc.results["energies"][1]/len(dyn.atoms))
        energies_3.append(dyn.atoms.calc.results["energies"][2]/len(dyn.atoms))
        variances.append(dyn.atoms.calc.results["energy_var"]/len(dyn.atoms))
        xtb_energies.append(xtb_energy/len(dyn.atoms))
        true_errors.append(np.var([dyn.atoms.calc.results["energy"],xtb_energy])/len(dyn.atoms))

        # plot the true error
        ax[0].plot(np.array(time_fs), np.array(true_errors), color="black")
        ax[0].set_ylabel(r'$\Delta$ E (eV$^2$/atom)')
        ax[0].legend(['Error w.r.t. XTB'], loc='upper left')

        # plot committee variance
        ax[1].plot(np.array(time_fs), np.array(variances), color="y")
        ax[1].set_ylabel(r'committee variance')
        ax[1].legend(['Estimated Error (committee variances)'], loc='upper left')

        # plot the temperature of the system as subplots
        ax[2].plot(np.array(time_fs), temperature, color="r", label='Temperature')
        ax[2].set_ylabel("T (K)")

        ax[3].plot(np.array(time_fs), energies_1, color="g")
        ax[3].plot(np.array(time_fs), energies_2, color="y")
        ax[3].plot(np.array(time_fs), energies_3, color="olive")
        ax[3].plot(np.array(time_fs), xtb_energies, color="black")
        ax[3].set_ylabel("E (eV/atom)")
        ax[3].set_xlabel('Time (fs)')
        ax[3].legend(['E mace1', 'E mace2', 'E mace3', 'E xtb'], loc='upper left')

        display.clear_output(wait=True)
        display.display(fig)
        time.sleep(0.01)

dyn.attach(write_frame, interval=10)
dyn.run(500)
print("MD finished!")

NOTE: if you get the error xtb could not evalute the config that means the dynamics went so crazy and gave such strange configurations, that xtb refused to run! thats to be expected if the model is really bad. Copy some of the code above to have a look at the trajectory if you’d like to see this.

Closely observe the dynamics. Notice how good the committee error is as a proxy for the true error. In this case the true is cheap to compute, but in most practical applications it won’t be. Therefore, we will need to rely on the committee error to identy configurations that should be added back to the training set. This is called active learning:

alt text

Active learning in practice#

The way to use active learning to improve the model is as follows:

run dynmics, track the uncertainty.
if the uncertainty breaches some predeterined value, stop the simulation and peform the ground truth calculation.
and the new config to the dataset, and retrain
repeat steps 1-3 until the uncertainty never crosses the threshold

This can be done without human supervision - you can write a program which loops this process.

As an exercise at the end of the notebook, try writing an active learing loop to gradually grow the dataset and produce a good model, without ever running XTB dynamics.

3 Foundational Models#

3.1 Molecular Dynamics with MACE-MP-0#

Foundation models changed everything. MACE-MP-0 is a model trained on >1 million DFT calcuations, and can run dynamics for the whole periodic table.

Mace provides a simple interface to load a foundational model, which we can use mow. Check the documentation for more details.

from mace.calculators import mace_mp

macemp = mace_mp(model="medium-0b3",device=device)
init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()
simpleMD(init_conf, temp=1200, calc=macemp, fname='moldyn/mace03_md.xyz', s=10, T=2000)

The dynamics looks good. It is stable with an out-of-the-box model. Let’s inspect the trajectory:

from weas_widget import WeasWidget

traj = read('moldyn/mace03_md.xyz', ':')
w = WeasWidget()
w.from_ase(traj)
w.avr.model_style = 1
w.avr.show_hydrogen_bonds = True
w

3.2 Compare to XTB#

Let’s compute the radial distribution functions of the trajectory and compare them to XTB. Remember MACE-MP was trained on PBE level of theory so we don’t necessarily expect them to match:

from matplotlib import pyplot as plt
from aseMolec import anaAtoms as aa

tag = 'CC_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'

for f in ['xtb_md', 'mace03_md']:
    traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames
    for at in traj:
        at.pbc = True #create a fake box for rdf compatibility
        at.cell = [100,100,100]
    rdf = aa.compute_rdfs_traj_avg(traj, rmax=3, nbins=200) #aseMolec provides functionality to compute RDFs
    plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)

plt.legend();
plt.yticks([]);
plt.xlabel(r'R ($\rm \AA$)');
plt.ylabel('RDF '+tag);

Notice there’s a substantial shift in the C-C RDF peak between XTB and MACE-MP-0. This is likely due to the different level of reference theory. Can we fix by fine tuning MACE-MP-0?

Task for you

Recompute the XTB energy on the MACE-MP-0 trajectory, how does it compare?

Depending on your application, the PBE reference using in MACE-MP-0 may not be appropriate, however it is already better than XTB. XTB was chosen here to minimize computational cost, so we can check the model’s true error very easily.

In practice, you might want to run dynamics of this small molecule at the MP2 or coupled cluster level (which is the ‘gold standard’ of molecular quantum chemistry). In that case, you would want to finetune MACE-MP-0 onto a small amount of very expensive couple cluster data.

3.3 Fine tune MACE-MP to XTB#

As described in the previous section, MACE-MP-0 offers qualitatively good performance across a wide range of chemistries and materials at the PBE+U level of theory. However, for specific applications, it may be beneficial to fine-tune the model to improve its accuracy. The fine-tuning process involves training a pre-trained model on a new dataset, called the fine-tuning dataset, which contains a limited amount of data that are only relevant to the specific application.

There exists two ways of finetuning a MACE mode:

standard or naive approach which just restarts training using the parameters of the pretrained model
multi-head approach which ensures that the model does not forget too much about its pretraining.

We will start with the standard aproach below by setting multiheads_finetuning: False in the config. You will try for yourself to train another model with the multi-head approach and compare performances. Remember to change the name both in the config and when running MD.

%%writefile config/config-07.yml

model: "MACE"
stress_weight: 0.0
forces_weight: 10.0
energy_weight: 1.0
foundation_model: "/home/jovyan/.cache/mace/macemp0b3mediummodel"
multiheads_finetuning: False
name: "finetuned_standard_MACE"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
pt_train_file: mp
heads:
   xtb:
     train_file: "data/solvent_xtb_train_23.xyz"
     valid_file: "data/solvent_xtb_valid_20.xyz"
     test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
batch_size: 10
max_num_epochs: 100
swa: False
seed: 345

dev = f'device: {device}'
%store dev >>"config/config-07.yml"

train_mace("config/config-07.yml")

from IPython.display import Image, display
display(Image("MACE_models/finetuned_standard_MACE_run-345_train_xtb_stage_one.png"))

Compare the final accuracy to the models trained from scratch. You should see that the errors are much better when doing finetuning.

from IPython import display

mace_calc = MACECalculator(model_paths=['MACE_models/finetuned_standard_MACE_compiled.model'], device=device, dtype="float32")

init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()
simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace_finetuned_standard_md.xyz', s=10, T=2000)

from matplotlib import pyplot as plt
from aseMolec import anaAtoms as aa

tag = 'CC_intra' # 'OO_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'

for f in ['xtb_md', 'mace_finetuned_standard_md']:
    traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames
    for at in traj:
        at.pbc = True #create a fake box for rdf compatibility
        at.cell = [100,100,100]
    rdf = aa.compute_rdfs_traj_avg(traj, rmax=3, nbins=200)
    plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)

plt.legend();
plt.yticks([]);
plt.xlabel(r'R ($\rm \AA$)');
plt.ylabel('RDF '+tag);

from weas_widget import WeasWidget

traj = read('moldyn/mace_finetuned_standard_md.xyz', ':')
w = WeasWidget()
w.from_ase(traj)
w.avr.model_style = 1
w.avr.show_hydrogen_bonds = True
w

What are the results - does it work?

Tasks for you

comparing more than one of the rdfs
compare XTB energy. does the fine-tuned model perform better than the model trained from scratch?
retry with the multi-head fine-tuning strategy, how does that perform?
notice the fine-tuned model has different hyperparameters inhreted from the foundational model, is this a fair comparison? Consider training a model from scracth with the same parameters.

%%writefile config/config-08.yml

model: "MACE"
stress_weight: 0.0
forces_weight: 10.0
energy_weight: 1.0
foundation_model: "/home/jovyan/.cache/mace/macemp0b3mediummodel"
multiheads_finetuning: True
skip_evaluate_heads: False
name: "finetuned_multihead_MACE"
model_dir: "MACE_models"
log_dir: "MACE_models"
checkpoints_dir: "MACE_models"
results_dir: "MACE_models"
pt_train_file: mp
heads:
   xtb:
     train_file: "data/solvent_xtb_train_23.xyz"
     valid_file: "data/solvent_xtb_valid_20.xyz"
     test_file: "data/solvent_xtb_test.xyz"
energy_key: "energy_xtb"
forces_key: "forces_xtb"
batch_size: 10
lr: 0.0001
scaling: rms_forces_scaling
force_mh_ft_lr: True
ema_decay: 0.99999
max_num_epochs: 50
num_samples_pt: 10000
swa: False
enable_cueq: True
seed: 345

dev = f'device: {device}'
%store dev >>"config/config-08.yml"

train_mace("config/config-08.yml")

MACE in Practice II

Contents

MACE in Practice II#

Learning Objectives for today:#

1. Iterative Training#

1.1 MD with a smaller MACE model#

1.2 Identify the problem and expand training#

1.3 Train a new MACE model and run MD again#

2. Active Learning with MACE#

2.1 Preparing a committee of models#

2.2 Running MD with the MACE committee#

Active learning in practice#

3 Foundational Models#

3.1 Molecular Dynamics with MACE-MP-0#

3.2 Compare to XTB#

3.3 Fine tune MACE-MP to XTB#