Permalink
58d497d May 4, 2017
@larsoner @almarklein @xoolive @nippoo
125 lines (96 sloc) 3.23 KB
# -*- coding: utf-8 -*-
# vispy: gallery 30
# vispy: testskip - because this example sometimes sets inactive attributes
"""Computing a Voronoi diagram on the GPU. Shows how to use uniform arrays.
Original version by Xavier Olive (xoolive).
"""
import numpy as np
from vispy import app
from vispy import gloo
# Voronoi shaders.
VS_voronoi = """
attribute vec2 a_position;
void main() {
gl_Position = vec4(a_position, 0., 1.);
}
"""
FS_voronoi = """
uniform vec2 u_seeds[32];
uniform vec3 u_colors[32];
uniform vec2 u_screen;
void main() {
float dist = distance(u_screen * u_seeds[0], gl_FragCoord.xy);
vec3 color = u_colors[0];
for (int i = 1; i < 32; i++) {
float current = distance(u_screen * u_seeds[i], gl_FragCoord.xy);
if (current < dist) {
color = u_colors[i];
dist = current;
}
}
gl_FragColor = vec4(color, 1.0);
}
"""
# Seed point shaders.
VS_seeds = """
attribute vec2 a_position;
uniform float u_ps;
void main() {
gl_Position = vec4(2. * a_position - 1., 0., 1.);
gl_PointSize = 10. * u_ps;
}
"""
FS_seeds = """
varying vec3 v_color;
void main() {
gl_FragColor = vec4(1., 1., 1., 1.);
}
"""
class Canvas(app.Canvas):
def __init__(self):
app.Canvas.__init__(self, size=(600, 600), title='Voronoi diagram',
keys='interactive')
self.ps = self.pixel_scale
self.seeds = np.random.uniform(0, 1.0 * self.ps,
size=(32, 2)).astype(np.float32)
self.colors = np.random.uniform(0.3, 0.8,
size=(32, 3)).astype(np.float32)
self.idx = 0
# Set Voronoi program.
self.program_v = gloo.Program(VS_voronoi, FS_voronoi)
self.program_v['a_position'] = [(-1, -1), (-1, +1), (+1, -1), (+1, +1)]
# HACK: work-around a bug related to uniform arrays until
# issue #345 is solved.
for i in range(32):
self.program_v['u_seeds[%d]' % i] = self.seeds[i, :]
self.program_v['u_colors[%d]' % i] = self.colors[i, :]
# Set seed points program.
self.program_s = gloo.Program(VS_seeds, FS_seeds)
self.program_s['a_position'] = self.seeds
self.program_s['u_ps'] = self.ps
self.activate_zoom()
self.show()
def on_draw(self, event):
gloo.clear()
self.program_v.draw('triangle_strip')
self.program_s.draw('points')
def on_resize(self, event):
self.activate_zoom()
def activate_zoom(self):
self.width, self.height = self.size
gloo.set_viewport(0, 0, *self.physical_size)
self.program_v['u_screen'] = self.physical_size
def on_mouse_move(self, event):
x, y = event.pos
x, y = x / float(self.width), 1 - y / float(self.height)
self.program_v['u_seeds[%d]' % self.idx] = x * self.ps, y * self.ps
# TODO: just update the first line in the VBO instead of uploading the
# whole array of seed points.
self.seeds[self.idx, :] = x, y
self.program_s['a_position'].set_data(self.seeds)
self.update()
def on_mouse_press(self, event):
self.idx = (self.idx + 1) % 32
if __name__ == "__main__":
c = Canvas()
app.run()