def get_grapher_hetero():
return HeteroMoleculeGraph(
atom_featurizer=AtomFeaturizerFull(),
bond_featurizer=BondAsNodeFeaturizerFull(),
global_featurizer=GlobalFeaturizer(),
self_loop=True,
)
def test_atom_featurizer():
m = make_a_mol()
species = list(set([a.GetSymbol() for a in m.GetAtoms()]))
featurizer = AtomFeaturizerFull()
feat = featurizer(m, dataset_species=species)
size = featurizer.feature_size
assert np.array_equal(feat["feat"].shape, (m.GetNumAtoms(), size))
assert len(featurizer.feature_name) == size
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 get_grapher_homo():
return HomoCompleteGraph(
atom_featurizer=AtomFeaturizerFull(),
bond_featurizer=BondAsEdgeCompleteFeaturizer(),
)