rdworks 0.25.7__py3-none-any.whl

rdworks/conf.py

@@ -0,0 +1,374 @@
+import io
+import copy
+import json
+import types
+import numpy as np
+
+import ase
+from ase.optimize import FIRE
+
+from collections.abc import Callable
+
+from rdkit import Chem
+from rdkit.Chem import rdMolTransforms, AllChem, rdMolAlign
+from rdkit.Chem.Draw import rdMolDraw2D
+
+from typing import List, Optional, Union, Self
+
+from .units import ev2kcalpermol
+from .element import radii
+
+class Conf:
+    """Container for 3D conformers.
+    """
+    
+    def __init__(self, molecular_input:Chem.Mol, name:str='') -> None:
+        """Create 3D conformers.
+
+        Args:
+            molecular_input (Chem.Mol): Molecule for conformer generation.
+            name (str, optional): Name prefix of the generated conformers. Defaults to ''.
+
+        Raises:
+            ValueError: if `molecular_input` is not rdkit.Chem.Mol object.
+        """
+        self.rdmol = None # has only one rdkit conformer
+        self.name = name
+        self.props = {}
+        if isinstance(molecular_input, Chem.Mol):
+            self.rdmol = molecular_input
+            self.natoms = self.rdmol.GetNumAtoms()
+            self.props.update({'atoms': self.natoms})
+        else:
+            raise ValueError(f'rdworks.Conf() takes Chem.Mol object')
+        
+
+    def __str__(self) -> str:
+        """Returns a string representation.
+
+        Returns:
+            str: string representation.
+        """
+        return f"<rdworks.Conf({self.rdmol} name={self.name} atoms={self.natoms})>"
+    
+
+    ##################################################
+    ### Cascading methods
+    ##################################################
+
+
+    def copy(self) -> Self:
+        """Returns a copy of self.
+
+        Returns:
+            Self: `rdworks.Conf` object.
+        """
+        return copy.deepcopy(self)
+    
+
+    def rename(self, name:str) -> Self:
+        """Rename and returns self.
+
+        Args:
+            name (str): a new name for conformers.
+
+        Raises:
+            ValueError: if `name` is not given.
+
+        Returns:
+            Self: `rdworks.Conf` object.
+        """
+        if not name:
+            raise ValueError('rdworks.Conf.rename() expects a name')
+        self.name = name
+        self.rdmol.SetProp('_Name', name)
+        return self
+    
+
+    def sync(self, coord:Union[np.ndarray, list]) -> Self:
+        """Synchronize the conformer coordinates with the provided `coord`.
+
+        Args:
+            coord (np.array): 3D coordinates.
+
+        Raises:
+            ValueError: if `coord` does not have the correct shape (natoms, 3).
+
+        Returns:
+            Self: `rdworks.Conf` object.
+        """
+        if isinstance(coord, np.ndarray) and coord.shape != (self.natoms, 3):
+            raise ValueError(f"`coord.shape` should be ({self.natoms},3)")
+        elif isinstance(coord, list) and len(coord) != self.natoms:
+            raise ValueError(f"`coord` should be length of {self.natoms}")
+        for i, a in enumerate(self.rdmol.GetAtoms()):
+            self.rdmol.GetConformer().SetAtomPosition(a.GetIdx(), coord[i])
+        
+        return self
+
+
+    def get_potential_energy(self, calculator: str | Callable = 'MMFF94') -> float:
+        """Get potential energy in kcal/mol.
+
+        Args:
+            calculator (str | Callable): MMFF94 (= MMFF), MMFF94s, UFF, or ASE calculator.
+                `MMFF94` or `MMFF` - Intended for general use, including organic molecules and proteins, 
+                    and primarily relies on data from quantum mechanical calculations. 
+                    It's often used in molecular dynamics simulations.
+                `MMFF94s` - A "static" variant of MMFF94, with adjusted parameters for out-of-plane 
+                    bending and dihedral torsions to favor planar geometries for specific nitrogen atoms. 
+                    This makes it better suited for geometry optimization studies where a static, 
+                    time-averaged structure is desired. The "s" stands for "static".
+                `UFF` - UFF refers to the "Universal Force Field," a force field model used for 
+                    molecular mechanics calculations. It's a tool for geometry optimization, 
+                    energy minimization, and exploring molecular conformations in 3D space. 
+                    UFF is often used to refine conformers generated by other methods, 
+                    such as random conformer generation, to produce more physically plausible 
+                    and stable structures.
+
+        Returns:
+            float: potential energy in kcal/mol.
+        """
+        PE = None
+        if isinstance(calculator, str):
+            if calculator == 'MMFF94' or calculator == 'MMFF':
+                mp = AllChem.MMFFGetMoleculeProperties(self.rdmol, mmffVariant='MMFF94')
+                ff = AllChem.MMFFGetMoleculeForceField(self.rdmol, mp)
+            elif calculator == 'MMFF94s':
+                mp = AllChem.MMFFGetMoleculeProperties(self.rdmol, mmffVariant='MMFF94s')
+                ff = AllChem.MMFFGetMoleculeForceField(self.rdmol, mp)
+            elif calculator == 'UFF':
+                ff = AllChem.UFFGetMoleculeForceField(self.rdmol)
+            else:
+                raise ValueError("Unsupported calculator")
+            PE = ff.CalcEnergy()
+            self.props.update({'E_tot(kcal/mol)': PE})
+        else:
+            try:
+                ase_atoms = ase.Atoms(symbols=self.symbols(), positions=self.positions())
+                ase_atoms.calc = calculator
+                PE = ase_atoms.get_potential_energy() # np.array
+                PE = ev2kcalpermol * float(PE[0]) # np.float64 to float
+                self.props.update({'E_tot(kcal/mol)': PE})
+            except:
+                raise RuntimeError("ASE calculator error")
+        return PE
+    
+
+    def optimize(self, calculator: str | Callable = 'MMFF94', fmax:float=0.05) -> Self:
+        """Optimize conformation using a callable.
+
+        Args:
+            calculator (str | Callable): MMFF94 (= MMFF), MMFF94s, UFF, or ASE calculator.
+                `MMFF94` or `MMFF` - Intended for general use, including organic molecules and proteins, 
+                    and primarily relies on data from quantum mechanical calculations. 
+                    It's often used in molecular dynamics simulations.
+                `MMFF94s` - A "static" variant of MMFF94, with adjusted parameters for out-of-plane 
+                    bending and dihedral torsions to favor planar geometries for specific nitrogen atoms. 
+                    This makes it better suited for geometry optimization studies where a static, 
+                    time-averaged structure is desired. The "s" stands for "static".
+                `UFF` - UFF refers to the "Universal Force Field," a force field model used for 
+                    molecular mechanics calculations. It's a tool for geometry optimization, 
+                    energy minimization, and exploring molecular conformations in 3D space. 
+                    UFF is often used to refine conformers generated by other methods, 
+                    such as random conformer generation, to produce more physically plausible 
+                    and stable structures.
+            fmax (float, optional): fmax for the calculator. Defaults to 0.05.
+
+        Returns:
+            Self: self
+        """
+        if isinstance(calculator, str) :
+            init = self.get_potential_energy(calculator)
+            if calculator == 'MMFF94' or calculator == 'MMFF':
+                retcode = AllChem.MMFFOptimizeMolecule(self.rdmol, mmffVariant='MMFF94')
+                # returns 0 if the optimization converged
+            elif calculator == 'MMFF94s':
+                retcode = AllChem.MMFFOptimizeMolecule(self.rdmol, mmffVariant='MMFF94s')
+                # returns 0 if the optimization converged
+            elif calculator == 'UFF':
+                retcode = AllChem.UFFOptimizeMolecule(self.rdmol)
+                # returns 0 if the optimization converged
+            final = self.get_potential_energy(calculator)
+            self.props.update({
+                'E_tot_init(kcal/mol)': init , # energy before optimization
+                'E_tot(kcal/mol)': final, # energy after optimization
+                'Converged' : retcode == 0, # True or False
+            })
+            return self
+
+        else:
+            with io.StringIO() as logfile:
+                ase_atoms = ase.Atoms(symbols=self.symbols(), positions=self.positions())
+                ase_atoms.calc = calculator
+                FIRE(ase_atoms, logfile=logfile).run(fmax=fmax)
+                lines = [l.strip().split()[1:] for l in logfile.getvalue().split('\n') if l.startswith('FIRE')]
+                data = [(float(e), float(f)) for (_, _, e, f) in lines]
+                self.props.update({
+                    'E_tot_init(kcal/mol)': data[0][0] * ev2kcalpermol, # energy before optimization
+                    'E_tot(kcal/mol)': data[-1][0] * ev2kcalpermol, # energy after optimization
+                    'Converged' : data[-1][1] < fmax, # True or False
+                })
+                # update atomic coordinates
+                return self.sync(ase_atoms.get_positions())
+ 
+
+    ##################################################
+    ### Endpoint methods
+    ##################################################
+
+    def check_bonds(self, tolerance:float=0.25) -> bool:
+        """Check bond lengths.
+
+        Args:
+            tolerance (float, optional): tolerance from the sum of 
+                van der Waals radii of bonded atoms. Defaults to 0.25 (A).
+
+        Returns:
+            bool: True if all bond lengths are accceptable.
+        """
+
+        for bond in self.rdmol.GetBonds():
+            idx1 = bond.GetBeginAtomIdx()
+            idx2 = bond.GetEndAtomIdx()
+            nuc1 = self.rdmol.GetAtomWithIdx(idx1).GetAtomicNum()
+            nuc2 = self.rdmol.GetAtomWithIdx(idx2).GetAtomicNum()
+            sum_radii = (radii[nuc1] + radii[nuc2])
+            bond_length = rdMolTransforms.GetBondLength(self.rdmol.GetConformer(), idx1, idx2)
+            if abs(bond_length - sum_radii) > tolerance:
+                return False
+            
+        return True
+
+
+    def positions(self) -> np.array:
+        """Returns the coordinates.
+
+        Returns:
+            np.array: the coordinates.
+        """
+        return np.array(self.rdmol.GetConformer().GetPositions().tolist())
+
+
+    def symbols(self) -> list[str]:
+        """Returns the element symbols.
+
+        Returns:
+            list: list of element symbols.
+        """
+        return [ a.GetSymbol() for a in self.rdmol.GetAtoms() ]
+
+
+    def numbers(self) -> list[int]:
+        """Returns the atomic numbers.
+
+        Returns:
+            list: list of atomic numbers.
+        """
+        return [ a.GetAtomicNum() for a in self.rdmol.GetAtoms() ]
+
+
+    def cog(self) -> np.array:
+        """Returns the center of geometry (COG).
+
+        Returns:
+            np.array: the center of geometry (COG).
+        """
+        xyz = []
+        for i in range(0, self.natoms):
+            pos = self.rdmol.GetConformer().GetAtomPositions(i)
+            xyz.append([pos.x, pos.y, pos.z])
+        return np.mean(xyz, axis=0)
+
+
+    def rg(self) -> float:
+        """Returns the radius of gyration (Rg).
+
+        Returns:
+            float: the radius of gyration (Rg).
+        """
+        xyz = []
+        for i in range(0, self.natoms):
+            pos = self.rdmol.GetConformer().GetAtomPositions(i)
+            xyz.append([pos.x, pos.y, pos.z])
+        xyz = np.array(xyz)
+        cog = np.mean(xyz, axis=0)
+        a = xyz-cog
+        b = np.einsum('ij,ij->i', a, a)
+        return np.sqrt(np.mean(b))
+
+
+    def serialize(self, key:str='') -> dict:
+        """Returns JSON dumps of the `props`.
+
+        Args:
+            key (str): a key for the `props` dictionary. Defaults to '' (all).
+
+        Returns:
+            dict: JSON dumps.
+        """
+        if key:
+            return json.dumps({key:self.props[key]})
+        else:
+            return json.dumps(self.props)
+    
+
+    def to_sdf(self, props:bool=True) -> str:
+        """Returns the SDF-formatted strings.
+
+        Args:
+            props (bool, optional): include `props as SDF properties. Defaults to True.
+
+        Returns:
+            str: strings in the SDF format.
+        """
+        in_memory = io.StringIO()
+        with Chem.SDWriter(in_memory) as f:
+            rdmol = Chem.Mol(self.rdmol)
+            rdmol.SetProp('_Name', self.name)
+            if props:
+                for k,v in self.props.items():
+                    rdmol.SetProp(k, str(v))
+            f.write(rdmol)
+        return in_memory.getvalue()
+    
+    
+    def to_svg(self, 
+               width:int=400, 
+               height:int=400,
+               legend:Optional[str]=None, 
+               atom_index:bool=False, 
+               highlight:Optional[List[int]]=None) -> str:
+        """Returns 2D SVG depiction of the molecule.
+
+        Examples:
+            >>> from IPython.display import SVG
+            >>> SVG(libr[0].confs[0].to_svg(atom_index=True))
+
+        Args:
+            width (int): width (default:400)
+            height (int): height (default:400)
+            legend (str, optional): title or Mol.name if not given
+            atom_index (bool): True/False whether to display atom index
+            highlight (list): list of atom indices to highlight
+            
+        Returns:
+            str: SVG text
+        """
+        rdmol_3d = Chem.Mol(self.rdmol) # a copy of self.rdmol (3D, with hydrogens)
+        AllChem.Compute2DCoords(rdmol_3d) # 2D depiction
+        for atom in rdmol_3d.GetAtoms():
+            for key in atom.GetPropsAsDict():
+                atom.ClearProp(key)
+        drawer = rdMolDraw2D.MolDraw2DSVG(width, height)
+        if not legend:
+            legend = self.name
+        if atom_index:
+            for atom in rdmol_3d.GetAtoms():
+                atom.SetProp("atomLabel", str(atom.GetIdx()))
+        if highlight:
+            drawer.DrawMolecule(rdmol_3d, legend=legend, highlightAtoms=highlight)
+        else:
+            drawer.DrawMolecule(rdmol_3d, legend=legend)
+        drawer.FinishDrawing()
+        return drawer.GetDrawingText()

