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 19:51:34    -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.382 eV
    E_kin = 0.025 eV (T = 195.5 K)
    E_tot = -1020.357 eV

Energy per atom:
    E_pot = -1021.898 eV
    E_kin = 1.542 eV (T = 11930.4 K)
    E_tot = -1020.356 eV

Energy per atom:
    E_pot = -1029.129 eV
    E_kin = 7.595 eV (T = 58759.9 K)
    E_tot = -1021.534 eV

Energy per atom:
    E_pot = -1032.069 eV
    E_kin = 9.142 eV (T = 70726.0 K)
    E_tot = -1022.927 eV

Energy per atom:
    E_pot = -1032.480 eV
    E_kin = 8.081 eV (T = 62516.4 K)
    E_tot = -1024.400 eV

Energy per atom:
    E_pot = -1032.552 eV
    E_kin = 6.890 eV (T = 53306.2 K)
    E_tot = -1025.662 eV

Energy per atom:
    E_pot = -1031.880 eV
    E_kin = 5.062 eV (T = 39159.5 K)
    E_tot = -1026.818 eV

Energy per atom:
    E_pot = -1032.781 eV
    E_kin = 4.891 eV (T = 37836.4 K)
    E_tot = -1027.890 eV

Energy per atom:
    E_pot = -1033.042 eV
    E_kin = 4.295 eV (T = 33227.3 K)
    E_tot = -1028.747 eV

Energy per atom:
    E_pot = -1033.482 eV
    E_kin = 3.944 eV (T = 30514.9 K)
    E_tot = -1029.538 eV


True