1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import absolute_import, print_function, division
8import os
9import re
11import numpy as num
13from pyrocko.guts import Object, Float, Tuple, List
14from pyrocko import util
17class CPTLevel(Object):
18 vmin = Float.T()
19 vmax = Float.T()
20 color_min = Tuple.T(3, Float.T())
21 color_max = Tuple.T(3, Float.T())
24class CPT(Object):
25 color_below = Tuple.T(3, Float.T(), optional=True)
26 color_above = Tuple.T(3, Float.T(), optional=True)
27 color_nan = Tuple.T(3, Float.T(), optional=True)
28 levels = List.T(CPTLevel.T())
30 @property
31 def vmin(self):
32 if self.levels:
33 return self.levels[0].vmin
34 else:
35 return None
37 @property
38 def vmax(self):
39 if self.levels:
40 return self.levels[-1].vmax
41 else:
42 return None
44 def scale(self, vmin, vmax):
45 vmin_old, vmax_old = self.levels[0].vmin, self.levels[-1].vmax
46 for level in self.levels:
47 level.vmin = (level.vmin - vmin_old) / (vmax_old - vmin_old) * \
48 (vmax - vmin) + vmin
49 level.vmax = (level.vmax - vmin_old) / (vmax_old - vmin_old) * \
50 (vmax - vmin) + vmin
52 def get_lut(self):
53 vs = []
54 colors = []
55 if self.color_below and self.levels:
56 vs.append(self.levels[0].vmin)
57 colors.append(self.color_below)
59 for level in self.levels:
60 vs.append(level.vmin)
61 vs.append(level.vmax)
62 colors.append(level.color_min)
63 colors.append(level.color_max)
65 if self.color_above and self.levels:
66 vs.append(self.levels[-1].vmax)
67 colors.append(self.color_above)
69 vs_lut = num.array(vs)
70 colors_lut = num.array(colors)
71 return vs_lut, colors_lut
73 def __call__(self, values):
74 vs_lut, colors_lut = self.get_lut()
76 colors = num.zeros((values.size, 3))
77 colors[:, 0] = num.interp(values, vs_lut, colors_lut[:, 0])
78 colors[:, 1] = num.interp(values, vs_lut, colors_lut[:, 1])
79 colors[:, 2] = num.interp(values, vs_lut, colors_lut[:, 2])
81 if self.color_nan:
82 cnan = num.zeros((1, 3))
83 cnan[0, :] = self.color_nan
84 colors[num.isnan(values), :] = cnan
86 return colors
88 @classmethod
89 def from_numpy(cls, colors):
90 nbins = colors.shape[0]
91 vs = num.linspace(0., 1., nbins)
92 cpt_data = num.hstack((num.atleast_2d(vs).T, colors))
93 return cls(
94 levels=[
95 CPTLevel(
96 vmin=a[0],
97 vmax=b[0],
98 color_min=[255*x for x in a[1:]],
99 color_max=[255*x for x in b[1:]])
100 for (a, b) in zip(cpt_data[:-1], cpt_data[1:])])
103def get_cpt_path(name):
104 if os.path.exists(name):
105 return name
107 if not re.match(r'[A-Za-z0-9_]+', name):
108 raise Exception('invalid cpt name')
110 fn = util.data_file(os.path.join('colortables', '%s.cpt' % name))
111 if not os.path.exists(fn):
112 raise Exception('cpt file does not exist: %s' % fn)
114 return fn
117def get_cpt(name):
118 return read_cpt(get_cpt_path(name))
121class CPTParseError(Exception):
122 pass
125def read_cpt(filename):
126 with open(filename) as f:
127 color_below = None
128 color_above = None
129 color_nan = None
130 levels = []
131 try:
132 for line in f:
133 line = line.strip()
134 toks = line.split()
136 if line.startswith('#'):
137 continue
139 elif line.startswith('B'):
140 color_below = tuple(map(float, toks[1:4]))
142 elif line.startswith('F'):
143 color_above = tuple(map(float, toks[1:4]))
145 elif line.startswith('N'):
146 color_nan = tuple(map(float, toks[1:4]))
148 else:
149 values = list(map(float, line.split()))
150 vmin = values[0]
151 color_min = tuple(values[1:4])
152 vmax = values[4]
153 color_max = tuple(values[5:8])
154 levels.append(CPTLevel(
155 vmin=vmin,
156 vmax=vmax,
157 color_min=color_min,
158 color_max=color_max))
160 except Exception:
161 raise CPTParseError()
163 return CPT(
164 color_below=color_below,
165 color_above=color_above,
166 color_nan=color_nan,
167 levels=levels)
170def color_to_int(color):
171 return tuple(max(0, min(255, int(round(x)))) for x in color)
174def write_cpt(cpt, filename):
175 with open(filename, 'w') as f:
176 for level in cpt.levels:
177 f.write(
178 '%e %i %i %i %e %i %i %i\n' %
179 ((level.vmin, ) + color_to_int(level.color_min) +
180 (level.vmax, ) + color_to_int(level.color_max)))
182 if cpt.color_below:
183 f.write('B %i %i %i\n' % color_to_int(cpt.color_below))
185 if cpt.color_above:
186 f.write('F %i %i %i\n' % color_to_int(cpt.color_above))
188 if cpt.color_nan:
189 f.write('N %i %i %i\n' % color_to_int(cpt.color_nan))
192def cpt_merge_wet_dry(wet, dry):
193 levels = []
194 for level in wet.levels:
195 if level.vmin < 0.:
196 if level.vmax > 0.:
197 level.vmax = 0.
199 levels.append(level)
201 for level in dry.levels:
202 if level.vmax > 0.:
203 if level.vmin < 0.:
204 level.vmin = 0.
206 levels.append(level)
208 combi = CPT(
209 color_below=wet.color_below,
210 color_above=dry.color_above,
211 color_nan=dry.color_nan,
212 levels=levels)
214 return combi