Use ASE interface with constraints

[sli259]

Hi, I am trying to use the ANI2x and ASE interface to do MD simulation with constraints. I found once I apply the bond length constraints, the simulation will fail with the converge issue.

Here is an example code I used. There won't be an issue if I use the native TIP3P calculator to replace the ANI2x. I do want to use the ANI2x for my small molecules, is there a way to work with the bond length constraints?

Thanks!

from ase import Atoms
from ase.constraints import FixBondLengths
from ase.md import Langevin
from ase.optimize import BFGS
import ase.units as units
from ase.io.trajectory import Trajectory
import numpy as np
import torchani

# Set up water box at 20 deg C density
x = angleHOH * np.pi / 180 / 2
pos = [[0, 0, 0],
       [0, rOH * np.cos(x), rOH * np.sin(x)],
       [0, rOH * np.cos(x), -rOH * np.sin(x)]]
atoms = Atoms('OH2', positions=pos)

vol = ((18.01528 / 6.022140857e23) / (0.9982 / 1e24))**(1 / 3.)
atoms.set_cell((vol, vol, vol))
atoms.center()

atoms = atoms.repeat((3, 3, 3))
atoms.set_pbc(True)

# RATTLE-type constraints on O-H1, O-H2, H1-H2.
atoms.constraints = FixBondLengths([(3 * i + j, 3 * i + (j + 1) % 3)
                                    for i in range(3**3)
                                    for j in [0, 1, 2]])


calculator = torchani.models.ANI2x().ase()
atoms.set_calculator(calculator)
# atoms.calc = TIP3P(rc=4.5)

print("Begin minimizing...")
opt = BFGS(atoms)
traj = Trajectory('opt.traj', 'w', atoms)
opt.attach(traj)
opt.run(fmax=0.001)
print()

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))
    
dyn = Langevin(atoms, 1 * units.fs, 300 * units.kB, 0.2)
dyn.attach(printenergy, interval=50)

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