rdworks 0.25.8__py3-none-any.whl → 0.36.1__py3-none-any.whl

rdworks/display.py

@@ -1,135 +1,296 @@
-import io
-import os
 import numpy as np
-from typing import Optional, List, Tuple
 
 from PIL import Image, ImageChops
+from io import BytesIO
 
-from rdkit import Chem, Geometry
+from rdkit import Chem
 from rdkit.Chem import AllChem, Draw, rdDepictor, rdMolTransforms
 from rdkit.Chem.Draw import rdMolDraw2D
+from rdkit.Chem.Draw import MolsMatrixToGridImage # new in RDKit 2023.09.1
+
+# SVG optimization
+from scour.scour import scourString
 
 
 # https://greglandrum.github.io/rdkit-blog/posts/2023-05-26-drawing-options-explained.html
 
 
-def twod_depictor(rdmol:Chem.Mol, index:bool=False, coordgen:bool=False) -> Chem.Mol:
-    """Sets up for 2D depiction.
+def trim_png(img:Image.Image) -> Image.Image:
+    """Removes white margin around molecular drawing.
 
     Args:
-        rdmol (Chem.Mol): input molecule.
-        index (bool, optional): whether to show atom index. Defaults to False.
-        coordgen (bool, optional): whether to set rdDepictor.SetPreferCoordGen. Defaults to False.
+        img (Image.Image): input PIL Image object.
+
+    Returns:
+        Image.Image: output PIL Image object.
+    """
+    bg = Image.new(img.mode, img.size, img.getpixel((0,0)))
+    diff = ImageChops.difference(img,bg)
+    diff = ImageChops.add(diff, diff, 2.0, -100)
+    bbox = diff.getbbox()
+    
+    if bbox:
+        return img.crop(bbox)
+    
+    return img
+
+
+def get_highlight_bonds(rdmol: Chem.Mol, atom_indices: list[int]) -> list[int] | None:
+    """Get bond indices for bonds between atom indices.
+
+    Args:
+        rdmol (Chem.Mol): rdkit Chem.Mol object.
+        atom_indices (list[int]): atom indices.
 
     Returns:
