glumpy
glumpy copied to clipboard
Qt integration
Here is an untested script for Qt/glumpy integration:
# -----------------------------------------------------------------------------
# All rights reserved. Distributed under the (new) BSD License.
# -----------------------------------------------------------------------------
from glumpy import app, gloo, gl, data
from glumpy.transforms import Position, Trackball
from glumpy.graphics.filter import Filter
import numpy as np
from PySide import QtCore, QtGui
import openbabel
import pybel
vertex = """
uniform vec3 light_position;
attribute vec3 position;
attribute vec3 color;
attribute float radius;
varying float v_size;
varying vec3 v_color;
varying float v_radius;
varying vec4 v_eye_position;
varying vec3 v_light_direction;
void main (void)
{
v_color = color;
v_radius = radius;
v_eye_position = <transform.trackball_view> *
<transform.trackball_model> *
vec4(position,1.0);
v_light_direction = normalize(light_position);
gl_Position = <transform(position)>;
// stackoverflow.com/questions/8608844/...
// ... resizing-point-sprites-based-on-distance-from-the-camera
vec4 p = <transform.trackball_projection> *
vec4(radius, radius, v_eye_position.z, v_eye_position.w);
v_size = 512.0 * p.x / p.w/2;
gl_PointSize = v_size + 0.0;
}
"""
fragment = """
#include "antialias/outline.glsl"
varying float v_size;
varying vec3 v_color;
varying float v_radius;
varying vec4 v_eye_position;
varying vec3 v_light_direction;
void main()
{
vec2 P = gl_PointCoord.xy - vec2(0.5,0.5);
float point_size = v_size + 0.0;
float distance = length(P*point_size) - v_size/2;
vec2 texcoord = gl_PointCoord* 2.0 - vec2(1.0);
float x = texcoord.x;
float y = texcoord.y;
float d = 1.0 - x*x - y*y;
if (d <= 0.0) discard;
float z = sqrt(d);
vec4 pos = v_eye_position;
pos.z += v_radius*z;
vec3 pos2 = pos.xyz;
pos = <transform.trackball_projection> * pos;
gl_FragDepth = 0.5*(pos.z / pos.w)+0.5;
vec3 normal = vec3(x,y,z);
float diffuse = clamp(dot(normal, v_light_direction), 0.0, 1.0);
vec4 color = vec4((0.5 + 0.5*diffuse)*v_color, 1.0);
//gl_FragColor = outline(distance, 1.0, 1.0, vec4(0,0,0,1), color);
gl_FragColor = color;
}
"""
vertex1 = """
attribute vec3 position;
void main (void)
{
gl_Position = <transform(position)>;
}
"""
fragment1 = """
void main() {
gl_FragColor = vec4(0.0, 1.0, 1.0, 1.0);
}
"""
def draw(mol):
radiusdict={'C':0.11,'O':0.095,'H':0.04,'N':0.1,'S':0.12,'I':0.12,'Cl':0.1,'P':0.1,'Du':0.1}
colordict={'C':[0.6,0.6,0.6],'O':[1.,0.,0.],'H':[1.,1.,1.],'N':[0.,0.,1.],'S':[1.,1.,0.],'I':[1.,0,1.],'Cl':[0.,1.,0.],'P':[0.,0.8,1.],'Du':[1.,0.6,0.]}
mymol = np.zeros(mol.NumAtoms(), dtype=[('position', '<f4', (3,)), ('color', '<f4', (3,)), ('radius', '<f4')])
i=0
for atom in openbabel.OBMolAtomIter(mol):
mymol[i]['position'][0]=atom.GetX ()/10.0
mymol[i]['position'][1]=atom.GetY ()/10.0
mymol[i]['position'][2]=atom.GetZ ()/10.0
if len(atom.GetType())>1:
if atom.GetType()[0:2] in radiusdict:
atomtype=atom.GetType()[0:2]
elif atom.GetType()[0] in radiusdict:
atomtype=atom.GetType()[0]
else:
atomtype='Du'
else:
if atom.GetType()in radiusdict:
atomtype=atom.GetType()
else:
atomtype='Du'
mymol[i]['color']=colordict[atomtype]
mymol[i]['radius']=radiusdict[atomtype]
i+=1
i=0
molbond = np.zeros(mol.NumBonds()*2, dtype=[('position', '<f4', (3,))])
for atom in openbabel.OBMolBondIter(mol):
a=atom.GetBeginAtom ()
b=atom.GetEndAtom ()
molbond[i*2]['position'][0]=a.GetX ()/10.0
molbond[i*2]['position'][1]=a.GetY ()/10.0
molbond[i*2]['position'][2]=a.GetZ ()/10.0
molbond[i*2+1]['position'][0]=b.GetX ()/10.0
molbond[i*2+1]['position'][1]=b.GetY ()/10.0
molbond[i*2+1]['position'][2]=b.GetZ ()/10.0
i+=1
molxyz.bind(mymol.view(gloo.VertexBuffer))
molbonds.bind(molbond.view(gloo.VertexBuffer))
def openfile():
global pybelmol,mol
dialog = QtGui.QFileDialog()
fileName, filtr = dialog.getOpenFileName()
if fileName=='':
return
fileformat=fileName.split('.')[-1]
obConversion = openbabel.OBConversion()
mol = openbabel.OBMol()
obConversion.SetInAndOutFormats(fileformat, "mol2")
obConversion.ReadFile(mol, fileName)
if mol.NumAtoms()==0:
print('No atoms in the file or something wrong in the format.')
return
mol.Center()
pybelmol = pybel.Molecule(mol)
draw(mol)
def buildfile():
global pybelmol,mol
gui = QtGui.QWidget()
text, result = QtGui.QInputDialog.getText(gui, "SMILES","""Formula: e.g. C1=CC=CC=C1 for benzene""")
if result and text:
#text="C1=CC=CS1"
obConversion = openbabel.OBConversion()
mol = openbabel.OBMol()
obConversion.SetInAndOutFormats("smi", "mol2")
obConversion.ReadString(mol, text)
mol.AddHydrogens()
pybelmol = pybel.Molecule(mol)
pybelmol.make3D(steps=1) #localopt()
mol.Center()
draw(mol)
def minimization():
global pybelmol,mol
pybelmol.localopt()
mol.Center()
draw(mol)
def savefile():
fileName, filtr = QtGui.QFileDialog.getSaveFileName(filter="mol2 files (*.mol2);;pdb files (*.pdb);;xyz files (*.xyz)")
if fileName=='':
return
fileformat=fileName.split('.')[-1]
pybelmol.write(fileformat, fileName,overwrite=True)
def about():
msgBox=QtGui.QMessageBox()
msgBox.setText("3D molecue viewer needs Glumpy, numpy, pyopenbabel and PySide \n\n")
msgBox.exec_()
app.use('pyside')
window = app.Window(width=800, height=800, color=(0,0,0,1),title='3D molecule viewer (Right-click to start)')
window.set_position(300,100)
mol = np.zeros(1, dtype=[('position', '<f4', (3,)), ('color', '<f4', (3,)), ('radius', '<f4')])
molbond = np.zeros(1, dtype=[('position', '<f4', (3,))])
#-----------------------------------
molxyz = gloo.Program(vertex, fragment)
molxyz['light_position'] = 0., 0., 2.
molxyz["transform"] = Trackball(Position())
molxyz.bind(mol.view(gloo.VertexBuffer))
molxyz['color'] *= .75
molxyz['color'] += .25
#-----------------------------------
molbonds = gloo.Program(vertex1, fragment1)
molbonds.bind(molbond.view(gloo.VertexBuffer))
molbonds["transform"] = Trackball(Position())
@window.event
def on_draw(dt):
window.clear()
molxyz.draw(gl.GL_POINTS)
molbonds.draw(gl.GL_LINES)
@window.event
def on_init():
gl.glEnable(gl.GL_DEPTH_TEST)
@window.event
def on_mouse_release(x, y, button):
if button==8:
windowx,windowy=window.get_position()
rightclickmenu = QtGui.QMenu()
rightclickmenu.addAction("Open",openfile)
rightclickmenu.addAction("Build SMILES",buildfile)
rightclickmenu.addAction("Structure Optimization",minimization)
rightclickmenu.addAction("Save",savefile)
rightclickmenu.addAction("About",about)
rightclickmenu.exec_(QtCore.QPoint(windowx+x,windowy+y))
window.attach(molxyz["transform"])
window.attach(molbonds["transform"])
app.run()