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

from torchani.models import ANI2x

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 = 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 19:46:33    -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)

/opt/hostedtoolcache/Python/3.11.15/x64/lib/python3.11/site-packages/ase/md/langevin.py:102: FutureWarning: The implementation of `fixcm=True` in `Langevin` does not strictly sample the correct NVT distributions. The deviations are typically small for large systems but can be more pronounced for small systems. Use `fixcm=False` together with `ase.constraints.FixCom`. `fixcm` is deprecated since ASE 3.28.0 and will be removed in a future release.
  warnings.warn(msg, FutureWarning)

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.381 eV
    E_kin = 0.027 eV (T = 212.7 K)
    E_tot = -1020.353 eV

Energy per atom:
    E_pot = -1021.747 eV
    E_kin = 1.427 eV (T = 11038.9 K)
    E_tot = -1020.320 eV

Energy per atom:
    E_pot = -1029.223 eV
    E_kin = 7.108 eV (T = 54989.4 K)
    E_tot = -1022.115 eV

Energy per atom:
    E_pot = -1031.547 eV
    E_kin = 7.381 eV (T = 57104.3 K)
    E_tot = -1024.166 eV

Energy per atom:
    E_pot = -1031.234 eV
    E_kin = 5.855 eV (T = 45293.6 K)
    E_tot = -1025.379 eV

Energy per atom:
    E_pot = -1031.895 eV
    E_kin = 5.552 eV (T = 42955.8 K)
    E_tot = -1026.342 eV

Energy per atom:
    E_pot = -1033.190 eV
    E_kin = 5.667 eV (T = 43843.0 K)
    E_tot = -1027.523 eV

Energy per atom:
    E_pot = -1032.544 eV
    E_kin = 3.972 eV (T = 30731.9 K)
    E_tot = -1028.572 eV

Energy per atom:
    E_pot = -1032.957 eV
    E_kin = 3.699 eV (T = 28615.6 K)
    E_tot = -1029.258 eV

Energy per atom:
    E_pot = -1032.551 eV
    E_kin = 2.775 eV (T = 21470.1 K)
    E_tot = -1029.776 eV


True