-        Chem.Mol: a copy of rdkit.Chem.Mol object.
+        list[int]: bond indices.
     """
-    if coordgen:
-        rdDepictor.SetPreferCoordGen(True)
+    bond_indices = []
+    for bond in rdmol.GetBonds():
+        if bond.GetBeginAtomIdx() in atom_indices and bond.GetEndAtomIdx() in atom_indices:
+            bond_indices.append(bond.GetIdx())
+    
+    if bond_indices:
+        return bond_indices
     else:
-        rdDepictor.SetPreferCoordGen(False)
+        return None
+
+
+def render_2D_mol(rdmol:Chem.Mol,
+                  moldrawer:rdMolDraw2D, 
+                  redraw: bool = False,
+                  coordgen: bool = False,
+                  legend: str = '', 
+                  atom_index: bool = False, 
+                  highlight_atoms: list[int] | None = None,
+                  highlight_bonds: list[int] | None = None,
+                  ) -> str:
     
     rdmol_2d = Chem.Mol(rdmol)
-    rdDepictor.Compute2DCoords(rdmol_2d)
-    rdDepictor.StraightenDepiction(rdmol_2d)
     
-    for atom in rdmol_2d.GetAtoms():
-        for key in atom.GetPropsAsDict():
-            atom.ClearProp(key)
-    
-    if index: # index hides polar hydrogens
-        for atom in rdmol_2d.GetAtoms():
-            atom.SetProp("atomLabel", str(atom.GetIdx()))
-            # atom.SetProp("atomNote", str(atom.GetIdx()))
-            # atom.SetProp("molAtomMapNumber", str(atom.GetIdx()))
+    if redraw or rdmol_2d.GetNumConformers() == 0:
+        rdDepictor.SetPreferCoordGen(coordgen)
+        rdmol_2d = Chem.RemoveHs(rdmol_2d)
+        rdDepictor.Compute2DCoords(rdmol_2d)
+
+    rdDepictor.StraightenDepiction(rdmol_2d)
+
+    if (highlight_bonds is None) and (highlight_atoms is not None):
+        # highlight bonds between the highlighted atoms
+        highlight_bonds = get_highlight_bonds(rdmol_2d, highlight_atoms)
+
+    draw_options = moldrawer.drawOptions()
+
+    draw_options.addAtomIndices = atom_index
+    # draw_options.setHighlightColour((0,.9,.9,.8)) # Cyan highlight
+    # draw_options.addBondIndices = True
+    # draw_options.noAtomLabels = True
+    draw_options.atomLabelDeuteriumTritium = True  # D, T
+    # draw_options.explicitMethyl = True
+    draw_options.singleColourWedgeBonds = True
+    draw_options.addStereoAnnotation = True
+    # draw_options.fillHighlights = False
+    # draw_options.highlightRadius = .4
+    # draw_options.highlightBondWidthMultiplier = 12
+    # draw_options.variableAtomRadius = 0.2
+    # draw_options.variableBondWidthMultiplier = 40
+    # draw_options.setVariableAttachmentColour((.5,.5,1))
+    # draw_options.baseFontSize = 1.0 # default is 0.6
+    # draw_options.annotationFontScale = 1
+    # draw_options.rotate = 30 # rotation angle in degrees
+    # draw_options.padding = 0.2 # default is 0.05
+
+    # for atom in rdmol_2d.GetAtoms():
+    #     for key in atom.GetPropsAsDict():
+    #         atom.ClearProp(key)
+    # if index: # index hides polar hydrogens
+    #     for atom in rdmol_2d.GetAtoms():
+    #        atom.SetProp("atomLabel", str(atom.GetIdx()))
+    #     #    # atom.SetProp("atomNote", str(atom.GetIdx()))
+    #     #    # atom.SetProp("molAtomMapNumber", str(atom.GetIdx()))        
+
+    moldrawer.DrawMolecule(rdmol_2d, 
+                           legend=legend, 
+                           highlightAtoms=highlight_atoms, 
+                           highlightBonds=highlight_bonds)
+    moldrawer.FinishDrawing()
     
-    return rdmol_2d
+    return moldrawer.GetDrawingText()
 
 
-def svg(rdmol:Chem.Mol, 
-        width:int=300, 
-        height:int=300, 
-        legend:str='',
-        index:bool=False, 
-        highlight:list[int] | None = None, 
-        coordgen:bool = False) -> str:
-    """Returns string SVG output of a molecule.
+def render_svg(rdmol: Chem.Mol, 
+               width: int = 300, 
+               height: int = 300,
+               legend: str = '', 
+               atom_index: bool = False, 
+               highlight_atoms: list[int] | None = None,
+               highlight_bonds: list[int] | None = None,
+               redraw: bool = False,
+               coordgen: bool = False,
+               optimize: bool = True) -> str:
+    """Draw 2D molecule in SVG format.
 
     Examples:
+        For Jupyternotebook, wrap the output with SVG:
+
         >>> from IPython.display import SVG
         >>> SVG(libr[0].to_svg())
 
     Args:
-        rdmol (Chem.Mol): input molecule.
-        width (int): width. Defaults to 300.
-        height (int): height. Defaults to 300.
-        legend (str): title of molecule. Defaults to ''.
-        index (bool): whether to show atom indexes. Defaults to False.
-        highlight (list[int]): list of atom indices to highlight. Defaults to None.
-        coordgen (bool): whether to use rdDepictor.SetPreferCoordGen. Defaults to False.
+        rdmol (Chem.Mol): rdkit Chem.Mol object.
+        width (int, optional): width. Defaults to 300.
+        height (int, optional): height. Defaults to 300.
+        legend (str, optional): legend. Defaults to ''.
+        atom_index (bool, optional): whether to show atom index. Defaults to False.
+        highlight_atoms (list[int] | None, optional): atom(s) to highlight. Defaults to None.
+        highlight_bonds (list[int] | None, optional): bond(s) to highlight. Defaults to None.
+        redraw (bool, optional): whether to redraw. Defaults to False.
+        coordgen (bool, optional): whether to use coordgen. Defaults to False.
+        optimize (bool, optional): whether to optimize SVG string. Defaults to True.
 
     Returns:
-        str: SVG text
+        str: SVG string
     """
