Structure minimization and constant temperature MD using ASE interface

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

To begin with, let’s first import the modules we will use:

from ase.lattice.cubic import Diamond
from ase.md.langevin import Langevin
from ase.optimize import BFGS
from ase import units
import torchani

Now let’s set up a crystal

atoms = Diamond(symbol="C", pbc=True)
print(len(atoms), "atoms in the cell")

8 atoms in the cell

Now let’s create a calculator from builtin models:

calculator = torchani.models.ANI1ccx().ase()

/opt/hostedtoolcache/Python/3.8.17/x64/lib/python3.8/site-packages/torchani-2.2.3-py3.8.egg/torchani/resources/
Downloading ANI model parameters ...

Note

Regardless of the dtype you use in your model, when converting it to ASE calculator, it always automatically the dtype to torch.float64. The reason for this behavior is, at many cases, the rounding error is too large for structure minimization. If you insist on using torch.float32, do the following instead:

calculator = torchani.models.ANI1ccx().ase(dtype=torch.float32)

Now let’s set the calculator for atoms:

atoms.set_calculator(calculator)

Now let’s minimize the structure:

print("Begin minimizing...")
opt = BFGS(atoms)
opt.run(fmax=0.001)
print()

Begin minimizing...
      Step     Time          Energy         fmax
BFGS:    0 00:04:06    -8311.226180        0.0000

Now create a callback function that print interesting physical quantities:

def printenergy(a=atoms):
    """Function to print the potential, kinetic and total energy."""
    epot = a.get_potential_energy() / len(a)
    ekin = a.get_kinetic_energy() / len(a)
    print('Energy per atom: Epot = %.3feV  Ekin = %.3feV (T=%3.0fK)  '
          'Etot = %.3feV' % (epot, ekin, ekin / (1.5 * units.kB), epot + ekin))

We want to run MD with constant energy using the Langevin algorithm with a time step of 1 fs, the temperature 300K and the friction coefficient to 0.02 atomic units.

dyn = Langevin(atoms, 1 * units.fs, 300 * units.kB, 0.2)
dyn.attach(printenergy, interval=50)

/opt/hostedtoolcache/Python/3.8.17/x64/lib/python3.8/site-packages/ase/md/md.py:48: FutureWarning: Specify the temperature in K using the 'temperature_K' argument
  warnings.warn(FutureWarning(w))

Now run the dynamics:

print("Beginning dynamics...")
printenergy()
dyn.run(500)

Beginning dynamics...
Energy per atom: Epot = -1038.903eV  Ekin = 0.000eV (T=  0K)  Etot = -1038.903eV
Energy per atom: Epot = -1038.903eV  Ekin = 0.000eV (T=  0K)  Etot = -1038.903eV
Energy per atom: Epot = -1038.885eV  Ekin = 0.016eV (T=125K)  Etot = -1038.869eV
Energy per atom: Epot = -1038.866eV  Ekin = 0.029eV (T=225K)  Etot = -1038.837eV
Energy per atom: Epot = -1038.871eV  Ekin = 0.038eV (T=292K)  Etot = -1038.833eV
Energy per atom: Epot = -1038.876eV  Ekin = 0.023eV (T=177K)  Etot = -1038.853eV
Energy per atom: Epot = -1038.868eV  Ekin = 0.023eV (T=177K)  Etot = -1038.845eV
Energy per atom: Epot = -1038.872eV  Ekin = 0.024eV (T=187K)  Etot = -1038.848eV
Energy per atom: Epot = -1038.882eV  Ekin = 0.036eV (T=275K)  Etot = -1038.846eV
Energy per atom: Epot = -1038.875eV  Ekin = 0.026eV (T=198K)  Etot = -1038.849eV
Energy per atom: Epot = -1038.876eV  Ekin = 0.033eV (T=256K)  Etot = -1038.843eV
Energy per atom: Epot = -1038.887eV  Ekin = 0.033eV (T=258K)  Etot = -1038.854eV

True