
import torch


import molgrid


from openbabel import pybel


# create simple molecule
m = pybel.readstring('smi','C')
m.addh()


# atom typing based on atomic number and valence, with a constant radius of 1.5
def mytyper(atom):
    if hasattr(atom, 'GetValence'):
        return ([atom.GetAtomicNum(),atom.GetValence()], 1.5)
    else:
        return ([atom.GetAtomicNum(),atom.GetExplicitDegree()], 1.5)


# create the typer; the explicit names may be omitted, in which case numerical names
# will be automatically created
t = molgrid.PythonCallbackVectorTyper(mytyper, 2, ["anum","valence"])


# get types for our simple molecule using our typer
types = [t.get_atom_type_vector(a.OBAtom) for a in m.atoms]

Custom atom typer example

Tutorials