pyfock.Graphics

  1import numpy as np
  2from OpenGL.GLUT import *
  3from OpenGL.GLU import *
  4from OpenGL.GL import *
  5from . import Data
  6import sys
  7
  8# Global variables
  9xCoord, yCoord, zCoord, atomicNumber, totalNumAtoms = [], [], [], [], 0
 10bondStart, bondEnd, bondLengths = [], [], []
 11maxDistance, selectedAtom = 0.0, -1
 12_mouse_dragging, _previous_mouse_position = False, None
 13zoom_factor = 1.0  # New zoom variable
 14min_zoom, max_zoom = 0.1, 5.0  # Zoom limits
 15
 16def initializeGlobalVars(mol):
 17    global xCoord, yCoord, zCoord, atomicNumber, totalNumAtoms, maxDistance
 18    totalNumAtoms = mol.natoms
 19    xCoord = [mol.coords[i][0] for i in range(mol.natoms)]
 20    yCoord = [mol.coords[i][1] for i in range(mol.natoms)]
 21    zCoord = [mol.coords[i][2] for i in range(mol.natoms)]
 22    atomicNumber = [mol.Zcharges[i] for i in range(mol.natoms)]
 23    
 24    # Center molecule and calculate max distance
 25    com = [sum(xCoord)/totalNumAtoms, sum(yCoord)/totalNumAtoms, sum(zCoord)/totalNumAtoms]
 26    for i in range(totalNumAtoms):
 27        xCoord[i] -= com[0]; yCoord[i] -= com[1]; zCoord[i] -= com[2]
 28    
 29    maxDistance = max(np.sqrt(x**2 + y**2 + z**2) for x,y,z in zip(xCoord, yCoord, zCoord)) * 2 + 5
 30    calculateBonds()
 31
 32def calculateBonds():
 33    global bondStart, bondEnd, bondLengths
 34    bondStart, bondEnd, bondLengths = [], [], []
 35    for i in range(totalNumAtoms):
 36        for j in range(i+1, totalNumAtoms):
 37            dist = np.sqrt((xCoord[i]-xCoord[j])**2 + (yCoord[i]-yCoord[j])**2 + (zCoord[i]-zCoord[j])**2)
 38            if dist <= (float(Data.covalentRadius[atomicNumber[i]]) + float(Data.covalentRadius[atomicNumber[j]])):
 39                bondStart.append(i); bondEnd.append(j); bondLengths.append(dist)
 40
 41def update_projection():
 42    """Update the projection matrix based on current zoom factor"""
 43    global zoom_factor, maxDistance
 44    glMatrixMode(GL_PROJECTION)
 45    glLoadIdentity()
 46    
 47    # Apply zoom by adjusting the orthographic viewing volume
 48    view_size = (maxDistance/2) / zoom_factor
 49    glOrtho(-view_size, view_size, -view_size, view_size, -100, 100)
 50    
 51    glMatrixMode(GL_MODELVIEW)
 52
 53def zoom_in():
 54    """Zoom in function"""
 55    global zoom_factor
 56    zoom_factor = min(zoom_factor * 1.2, max_zoom)
 57    update_projection()
 58    glutPostRedisplay()
 59
 60def zoom_out():
 61    """Zoom out function"""
 62    global zoom_factor
 63    zoom_factor = max(zoom_factor / 1.2, min_zoom)
 64    update_projection()
 65    glutPostRedisplay()
 66
 67def mouse_wheel(wheel, direction, x, y):
 68    """Handle mouse wheel for zooming"""
 69    if direction > 0:
 70        zoom_in()
 71    else:
 72        zoom_out()
 73
 74def draw_sphere(xyz, radius, color):
 75    glPushMatrix()
 76    color = [c/255.0 for c in color] + [1.0]
 77    glMaterialfv(GL_FRONT, GL_AMBIENT_AND_DIFFUSE, color)
 78    glMaterialfv(GL_FRONT, GL_SPECULAR, [1.0, 1.0, 1.0, 1.0])
 79    glMaterialf(GL_FRONT, GL_SHININESS, 30.0)
 80    glTranslate(*xyz)
 81    glutSolidSphere(radius, 20, 20)
 82    glPopMatrix()
 83
 84def cylinder_between(x1, y1, z1, x2, y2, z2, length, radius):
 85    v = [x2-x1, y2-y1, z2-z1]
 86    axis = (1, 0, 0) if np.hypot(v[0], v[1]) < 0.001 else np.cross(v, (0, 0, 1))
 87    angle = -np.arctan2(np.hypot(v[0], v[1]), v[2]) * 180/np.pi
 88    glPushMatrix()
 89    glTranslate(x1, y1, z1)
 90    glRotate(angle, *axis)
 91    gluCylinder(gluNewQuadric(), radius, radius, length, 10, 10)
 92    glPopMatrix()
 93
 94def display():
 95    global selectedAtom
 96    glClear(GL_COLOR_BUFFER_BIT | GL_DEPTH_BUFFER_BIT)
 97    
 98    # Draw atoms
 99    for i in range(totalNumAtoms):