rdworks/descriptor.py

@@ -0,0 +1,36 @@
+from rdkit.Chem import Descriptors, rdMolDescriptors, QED
+
+rd_descriptor = {
+    'QED'         : "Quantitative estimate of drug-likeness.",
+    'MolWt'       : 'Molecular weight',
+    'LogP'        : 'Predicted octanol/water partition coefficient',
+    'TPSA'        : 'Topological polar surface area',
+    'HBD'         : 'Number of hydrogen bonding donors',
+    'HBA'         : 'Number of hydrogen bonding acceptors',
+    'RotBonds'    : 'Number of rotatable bonds',
+    'RingCount'   : 'Number of rings',
+    'FCsp3'       : 'Fraction of SP3 carbons',
+    'HAC'         : 'Number of heavy atoms',
+    'Hetero'      : 'Number of hetero atoms (not H or C) [B,N,O,P,S,F,Cl,Br,I]',
+    'LipinskiHBA' : 'Number of hydrogen bonding acceptors according to the Lipinski definition',
+    'LipinskiHBD' : 'Number of hydrogen bonding donors according to the Lipinski definition',
+}
+
+rd_descriptor_f = {
+    "QED"           : QED.qed,
+    "MolWt"         : Descriptors.MolWt,
+    "HAC"           : Descriptors.HeavyAtomCount,
+    "LogP"          : Descriptors.MolLogP,                      # == Crippen.MolLogP
+    "TPSA"          : Descriptors.TPSA,                         # == MolSurf.TPSA
+    "HBA"           : rdMolDescriptors.CalcNumHBA,              # == Descriptors.NumHAcceptors
+    "HBD"           : rdMolDescriptors.CalcNumHBD,              # == Descriptors.NumHDonors
+    "RotBonds"      : rdMolDescriptors.CalcNumRotatableBonds,   # == Descriptors.NumRotatableBonds
+    "RingCount"     : rdMolDescriptors.CalcNumRings,            # == Descriptors.RingCount
+    "FCsp3"         : rdMolDescriptors.CalcFractionCSP3,        # == Descriptors.FractionCSP3
+    "Hetero"        : rdMolDescriptors.CalcNumHeteroatoms,      # not (H or C) [B,N,O,P,S,F,Cl,Br,I]
+    "LipinskiHBA"   : rdMolDescriptors.CalcNumLipinskiHBA,
+    "LipinskiHBD"   : rdMolDescriptors.CalcNumLipinskiHBD,
+    # "StereoCenters"     : rdMolDescriptors.CalcNumAtomStereoCenters,
+    # props_dict[k] = rd_descriptor_f[k](self.rdmol)
+    # ValueError: numStereoCenters called without stereo being assigned
+}

