Fix drawing errors for charged species by moving to rdkit geometry generation

rmgpy/molecule/draw.py

@@ -155,7 +155,7 @@ def clear(self):
         self.surface = None
         self.cr = None
 
-    def draw(self, molecule, file_format, target=None):
+    def draw(self, molecule, file_format, target=None, use_rdkit=True):
         """
         Draw the given `molecule` using the given image `file_format` - pdf, svg, ps, or
         png. If `path` is given, the drawing is saved to that location on disk. The
@@ -165,6 +165,9 @@ def draw(self, molecule, file_format, target=None):
         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`).
+
+        If `use_rdkit` is True, then the RDKit 2D coordinate generation is used to generate the coordinates.
+        If `use_rdkit` is False, then the molecule is drawn using our (deprecated) original algorithm.
         """
 
         # The Cairo 2D graphics library (and its Python wrapper) is required for
@@ -219,13 +222,13 @@ def draw(self, molecule, file_format, target=None):
                 if molecule.contains_surface_site():
                     try:
                         self._connect_surface_sites()
-                        self._generate_coordinates()
+                        self._generate_coordinates(use_rdkit=use_rdkit)
                         self._disconnect_surface_sites()
                     except AdsorbateDrawingError as e:
                         self._disconnect_surface_sites()
-                        self._generate_coordinates(fix_surface_sites=False)
+                        self._generate_coordinates(fix_surface_sites=False, use_rdkit=use_rdkit)
                 else:
-                    self._generate_coordinates()
+                    self._generate_coordinates(use_rdkit=use_rdkit)
                 self._replace_bonds(old_bond_dictionary)
 
                 # Generate labels to use
@@ -341,7 +344,7 @@ def _find_ring_groups(self):
                 if not found:
                     self.ringSystems.append([cycle])
 
