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 05:03:04    -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.16/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.375 eV
    E_kin = 0.022 eV (T = 173.3 K)
    E_tot = -1020.353 eV

Energy per atom:
    E_pot = -1021.331 eV
    E_kin = 0.964 eV (T = 7454.7 K)
    E_tot = -1020.367 eV

Energy per atom:
    E_pot = -1031.305 eV
    E_kin = 9.790 eV (T = 75738.0 K)
    E_tot = -1021.515 eV

Energy per atom:
    E_pot = -1030.806 eV
    E_kin = 8.391 eV (T = 64918.3 K)
    E_tot = -1022.415 eV

Energy per atom:
    E_pot = -1032.091 eV
    E_kin = 8.511 eV (T = 65844.0 K)
    E_tot = -1023.580 eV

Energy per atom:
    E_pot = -1031.118 eV
    E_kin = 6.252 eV (T = 48366.4 K)
    E_tot = -1024.866 eV

Energy per atom:
    E_pot = -1029.213 eV
    E_kin = 3.615 eV (T = 27964.7 K)
    E_tot = -1025.598 eV

Energy per atom:
    E_pot = -1031.826 eV
    E_kin = 5.740 eV (T = 44403.0 K)
    E_tot = -1026.086 eV

Energy per atom:
    E_pot = -1028.912 eV
    E_kin = 3.902 eV (T = 30190.6 K)
    E_tot = -1025.010 eV

Energy per atom:
    E_pot = -1031.529 eV
    E_kin = 5.589 eV (T = 43235.5 K)
    E_tot = -1025.940 eV


True