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.16/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 18:14:04    -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.16/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.029eV (T=223K)  Etot = -1038.856eV
Energy per atom: Epot = -1038.878eV  Ekin = 0.018eV (T=141K)  Etot = -1038.859eV
Energy per atom: Epot = -1038.875eV  Ekin = 0.025eV (T=194K)  Etot = -1038.850eV
Energy per atom: Epot = -1038.875eV  Ekin = 0.024eV (T=187K)  Etot = -1038.850eV
Energy per atom: Epot = -1038.863eV  Ekin = 0.026eV (T=202K)  Etot = -1038.837eV
Energy per atom: Epot = -1038.859eV  Ekin = 0.024eV (T=186K)  Etot = -1038.835eV
Energy per atom: Epot = -1038.881eV  Ekin = 0.023eV (T=180K)  Etot = -1038.858eV
Energy per atom: Epot = -1038.871eV  Ekin = 0.054eV (T=416K)  Etot = -1038.817eV
Energy per atom: Epot = -1038.877eV  Ekin = 0.059eV (T=458K)  Etot = -1038.818eV
Energy per atom: Epot = -1038.866eV  Ekin = 0.039eV (T=301K)  Etot = -1038.827eV

True