Skip to content

Add add_bond_type_one_hot as edge features for molecules #430

@chris-clem

Description

@chris-clem

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

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions