<commit_before>#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2009-2011 by the RMG Team (<EMAIL>)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module provides functionality for automatic two-dimensional drawing of the
`skeletal formulae <http://en.wikipedia.org/wiki/Skeletal_formula>`_ of a wide
variety of organic and inorganic molecules. The general method for creating
these drawings is to utilize the :meth:`draw()` method of the :class:`Molecule`
you wish to draw; this wraps a call to :meth:`MoleculeDrawer.draw()`, where the
molecule drawing algorithm begins. Advanced use may require use of the
:class:`MoleculeDrawer` class directly.

The `Cairo <http://cairographics.org/>`_ 2D graphics library is used to create
the drawings. The :class:`MoleculeDrawer` class module will fail gracefully if
Cairo is not installed.

The general procedure for creating drawings of skeletal formula is as follows:

1.  **Find the molecular backbone.** If the molecule contains no cycles, the
    longest straight chain of heavy atoms is used as the backbone. If the 
    molecule contains cycles, the largest independent cycle group is used as the
    backbone. The :meth:`findBackbone()` method is used for this purpose.

2.  **Generate coordinates for the backbone atoms.** Straight-chain backbones
    are laid out in a horizontal seesaw pattern. Cyclic backbones are laid out
    as regular polygons (or as close to this as is possible). The
    :meth:`generateStraightChainCoordinates()` and 
    :meth:`generateRingSystemCoordinates()` methods are used for this purpose.

3.  **Generate coordinates for immediate neighbors to backbone.** Each neighbor
    atom represents the start of a functional group attached to the backbone.
    Generating coordinates for these means that we have determined the bonds
    for all backbone atoms. The :meth:`generateNeighborCoordinates()` method is
    used for this purpose.

4.  **Continue generating coordinates for atoms in functional groups.** Moving
    away from the molecular backbone and its immediate neighbors, the
    coordinates for each atom in each functional group are determined such that
    the functional groups tend to radiate away from the center of the backbone
    (to reduce chances of overlap). If cycles are encountered in the functional
    groups, their coordinates are processed as a unit. This continues until
    the coordinates of all atoms in the molecule have been assigned. The
    :meth:`generateFunctionalGroupCoordinates()` recursive method is used for
    this.

5.  **Use the generated coordinates and the atom and bond types to render the
    skeletal formula.** The :meth:`render()`,  and :meth:`renderBond()`, and
    :meth:`renderAtom()` methods are used for this.

The developed procedure seems to be rather robust, but occasionally it will
encounter a molecule that it renders incorrectly. In particular, features which
have not yet been implemented by this drawing algorithm include:

* cis-trans isomerism

* stereoisomerism

* bridging atoms in fused rings

