glumpy icon indicating copy to clipboard operation
glumpy copied to clipboard

Qt integration

Open rougier opened this issue 6 years ago • 0 comments

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()

rougier avatar Feb 24 '20 06:02 rougier