graphein icon indicating copy to clipboard operation
graphein copied to clipboard

Add add_bond_type_one_hot as edge features for molecules

Open chris-clem opened this issue 5 months ago • 0 comments

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

chris-clem avatar Jun 18 '25 09:06 chris-clem