{
"cells": [
{
"cell_type": "markdown",
"metadata": {
"id": "XXosxKoCxNVG"
},
"source": [
"# MACE in Practice II"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"[](https://colab.research.google.com/github/ddmms/camml-tutorials/blob/main/notebooks/day-4/T02-MACE-Practice-II.ipynb)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "XYamp68VxNVI"
},
"source": [
"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](https://matbench-discovery.materialsproject.org/) - the latest development in the field of MLIPs. These models are trained on massive training sets of [inorganic](https://doi.org/10.48550/arXiv.2401.00096) and [organic](https://doi.org/10.48550/arXiv.2312.15211) databases and show a great deal of `out-of-the-box` MD stability in an extensive variety of [applications](https://doi.org/10.48550/arXiv.2401.00096). We will discuss [fine-tunning](https://doi.org/10.48550/arXiv.2405.20217) which is an actively-researched technique to tweak these foundational models to new systems (out of training) and/or new levels of reference theory.\n",
"\n",
"## Learning Objectives for today:\n",
"\n",
"1. **Iterative Training: improving stability and accuracy**\n",
"2. **Active learning: committee models**\n",
"3. **Foundational models and fine-tuning**"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KSAoZccQxNVJ"
},
"source": [
"## 1. Iterative Training"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8ieIGRPOxNVK"
},
"source": [
"### 1.1 MD with a smaller MACE model"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "UTK__Hd8xNVK"
},
"source": [
"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.\n",
"\n",
"**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.**"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "9FRg6XJ0D2yG"
},
"outputs": [],
"source": [
"from ase.io import read, write\n",
"\n",
"device = 'cuda'\n",
"# device = 'cpu'\n",
"\n",
"db = read('data/solvent_xtb.xyz', ':')\n",
"write('data/solvent_xtb_train_23.xyz', db[:23]) #first 20 configs plus the 3 E0s will be used for training (parameter optimization)\n",
"write('data/solvent_xtb_valid_20.xyz', db[23:43]) #next 20 configs will be used for validation (independent, stop condition)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "VleNMLWZD2yH"
},
"source": [
"Let's make an input config."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Ek3iJ6rFD2yH",
"outputId": "138ddcc9-7a15-45ce-fe2a-ca272d612725"
},
"outputs": [],
"source": [
"%%writefile config/config-03.yml\n",
"\n",
"model: \"MACE\"\n",
"num_channels: 32\n",
"max_L: 0\n",
"r_max: 4.0\n",
"max_ell: 2\n",
"name: \"mace02_com1\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"train_file: \"data/solvent_xtb_train_23.xyz\"\n",
"valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
"test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"batch_size: 10\n",
"max_num_epochs: 200\n",
"swa: True\n",
"seed: 123\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-03.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "QiAR1bg1D2yH",
"outputId": "6a4d18b1-2fae-4b62-b1a1-597b434a1f62"
},
"outputs": [],
"source": [
"import warnings\n",
"warnings.filterwarnings(\"ignore\")\n",
"from mace.cli.run_train import main as mace_run_train_main\n",
"import sys\n",
"import logging\n",
"\n",
"def train_mace(config_file_path):\n",
" logging.getLogger().handlers.clear()\n",
" sys.argv = [\"program\", \"--config\", config_file_path]\n",
" mace_run_train_main()"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"colab": {
"base_uri": "https://localhost:8080/"
},
"id": "Cl9JEPvuD2yI",
"outputId": "166104df-10db-47a6-e970-ea361ae55868"
},
"outputs": [],
"source": [
"train_mace(\"config/config-03.yml\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "gyJtdwaIxNVK"
},
"outputs": [],
"source": [
"#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n",
"import glob\n",
"import os\n",
"for file in glob.glob(\"MACE_models/*.pt\"):\n",
" os.remove(file)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "30ij5FnzxNVL"
},
"source": [
"Notice, we are getting substantially larger errors than before. Now, let's run some dynamics:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "l3Y0XqKCxNVL"
},
"outputs": [],
"source": [
"from ase.io import read, write\n",
"from ase import units\n",
"from ase.md.langevin import Langevin\n",
"from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution\n",
"from aseMolec import extAtoms as ea\n",
"\n",
"import os\n",
"import time\n",
"import numpy as np\n",
"import pylab as pl\n",
"from IPython import display\n",
"np.random.seed(701) #just making sure the MD failure is reproducible\n",
"\n",
"def simpleMD(init_conf, temp, calc, fname, s, T):\n",
" init_conf.set_calculator(calc)\n",
"\n",
" #initialize the temperature\n",
"\n",
" MaxwellBoltzmannDistribution(init_conf, temperature_K=300) #initialize temperature at 300\n",
" Stationary(init_conf)\n",
" ZeroRotation(init_conf)\n",
"\n",
" dyn = Langevin(init_conf, 1.0*units.fs, temperature_K=temp, friction=0.1) #drive system to desired temperature\n",
"\n",
" %matplotlib inline\n",
"\n",
" time_fs = []\n",
" temperature = []\n",
" energies = []\n",
"\n",
" #remove previously stored trajectory with the same name\n",
" os.system('rm -rfv '+fname)\n",
"\n",
" fig, ax = pl.subplots(2, 1, figsize=(6,6), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})\n",
"\n",
" def write_frame():\n",
" dyn.atoms.info['energy_mace'] = dyn.atoms.get_potential_energy()\n",
" dyn.atoms.arrays['force_mace'] = dyn.atoms.calc.get_forces()\n",
" dyn.atoms.write(fname, append=True)\n",
" time_fs.append(dyn.get_time()/units.fs)\n",
" temperature.append(dyn.atoms.get_temperature())\n",
" energies.append(dyn.atoms.get_potential_energy()/len(dyn.atoms))\n",
"\n",
" ax[0].plot(np.array(time_fs), np.array(energies), color=\"b\")\n",
" ax[0].set_ylabel('E (eV/atom)')\n",
"\n",
" # plot the temperature of the system as subplots\n",
" ax[1].plot(np.array(time_fs), temperature, color=\"r\")\n",
" ax[1].set_ylabel('T (K)')\n",
" ax[1].set_xlabel('Time (fs)')\n",
"\n",
" display.clear_output(wait=True)\n",
" display.display(pl.gcf())\n",
" time.sleep(0.01)\n",
"\n",
" dyn.attach(write_frame, interval=s)\n",
" t0 = time.time()\n",
" dyn.run(T)\n",
" t1 = time.time()\n",
" print(\"MD finished in {0:.2f} minutes!\".format((t1-t0)/60))"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"#let us start with a single molecule\n",
"init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n",
"\n",
"#we can use MACE as a calculator in ASE!\n",
"from mace.calculators import MACECalculator\n",
"mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_run-123_stagetwo.model'], device=device)\n",
"\n",
"simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md.xyz', s=10, T=2000)"
]
},