Note
Go to the end to download the full example code.
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
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
8
After, we create a calculator from an ANI model and attach it to our atoms
atoms.calc = torchani.models.ANI2x().ase()
/home/ipickering/Repos/ani/torchani/arch.py:1196: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
dict_ = torch.load(path, map_location=torch.device("cpu"))
Then we minimize our system using the
L-BFGS optimizer,
which is included in ASE, under ase.optimize.LBFGS.
Step Time Energy fmax
LBFGS: 0 04:37:38 -8162.894043 0.000035
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)
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.381 eV
E_kin = 0.027 eV (T = 210.5 K)
E_tot = -1020.354 eV
Energy per atom:
E_pot = -1022.199 eV
E_kin = 1.865 eV (T = 14426.4 K)
E_tot = -1020.335 eV
Energy per atom:
E_pot = -1030.578 eV
E_kin = 8.615 eV (T = 66646.7 K)
E_tot = -1021.963 eV
Energy per atom:
E_pot = -1032.780 eV
E_kin = 9.074 eV (T = 70198.8 K)
E_tot = -1023.706 eV
Energy per atom:
E_pot = -1033.242 eV
E_kin = 7.803 eV (T = 60367.3 K)
E_tot = -1025.439 eV
Energy per atom:
E_pot = -1033.074 eV
E_kin = 6.481 eV (T = 50139.4 K)
E_tot = -1026.593 eV
Energy per atom:
E_pot = -1032.212 eV
E_kin = 4.428 eV (T = 34259.7 K)
E_tot = -1027.783 eV
Energy per atom:
E_pot = -1032.446 eV
E_kin = 3.842 eV (T = 29719.3 K)
E_tot = -1028.605 eV
Energy per atom:
E_pot = -1032.579 eV
E_kin = 3.388 eV (T = 26213.9 K)
E_tot = -1029.191 eV
Energy per atom:
E_pot = -1033.502 eV
E_kin = 3.702 eV (T = 28637.2 K)
E_tot = -1029.800 eV
True