-    def _generate_coordinates(self, fix_surface_sites=True):
+    def _generate_coordinates(self, fix_surface_sites=True, use_rdkit=True):
         """
         Generate the 2D coordinates to be used when drawing the current
         molecule. The function uses rdKits 2D coordinate generation.
@@ -372,15 +375,34 @@ def _generate_coordinates(self, fix_surface_sites=True):
                 self.coordinates[1, :] = [0.5, 0.0]
             return self.coordinates
 
-        # Decide whether we can use RDKit or have to generate coordinates ourselves
-        for atom in self.molecule.atoms:
-            if atom.charge != 0:
-                use_rdkit = False
-                break
-        else:  # didn't break
-            use_rdkit = True
+        if use_rdkit == True:
+            # Use RDKit 2D coordinate generation:
+
+            # Generate the RDkit molecule from the RDkit molecule, use geometry
+            # in order to match the atoms in the rdmol with the atoms in the
+            # RMG molecule (which is required to extract coordinates).
+            self.geometry = Geometry(None, None, self.molecule, None)
+
+            rdmol, rd_atom_idx = self.geometry.rd_build()
+            AllChem.Compute2DCoords(rdmol)
+
+            # Extract the coordinates from each atom.
+            for atom in atoms:
+                index = rd_atom_idx[atom]
+                point = rdmol.GetConformer(0).GetAtomPosition(index)
+                coordinates[index, :] = [point.x * 0.6, point.y * 0.6]
+
+            # RDKit generates some molecules more vertically than horizontally,
+            # Especially linear ones. This will reflect any molecule taller than
+            # it is wide across the line y=x
+            ranges = np.ptp(coordinates, axis=0)
+            if ranges[1] > ranges[0]:
+                temp = np.copy(coordinates)
+                coordinates[:, 0] = temp[:, 1]
+                coordinates[:, 1] = temp[:, 0]
 
-        if not use_rdkit:
+        else:
+            logging.warning("Using deprecated molecule drawing algorithm; undesired behavior may occur. Consider using use_rdkit=True.")
             if len(self.cycles) > 0:
                 # Cyclic molecule
                 backbone = self._find_cyclic_backbone()
@@ -438,32 +460,6 @@ def _generate_coordinates(self, fix_surface_sites=True):
             # minimize likelihood of overlap
             self._generate_neighbor_coordinates(backbone)
 
-        else:
-            # Use RDKit 2D coordinate generation:
-
-            # Generate the RDkit molecule from the RDkit molecule, use geometry
-            # in order to match the atoms in the rdmol with the atoms in the
-            # RMG molecule (which is required to extract coordinates).
-            self.geometry = Geometry(None, None, self.molecule, None)
-
-            rdmol, rd_atom_idx = self.geometry.rd_build()
-            AllChem.Compute2DCoords(rdmol)
-
-            # Extract the coordinates from each atom.
-            for atom in atoms:
-                index = rd_atom_idx[atom]
-                point = rdmol.GetConformer(0).GetAtomPosition(index)
-                coordinates[index, :] = [point.x * 0.6, point.y * 0.6]
-
-            # RDKit generates some molecules more vertically than horizontally,
-            # Especially linear ones. This will reflect any molecule taller than
-            # it is wide across the line y=x
-            ranges = np.ptp(coordinates, axis=0)
-            if ranges[1] > ranges[0]:
-                temp = np.copy(coordinates)
-                coordinates[:, 0] = temp[:, 1]
-                coordinates[:, 1] = temp[:, 0]
-
         # For surface species
         if fix_surface_sites and self.molecule.contains_surface_site():
             if len(self.molecule.atoms) == 1:

test/rmgpy/molecule/drawTest.py

@@ -297,4 +297,102 @@ def test_draw_bidentate_adsorbate(self):
             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)
+            assert isinstance(surface, ImageSurface)
+
+    def test_draw_bidentate_with_charge_separation(self):
+        molecule = Molecule().from_adjacency_list(
+            """
+1 X u0 p0 c0 {3,S}
+2 X u0 p0 c0 {4,D}
+3 O u0 p2 c0 {1,S} {4,S}
+4 N u0 p0 c+1 {3,S} {2,D} {5,S}
+5 O u0 p3 c-1 {4,S}
+        """
+        )
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        surface, _cr, (_xoff, _yoff, _width, _height) = self.drawer.draw(molecule, file_format="pdf")
+        assert isinstance(surface, PDFSurface)
+
+    def test_draw_cation(self):
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        path = "test_molecule.pdf"
+        if os.path.exists(path):
+            os.unlink(path)
+        polycycle = Molecule(smiles="C1=NC2=C(N1)C(=O)[NH2+]C(=N2)N")
+        surface, _cr, (_xoff, _yoff, width, height) = self.drawer.draw(polycycle, file_format="pdf", target=path)
+        assert isinstance(surface, PDFSurface)
+        assert width > height
+        os.unlink(path)
+
+    def test_draw_anion(self):
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        path = "test_molecule.pdf"
+        if os.path.exists(path):
+            os.unlink(path)
+        polycycle = Molecule(smiles="c1ccc2c3ccccc3[CH-]c2c1")
+        surface, _cr, (_xoff, _yoff, width, height) = self.drawer.draw(polycycle, file_format="pdf", target=path)
+        assert isinstance(surface, PDFSurface)
+        assert width > height
+        os.unlink(path)
+
+    def test_draw_zwitterion(self):
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        path = "test_molecule.pdf"
+        if os.path.exists(path):
+            os.unlink(path)
+        polycycle = Molecule(smiles="[NH3+]CC(=O)[O-]")
+        surface, _cr, (_xoff, _yoff, width, height) = self.drawer.draw(polycycle, file_format="pdf", target=path)
+        assert isinstance(surface, PDFSurface)
+        assert width > height
+        os.unlink(path)
+
+    def test_draw_cation_on_surface(self):
+        molecule = Molecule().from_adjacency_list(
+            """
+1 X u0 p0 c0 {3,S}
+2 X u0 p0 c0 {3,S}
+3 O u0 p1 c+1 {1,S} {2,S} {4,S}
+4 H u0 p0 c0 {3,S}
+        """
+        )
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        path = "test_molecule.pdf"
+        if os.path.exists(path):
+            os.unlink(path)
+        surface, _cr, (_xoff, _yoff, _width, _height) = self.drawer.draw(molecule, file_format="pdf", target=path)
+        assert isinstance(surface, PDFSurface)
+        os.unlink(path)
+
+
+    def test_draw_anion_on_surface(self):
+        molecule = Molecule().from_adjacency_list(
+            """
+1 X u0 p0 c0 {2,S}
+2 O u0 p3 c-1 {1,S}
+        """
+        )
+        try:
+            from cairocffi import PDFSurface
+        except ImportError:
+            from cairo import PDFSurface
+        path = "test_molecule.pdf"
+        if os.path.exists(path):
+            os.unlink(path)
+        surface, _cr, (_xoff, _yoff, _width, _height) = self.drawer.draw(molecule, file_format="pdf", target=path)
+        assert isinstance(surface, PDFSurface)
+        os.unlink(path)