Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/plot/cpt.py: 95%

1# https://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6import os 

7import re 

8 

9import numpy as num 

10 

11from pyrocko.guts import Object, Float, Tuple, List 

12from pyrocko import util 

13 

14 

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()) 

20 

21 

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()) 

27 

28 @property 

29 def vmin(self): 

30 if self.levels: 

31 return self.levels[0].vmin 

32 else: 

33 return None 

34 

35 @property 

36 def vmax(self): 

37 if self.levels: 

38 return self.levels[-1].vmax 

39 else: 

40 return None 

41 

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 

49 

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) 

56 

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) 

62 

63 if self.color_above and self.levels: 

64 vs.append(self.levels[-1].vmax) 

65 colors.append(self.color_above) 

66 

67 vs_lut = num.array(vs) 

68 colors_lut = num.array(colors) 

69 return vs_lut, colors_lut 

70 

71 def __call__(self, values): 

72 vs_lut, colors_lut = self.get_lut() 

73 

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]) 

78 

79 if self.color_nan: 

80 cnan = num.zeros((1, 3)) 

81 cnan[0, :] = self.color_nan 

82 colors[num.isnan(values), :] = cnan 

83 

84 return colors 

85 

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:])]) 

99 

100 

101def get_cpt_path(name): 

102 if os.path.exists(name): 

103 return name 

104 

105 if not re.match(r'[A-Za-z0-9_]+', name): 

106 raise Exception('invalid cpt name') 

107 

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) 

111 

112 return fn 

113 

114 

115def get_cpt(name): 

116 return read_cpt(get_cpt_path(name)) 

117 

118 

119class CPTParseError(Exception): 

120 pass 

121 

122 

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() 

133 

134 if line.startswith('#'): 

135 continue 

136 

137 elif line.startswith('B'): 

138 color_below = tuple(map(float, toks[1:4])) 

139 

140 elif line.startswith('F'): 

141 color_above = tuple(map(float, toks[1:4])) 

142 

143 elif line.startswith('N'): 

144 color_nan = tuple(map(float, toks[1:4])) 

145 

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)) 

157 

158 except Exception: 

159 raise CPTParseError() 

160 

161 return CPT( 

162 color_below=color_below, 

163 color_above=color_above, 

164 color_nan=color_nan, 

165 levels=levels) 

166 

167 

168def color_to_int(color): 

169 return tuple(max(0, min(255, int(round(x)))) for x in color) 

170 

171 

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))) 

179 

180 if cpt.color_below: 

181 f.write('B %i %i %i\n' % color_to_int(cpt.color_below)) 

182 

183 if cpt.color_above: 

184 f.write('F %i %i %i\n' % color_to_int(cpt.color_above)) 

185 

186 if cpt.color_nan: 

187 f.write('N %i %i %i\n' % color_to_int(cpt.color_nan)) 

188 

189 

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. 

196 

197 levels.append(level) 

198 

199 for level in dry.levels: 

200 if level.vmax > 0.: 

201 if level.vmin < 0.: 

202 level.vmin = 0. 

203 

204 levels.append(level) 

205 

206 combi = CPT( 

207 color_below=wet.color_below, 

208 color_above=dry.color_above, 

209 color_nan=dry.color_nan, 

210 levels=levels) 

211 

212 return combi