-    d2d_svg = rdMolDraw2D.MolDraw2DSVG(width, height)
-    rdmol_2d = twod_depictor(rdmol, index, coordgen)
-    if highlight:
-        d2d_svg.DrawMolecule(rdmol_2d, legend=legend, highlightAtoms=highlight)
-    else:
-        d2d_svg.DrawMolecule(rdmol_2d, legend=legend)
-    #rdMolDraw2D.PrepareAndDrawMolecule(d2d_svg, rdmol_2d, highlightAtoms=highlight, legend=legend)
-    d2d_svg.FinishDrawing()
-    return d2d_svg.GetDrawingText()
+
+    svg_string = render_2D_mol(rdmol,
+                                moldrawer = rdMolDraw2D.MolDraw2DSVG(width, height),
+                                redraw = redraw, 
+                                coordgen = coordgen,
+                                legend = legend, 
+                                atom_index = atom_index, 
+                                highlight_atoms = highlight_atoms, 
+                                highlight_bonds = highlight_bonds,
+                                )
+
+    if optimize:
+        scour_options = {
+            'strip_comments': True,
+            'strip_ids': True,
+            'shorten_ids': True,
+            'compact_paths': True,
+            'indent_type': 'none',
+        }
+        svg_string = scourString(svg_string, options=scour_options)
+
+    return svg_string
 
 
-def png(rdmol:Chem.Mol, width:int=300, height:int=300, legend:str='', 
-        index:bool=False, highlight:Optional[List[int]]=None, coordgen:bool=False) -> Image.Image:
-    """Returns a trimmed PIL Image object of a molecule.
+def render_png(rdmol: Chem.Mol, 
+               width: int = 300, 
+               height: int = 300,
+               legend: str = '', 
+               atom_index: bool = False, 
+               highlight_atoms: list[int] | None = None,
+               highlight_bonds: list[int] | None = None,
+               redraw: bool = False,
+               coordgen: bool = False,
+               trim: bool = True) -> Image.Image:
+    """Draw 2D molecule in PNG format.
 
     Args:
