Commit 2efae2161771ab70db9e355726936b3256f7fb96
-
- Michał Januszewski <spock @gen…oo.org>
- Diff rendering mode:
- inline
- side by side
sailfish/vis_surf.py
(26 / 7)
| 32 | 32 | glDepthFunc(GL_LEQUAL) | |
| 33 | 33 | glHint(GL_PERSPECTIVE_CORRECTION_HINT, GL_NICEST) | |
| 34 | 34 | ||
| 35 | def _vis_rgb(v): | ||
| 36 | r = numpy.sqrt(v) | ||
| 37 | g = numpy.power(v, 3) | ||
| 38 | b = numpy.sin(v * math.pi) | ||
| 39 | return numpy.vstack((r,g,b)) | ||
| 40 | |||
| 41 | def _vis_scl(v): | ||
| 42 | col = numpy.zeros((3, len(v)), dtype=numpy.float32) | ||
| 43 | col[1,:] = v | ||
| 44 | return col | ||
| 45 | |||
| 46 | colormaps = { | ||
| 47 | 'rgb': _vis_rgb, | ||
| 48 | 'scl': _vis_scl, | ||
| 49 | } | ||
| 50 | |||
| 35 | 51 | class FluidSurfaceVis(object): | |
| 36 | 52 | display_flags = pygame.OPENGL | pygame.DOUBLEBUF | pygame.RESIZABLE | |
| 37 | 53 | display_depth = 24 | |
| … | … | ||
| 86 | 86 | self._paused = False | |
| 87 | 87 | self._reset() | |
| 88 | 88 | ||
| 89 | self._minh = -0.2 | ||
| 90 | self._maxh = 0.2 | ||
| 89 | self._colormap = colormaps.keys()[0] | ||
| 90 | self._minh = -0.1 | ||
| 91 | self._maxh = 0.1 | ||
| 91 | 92 | ||
| 92 | 93 | pygame.key.set_repeat(100,50) | |
| 93 | 94 | ||
| … | … | ||
| 141 | 141 | res = numpy.dstack((a, b, c, d)) | |
| 142 | 142 | vtx = numpy.ravel(res) | |
| 143 | 143 | ||
| 144 | col = numpy.zeros((3, self.mesh_n), dtype=numpy.float32) | ||
| 145 | 144 | min_ = numpy.min(mesh_z) | |
| 146 | 145 | max_ = numpy.max(mesh_z) | |
| 147 | 146 | ||
| 148 | 147 | self._minh = min(self._minh, numpy.min(mesh_z)) | |
| 149 | 148 | self._maxh = max(self._maxh, numpy.max(mesh_z)) | |
| 150 | 149 | ||
| 151 | tmp = vtx[2::3] | ||
| 152 | col[1,:] = (tmp[:]-self._minh)/(self._maxh-self._minh) | ||
| 150 | col = colormaps[self._colormap]((vtx[2::3]-self._minh)/(self._maxh-self._minh)) | ||
| 153 | 151 | col = numpy.ravel(numpy.transpose(col)) | |
| 154 | 152 | ||
| 155 | 153 | vtx.shape = (self.mesh_n, 3) | |
| … | … | ||
| 189 | 189 | self.display_flags, self.display_depth) | |
| 190 | 190 | _GL_resize(*event.size) | |
| 191 | 191 | elif event.type == pygame.MOUSEBUTTONDOWN: | |
| 192 | if event.button == 4: | ||
| 192 | if event.button == 5: | ||
| 193 | 193 | self._zoom -= 0.05 | |
| 194 | elif event.button == 5: | ||
| 194 | elif event.button == 4: | ||
| 195 | 195 | self._zoom += 0.05 | |
| 196 | 196 | elif event.type == pygame.KEYDOWN: | |
| 197 | 197 | if event.key == pygame.K_v: | |
| … | … | ||
| 249 | 249 | self._angle_z += 3.0 | |
| 250 | 250 | else: | |
| 251 | 251 | self._angle_z += 1.0 | |
| 252 | elif event.key == pygame.K_c: | ||
| 253 | idx = colormaps.keys().index(self._colormap) + 1 | ||
| 254 | idx %= len(colormaps.keys()) | ||
| 255 | self._colormap = colormaps.keys()[idx] | ||
| 252 | 256 | ||
| 253 | 257 | def main(self): | |
| 254 | 258 | t_prev = time.time() |
