graphein
graphein copied to clipboard
Add add_bond_type_one_hot as edge features for molecules
It would be great to have add_bond_type_one_hot as an edge features for molecules.
Something similar to
def add_bond_type_one_hot(
u: str,
v: str,
d: Dict[str, Any],
return_array: bool = True,
allowable_set: Optional[List[str]] = None,
) -> np.ndarray:
if allowable_set is None:
allowable_set = [
rdkit.Chem.rdchem.BondType.SINGLE,
rdkit.Chem.rdchem.BondType.DOUBLE,
rdkit.Chem.rdchem.BondType.TRIPLE,
rdkit.Chem.rdchem.BondType.AROMATIC,
]
features = onek_encoding_unk(d["bond"].GetBondType(), allowable_set)
if return_array:
features = np.array(features).astype(int)
else:
features = pd.Series(features).astype(int)
features.index = allowable_set
d["bond_type_one_hot"] = features
return features
Happy to open a PR