A quick molecule in Processing.js and Python

Code Science Visualization

As an extension to my last post on rendering Processing code in an IPython notebook, I thought it might be fun to play a bit with the 3D functionality and see how easy it would be to build an extremely basic molecule viewer! I didn’t spend much time polishing it, but the basic rendering turned out to be straight forward, and linking it with Python made pre-processing a breeze. Perhaps someone out there will find it useful! So let’s get into the anatomy of rendering a molecule in 3D with a combination of IPython and Processing.js.

We’ll use the same format as before, that is a Python string of Processing code, that we modify using the format method.  We’ll start with a little bit of Processing setup,

int screen_size = 320; //Screen size in pixels
float rx=0; //Rotation angle for view
float ry=0; //Rotation angle for view

void setup() {{
  size(320, 320, P3D);
}}

void draw() {{
  translate(width / 4, height / 4, -100);
  lights();
  background(0);
  rotateX(rx);
  rotateY(ry);
{}
}}

void mouseDragged() {{
  //Add basic click and drag interactivity
  rx += (pmouseY-mouseY)*.005; 
  ry += (pmouseX-mouseX)*-.005; 
}}

Which sets us up with a basic 3D scene, and a bit of mouse interactivity. Note that the double braces are for the Python format method to ignore them in replacement. Now we just need to place something functional into the draw() method. For example, we’ll need to be able to draw an atom, which is simple enough

void atom(float x, float y, float z) {{
 //Draw an atom centered at (x, y, z)
 pushMatrix();
 translate(screen_size * x, screen_size * y, screen_size * z);
 fill(250);
 //noFill();
 noStroke();
 //stroke(255);
 sphere(screen_size * .05);
 popMatrix();
}}

and a bond between two points, which takes a bit more code in the Processing style

void bond(float x1, float y1, float z1,
          float x2, float y2, float z2) {{
  //Draw a bond between (x1,y1,z1) and (x2, y2, z2)
  float rho = screen_size * sqrt((x1-x2) * (x1-x2) + 
                                 (y1-y2) * (y1-y2) + 
                                 (z1-z2) * (z1-z2));
  float theta = atan2((y2-y1),(x2-x1));
  float phi = acos(screen_size * (z2-z1) / rho);
  int sides = 5; //Bond resolution
  float bond_radius = .015*screen_size; //Bond radius
  float angle = 2 * PI / sides;  
  
  pushMatrix();
  noStroke();
  fill(255);
  translate(screen_size * x1, screen_size * y1, screen_size * z1);
  rotateZ(theta);
  rotateY(phi); 
  beginShape(TRIANGLE_STRIP);
  for (int i = 0; i < sides + 1; i++) {{
    float x = cos( i * angle ) * bond_radius;
    float y = sin( i * angle ) * bond_radius;
    vertex( x, y, 0.0);
    vertex( x, y, rho);    
  }}
  endShape(CLOSE);
  popMatrix();
}}

With these primitive rendering functions defined, let’s turn to where we get our data and how to prepare it. I’m going to use a simple XYZ list of coordinates for adenine as an example. That is, I have a file adenine.xyz that reads as follows

15
7H-Purin-6-amine
  N     -2.2027     -0.0935     -0.0068
  C     -0.9319      0.4989     -0.0046
  C      0.0032     -0.5640     -0.0064
  N     -0.7036     -1.7813     -0.0164
  C     -2.0115     -1.4937     -0.0151
  C      1.3811     -0.2282     -0.0017
  N      1.7207      1.0926     -0.0252
  C      0.7470      2.0612     -0.0194
  N     -0.5782      1.8297     -0.0072
  H      1.0891      3.1044     -0.0278
  N      2.4103     -1.1617      0.1252
  H     -3.0709      0.3774     -0.0035
  H     -2.8131     -2.2379     -0.0191
  H      2.1765     -2.0715     -0.1952
  H      3.3110     -0.8521     -0.1580

To render something in Processing3D, the coordinate system starts from 0, and is oriented perhaps in the opposite way that you’d expect. As a result, it’s prudent to load this data in, shift, and rescale it to fit inside a unit cube. Moreover, xyz files don’t contain information about bonds, so we can quickly infer if they might exist based on a distance calculation. There are certainly better ways to do this, and bond distance should depend on atom types, but this works for a quick pass. We can do this in Python as follows:

from IPython.display import HTML
import numpy as np

xyz_filename = 'adenine.xyz'

#Load molecule data from a file
atom_ids = [line.split()[0] for line in open(xyz_filename, 'r') if len(line.split()) > 1]
atom_xyz = np.array([[float(line.split()[1]), float(line.split()[2]), float(line.split()[3])]
                     for line in open(xyz_filename, 'r') if len(line.split()) > 1])

#Find potential bonds between nearby atoms
bonds = []
bond_cutoff = 1.5**2
for i, atom1 in enumerate(atom_xyz):
    for j, atom2 in enumerate(atom_xyz[:i]):
        if (np.sum((atom1 - atom2)**2) < bond_cutoff):
            bonds.append([i,j])
            
#Rescale so largest distance between atoms is bounded by 1, and all coordinates start from 0 reference
min_coord, max_coord = np.amin(atom_xyz, axis=0), np.amax(atom_xyz, axis=0)
scale_factor = 1.0/np.sqrt(np.sum((max_coord-min_coord)**2))
translation_factor = np.repeat([min_coord], atom_xyz.shape[0], axis=0)
atom_xyz -= translation_factor
atom_xyz *= scale_factor

From here, all we have to do is prepare a small template string, loop through the atoms and bonds, and render:

atom_template = """
atom({}, {}, {});
"""

bond_template = """
bond({}, {}, {}, {}, {}, {});
"""

#Render bonds that will sit between atoms
bond_code = ""
for bond in bonds:
    bond_code += bond_template.format(*np.concatenate((atom_xyz[bond[0]], atom_xyz[bond[1]])))
    
#Render atoms
atom_code = ""
for i, atom in enumerate(atom_xyz):
    atom_code += atom_template.format(*atom)

processing_code = processing_template.format(bond_code + atom_code)

Now we make an HTML template, and ask IPython to render it:

html_template = """
<script type="text/javascript" src="processing.min.js"></script> 
<script type="text/javascript">
  var processingCode = `{}`;
  var myCanvas = document.getElementById("canvas1");
  var jsCode = Processing.compile(processingCode);
  var processingInstance = new Processing(myCanvas, jsCode);
 </script>
<canvas id="canvas1"> </canvas>    
"""
html_code = html_template.format(processing_code)
HTML(html_code)

If everything went as planned, hopefully you’ll now have a somewhat crude rendering of adenine to play with and rotate. You could even put it on your webpage with Processing.js if you dump the code to a file like

with open('adenine.pde', 'w') as processing_file:
    processing_file.write(processing_code)

which hopefully gives you the crude rendering below (try clicking a dragging)

Your browser does not currently support the canvas element.