Adding functions to display HoukMol figures in Matplotlib plots

Author: swheele2

atoms.py

@@ -16,6 +16,7 @@
     SATURATION,
     TMETAL,
     VDW_RADII,
+    COLORS,
 )
 
 warn_LJ = set([])
@@ -1248,3 +1249,40 @@ def get_vsepr(self):
                 best_shape = shape
 
         return best_shape, best_score
+
+
+    def draw_atom(self, ax, fp = 40, ascale = 0.5, linewidth=0.1):
+        """
+        Draw HoukMol style atom and add to Matplotlib axis ax
+
+        :param matplotlib.pyplot.Axis ax: Matplotlib axis object
+        :param float fp: z-value for focal point to add simple perspective (Default: 40)
+        :param float ascale: scaling factor for covalent radii (Default: 0.5)
+        :param float linewidth: linewidth (in points) for atom details (default: 0.1)
+        """
+
+        try:
+            import matplotlib.patches as patches
+        except:
+            ls.LOG.error("Must install matplot lib")
+            return None
+
+        x, y, z = self.coords
+        r = RADII[self.element]*ascale*(fp + z)/fp
+        color=COLORS[self.element]
+
+        # draw circle and two ellipses for atom plus two circles for specular highlight
+        circle = patches.Circle((x, y), r, facecolor=color, edgecolor='black', lw=linewidth, zorder=z)
+        acc = patches.Circle((x+r/6, y+r/4), 0.2*r, alpha=0.7, facecolor='white', edgecolor="None", zorder=z)
+        acc2 = patches.Circle((x+r/6, y+r/4), 0.3*r, alpha=0.3, facecolor='white', edgecolor="None", zorder=z)
+        arc1 = patches.Arc((x, y), r, 2*r, theta1=270, theta2=90, edgecolor='black', lw=linewidth, zorder=z)
+        arc2 = patches.Arc((x, y), 2*r, r, theta1=180, theta2=360, edgecolor='black', lw=linewidth, zorder=z)
+
+        # add atom componets to axis
+        ax.add_patch(circle)
+        ax.add_patch(acc)
+        ax.add_patch(acc2)
+        ax.add_patch(arc1)
+        ax.add_patch(arc2)
+
+

const.py

@@ -2796,6 +2796,27 @@
     "Au",
 )
 
+# Rasmol colors from jmol up to Cl
+COLORS={
+        "H":"#FFFFFF",
+        "He": "#FFC0CB",
+        "Li": "#B22222",
+        "Be": "#FF1493",
+        "B": "#00FF00",
+        "C": "#C8C8C8",
+        "N": "#8F8FFF",
+        "O": "#F00000",
+        "F": "#DAA520",
+        "Ne": "#FF1493",
+        "Na": "#0000FF",
+        "Mg": "#228B22",
+        "Al": "#808090",
+        "Si": "#DAA520",
+        "P": "#FFA500",
+        "S": "#FFC832",
+        "Cl": "00FF00"
+}
+
 class PHYSICAL:
     # Physical constants
     BOLTZMANN = 0.001987204  # kcal/mol

geometry.py

@@ -16,7 +16,7 @@
 from AaronTools import addlogger
 from AaronTools.atoms import Atom, BondOrder
 from AaronTools.config import Config
-from AaronTools.const import AARONLIB, AARONTOOLS, BONDI_RADII, D_CUTOFF, ELEMENTS, TMETAL, VDW_RADII
+from AaronTools.const import AARONLIB, AARONTOOLS, BONDI_RADII, D_CUTOFF, ELEMENTS, TMETAL, VDW_RADII, RADII
 from AaronTools.fileIO import FileReader, FileWriter, read_types
 from AaronTools.finders import Finder, OfType, WithinRadiusFromPoint, WithinRadiusFromAtom 
 from AaronTools.utils.prime_numbers import Primes
@@ -1013,6 +1013,92 @@ def convert_to_Psi4(self, charge=0, mult=1, fix_com=True, fix_orientation=True):
         return mol 
 
 