"""

import math
import numpy
import os.path
import re
import logging

from rmgpy.qm.molecule import Geometry
from rdkit.Chem import AllChem

from numpy.linalg import LinAlgError

################################################################################

def createNewSurface(format, path=None, width=1024, height=768):
    """
    Create a new surface of the specified `type`: "png" for
    :class:`ImageSurface`, "svg" for :class:`SVGSurface`, "pdf" for
    :class:`PDFSurface`, or "ps" for :class:`PSSurface`. If the surface is to
    be saved to a file, use the `path` parameter to give the path to the file.
    You can also optionally specify the `width` and `height` of the generated
    surface if you know what it is; otherwise a default size of 1024 by 768 is
    used.
    """
    import cairo
    format = format.lower()
    if format == 'png':
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, int(width), int(height))
    elif format == 'svg':
        surface = cairo.SVGSurface(path, width, height)
    elif format == 'pdf':
        surface = cairo.PDFSurface(path, width, height)
    elif format == 'ps':
        surface = cairo.PSSurface(path, width, height)
    else:
        raise ValueError('Invalid value "{0}" for type parameter; valid values are "png", "svg", "pdf", and "ps".'.format(type))
    return surface
    
################################################################################

class MoleculeDrawer:
    """
    This class provides functionality for drawing the skeletal formula of
    molecules using the Cairo 2D graphics engine. The most common use case is
    simply::
    
        MoleculeDrawer().draw(molecule, format='png', path='molecule.png')
    
    where ``molecule`` is the :class:`Molecule` object to draw. You can also
    pass a dict of options to the constructor to affect how the molecules are
    drawn.
    """
    
    def __init__(self, options=None):
        self.options = {
            'fontFamily': 'sans',
            'fontSizeNormal': 12,
            'fontSizeSubscript': 8,
            'bondLength': 24,
            'padding': 2,
        }
        if options: self.options.update(options)
        self.clear()
    
    def clear(self):
        self.molecule = None
        self.cycles = None
        self.ringSystems = None
        self.coordinates = None
        self.symbols = None
        self.implicitHydrogens = None
        self.left = 0.0
        self.top = 0.0
        self.right = 0.0
        self.bottom = 0.0
        self.surface = None
        self.cr = None
        
    def draw(self, molecule, format, path=None):
        """
        Draw the given `molecule` using the given image `format` - pdf, svg, ps, or
        png. If `path` is given, the drawing is saved to that location on disk. The
        `options` dict is an optional set of key-value pairs that can be used to
        control the generated drawing.
        
        This function returns the Cairo surface and context used to create the
        drawing, as well as a bounding box for the molecule being drawn as the
        tuple (`left`, `top`, `width`, `height`).
        """
        
        # The Cairo 2D graphics library (and its Python wrapper) is required for
        # the molecule drawing algorithm
        try:
            import cairo
        except ImportError:
            print 'Cairo not found; molecule will not be drawn.'
            return
        
        # Make a copy of the molecule so we don't modify the original
        self.molecule = molecule.copy(deep=True)
        
        # Remove all unlabeled hydrogen atoms from the copied atoms and bonds, as
        # they are not drawn
        # However, if this would remove all atoms, then don't remove any
        atomsToRemove = []
        self.implicitHydrogens = {}
        for atom in self.molecule.atoms:
            if atom.isHydrogen() and atom.label == '': atomsToRemove.append(atom)
        if len(atomsToRemove) < len(self.molecule.atoms):
            for atom in atomsToRemove:
                for atom2 in atom.bonds:
                    try:
                        self.implicitHydrogens[atom2] += 1
                    except KeyError:
                        self.implicitHydrogens[atom2] = 1
                self.molecule.removeAtom(atom)
    
        # Generate information about any cycles present in the molecule, as
        # they will need special attention
        self.__findRingGroups()
        
        # Generate the coordinates to use to draw the molecule
        try:
            self.__generateCoordinates()
        except (ValueError, numpy.linalg.LinAlgError), e:
            logging.error('Error while drawing molecule {0}: {1}'.format(molecule.toSMILES(), e))
            import sys, traceback
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback.print_exc()
            return None, None, None

        self.coordinates[:,1] *= -1
        self.coordinates *= self.options['bondLength']
        
        # Generate labels to use
        self.__generateAtomLabels()
        
        # Handle some special cases
        if self.symbols == ['H','H']:
            # Render as H2 instead of H-H
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['H2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['O', 'O']:
            # Render as O2 instead of O-O
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.molecule.atoms[0].radicalElectrons = 0
            self.symbols = ['O2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['OH', 'O'] or self.symbols == ['O', 'OH']:
            # Render as HO2 instead of HO-O or O-OH
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['O2H']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['OH', 'OH']:
            # Render as H2O2 instead of HO-OH or O-OH
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['O2H2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['C', 'O'] or self.symbols == ['O', 'C']:
            # Render as CO instead of C=O
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['CO']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['O', 'C', 'O']:
            # Render as CO2 instead of O=C=O
            self.molecule.removeAtom(self.molecule.atoms[0])
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['CO2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        
        # Create a dummy surface to draw to, since we don't know the bounding rect
        # We will copy this to another surface with the correct bounding rect
        surface0 = createNewSurface(format=format, path=None)
        cr0 = cairo.Context(surface0)
    
        # Render using Cairo
        self.render(cr0)
        
        # Create the real surface with the appropriate size
        xoff = self.left
        yoff = self.top
        width = self.right - self.left
        height = self.bottom - self.top
        self.surface = createNewSurface(format=format, path=path, width=width, height=height)
        self.cr = cairo.Context(self.surface)

        # Draw white background
        self.cr.set_source_rgba(1.0, 1.0, 1.0, 1.0)
        self.cr.paint()        

        self.render(self.cr, offset=(-xoff,-yoff))

        if path is not None:
            # Finish Cairo drawing
            # Save PNG of drawing if appropriate
            ext = os.path.splitext(path)[1].lower()
            if ext == '.png':
                self.surface.write_to_png(path)
            else:
                self.surface.finish()
    
        return self.surface, self.cr, (xoff, yoff, width, height)

    def __findRingGroups(self):
        """
        Find all of the cycles in the current molecule, and group them into
        sets of adjacent cycles.
        """
                
        # Find all of the cycles in the molecule
        self.cycles = self.molecule.getSmallestSetOfSmallestRings()
        self.ringSystems = []
        
        # If the molecule contains cycles, find them and group them
        if len(self.cycles) > 0:
            # Split the list of cycles into groups
            # Each atom in the molecule should belong to exactly zero or one such groups
            for cycle in self.cycles:
                found = False
                for ringSystem in self.ringSystems:
                    for ring in ringSystem:
                        if any([atom in ring for atom in cycle]) and not found:
                            ringSystem.append(cycle)
                            found = True
                if not found:
                    self.ringSystems.append([cycle])
    
    def __generateCoordinates(self):
        """
        Generate the 2D coordinates to be used when drawing the current 
        molecule. The vertices are arranged based on a standard bond length of
        unity, and can be scaled later for longer bond lengths. The coordinates
        are arranged such that the "center" of the molecule is at the origin.
        This function ignores any previously-existing coordinate information.
        """
        atoms = self.molecule.atoms
        Natoms = len(atoms)
        
        # Initialize array of coordinates
        self.coordinates = coordinates = numpy.zeros((Natoms, 2))
        <commit_msg>Cleaned up comments to reflect changes.<commit_after>#!/usr/bin/env python
# encoding: utf-8

################################################################################
#
#   RMG - Reaction Mechanism Generator
#
#   Copyright (c) 2009-2011 by the RMG Team (<EMAIL>)
#
#   Permission is hereby granted, free of charge, to any person obtaining a
#   copy of this software and associated documentation files (the 'Software'),
#   to deal in the Software without restriction, including without limitation
#   the rights to use, copy, modify, merge, publish, distribute, sublicense,
#   and/or sell copies of the Software, and to permit persons to whom the
#   Software is furnished to do so, subject to the following conditions:
#
#   The above copyright notice and this permission notice shall be included in
#   all copies or substantial portions of the Software.
#
#   THE SOFTWARE IS PROVIDED 'AS IS', WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
#   IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
#   FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
#   AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
#   LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
#   FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
#   DEALINGS IN THE SOFTWARE.
#
################################################################################

"""
This module provides functionality for automatic two-dimensional drawing of the
`skeletal formulae <http://en.wikipedia.org/wiki/Skeletal_formula>`_ of a wide
variety of organic and inorganic molecules. The general method for creating
these drawings is to utilize the :meth:`draw()` method of the :class:`Molecule`
you wish to draw; this wraps a call to :meth:`MoleculeDrawer.draw()`, where the
molecule drawing algorithm begins. Advanced use may require use of the
:class:`MoleculeDrawer` class directly.

