fix energy force example (#67)

benchmark/benchmark.py

-import torchani
-
-
-class ANIBenchmark:
-
-    def __init__(self, device):
-        super(ANIBenchmark, self).__init__(device)
-        self.aev_computer = torchani.SortedAEV(device=device)
-        self.model = torchani.ModelOnAEV(
-            self.aev_computer, benchmark=True, derivative=True, from_nc=None)
-
-    def oneByOne(self, coordinates, species):
-        conformations = coordinates.shape[0]
-        coordinates = coordinates.to(self.device)
-        for i in range(conformations):
-            c = coordinates[i:i+1, :, :]
-            self.model(c, species)
-        ret = {
-            'aev': self.model.timers['aev'],
-            'energy': self.model.timers['nn'],
-            'force': self.model.timers['derivative']
-        }
-        self.model.reset_timers()
-        return ret
-
-    def inBatch(self, coordinates, species):
-        coordinates = coordinates.to(self.device)
-        self.model(coordinates, species)
-        ret = {
-            'aev': self.model.timers['aev'],
-            'energy': self.model.timers['nn'],
-            'force': self.model.timers['derivative']
-        }
-        self.model.reset_timers()
-        return ret

benchmark/generate_data.py

-from ase import Atoms
-from ase.calculators.tip3p import TIP3P, rOH, angleHOH
-from ase.md import Langevin
-import ase.units as units
-import numpy
-import h5py
-from rdkit import Chem
-from rdkit.Chem import AllChem
-# from asap3 import EMT
-from ase.calculators.emt import EMT
-from multiprocessing import Pool
-from tqdm import tqdm, trange
-from selected_system import mols, mol_file
-
-conformations = 1024
-T = 30
-tqdm.monitor_interval = 0
-
-fw = h5py.File("waters.hdf5", "w")
-fm = h5py.File(mol_file, "w")
-
-
-def save(h5file, name, species, coordinates):
-    h5file[name] = coordinates
-    h5file[name].attrs['species'] = ' '.join(species)
-
-
-def waterbox(x, y, z, tqdmpos):
-    name = '{}_waters'.format(x*y*z)
-    # Set up water box at 20 deg C density
-    a = angleHOH * numpy.pi / 180 / 2
-    pos = [[0, 0, 0],
-           [0, rOH * numpy.cos(a), rOH * numpy.sin(a)],
-           [0, rOH * numpy.cos(a), -rOH * numpy.sin(a)]]
-    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((x, y, z))
-    atoms.set_pbc(False)
-    species = atoms.get_chemical_symbols()
-
-    atoms.calc = TIP3P()
-    md = Langevin(atoms, 1 * units.fs, temperature=T *
-                  units.kB, friction=0.01)
-
-    def generator(n):
-        for _ in trange(n, desc=name, position=tqdmpos):
-            md.run(1)
-            positions = atoms.get_positions()
-            yield positions
-
-    save(fw, name, species, numpy.stack(generator(conformations)))
-
-
-def compute(smiles):
-    m = Chem.MolFromSmiles(smiles)
-    m = Chem.AddHs(m)
-    AllChem.EmbedMolecule(m, useRandomCoords=True)
-    AllChem.UFFOptimizeMolecule(m)
-    pos = m.GetConformer().GetPositions()
-    natoms = m.GetNumAtoms()
-    species = [m.GetAtomWithIdx(j).GetSymbol() for j in range(natoms)]
-
-    atoms = Atoms(species, positions=pos)
-
-    atoms.calc = EMT()
-    md = Langevin(atoms, 1 * units.fs, temperature=T *
-                  units.kB, friction=0.01)
-
-    def generator(n):
-        for _ in range(n):
-            md.run(1)
-            positions = atoms.get_positions()
-            yield positions
-
-    c = numpy.stack(generator(conformations))
-    return smiles.replace('/', '_'), species, c
-
-
-def molecules():
-    smiles = [s for atoms in mols for s in mols[atoms]]
-    with Pool() as p:
-        return p.map(compute, smiles)
-
-
-if __name__ == '__main__':
-    for i in molecules():
-        save(fm, *i)
-    print(list(fm.keys()))
-    print('done with molecules')
-
-    with Pool() as p:
-        p.starmap(waterbox, [(10, 10, 10, 0), (20, 20, 10, 1),
-                             (30, 30, 30, 2), (40, 40, 40, 3)])
-        print(list(fw.keys()))
-        print('done with water boxes')

examples/energy_force.py

@@ -20,7 +20,7 @@ coordinates = torch.tensor([[[0.03192167,  0.00638559,  0.01301679],
                              [0.45554739,   0.54289633,  0.81170881],
                              [0.66091919,  -0.16799635, -0.91037834]]],
                            requires_grad=True)
-species = torch.LongTensor([[2, 1, 1, 1, 1]])  # 1 = H, 2 = C, 3 = N, 4 = O
+species = torch.LongTensor([[1, 0, 0, 0, 0]])  # 0 = H, 1 = C, 2 = N, 3 = O
 
 _, energy = model((species, coordinates))
 derivative = torch.autograd.grad(energy.sum(), coordinates)[0]