
{
"cell_type": "raw",
"metadata": {
"id": "O9elyP0FxNVM"
},
"source": [
"Depending on the random seed for the MD, you may see different things. At first sight, the energy against time doesn't look too bad. Let's look at the trajectory."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "22uVfg2UxNVM"
},
"outputs": [],
"source": [
"from ase.io import read, write\n",
"import nglview as nv\n",
"\n",
"traj = read('moldyn/mace02_md.xyz', ':')\n",
"view = nv.show_asetraj(traj)\n",
"view._set_size(w='100%', h='500px')\n",
"view"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "hBbdTtPTxNVN"
},
"source": [
"If you go to the end of the trajectory, you should find that the bond angles are actually very strange - it looks unphsyical."
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"**Task for you**\n",
"\n",
"\n",
"\n",
"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.\n",
"
"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "lwFdVZGSxNVN"
},
"source": [
"### 1.2 Identify the problem and expand training\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "WdwunhD4xNVN"
},
"outputs": [],
"source": [
"from ase.io import read, write\n",
"from tqdm import tqdm\n",
"from tblite.ase import TBLite\n",
"xtb_calc = TBLite(method=\"GFN2-xTB\",verbosity=0,max_iterations=1000)\n",
"\n",
"#compute true reference XTB values\n",
"print(\"Evaluating MACE configurations with XTB\")\n",
"traj = read('moldyn/mace02_md.xyz', ':')\n",
"for at in tqdm(traj[:100]):\n",
" at.calc = TBLite(method=\"GFN2-xTB\",verbosity=0)\n",
" at.info['energy_xtb'] = at.get_potential_energy()\n",
" at.arrays['forces_xtb'] = at.get_forces()\n",
" at.calc = None\n",
"\n",
"write('data/mace02_md_100_xtb.xyz', traj[:100]) #save full result"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "a9lNJVNbxNVN"
},
"outputs": [],
"source": [
"import numpy as np\n",
"from aseMolec import extAtoms as ea\n",
"from matplotlib import pyplot as plt\n",
"\n",
"traj = read('data/mace02_md_100_xtb.xyz', ':')\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB')\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace', peratom=True), label='MACE')\n",
"plt.legend()\n",
"plt.xlabel('Time (fs)')\n",
"plt.ylabel('Total Energy per Atom (eV)')"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "V835Pt7KxNVN"
},
"source": [
"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:\n",
""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "B0qytkOCxNVO"
},
"outputs": [],
"source": [
"db = read('data/solvent_xtb_train_23.xyz', ':')\n",
"db += traj[40:100:5] #add failed configs to the training set\n",
"write('data/solvent_xtb_train_35_gen1.xyz', db)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "TU0SLvwFD2yK"
},
"source": [
"A good exercise is to try to pick the configs that have the largest error on forces and energies. Try coding it yourself."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "7Gbxj436xNVO"
},
"source": [
"### 1.3 Train a new MACE model and run MD again"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "TirW5fX-D2yK"
},
"outputs": [],
"source": [
"%%writefile config/config-04.yml\n",
"\n",
"model: \"MACE\"\n",
"num_channels: 32\n",
"max_L: 0\n",
"r_max: 4.0\n",
"max_ell: 2\n",
"name: \"mace02_com1_gen1\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"train_file: \"data/solvent_xtb_train_35_gen1.xyz\"\n",
"valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
"test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"batch_size: 10\n",
"max_num_epochs: 200\n",
"swa: True\n",
"seed: 123\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-04.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "SR1oyfPfD2yK"
},
"outputs": [],
"source": [
"train_mace(\"config/config-04.yml\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "knsX6u-ZxNVO"
},
"outputs": [],
"source": [
"#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n",
"import glob\n",
"import os\n",
"for file in glob.glob(\"MACE_models/*.pt\"):\n",
" os.remove(file)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "yfYqcls4xNVO"
},
"source": [
"For some reason the energy error on the training set is now huge - can you work out why this is?\n",
"\n",
"What does this imply about how we do iterative training?"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "VyJ_sBJTxNVO"
},
"outputs": [],
"source": [
"from mace.calculators import MACECalculator\n",
"\n",
"mace_calc = MACECalculator(model_paths=['MACE_models/mace02_com1_gen1_run-123_stagetwo.model'], device=device)\n",
"init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n",
"simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace02_md_gen1.xyz', s=10, T=2000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "RDR3fOAVxNVO"
},
"outputs": [],
"source": [
"from weas_widget import WeasWidget\n",
"\n",
"traj = read('moldyn/mace02_md_gen1.xyz', ':')\n",
"w = WeasWidget()\n",
"w.from_ase(traj)\n",
"w.avr.model_style = 1\n",
"w.avr.show_hydrogen_bonds = True\n",
"w"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "EuaTYEsjxNVP"
},
"source": [
"Great! The dynamics is already looking better, however it can be difficult to tell if it is really correct or not.\n",
"\n",
"**Task for you**\n",
"\n",
" \n",
"- Compute the true XTB enery and check if they agree with the MACE energy predicted along the trajectory\n",
"- Compute radial distribution functions (RDFs) and compared to a ground truth RDFs\n",
"- If the MACE trajectory does not agree with the ground truth, continue the iterative process and gradually improve the potential\n",
"- For every new potential, repeat the dynamics and check against grount truth observables.\n",
"
\n",
"\n",
"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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "DpH3OgAyxNVP"
},
"source": [
"## 2. Active Learning with MACE"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "IFMwXevOxNVP"
},
"source": [
"### 2.1 Preparing a committee of models"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "PzmSnYbAxNVP"
},
"source": [
"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:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "-C7jn1B8D2yM"
},
"outputs": [],
"source": [
"%%writefile config/config-05.yml\n",
"\n",
"model: \"MACE\"\n",
"num_channels: 32\n",
"max_L: 0\n",
"r_max: 4.0\n",
"name: \"mace02_com2\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"train_file: \"data/solvent_xtb_train_23.xyz\"\n",
"valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
"test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"batch_size: 10\n",
"max_num_epochs: 200\n",
"swa: True\n",
"seed: 345"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-05.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "JMDC9wBWD2yM"
},
"outputs": [],
"source": [
"%%writefile config/config-06.yml\n",
"\n",
"model: \"MACE\"\n",
"num_channels: 32\n",
"max_L: 0\n",
"r_max: 4.0\n",
"name: \"mace02_com3\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"train_file: \"data/solvent_xtb_train_23.xyz\"\n",
"valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
"test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"device: cuda\n",
"batch_size: 10\n",
"max_num_epochs: 200\n",
"swa: True\n",
"seed: 567"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-06.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "2-d7doFND2yM"
},
"outputs": [],
"source": [
"train_mace(\"config/config-05.yml\")\n",
"train_mace(\"config/config-06.yml\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "FoL5JoLExNVP"
},
"outputs": [],
"source": [
"#remove checkpoints since they may cause errors on retraining a model with the same name but a different architecture\n",
"import glob\n",
"import os\n",
"for file in glob.glob(\"MACE_models/*.pt\"):\n",
" os.remove(file)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "09GP4CZVxNVP"
},
"source": [
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "U2SaupbgxNVP"
},
"outputs": [],
"source": [
"import numpy as np\n",
"from ase.io import read\n",
"from aseMolec import extAtoms as ea\n",
"from matplotlib import pyplot as plt\n",
"from mace.calculators import MACECalculator\n",
"from tqdm import tqdm\n",
"import warnings\n",
"warnings.filterwarnings(\"ignore\")\n",
"\n",
"model_paths = ['MACE_models/mace02_com1_run-123_stagetwo.model',\n",
" 'MACE_models/mace02_com2_run-345_stagetwo.model',\n",
" 'MACE_models/mace02_com3_run-567_stagetwo.model',]\n",
"mace_calcs = MACECalculator(model_paths=model_paths, device=device)\n",
"\n",
"traj = read('data/mace02_md_100_xtb.xyz', ':')\n",
"for at in tqdm(traj):\n",
" at.calc = mace_calcs\n",
" engs = at.get_potential_energies()\n",
" at.info['energy_mace_1'] = at.info.pop('energy_mace') #rename value obtained with first member of the committee\n",
" at.info['energy_mace_2'] = engs[1]\n",
" at.info['energy_mace_3'] = engs[2]\n",
" at.info['variance'] = np.std(engs)\n",
" at.info['average_mace_energy'] = np.average(engs)\n",
" at.info['true_error'] = np.abs(at.info['average_mace_energy'] - at.info['energy_xtb'])\n",
"\n",
"#Let's check the energies of the MACE committee vs XTB energy\n",
"plt.figure(figsize=(12,6))\n",
"plt.subplot(1,2,1)\n",
"\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_xtb', peratom=True), label='XTB');\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_1', peratom=True), label='MACE_1');\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_2', peratom=True), label='MACE_2');\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'energy_mace_3', peratom=True), label='MACE_3');\n",
"plt.legend()\n",
"plt.xlabel('Time (fs)');\n",
"plt.ylabel('Energy per Atom (eV)');\n",
"\n",
"plt.subplot(1,2,2)\n",
"plt.plot(np.arange(len(traj)), ea.get_prop(traj, 'info', 'variance', peratom=True), label='committee variance', color='black');\n",
"plt.xlabel('Time (fs)');\n",
"plt.ylabel('Committee Energy Variance');\n"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "5yXhlCtVxNVQ"
},
"source": [
"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.\n",
"\n",
"**Task for you**\n",
"\n",
"\n",
"- Make a correlation plot between true error and estimated error. How well does the estimation perform?\n",
"- Try a commitee of models with both different seeds and trained on different selections of data. Does it perform better?\n",
"
\n",
"\n",
"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?"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "ZPX9TP7fxNVQ"
},
"source": [
"### 2.2 Running MD with the MACE committee"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yOSnpIJUxNVQ"
},
"outputs": [],
"source": [
"from aseMolec import extAtoms as ea\n",
"from ase import units\n",
"from ase.md.langevin import Langevin\n",
"from ase.md.velocitydistribution import Stationary, ZeroRotation, MaxwellBoltzmannDistribution\n",
"from ase.io import read, write\n",
"\n",
"import random\n",
"import numpy as np\n",
"import time\n",
"import pylab as pl\n",
"from IPython import display\n",
"\n",
"from tblite.ase import TBLite\n",
"from mace.calculators import MACECalculator\n",
"\n",
"model_paths = ['MACE_models/mace02_com1_run-123_stagetwo.model',\n",
" 'MACE_models/mace02_com2_run-345_stagetwo.model',\n",
" 'MACE_models/mace02_com3_run-567_stagetwo.model']\n",
"\n",
"xtb_calc = TBLite(method=\"GFN2-xTB\",verbosity=0)\n",
"mace_calc = MACECalculator(model_paths=model_paths, device=device)\n",
"\n",
"init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n",
"init_conf.calc = mace_calc\n",
"\n",
"#initialize the temperature\n",
"np.random.seed(701)\n",
"MaxwellBoltzmannDistribution(init_conf, temperature_K=300)\n",
"Stationary(init_conf)\n",
"ZeroRotation(init_conf)\n",
"\n",
"dyn = Langevin(init_conf, 1*units.fs, temperature_K=1200, friction=0.1)\n",
"\n",
"%matplotlib inline\n",
"\n",
"time_fs = []\n",
"temperature = []\n",
"energies_1 = []\n",
"energies_2 = []\n",
"energies_3 = []\n",
"variances = []\n",
"xtb_energies = []\n",
"true_errors = []\n",
"\n",
"!rm -rfv moldyn/mace02_md_committee.xyz\n",
"fig, ax = pl.subplots(4, 1, figsize=(8,8), sharex='all', gridspec_kw={'hspace': 0, 'wspace': 0})\n",
"\n",
"\n",
"def write_frame():\n",
" at = dyn.atoms.copy()\n",
" at.calc = xtb_calc\n",
" xtb_energy = at.get_potential_energy()\n",
"\n",
" dyn.atoms.write('moldyn/mace02_md_committee.xyz', append=True, write_results=False)\n",
" time_fs.append(dyn.get_time()/units.fs)\n",
" temperature.append(dyn.atoms.get_temperature())\n",
" energies_1.append(dyn.atoms.calc.results[\"energies\"][0]/len(dyn.atoms))\n",
" energies_2.append(dyn.atoms.calc.results[\"energies\"][1]/len(dyn.atoms))\n",
" energies_3.append(dyn.atoms.calc.results[\"energies\"][2]/len(dyn.atoms))\n",
" variances.append(dyn.atoms.calc.results[\"energy_var\"]/len(dyn.atoms))\n",
" xtb_energies.append(xtb_energy/len(dyn.atoms))\n",
" true_errors.append(np.var([dyn.atoms.calc.results[\"energy\"],xtb_energy])/len(dyn.atoms))\n",
"\n",
" # plot the true error\n",
" ax[0].plot(np.array(time_fs), np.array(true_errors), color=\"black\")\n",
" ax[0].set_ylabel(r'$\\Delta$ E (eV$^2$/atom)')\n",
" ax[0].legend(['Error w.r.t. XTB'], loc='upper left')\n",
"\n",
" # plot committee variance\n",
" ax[1].plot(np.array(time_fs), np.array(variances), color=\"y\")\n",
" ax[1].set_ylabel(r'committee variance')\n",
" ax[1].legend(['Estimated Error (committee variances)'], loc='upper left')\n",
"\n",
" # plot the temperature of the system as subplots\n",
" ax[2].plot(np.array(time_fs), temperature, color=\"r\", label='Temperature')\n",
" ax[2].set_ylabel(\"T (K)\")\n",
"\n",
" ax[3].plot(np.array(time_fs), energies_1, color=\"g\")\n",
" ax[3].plot(np.array(time_fs), energies_2, color=\"y\")\n",
" ax[3].plot(np.array(time_fs), energies_3, color=\"olive\")\n",
" ax[3].plot(np.array(time_fs), xtb_energies, color=\"black\")\n",
" ax[3].set_ylabel(\"E (eV/atom)\")\n",
" ax[3].set_xlabel('Time (fs)')\n",
" ax[3].legend(['E mace1', 'E mace2', 'E mace3', 'E xtb'], loc='upper left')\n",
"\n",
" display.clear_output(wait=True)\n",
" display.display(fig)\n",
" time.sleep(0.01)\n",
"\n",
"dyn.attach(write_frame, interval=10)\n",
"dyn.run(500)\n",
"print(\"MD finished!\")"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "8MNDXXavxNVQ"
},
"source": [
"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.\n",
"\n",
"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:"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "aqbEUoAPxNVR"
},
"source": [
""
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "2kkq8XlAxNVR"
},
"source": [
"### Active learning in practice\n",
"\n",
"The way to use active learning to improve the model is as follows:\n",
"1. run dynmics, track the uncertainty.\n",
"2. if the uncertainty breaches some predeterined value, stop the simulation and peform the ground truth calculation.\n",
"3. and the new config to the dataset, and retrain\n",
"4. repeat steps 1-3 until the uncertainty never crosses the threshold\n",
"\n",
"This can be done without human supervision - you can write a program which loops this process.\n",
"\n",
"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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "P9x25o-PxNVR"
},
"source": [
"## 3 Foundational Models"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qoh5bDxqxNVR"
},
"source": [
"### 3.1 Molecular Dynamics with MACE-MP-0"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "maotXGR3xNVR"
},
"source": [
"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.\n",
"\n",
"Mace provides a simple interface to load a foundational model, which we can use mow. Check the [documentation](https://mace-docs.readthedocs.io/en/latest/guide/foundation_models.html) for more details."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "ySF0cKm1xNVR"
},
"outputs": [],
"source": [
"from mace.calculators import mace_mp\n",
"\n",
"macemp = mace_mp(model=\"medium-0b3\",device=device)\n",
"init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n",
"simpleMD(init_conf, temp=1200, calc=macemp, fname='moldyn/mace03_md.xyz', s=10, T=2000)"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "KRMKjfYBxNVW"
},
"source": [
"The dynamics looks good. It is stable with an out-of-the-box model. Let's inspect the trajectory:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "PK9nLfWPxNVW"
},
"outputs": [],
"source": [
"from weas_widget import WeasWidget\n",
"\n",
"traj = read('moldyn/mace03_md.xyz', ':')\n",
"w = WeasWidget()\n",
"w.from_ase(traj)\n",
"w.avr.model_style = 1\n",
"w.avr.show_hydrogen_bonds = True\n",
"w"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "idQExiN2xNVW"
},
"source": [
"### 3.2 Compare to XTB"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "nlh6ZyTGxNVW"
},
"source": [
"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:"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "3W2GjUexxNVX"
},
"outputs": [],
"source": [
"from matplotlib import pyplot as plt\n",
"from aseMolec import anaAtoms as aa\n",
"\n",
"tag = 'CC_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'\n",
"\n",
"for f in ['xtb_md', 'mace03_md']:\n",
" traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames\n",
" for at in traj:\n",
" at.pbc = True #create a fake box for rdf compatibility\n",
" at.cell = [100,100,100]\n",
" rdf = aa.compute_rdfs_traj_avg(traj, rmax=3, nbins=200) #aseMolec provides functionality to compute RDFs\n",
" plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)\n",
"\n",
"plt.legend();\n",
"plt.yticks([]);\n",
"plt.xlabel(r'R ($\\rm \\AA$)');\n",
"plt.ylabel('RDF '+tag);"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "16UZzmpexNVX"
},
"source": [
"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?\n",
"\n",
"**Task for you**\n",
"\n",
"\n",
"Recompute the XTB energy on the MACE-MP-0 trajectory, how does it compare?\n",
"
\n",
"\n",
"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.\n",
"\n",
"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."
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "H_kSVlzxxNVX"
},
"source": [
"### 3.3 Fine tune MACE-MP to XTB"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "vF-sDD9ED2yP"
},
"source": [
"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.\n",
"However, for specific applications, it may be beneficial to fine-tune the model to improve its accuracy.\n",
"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.\n",
"\n",
"There exists two ways of finetuning a MACE mode:\n",
"1. ***standard or naive*** approach which just restarts training using the parameters of the pretrained model\n",
"2. ***multi-head*** approach which ensures that the model does not forget too much about its pretraining.\n",
"\n",
"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."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "0mTww8OmD2yP"
},
"outputs": [],
"source": [
"%%writefile config/config-07.yml\n",
"\n",
"model: \"MACE\"\n",
"stress_weight: 0.0\n",
"forces_weight: 10.0\n",
"energy_weight: 1.0\n",
"foundation_model: \"/home/jovyan/.cache/mace/macemp0b3mediummodel\"\n",
"multiheads_finetuning: False\n",
"name: \"finetuned_standard_MACE\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"pt_train_file: mp\n",
"heads:\n",
" xtb:\n",
" train_file: \"data/solvent_xtb_train_23.xyz\"\n",
" valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
" test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"batch_size: 10\n",
"max_num_epochs: 100\n",
"swa: False\n",
"seed: 345"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-07.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "yc4sneynxNVX"
},
"outputs": [],
"source": [
"train_mace(\"config/config-07.yml\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"from IPython.display import Image, display\n",
"display(Image(\"MACE_models/finetuned_standard_MACE_run-345_train_xtb_stage_one.png\"))"
]
},
{
"cell_type": "markdown",
"metadata": {
"id": "qT23YJaBD2yP"
},
"source": [
"Compare the final accuracy to the models trained from scratch. You should see that the errors are much better when doing finetuning."
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "iBaqOL9_xNVX"
},
"outputs": [],
"source": [
"from IPython import display\n",
"\n",
"mace_calc = MACECalculator(model_paths=['MACE_models/finetuned_standard_MACE_compiled.model'], device=device, dtype=\"float32\")\n",
"\n",
"init_conf = ea.sel_by_info_val(read('data/solvent_molecs.xyz',':'), 'Nmols', 1)[0].copy()\n",
"simpleMD(init_conf, temp=1200, calc=mace_calc, fname='moldyn/mace_finetuned_standard_md.xyz', s=10, T=2000)"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "61qwBSW3xNVY"
},
"outputs": [],
"source": [
"from matplotlib import pyplot as plt\n",
"from aseMolec import anaAtoms as aa\n",
"\n",
"tag = 'CC_intra' # 'OO_intra' #choose one of 'HH_intra', 'HC_intra', 'HO_intra', 'CC_intra', 'CO_intra', 'OO_intra'\n",
"\n",
"for f in ['xtb_md', 'mace_finetuned_standard_md']:\n",
" traj = read('moldyn/'+f+'.xyz', '50:') #ignore first 50 frames\n",
" for at in traj:\n",
" at.pbc = True #create a fake box for rdf compatibility\n",
" at.cell = [100,100,100]\n",
" rdf = aa.compute_rdfs_traj_avg(traj, rmax=3, nbins=200)\n",
" plt.plot(rdf[1], rdf[0][tag], '.-', label=f, alpha=0.7, linewidth=3)\n",
"\n",
"plt.legend();\n",
"plt.yticks([]);\n",
"plt.xlabel(r'R ($\\rm \\AA$)');\n",
"plt.ylabel('RDF '+tag);"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {
"id": "nGjLB6hvxNVY"
},
"outputs": [],
"source": [
"from weas_widget import WeasWidget\n",
"\n",
"traj = read('moldyn/mace_finetuned_standard_md.xyz', ':')\n",
"w = WeasWidget()\n",
"w.from_ase(traj)\n",
"w.avr.model_style = 1\n",
"w.avr.show_hydrogen_bonds = True\n",
"w"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"What are the results - does it work?\n",
"\n",
"**Tasks for you**\n",
"\n",
"\n",
"- comparing more than one of the rdfs\n",
"- compare XTB energy. does the fine-tuned model perform better than the model trained from scratch?\n",
"- retry with the multi-head fine-tuning strategy, how does that perform?\n",
"- 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.\n",
"
"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"%%writefile config/config-08.yml\n",
"\n",
"model: \"MACE\"\n",
"stress_weight: 0.0\n",
"forces_weight: 10.0\n",
"energy_weight: 1.0\n",
"foundation_model: \"/home/jovyan/.cache/mace/macemp0b3mediummodel\"\n",
"multiheads_finetuning: True\n",
"skip_evaluate_heads: False\n",
"name: \"finetuned_multihead_MACE\"\n",
"model_dir: \"MACE_models\"\n",
"log_dir: \"MACE_models\"\n",
"checkpoints_dir: \"MACE_models\"\n",
"results_dir: \"MACE_models\"\n",
"pt_train_file: mp\n",
"heads:\n",
" xtb:\n",
" train_file: \"data/solvent_xtb_train_23.xyz\"\n",
" valid_file: \"data/solvent_xtb_valid_20.xyz\"\n",
" test_file: \"data/solvent_xtb_test.xyz\"\n",
"energy_key: \"energy_xtb\"\n",
"forces_key: \"forces_xtb\"\n",
"batch_size: 10\n",
"lr: 0.0001\n",
"scaling: rms_forces_scaling\n",
"force_mh_ft_lr: True\n",
"ema_decay: 0.99999\n",
"max_num_epochs: 50\n",
"num_samples_pt: 10000\n",
"swa: False\n",
"enable_cueq: True\n",
"seed: 345"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"dev = f'device: {device}'\n",
"%store dev >>\"config/config-08.yml\""
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": [
"train_mace(\"config/config-08.yml\")"
]
}
],
"metadata": {
"accelerator": "GPU",
"colab": {
"gpuType": "T4",
"machine_shape": "hm",
"provenance": []
},
"deepnote": {
"is_reactive": false
},
"deepnote_execution_queue": [],
"deepnote_notebook_id": "2a16e1e8a7ad42a8825007f5cff15d19",
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.12.9"
}
},
"nbformat": 4,
"nbformat_minor": 0
}