Coverage for model_workflow/tools/get_screenshot.py: 87%
205 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-23 10:54 +0000
1from subprocess import run, PIPE
2from os import environ, remove
3from os.path import exists
4from PIL import Image
5import math
6import numpy as np
8from model_workflow.utils.auxiliar import ToolError
9from model_workflow.utils.type_hints import *
11# The Convex Hull is a polygon which covers all the given point and Convex Hull is the smallest polygon.
12# SciPy provides a method "scipy. spatial. ConvexHull()" to create a Convex Hull
13from scipy.spatial import ConvexHull
14from scipy.spatial.distance import cdist
16# Auxiliar files
17AUXILIAR_PDB_FILENAME = '.screenshot_structure.pdb'
18AUXILIAR_TGA_FILENAME = '.transition_screenshot.tga'
20# Use this to draw some shapes as references in the screenshot
21# This is useful for debugging purposes only
22debug = False
25#
26def get_screenshot (
27 structure : 'Structure',
28 output_filepath : str,
29 # You may pass the camera rotation, translation and zoom parameters so they are not calculated again
30 # This is useful to keep screenshots coherent between different clusters/markov states
31 # Note that a slight movement in the molecule may make the rotation logic here use a different angle
32 # Thus the image could be radically different and misleading, since the change could be minimal
33 parameters : Optional[dict] = None,
34) -> dict:
35 """ Obtain a screenshot from the pdb file using VMD. This screenshot of the system is uploaded to the database.
36 Returns the rotation values used to take the photo so they can be saved and reused."""
37 # Check the output screenshot file extension is JPG
38 if output_filepath.split('.')[-1] != 'jpg':
39 raise SystemExit('You must provide a .jpg file name!')
41 # Produce a PDB file to feed VMD
42 structure.generate_pdb_file(AUXILIAR_PDB_FILENAME)
44 # Number of pixels to scale in x
45 x_number_pixels = 350
46 # Number of pixels to scale in y
47 y_number_pixels = 350
49 # Save the current value of the VMDSCRSIZE enviornmental variable
50 VMDSCRSIZE_backup = environ.get('VMDSCRSIZE', '')
52 # Set an enviornment variable to handle the window size
53 # WARNING: Use this system instead of the 'display resize' command
54 # This command silently kills VMD when it is run in sbatch (but not in salloc)
55 # The cause is not clear but this command may depend on OpenGL rendering features which are not supported in sbatch
56 # https://www.ks.uiuc.edu/Training/Tutorials/vmd/tutorial-html/node8.html
57 environ['VMDSCRSIZE'] = str(x_number_pixels) + " " + str(y_number_pixels)
59 # Prepare a Tcl script in order to execute different commands in the terminal of VMD
61 # Generate a file name for the commands file
62 commands_filename_1 = '.commands_1.vmd'
64 # Set a file to save the result obtained when calculating the center
65 center_filename = '.center_point_filename.txt'
66 with open(commands_filename_1, "w") as file:
67 # Select the whole molecule
68 file.write('set sel [atomselect 0 all] \n')
69 # Find the center
70 file.write('set center [measure center $sel] \n')
71 # Export the center to a file
72 file.write(f'set center_file [open {center_filename} w] \n')
73 file.write('puts $center_file $center \n')
74 # Exit VMD
75 file.write('exit \n')
77 # Run VMD
78 process = run([
79 "vmd",
80 AUXILIAR_PDB_FILENAME,
81 "-e",
82 commands_filename_1,
83 "-dispdev",
84 "none"
85 ], stdout=PIPE, stderr=PIPE)
86 logs = process.stdout.decode()
88 # If the center file was not generated then something went wrong with VMD
89 if not exists(center_filename):
90 print(logs)
91 error_logs = process.stderr.decode()
92 print(error_logs)
93 raise ToolError('Something went wrong with VMD while generating the center point file')
95 # Read the generated file to get the center
96 with open(center_filename,"r") as file:
97 line = file.readline().split()
98 line = [float(i) for i in line]
99 vmd_center_coordinates = line
101 # Set the camera rotation, translation and zoom values to get the optimal picture
102 angle = None
103 angle2 = None
104 y_axis_difference_vector = None
105 x_axis_difference_vector = None
106 scale = None
108 # If precalculated values are passed then use them
109 if parameters:
110 angle = parameters['angle']
111 angle2 = parameters['angle2']
112 y_axis_difference_vector = parameters['y_axis_difference_vector']
113 x_axis_difference_vector = parameters['x_axis_difference_vector']
114 scale = parameters['scale']
116 # We must calculate these values otherwise
117 else:
119 # Obtain all coordinates from each atom of the filtered structure
120 coordinates = [list(atom.coords) for atom in structure.atoms]
121 # Convert the list into a Numpy Array, since Scipy library just works with this type of data structure
122 coordinates_np = np.array(coordinates)
124 # Iterate over all coordinates to store in a dictionary as a key a tuple with the coordinates with the index as its value
125 d = {}
126 for index,coordinate in enumerate(coordinates):
127 d[tuple(coordinate)] = index
129 # Compute the hull of the molecule
130 hull = ConvexHull(coordinates_np)
131 # Extract the points forming the hull
132 hullpoints = coordinates_np[hull.vertices,:]
133 # Naive way of finding the best pair in O(H^2) time if H is number of points on hull
134 hull_best_points = cdist(hullpoints, hullpoints, metric='euclidean')
135 # Get the farthest apart points
136 bestpair = np.unravel_index(hull_best_points.argmax(), hull_best_points.shape)
137 # Convert the first point coordinates from Numpy array to list
138 first_point = hullpoints[bestpair[0]].tolist()
139 # Now into tuple
140 first_point = tuple(first_point)
141 # Convert the second point coordinates from Numpy array to list
142 second_point = hullpoints[bestpair[1]].tolist()
143 # Now into tuple
144 second_point = tuple(second_point)
146 ### TRIGONOMETRY TO COMPUTE THE ANGLE WE NEED TO ROTATE
148 # FIRST ROTATION
149 # Looking at the x-z plane, consider we have a rectangle triangle where the segment between the
150 # first and second points is the hypotenuse and the sides of traingle are paralel to x and z axes
151 # Calculate the grades we must rotate the molecule to align the two points in the x axis
153 # Get the hypotenuse
154 # Note that the hypotenuse is calculated as a projection of the segment in the x-z plane
155 # This is lower than the actual distance between the points
156 hypotenuse = calculate_distance(first_point, second_point, ['x', 'z'])
158 # Get the z-side which it is just the difference in the z coordinates
159 z_side = abs(first_point[2] - second_point[2])
161 # Obtain the angle that we are interested in order to rotate the molecule
162 # Apply trigonometry, compute the arcsinus of the division between the opposite side and the hypothenuse
163 angle = math.degrees(math.asin(z_side/hypotenuse)) if hypotenuse > 0 else 0
165 # Set the right rotation direction
166 most_negative_x_point = first_point if first_point[0] <= second_point[0] else second_point
167 most_negative_z_point = first_point if first_point[2] <= second_point[2] else second_point
168 if most_negative_x_point != most_negative_z_point:
169 angle = -angle
171 # SECOND ROTATION
172 # Now looking at the x-y plane, consider we have a rectangle triangle where the segment between the
173 # first and second points is the hypotenuse and the sides of traingle are paralel to x and y axes
174 # Calculate the grades we must rotate the molecule to align the two points in the x axis
176 # Note that once the rotation for this triangle is solved we add an extra 45 grades
177 # This is becase we want the molecule to be aligned with the diagonal of the image, not the x axis
179 # Get the hypotenuse
180 # Note that here we use all dimensions and not the projection in x-y plane
181 # This is because VMD rotates using its current rotation as reference, not the absolute
182 # Thus the segment is already in "our" x-y plane after the previous rotation
183 # Actually removing the z would be wrong since this z is not "our" z
184 # It is hard to understand only with words but trust
185 hypotenuse2 = calculate_distance(first_point, second_point, ['x', 'y', 'z'])
187 # Get the y-side which it is just the difference in the y coordinates
188 y_side = abs(first_point[1] - second_point[1])
190 # Obtain the angle that we are interested in order to rotate the molecule
191 # Apply trigonometry, compute the arcsinus of the division between the opposite side and the hypothenuse
192 angle2 = math.degrees(math.asin(y_side/hypotenuse2)) if hypotenuse2 > 0 else 0
194 # Set the right rotation direction
195 most_negative_x_point = first_point if first_point[0] <= second_point[0] else second_point
196 most_positive_y_point = first_point if first_point[1] >= second_point[1] else second_point
197 if most_negative_x_point != most_positive_y_point:
198 angle2 = -angle2
199 # As we want it diagonal with respect y and x we add 45 degrees
200 angle2 += 45
202 # Vector that we are going to use as x axis
203 absolute_x_axis = (1,0,0)
204 # Vector that we are going to use as y axis
205 absolute_y_axis = (0,1,0)
207 # Initial normal vector representing the direction of the camera view
208 # Note that when we enter VMD the z axis is pointing towards us
209 initial_normal_vector = (0,0,-1)
211 # Rotate normal vector in order to have it in the same direction as the camera after the first rotation
212 # Note that the angle here and below is negative
213 # Rotating in VMD by the y axis is counter intuitive: negative means clockwise
214 rotated_normal_vector = rotate_vector(initial_normal_vector, -angle, absolute_y_axis)
216 # Rotate x axis around y axis
217 first_rotated_x_axis = rotate_vector(absolute_x_axis, -angle, absolute_y_axis)
218 # Note that there is no need to apply the first rotation to the y axis since we rotate around the y axis
220 # Now, for the second rotation we must use the rotated normal vector as pivot
221 # This emulates how VMD rotates around the recently moved camera axis, not the global axis
222 # Obtain rotation matrix to perform the second rotation we have done previously
223 rotated_x_axis = rotate_vector(first_rotated_x_axis, angle2, rotated_normal_vector)
224 # Perform just the same procedure as before
225 rotated_y_axis = rotate_vector(absolute_y_axis, angle2, rotated_normal_vector)
227 ##########################################################
229 # Project all atom coordinates in each rotated axis to find the range of coordinates and thus its center
230 # Formula from https://gamedev.stackexchange.com/questions/72528/how-can-i-project-a-3d-point-onto-a-3d-line
231 # A + dot(AP,AB) / dot(AB,AB) * AB
233 # The point to set the line does not matter. It could be (0,0,0)
234 # However, we use the vmd center so then we can take the projected points as references for debugging
235 A_point = vmd_center_coordinates
237 # As we are going to compute a line in Y axis we just set AB_vector equal to the rotated y axis
238 AB_vector = rotated_y_axis
240 # Function to project a point into a line
241 # WARNING: It has been observed experimentally that this function is not working as expected
242 # WARNING: Points are projected in the line but very shifted
243 # WARNING: However the error is the same for the molecule points and the view center
244 # WARNING: For this reason the resulting difference is correct and it works
245 def get_projected_point (P_point : Tuple[float, float, float]) -> Tuple[float, float, float]:
246 AP_vector = [A_point[i] + P_point[i] for i in range(3)]
247 AP_AB_dot_product = np.dot(AP_vector, AB_vector)
248 AB_AB_dot_product = np.dot(AB_vector, AB_vector)
249 scalar = AP_AB_dot_product / AB_AB_dot_product
250 projected_vector = [AB_vector[i] * scalar for i in range(3)]
251 projected_point = [A_point[i] + projected_vector[i] for i in range(3)]
252 return projected_point
254 # Project all atom coordinates in the rotated y axis
255 projected_points = [ get_projected_point(P_point) for P_point in coordinates ]
256 projected_points.sort(key=lambda projected_points: projected_points[0])
257 projected_points.sort(key=lambda projected_points: projected_points[1])
258 projected_points.sort(key=lambda projected_points: projected_points[2])
259 # Get the most distant projected points
260 max_ypoint = projected_points[0]
261 min_ypoint = projected_points[-1]
262 # Find the center between the most distant points
263 y_axis_center = [(max_ypoint[i] + min_ypoint[i]) / 2 for i in range(3)]
265 # Project the center of the view in the rotated y axis
266 y_axis_projected_center = get_projected_point(vmd_center_coordinates)
267 # Calculate the difference between the molecule center and the view center in the rotated y axis
268 y_axis_difference_vector = [y_axis_projected_center[i] - y_axis_center[i] for i in range(3)]
270 # Repeat the whole process with the rotated x axis
272 # Project all atom coordinates in the rotated x axis
273 AB_vector = rotated_x_axis
274 projected_points2 = [ get_projected_point(P_point) for P_point in coordinates ]
275 projected_points2.sort(key=lambda projected_points2: projected_points2[0])
276 projected_points2.sort(key=lambda projected_points2: projected_points2[1])
277 projected_points2.sort(key=lambda projected_points2: projected_points2[2])
278 # Get the most distant projected points
279 max_xpoint = projected_points2[0]
280 min_xpoint = projected_points2[-1]
281 # Find the center between the most distant points
282 x_axis_center = [(max_xpoint[i] + min_xpoint[i]) / 2 for i in range(3)]
284 # Project the center of the view in the rotated x axis
285 x_axis_projected_center = get_projected_point(vmd_center_coordinates)
286 # Calculate the difference between the molecule center and the view center in the rotated x axis
287 x_axis_difference_vector = [x_axis_projected_center[i] - x_axis_center[i] for i in range(3)]
289 # Calculate height and width and get the widest dimension
290 width = calculate_distance(max_xpoint, min_xpoint, ['x','y','z'])
291 height = calculate_distance(max_ypoint, min_ypoint, ['x','y','z'])
292 widest = max(width, height)
293 # Set how close to the molecule we want the camera to be
294 # This value has been found experimentally and it keeps a small white margin
295 zoom = 2.8
296 # Set the scale
297 scale = zoom / widest
299 # We must find also what is representable through cartoon and what is not
300 # Relying in 'protein or nucleic' is not safe enough
301 # For some structures large regions could remain invisible
302 # e.g. a large peptide noted as a single residue
303 # Note that we also include terminals in the cartoon selection although they are not representable
304 # This is because terminals are better hidden than represented as ligands, this would be missleading
305 cartoon_selection = structure.select_cartoon(include_terminals=True)
306 non_cartoon_selection = structure.invert_selection(cartoon_selection)
307 # Also coarse grain beads have to be considered
308 # We cannot paint them by their elements so we must rely in atom names or chains
309 cg_selection = structure.select_cg()
310 non_cartoon_selection -= cg_selection
311 # Set a file name for the VMD script file
312 commands_filename_2 = '.commands_2.vmd'
314 # Now write the VMD script for the rotation
315 with open(commands_filename_2, "w") as file:
316 # Set the Background of the molecule to white color
317 file.write('color Display Background white \n')
318 # Delete the axes drawing in the VMD window
319 file.write('axes location Off \n')
320 # Eliminate the molecule to perform changes in the representation, color and material
321 file.write('mol delrep 0 top \n')
322 # First add a spcific representation for polymers (protein and nucleic acids)
323 if cartoon_selection:
324 # Change the default representation model to Newcartoon
325 file.write('mol representation Newcartoon \n')
326 # Change the default atom coloring method setting to Chain
327 file.write('mol color Chain \n')
328 # Set the default atom selection to atoms to be represented as cartoon
329 file.write(f'mol selection "{cartoon_selection.to_vmd()}" \n')
330 # Change the current material of the representation of the molecule
331 file.write('mol material Opaque \n')
332 # Using the new changes performed previously add a new representation to the new molecule
333 file.write('mol addrep top \n')
334 # In case we have any non-cartoon selection to represent...
335 if non_cartoon_selection:
336 # Change the default representation model to CPK (ball and stick)
337 file.write('mol representation cpk \n')
338 # Change the default atom coloring method setting to Chain
339 file.write('mol color element \n')
340 # Set the default atom selection to atoms to be represented as CPK
341 file.write(f'mol selection "{non_cartoon_selection.to_vmd()}" \n')
342 # Change the current material of the representation of the molecule
343 file.write('mol material Opaque \n')
344 # Using the new changes performed previously add a new representation to the new molecule
345 file.write('mol addrep top \n')
346 # In case we have coarse grain beads
347 if cg_selection:
348 # Change the default representation model to CPK (ball and stick)
349 file.write('mol representation cpk \n')
350 # Change the default atom coloring method setting to Chain
351 file.write('mol color chain \n')
352 # Set the default atom selection to atoms to be represented as CPK
353 file.write(f'mol selection "{cg_selection.to_vmd()}" \n')
354 # Change the current material of the representation of the molecule
355 file.write('mol material Opaque \n')
356 # Using the new changes performed previously add a new representation to the new molecule
357 file.write('mol addrep top \n')
358 # Change projection from perspective (used by VMD by default) to orthographic
359 file.write('display projection orthographic \n')
360 # Select all atoms
361 file.write('set sel [atomselect 0 all] \n')
362 # First rotation of the molecule to set it perpendicular with respect to z axis
363 file.write(f'rotate y by {angle} \n')
364 # Second rotatoin of the molecule to set it diagonal with respect to z axis
365 file.write(f'rotate z by {angle2} \n')
366 # Move to rectify the difference
367 file.write('$sel moveby { ' + tuple_to_vmd(y_axis_difference_vector) + ' } \n')
368 # Move to rectify the difference
369 file.write('$sel moveby { ' + tuple_to_vmd(x_axis_difference_vector) + ' } \n')
370 # Set the scale
371 file.write(f'scale to {scale} \n')
373 # Show the theoretical view center if we are debugging
374 if debug:
375 # Note that all draw commands work with coordinates and thus they are absolute, not relative to camera
376 # Draw the center
377 file.write('draw sphere {' + tuple_to_vmd(vmd_center_coordinates) + '} radius 5\n')
378 # Set a function to display vectors
379 def display_vector (vector : tuple, color : str):
380 ref_right_vector = multiply_vector(vector, 100)
381 ref_right_point = point_adds_vector(vmd_center_coordinates, ref_right_vector)
382 ref_left_vector = multiply_vector(vector, -100)
383 ref_left_point = point_adds_vector(vmd_center_coordinates, ref_left_vector)
384 file.write('draw color ' + color + '\n')
385 file.write('draw line {' + tuple_to_vmd(ref_right_point) + '} {' + tuple_to_vmd(ref_left_point) + '}\n')
386 # Draw the rotated axes
387 #display_vector(first_rotated_x_axis, 'orange')
388 display_vector(rotated_x_axis, 'red')
389 display_vector(rotated_y_axis, 'green')
390 # Draw projected point ranges
391 # file.write('draw color red\n')
392 # file.write('draw line {' + tuple_to_vmd(min_xpoint) + '} {' + tuple_to_vmd(max_xpoint) + '}\n')
393 # file.write('draw color green\n')
394 # file.write('draw line {' + tuple_to_vmd(min_ypoint) + '} {' + tuple_to_vmd(max_ypoint) + '}\n')
396 # Finally generate the image from the current view
397 file.write(f'render TachyonInternal {AUXILIAR_TGA_FILENAME} \n')
398 # Exit VMD
399 file.write('exit\n')
401 # Run VMD
402 process = run([
403 "vmd",
404 AUXILIAR_PDB_FILENAME,
405 "-e",
406 commands_filename_2,
407 "-dispdev",
408 "none"
409 ], stdout=PIPE, stderr=PIPE)
410 logs = process.stdout.decode()
411 # If the output file does not exist at this point then it means something went wrong with VMD
412 if not exists(AUXILIAR_TGA_FILENAME):
413 print(logs)
414 error_logs = process.stderr.decode()
415 print(error_logs)
416 raise SystemExit('Something went wrong with VMD while taking the screenshot')
418 im = Image.open(AUXILIAR_TGA_FILENAME)
419 # converting to jpg
420 rgb_im = im.convert("RGB")
421 # exporting the image
422 rgb_im.save(output_filepath)
424 # Remove trash files
425 trash_files = [ AUXILIAR_PDB_FILENAME, commands_filename_1, commands_filename_2, AUXILIAR_TGA_FILENAME, center_filename ]
426 for trash_file in trash_files:
427 remove(trash_file)
429 # Restore the environment variable to not cause problem in possible future uses of VMD
430 environ['VMDSCRSIZE'] = VMDSCRSIZE_backup
432 # Return the camera rotation, translation and zoom values we just used to get the picture
433 return {
434 'angle': angle,
435 'angle2': angle2,
436 'y_axis_difference_vector': y_axis_difference_vector,
437 'x_axis_difference_vector': x_axis_difference_vector,
438 'scale': scale
439 }
441# Set a function to write a tuple of numbers (point coorinates) as VMD expect these value: a string separated by space
442def tuple_to_vmd (point : tuple) -> str:
443 return str(point[0]) + ' ' + str(point[1]) + ' ' + str(point[2])
445# Return the rotated vector
446# https://stackoverflow.com/questions/6802577/rotation-of-3d-vector
447def rotate_vector (vector : tuple, angle : float, axis : tuple) -> tuple:
448 radians = math.radians(angle)
449 axis = np.asarray(axis)
450 axis = axis / math.sqrt(np.dot(axis, axis))
451 a = math.cos(radians / 2.0)
452 b, c, d = -axis * math.sin(radians / 2.0)
453 aa, bb, cc, dd = a * a, b * b, c * c, d * d
454 bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
455 rotation_matrix = np.array([[aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
456 [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
457 [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]])
458 return np.dot(rotation_matrix, vector)
461# Calculate the distance between two atoms
462def calculate_distance (first_coordinates : tuple, second_coordinates : tuple, dimensions : List[str]) -> float:
463 squared_distances_sum = 0
464 for dimension, i in { 'x': 0, 'y': 1, 'z': 2 }.items():
465 if dimension not in dimensions:
466 continue
467 squared_distances_sum += (first_coordinates[i] - second_coordinates[i])**2
468 return math.sqrt(squared_distances_sum)
470# Set a function to gen a new point from a point and vector
471def point_adds_vector (point : tuple, vector : tuple) -> tuple:
472 return tuple([ point[d] + vector[d] for d in range(3) ])
474def multiply_vector (vector : tuple, multiplier : float) -> tuple:
475 return tuple([vector[d] * multiplier for d in range(3)])