1# https://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6import os
7import re
9import numpy as num
11from pyrocko.guts import Object, Float, Tuple, List
12from pyrocko import util
15class CPTLevel(Object):
16 vmin = Float.T()
17 vmax = Float.T()
18 color_min = Tuple.T(3, Float.T())
19 color_max = Tuple.T(3, Float.T())
22class CPT(Object):
23 color_below = Tuple.T(3, Float.T(), optional=True)
24 color_above = Tuple.T(3, Float.T(), optional=True)
25 color_nan = Tuple.T(3, Float.T(), optional=True)
26 levels = List.T(CPTLevel.T())
28 @property
29 def vmin(self):
30 if self.levels:
31 return self.levels[0].vmin
32 else:
33 return None
35 @property
36 def vmax(self):
37 if self.levels:
38 return self.levels[-1].vmax
39 else:
40 return None
42 def scale(self, vmin, vmax):
43 vmin_old, vmax_old = self.levels[0].vmin, self.levels[-1].vmax
44 for level in self.levels:
45 level.vmin = (level.vmin - vmin_old) / (vmax_old - vmin_old) * \
46 (vmax - vmin) + vmin
47 level.vmax = (level.vmax - vmin_old) / (vmax_old - vmin_old) * \
48 (vmax - vmin) + vmin
50 def get_lut(self):
51 vs = []
52 colors = []
53 if self.color_below and self.levels:
54 vs.append(self.levels[0].vmin)
55 colors.append(self.color_below)
57 for level in self.levels:
58 vs.append(level.vmin)
59 vs.append(level.vmax)
60 colors.append(level.color_min)
61 colors.append(level.color_max)
63 if self.color_above and self.levels:
64 vs.append(self.levels[-1].vmax)
65 colors.append(self.color_above)
67 vs_lut = num.array(vs)
68 colors_lut = num.array(colors)
69 return vs_lut, colors_lut
71 def __call__(self, values):
72 vs_lut, colors_lut = self.get_lut()
74 colors = num.zeros((values.size, 3))
75 colors[:, 0] = num.interp(values, vs_lut, colors_lut[:, 0])
76 colors[:, 1] = num.interp(values, vs_lut, colors_lut[:, 1])
77 colors[:, 2] = num.interp(values, vs_lut, colors_lut[:, 2])
79 if self.color_nan:
80 cnan = num.zeros((1, 3))
81 cnan[0, :] = self.color_nan
82 colors[num.isnan(values), :] = cnan
84 return colors
86 @classmethod
87 def from_numpy(cls, colors):
88 nbins = colors.shape[0]
89 vs = num.linspace(0., 1., nbins)
90 cpt_data = num.hstack((num.atleast_2d(vs).T, colors))
91 return cls(
92 levels=[
93 CPTLevel(
94 vmin=a[0],
95 vmax=b[0],
96 color_min=[255*x for x in a[1:]],
97 color_max=[255*x for x in b[1:]])
98 for (a, b) in zip(cpt_data[:-1], cpt_data[1:])])
101def get_cpt_path(name):
102 if os.path.exists(name):
103 return name
105 if not re.match(r'[A-Za-z0-9_]+', name):
106 raise Exception('invalid cpt name')
108 fn = util.data_file(os.path.join('colortables', '%s.cpt' % name))
109 if not os.path.exists(fn):
110 raise Exception('cpt file does not exist: %s' % fn)
112 return fn
115def get_cpt(name):
116 return read_cpt(get_cpt_path(name))
119class CPTParseError(Exception):
120 pass
123def read_cpt(filename):
124 with open(filename) as f:
125 color_below = None
126 color_above = None
127 color_nan = None
128 levels = []
129 try:
130 for line in f:
131 line = line.strip()
132 toks = line.split()
134 if line.startswith('#'):
135 continue
137 elif line.startswith('B'):
138 color_below = tuple(map(float, toks[1:4]))
140 elif line.startswith('F'):
141 color_above = tuple(map(float, toks[1:4]))
143 elif line.startswith('N'):
144 color_nan = tuple(map(float, toks[1:4]))
146 else:
147 values = list(map(float, line.split()))
148 vmin = values[0]
149 color_min = tuple(values[1:4])
150 vmax = values[4]
151 color_max = tuple(values[5:8])
152 levels.append(CPTLevel(
153 vmin=vmin,
154 vmax=vmax,
155 color_min=color_min,
156 color_max=color_max))
158 except Exception:
159 raise CPTParseError()
161 return CPT(
162 color_below=color_below,
163 color_above=color_above,
164 color_nan=color_nan,
165 levels=levels)
168def color_to_int(color):
169 return tuple(max(0, min(255, int(round(x)))) for x in color)
172def write_cpt(cpt, filename):
173 with open(filename, 'w') as f:
174 for level in cpt.levels:
175 f.write(
176 '%e %i %i %i %e %i %i %i\n' %
177 ((level.vmin, ) + color_to_int(level.color_min) +
178 (level.vmax, ) + color_to_int(level.color_max)))
180 if cpt.color_below:
181 f.write('B %i %i %i\n' % color_to_int(cpt.color_below))
183 if cpt.color_above:
184 f.write('F %i %i %i\n' % color_to_int(cpt.color_above))
186 if cpt.color_nan:
187 f.write('N %i %i %i\n' % color_to_int(cpt.color_nan))
190def cpt_merge_wet_dry(wet, dry):
191 levels = []
192 for level in wet.levels:
193 if level.vmin < 0.:
194 if level.vmax > 0.:
195 level.vmax = 0.
197 levels.append(level)
199 for level in dry.levels:
200 if level.vmax > 0.:
201 if level.vmin < 0.:
202 level.vmin = 0.
204 levels.append(level)
206 combi = CPT(
207 color_below=wet.color_below,
208 color_above=dry.color_above,
209 color_nan=dry.color_nan,
210 levels=levels)
212 return combi