forked from sympy/sympy
-
Notifications
You must be signed in to change notification settings - Fork 1
/
color_scheme.py
300 lines (266 loc) · 11.6 KB
/
color_scheme.py
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
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
from sympy import Basic, Symbol, symbols, lambdify
from util import interpolate, rinterpolate, create_bounds, update_bounds
class ColorGradient(object):
colors = [0.4,0.4,0.4], [0.9,0.9,0.9]
intervals = 0.0, 1.0
def __init__(self, *args):
if len(args) == 2:
self.colors = list(args)
self.intervals = [0.0, 1.0]
elif len(args) > 0:
assert len(args) % 2 == 0
self.colors = [args[i] for i in xrange(1, len(args), 2)]
self.intervals = [args[i] for i in xrange(0, len(args), 2)]
assert len(self.colors) == len(self.intervals)
def copy(self):
c = ColorGradient()
c.colors = [e[::] for e in self.colors]
c.intervals = self.intervals[::]
return c
def _find_interval(self, v):
m = len(self.intervals)
i = 0
while i < m-1 and self.intervals[i] <= v: i+=1
return i
def _interpolate_axis(self, axis, v):
i = self._find_interval(v)
v = rinterpolate(self.intervals[i-1], self.intervals[i], v)
return interpolate(self.colors[i-1][axis], self.colors[i][axis], v)
def __call__(self, r, g, b):
c = self._interpolate_axis
return c(0, r), c(1, g), c(2, b)
default_color_schemes = {} # defined at the bottom of this file
class ColorScheme(object):
def __init__(self, *args, **kwargs):
self.args = args
self.f, self.gradient = None, ColorGradient()
if len(args) == 1 and not isinstance(args[0],Basic) and callable(args[0]):
self.f = args[0]
elif len(args) == 1 and isinstance(args[0], str):
if args[0] in default_color_schemes:
cs = default_color_schemes[args[0]]
self.f, self.gradient = cs.f, cs.gradient.copy()
else:
self.f = lambdify('x,y,z,u,v', args[0])
else:
self.f, self.gradient = self._interpret_args(args, kwargs)
self._test_color_function()
if not isinstance(self.gradient, ColorGradient):
raise ValueError("Color gradient not properly initialized. "
"(Not a ColorGradient instance.)")
def _interpret_args(self, args, kwargs):
f, gradient = None, self.gradient
atoms, lists = self._sort_args(args)
s = self._pop_symbol_list(lists)
s = self._fill_in_vars(s)
# prepare the error message for lambdification failure
f_str = ', '.join(str(fa) for fa in atoms)
s_str = (str(sa) for sa in s)
s_str = ', '.join(sa for sa in s_str if sa.find('unbound') < 0)
f_error = ValueError("Could not interpret arguments "
"%s as functions of %s." % (f_str, s_str))
# try to lambdify args
if len(atoms) == 1:
fv = atoms[0]
try: f = lambdify(s, [fv,fv,fv])
except: raise f_error
elif len(atoms) == 3:
fr, fg, fb = atoms
try: f = lambdify(s, [fr,fg,fb])
except: raise f_error
else: raise ValueError("A ColorScheme must provide 1 or 3 "
"functions in x, y, z, u, and/or v.")
# try to intrepret any given color information
if len(lists) == 0:
gargs = []
elif len(lists) == 1:
gargs = lists[0]
elif len(lists) == 2:
try: (r1,g1,b1), (r2,g2,b2) = lists
except: raise ValueError("If two color arguments are given, "
"they must be given in the format "
"(r1,g1,b1), (r2,g2,b2).")
gargs = lists
elif len(lists) == 3:
try: (r1,r2), (g1,g2), (b1,b2) = lists
except: raise ValueError("If three color arguments are given, "
"they must be given in the format "
"(r1,r2), (g1,g2), (b1,b2). To create "
"a multi-step gradient, use the syntax "
"[0, colorStart, step1, color1, ..., 1, "
"colorEnd].")
gargs = [[r1,g1,b1], [r2,g2,b2]]
else: raise ValueError("Don't know what to do with collection "
"arguments %s." % (', '.join(str(l) for l in lists)))
if gargs:
try: gradient = ColorGradient(*gargs)
except Exception, ex:
raise ValueError(("Could not initialize a gradient "
"with arguments %s. Inner "
"exception: %s") % (gargs, str(ex)))
return f, gradient
def _pop_symbol_list(self, lists):
symbol_lists = []
for l in lists:
mark = True
for s in l:
if s is not None and not isinstance(s, Symbol):
mark = False
break
if mark:
lists.remove(l)
symbol_lists.append(l)
if len(symbol_lists) == 1:
return symbol_lists[0]
elif len(symbol_lists) == 0:
return []
else:
raise ValueError("Only one list of Symbols "
"can be given for a color scheme.")
def _fill_in_vars(self, args):
defaults = symbols('x,y,z,u,v')
if len(args) == 0: return defaults
if not isinstance(args, (tuple, list)):
raise v_error
if len(args) == 0: return defaults
for s in args:
if s is not None and not isinstance(s, Symbol):
raise v_error
# when vars are given explicitly, any vars
# not given are marked 'unbound' as to not
# be accidentally used in an expression
vars = [Symbol('unbound%i'%(i)) for i in xrange(1,6)]
# interpret as t
if len(args) == 1:
vars[3] = args[0]
# interpret as u,v
elif len(args) == 2:
if args[0] is not None:
vars[3] = args[0]
if args[1] is not None:
vars[4] = args[1]
# interpret as x,y,z
elif len(args) >= 3:
# allow some of x,y,z to be
# left unbound if not given
if args[0] is not None: vars[0] = args[0]
if args[1] is not None: vars[1] = args[1]
if args[2] is not None: vars[2] = args[2]
# interpret the rest as t
if len(args) >= 4:
vars[3] = args[3]
# ...or u,v
if len(args) >= 5:
vars[4] = args[4]
return vars
def _sort_args(self, args):
atoms, lists = [], []
for a in args:
if isinstance(a, (tuple, list)):
lists.append(a)
else: atoms.append(a)
return atoms, lists
def _test_color_function(self):
if not callable(self.f):
raise ValueError("Color function is not callable.")
try:
result = self.f(0,0,0,0,0)
assert len(result) == 3
except TypeError, te:
raise ValueError("Color function needs to accept x,y,z,u,v, "
"as arguments even if it doesn't use all of them.")
except AssertionError, ae:
raise ValueError("Color function needs to return 3-tuple r,g,b.")
except Exception, ie:
pass # color function probably not valid at 0,0,0,0,0
def __call__(self, x,y,z,u,v):
try: return self.f(x,y,z,u,v)
except Exception, e:
#print e
return None
def apply_to_curve(self, verts, u_set, set_len=None, inc_pos=None):
"""
Apply this color scheme to a
set of vertices over a single
independent variable u.
"""
bounds = create_bounds()
cverts = list()
if callable(set_len): set_len(len(u_set)*2)
# calculate f() = r,g,b for each vert
# and find the min and max for r,g,b
for _u in xrange(len(u_set)):
if verts[_u] is None:
cverts.append(None)
else:
x,y,z = verts[_u]
u,v = u_set[_u], None
c = self(x,y,z,u,v)
if c is not None:
c = list(c)
update_bounds(bounds, c)
cverts.append(c)
if callable(inc_pos): inc_pos()
# scale and apply gradient
for _u in xrange(len(u_set)):
if cverts[_u] is not None:
for _c in xrange(3):
# scale from [f_min, f_max] to [0,1]
cverts[_u][_c] = rinterpolate(bounds[_c][0], bounds[_c][1], cverts[_u][_c])
# apply gradient
cverts[_u] = self.gradient(*cverts[_u])
if callable(inc_pos): inc_pos()
return cverts
def apply_to_surface(self, verts, u_set, v_set, set_len=None, inc_pos=None):
"""
Apply this color scheme to a
set of vertices over two
independent variables u and v.
"""
bounds = create_bounds()
cverts = list()
if callable(set_len): set_len(len(u_set)*len(v_set)*2)
# calculate f() = r,g,b for each vert
# and find the min and max for r,g,b
for _u in xrange(len(u_set)):
column = list()
for _v in xrange(len(v_set)):
if verts[_u][_v] is None:
column.append(None)
else:
x,y,z = verts[_u][_v]
u,v = u_set[_u], v_set[_v]
c = self(x,y,z,u,v)
if c is not None:
c = list(c)
update_bounds(bounds, c)
column.append(c)
if callable(inc_pos): inc_pos()
cverts.append(column)
# scale and apply gradient
for _u in xrange(len(u_set)):
for _v in xrange(len(v_set)):
if cverts[_u][_v] is not None:
# scale from [f_min, f_max] to [0,1]
for _c in xrange(3):
cverts[_u][_v][_c] = rinterpolate(bounds[_c][0], bounds[_c][1], cverts[_u][_v][_c])
# apply gradient
cverts[_u][_v] = self.gradient(*cverts[_u][_v])
if callable(inc_pos): inc_pos()
return cverts
def str_base(self): return ", ".join(str(a) for a in self.args)
def __repr__(self): return "%s" % (self.str_base())
x,y,z,t,u,v = symbols('x,y,z,t,u,v')
default_color_schemes['rainbow'] = ColorScheme( z, y, x )
default_color_schemes['zfade'] = ColorScheme( z, (0.4,0.4,0.97), (0.97,0.4,0.4), (None, None, z) )
default_color_schemes['zfade3'] = ColorScheme( z, (None, None, z), [ 0.00, (0.2,0.2,1.0),
0.35, (0.2,0.8,0.4),
0.50, (0.3,0.9,0.3),
0.65, (0.4,0.8,0.2),
1.00, (1.0,0.2,0.2) ] )
default_color_schemes['zfade4'] = ColorScheme( z, (None, None, z), [ 0.0, (0.3, 0.3, 1.0),
0.30, (0.3, 1.0, 0.3),
0.55, (0.95,1.0, 0.2),
0.65, (1.0,0.95, 0.2),
0.85, (1.0, 0.7, 0.2),
1.0, (1.0, 0.3, 0.2) ] )