def test_bond_as_node_featurizer():
m = make_a_mol()
featurizer = BondAsNodeFeaturizerFull(length_featurizer="bin")
feat = featurizer(m)
size = featurizer.feature_size
assert np.array_equal(feat["feat"].shape, (m.GetNumBonds(), size))
assert len(featurizer.feature_name) == size
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