1import logging 

2 

3import numpy as num 

4from matplotlib import cm, patches 

5 

6from pyrocko import orthodrome as pod 

7from pyrocko.guts import Float, Bool, Tuple 

8 

9from pyrocko.plot import mpl_color, mpl_init, automap 

10 

11from grond.plot.section import SectionPlotConfig, SectionPlot 

12from grond.plot.collection import PlotItem 

13from grond.plot.config import PlotConfig 

14from grond.problems.plot import fixlim 

15from matplotlib import pyplot as plt 

16 

17logger = logging.getLogger('grond.problem.singleforce.plot') 

18 

19guts_prefix = 'grond' 

20 

21km = 1e3 

22 

23 

24class SFLocationPlot(SectionPlotConfig): 

25 ''' MT location plot of the best solutions in three cross-sections. ''' 

26 name = 'location_sf' 

27 normalisation_gamma = Float.T( 

28 default=3., 

29 help='Normalisation of colors and alpha as :math:`x^\\gamma`.' 

30 'A linear colormap/alpha with :math:`\\gamma=1`.') 

31 

32 def make(self, environ): 

33 environ.setup_modelling() 

34 cm = environ.get_plot_collection_manager() 

35 history = environ.get_history(subset='harvest') 

36 mpl_init(fontsize=self.font_size) 

37 self._to_be_closed = [] 

38 cm.create_group_mpl( 

39 self, 

40 self.draw_figures(history), 

41 title=u'SF Location', 

42 section='solution', 

43 feather_icon='target', 

44 description=u''' 

45Location plot of the ensemble of best solutions in three cross-sections. 

46 

47The coordinate range is defined by the search space given in the config file. 

48Dots show best single force mechanisms, and colors indicate low (red) and 

49high (blue) misfit. 

50''') 

51 for obj in self._to_be_closed: 

52 obj.close() 

53 

54 def draw_figures(self, history, color_p_axis=False): 

55 from matplotlib import colors 

56 

57 color = 'black' 

58 fontsize = self.font_size 

59 markersize = fontsize * 1.5 

60 

61 problem = history.problem 

62 sp = SectionPlot(config=self) 

63 self._to_be_closed.append(sp) 

64 

65 fig = sp.fig 

66 axes_en = sp.axes_xy 

67 axes_dn = sp.axes_zy 

68 axes_ed = sp.axes_xz 

69 

70 bounds = problem.get_combined_bounds() 

71 

72 models = history.get_sorted_primary_models()[::-1] 

73 

74 iorder = num.arange(history.nmodels) 

75 

76 for parname, set_label, set_lim in [ 

77 ['east_shift', sp.set_xlabel, sp.set_xlim], 

78 ['north_shift', sp.set_ylabel, sp.set_ylim], 

79 ['depth', sp.set_zlabel, sp.set_zlim]]: 

80 

81 ipar = problem.name_to_index(parname) 

82 par = problem.combined[ipar] 

83 set_label(par.get_label()) 

84 xmin, xmax = fixlim(*par.scaled(bounds[ipar])) 

85 set_lim(xmin, xmax) 

86 

87 def scale_size(source): 

88 return markersize * 1.5 

89 

90 for axes, xparname, yparname in [ 

91 (axes_en, 'east_shift', 'north_shift'), 

92 (axes_dn, 'depth', 'north_shift'), 

93 (axes_ed, 'east_shift', 'depth')]: 

94 

95 ixpar = problem.name_to_index(xparname) 

96 iypar = problem.name_to_index(yparname) 

97 

98 xpar = problem.combined[ixpar] 

99 ypar = problem.combined[iypar] 

100 

101 xmin, xmax = fixlim(*xpar.scaled(bounds[ixpar])) 

102 ymin, ymax = fixlim(*ypar.scaled(bounds[iypar])) 

103 

104 try: 

105 axes.set_facecolor(mpl_color('aluminium1')) 

106 except AttributeError: 

107 axes.patch.set_facecolor(mpl_color('aluminium1')) 

108 

109 rect = patches.Rectangle( 

110 (xmin, ymin), xmax-xmin, ymax-ymin, 

111 facecolor=mpl_color('white'), 

112 edgecolor=mpl_color('aluminium2'), 

113 zorder=1) 

114 

115 axes.add_patch(rect) 

116 

117 cmap = cm.ScalarMappable( 

118 norm=colors.PowerNorm( 

119 gamma=self.normalisation_gamma, 

120 vmin=iorder.min(), 

121 vmax=iorder.max()), 

122 

123 cmap=plt.get_cmap('coolwarm')) 

124 

125 for ix, x in enumerate(models): 

126 

127 source = problem.get_source(x) 

128 fx = problem.extract(x, ixpar) 

129 fy = problem.extract(x, iypar) 

130 sx, sy = xpar.scaled(fx), ypar.scaled(fy) 

131 

132 # TODO: Add rotation in cross-sections 

133 color = cmap.to_rgba(iorder[ix]) 

134 

135 alpha = (iorder[ix] - iorder.min()) / \ 

136 float(iorder.max() - iorder.min()) 

137 alpha = alpha**self.normalisation_gamma 

138 

139 axes.scatter( 

140 [sx], [sy], 

141 c=[color], 

142 s=[scale_size(source)], 

143 alpha=alpha, 

144 zorder=2) 

145 

146 item = PlotItem(name='main') 

147 return [[item, fig]] 

148 

149 

150class SFForcePlot(PlotConfig): 

151 ''' Maps showing horizontal and vertical force 

152 of the best single force model ''' 

153 

154 name = 'forces_singleforce' 

155 

156 size_cm = Tuple.T( 

157 2, Float.T(), 

158 default=(15., 15.), 

159 help='width and length of the figure in cm') 

160 show_topo = Bool.T( 

161 default=False, 

162 help='show topography') 

163 show_grid = Bool.T( 

164 default=True, 

165 help='show the lat/lon grid') 

166 show_rivers = Bool.T( 

167 default=True, 

168 help='show rivers on the map') 

169 radius = Float.T( 

170 optional=True, 

171 help='radius of the map around campaign center lat/lon') 

172 

173 def make(self, environ): 

174 cm = environ.get_plot_collection_manager() 

175 history = environ.get_history(subset='harvest') 

176 optimiser = environ.get_optimiser() 

177 ds = environ.get_dataset() 

178 

179 environ.setup_modelling() 

180 

181 cm.create_group_automap( 

182 self, 

183 self.draw_best_sf(ds, history, optimiser), 

184 title=u'Single Force Source Forces', 

185 section='solution', 

186 feather_icon='map', 

187 description=u''' 

188Maps showing location and force vectors of the best Single Force Source model. 

189 

190Arrows show the modelled forces (red arrows). The top plot shows the horizontal 

191forces and the bottom plot the vertical force. The dot indicates the location 

192of the best single force source model. 

193''') 

194 

195 def draw_best_sf(self, ds, history, optimiser, vertical=False): 

196 from grond.core import make_stats 

197 

198 source = history.get_best_source() 

199 

200 problem = history.problem 

201 models = history.models 

202 

203 stats = make_stats( 

204 problem, models, history.get_primary_chain_misfits()) 

205 

206 def plot_sf(source, stats, ifig, vertical=False): 

207 orient = 'vertical' if vertical else 'horizontal' 

208 

209 item = PlotItem( 

210 name='fig_%i' % ifig, 

211 attributes={}, 

212 title=u'Best %s single force model force vector' % orient, 

213 description=u''' 

214Single force source %s force vector for the best model (red). The circle shows 

215the 95%% confidence ellipse. 

216''' % orient) 

217 

218 event = source.pyrocko_event() 

219 

220 radius = self.radius 

221 if radius is None or radius < 30.*km: 

222 logger.warn( 

223 'Radius too small, defaulting to 30 km') 

224 radius = 30*km 

225 

226 m = automap.Map( 

227 width=self.size_cm[0], 

228 height=self.size_cm[1], 

229 lat=event.lat, 

230 lon=event.lon, 

231 radius=radius, 

232 show_topo=self.show_topo, 

233 show_grid=self.show_grid, 

234 show_rivers=self.show_rivers, 

235 color_wet=(216, 242, 254), 

236 color_dry=(238, 236, 230)) 

237 

238 offset_scale = num.abs([source.fn, source.fe, source.fd]).sum() 

239 size = num.linalg.norm(self.size_cm) 

240 

241 scale = (size / 5.) / offset_scale 

242 

243 lat, lon = pod.ne_to_latlon( 

244 event.lat, 

245 event.lon, 

246 source.north_shift, 

247 source.east_shift) 

248 

249 stats_dict = stats.get_values_dict() 

250 

251 if vertical: 

252 rows = [lon, lat, 

253 0., -source.fd * scale, 

254 (stats_dict['fn.std'] + stats_dict['fe.std']) * scale, 

255 stats_dict['fd.std'] * scale, 

256 0.] 

257 

258 else: 

259 rows = [lon, lat, 

260 source.fe * scale, source.fn * scale, 

261 stats_dict['fe.std'] * scale, 

262 stats_dict['fn.std'] * scale, 

263 0.] 

264 

265 fontsize = 10. 

266 

267 default_psxy_style = { 

268 'h': 0, 

269 'W': '2.0p,red', 

270 'A': '+p4p,black+e+a40', 

271 'G': 'red', 

272 't': 30, 

273 'L': True, 

274 'S': 'e1c/0.95/%d' % fontsize, 

275 } 

276 

277 m.gmt.psvelo( 

278 in_rows=[rows], 

279 *m.jxyr, 

280 **default_psxy_style) 

281 

282 m.gmt.psxy( 

283 S='c10p', 

284 in_rows=[[lon, lat]], 

285 W='1p,black', 

286 G='orange3', 

287 *m.jxyr) 

288 

289 return (item, m) 

290 

291 ifig = 0 

292 for vertical in (False, True): 

293 yield plot_sf(source, stats, ifig, vertical) 

294 ifig += 1 

295 

296# TODO Fuzzy Single Force Plot