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

1# https://pyrocko.org/grond - GPLv3 

2# 

3# The Grond Developers, 21st Century 

4import logging 

5from collections import defaultdict 

6 

7import numpy as num 

8 

9from matplotlib import pyplot as plt 

10from mpl_toolkits.axes_grid1 import ImageGrid 

11 

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

13from grond import dataset 

14from grond.meta import store_t 

15 

16km = 1000. 

17 

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

19 

20 

21def darken(c): 

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

23 

24 

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

26 from matplotlib.colors import LinearSegmentedColormap 

27 from matplotlib.collections import LineCollection 

28 

29 cmap = LinearSegmentedColormap.from_list( 

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

31 

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

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

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

35 lc.set_array(t) 

36 axes.add_collection(lc, autolim=True) 

37 

38 

39def polarization( 

40 ds, store, timing, fmin, fmax, ffactor, 

41 time_factor_pre=2., 

42 time_factor_post=2., 

43 distance_min=None, 

44 distance_max=None, 

45 depth_min=None, 

46 depth_max=None, 

47 size_factor=0.05, 

48 nsl_to_time=None, 

49 output_filename=None, 

50 output_format=None, 

51 output_dpi=None): 

52 

53 event = ds.get_event() 

54 stations = ds.get_stations() 

55 

56 source = gf.Source.from_pyrocko_event(event) 

57 

58 trs = [] 

59 for station in stations: 

60 

61 nsl = station.nsl() 

62 

63 dist = source.distance_to(station) 

64 

65 if distance_min is not None and dist < distance_min: 

66 continue 

67 

68 if distance_max is not None and distance_max < dist: 

69 continue 

70 

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

72 continue 

73 

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

75 continue 

76 

77 if nsl_to_time is None: 

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

79 

80 else: 

81 if nsl not in nsl_to_time: 

82 continue 

83 

84 tp = nsl_to_time[nsl] 

85 

86 for component in 'ZNE': 

87 

88 tmin = tp - time_factor_pre / fmin 

89 tmax = tp + time_factor_post / fmin 

90 

91 nslc = nsl + (component,) 

92 

93 freqlimits = [ 

94 fmin / ffactor, 

95 fmin, 

96 fmax, 

97 fmax * ffactor] 

98 

99 tfade = 1.0 / (fmin / ffactor) 

100 

101 try: 

102 trs_projected, trs_restituted, trs_raw, _ = \ 

103 ds.get_waveform( 

104 nslc, 

105 tmin=tmin, 

106 tmax=tmax, 

107 tfade=tfade, 

108 freqlimits=freqlimits, 

109 debug=True) 

110 

111 for tr in trs_projected: 

112 tr.shift(-tp) 

113 

114 trs.extend(trs_projected) 

115 

116 except dataset.NotFound as e: 

117 logger.warning(str(e)) 

118 continue 

119 

120 trace.snuffle(trs, stations=stations) 

121 

122 plot_polarizations( 

123 stations, trs, 

124 event=event, 

125 size_factor=size_factor, 

126 output_filename=output_filename, 

127 output_format=output_format, 

128 output_dpi=output_dpi) 

129 

130 

131def plot_polarizations( 

132 stations, trs, 

133 event=None, 

134 size_factor=0.05, 

135 fontsize=10., 

136 output_filename=None, 

137 output_format=None, 

138 output_dpi=None): 

139 

140 if event is None: 

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

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

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

144 event = od.Loc(clat, clon) 

145 

146 nsl_c_to_trs = defaultdict(dict) 

147 for tr in trs: 

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

149 

150 nsl_to_station = dict( 

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

152 

153 plot.mpl_init(fontsize=fontsize) 

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

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

156 

157 grid = ImageGrid( 

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

159 axes_pad=0.5, 

160 add_all=True, 

161 label_mode='L', 

162 aspect=True) 

163 

164 axes_en = grid[0] 

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

166 

167 axes_dn = grid[1] 

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

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

170 

171 axes_ed = grid[2] 

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

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

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

175 

176 if isinstance(event, model.Event): 

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

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

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

180 

181 grid[3].set_axis_off() 

182 

183 locations = [] 

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

185 station = nsl_to_station[nsl] 

186 n, e = od.latlon_to_ne( 

187 event.lat, event.lon, station.lat, station.lon) 

188 

189 locations.append((n, e)) 

190 

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

192 

193 n_min = num.min(ns) 

194 n_max = num.max(ns) 

195 e_min = num.min(es) 

196 e_max = num.max(es) 

197 

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

199 

200 fontsize_annot = fontsize * 0.7 

201 

202 data = {} 

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

204 

205 color = plot.mpl_graph_color(insl) 

206 

207 try: 

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

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

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

211 

212 except KeyError: 

213 continue 

214 

215 station = nsl_to_station[nsl] 

216 

217 n, e = od.latlon_to_ne( 

218 event.lat, event.lon, station.lat, station.lon) 

219 

220 d = station.depth 

221 

222 axes_en.annotate( 

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

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

225 xycoords='data', 

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

227 textcoords='offset points', 

228 verticalalignment='bottom', 

229 horizontalalignment='left', 

230 rotation=0., 

231 size=fontsize_annot) 

232 

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

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

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

236 

237 arr_e = tr_e.ydata 

238 arr_n = tr_n.ydata 

239 arr_z = tr_z.ydata 

240 arr_t = tr_z.get_xdata() 

241 

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

243 

244 amaxs = [] 

245 amax_hors = [] 

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

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

248 amaxs.append( 

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

250 amax_hors.append( 

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

252 

253 amax = num.median(amaxs) 

254 amax_hor = num.median(amax_hors) 

255 

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

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

258 tmin = arr_t.min() 

259 tmax = arr_t.max() 

260 plot_color_line( 

261 axes_en, 

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

263 arr_t, color, tmin, tmax) 

264 plot_color_line( 

265 axes_dn, 

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

267 arr_t, color, tmin, tmax) 

268 plot_color_line( 

269 axes_ed, 

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

271 arr_t, color, tmin, tmax) 

272 

273 axes_ed.invert_yaxis() 

274 

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

276 axes.autoscale_view(tight=True) 

277 

278 if output_filename is None: 

279 plt.show() 

280 else: 

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