Molecular dynamics using ASE

This example is modified from the official home page and Constant temperature MD to use the ASE interface of TorchANI as energy calculator.

As always, we start by importing the modules we need

import ase
from ase.lattice.cubic import Diamond
from ase.md.langevin import Langevin
from ase.optimize import LBFGS

import torchani

First we set up our system (in this case a diamond crystal, with PBC enabled)

atoms = Diamond(symbol="C", pbc=True)
len(atoms)  # The number of atoms in the system

8

After, we create a calculator from an ANI model and attach it to our atoms

atoms.calc = torchani.models.ANI2x().ase()

Then we minimize our system using the L-BFGS optimizer, which is included in ASE, under ase.optimize.LBFGS.

opt = LBFGS(atoms)
opt.run(fmax=0.0002)

       Step     Time          Energy          fmax
LBFGS:    0 05:33:37    -8162.894043        0.000034

np.True_

We want to run constant temperature MD, and print some quantities throughout the MD. For this we need to create a callback (function) that prints the quantities we are interested in. For example, to print the energy of the system we can use:

def print_energy(atoms: ase.Atoms):
    pot_energy = atoms.get_potential_energy() / len(atoms)
    kin_energy = atoms.get_kinetic_energy() / len(atoms)
    temperature = kin_energy / (1.5 * ase.units.kB)
    tot_energy = pot_energy + kin_energy
    print(
        "Energy per atom: \n"
        f"    E_pot = {pot_energy:.3f} eV\n"
        f"    E_kin = {kin_energy:.3f} eV (T = {temperature:.1f} K)\n"
        f"    E_tot = {tot_energy:.3f} eV\n"
    )

We will use the Langevin thermostat to control the temperature. To do this we need to first construct an ase.md.langevin.Langevin object. Here we use with a time step of 1 fs, a temperature of 300 K and a friction coefficient of 0.2

dyn = Langevin(
    atoms,
    timestep=1 * ase.units.fs,
    temperature_K=300,
    friction=0.2,
)
dyn.attach(print_energy, interval=5, atoms=atoms)

Finally we run the dynamics using dyn.run(steps)

print_energy(atoms)  # Print the initial energy
dyn.run(50)

Energy per atom:
    E_pot = -1020.362 eV
    E_kin = 0.000 eV (T = 0.0 K)
    E_tot = -1020.362 eV

Energy per atom:
    E_pot = -1020.362 eV
    E_kin = 0.000 eV (T = 0.0 K)
    E_tot = -1020.362 eV

Energy per atom:
    E_pot = -1020.383 eV
    E_kin = 0.032 eV (T = 244.1 K)
    E_tot = -1020.351 eV

Energy per atom:
    E_pot = -1022.672 eV
    E_kin = 2.339 eV (T = 18097.7 K)
    E_tot = -1020.333 eV

Energy per atom:
    E_pot = -1027.272 eV
    E_kin = 5.427 eV (T = 41986.0 K)
    E_tot = -1021.845 eV

Energy per atom:
    E_pot = -1028.548 eV
    E_kin = 7.213 eV (T = 55805.0 K)
    E_tot = -1021.335 eV

Energy per atom:
    E_pot = -1027.164 eV
    E_kin = 5.131 eV (T = 39693.5 K)
    E_tot = -1022.034 eV

Energy per atom:
    E_pot = -1030.740 eV
    E_kin = 7.664 eV (T = 59291.7 K)
    E_tot = -1023.076 eV

Energy per atom:
    E_pot = -1028.293 eV
    E_kin = 4.195 eV (T = 32452.8 K)
    E_tot = -1024.098 eV

Energy per atom:
    E_pot = -1031.304 eV
    E_kin = 6.085 eV (T = 47077.8 K)
    E_tot = -1025.219 eV

Energy per atom:
    E_pot = -1031.107 eV
    E_kin = 4.973 eV (T = 38469.4 K)
    E_tot = -1026.134 eV

Energy per atom:
    E_pot = -1031.603 eV
    E_kin = 4.413 eV (T = 34140.7 K)
    E_tot = -1027.190 eV


True