Create a Povray image of a molecule

povraymolecules.py

import math
import sys
import numpy as np
import pybel
ob = pybel.ob
from collections import defaultdict

header = """#include "textures.inc"
#include "colors.inc"
#include "functions.inc"
#include "glass.inc"
#include "stones.inc"

camera {location <0,1,-5> look_at <%.3f, %.3f, %.3f>}

light_source {<0,6,-10> White}
#declare Fluorine = texture{ T_Green_Glass}
#declare Oxygen = texture{ T_Ruby_Glass}
#declare Nitrogen =  texture {
   pigment { rgbf <0.2,0.1,1,0.8> } // A blue-tinted glass
    finish { F_Glass4 }
}   
                 
#declare Hydrogen = texture {
pigment {color White}
}

#declare Carbon = texture{ T_Stone10    
                normal { agate 0.25 scale 0.15 rotate<0,0,0> }
                finish { phong 1 } 
                rotate<0,0,0> scale 0.5 translate<0,0,0>
              } // end of texture 
 
"""

pov_sphere = """
sphere {
  <%f, %f, %f>
  %f
  %s
  %s
  %s
}
"""

class Scene:
    def __init__(self, lookat=None):
        if lookat is None:
            lookat = (0, 0, 0)
        self.lookat = lookat
        self.spheres = defaultdict(list)
    def add_sphere(self, id, show, center, radius, texture, transform=None):
        self.spheres[id].append( (center, radius, texture, transform, show) )
    def write(self, filename):
        f = open(filename, "w")
        f.write(header % (self.lookat[0], self.lookat[1], self.lookat[2]))
        sphere_output = get_sphere_output(self.spheres)
        f.write(sphere_output)
        f.close()

def dist(a, b):
    tot = 0
    for i in range(3):
        tot += (a[i]-b[i])**2
    return math.sqrt(tot)

def get_sphere_output(spheres):
    text = []
    for id in spheres.keys():
        mspheres = spheres[id]
        clip_planes = [[] for x in mspheres]
        for i, a in enumerate(mspheres[:-1]):
            for j in range(i+1, len(mspheres)):
                if i==j: continue
                b = mspheres[j]
                if dist(a[0], b[0]) < (a[1]+b[1]): # Do they overlap?
                    clip_planes[i].append(get_clip_plane(b, a))
                    clip_planes[j].append(get_clip_plane(a, b))
        for sphere, cps in zip(mspheres, clip_planes):
            if sphere[-1]: # If show
                text.append(pov_sphere % (sphere[0][0], sphere[0][1], sphere[0][2],
                    sphere[1], sphere[2], "\n".join(cps), "" if not sphere[3] else "translate <%.3f, %.3f, %.3f>" % (sphere[3][0], sphere[3][1], sphere[3][2])))
    return "".join(text)

def unit_vector(vector):
    """ Returns the unit vector of the vector.  """
    norm = np.linalg.norm(vector)
    if abs(norm) < 0.001:
        print vector
        print norm
    ans = vector / norm
    return ans

def angle_between(v1, v2):
    """ Returns the angle in radians between vectors 'v1' and 'v2'::

            >>> angle_between((1, 0, 0), (0, 1, 0))
            1.5707963267948966
            >>> angle_between((1, 0, 0), (1, 0, 0))
            0.0
            >>> angle_between((1, 0, 0), (-1, 0, 0))
            3.141592653589793
    """
    v1_u = unit_vector(v1)
    v2_u = unit_vector(v2)
    angle = np.arccos(np.dot(v1_u, v2_u))
    if np.isnan(angle):
        if (v1_u == v2_u).all():
            return 0.0
        else:
            return np.pi
    return angle

def t(v): # Convert vector3 to np.array
    return np.array([v.GetX(), v.GetY(), v.GetZ()])
def u(ar): # Convert np.array to vector3
    return ob.vector3(float(ar[0]), float(ar[1]), float(ar[2]))

def m(mat): # Convert matrix3x3 to np.array
    ans = np.zeros((3,3), "d")
    for i in range(3):
        for j in range(3):
            ans[i,j] = mat.Get(i, j)
    return ans

def get_clip_plane(a, b):
    d = dist(a[0], b[0])
    x = (d*d - b[1]*b[1] + a[1]*a[1]) / (2*d)

    c = np.array(a[0])
    d = np.array(b[0])
    v = d-c
    nplane = v/np.linalg.norm(v) # Normal to plane

    # How to rotate the x axis onto the normal
    xaxis = np.array( (1.0, 0.0, 0.0) )
    angle = angle_between(nplane, xaxis)
    rotaxis = u(np.cross(xaxis, nplane))

    if angle < 0.001:
        mmat = np.array( [[1.0, 0, 0], [0, 1.0, 0] ,[0, 0, 1.0]])
    else:
        # Get transformation matrix
        matrix = ob.matrix3x3()
        matrix.RotAboutAxisByAngle(rotaxis, angle*180./math.pi)
        mmat = m(matrix)

    return """clipped_by {
  plane {x, %.3f
    inverse}
  matrix <%f, %f, %f,
          %f, %f, %f,
          %f, %f, %f,
          0, 0, 0>
  translate <%.3f, %.3f, %.3f>
  }
""" % (x, mmat[0][0], mmat[0][1], mmat[0][2], mmat[1][0], mmat[1][1], mmat[1][2], mmat[2][0], mmat[2][1], mmat[2][2], c[0], c[1], c[2])

def transform_mol(mol, match):

    pat = ("c" + match)
    smart = pybel.Smarts(pat)
    match = smart.findall(mol)[0]

    a = mol.OBMol.GetAtom(match[0]).GetVector() # the 'c'
    b = mol.OBMol.GetAtom(match[1]).GetVector() # the 'F'
    v = t(b)-t(a) # From the c to the F

    xaxis = np.array( (1.0, 0.0, 0.0) )
    angle = angle_between(v, xaxis)
    rotaxis = u(np.cross(v, xaxis))

    # Get transformation matrix
    matrix = ob.matrix3x3()
    matrix.RotAboutAxisByAngle(rotaxis, angle*180./math.pi)

    # Apply transformation
    for atom in mol:
        v = t(atom.OBAtom.GetVector()) - t(a) # Move the 'c' to (0, 0, )
        v = u(v)
        v *= matrix # Rotate so that cF is along xaxis
        atom.OBAtom.SetVector(v)

if __name__ == "__main__":
    mol = pybel.readfile("mol", "sunitinib.mol").next()

    transform_mol(mol, "F")

    et = ob.OBElementTable()

    colorlookup = {1: "Hydrogen", 6: "Carbon", 8: "Oxygen", 7: "Nitrogen", 9: "Fluorine"}

    for atom in mol:
        if atom.atomicnum == 9:
            fluorine_coords = atom.coords
            break

    pov = Scene(lookat=fluorine_coords)
    for atom in mol:
        transform = (-5, 0, 0) if atom.atomicnum == 9 else None
        pov.add_sphere( 0, True, atom.coords, et.GetVdwRad(atom.atomicnum), "texture{%s}" % colorlookup[atom.atomicnum], transform=transform)


    #for mol in pybel.readfile("sdf", "substs.sdf"):
    #    transform_mol(mol, mol.title)

    pov.write("frontcover.pov")