@@ -20,19 +20,26 @@ def get_cdk_depiction(
2020 unicolor = False ,
2121 highlight = "" ,
2222 highlight_atoms = None ,
23+ showAtomNumbers = False ,
2324):
2425 """This function takes the user input SMILES and Depicts it.
2526
2627 using the CDK Depiction Generator.
2728
2829 Args:
2930 molecule (any): CDK IAtomContainer parsed from SMILES string given by the user.
31+ molSize (tuple, optional): Size of the output image. Defaults to (512, 512).
32+ rotate (int, optional): Rotation angle in degrees. Defaults to 0.
33+ kekulize (bool, optional): Whether to kekulize the molecule. Defaults to True.
34+ CIP (bool, optional): Whether to annotate CIP stereochemistry. Defaults to True.
35+ unicolor (bool, optional): Whether to use black and white colors. Defaults to False.
36+ highlight (str, optional): SMARTS pattern to highlight. Defaults to empty.
37+ highlight_atoms (list, optional): List of atom indices to highlight. Defaults to None.
38+ showAtomNumbers (bool, optional): Whether to display atom numbers. Defaults to False.
3039
3140 Returns:
3241 image (SVG): CDK Structure Depiction as an SVG image.
3342 """
34- print (unicolor )
35-
3643 cdk_base = "org.openscience.cdk"
3744 StandardGenerator = JClass (
3845 cdk_base + ".renderer.generators.standard.StandardGenerator" ,
@@ -75,7 +82,7 @@ def get_cdk_depiction(
7582 try :
7683 Kekulization .kekulize (SDGMol )
7784 except Exception as e :
78- print (e + "Can't Kekulize molecule" )
85+ print (str ( e ) + " Can't Kekulize molecule" )
7986
8087 point = JClass (
8188 cdk_base + ".geometry.GeometryTools" ,
@@ -86,22 +93,65 @@ def get_cdk_depiction(
8693 (rotate * JClass ("java.lang.Math" ).PI / 180.0 ),
8794 )
8895
96+ # Add atom numbers if requested
97+ if showAtomNumbers :
98+ DepictionGenerator = DepictionGenerator .withAtomNumbers ()
99+
89100 # Handle highlighting: prioritize atom indices over SMARTS patterns
90101 if highlight_atoms and len (highlight_atoms ) > 0 :
91- # For CDK, we need to create substructures from atom indices
92- # This is more complex and would require additional CDK classes
93- # For now, fall back to SMARTS pattern if available
94- if highlight and highlight .strip ():
95- tmpPattern = SmartsPattern .create (highlight , SCOB .getInstance ())
96- SmartsPattern .prepare (SDGMol )
97- tmpMappings = tmpPattern .matchAll (SDGMol )
98- tmpSubstructures = tmpMappings .toSubstructures ()
99- lightBlue = Color (173 , 216 , 230 )
100- DepictionGenerator = DepictionGenerator .withHighlight (
101- tmpSubstructures , lightBlue
102- ).withOuterGlowHighlight ()
103- # Note: Direct atom index highlighting in CDK requires more complex implementation
104- # This would need creating IAtomContainerSet from specific atoms
102+ # Create atom sets from indices for highlighting
103+ AtomContainer = JClass (cdk_base + ".AtomContainer" )
104+ AtomContainerSet = JClass (cdk_base + ".AtomContainerSet" )
105+
106+ # Create a set of substructures from atom indices
107+ tmpSubstructures = AtomContainerSet ()
108+
109+ # If highlight_atoms is a list of lists, each list is a separate substructure
110+ if isinstance (highlight_atoms [0 ], list ):
111+ # Multiple substructures (e.g., multiple sugars)
112+ for atom_indices in highlight_atoms :
113+ if len (atom_indices ) > 0 :
114+ subset = AtomContainer ()
115+ for idx in atom_indices :
116+ if idx < SDGMol .getAtomCount ():
117+ subset .addAtom (SDGMol .getAtom (idx ))
118+ # Add bonds between highlighted atoms
119+ for i , idx1 in enumerate (atom_indices ):
120+ for idx2 in atom_indices [i + 1 :]:
121+ if (
122+ idx1 < SDGMol .getAtomCount ()
123+ and idx2 < SDGMol .getAtomCount ()
124+ ):
125+ bond = SDGMol .getBond (
126+ SDGMol .getAtom (idx1 ), SDGMol .getAtom (idx2 )
127+ )
128+ if bond is not None :
129+ subset .addBond (bond )
130+ tmpSubstructures .addAtomContainer (subset )
131+ else :
132+ # Single substructure
133+ subset = AtomContainer ()
134+ for idx in highlight_atoms :
135+ if idx < SDGMol .getAtomCount ():
136+ subset .addAtom (SDGMol .getAtom (idx ))
137+ # Add bonds between highlighted atoms
138+ for i , idx1 in enumerate (highlight_atoms ):
139+ for idx2 in highlight_atoms [i + 1 :]:
140+ if (
141+ idx1 < SDGMol .getAtomCount ()
142+ and idx2 < SDGMol .getAtomCount ()
143+ ):
144+ bond = SDGMol .getBond (
145+ SDGMol .getAtom (idx1 ), SDGMol .getAtom (idx2 )
146+ )
147+ if bond is not None :
148+ subset .addBond (bond )
149+ tmpSubstructures .addAtomContainer (subset )
150+
151+ lightBlue = Color (173 , 216 , 230 )
152+ DepictionGenerator = DepictionGenerator .withHighlight (
153+ tmpSubstructures , lightBlue
154+ ).withOuterGlowHighlight ()
105155 elif highlight and highlight .strip ():
106156 tmpPattern = SmartsPattern .create (highlight , SCOB .getInstance ())
107157 SmartsPattern .prepare (SDGMol )
@@ -138,6 +188,7 @@ def get_rdkit_depiction(
138188 unicolor = False ,
139189 highlight : str = "" ,
140190 highlight_atoms = None ,
191+ showAtomNumbers = False ,
141192) -> str :
142193 """
143194 Generate a 2D depiction of the input molecule using RDKit.
@@ -150,6 +201,8 @@ def get_rdkit_depiction(
150201 CIP (bool, optional): Whether to assign CIP stereochemistry. Defaults to False.
151202 unicolor (bool, optional): Whether to use a unicolor palette. Defaults to False.
152203 highlight (str, optional): SMARTS pattern to highlight atoms/bonds. Defaults to empty.
204+ highlight_atoms (list, optional): List of atom indices to highlight. Defaults to None.
205+ showAtomNumbers (bool, optional): Whether to display atom numbers. Defaults to False.
153206
154207 Returns:
155208 str: RDKit Structure Depiction as an SVG image.
@@ -176,6 +229,12 @@ def get_rdkit_depiction(
176229 if unicolor :
177230 drawer .drawOptions ().useBWAtomPalette ()
178231
232+ # Add atom numbers if requested
233+ if showAtomNumbers :
234+ # Set atom numbers as notes on each atom
235+ for atom in mc .GetAtoms ():
236+ atom .SetProp ("atomNote" , str (atom .GetIdx ()))
237+
179238 # Handle highlighting based on priority: anchor atoms + SMARTS pattern, then atom indices, then SMARTS pattern alone
180239 if highlight_atoms and len (highlight_atoms ) > 0 and highlight :
181240 # Combined approach: Use SMARTS pattern but only highlight the match that contains the anchor atoms
@@ -222,12 +281,11 @@ def get_rdkit_depiction(
222281 hit_ats = tuple (highlight_atoms )
223282 # Find ALL bonds that connect atoms within the functional group
224283 hit_bonds = []
225- for i , atom1_idx in enumerate (hit_ats ):
226- for j , atom2_idx in enumerate (hit_ats ):
227- if i < j : # Avoid duplicate bonds
228- bond = mc .GetBondBetweenAtoms (atom1_idx , atom2_idx )
229- if bond :
230- hit_bonds .append (bond .GetIdx ())
284+ for i in range (len (hit_ats )):
285+ for j in range (i + 1 , len (hit_ats )):
286+ bond = mc .GetBondBetweenAtoms (hit_ats [i ], hit_ats [j ])
287+ if bond :
288+ hit_bonds .append (bond .GetIdx ())
231289
232290 rdMolDraw2D .PrepareAndDrawMolecule (
233291 drawer , mc , highlightAtoms = hit_ats , highlightBonds = hit_bonds
0 commit comments