From e260b01349a8b4b229f3b35d4c2fb76c7c9981fa Mon Sep 17 00:00:00 2001 From: Timotej Bernat Date: Wed, 20 Mar 2024 16:16:53 -0600 Subject: [PATCH] Added functions for labelling and unlabelling bond indices --- polymerist/rdutils/labeling/bondwise.py | 22 ++++++++++++++++++++-- 1 file changed, 20 insertions(+), 2 deletions(-) diff --git a/polymerist/rdutils/labeling/bondwise.py b/polymerist/rdutils/labeling/bondwise.py index 650dbc3..96c5a06 100644 --- a/polymerist/rdutils/labeling/bondwise.py +++ b/polymerist/rdutils/labeling/bondwise.py @@ -1,6 +1,6 @@ '''For obtaining info from and for labelling individual RDKit Bonds''' -from typing import Any, Callable, Iterable +from typing import Any, Callable, Iterable, Optional from itertools import combinations from rdkit import Chem @@ -8,11 +8,13 @@ from ..rdtypes import RDMol, RDBond from .molwise import atom_ids_by_map_nums + from ...genutils.typetools import Args, KWArgs from ...genutils.iteration import sliding_window +from ...genutils.decorators.functional import optional_in_place -# BOND ID QUERY FUNCTIONS +# BOND ID QUERYING def are_bonded_atoms(rdmol : RDMol, atom_id_1 : int, atom_id_2 : int) -> bool: '''Check if pair of atoms in an RDMol have a bond between them''' return (rdmol.GetBondBetweenAtoms(atom_id_1, atom_id_2) is not None) @@ -56,3 +58,19 @@ def get_shortest_path_bonds(rdmol : RDMol, start_idx : int, end_idx : int) -> li ] +# BOND ID LABELING +@optional_in_place +def assign_bond_id_labels(rdmol : RDMol, bond_id_remap : Optional[dict[int, int]]=None) -> None: + '''Draws bond indices over their positions when displaying a Mol. + Can optionally provide a dict mapping bond indices to some other integers''' + if bond_id_remap is None: + bond_id_remap = {} # avoid mutable default + + for bond in rdmol.GetBonds(): + bond.SetIntProp('bondNote', bond_id_remap.get(bond.GetIdx(), bond.GetIdx())) # check if map value exists, if not default to index + +@optional_in_place +def clear_bond_id_labels(rdmol : RDMol) -> None: + '''Removes bond indices over their positions when displaying a Mol''' + for bond in rdmol.GetBonds(): + bond.ClearProp('bondNote') \ No newline at end of file