Open Babel to VRML script

pywrl.py

#!/usr/bin/env python

import sys, os, math, random

from openbabel import pybel
from openbabel import openbabel as ob

colors = [ [ 0, 0, 0 ], # Du
            [0.75, 0.75, 0.75], # H
            [0.85, 1.00, 1.00], # He
            [0.80, 0.50, 1.00], # Li
            [0.76, 1.00, 0.00], # Be
            [1.00, 0.71, 0.71], # B
            [0.40, 0.40, 0.40], # C
            [0.05, 0.05, 1.00], # N
            [1.00, 0.05, 0.05], # O
            [0.50, 0.70, 1.00], # F
            [0.70, 0.89, 0.96], # Ne
            [0.67, 0.36, 0.95], # Na
            [0.54, 1.00, 0.00], # Mg
            [0.75, 0.65, 0.65], # Al
            [0.50, 0.60, 0.60], # Si
            [1.00, 0.50, 0.00], # P
            [0.70, 0.70, 0.00], # S
            [0.12, 0.94, 0.12], # Cl
            [0.50, 0.82, 0.89], # Ar
]

def calcBond(a1, a2):
    (x1, y1, z1) = (a1.x(), a1.y(), a1.z())
    (x2, y2, z2) = (a2.x(), a2.y(), a2.z())

    dx = x2 - x1
    dy = y2 - y1
    dz = z2 - z1

    length = math.sqrt(dx*dx + dy*dy + dz*dz)
    # center (for translation)
    (tx, ty, tz) = (dx / 2.0 + x1, dy / 2.0 + y1, dz / 2.0 + z1)
    # normalize
    dx = dx/length
    dy = dy/length
    dz = dz/length

    ax = ay = az = angle = 0.0
    if (dy > 0.999):
        (ax, ay, az, angle) = (1.0, 0.0, 0.0, 0.0)
    elif (dy < -0.999):
        (ax, ay, az, angle) = (1.0, 0.0, 0.0, math.pi)
    else:
        (ax, ay, az, angle) = (dz, 0.0, -dx, math.acos(dy))

    return (tx, ty, tz, length / 2.0, ax, ay, az, angle)

def fibonacci_sphere(samples=1,randomize=True):
    rnd = 1.
    if randomize:
        rnd = random.random() * samples

    points = []
    offset = 2./samples
    increment = math.pi * (3. - math.sqrt(5.));

    for i in range(samples):
        y = ((i * offset) - 1) + (offset / 2);
        r = math.sqrt(1 - pow(y,2))

        phi = ((i + rnd) % samples) * increment

        x = math.cos(phi) * r
        z = math.sin(phi) * r

        points.append([x,y,z])

    return points

c = -1
for argument in sys.argv[1:]:
    fileName, fileExt = os.path.splitext(argument)
    #print("Reading", fileName, fileExt[1:])
    mol = next(pybel.readfile(fileExt[1:], argument))

    # inertial frame
    mol.OBMol.Center()
    mol.OBMol.ToInertialFrame()

    # translate for a grid
    c += 1
    i = c % 7
    j = c / 7
    mol.OBMol.Translate(ob.vector3(i*10.0, j*10.0, 0.0))

    print("""#VRML V2.0 utf8
WorldInfo {
title "%s"
info [ "By Avogadro" ]
}
NavigationInfo {
 type ["EXAMINE","ANY"]
 headlight TRUE
 visibilityLimit 0.0
 speed 1.0
}""" % (mol.title))
    print("""DEF DefaultView Viewpoint {
 position 0 0 57.558
 description "Default view"
 fieldOfView 0.785398
}""")

    minX = 1000
    minY = 1000
    minZ = 1000
    maxX = -1000
    maxY = -1000
    maxZ = -1000
    for atom in mol.atoms:
        (x, y, z) = atom.coords
        minX = min(x, minX)
        minY = min(y, minY)
        minZ = min(z, minZ)
        maxX = max(x, maxX)
        maxY = max(y, maxY)
        maxZ = max(z, maxZ)

    # set the scale
    maxDim = max( maxX - minX, maxY - minY, maxZ - minZ)
    # TODO - use as an option for scale
    # 1.0 = 1 mm, so set range to 5 cm
    scale = 6.67
    # 10 mm = 1.0A

    for atom in mol.atoms:
        (x, y, z) = atom.coords
        x = x*scale
        y = y*scale
        z = z*scale
        radius = ob.GetVdwRad(atom.atomicnum) * scale * 0.45
#        radius = ob.etab.GetVdwRad(atom.atomicnum) * scale * 0.45
        if radius < 1.0:
            radius = 1.0
        red = green = blue = 0.0
        try:
            (red, green, blue) = colors[atom.atomicnum]
        except:
            print("Color ", atom.atomicnum, colors[atom.atomicnum])
#        (red, green, blue) = ob.GetRGB(atom.atomicnum)
#        (red, green, blue) = ob.etab.GetRGB(atom.atomicnum)
        print("""Transform {
 translation %8.3f %8.3f %8.3f
 children Shape {
  geometry Sphere {
   radius %8.3f
  }
  appearance Appearance {
   material Material {
    diffuseColor %8.2f %8.2f %8.2f
   }
  }
 }
}""" % (x, y, z, radius, red, green, blue))

    for bond in ob.OBMolBondIter(mol.OBMol):
        a1 = bond.GetBeginAtom()
        a2 = bond.GetEndAtom()

        (tx, ty, tz, length, ax, ay, az, angle) = calcBond(a1, a2)
        radius = 2.0 + (bond.GetBondOrder() / 2.0)
        if bond.IsAromatic():
            radius = 2.0 + (0.75)
        print("""Transform {
 translation %8.3f %8.3f %8.3f
 scale 1 %8.3f 1
 rotation %8.3f %8.3f %8.3f %8.3f
 children Shape {
  geometry Cylinder {
   radius %8.3f
  }
  appearance Appearance {
   material Material {
    diffuseColor 0.3 0.3 0.3
   }
  }
 }
}""" % (tx*scale, ty*scale, tz*scale, length*scale, ax, ay, az, angle, radius))