feat: include CIP, unicolor configurations rdkit

Steinbeck-Lab · May 28, 2024 · d38a75b · d38a75b
1 parent c71b609
commit d38a75b
Show file tree

Hide file tree

Showing 3 changed files with 33 additions and 22 deletions.
diff --git a/app/modules/depiction.py b/app/modules/depiction.py
@@ -30,6 +30,7 @@ def get_cdk_depiction(
     Returns:
         image (SVG): CDK Structure Depiction as an SVG image.
     """
+    print(unicolor)
 
     cdk_base = "org.openscience.cdk"
     StandardGenerator = JClass(
@@ -89,8 +90,9 @@ def get_cdk_depiction(
             SmartsPattern.prepare(SDGMol)
             tmpMappings = tmpPattern.matchAll(SDGMol)
             tmpSubstructures = tmpMappings.toSubstructures()
+            lightBlue = Color(173, 216, 230)
             DepictionGenerator = DepictionGenerator.withHighlight(
-                tmpSubstructures, Color.RED
+                tmpSubstructures, lightBlue
             ).withOuterGlowHighlight()
 
         mol_imageSVG = (
@@ -115,6 +117,8 @@ def get_rdkit_depiction(
     mol_size=(512, 512),
     rotate=0,
     kekulize=True,
+    CIP=False,
+    unicolor=False,
     highlight: str = "",
 ) -> str:
     """
@@ -125,44 +129,46 @@ def get_rdkit_depiction(
         mol_size (tuple, optional): Size of the output image. Defaults to (512, 512).
         rotate (int, optional): Rotation angle of the molecule. Defaults to 0.
         kekulize (bool, optional): Whether to kekulize the molecule. Defaults to True.
+        CIP (bool, optional): Whether to assign CIP stereochemistry. Defaults to False.
+        unicolor (bool, optional): Whether to use a unicolor palette. Defaults to False.
         highlight (str, optional): SMARTS pattern to highlight atoms/bonds. Defaults to empty.
 
     Returns:
         str: RDKit Structure Depiction as an SVG image.
     """
-    if not molecule:
-        return "Error reading SMILES string, check again."
 
     mc = Chem.Mol(molecule.ToBinary())
 
     if kekulize:
         try:
             Chem.Kekulize(mc)
-        except Exception:
+        except Chem.KekulizeException:
             mc = Chem.Mol(molecule.ToBinary())
 
     if not mc.GetNumConformers():
         rdDepictor.Compute2DCoords(mc)
 
+    if CIP:
+        Chem.AssignStereochemistry(mc, force=True, cleanIt=True)
+
     drawer = rdMolDraw2D.MolDraw2DSVG(mol_size[0], mol_size[1])
     drawer.drawOptions().rotate = rotate
+    drawer.drawOptions().addStereoAnnotation = CIP
+
+    if unicolor:
+        drawer.drawOptions().useBWAtomPalette()
 
     if highlight:
         patt = Chem.MolFromSmarts(highlight)
         if patt:
-            hit_ats = list(mc.GetSubstructMatch(patt))
-            if hit_ats:
-                hit_bonds = [
-                    mc.GetBondBetweenAtoms(
-                        hit_ats[bond.GetBeginAtomIdx()], hit_ats[bond.GetEndAtomIdx()]
-                    ).GetIdx()
-                    for bond in patt.GetBonds()
-                ]
-                rdMolDraw2D.PrepareAndDrawMolecule(
-                    drawer, mc, highlightAtoms=hit_ats, highlightBonds=hit_bonds
-                )
-            else:
-                drawer.DrawMolecule(mc)
+            hit_ats = mc.GetSubstructMatch(patt)
+            hit_bonds = [
+                mc.GetBondBetweenAtoms(at1, at2).GetIdx()
+                for at1, at2 in zip(hit_ats[:-1], hit_ats[1:])
+            ]
+            rdMolDraw2D.PrepareAndDrawMolecule(
+                drawer, mc, highlightAtoms=hit_ats, highlightBonds=hit_bonds
+            )
         else:
             drawer.DrawMolecule(mc)
     else:

diff --git a/app/modules/toolkits/rdkit_wrapper.py b/app/modules/toolkits/rdkit_wrapper.py
@@ -58,7 +58,7 @@ def get_MolVolume(molecule: any) -> float:
         float: The volume of the molecule.
     """
     molecule = Chem.AddHs(molecule)
-    AllChem.EmbedMolecule(molecule)
+    AllChem.EmbedMolecule(molecule, useRandomCoords=True)
     volume = AllChem.ComputeMolVolume(molecule, gridSpacing=0.2)
     return volume
 

diff --git a/app/routers/depict.py b/app/routers/depict.py
@@ -108,12 +108,12 @@ async def depict_2d_molecule(
         title="Rotate",
         description="The rotation angle of the molecule in degrees.",
     ),
-    CIP: Optional[bool] = Query(
+    CIP: bool = Query(
         False,
         title="CIP",
         description="Whether to include Cahn-Ingold-Prelog (CIP) stereochemistry information.",
     ),
-    unicolor: Optional[bool] = Query(
+    unicolor: bool = Query(
         False,
         title="Unicolor",
         description="Whether to use a single colour for the molecule.",
@@ -158,12 +158,17 @@ async def depict_2d_molecule(
         if toolkit == "cdk":
             mol = parse_input(smiles, "cdk", False)
             depiction = get_cdk_depiction(
-                mol, [width, height], rotate, CIP, unicolor, highlight=highlight
+                mol,
+                [width, height],
+                rotate,
+                CIP=CIP,
+                unicolor=unicolor,
+                highlight=highlight,
             )
         elif toolkit == "rdkit":
             mol = parse_input(smiles, "rdkit", False)
             depiction = get_rdkit_depiction(
-                mol, [width, height], rotate, highlight=highlight
+                mol, [width, height], rotate, unicolor=unicolor, highlight=highlight
             )
         else:
             raise HTTPException(