1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
|
# -*- coding: utf-8 -*-
# -----------------------------------------------------------------------------
# Copyright (c) Vispy Development Team. All Rights Reserved.
# Distributed under the (new) BSD License. See LICENSE.txt for more info.
# -----------------------------------------------------------------------------
"""
Show vector field flow
"""
from __future__ import division
from vispy import app, scene, visuals, gloo
from vispy.util import ptime
import numpy as np
class VectorFieldVisual(visuals.Visual):
vertex = """
uniform sampler2D field;
attribute vec2 index;
uniform vec2 shape;
uniform vec2 field_shape;
uniform float spacing;
varying float dist; // distance along path for this vertex
varying vec2 ij;
uniform sampler2D offset;
uniform float seg_len;
uniform int n_iter; // iterations to integrate along field per vertex
uniform vec2 attractor;
varying vec4 base_color;
uniform sampler2D color;
void main() {
// distance along one line
dist = index.y * seg_len;
vec2 local;
ij = vec2(mod(index.x, shape.x), floor(index.x / shape.x));
// *off* is a random offset to the starting location, which prevents
// the appearance of combs in the field
vec2 off = texture2D(offset, ij / shape).xy - 0.5;
local = spacing * (ij + off);
vec2 uv;
vec2 dir;
vec2 da;
int index_y = int(index.y);
for( int i=0; i<index.y; i+=1 ) {
for ( int j=0; j<n_iter; j += 1 ) {
uv = local / field_shape;
dir = texture2D(field, uv).xy;
// add influence of variable attractor (mouse)
da = attractor - local;
float al = 0.1 * length(da);
da /= 0.5 * (1 + al*al);
dir += da;
// maybe pick a more accurate integration method?
local += seg_len * dir / n_iter;
}
}
base_color = texture2D(color, uv);
gl_Position = $transform(vec4(local, 0, 1));
}
"""
fragment = """
uniform float time;
uniform float speed;
varying float dist;
varying vec2 ij;
uniform sampler2D offset;
uniform vec2 shape;
uniform float nseg;
uniform float seg_len;
varying vec4 base_color;
void main() {
float totlen = nseg * seg_len;
float phase = texture2D(offset, ij / shape).b;
float alpha;
// vary alpha along the length of the line to give the appearance of
// motion
alpha = mod((dist / totlen) + phase - time * speed, 1);
// add a cosine envelope to fade in and out smoothly at the ends
alpha *= (1 - cos(2 * 3.141592 * dist / totlen)) * 0.5;
gl_FragColor = vec4(base_color.rgb, base_color.a * alpha);
}
"""
def __init__(self, field, spacing=10, segments=3, seg_len=0.5,
color=(1, 1, 1, 0.3)):
self._time = 0.0
self._last_time = ptime.time()
rows = int(field.shape[0] / spacing)
cols = int(field.shape[1] / spacing)
index = np.empty((rows * cols, int(segments) * 2, 2), dtype=np.float32)
# encodes starting position within vector field
index[:, :, 0] = np.arange(rows * cols)[:, np.newaxis]
# encodes distance along length of line
index[:, ::2, 1] = np.arange(segments)[np.newaxis, :]
index[:, 1::2, 1] = np.arange(segments)[np.newaxis, :] + 1
self._index = gloo.VertexBuffer(index)
if not isinstance(color, np.ndarray):
color = np.array([[list(color)]], dtype='float32')
self._color = gloo.Texture2D(color)
offset = np.random.uniform(256, size=(rows, cols, 3)).astype(np.ubyte)
self._offset = gloo.Texture2D(offset, format='rgb')
self._field = gloo.Texture2D(field, format='rg',
internalformat='rg32f',
interpolation='linear')
self._field_shape = field.shape[:2]
visuals.Visual.__init__(self, vcode=self.vertex, fcode=self.fragment)
self.timer = app.Timer(interval='auto', connect=self.update_time,
start=False)
self.freeze()
self.shared_program['field'] = self._field
self.shared_program['field_shape'] = self._field.shape[:2]
self.shared_program['shape'] = (rows, cols)
self.shared_program['index'] = self._index
self.shared_program['spacing'] = spacing
self.shared_program['t'] = self._time
self.shared_program['offset'] = self._offset
self.shared_program['speed'] = 1
self.shared_program['color'] = self._color
self.shared_program['seg_len'] = seg_len
self.shared_program['nseg'] = segments
self.shared_program['n_iter'] = 1
self.shared_program['attractor'] = (0, 0)
self.shared_program['time'] = 0
self._draw_mode = 'lines'
self.set_gl_state('translucent', depth_test=False)
self.timer.start()
def _prepare_transforms(self, view):
view.view_program.vert['transform'] = view.get_transform()
def _prepare_draw(self, view):
pass
def _compute_bounds(self, axis, view):
if axis > 1:
return (0, 0)
return (0, self._field_shape[axis])
def update_time(self, ev):
t = ptime.time()
self._time += t - self._last_time
self._last_time = t
self.shared_program['time'] = self._time
self.update()
VectorField = scene.visuals.create_visual_node(VectorFieldVisual)
def fn(y, x):
dx = x-50
dy = y-30
hyp = (dx**2 + dy**2)**0.5 + 0.01
return np.array([100 * dy / hyp**1.7, -100 * dx / hyp**1.8])
field = np.fromfunction(fn, (100, 100)).transpose(1, 2, 0).astype('float32')
field[..., 0] += 10 * np.cos(np.linspace(0, 2 * 3.1415, 100))
color = np.zeros((100, 100, 4), dtype='float32')
color[..., :2] = (field + 5) / 10.
color[..., 2] = 0.5
color[..., 3] = 0.5
canvas = scene.SceneCanvas(keys='interactive', show=True)
view = canvas.central_widget.add_view(camera='panzoom')
vfield = VectorField(field[..., :2], spacing=0.5, segments=30, seg_len=0.05,
parent=view.scene, color=color)
view.camera.set_range()
@canvas.connect
def on_mouse_move(event):
if 3 in event.buttons:
tr = canvas.scene.node_transform(vfield)
vfield.shared_program['attractor'] = tr.map(event.pos)[:2]
if __name__ == '__main__':
app.run()
|
