Skip to content

Commit

Permalink
Account for review comments
Browse files Browse the repository at this point in the history
  • Loading branch information
@PythonFZ
PythonFZ committed Aug 29, 2024
1 parent 6f0494c commit e49ab3c
Showing 1 changed file with 55 additions and 27 deletions.
82 changes: 55 additions & 27 deletions doc/tutorials/mlip/mlip.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
"cell_type": "markdown",
"metadata": {},
"source": [
"# Machine-Learning Interatomic Potentials\n",
"# Machine-learned interatomic potential\n",
"\n",
"Running simulations with energies and forces from *ab initio* methods, such as Density Functional Theory (DFT) allows for higher accuracy and the investigation of different systems. E.g., *ab initio* MD can be used to study chemical reactions on an atomistic scale. With the recent advancements in machine-learning interatomic potentials (MLIPs) it is possible to scale up these types of simulations to large systems [Kozinsky et al.](https://doi.org/10.1145/3581784.3627041) and long simulation times [Zills et al.](https://doi.org/10.1039/d4fd00025k) while retaining the accuracy of the underlying *ab initio* methods.\n",
"\n",
Expand All @@ -16,8 +16,10 @@
"- TorchANI https://github.com/aiqm/torchani\n",
"\n",
"In this tutorial, we will focus on the [MACE-MP-0](https://doi.org/10.48550/arXiv.2401.00096) model, which has been trained on the [materials project dataset](https://doi.org/10.1063/1.4812323) and is able to accurately predict energies and forces for many organic and inorganic systems.\n",
"The MACE-MP-0 model uses the state-of-the-art MACE equivariant message passing model architecture.\n",
"\n",
"Most MLIPs provide a Python interface through the ASE package. ESPResSo also utilizes this interface, and thus we will have a quick overview of the ASE features. The main object used is the `ase.Atoms`, which represents a single frame of an atomistic simulation."
"Most MLIPs provide a Python interface through the ASE package. ESPResSo also utilizes this interface, and thus we will have a quick overview of the ASE features. The main object used is the `ase.Atoms`, which represents a single frame of an atomistic simulation.\n",
"We will attach the MACE-MP-0 `Calculator` to the `ase.Atoms` object for the computation of energies (`atoms.get_potential_energy() -> float`) and forces (`atoms.get_forces() -> np.ndarray`). More information can be found in the ASE documentation https://wiki.fysik.dtu.dk/ase/ase/atoms.html."
]
},
{
Expand All @@ -33,23 +35,30 @@
"metadata": {},
"outputs": [],
"source": [
"# ESPResSo imports\n",
"import espressomd\n",
"from espressomd.plugins.ase import ASEInterface\n",
"from espressomd.zn import Visualizer\n",
"\n",
"# MACE-MP-0\n",
"from mace.calculators import mace_mp\n",
"\n",
"# Simulation box\n",
"from rdkit2ase import smiles2atoms, pack\n",
"\n",
"# Miscellaneous\n",
"import numpy as np\n",
"import pandas as pd\n",
"import pint\n",
"import plotly.express as px\n",
"import tqdm\n",
"from espressomd.plugins.ase import ASEInterface\n",
"from espressomd.zn import Visualizer\n",
"from mace.calculators import mace_mp\n",
"from rdkit2ase import smiles2atoms, pack"
"import tqdm"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"For MLIP the typical unit system is energies in eV, distances in Angstrom, time in fs and mass in atomic units. This is also the default unit system within ASE. We will use pint to ensure correct unit usage within ESPResSo."
"For MLIP the typical unit system is energies in eV, distances in Angstrom, time in fs and mass in atomic units. This is also the default unit system within ASE. We will use pint to ensure correct unit usage within ESPResSo. For setting the simulation temperature, we will also convert the boltzmann constant `boltzmann_k` to our desired unit system."
]
},
{
Expand All @@ -59,8 +68,8 @@
"outputs": [],
"source": [
"ureg = pint.UnitRegistry()\n",
"boltmk = 1.380649e-23 * ureg.J / ureg.K\n",
"boltmk = boltmk.to(ureg.eV / ureg.K)"
"boltzmann_k = 1.380649e-23 * ureg.J / ureg.K\n",
"boltzmann_k = boltzmann_k.to(ureg.eV / ureg.K)"
]
},
{
Expand All @@ -69,8 +78,15 @@
"source": [
"In this tutorial we will use [Simplified Molecular Input Line Entry System](https://doi.org/10.1021/ci00057a005) (SMILES) representations and [RDKit](https://github.com/rdkit/rdkit) to generate starting structures.\n",
"RDKit provides us with a powerful utilities to create 3D structures from these string representations of molecules.\n",
"\n",
"![ethanol structure](https://upload.wikimedia.org/wikipedia/commons/thumb/b/be/Ethanol_Lewis.svg/320px-Ethanol_Lewis.svg.png)\n",
"\n",
"Looking at the structure of ethanol, the respective SMILES string is given, by removing all the hydrogens and writing down the atomic symbols as `CCO`.\n",
"Note that SMILES are also capable of representing e.g. cyclic structures, double bonds and other chemically relevant properties, which are not required when representing ethanol.\n",
"\n",
"\n",
"In this first part, we will look at a gas phase system of Ethanol, run a geometry optimization followed by a short MD simulation.\n",
"We will start by instantiate our MLIP."
"Let us first create our ASE atoms object to define the sysyem we want to investigate."
]
},
{
Expand All @@ -79,14 +95,15 @@
"metadata": {},
"outputs": [],
"source": [
"mace_mp_calc = mace_mp()"
"ethanol_atoms = smiles2atoms(\"CCO\")\n",
"ethanol_atoms"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can now create our ASE atoms object and attach the aforementioned MLIP calculator to it. "
"We can no initialize the MACE-MP-0 calculator and attach it to the `ase.Atoms` object."
]
},
{
Expand All @@ -95,16 +112,15 @@
"metadata": {},
"outputs": [],
"source": [
"etoh = smiles2atoms(\"CCO\")\n",
"etoh.calc = mace_mp_calc\n",
"etoh"
"mace_mp_calc = mace_mp()\n",
"ethanol_atoms.calc = mace_mp_calc"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We can now levarage the ase interface to compute energies and forces."
"We can now leverage the ase interface to compute energies and forces."
]
},
{
Expand All @@ -113,8 +129,8 @@
"metadata": {},
"outputs": [],
"source": [
"print(etoh.get_potential_energy() * ureg.eV)\n",
"print(etoh.get_forces() * ureg.eV / ureg.angstrom)"
"print(ethanol_atoms.get_potential_energy() * ureg.eV)\n",
"print(ethanol_atoms.get_forces() * ureg.eV / ureg.angstrom)"
]
},
{
Expand All @@ -130,20 +146,21 @@
"metadata": {},
"outputs": [],
"source": [
"system = espressomd.System(box_l=[16] * 3)\n",
"system = espressomd.System(box_l=[16] * 3) # A cubic box of 16 Angstroms\n",
"system.time_step = (\n",
" (0.05 * ureg.fs)\n",
" .to(((1 * ureg.u * ureg.angstrom**2) / ureg.electron_volt) ** 0.5)\n",
" .magnitude\n",
")\n",
") # we convert the time units from fs to be in line with our other units for mass, distance and energy.\n",
"system.cell_system.skin = 0.4"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We add the ASE structure to the ESPResSo system and use the provided ASEInterface to go back to ASE, such that we can use the calculator to compute forces."
"We add the ASE structure to the ESPResSo system and use the provided ASEInterface to go back to ASE, such that we can use the calculator to compute forces.\n",
"We need to configure the `ASEInterface` to go back and forth between the ESPResSo system and the `ase.Atoms` objects. Because `ase.Atoms` is built around the atomic numbers, we use them for the ESPResSo type mapping."
]
},
{
Expand All @@ -152,17 +169,17 @@
"metadata": {},
"outputs": [],
"source": [
"for atom in etoh:\n",
"for atom in ethanol_atoms:\n",
" system.part.add(pos=atom.position, type=atom.number, mass=atom.mass)\n",
"system.ase = ASEInterface(type_mapping={x: x for x in set(etoh.numbers)})\n",
"system.ase = ASEInterface(type_mapping={x: x for x in set(ethanol_atoms.numbers)})\n",
"system.ase.get()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We will use the [ZnDraw](https://github.com/zincware/ZnDraw) visualizer integrated with ESPResSo to follow the simulation. To see the effect of the geometry optimization, we display the atomic forces."
"We will use the [ZnDraw](https://github.com/zincware/ZnDraw) visualizer integrated with ESPResSo to follow the simulation. To see the effect of the geometry optimization, we display the atomic forces. You can either follow the simulation inside this Notebook or open the app in a dedicated window to the side, following the printed URL."
]
},
{
Expand Down Expand Up @@ -350,7 +367,7 @@
"source": [
"# Langevin dynamics at 400 K\n",
"system.integrator.set_vv()\n",
"system.thermostat.set_langevin(kT=(400 * ureg.K * boltmk).magnitude, gamma=2, seed=42)"
"system.thermostat.set_langevin(kT=(400 * ureg.K * boltzmann_k).magnitude, gamma=2, seed=42)"
]
},
{
Expand Down Expand Up @@ -406,7 +423,18 @@
"metadata": {},
"outputs": [],
"source": [
"vis.zndraw.figure = fig.to_json()"
"vis.zndraw.figures = {\"distance\": fig.to_json()}"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Conclusion\n",
"\n",
"In this tutorial, you have created simulations from simple SMILES representations and were able to run geometry optimization and MD simulations with (almost) DFT accuracy but much faster performance using machine-learned interatomic potentials (MLIP). For the usage of MLIP you did not have to define bonds or dihedrals. You have seen this in the second part of the tutorial, where you observed bond breaking and bond formation in an MD simulation.\n",
"\n",
"MLIP can be used to study systems with much higher **ab initio** accuracy or to investigate phenomena on an atomistic scale which are impossible to study using classical force fields, such as chemical reactions."
]
},
{
Expand Down

0 comments on commit e49ab3c

Please sign in to comment.