The `Cairo <http://cairographics.org/>`_ 2D graphics library is used to create
the drawings. The :class:`MoleculeDrawer` class module will fail gracefully if
Cairo is not installed.

The implementation uses the 2D coordinate generation of rdKit to find coordinates,
then uses Cairo to render the atom.

"""

import math
import numpy
import os.path
import re
import logging

from rmgpy.qm.molecule import Geometry
from rdkit.Chem import AllChem

from numpy.linalg import LinAlgError

################################################################################

def createNewSurface(format, path=None, width=1024, height=768):
    """
    Create a new surface of the specified `type`: "png" for
    :class:`ImageSurface`, "svg" for :class:`SVGSurface`, "pdf" for
    :class:`PDFSurface`, or "ps" for :class:`PSSurface`. If the surface is to
    be saved to a file, use the `path` parameter to give the path to the file.
    You can also optionally specify the `width` and `height` of the generated
    surface if you know what it is; otherwise a default size of 1024 by 768 is
    used.
    """
    import cairo
    format = format.lower()
    if format == 'png':
        surface = cairo.ImageSurface(cairo.FORMAT_ARGB32, int(width), int(height))
    elif format == 'svg':
        surface = cairo.SVGSurface(path, width, height)
    elif format == 'pdf':
        surface = cairo.PDFSurface(path, width, height)
    elif format == 'ps':
        surface = cairo.PSSurface(path, width, height)
    else:
        raise ValueError('Invalid value "{0}" for type parameter; valid values are "png", "svg", "pdf", and "ps".'.format(type))
    return surface
    
################################################################################

class MoleculeDrawer:
    """
    This class provides functionality for drawing the skeletal formula of
    molecules using the Cairo 2D graphics engine. The most common use case is
    simply::
    
        MoleculeDrawer().draw(molecule, format='png', path='molecule.png')
    
    where ``molecule`` is the :class:`Molecule` object to draw. You can also
    pass a dict of options to the constructor to affect how the molecules are
    drawn.
    """
    
    def __init__(self, options=None):
        self.options = {
            'fontFamily': 'sans',
            'fontSizeNormal': 12,
            'fontSizeSubscript': 8,
            'bondLength': 24,
            'padding': 2,
        }
        if options: self.options.update(options)
        self.clear()
    
    def clear(self):
        self.molecule = None
        self.cycles = None
        self.ringSystems = None
        self.coordinates = None
        self.symbols = None
        self.implicitHydrogens = None
        self.left = 0.0
        self.top = 0.0
        self.right = 0.0
        self.bottom = 0.0
        self.surface = None
        self.cr = None
        
    def draw(self, molecule, format, path=None):
        """
        Draw the given `molecule` using the given image `format` - pdf, svg, ps, or
        png. If `path` is given, the drawing is saved to that location on disk. The
        `options` dict is an optional set of key-value pairs that can be used to
        control the generated drawing.
        
        This function returns the Cairo surface and context used to create the
        drawing, as well as a bounding box for the molecule being drawn as the
        tuple (`left`, `top`, `width`, `height`).
        """
        
        # The Cairo 2D graphics library (and its Python wrapper) is required for
        # the molecule drawing algorithm
        try:
            import cairo
        except ImportError:
            print 'Cairo not found; molecule will not be drawn.'
            return
        
        # Make a copy of the molecule so we don't modify the original
        self.molecule = molecule.copy(deep=True)
        
        # Remove all unlabeled hydrogen atoms from the copied atoms and bonds, as
        # they are not drawn
        # However, if this would remove all atoms, then don't remove any
        atomsToRemove = []
        self.implicitHydrogens = {}
        for atom in self.molecule.atoms:
            if atom.isHydrogen() and atom.label == '': atomsToRemove.append(atom)
        if len(atomsToRemove) < len(self.molecule.atoms):
            for atom in atomsToRemove:
                for atom2 in atom.bonds:
                    try:
                        self.implicitHydrogens[atom2] += 1
                    except KeyError:
                        self.implicitHydrogens[atom2] = 1
                self.molecule.removeAtom(atom)
    
        # Generate information about any cycles present in the molecule, as
        # they will need special attention
        self.__findRingGroups()
        
        # Generate the coordinates to use to draw the molecule
        try:
            self.__generateCoordinates()
        except (ValueError, numpy.linalg.LinAlgError), e:
            logging.error('Error while drawing molecule {0}: {1}'.format(molecule.toSMILES(), e))
            import sys, traceback
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback.print_exc()
            return None, None, None

        self.coordinates[:,1] *= -1
        self.coordinates *= self.options['bondLength']
        
        # Generate labels to use
        self.__generateAtomLabels()
        
        # Handle some special cases
        if self.symbols == ['H','H']:
            # Render as H2 instead of H-H
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['H2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['O', 'O']:
            # Render as O2 instead of O-O
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.molecule.atoms[0].radicalElectrons = 0
            self.symbols = ['O2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['OH', 'O'] or self.symbols == ['O', 'OH']:
            # Render as HO2 instead of HO-O or O-OH
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['O2H']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['OH', 'OH']:
            # Render as H2O2 instead of HO-OH or O-OH
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['O2H2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['C', 'O'] or self.symbols == ['O', 'C']:
            # Render as CO instead of C=O
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['CO']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        elif self.symbols == ['O', 'C', 'O']:
            # Render as CO2 instead of O=C=O
            self.molecule.removeAtom(self.molecule.atoms[0])
            self.molecule.removeAtom(self.molecule.atoms[-1])
            self.symbols = ['CO2']
            self.coordinates = numpy.array([[0,0]], numpy.float64)
        
        # Create a dummy surface to draw to, since we don't know the bounding rect
        # We will copy this to another surface with the correct bounding rect
        surface0 = createNewSurface(format=format, path=None)
        cr0 = cairo.Context(surface0)
    
        # Render using Cairo
        self.render(cr0)
        
        # Create the real surface with the appropriate size
        xoff = self.left
        yoff = self.top
        width = self.right - self.left
        height = self.bottom - self.top
        self.surface = createNewSurface(format=format, path=path, width=width, height=height)
        self.cr = cairo.Context(self.surface)

        # Draw white background
        self.cr.set_source_rgba(1.0, 1.0, 1.0, 1.0)
        self.cr.paint()        

        self.render(self.cr, offset=(-xoff,-yoff))

        if path is not None:
            # Finish Cairo drawing
            # Save PNG of drawing if appropriate
            ext = os.path.splitext(path)[1].lower()
            if ext == '.png':
                self.surface.write_to_png(path)
            else:
                self.surface.finish()
    
        return self.surface, self.cr, (xoff, yoff, width, height)

    def __findRingGroups(self):
        """
        Find all of the cycles in the current molecule, and group them into
        sets of adjacent cycles.
        """
                
        # Find all of the cycles in the molecule
        self.cycles = self.molecule.getSmallestSetOfSmallestRings()
        self.ringSystems = []
        
        # If the molecule contains cycles, find them and group them
        if len(self.cycles) > 0:
            # Split the list of cycles into groups
            # Each atom in the molecule should belong to exactly zero or one such groups
            for cycle in self.cycles:
                found = False
                for ringSystem in self.ringSystems:
                    for ring in ringSystem:
                        if any([atom in ring for atom in cycle]) and not found:
                            ringSystem.append(cycle)
                            found = True
                if not found:
                    self.ringSystems.append([cycle])
    
    def __generateCoordinates(self):
        """
        Generate the 2D coordinates to be used when drawing the current 
        molecule. The function uses rdKits 2D coordinate generation.
        """
        atoms = self.molecule.atoms
        Natoms = len(atoms)
        
        # Initialize array of coordinates
        self.coordinates = coordinates = numpy.zeros((Natoms, 2))
        
{'fraction_non_alphanumeric': 0.11035746306474876, 'fraction_numerical': 0.009082758351849686, 'mean_word_length': 2.3790832455216018, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4516575', 'n_tokens_mistral': 7148, 'n_tokens_neox': 6711, 'n_words': 2559}