def get_grapher_hetero():
return HeteroMoleculeGraph(
atom_featurizer=AtomFeaturizerFull(),
bond_featurizer=BondAsNodeFeaturizerFull(),
global_featurizer=GlobalFeaturizer(),
self_loop=True,
)
def get_grapher():
# atom_featurizer = AtomFeaturizerFull()
# bond_featurizer = BondAsNodeFeaturizerFull(length_featurizer=None, dative=False)
# global_featurizer = GlobalFeaturizer(allowed_charges=None)
atom_featurizer = AtomFeaturizerMinimum()
bond_featurizer = BondAsNodeFeaturizerMinimum(length_featurizer=None)
global_featurizer = GlobalFeaturizer(allowed_charges=[-1, 0, 1])
grapher = HeteroMoleculeGraph(
atom_featurizer=atom_featurizer,
bond_featurizer=bond_featurizer,
global_featurizer=global_featurizer,
self_loop=True,
)
return grapher
def _get_grapher(model_path):
model_info = get_model_info(model_path)
allowed_charge = model_info["allowed_charge"]
featurizer_set = model_info["featurizer_set"]
if featurizer_set == "full":
atom_featurizer = AtomFeaturizerFull()
bond_featurizer = BondAsNodeFeaturizerFull(length_featurizer=None,
dative=False)
# This is used by the pubchem dataset, which only allows charges of 0.
# We still pass None to allowed_charges because we do not use any charge info
# in training.
global_featurizer = GlobalFeaturizer(allowed_charges=None)
elif featurizer_set == "minimum":
atom_featurizer = AtomFeaturizerMinimum()
bond_featurizer = BondAsNodeFeaturizerMinimum(length_featurizer=None)
global_featurizer = GlobalFeaturizer(allowed_charges=allowed_charge)
elif featurizer_set == "mg_thf_g2":
atom_featurizer = AtomFeaturizerMinimum()
bond_featurizer = BondAsNodeFeaturizerMinimum(length_featurizer=None)
global_featurizer = GlobalFeaturizer(
allowed_charges=allowed_charge,
solvent_environment=[
"smd_thf", "smd_7.23,1.4097,0,0.859,36.83,0.00,0.00"
],
)
else:
raise ValueError(
f"Unsupported featurizer set: {featurizer_set}. Cannot load grapher."
)
grapher = HeteroMoleculeGraph(
atom_featurizer=atom_featurizer,
bond_featurizer=bond_featurizer,
global_featurizer=global_featurizer,
self_loop=True,
)
return grapher
def assert_graph(self_loop):
m = make_a_mol()
grapher = HeteroMoleculeGraph(self_loop=self_loop)
g = grapher.build_graph(m)
# number of atoms
na = 7
# number of bonds
nb = 7
# number of edges between atoms and bonds
ne = 2 * nb
nodes = ["atom", "bond", "global"]
num_nodes = [g.number_of_nodes(n) for n in nodes]
ref_num_nodes = [na, 7, 1]
if self_loop:
edges = [
"a2b", "b2a", "a2g", "g2a", "b2g", "g2b", "a2a", "b2b", "g2g"
]
num_edges = [g.number_of_edges(e) for e in edges]
ref_num_edges = [ne, ne, na, na, nb, nb, na, nb, 1]
else:
edges = ["a2b", "b2a", "a2g", "g2a", "b2g", "g2b"]
num_edges = [g.number_of_edges(e) for e in edges]
ref_num_edges = [ne, ne, na, na, nb, nb]
assert set(g.ntypes) == set(nodes)
assert set(g.etypes) == set(edges)
assert num_nodes == ref_num_nodes
assert num_edges == ref_num_edges
bond_to_atom_map = {
0: [0, 1],
1: [0, 4],
2: [1, 2],
3: [1, 6],
4: [2, 3],
5: [3, 4],
6: [0, 5],
}
atom_to_bond_map = defaultdict(list)
for b, atoms in bond_to_atom_map.items():
for a in atoms:
atom_to_bond_map[a].append(b)
atom_to_bond_map = {
a: sorted(bonds)
for a, bonds in atom_to_bond_map.items()
}
ref_b2a_map = get_bond_to_atom_map(g)
ref_a2b_map = get_atom_to_bond_map(g)
assert bond_to_atom_map == ref_b2a_map
assert atom_to_bond_map == ref_a2b_map
if self_loop:
for nt, n in zip(nodes, num_nodes):
assert get_hetero_self_loop_map(
g, nt) == {i: [i]
for i in range(n)}