-        rdmol (Chem.Mol): input molecule.
-        width (int): width. Defaults to 300.
-        height (int): height. Defaults to 300.
-        legend (str): title of molecule. Defaults to ''.
-        index (bool): whether to show atom indexes. Defaults to False.
-        highlight (list): list of atom indices to highlight. Defaults to None.
-        coordgen (bool): whether to use rdDepictor.SetPreferCoordGen. Defaults to False.
+        rdmol (Chem.Mol): rdkit Chem.Mol object.
+        width (int, optional): width. Defaults to 300.
+        height (int, optional): height. Defaults to 300.
+        legend (str, optional): legend. Defaults to ''.
+        atom_index (bool, optional): whether to show atom index. Defaults to False.
+        highlight_atoms (list[int] | None, optional): atom(s) to highlight. Defaults to None.
+        highlight_bonds (list[int] | None, optional): bond(s) to highlight. Defaults to None.
+        redraw (bool, optional): whether to redraw. Defaults to False.
+        coordgen (bool, optional): whether to use coordgen. Defaults to False.
     
     Returns:
         Image.Image: output PIL Image object.
     """
-    rdmol_2d = twod_depictor(rdmol, index, coordgen)
-    img = Draw.MolToImage(rdmol_2d, 
-                          size=(width,height),
-                          highlightAtoms=highlight,
-                          kekulize=True,
-                          wedgeBonds=True,
-                          fitImage=False,
-                          )
-    # highlightAtoms: list of atoms to highlight (default [])
-    # highlightBonds: list of bonds to highlight (default [])
-    # highlightColor: RGB color as tuple (default [1, 0, 0])
 
-    return trim_png(img)
+    png_string = render_2D_mol(rdmol,
+                                moldrawer = rdMolDraw2D.MolDraw2DCairo(width, height),
+                                redraw = redraw, 
+                                coordgen = coordgen,
+                                legend = legend, 
+                                atom_index = atom_index, 
+                                highlight_atoms = highlight_atoms, 
+                                highlight_bonds = highlight_bonds,
+                                )
 
+    img = Image.open(BytesIO(png_string))
+
+    if trim:
+        img = trim_png(img)
+
+    return img
 
-def trim_png(img:Image.Image) -> Image.Image:
-    """Removes white margin around molecular drawing.
+
+def render_matrix_grid(rdmol: list[Chem.Mol],
+                        legend: list[str] | None,
+                        highlight_atoms: list[list[int]] | None = None,
+                        highlight_bonds: list[list[int]] | None = None,
+                        mols_per_row: int = 5,
+                        width: int = 200,
+                        height: int = 200,
+                        atom_index: bool = False,
+                        redraw: bool = False,
+                        coordgen: bool = False,
+                        svg: bool = True,
+                        ) -> str | Image.Image:
+    """Rendering a grid image from a list of molecules.
 
     Args:
-        img (Image.Image): input PIL Image object.
+        rdmol (list[Chem.Mol]): list of rdkit Chem.Mol objects.
+        legend (list[str]): list of legends
+        highlight_atoms (list[list[int]] | None, optional): list of atom(s) to highlight. Defaults to None.
+        highlight_bonds (list[list[int]] | None, optional): list of bond(s) to highlight. Defaults to None.
+        mols_per_row (int, optional): molecules per row. Defaults to 5.
+        width (int, optional): width. Defaults to 200.
+        height (int, optional): height. Defaults to 200.
+        atom_index (bool, optional): whether to show atom index. Defaults to False.
+        redraw (bool, optional): whether to redraw 2D. Defaults to False.
+        coordgen (bool, optional): whether to use coordgen to depict. Defaults to False.
 
     Returns:
-        Image.Image: output PIL Image object.
+        str | Image.Image: SVG string or PIL Image object.
+
+    Reference:
+        https://greglandrum.github.io/rdkit-blog/posts/2023-10-25-molsmatrixtogridimage.html
     """
-    bg = Image.new(img.mode, img.size, img.getpixel((0,0)))
-    diff = ImageChops.difference(img,bg)
-    diff = ImageChops.add(diff, diff, 2.0, -100)
-    bbox = diff.getbbox()
-    if bbox:
-        return img.crop(bbox)
-    return img
 
+    n = len(rdmol)
+
+    if isinstance(legend, list):
+        assert len(legend) == n, "number of legends and molecules must be the same"
+    elif legend is None:
+        legend = ['',] * n
+    
+    if isinstance(highlight_atoms, list):
+        assert len(highlight_atoms) == n, "number of highlights and molecules must be the same"
+    elif highlight_atoms is None:
+        highlight_atoms = [ (), ] * n
+
+    if isinstance(highlight_bonds, list):
+        assert len(highlight_bonds) == n, "number of highlights and molecules must be the same"    
+    elif highlight_bonds is None:
+        highlight_bonds = [ (), ] * n
+
+    rdmol_matrix = []
+    legend_matrix = []
+    highlight_atoms_matrix = []
+    highlight_bonds_matrix = []
+
+    for i in range(0, n, mols_per_row):
+        rdmol_matrix.append(rdmol[i:(i + mols_per_row)])
+        legend_matrix.append(legend[i:(i + mols_per_row)])
+        highlight_atoms_matrix.append(highlight_atoms[i:(i + mols_per_row)])
+        highlight_bonds_matrix.append(highlight_bonds[i:(i + mols_per_row)])
+
+    return MolsMatrixToGridImage(
+            molsMatrix = rdmol_matrix,
+            subImgSize = (width, height),
+            legendsMatrix = legend_matrix,
+            highlightAtomListsMatrix = highlight_atoms_matrix,
+            highlightBondListsMatrix = highlight_bonds_matrix,
+            useSVG = svg,
+            returnPNG = False # whether to return PNG data (True) or a PIL object (False)
+            )
+    
 
 def rescale(rdmol:Chem.Mol, factor:float=1.5) -> Chem.Mol:
     """Returns a copy of `rdmol` by a `factor`.