Unverified Commit f4c703bb authored by Gao, Xiang's avatar Gao, Xiang Committed by GitHub
Browse files

fix energy force example (#67)

parent 99c2c3dc
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
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')
{
"cells": [
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [
{
"name": "stderr",
"output_type": "stream",
"text": [
"/home/gaoxiang/pytorch-dev/pytorchdev/lib/python3.6/site-packages/h5py/__init__.py:36: FutureWarning: Conversion of the second argument of issubdtype from `float` to `np.floating` is deprecated. In future, it will be treated as `np.float64 == np.dtype(float).type`.\n",
" from ._conv import register_converters as _register_converters\n",
"WARNING:root:Unable to import NeuroChemAEV, please check your pyNeuroChem installation.\n"
]
}
],
"source": [
"import h5py\n",
"import torch\n",
"from selected_system import mols, mol_file\n",
"from ani_benchmark import ANIBenchmark\n",
"import pandas\n",
"import os\n",
"import tqdm\n",
"from IPython.display import display\n",
"import itertools\n",
"tqdm.monitor_interval = 0"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"torch.set_num_threads(1)\n",
"fm = h5py.File(os.path.join('../',mol_file), \"r\")\n",
"\n",
"benchmarks = {\n",
" 'C': ANIBenchmark(device=torch.device(\"cpu\")),\n",
"}\n",
"\n",
"if torch.cuda.is_available():\n",
" benchmarks.update({\n",
" 'G': ANIBenchmark(device=torch.device(\"cuda\")),\n",
" })"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 20\n",
"Running benchmark on molecule COC(=O)c1ccc([N+](=O)[O-])cc1\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "9ff5dae10b0e4a85b250095cc1c40f64",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>0.700781</td>\n",
" <td>0.537888</td>\n",
" <td>1.095591</td>\n",
" <td>0.022238</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>0.422604</td>\n",
" <td>0.183495</td>\n",
" <td>0.564654</td>\n",
" <td>0.002251</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>1.253797</td>\n",
" <td>0.561175</td>\n",
" <td>1.283622</td>\n",
" <td>0.010574</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>1.123385</td>\n",
" <td>0.721383</td>\n",
" <td>1.660245</td>\n",
" <td>0.024489</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>2.377182</td>\n",
" <td>1.282558</td>\n",
" <td>2.943866</td>\n",
" <td>0.035064</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 0.700781 0.537888 1.095591 0.022238\n",
"energy 0.422604 0.183495 0.564654 0.002251\n",
"force 1.253797 0.561175 1.283622 0.010574\n",
"forward 1.123385 0.721383 1.660245 0.024489\n",
"total 2.377182 1.282558 2.943866 0.035064"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 50\n",
"Running benchmark on molecule O=[N+]([O-])c1ccc(NN=Cc2ccc(C=NNc3ccc([N+](=O)[O-])cc3[N+](=O)[O-])cc2)c([N+](=O)[O-])c1\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "dbb459e25032453e9d6d5a8dc84a08de",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>1.680758</td>\n",
" <td>2.910254</td>\n",
" <td>1.120638</td>\n",
" <td>0.033106</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>0.546913</td>\n",
" <td>0.406957</td>\n",
" <td>0.474391</td>\n",
" <td>0.002537</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>3.637797</td>\n",
" <td>2.189924</td>\n",
" <td>1.495740</td>\n",
" <td>0.036373</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>2.227671</td>\n",
" <td>3.317210</td>\n",
" <td>1.595029</td>\n",
" <td>0.035643</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>5.865469</td>\n",
" <td>5.507135</td>\n",
" <td>3.090769</td>\n",
" <td>0.072016</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 1.680758 2.910254 1.120638 0.033106\n",
"energy 0.546913 0.406957 0.474391 0.002537\n",
"force 3.637797 2.189924 1.495740 0.036373\n",
"forward 2.227671 3.317210 1.595029 0.035643\n",
"total 5.865469 5.507135 3.090769 0.072016"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 10\n",
"Running benchmark on molecule N#CCC(=O)N\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "4e04277c06f5497384bed8e910669212",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>0.420929</td>\n",
" <td>0.149769</td>\n",
" <td>1.081606</td>\n",
" <td>0.006245</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>0.288124</td>\n",
" <td>0.093577</td>\n",
" <td>0.446744</td>\n",
" <td>0.002521</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>0.755610</td>\n",
" <td>0.216397</td>\n",
" <td>1.218433</td>\n",
" <td>0.007300</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>0.709054</td>\n",
" <td>0.243346</td>\n",
" <td>1.528349</td>\n",
" <td>0.008766</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>1.464663</td>\n",
" <td>0.459743</td>\n",
" <td>2.746782</td>\n",
" <td>0.016067</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 0.420929 0.149769 1.081606 0.006245\n",
"energy 0.288124 0.093577 0.446744 0.002521\n",
"force 0.755610 0.216397 1.218433 0.007300\n",
"forward 0.709054 0.243346 1.528349 0.008766\n",
"total 1.464663 0.459743 2.746782 0.016067"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 4,5,6\n",
"Running benchmark on molecule C\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "65906ab1bdf1470aaf7515189c163f09",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>0.207580</td>\n",
" <td>0.017271</td>\n",
" <td>0.775032</td>\n",
" <td>0.004428</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>0.149350</td>\n",
" <td>0.048953</td>\n",
" <td>0.219576</td>\n",
" <td>0.001350</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>0.390866</td>\n",
" <td>0.121347</td>\n",
" <td>0.768956</td>\n",
" <td>0.005004</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>0.356931</td>\n",
" <td>0.066225</td>\n",
" <td>0.994608</td>\n",
" <td>0.005778</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>0.747796</td>\n",
" <td>0.187572</td>\n",
" <td>1.763565</td>\n",
" <td>0.010781</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 0.207580 0.017271 0.775032 0.004428\n",
"energy 0.149350 0.048953 0.219576 0.001350\n",
"force 0.390866 0.121347 0.768956 0.005004\n",
"forward 0.356931 0.066225 0.994608 0.005778\n",
"total 0.747796 0.187572 1.763565 0.010781"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 100\n",
"Running benchmark on molecule CC(C)C[C@@H](C(=O)O)NC(=O)C[C@@H]([C@H](CC1CCCCC1)NC(=O)CC[C@@H]([C@H](Cc2ccccc2)NC(=O)OC(C)(C)C)O)O\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "bb3a02673cc14b468839c8b3a0e84edc",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>6.933071</td>\n",
" <td>11.386394</td>\n",
" <td>1.226802</td>\n",
" <td>0.118203</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>1.004400</td>\n",
" <td>0.848959</td>\n",
" <td>0.517654</td>\n",
" <td>0.001794</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>8.206886</td>\n",
" <td>6.606211</td>\n",
" <td>1.827498</td>\n",
" <td>0.108411</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>7.937472</td>\n",
" <td>12.235354</td>\n",
" <td>1.744456</td>\n",
" <td>0.119997</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>16.144357</td>\n",
" <td>18.841565</td>\n",
" <td>3.571954</td>\n",
" <td>0.228409</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 6.933071 11.386394 1.226802 0.118203\n",
"energy 1.004400 0.848959 0.517654 0.001794\n",
"force 8.206886 6.606211 1.827498 0.108411\n",
"forward 7.937472 12.235354 1.744456 0.119997\n",
"total 16.144357 18.841565 3.571954 0.228409"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"number of atoms: 305\n",
"Running benchmark on molecule [H]/N=C(/N)\\NCCC[C@H](C(=O)N[C@H]([C@@H](C)O)C(=O)N[C@H](Cc1ccc(cc1)O)C(=O)NCCCC[C@@H](C(=O)NCCCC[C@@H](C(=O)NCC(=O)O)NC(=O)[C@H](CCCCNC(=O)[C@@H](Cc2ccc(cc2)O)NC(=O)[C@@H]([C@@H](C)O)NC(=O)[C@@H](CCCN/C(=N\\[H])/N)N)NC(=O)[C@@H](Cc3ccc(cc3)O)NC(=O)[C@@H]([C@@H](C)O)NC(=O)[C@@H](CCCN/C(=N\\[H])/N)N)NC(=O)[C@@H](Cc4ccc(cc4)O)NC(=O)[C@@H]([C@@H](C)O)NC(=O)[C@@H](CCCN/C(=N\\[H])/N)N)N\n"
]
},
{
"data": {
"application/vnd.jupyter.widget-view+json": {
"model_id": "0a01e82e74094812b0a041a95e42ae38",
"version_major": 2,
"version_minor": 0
},
"text/plain": [
"HBox(children=(IntProgress(value=0, max=4), HTML(value='')))"
]
},
"metadata": {},
"output_type": "display_data"
},
{
"name": "stdout",
"output_type": "stream",
"text": [
"cuda runtime error (2) : out of memory at /home/gaoxiang/pytorch/aten/src/THC/generic/THCStorage.cu:58\n",
"\n"
]
},
{
"data": {
"text/html": [
"<div>\n",
"<style scoped>\n",
" .dataframe tbody tr th:only-of-type {\n",
" vertical-align: middle;\n",
" }\n",
"\n",
" .dataframe tbody tr th {\n",
" vertical-align: top;\n",
" }\n",
"\n",
" .dataframe thead th {\n",
" text-align: right;\n",
" }\n",
"</style>\n",
"<table border=\"1\" class=\"dataframe\">\n",
" <thead>\n",
" <tr style=\"text-align: right;\">\n",
" <th></th>\n",
" <th>C,1</th>\n",
" <th>C,B</th>\n",
" <th>G,1</th>\n",
" <th>G,B</th>\n",
" </tr>\n",
" </thead>\n",
" <tbody>\n",
" <tr>\n",
" <th>aev</th>\n",
" <td>22.462832</td>\n",
" <td>38.629616</td>\n",
" <td>1.459434</td>\n",
" <td>NaN</td>\n",
" </tr>\n",
" <tr>\n",
" <th>energy</th>\n",
" <td>2.698251</td>\n",
" <td>2.419860</td>\n",
" <td>0.525845</td>\n",
" <td>NaN</td>\n",
" </tr>\n",
" <tr>\n",
" <th>force</th>\n",
" <td>21.112948</td>\n",
" <td>20.788418</td>\n",
" <td>3.065523</td>\n",
" <td>NaN</td>\n",
" </tr>\n",
" <tr>\n",
" <th>forward</th>\n",
" <td>25.161083</td>\n",
" <td>41.049476</td>\n",
" <td>1.985280</td>\n",
" <td>NaN</td>\n",
" </tr>\n",
" <tr>\n",
" <th>total</th>\n",
" <td>46.274031</td>\n",
" <td>61.837895</td>\n",
" <td>5.050802</td>\n",
" <td>NaN</td>\n",
" </tr>\n",
" </tbody>\n",
"</table>\n",
"</div>"
],
"text/plain": [
" C,1 C,B G,1 G,B\n",
"aev 22.462832 38.629616 1.459434 NaN\n",
"energy 2.698251 2.419860 0.525845 NaN\n",
"force 21.112948 20.788418 3.065523 NaN\n",
"forward 25.161083 41.049476 1.985280 NaN\n",
"total 46.274031 61.837895 5.050802 NaN"
]
},
"metadata": {},
"output_type": "display_data"
}
],
"source": [
"for i in mols:\n",
" print('number of atoms:', i)\n",
" smiles = mols[i]\n",
" for s in smiles:\n",
" print('Running benchmark on molecule', s)\n",
" key = s.replace('/', '_')\n",
" coordinates = torch.from_numpy(fm[key][()])\n",
" coordinates = coordinates[:200]\n",
" species = fm[key].attrs['species'].split()\n",
" results = {}\n",
" for b,m in tqdm.tqdm_notebook(list(itertools.product(benchmarks, ['1','B']))):\n",
" bench = benchmarks[b]\n",
" coordinates = coordinates.type(bench.aev_computer.dtype)\n",
" try:\n",
" if m == '1':\n",
" result = bench.oneByOne(coordinates, species)\n",
" elif m == 'B':\n",
" result = bench.inBatch(coordinates, species)\n",
" else:\n",
" raise ValueError('BUG here')\n",
" result['forward'] = result['aev'] + result['energy']\n",
" result['total'] = result['forward'] + result['force']\n",
" except RuntimeError as e:\n",
" print(e)\n",
" result = {'aev': None, 'energy': None, 'force': None, 'total': None }\n",
" results[b + ',' + m] = result\n",
" df = pandas.DataFrame(results)\n",
" display(df)\n",
" break"
]
},
{
"cell_type": "code",
"execution_count": null,
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.6.5"
}
},
"nbformat": 4,
"nbformat_minor": 2
}
...@@ -20,7 +20,7 @@ coordinates = torch.tensor([[[0.03192167, 0.00638559, 0.01301679], ...@@ -20,7 +20,7 @@ coordinates = torch.tensor([[[0.03192167, 0.00638559, 0.01301679],
[0.45554739, 0.54289633, 0.81170881], [0.45554739, 0.54289633, 0.81170881],
[0.66091919, -0.16799635, -0.91037834]]], [0.66091919, -0.16799635, -0.91037834]]],
requires_grad=True) 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)) _, energy = model((species, coordinates))
derivative = torch.autograd.grad(energy.sum(), coordinates)[0] derivative = torch.autograd.grad(energy.sum(), coordinates)[0]
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment