Source code for luna.mol.depiction

from math import cos, sin, radians

from rdkit import Chem
from rdkit.Chem.Draw import rdMolDraw2D
from rdkit.Chem.AllChem import Compute2DCoords

from luna.wrappers.base import MolWrapper
from luna.util.default_values import ATOM_TYPES_COLOR


[docs]class PharmacophoreDepiction: """Draw molecules and depict pharmacophoric properties as colored circles. Parameters ---------- feature_extractor : :class:`~luna.mol.features.FeatureExtractor` Perceive pharmacophoric properties from molecules. colors : :class:`~luna.util.ColorPallete` Color scheme for pharmacophoric properties perceived by ``feature_extractor``. The default value is :const:`~luna.util.default_values.ATOM_TYPES_COLOR`. format : {'png', 'svg'} The output file format. The default value is 'png'. figsize : tuple of (float, float) Width and height in inches. The default value is (800, 800). font_size : float The font size. The units are, roughly, pixels. The default value is 0.5. circle_dist : float Distance between circles (pharmacophoric properties). The default value is 0.2. circle_radius : Circles' radius size of pharmacophoric properties. The default value is 0.3. use_bw_atom_palette : bool Use a black & white palette for atoms and bonds. Examples -------- First, let's read a molecule (glutamine). >>> from luna.wrappers.base import MolWrapper >>> mol = MolWrapper.from_smiles("N[C@@H](CCC(N)=O)C(O)=O") Now, create a feature factory and instantiate a new FeatureExtractor object. >>> from luna.util.default_values import ATOM_PROP_FILE >>> from rdkit.Chem import ChemicalFeatures >>> from luna.mol.features import FeatureExtractor >>> feature_factory = ChemicalFeatures.BuildFeatureFactory(ATOM_PROP_FILE) >>> feature_extractor = FeatureExtractor(feature_factory) Instantiate a new PharmacophoreDepiction object with the desired configuration. For example, you can provide a color scheme for pharmacophoric properties, the image size, and its format. >>> from luna.util.default_values import ATOM_TYPES_COLOR >>> from luna.mol.depiction import PharmacophoreDepiction pd = PharmacophoreDepiction(feature_extractor=feature_extractor, colors=ATOM_TYPES_COLOR, fig_size=(500, 500), format="svg") Finally, you can draw the molecule with annotated pharmacophoric properties. >>> pd.plot_fig(mol, "output.svg") """ def __init__(self, feature_extractor=None, colors=ATOM_TYPES_COLOR, format="png", fig_size=(800, 800), font_size=0.5, circle_dist=0.2, circle_radius=0.3, use_bw_atom_palette=True): self.feature_extractor = feature_extractor self.colors = colors self.format = format self.fig_size = fig_size self.font_size = font_size self.circle_dist = circle_dist self.circle_radius = circle_radius self.use_bw_atom_palette = use_bw_atom_palette def _perceive_atm_types(self, rdmol): if self.feature_extractor is not None: return self.feature_extractor.get_features_by_atoms(rdmol) return {}
[docs] def plot_fig(self, mol_obj, output=None, atm_types=None, legend=None): """Draw the molecule ``mol_obj`` and depict its pharmacophoric properties. Parameters ---------- mol_obj : :class:`~luna.wrappers.base.MolWrapper`, :class:`rdkit.Chem.rdchem.Mol`, or :class:`openbabel.pybel.Molecule` The molecule. output : str The output file where the molecule will be drawn. If None, returns a drawing object (:class:`~rdkit.Chem.Draw.rdMolDraw2D.MolDraw2DCairo` or :class:`~rdkit.Chem.Draw.rdMolDraw2D.MolDraw2DSVG`). atm_types : dict or None A pre-annotated dictionary for mapping atoms and pharmacophoric properties. If None, try to perceive properties with ``feature_extractor``. legend : str A title for the figure. Returns ------- drawer : None or a drawing object (:class:`~rdkit.Chem.Draw.rdMolDraw2D.MolDraw2DCairo` or \ :class:`~rdkit.Chem.Draw.rdMolDraw2D.MolDraw2DSVG`) """ rdmol = MolWrapper(mol_obj).as_rdkit() # Make a copy of the molecule rwm = Chem.RWMol(rdmol) Compute2DCoords(rwm) if atm_types is None: atm_types = self._perceive_atm_types(rdmol) if self.format == "png": drawer = rdMolDraw2D.MolDraw2DCairo(*self.fig_size) elif self.format == "svg": drawer = rdMolDraw2D.MolDraw2DSVG(*self.fig_size) else: # raise pass opts = drawer.drawOptions() highlight = {} for atm_id in atm_types: centroid = list(rwm.GetConformer().GetAtomPosition(atm_id)) valid_features = [f for f in atm_types[atm_id] if f.name in self.colors] if valid_features: if len(valid_features) == 1: pos = centroid atmIdx = self._add_dummy_atom(rwm, centroid) highlight[atmIdx] = self.colors.get_normalized_color(valid_features[0].name) opts.atomLabels[atmIdx] = '' else: sliceRad = radians(360 / len(valid_features)) for i, feature in enumerate(valid_features): rad = i * sliceRad pos = [self.circle_dist * cos(rad), self.circle_dist * sin(rad), 0] adj_Pos = [x + y for x, y in zip(centroid, pos)] new_atm_id = self._add_dummy_atom(rwm, adj_Pos) highlight[new_atm_id] = self.colors.get_normalized_color(feature.name) opts.atomLabels[new_atm_id] = '' atoms = [x for x in highlight] radius = {a: self.circle_radius for a in atoms} opts.flagCloseContactsDist = -1000 opts.legendFontSize = 20 opts.padding = 0.02 if self.use_bw_atom_palette: opts.useBWAtomPalette() else: opts.useDefaultAtomPalette() legend = legend or "" drawer.SetFontSize(self.font_size) drawer.DrawMolecule(rwm, highlightAtoms=atoms, highlightAtomColors=highlight, highlightBonds=[], highlightAtomRadii=radius, legend=legend) drawer.FinishDrawing() if output: if self.format == "png": drawer.WriteDrawingText(output) elif self.format == "svg": svg = drawer.GetDrawingText().replace('svg:', '') with open(output, "w") as fh: fh.write(svg) else: return drawer
def _add_dummy_atom(self, mol, pos=None): new_atm = Chem.rdchem.Atom(10) atm_id = mol.AddAtom(new_atm) mol.GetConformer().SetAtomPosition(atm_id, pos) new_atm.SetNoImplicit(True) return atm_id