100        draw_sphere([xCoord[i], yCoord[i], zCoord[i]], 
101                   float(Data.atomicRadius[atomicNumber[i]])/2, 
102                   Data.CPKcolorRGB[atomicNumber[i]])
103    
104    # Draw bonds
105    for i in range(len(bondStart)):
106        start, end = bondStart[i], bondEnd[i]
107        x1, y1, z1 = xCoord[start], yCoord[start], zCoord[start]
108        x2, y2, z2 = xCoord[end], yCoord[end], zCoord[end]
109        
110        # First half with start atom color
111        color = [c/255.0 for c in Data.CPKcolorRGB[atomicNumber[start]]] + [1.0]
112        glMaterialfv(GL_FRONT, GL_AMBIENT_AND_DIFFUSE, color)
113        cylinder_between(x1, y1, z1, (x1+x2)/2, (y1+y2)/2, (z1+z2)/2, bondLengths[i]/2, 0.1)
114        
115        # Second half with end atom color
116        color = [c/255.0 for c in Data.CPKcolorRGB[atomicNumber[end]]] + [1.0]
117        glMaterialfv(GL_FRONT, GL_AMBIENT_AND_DIFFUSE, color)
118        cylinder_between((x1+x2)/2, (y1+y2)/2, (z1+z2)/2, x2, y2, z2, bondLengths[i]/2, 0.1)
119    
120    # Display selected atom info and zoom level
121    if selectedAtom >= 0:
122        glDisable(GL_LIGHTING)
123        glColor3f(1, 1, 1)
124        glRasterPos3f(maxDistance/3/zoom_factor, maxDistance/3/zoom_factor, 0)
125        info = f"Atom {selectedAtom}: {Data.elementSymbols[atomicNumber[selectedAtom]]} ({xCoord[selectedAtom]:.3f}, {yCoord[selectedAtom]:.3f}, {zCoord[selectedAtom]:.3f})"
126        for char in info:
127            glutBitmapCharacter(GLUT_BITMAP_HELVETICA_12, ord(char))
128        glEnable(GL_LIGHTING)
129    
130    # Display zoom level
131    glDisable(GL_LIGHTING)
132    glColor3f(0.8, 0.8, 0.8)
133    glRasterPos3f(-maxDistance/2.5/zoom_factor, maxDistance/2.5/zoom_factor, 0)
134    zoom_info = f"Zoom: {zoom_factor:.2f}x"
135    for char in zoom_info:
136        glutBitmapCharacter(GLUT_BITMAP_HELVETICA_10, ord(char))
137    glEnable(GL_LIGHTING)
138    
139    glutSwapBuffers()
140
141def mouse_click(button, state, x, y):
142    global _mouse_dragging, _previous_mouse_position, selectedAtom
143    if button == GLUT_LEFT_BUTTON:
144        if state == GLUT_DOWN:
145            _mouse_dragging = True
146            _previous_mouse_position = (x, y)
147            
148            # Check for atom selection
149            viewport = glGetIntegerv(GL_VIEWPORT)
150            modelview = glGetDoublev(GL_MODELVIEW_MATRIX)
151            projection = glGetDoublev(GL_PROJECTION_MATRIX)
152            
153            winY = viewport[3] - y
154            _, _, winZ = gluUnProject(x, winY, 0.5, modelview, projection, viewport)
155            
156            min_dist, closest = float('inf'), -1
157            for i in range(totalNumAtoms):
158                winX, winY, _ = gluProject(xCoord[i], yCoord[i], zCoord[i], modelview, projection, viewport)
159                dist = np.sqrt((x - winX)**2 + (y - (viewport[3] - winY))**2)
160                if dist < 20 and dist < min_dist:  # 20 pixel tolerance
161                    min_dist, closest = dist, i
162            
163            selectedAtom = closest
164            print(f"Selected atom {selectedAtom}: {Data.elementSymbols[atomicNumber[selectedAtom]] if selectedAtom >= 0 else 'None'}")
165            glutPostRedisplay()
166        else:
167            _mouse_dragging = False
168
169def mouse_motion(x, y):
170    global _previous_mouse_position
171    if _mouse_dragging and _previous_mouse_position:
172        dx, dy = (x - _previous_mouse_position[0])/100.0, (y - _previous_mouse_position[1])/100.0
173        glRotatef(dx * 50, 0, 1, 0)
174        glRotatef(dy * 50, 1, 0, 0)
175        _previous_mouse_position = (x, y)
176        glutPostRedisplay()
177
178def keyboard(key, x, y):
179    global zoom_factor
180    if key == b'\x1b':  # Escape key
181        glutLeaveMainLoop()
182    elif key == b'+' or key == b'=':  # Plus key for zoom in
183        zoom_in()
184        print(f"Zoom: {zoom_factor:.2f}x")
185    elif key == b'-':  # Minus key for zoom out
186        zoom_out()
187        print(f"Zoom: {zoom_factor:.2f}x")
188    elif key == b'r' or key == b'R':  # Reset zoom
189        zoom_factor = 1.0
190        update_projection()
191        glutPostRedisplay()
192        print("Zoom reset to 1.0x")
193
194def special_keys(key, x, y):
195    """Handle special keys like Page Up/Down for zooming"""
196    if key == GLUT_KEY_PAGE_UP:
197        zoom_in()
198        print(f"Zoom: {zoom_factor:.2f}x")
199    elif key == GLUT_KEY_PAGE_DOWN:
200        zoom_out()
201        print(f"Zoom: {zoom_factor:.2f}x")
202
203def visualize(mol, title='PyFock | CrysX - 3D Viewer'):
204    initializeGlobalVars(mol)
205    
206    glutInit(sys.argv)
207    glutInitDisplayMode(GLUT_DOUBLE | GLUT_RGB | GLUT_DEPTH)
208    glutInitWindowSize(700, 700)
209    glutCreateWindow(title)
210    
211    glClearColor(0.1, 0.1, 0.1, 1.0)
212    glEnable(GL_DEPTH_TEST)
213    glEnable(GL_LIGHTING)
214    glEnable(GL_LIGHT0)
215    glLightfv(GL_LIGHT0, GL_POSITION, [10, 10, 10, 1])
216    glLightfv(GL_LIGHT0, GL_DIFFUSE, [1, 1, 1, 1])
217    
218    # Set initial projection with zoom support
219    update_projection()
220    glMatrixMode(GL_MODELVIEW)
221    glLoadIdentity()
222    
223    glutDisplayFunc(display)
224    glutMouseFunc(mouse_click)
225    glutMotionFunc(mouse_motion)
226    glutKeyboardFunc(keyboard)
227    glutSpecialFunc(special_keys)
228    
229    # Register mouse wheel callback if available
230    try:
231        glutMouseWheelFunc(mouse_wheel)
232    except:
233        pass  # Mouse wheel not supported in older PyOpenGL versions
234    
235    print("Controls:")
236    print("- Click atoms to see coordinates")
237    print("- Drag to rotate")
238    print("- Mouse wheel: Zoom in/out")
239    print("- '+' or '=': Zoom in")
240    print("- '-': Zoom out") 
241    print("- 'R': Reset zoom")
242    print("- Page Up/Down: Zoom in/out")
243    print("- ESC: Exit")
244    print(f"Current zoom: {zoom_factor:.2f}x")
245    
246    glutMainLoop()