Ticket #1754: chromadepth.py

File chromadepth.py, 3.2 KB (added by Tristan Croll, 7 years ago)
Line 
1
2class ChromaDepth_Mgr:
3 def __init__(self, session, model):
4 self.session = session
5 self.model = model
6 self._colormap = None
7 self._handler = None
8 self._last_camera_origin = [0,0,0]
9 self._last_view_direction = [0,0,0]
10
11 def start_auto_coloring(self):
12 if self._handler is None:
13 self._handler = self.session.triggers.add_handler('new frame',
14 self._update_colors)
15
16 def stop_auto_coloring(self):
17 if self._handler is not None:
18 self.session.triggers.remove_handler(self._handler)
19 self._handler = None
20
21 @property
22 def colormap(self):
23 if self._colormap is None:
24 from chimerax.core.colors import Colormap
25 self._colormap = Colormap([0,0.5,1],[[1,0,0,1],[0,1,0,1],[0,0,1,1]])
26 return self._colormap
27
28 @property
29 def far_clip(self):
30 v = self.session.view
31 cp = v.clip_planes
32 fc = cp.find_plane('far')
33 if fc is not None:
34 return fc
35 else:
36 from chimerax.core.graphics.view import ClipPlane
37 c = v.camera
38 cofr = v.center_of_rotation
39 cpos = c.position.origin()
40 clip_point = cpos + (1+self._far_clip_multiplier)* (cofr-cpos)
41 fc = ClipPlane('far', c.view_direction(),
42 clip_point, camera_normal = (0,0,1))
43 # Put the near clip at the camera position for depth cueing
44 v.clip_planes.add_plane(fc)
45 return fc
46
47 @property
48 def near_clip(self):
49 v = self.session.view
50 cp = v.clip_planes
51 nc = cp.find_plane('near')
52 if nc is not None:
53 return nc
54 else:
55 from chimerax.core.graphics.view import ClipPlane
56 c = v.camera
57 cofr = v.center_of_rotation
58 cpos = c.position.origin()
59 clip_point = cpos + (1-self._near_clip_multiplier) * (cofr-cpos)
60 nc = ClipPlane('near', c.view_direction(),
61 clip_point, camera_normal = (0,0,-1))
62 # Put the near clip at the camera position for depth cueing
63 v.clip_planes.add_plane(nc)
64 return nc
65
66 def _update_colors(self, *_):
67 import numpy
68 v = self.session.view
69 camera = v.camera
70 o = camera.position.origin()
71 d = camera.view_direction()
72 if numpy.allclose(o,self._last_camera_origin) and numpy.allclose(d, self._last_view_direction):
73 return
74 self._last_camera_origin = o
75 self._last_view_direction = d
76
77
78
79
80 v = self.session.view
81 camera = v.camera
82 o = camera.position.origin()
83 d = camera.view_direction()
84
85 m = self.model
86 b = m.bounds()
87 distances = numpy.dot(b.box_corners()-o, d)
88 nd = min(distances)
89 fd = max(distances)
90
91 atoms = m.atoms
92 residues = m.residues
93
94 coords = atoms.scene_coords
95 distances = numpy.dot(coords-o, d)
96 normalised_distances = (distances-nd)/(fd-nd)
97 atoms.colors = self.colormap.interpolated_rgba8(normalised_distances)
98
99 residues.ribbon_colors = [r.principal_atom.color if r.principal_atom is not None else [0,0,0,0] for r in residues]