ASE examples (#122)

codefresh.yml

@@ -21,21 +21,23 @@ steps:
       - python setup.py test
       # - python2 setup.py test
 
-  Examples:
+  Tools:
     image: '${{BuildTorchANI}}'
     commands:
       - rm -rf *.pt
-      - python examples/nnp_training.py dataset/ani_gdb_s01.h5 dataset/ani_gdb_s01.h5
-      - python examples/nnp_training.py dataset/ani_gdb_s01.h5 dataset/ani_gdb_s01.h5  # run twice to test if checkpoint is working
-      - python examples/energy_force.py
       - python tools/training-benchmark.py ./dataset/ani_gdb_s01.h5
       - python tools/neurochem-test.py ./dataset/ani_gdb_s01.h5
       - python tools/inference-benchmark.py --tqdm ./xyz_files/CH4-5.xyz
+
+  ModuleMain:
+    image: '${{BuildTorchANI}}'
+    commands:
+      - rm -rf *.pt
       - python -m torchani.neurochem.trainer --tqdm tests/test_data/inputtrain.ipt  dataset/ani_gdb_s01.h5 dataset/ani_gdb_s01.h5
       - python -m torchani.data.cache_aev tmp dataset/ani_gdb_s01.h5 256
 
   Docs:
     image: '${{BuildTorchANI}}'
     commands:
       - find . -name '*.pt' -delete
-      - sphinx-build -D plot_gallery=0 docs build
+      - sphinx-build docs build

docs/index.rst

@@ -18,6 +18,7 @@ Welcome to TorchANI's documentation!
     examples/nnp_training
     examples/cache_aev
     examples/neurochem_trainer
+    examples/ase_langevin
 
 .. toctree::
     :maxdepth: 2

examples/ase_langevin.py

@@ -0,0 +1,55 @@
+# -*- coding: utf-8 -*-
+"""
+Constant temperature MD using ASE interface
+===========================================
+
+This example is modified from the official `Constant temperature MD`_ to use
+the ASE interface of TorchANI as energy calculator.
+
+.. _Constant temperature MD:
+    https://wiki.fysik.dtu.dk/ase/tutorials/md/md.html#constant-temperature-md
+"""
+
+
+###############################################################################
+# 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 import units
+import torchani
+
+
+###############################################################################
+# Now let's set up a crystal
+atoms = Diamond(symbol="C", pbc=True)
+
+###############################################################################
+# Now let's create a calculator from builtin models:
+builtin = torchani.neurochem.Builtins()
+calculator = torchani.ase.Calculator(builtin.species, builtin.aev_computer,
+                                     builtin.models, builtin.energy_shifter)
+atoms.set_calculator(calculator)
+
+###############################################################################
+# We want to run MD with constant energy using the Langevin algorithm
+# with a time step of 5 fs, the temperature 50K and the friction
+# coefficient to 0.02 atomic units.
+dyn = Langevin(atoms, 5 * units.fs, 50 * units.kB, 0.002)
+
+
+###############################################################################
+# Let's print energies every 50 steps:
+def printenergy(a=atoms):  # store a reference to atoms in the definition.
+    """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.attach(printenergy, interval=50)
+
+###############################################################################
+# Now run the dynamics:
+printenergy()
+dyn.run(500)

tests/test_ase.py

@@ -0,0 +1,37 @@
+from ase.lattice.cubic import Diamond
+from ase.md.langevin import Langevin
+from ase import units
+from ase.calculators.test import numeric_force
+import torch
+import torchani
+import unittest
+
+
+def get_numeric_force(atoms, eps):
+    fn = torch.zeros((len(atoms), 3))
+    for i in range(len(atoms)):
+        for j in range(3):
+            fn[i, j] = numeric_force(atoms, i, j, eps)
+    return fn
+
+
+class TestASE(unittest.TestCase):
+
+    def testForceWithPBCEnabled(self):
+        atoms = Diamond(symbol="C", pbc=True)
+        builtin = torchani.neurochem.Builtins()
+        calculator = torchani.ase.Calculator(
+            builtin.species, builtin.aev_computer,
+            builtin.models, builtin.energy_shifter)
+        atoms.set_calculator(calculator)
+        dyn = Langevin(atoms, 5 * units.fs, 30000000 * units.kB, 0.002)
+        dyn.run(100)
+        f = torch.from_numpy(atoms.get_forces())
+        fn = get_numeric_force(atoms, 0.001)
+        df = (f - fn).abs().max()
+        avgf = f.abs().mean()
+        self.assertLess(df / avgf, 0.1)
+
+
+if __name__ == '__main__':
+    unittest.main()

torchani/ase.py

@@ -22,8 +22,10 @@ class NeighborList:
     """
 
     def __init__(self, cell=None, pbc=None):
-        self.pbc = pbc
-        self.cell = cell
+        # wrap `cell` and `pbc` with `ase.Atoms`
+        a = ase.Atoms('He', [[0, 0, 0]], cell=cell, pbc=pbc)
+        self.pbc = a.get_pbc()
+        self.cell = a.get_cell(complete=True)
 
     def __call__(self, species, coordinates, cutoff):
         conformations = species.shape[0]
@@ -39,18 +41,23 @@ def __call__(self, species, coordinates, cutoff):
             c = c.squeeze()
             atoms = s.shape[0]
             atoms_object = ase.Atoms(
-                'C'*atoms,  # chemical symbols are not important here
+                ['He'] * atoms,  # chemical symbols are not important here
                 positions=c.detach().numpy(),
                 pbc=self.pbc,
                 cell=self.cell)
-            idx1, idx2 = ase.neighborlist.neighbor_list(
-                'ij', atoms_object, cutoff)
+            idx1, idx2, shift = ase.neighborlist.neighbor_list(
+                'ijS', atoms_object, cutoff)
             # NB: The absolute distance and distance vectors computed by
             # `neighbor_list`can not be used since it does not preserve
             # gradient information
             idx1 = torch.from_numpy(idx1).to(coordinates.device)
             idx2 = torch.from_numpy(idx2).to(coordinates.device)
             D = c.index_select(0, idx2) - c.index_select(0, idx1)
+            shift = torch.from_numpy(shift).to(coordinates.device) \
+                         .to(coordinates.dtype)
+            cell = torch.from_numpy(self.cell).to(coordinates.device) \
+                        .to(coordinates.dtype)
+            D += shift @ cell
             d = D.norm(2, -1)
             neighbor_species1 = []
             neighbor_distances1 = []
@@ -97,7 +104,10 @@ class Calculator(ase.calculators.calculator.Calculator):
         energy_shifter (:class:`torchani.EnergyShifter`): Energy shifter.
     """
 
+    implemented_properties = ['energy', 'forces']
+
     def __init__(self, species, aev_computer, model, energy_shifter):
+        super(Calculator, self).__init__()
         self.species_to_tensor = utils.ChemicalSymbolsToInts(species)
         self.aev_computer = aev_computer
         self.model = model
@@ -115,16 +125,16 @@ def __init__(self, species, aev_computer, model, energy_shifter):
     def calculate(self, atoms=None, properties=['energy'],
                   system_changes=ase.calculators.calculator.all_changes):
         super(Calculator, self).calculate(atoms, properties, system_changes)
-        self.aev_computer.neighbor_list = NeighborList(
-            cell=self.atoms.get_cell(), pbc=self.atoms.get_pbc())
+        self.aev_computer.neighborlist = NeighborList(
+            cell=self.atoms.get_cell(complete=True), pbc=self.atoms.get_pbc())
         species = self.species_to_tensor(self.atoms.get_chemical_symbols())
-        coordinates = self.atoms.get_positions(wrap=True).unsqueeze(0)
-        coordinates = torch.tensor(coordinates,
-                                   device=self.device,
-                                   dtype=self.dtype,
-                                   requires_grad=('forces' in properties))
-        _, energy = self.whole((species, coordinates)) * ase.units.Hartree
+        species = species.unsqueeze(0)
+        coordinates = torch.tensor(self.atoms.get_positions(wrap=True))
+        coordinates = coordinates.unsqueeze(0).to(self.device).to(self.dtype) \
+                                 .requires_grad_('forces' in properties)
+        _, energy = self.whole((species, coordinates))
+        energy *= ase.units.Hartree
         self.results['energy'] = energy.item()
         if 'forces' in properties:
             forces = -torch.autograd.grad(energy.squeeze(), coordinates)[0]
-            self.results['forces'] = forces.item()
+            self.results['forces'] = forces.squeeze().numpy()

torchani/neurochem/__init__.py

@@ -282,6 +282,7 @@ def __init__(self):
             parent_name,
             'resources/ani-1x_dft_x8ens/rHCNO-5.2R_16-3.5A_a4-8.params')
         self.consts = Constants(self.const_file)
+        self.species = self.consts.species
         self.aev_computer = AEVComputer(**self.consts)
 
         self.sae_file = pkg_resources.resource_filename(