#
# This code is Copyright (C) 2015 The Cambridge Crystallographic Data Centre
# (CCDC) of 12 Union Road, Cambridge CB2 1EZ, UK and a proprietary work of CCDC.
# This code may not be used, reproduced, translated, modified, disassembled or
# copied, except in accordance with a valid licence agreement with CCDC and may
# not be disclosed or redistributed in any form, either in whole or in part, to
# any third party. All copies of this code made in accordance with a valid
# licence agreement as referred to above must contain this copyright notice.
#
# No representations, warranties, or liabilities are expressed or implied in the
# supply of this code by CCDC, its servants or agents, except where such
# exclusion or limitation is prohibited, void or unenforceable under governing
# law.
#
'''
The :mod:`ccdc.diagram` module has functionality for generating 2D diagrams.
The main class of the :mod:`ccdc.diagram` module is
:class:`ccdc.diagram.DiagramGenerator`.
>>> from ccdc.io import EntryReader
>>> from ccdc.diagram import DiagramGenerator
>>> entry_reader = EntryReader('CSD')
>>> mol = entry_reader.molecule('ABEBUF')
>>> diagram_generator = DiagramGenerator()
>>> img = diagram_generator.image(mol) # img is a PIL image
'''
#############################################################################
try:
from PIL import Image
except ImportError:
import Image
import pathlib
from ccdc.entry import Entry
from ccdc.crystal import Crystal
from ccdc.molecule import Molecule
from ccdc.utilities import _argb2rgba, bidirectional_dict
from ccdc.search import QuerySubstructure, QueryAtom, MoleculeSubstructure
from ccdc.io import MoleculeWriter
from ccdc.utilities import _private_importer
with _private_importer() as pi:
pi.import_ccdc_module('ChemistryLib')
pi.import_ccdc_module('VirtualScreening')
pi.import_ccdc_module('DatabaseEntryLib')
pi.import_ccdc_module('FileFormatsLib')
#############################################################################
[docs]class DiagramGenerator(object):
'''Diagram generator.'''
[docs] class Settings(object):
'''Settings base class for diagram generation.'''
_font_weights = bidirectional_dict(
light=300,
normal=400,
demibold=600,
bold=700,
black=900
)
def __init__(self, _settings=None):
if _settings is None:
_settings = VirtualScreening.ChemicalDiagramImageGeneratorSettings()
self._settings = _settings
self._return_type = 'PIL' #: Format of returned image, either 'PIL' for a PIL image object or 'SVG' for an SVG string.
self.highlight_color = 'red' #: Color used for highlighting atoms and bonds.
self.element_coloring = True #: Color atomic labels.
@property
def return_type(self):
'''The image method return type.'''
return self._return_type
@return_type.setter
def return_type(self, value):
'''Set the return type of the image method, the following values are acceptable.
"PIL" - The image method returns the diagram image as a PIL image object.
"SVG" - The image method returns the diagram image as an SVG string.'''
self._return_type = value.upper()
@property
def image_width(self):
'''The image width.'''
return self._settings.image_width()
@image_width.setter
def image_width(self, value):
'''Set the image width.'''
self._settings.set_image_width(value)
@property
def image_height(self):
'''The image height.'''
return self._settings.image_height()
@image_height.setter
def image_height(self, value):
'''Set the image height.'''
self._settings.set_image_height(value)
@property
def line_width(self):
'''The image line width.'''
return self._settings.image_line_width()
@line_width.setter
def line_width(self, value):
'''Set the image line width.'''
self._settings.set_image_line_width(value)
@property
def font_family(self):
'''The name of the font family.'''
return self._settings.image_font_family()
@font_family.setter
def font_family(self, value):
self._settings.set_image_font_family(value)
@property
def font_weight(self):
'''The weight of the font.'''
return self._font_weights.inverse_lookup(self._settings.image_font_weight())
@font_weight.setter
def font_weight(self, value):
self._settings.set_image_font_weight(
self._font_weights.prefix_lookup(value)
)
@property
def font_size(self):
'''The image font size.'''
s = self._settings.image_font_size()
if s == -1:
return None
return s
@font_size.setter
def font_size(self, value):
'''Set the image font size.'''
if not value:
value = -1
self._settings.set_image_font_size(value)
@property
def font_italic(self):
'''Whether the font is italic or not.'''
return self._settings.image_font_italic()
@font_italic.setter
def font_italic(self, tf):
self._settings.set_image_font_italic(tf)
@property
def foreground_color(self):
'''The foreground color.'''
return self._settings.image_foreground_color()
@foreground_color.setter
def foreground_color(self, value):
self._settings.set_image_foreground_color(VirtualScreening.QColor(value));
@property
def background_color(self):
'''The background color.'''
return self._settings.image_background_color()
@foreground_color.setter
def background_color(self, value):
self._settings.set_image_background_color(VirtualScreening.QColor(value));
@property
def highlight_color(self):
'''The highlight color to be used.'''
return self._settings.image_highlight_color().name()
@highlight_color.setter
def highlight_color(self, value):
self._settings.set_image_highlight_color(VirtualScreening.QColor(value))
@property
def element_coloring(self):
'''Whether or not elements will be coloured.'''
return self._settings.image_element_color_policy() == VirtualScreening.ChemicalDiagramImageGenerator.ELEMENT_COLOR
@element_coloring.setter
def element_coloring(self, tf):
if tf:
self._settings.set_image_element_color_policy(VirtualScreening.ChemicalDiagramImageGenerator.ELEMENT_COLOR)
else:
self._settings.set_image_element_color_policy(VirtualScreening.ChemicalDiagramImageGenerator.FIXED_COLOR)
shrink_symbols = True #: Whether to replace particular substructures with functional group labels.
explicit_polar_hydrogens = False #: Whether to use explicit polar hydrogens.
detect_intra_hbonds = False #: Whether to detect intra-molecular hbonds.
_detect_inter_hbonds = False # Whether to detect inter-molecular hbonds.
overwrite_existing_image = False #: Whether to override the existing image of an entry.
def __init__(self, settings=None):
if settings is None:
settings = self.Settings()
self.settings = settings
self._ff_gen = VirtualScreening.ForceField2DGenerator()
def _generate_image(self, molecule, highlight_atoms, label_atoms, image_type):
generate_new_diagram = True
if isinstance(molecule, Entry):
entry = molecule
if not self.settings.overwrite_existing_image:
generate_new_diagram = False
else:
entry = Entry.from_molecule(molecule)
if generate_new_diagram:
cv = ChemistryLib.CrystalStructureView.instantiate(entry.crystal._crystal)
self._ff_gen.set_seed(10000)
m2d = VirtualScreening.Molecules2Diagram(self._ff_gen,
cv)
cd = m2d.diagram()
if label_atoms is not None:
for a in label_atoms:
if a.index < cd.natoms():
cd.atom(a.index).set_centre_label(a.label)
cdv = ChemistryLib.ChemicalDiagramViews2D(cd)
entry._entry.set_chemical_diagram_views(cdv)
cdmatcher = DatabaseEntryLib.CrystalChemicalDiagramMatcher(
entry.crystal._crystal, cdv)
cdm = cdmatcher.crystal_diagram_match()
entry._entry.set_crystal_diagram_match(cdm)
atom_labels = []
if highlight_atoms:
atom_labels = [atom.label for atom in highlight_atoms]
if image_type == 'PIL':
image = VirtualScreening.ChemicalDiagramImageGenerator().write(
entry._entry,
self.settings._settings,
atom_labels
)
return DiagramGenerator._extract_pil(image)
elif image_type == 'SVG':
image = VirtualScreening.ChemicalDiagramImageGenerator().write_svg(
entry._entry,
self.settings._settings,
atom_labels
)
return image
@staticmethod
def _errors():
return VirtualScreening.ChemicalDiagramImageGenerator().errors()
@staticmethod
def _extract_pil(image):
if image:
try:
pil = Image.frombytes(
'RGBA',
(image.width(), image.height()),
_argb2rgba(image.data())
)
except AttributeError:
pil = Image.fromstring(
'RGBA',
(image.width(), image.height()),
_argb2rgba(image.data())
)
return pil
[docs] def image(self, argument, highlight_atoms=None, label_atoms=None):
'''Return an image or list of images according to arguments.
:param argument: a :class:`ccdc.crystal.Crystal`, a
:class:`ccdc.molecule.Molecule`, a
:class:`ccdc.entry.Entry`, or a list of them.
:param highlight_atoms: list of :class:`ccdc.molecule.Atom`,
or a list of atom lists
:returns: a PIL image, or a list of PIL images. It may return ``None`` if the
diagram generation failed.
'''
def _check_crystal(c):
try:
c.molecule
return c
except TypeError:
return None
def _get_highlight_atoms_list(atoms, i):
if not atoms:
return None
try:
a = atoms[i]
if isinstance(a, (list, tuple)):
return a
else:
return None
except IndexError:
return None
def _check_highlight_atoms(atoms):
if not atoms or not len(atoms):
return None
elif isinstance(atoms[0], (list, tuple)):
return atoms[0]
else:
return atoms
# Settings for ForceField2DGenerator
self._ff_gen.set_shrink_known_symbols(self.settings.shrink_symbols)
self._ff_gen.set_explicit_polar_hydrogens(self.settings.explicit_polar_hydrogens)
self._ff_gen.set_detect_intra_molecular_hbonds(self.settings.detect_intra_hbonds)
self._ff_gen.set_detect_inter_molecular_hbonds(self.settings._detect_inter_hbonds)
if isinstance(argument, Crystal):
if _check_crystal(argument):
h_atoms = _check_highlight_atoms(highlight_atoms)
l_atoms = _check_highlight_atoms(label_atoms)
return self._generate_image(argument.molecule,
h_atoms,
l_atoms,
image_type=self.settings.return_type)
else:
return None
elif isinstance(argument, Molecule) or isinstance(argument, Entry):
h_atoms = _check_highlight_atoms(highlight_atoms)
l_atoms = _check_highlight_atoms(label_atoms)
return self._generate_image(argument, h_atoms, l_atoms, image_type=self.settings.return_type)
elif isinstance(argument, QuerySubstructure):
if highlight_atoms is None:
highlight_atoms = []
if all(isinstance(i, QueryAtom) for i in highlight_atoms):
highlight_atoms = [a.index for a in highlight_atoms]
if not all(isinstance(i, int) for i in highlight_atoms):
raise RuntimeError('highlight atoms of a QuerySubstructure must be indices')
if self.settings.return_type == 'PIL':
return DiagramGenerator._extract_pil(
VirtualScreening.substructure_to_diagram(
argument._substructure, highlight_atoms, self.settings._settings
)
)
else:
return VirtualScreening.substructure_to_svg(
argument._substructure, highlight_atoms, self.settings._settings
)
elif isinstance(argument, (list, tuple)) and len(argument):
images = []
cur_s = 0
if isinstance(argument[0], Crystal):
l = [_check_crystal(c) for c in argument]
for c in l:
if c is not None:
h_atoms = _get_highlight_atoms_list(highlight_atoms, cur_s)
l_atoms = _get_highlight_atoms_list(label_atoms, cur_s)
img = self._generate_image(c.molecule, h_atoms, l_atoms, image_type=self.settings.return_type)
images.append(img)
else:
images.append(None)
cur_s += 1
elif isinstance(argument[0], Molecule) or isinstance(argument[0], Entry):
for m in argument:
h_atoms = _get_highlight_atoms_list(highlight_atoms, cur_s)
l_atoms = _get_highlight_atoms_list(label_atoms, cur_s)
img = self._generate_image(m, h_atoms, l_atoms, image_type=self.settings.return_type)
images.append(img)
cur_s += 1
return images
[docs] def chemdraw_xml(self, molecule, file_name):
'''Write a chemdraw XML file containing the 2D diagram of the molecule passed in.
:param molecule: a :class:`ccdc.molecule.Molecule` or a :class:`ccdc.entry.Entry`
:param file_name: a file name to write
:returns: the text written to the file, or None
'''
if isinstance(molecule, Molecule) or isinstance(molecule, Entry):
entry = Entry.from_molecule(molecule)
p = pathlib.Path(file_name)
writer = FileFormatsLib.CdxmlFile()
writer.write(entry._entry.chemical_diagram_views(), str(p.absolute()))
with open(str(p.absolute()), 'r', newline='') as cdxml_file:
cdxml = cdxml_file.read()
return cdxml
else:
return None
#############################################################################