def test_callbacktyping():
m = pybel.readstring('smi', 'c1c(Cl)cccc1CO')
m.addh()
def mytyper(atom):
return (atom.GetAtomicNum(), 2.0)
t = molgrid.PythonCallbackIndexTyper(mytyper, 16)
names = list(t.get_type_names())
assert len(names) == 16
assert names[3] == '3' #numerical names
typs = [t.get_atom_type_index(a.OBAtom) for a in m.atoms]
assert len(typs) == 16
ccnt = 0
ocnt = 0
hcnt = 0
other = 0
for t, r in typs:
if t < 0:
other += 1
else:
if names[t] == '6':
ccnt += 1
assert r == approx(2)
if names[t] == '8':
ocnt += 1
assert r == approx(2)
if names[t] == '1':
hcnt += 1
assert r == approx(2)
assert ccnt == 7
assert ocnt == 1
assert hcnt == 7
assert other == 1
for molfile in args.molfile:
try:
ext = os.path.splitext(molfile)[1].lstrip('.')
mol = next(pybel.readfile(ext, molfile))
mol.OBMol.Center()
elemmap = {1: 0, 6: 1, 7: 2, 8: 3} #type indices
typeradii = [1.0, 1.6, 1.5, 1.4]
def mytyper(atom):
i = elemmap[atom.GetAtomicNum()]
r = typeradii[i]
return (i, r)
typer = molgrid.PythonCallbackIndexTyper(mytyper, 4,
['H', 'C', 'N', 'O'])
tensor_shape = (1, ) + dims
input_tensor = torch.zeros(tensor_shape,
dtype=torch.float32,
device=device)
predictions = []
with torch.no_grad():
labelvec = torch.zeros(1, dtype=torch.float32, device=device)
c = molgrid.CoordinateSet(mol, typer)
ex = molgrid.Example()
ex.coord_sets.append(c)
batch = molgrid.ExampleVec([ex])
types = c.type_index.tonumpy()