Coverage for /usr/local/lib/python3.11/dist-packages/grond/qc.py: 0%

1import logging 

2from collections import defaultdict 

3 

4import numpy as num 

5 

6from matplotlib import pyplot as plt 

7from mpl_toolkits.axes_grid1 import ImageGrid 

8 

9from pyrocko import gf, orthodrome as od, plot, model, trace 

10from grond import dataset 

11from grond.meta import store_t 

12 

13km = 1000. 

14 

15logger = logging.getLogger('grond.qc') 

16 

17 

18def darken(c): 

19 return (c[0]*0.5, c[1]*0.5, c[2]*0.5) 

20 

21 

22def plot_color_line(axes, x, y, t, color, tmin, tmax): 

23 from matplotlib.colors import LinearSegmentedColormap 

24 from matplotlib.collections import LineCollection 

25 

26 cmap = LinearSegmentedColormap.from_list( 

27 'noname', [color, (1., 1., 1.)], 100) 

28 

29 points = num.array([x, y], dtype=float).T.reshape(-1, 1, 2) 

30 segments = num.concatenate([points[:-1], points[1:]], axis=1) 

31 lc = LineCollection(segments, cmap=cmap, norm=plt.Normalize(tmin, tmax)) 

32 lc.set_array(t) 

33 axes.add_collection(lc, autolim=True) 

34 

35 

36def polarization( 

37 ds, store, timing, fmin, fmax, ffactor, 

38 time_factor_pre=2., 

39 time_factor_post=2., 

40 distance_min=None, 

41 distance_max=None, 

42 depth_min=None, 

43 depth_max=None, 

44 size_factor=0.05, 

45 nsl_to_time=None, 

46 output_filename=None, 

47 output_format=None, 

48 output_dpi=None): 

49 

50 event = ds.get_event() 

51 stations = ds.get_stations() 

52 

53 source = gf.Source.from_pyrocko_event(event) 

54 

55 trs = [] 

56 for station in stations: 

57 

58 nsl = station.nsl() 

59 

60 dist = source.distance_to(station) 

61 

62 if distance_min is not None and dist < distance_min: 

63 continue 

64 

65 if distance_max is not None and distance_max < dist: 

66 continue 

67 

68 if depth_min is not None and station.depth < depth_min: 

69 continue 

70 

71 if depth_max is not None and depth_max < station.depth: 

72 continue 

73 

74 if nsl_to_time is None: 

75 tp = event.time + store_t(store, timing, source, station) 

76 

77 else: 

78 if nsl not in nsl_to_time: 

79 continue 

80 

81 tp = nsl_to_time[nsl] 

82 

83 for component in 'ZNE': 

84 

85 tmin = tp - time_factor_pre / fmin 

86 tmax = tp + time_factor_post / fmin 

87 

88 nslc = nsl + (component,) 

89 

90 freqlimits = [ 

91 fmin / ffactor, 

92 fmin, 

93 fmax, 

94 fmax * ffactor] 

95 

96 tfade = 1.0 / (fmin / ffactor) 

97 

98 try: 

99 trs_projected, trs_restituted, trs_raw, _ = \ 

100 ds.get_waveform( 

101 nslc, 

102 tmin=tmin, 

103 tmax=tmax, 

104 tfade=tfade, 

105 freqlimits=freqlimits, 

106 debug=True) 

107 

108 for tr in trs_projected: 

109 tr.shift(-tp) 

110 

111 trs.extend(trs_projected) 

112 

113 except dataset.NotFound as e: 

114 logger.warning(str(e)) 

115 continue 

116 

117 trace.snuffle(trs, stations=stations) 

118 

119 plot_polarizations( 

120 stations, trs, 

121 event=event, 

122 size_factor=size_factor, 

123 output_filename=output_filename, 

124 output_format=output_format, 

125 output_dpi=output_dpi) 

126 

127 

128def plot_polarizations( 

129 stations, trs, 

130 event=None, 

131 size_factor=0.05, 

132 fontsize=10., 

133 output_filename=None, 

134 output_format=None, 

135 output_dpi=None): 

136 

137 if event is None: 

138 slats = num.array([s.lat for s in stations], dtype=float) 

139 slons = num.array([s.lon for s in stations], dtype=float) 

140 clat, clon = od.geographic_midpoint(slats, slons) 

141 event = od.Loc(clat, clon) 

142 

143 nsl_c_to_trs = defaultdict(dict) 

144 for tr in trs: 

145 nsl_c_to_trs[tr.nslc_id[:3]][tr.nslc_id[3]] = tr 

146 

147 nsl_to_station = dict( 

148 (s.nsl(), s) for s in stations) 

149 

150 plot.mpl_init(fontsize=fontsize) 

151 fig = plt.figure(figsize=plot.mpl_papersize('a4', 'landscape')) 

152 plot.mpl_margins(fig, w=7., h=6., units=fontsize) 

153 

154 grid = ImageGrid( 

155 fig, 111, nrows_ncols=(2, 2), 

156 axes_pad=0.5, 

157 add_all=True, 

158 label_mode='L', 

159 aspect=True) 

160 

161 axes_en = grid[0] 

162 axes_en.set_ylabel('Northing [km]') 

163 

164 axes_dn = grid[1] 

165 axes_dn.locator_params(axis='x', nbins=4) 

166 axes_dn.set_xlabel('Depth [km]') 

167 

168 axes_ed = grid[2] 

169 axes_ed.locator_params(axis='y', nbins=4) 

170 axes_ed.set_ylabel('Depth [km]') 

171 axes_ed.set_xlabel('Easting [km]') 

172 

173 if isinstance(event, model.Event): 

174 axes_en.plot(0., 0., '*') 

175 axes_dn.plot(event.depth/km, 0., '*') 

176 axes_ed.plot(0., event.depth/km, '*') 

177 

178 grid[3].set_axis_off() 

179 

180 locations = [] 

181 for nsl in sorted(nsl_c_to_trs.keys()): 

182 station = nsl_to_station[nsl] 

183 n, e = od.latlon_to_ne( 

184 event.lat, event.lon, station.lat, station.lon) 

185 

186 locations.append((n, e)) 

187 

188 ns, es = num.array(locations, dtype=float).T 

189 

190 n_min = num.min(ns) 

191 n_max = num.max(ns) 

192 e_min = num.min(es) 

193 e_max = num.max(es) 

194 

195 factor = max((n_max - n_min) * size_factor, (e_max - e_min) * size_factor) 

196 

197 fontsize_annot = fontsize * 0.7 

198 

199 data = {} 

200 for insl, nsl in enumerate(sorted(nsl_c_to_trs.keys())): 

201 

202 color = plot.mpl_graph_color(insl) 

203 

204 try: 

205 tr_e = nsl_c_to_trs[nsl]['E'] 

206 tr_n = nsl_c_to_trs[nsl]['N'] 

207 tr_z = nsl_c_to_trs[nsl]['Z'] 

208 

209 except KeyError: 

210 continue 

211 

212 station = nsl_to_station[nsl] 

213 

214 n, e = od.latlon_to_ne( 

215 event.lat, event.lon, station.lat, station.lon) 

216 

217 d = station.depth 

218 

219 axes_en.annotate( 

220 '.'.join(x for x in nsl if x), 

221 xy=(e/km, n/km), 

222 xycoords='data', 

223 xytext=(fontsize_annot/3., fontsize_annot/3.), 

224 textcoords='offset points', 

225 verticalalignment='bottom', 

226 horizontalalignment='left', 

227 rotation=0., 

228 size=fontsize_annot) 

229 

230 axes_en.plot(e/km, n/km, '^', mfc=color, mec=darken(color)) 

231 axes_dn.plot(d/km, n/km, '^', mfc=color, mec=darken(color)) 

232 axes_ed.plot(e/km, d/km, '^', mfc=color, mec=darken(color)) 

233 

234 arr_e = tr_e.ydata 

235 arr_n = tr_n.ydata 

236 arr_z = tr_z.ydata 

237 arr_t = tr_z.get_xdata() 

238 

239 data[nsl] = (arr_e, arr_n, arr_z, arr_t, n, e, d, color) 

240 

241 amaxs = [] 

242 amax_hors = [] 

243 for nsl in sorted(data.keys()): 

244 arr_e, arr_n, arr_z, arr_t, n, e, d, color = data[nsl] 

245 amaxs.append( 

246 num.max(num.abs(num.sqrt(arr_e**2 + arr_n**2 + arr_z**2)))) 

247 amax_hors.append( 

248 num.max(num.abs(num.sqrt(arr_e**2 + arr_n**2)))) 

249 

250 amax = num.median(amaxs) 

251 amax_hor = num.median(amax_hors) 

252 

253 for nsl in sorted(data.keys()): 

254 arr_e, arr_n, arr_z, arr_t, n, e, d, color = data[nsl] 

255 tmin = arr_t.min() 

256 tmax = arr_t.max() 

257 plot_color_line( 

258 axes_en, 

259 (e + arr_e/amax_hor * factor)/km, (n + arr_n/amax_hor * factor)/km, 

260 arr_t, color, tmin, tmax) 

261 plot_color_line( 

262 axes_dn, 

263 (d - arr_z/amax * factor)/km, (n + arr_n/amax * factor)/km, 

264 arr_t, color, tmin, tmax) 

265 plot_color_line( 

266 axes_ed, 

267 (e + arr_e/amax * factor)/km, (d - arr_z/amax * factor)/km, 

268 arr_t, color, tmin, tmax) 

269 

270 axes_ed.invert_yaxis() 

271 

272 for axes in (axes_dn, axes_ed, axes_en): 

273 axes.autoscale_view(tight=True) 

274 

275 if output_filename is None: 

276 plt.show() 

277 else: 

278 fig.savefig(output_filename, format=output_format, dpi=output_dpi)