Commit c06df0c8caddbd1da1c41f35b8d9e88c3710c547
-
- Michał Januszewski <spock @gen…oo.org>
- Diff rendering mode:
- inline
- side by side
sailfish/vis2d.py
(22 / 19)
| 110 | 110 | ||
| 111 | 111 | def _reset(self): | |
| 112 | 112 | self._maxv = 0.000001 | |
| 113 | self._vscale = 0.005 | ||
| 113 | self._vscale = 0.005 | ||
| 114 | 114 | ||
| 115 | 115 | @property | |
| 116 | 116 | def velocity_norm(self): | |
| … | … | ||
| 148 | 148 | ret.append('max_v: %.3f' % maxv) | |
| 149 | 149 | ret.append('rho_avg: %.3f' % numpy.average(self.density)) | |
| 150 | 150 | ||
| 151 | b = (self.sim.geo._decode_node_type(self.geo_map) == geo.LBMGeo.NODE_WALL) | ||
| 151 | wall_map = (self.sim.geo._decode_node_type(self.geo_map) == geo.LBMGeo.NODE_WALL) | ||
| 152 | 152 | ||
| 153 | 153 | if self._vismode == 0: | |
| 154 | 154 | drw = self.velocity_norm / self._maxv | |
| … | … | ||
| 157 | 157 | elif self._vismode == 2: | |
| 158 | 158 | drw = self.vy / self._maxv | |
| 159 | 159 | elif self._vismode == 3: | |
| 160 | mrho = numpy.ma.array(self.density, mask=(b)) | ||
| 160 | mrho = numpy.ma.array(self.density, mask=(wall_map)) | ||
| 161 | 161 | rho_min = numpy.min(mrho) | |
| 162 | 162 | rho_max = numpy.max(mrho) | |
| 163 | 163 | drw = ((self.density - rho_min) / (rho_max - rho_min)) | |
| … | … | ||
| 169 | 169 | g = gauss_kernel(2, sizey=2) | |
| 170 | 170 | drw = signal.convolve(drw,g, mode='same') | |
| 171 | 171 | ||
| 172 | # Rotate the field to the correct position. | ||
| 173 | drw = numpy.rot90(drw.astype(numpy.float32), 3) | ||
| 174 | a = pygame.surfarray.pixels3d(srf) | ||
| 175 | b = numpy.rot90(b, 3) | ||
| 176 | # Draw the walls. | ||
| 177 | a[b] = (255, 255, 255) | ||
| 178 | |||
| 179 | # Draw the data field for all sites which are not marked as a wall. | ||
| 180 | b = numpy.logical_not(b) | ||
| 181 | drw = vis_map[vismode](drw, width, height) | ||
| 182 | a[b] = drw[b] | ||
| 183 | |||
| 184 | # Unlock the surface and put the picture on screen. | ||
| 185 | del a | ||
| 186 | pygame.transform.scale(srf, self._screen.get_size(), self._screen) | ||
| 187 | |||
| 172 | srf2 = self._draw_field(drw, srf, wall_map, vismode, width, height) | ||
| 173 | pygame.transform.scale(srf2, self._screen.get_size(), self._screen) | ||
| 188 | 174 | sw, sh = self._screen.get_size() | |
| 189 | 175 | ||
| 190 | 176 | # Draw the velocity field | |
| … | … | ||
| 190 | 190 | ||
| 191 | 191 | self._draw_tracers(tx, ty, sw, sh, width, height) | |
| 192 | 192 | return ret | |
| 193 | |||
| 194 | def _draw_field(self, field, srf, wall_map, vismode, width, height): | ||
| 195 | # Rotate the field to the correct position. | ||
| 196 | field = numpy.rot90(field.astype(numpy.float32), 3) | ||
| 197 | a = pygame.surfarray.pixels3d(srf) | ||
| 198 | wall_map = numpy.rot90(wall_map, 3) | ||
| 199 | # Draw the walls. | ||
| 200 | a[wall_map] = (255, 255, 255) | ||
| 201 | |||
| 202 | # Draw the data field for all sites which are not marked as a wall. | ||
| 203 | wall_map = numpy.logical_not(wall_map) | ||
| 204 | field = vis_map[vismode](field, width, height) | ||
| 205 | a[wall_map] = field[wall_map] | ||
| 206 | |||
| 207 | # Unlock the surface and put the picture on screen. | ||
| 208 | del a | ||
| 209 | return srf | ||
| 193 | 210 | ||
| 194 | 211 | def _draw_tracers(self, tx, ty, sw, sh, width, height): | |
| 195 | 212 | # Draw the tracer particles |
