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

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 17:44:51    -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:110: 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.377 eV
    E_kin = 0.026 eV (T = 204.8 K)
    E_tot = -1020.350 eV

Energy per atom:
    E_pot = -1021.225 eV
    E_kin = 0.890 eV (T = 6882.9 K)
    E_tot = -1020.335 eV

Energy per atom:
    E_pot = -1029.342 eV
    E_kin = 8.163 eV (T = 63155.1 K)
    E_tot = -1021.179 eV

Energy per atom:
    E_pot = -1030.774 eV
    E_kin = 8.017 eV (T = 62018.8 K)
    E_tot = -1022.758 eV

Energy per atom:
    E_pot = -1030.880 eV
    E_kin = 6.628 eV (T = 51275.2 K)
    E_tot = -1024.253 eV

Energy per atom:
    E_pot = -1030.971 eV
    E_kin = 5.708 eV (T = 44156.5 K)
    E_tot = -1025.264 eV

Energy per atom:
    E_pot = -1030.660 eV
    E_kin = 4.297 eV (T = 33241.6 K)
    E_tot = -1026.363 eV

Energy per atom:
    E_pot = -1033.582 eV
    E_kin = 5.998 eV (T = 46402.9 K)
    E_tot = -1027.584 eV

Energy per atom:
    E_pot = -1032.497 eV
    E_kin = 4.001 eV (T = 30954.0 K)
    E_tot = -1028.496 eV

Energy per atom:
    E_pot = -1032.316 eV
    E_kin = 3.172 eV (T = 24539.5 K)
    E_tot = -1029.144 eV


True