rdworks/display.py

@@ -0,0 +1,206 @@
+import io
+import os
+import numpy as np
+from typing import Optional, List, Tuple
+
+from PIL import Image, ImageChops
+
+from rdkit import Chem, Geometry
+from rdkit.Chem import AllChem, Draw, rdDepictor, rdMolTransforms
+from rdkit.Chem.Draw import rdMolDraw2D
+
+
+# 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.
+
+    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.
+
+    Returns:
+        Chem.Mol: a copy of rdkit.Chem.Mol object.
+    """
+    if coordgen:
+        rdDepictor.SetPreferCoordGen(True)
+    else:
+        rdDepictor.SetPreferCoordGen(False)
+    
+    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()))
+    
+    return rdmol_2d
+
+
+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.
+
+    Examples:
+        >>> 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.
+
+    Returns:
+        str: SVG text
+    """
+    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()
+
+
+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.
+
+    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.
+    
+    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)
+
+
+def trim_png(img:Image.Image) -> Image.Image:
+    """Removes white margin around molecular drawing.
+
+    Args:
+        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 rescale(rdmol:Chem.Mol, factor:float=1.5) -> Chem.Mol:
+    """Returns a copy of `rdmol` by a `factor`.
+
+    Args:
+        rdmol (Chem.Mol): input molecule.
+        factor (float): scaling factor.
+    
+    Returns:
+        Chem.Mol: a copy of rescaled rdkit.Chem.Mol object.
+    """
+    transformed_rdmol = Chem.Mol(rdmol)
+    center = AllChem.ComputeCentroid(transformed_rdmol.GetConformer())
+    tf = np.identity(4, np.float)
+    tf[0][3] -= center[0]
+    tf[1][3] -= center[1]
+    tf[0][0] = tf[1][1] = tf[2][2] = factor
+    AllChem.TransformMol(transformed_rdmol, tf)
+    return transformed_rdmol
+
+
+def rotation_matrix(axis:str, degree:float) -> np.ndarray:
+    """Returns a numpy rotation matrix of shape (4,4).
+    
+    Args:
+        axis (str): 'x' or 'y' or 'z'.
+        degree (float): degree of rotation.
+
+    Returns:
+        np.ndarray: a numpy array of shape (4,4).
+    """
+    rad = (np.pi/180.0) * degree
+    c = np.cos(rad)
+    s = np.sin(rad)
+    if axis.lower() == 'x':
+        return np.array([
+            [1., 0., 0., 0.],
+            [0., c, -s,  0.],
+            [0., s,  c,  0.],
+            [0., 0., 0., 1.],
+            ])
+    elif axis.lower() == 'y':
+        return np.array([
+            [ c,  0., s,  0.],
+            [ 0., 1., 0., 0.],
+            [-s,  0., c,  0.],
+            [ 0., 0., 0., 1.],
+            ])
+    elif axis.lower() == 'z':
+        return np.array([
+            [c, -s,  0., 0.],
+            [s,  c,  0., 0.],
+            [0., 0., 1., 0.],
+            [0., 0., 0., 1.],
+            ])
+
+
+def rotate(rdmol:Chem.Mol, axis:str, degree:float) -> None:
+    """Rotate `rdmol` around given axis and degree.
+
+    Input `rdmol` will be modified.
+
+    Args:
+        rdmol (Chem.Mol): input molecule.
+        axis (str): axis of rotation, 'x' or 'y' or 'z'.
+        degree (float): degree of rotation.
+    """
+    try:
+        conf = rdmol.GetConformer()
+    except:
+        AllChem.Compute2DCoords(rdmol)
+        conf = rdmol.GetConformer()
+    R = rotation_matrix(axis, degree)
+    rdMolTransforms.TransformConformer(conf, R)