diff --git a/rmgpy/molecule/draw.py b/rmgpy/molecule/draw.py index e18611b92a..3792772af3 100644 --- a/rmgpy/molecule/draw.py +++ b/rmgpy/molecule/draw.py @@ -49,6 +49,7 @@ import math import os.path import re +import itertools try: import cairocffi as cairo @@ -60,7 +61,8 @@ import numpy as np from rdkit.Chem import AllChem -from rmgpy.molecule.molecule import Atom, Molecule +from rmgpy.molecule.molecule import Atom, Molecule, Bond +from rmgpy.molecule.pathfinder import find_shortest_path from rmgpy.qm.molecule import Geometry @@ -96,6 +98,12 @@ def create_new_surface(file_format, target=None, width=1024, height=768): ################################################################################ +class AdsorbateDrawingError(Exception): + """ + When something goes wrong trying to draw an adsorbate. + """ + pass + class MoleculeDrawer(object): """ This class provides functionality for drawing the skeletal formula of @@ -207,7 +215,16 @@ def draw(self, molecule, file_format, target=None): # replace the bonds after generating coordinates. This avoids # bugs with RDKit old_bond_dictionary = self._make_single_bonds() - self._generate_coordinates() + if molecule.contains_surface_site(): + try: + self._connect_surface_sites() + self._generate_coordinates() + self._disconnect_surface_sites() + except AdsorbateDrawingError as e: + self._disconnect_surface_sites() + self._generate_coordinates(fix_surface_sites=False) + else: + self._generate_coordinates() self._replace_bonds(old_bond_dictionary) # Generate labels to use @@ -323,11 +340,13 @@ def _find_ring_groups(self): if not found: self.ringSystems.append([cycle]) - def _generate_coordinates(self): + def _generate_coordinates(self, fix_surface_sites=True): """ Generate the 2D coordinates to be used when drawing the current molecule. The function uses rdKits 2D coordinate generation. Updates the self.coordinates Array in place. + If `fix_surface_sites` is True, then the surface sites are placed + at the bottom of the molecule. """ atoms = self.molecule.atoms natoms = len(atoms) @@ -390,7 +409,6 @@ def _generate_coordinates(self): # If two atoms lie on top of each other, push them apart a bit # This is ugly, but at least the mess you end up with isn't as misleading # as leaving everything piled on top of each other at the origin - import itertools for atom1, atom2 in itertools.combinations(backbone, 2): i1, i2 = atoms.index(atom1), atoms.index(atom2) if np.linalg.norm(coordinates[i1, :] - coordinates[i2, :]) < 0.5: @@ -402,7 +420,6 @@ def _generate_coordinates(self): # If two atoms lie on top of each other, push them apart a bit # This is ugly, but at least the mess you end up with isn't as misleading # as leaving everything piled on top of each other at the origin - import itertools for atom1, atom2 in itertools.combinations(backbone, 2): i1, i2 = atoms.index(atom1), atoms.index(atom2) if np.linalg.norm(coordinates[i1, :] - coordinates[i2, :]) < 0.5: @@ -457,26 +474,59 @@ def _generate_coordinates(self): coordinates[:, 0] = temp[:, 1] coordinates[:, 1] = temp[:, 0] - # For surface species, rotate them so the site is at the bottom. - if self.molecule.contains_surface_site(): + # For surface species + if fix_surface_sites and self.molecule.contains_surface_site(): if len(self.molecule.atoms) == 1: return coordinates - for site in self.molecule.atoms: - if site.is_surface_site(): - break - else: - raise Exception("Can't find surface site") - if site.bonds: - adsorbate = next(iter(site.bonds)) - vector0 = coordinates[atoms.index(site), :] - coordinates[atoms.index(adsorbate), :] - angle = math.atan2(vector0[0], vector0[1]) - math.pi + sites = [atom for atom in self.molecule.atoms if atom.is_surface_site()] + if len(sites) == 1: + # rotate them so the site is at the bottom. + site = sites[0] + if site.bonds: + adatom = next(iter(site.bonds)) + vector0 = coordinates[atoms.index(site), :] - coordinates[atoms.index(adatom), :] + angle = math.atan2(vector0[0], vector0[1]) - math.pi + rot = np.array([[math.cos(angle), math.sin(angle)], [-math.sin(angle), math.cos(angle)]], float) + self.coordinates = coordinates = np.dot(coordinates, rot) + else: + # van der Waals + index = atoms.index(site) + coordinates[index, 1] = min(coordinates[:, 1]) - 0.8 # just move the site down a bit + coordinates[index, 0] = coordinates[:, 0].mean() # and center it + elif len(sites) <= 4: + # Rotate so the line of best fit through the adatoms is horizontal. + # find atoms bonded to sites + adatoms = [next(iter(site.bonds)) for site in sites] + adatom_indices = [atoms.index(a) for a in adatoms] + # find the best fit line through the bonded atoms + x = coordinates[adatom_indices, 0] + y = coordinates[adatom_indices, 1] + A = np.vstack([x, np.ones(len(x))]).T + m, c = np.linalg.lstsq(A, y, rcond=None)[0] + # rotate so the line is horizontal + angle = -math.atan(m) rot = np.array([[math.cos(angle), math.sin(angle)], [-math.sin(angle), math.cos(angle)]], float) self.coordinates = coordinates = np.dot(coordinates, rot) + # if the line is above the middle, flip it + not_site_indices = [atoms.index(a) for a in atoms if not a.is_surface_site()] + if coordinates[adatom_indices, 1].mean() > coordinates[not_site_indices, 1].mean(): + coordinates[:, 1] *= -1 + x = coordinates[adatom_indices, 0] + y = coordinates[adatom_indices, 1] + site_y_pos = min(min(y) - 0.8, min(coordinates[not_site_indices, 1]) - 0.5) + if max(y) - site_y_pos > 1.5: + raise AdsorbateDrawingError("Adsorbate bond too long") + for x1, x2 in itertools.combinations(x, 2): + if abs(x1 - x2) < 0.2: + raise AdsorbateDrawingError("Sites overlapping") + for site, x_pos in zip(sites, x): + index = atoms.index(site) + coordinates[index, 1] = site_y_pos + coordinates[index, 0] = x_pos + else: - # van der waals - index = atoms.index(site) - coordinates[index, 1] = min(coordinates[:, 1]) - 0.8 # just move the site down a bit - coordinates[index, 0] = coordinates[:, 0].mean() # and center it + # more than 4 surface sites? leave them alone + pass def _find_cyclic_backbone(self): """ @@ -854,7 +904,7 @@ def _generate_functional_group_coordinates(self, atom0, atom1): # Check to see if atom1 is in any cycles in the molecule ring_system = None for ring_sys in self.ringSystems: - if any([atom1 in ring for ring in ring_sys]): + if any(atom1 in ring for ring in ring_sys): ring_system = ring_sys if ring_system is not None: @@ -1624,6 +1674,40 @@ def _replace_bonds(self, bond_order_dictionary): for bond, order in bond_order_dictionary.items(): bond.set_order_num(order) + def _connect_surface_sites(self): + """ + Creates single bonds between atoms that are surface sites. + This is to help make multidentate adsorbates look better. + """ + sites = [a for a in self.molecule.atoms if a.is_surface_site()] + if len(sites) > 4: + return + for site1 in sites: + other_sites = [a for a in sites if a != site1] + if not other_sites: break + # connect to the nearest site + site2 = min(other_sites, key=lambda a: len(find_shortest_path(site1, a))) + if len(find_shortest_path(site1, site2)) > 2 and len(sites) > 3: + # if there are more than 3 sites, don't connect sites that aren't neighbors + continue + + bond = site1.bonds.get(site2) + if bond is None: + bond = Bond(site1, site2, 1) + site1.bonds[site2] = bond + site2.bonds[site1] = bond + + def _disconnect_surface_sites(self): + """ + Removes all bonds between atoms that are surface sites. + """ + for site1 in self.molecule.atoms: + if site1.is_surface_site(): + for site2 in list(site1.bonds.keys()): # make a list copy so we can delete from the dict + if site2.is_surface_site(): + del site1.bonds[site2] + del site2.bonds[site1] + ################################################################################ diff --git a/test/rmgpy/molecule/drawTest.py b/test/rmgpy/molecule/drawTest.py index 03e68a0c21..6c416f171d 100644 --- a/test/rmgpy/molecule/drawTest.py +++ b/test/rmgpy/molecule/drawTest.py @@ -33,9 +33,9 @@ import os import os.path +import itertools - -from rmgpy.molecule import Molecule +from rmgpy.molecule import Molecule, Atom, Bond from rmgpy.molecule.draw import MoleculeDrawer from rmgpy.species import Species @@ -144,3 +144,157 @@ def test_draw_hydrogen_bond_adsorbate(self): from cairo import PDFSurface surface, _cr, (_xoff, _yoff, _width, _height) = self.drawer.draw(molecule, file_format="pdf") assert isinstance(surface, PDFSurface) + + def test_draw_bidentate_adsorbate(self): + try: + from cairocffi import ImageSurface + except ImportError: + from cairo import ImageSurface + + test_molecules = [ + Molecule().from_adjacency_list( + """ +1 C u0 p0 c0 {2,S} {4,S} {5,S} {6,S} +2 C u0 p0 c0 {1,S} {3,S} {7,S} {8,S} +3 C u0 p0 c0 {2,S} {9,S} {10,S} {11,S} +4 X u0 p0 c0 {1,S} +5 H u0 p0 c0 {1,S} +6 H u0 p0 c0 {1,S} +7 H u0 p0 c0 {2,S} +8 X u0 p0 c0 {2,S} +9 H u0 p0 c0 {3,S} +10 H u0 p0 c0 {3,S} +11 H u0 p0 c0 {3,S} + """), + Molecule().from_adjacency_list( +""" +1 C u0 p0 c0 {2,S} {4,S} {5,S} {6,S} +2 C u0 p0 c0 {1,S} {3,S} {7,S} {8,S} +3 C u0 p0 c0 {2,S} {9,S} {10,S} {11,S} +4 X u0 p0 c0 {1,S} +5 H u0 p0 c0 {1,S} +6 H u0 p0 c0 {1,S} +7 H u0 p0 c0 {2,S} +8 H u0 p0 c0 {2,S} +9 H u0 p0 c0 {3,S} +10 H u0 p0 c0 {3,S} +11 X u0 p0 c0 {3,S} +"""), + ] + for molecule in [Molecule(smiles="CC(=O)CCO"), Molecule(smiles="C1CC=CC1CC"), Molecule(smiles="C=CCC(O)=O")]: + bondable = [a for a in molecule.atoms if a.is_non_hydrogen() and any(b.is_hydrogen() for b in a.bonds)] + for b1, b2 in itertools.combinations(bondable, 2): + # find a hydrogen atom bonded to each of the two atoms + for h1 in b1.bonds: + if h1.is_hydrogen(): + break + for h2 in b2.bonds: + if h2.is_hydrogen(): + break + molecule.remove_atom(h1) + molecule.remove_atom(h2) + x1 = Atom(element='X', radical_electrons=0, charge=0, label='', lone_pairs=0) + x2 = Atom(element='X', radical_electrons=0, charge=0, label='', lone_pairs=0) + molecule.add_atom(x1) + molecule.add_atom(x2) + molecule.add_bond(Bond(b1, x1, order=1)) + molecule.add_bond(Bond(b2, x2, order=1)) + test_molecules.append(molecule.copy(deep=True)) + molecule.remove_atom(x1) + molecule.remove_atom(x2) + molecule.add_atom(h1) + molecule.add_atom(h2) + molecule.add_bond(Bond(b1, h1, order=1)) + molecule.add_bond(Bond(b2, h2, order=1)) + + for b1, b2, b3 in itertools.combinations(bondable, 3): + # find a hydrogen atom bonded to each of the two atoms + for h1 in b1.bonds: + if h1.is_hydrogen(): + break + for h2 in b2.bonds: + if h2.is_hydrogen(): + break + for h3 in b3.bonds: + if h3.is_hydrogen(): + break + molecule.remove_atom(h1) + molecule.remove_atom(h2) + molecule.remove_atom(h3) + x1 = Atom(element='X', radical_electrons=0, charge=0, label='', lone_pairs=0) + x2 = Atom(element='X', radical_electrons=0, charge=0, label='', lone_pairs=0) + x3 = Atom(element='X', radical_electrons=0, charge=0, label='', lone_pairs=0) + molecule.add_atom(x1) + molecule.add_atom(x2) + molecule.add_atom(x3) + molecule.add_bond(Bond(b1, x1, order=1)) + molecule.add_bond(Bond(b2, x2, order=1)) + molecule.add_bond(Bond(b3, x3, order=1)) + test_molecules.append(molecule.copy(deep=True)) + molecule.remove_atom(x1) + molecule.remove_atom(x2) + molecule.remove_atom(x3) + molecule.add_atom(h1) + molecule.add_atom(h2) + molecule.add_atom(h3) + molecule.add_bond(Bond(b1, h1, order=1)) + molecule.add_bond(Bond(b2, h2, order=1)) + molecule.add_bond(Bond(b3, h3, order=1)) + + test_molecules.append(Molecule().from_adjacency_list( +""" +1 C u0 p0 c0 {2,S} {3,S} {9,S} {10,S} +2 X u0 p0 c0 {1,S} +3 C u0 p0 c0 {1,S} {4,S} {5,S} {11,S} +4 X u0 p0 c0 {3,S} +5 C u0 p0 c0 {3,S} {6,S} {7,S} {12,S} +6 X u0 p0 c0 {5,S} +7 C u0 p0 c0 {5,S} {8,S} {13,S} {14,S} +8 X u0 p0 c0 {7,S} +9 H u0 p0 c0 {1,S} +10 H u0 p0 c0 {1,S} +11 H u0 p0 c0 {3,S} +12 H u0 p0 c0 {5,S} +13 H u0 p0 c0 {7,S} +14 H u0 p0 c0 {7,S} +""")) + test_molecules.append(Molecule().from_adjacency_list( +""" +1 O u0 p2 c0 {4,S} {5,S} +2 O u0 p2 c0 {5,S} {6,S} +3 O u0 p2 c0 {6,S} {26,S} +4 C u0 p0 c0 {1,S} {7,S} {8,S} {19,S} +5 C u0 p0 c0 {1,S} {2,S} {10,S} {21,S} +6 C u0 p0 c0 {2,S} {3,S} {9,S} {20,S} +7 C u0 p0 c0 {4,S} {15,S} {16,S} {24,S} +8 C u0 p0 c0 {4,S} {17,S} {18,S} {25,S} +9 C u0 p0 c0 {6,S} {11,S} {12,S} {22,S} +10 C u0 p0 c0 {5,S} {13,S} {14,S} {23,S} +11 H u0 p0 c0 {9,S} +12 H u0 p0 c0 {9,S} +13 H u0 p0 c0 {10,S} +14 H u0 p0 c0 {10,S} +15 H u0 p0 c0 {7,S} +16 H u0 p0 c0 {7,S} +17 H u0 p0 c0 {8,S} +18 H u0 p0 c0 {8,S} +19 X u0 p0 c0 {4,S} +20 X u0 p0 c0 {6,S} +21 X u0 p0 c0 {5,S} +22 X u0 p0 c0 {9,S} +23 X u0 p0 c0 {10,S} +24 X u0 p0 c0 {7,S} +25 X u0 p0 c0 {8,S} +26 X u0 p0 c0 {3,S} +""")) + test_molecules.append(Molecule(smiles="*CC(*)(C*)OCC#*")) + test_molecules.append(Molecule(smiles="*CC(*)(C*)C*")) + for number, molecule in enumerate(test_molecules, 1): + path = f"test_polydentate_{number}.png" + if os.path.exists(path): + os.unlink(path) + self.drawer.clear() + surface, _cr, (_xoff, _yoff, width, height) = self.drawer.draw(molecule, file_format="png", target=path) + assert os.path.exists(path), "File doesn't exist" + os.unlink(path) + assert isinstance(surface, ImageSurface) \ No newline at end of file