+    # quick and dirty code to display HoukMol style figures in Matplotlib
+    # bond ends could be handled more elegently, but this looks fine for most molecules
+
+    def plot(self, ax, fig, fp=40, ascale=0.5, bwidth=0.15):
+        """
+        displays HoukMol style molecule in Matplotlib
+
+        :param matplotlib.pyplot.Axis ax: Matplotlib Axis object
+        :param matplotlib.pyplot.Figure fig: Matplotlib Figure object:w
+        :param float fp: z-value (Angstroms) for focal point for adding perspective (Default: 40)
+        :param float ascale: scaling factor for covalent radii (default: 0.5)
+        :param float bwidth: scaling factor for bond radii (default: 0.15)
+        """
+
+        try:
+            import matplotlib.patches as patches
+        except:
+            self.LOG.error("Must install matplot lib")
+            return None
+
+        def intersect(atom1, atom2, fp):
+            # returns intersection of line from atom1 to edge of scaled sphere for atom2
+            # from https://en.wikipedia.org/wiki/Line%E2%80%93sphere_intersection
+        
+            c = atom2.coords             # center of sphere
+            r = RADII[atom2.element]*0.5*(fp + atom2.coords[2])/fp # radius of sphere, adjusted for perspective
+            o = atom1.coords             # starting point for line
+            u = atom1.bond(atom2)        # vector along line
+            unorm = u/np.sqrt(np.linalg.norm(u)) # normalized vector
+        
+            dot = np.dot(unorm, o - c)
+            delta = dot**2 - (np.linalg.norm(o - c)**2 - r**2)
+        
+            if delta <= 0:
+                # catch cases where there is no intersection for some reason
+                return c
+            else:
+                d = dot + np.sqrt(delta)
+                return o - d*u*1.0 # adjust endpoints slightly to account for rounded capstyle
+
+
+
+        def draw_bond(x1, y1, x2, y2, z, fp, scale, ax, bwidth=0.15):
+            """
+            Draw HoukMol style bond as black line with rounded ends from (x1, y1) to (x2, y2)
+            width of bond is controlled by width and scaled to z-value for perspective
+            """
+        
+            w=(bwidth*scale) * (fp + z)/fp
+            ax.plot([x1, x2], [y1, y2], lw=w, color='black', solid_capstyle='round', zorder=z)
+
+        
+        # if xlim not set manually then I can't determine the overall size of the plot until
+        # after the molecule is drawn, but I need overall size to determine scale for bond
+        # widths
+
+        if ax.get_xaxis()._get_autoscale_on():
+            self.LOG.warning("You will probalby need to set xlim manually for correct bond widths.")
+
+        # get scale and size of plot set linewidths in terms of pts (1/72 inch per pt)
+        xmin, xmax = ax.get_xlim()
+        dx = xmax - xmin
+        ymin, ymax = ax.get_ylim()
+        dy = ymax - ymin
+        fw, fh = fig.get_size_inches()
+        bbox = ax.get_position()
+        ax_width = fw*bbox.width
+        scale = 72*ax_width/dx
+        
+        # sort atoms by z-value
+        sorted_atoms = [atom for atom in sorted(self.atoms, key=lambda a: a.coords[2])]
+
+        for atom in sorted_atoms:
+            # note that I draw each bond twice to get both ends correct
+            # this avoids having to do the logic of figuring out the order of
+            # drawing bonds to a given atom--each atom obscures bonds originating
+            # from that atom, but this is corrected when the bond is drawn
+            # from the connected atom. Any way I've tried to fix this looks wrong
+            # for planar molecules, which are the only ones I actually care about.
+        
+            for connected in atom.connected:
+                endpoint = intersect(atom, connected, fp)
+                draw_bond(atom.coords[0], atom.coords[1], endpoint[0], endpoint[1], atom.coords[2], fp, scale, ax, bwidth)
+
+            atom.draw_atom(ax, fp, ascale=ascale, linewidth=0.01*scale)
+
     def copy(self, atoms=None, name=None, comment=None, copy_atoms=True):
         """
         creates a new copy of the geometry