Coverage for /usr/local/lib/python3.11/dist-packages/grond/plot/section.py: 24%

229 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-10-26 16:25 +0000

1from __future__ import print_function 

2import numpy as num 

3from matplotlib.axes import Axes 

4from matplotlib.ticker import MultipleLocator 

5 

6from pyrocko.guts import Tuple, Float 

7from pyrocko import plot 

8 

9from .config import PlotConfig 

10 

11guts_prefix = 'grond' 

12 

13 

14def get_callbacks(obj): 

15 try: 

16 return obj.callbacks 

17 except AttributeError: 

18 return obj._callbacks 

19 

20 

21def limits(points): 

22 lims = num.zeros((3, 2)) 

23 if points.size != 0: 

24 lims[:, 0] = num.min(points, axis=0) 

25 lims[:, 1] = num.max(points, axis=0) 

26 

27 return lims 

28 

29 

30class NotEnoughSpace(Exception): 

31 pass 

32 

33 

34class SectionPlotConfig(PlotConfig): 

35 

36 size_cm = Tuple.T( 

37 2, Float.T(), default=(20., 20.)) 

38 

39 margins_em = Tuple.T( 

40 4, Float.T(), default=(7., 5., 7., 5.)) 

41 

42 separator_em = Float.T(default=1.0) 

43 

44 

45class SectionPlot(object): 

46 

47 def __init__(self, config=None): 

48 if config is None: 

49 config = SectionPlotConfig() 

50 

51 self.config = config 

52 self._disconnect_data = [] 

53 self._width = self._height = self._pixels = None 

54 self._plt = plot.mpl_init(self.config.font_size) 

55 self._fig = fig = self._plt.figure(figsize=self.config.size_inch) 

56 

57 rect = [0., 0., 1., 1.] 

58 self._axes_xy = Axes(fig, rect) 

59 self._axes_xz = Axes(fig, rect) 

60 self._axes_zy = Axes(fig, rect) 

61 

62 self._view_limits = num.zeros((3, 2)) 

63 

64 self._view_limits[:, :] = num.nan 

65 

66 self._update_geometry() 

67 

68 for axes in self.axes_list: 

69 fig.add_axes(axes) 

70 self._connect(axes, 'xlim_changed', self.lim_changed_handler) 

71 self._connect(axes, 'ylim_changed', self.lim_changed_handler) 

72 

73 self._cid_resize = fig.canvas.mpl_connect( 

74 'resize_event', self.resize_handler) 

75 

76 try: 

77 self._connect(fig, 'dpi_changed', self.dpi_changed_handler) 

78 except ValueError: 

79 # 'dpi_changed' event has been removed in MPL 3.8. 

80 # canvas 'resize_event' may be sufficient but needs to be checked. 

81 # https://matplotlib.org/stable/api/prev_api_changes/api_changes_3.8.0.html#text-get-rotation 

82 pass 

83 

84 self._lim_changed_depth = 0 

85 

86 def _connect(self, obj, sig, handler): 

87 cid = get_callbacks(obj).connect(sig, handler) 

88 self._disconnect_data.append((obj, cid)) 

89 

90 def _disconnect_all(self): 

91 for obj, cid in self._disconnect_data: 

92 get_callbacks(obj).disconnect(cid) 

93 

94 self._fig.canvas.mpl_disconnect(self._cid_resize) 

95 

96 def dpi_changed_handler(self, fig): 

97 self._update_geometry() 

98 

99 def resize_handler(self, event): 

100 self._update_geometry() 

101 

102 def lim_changed_handler(self, axes): 

103 self._lim_changed_depth += 1 

104 if self._lim_changed_depth < 2: 

105 self._update_layout() 

106 

107 self._lim_changed_depth -= 1 

108 

109 def _update_geometry(self): 

110 w, h = self._fig.canvas.get_width_height() 

111 p = self.get_pixels_factor() 

112 

113 if (self._width, self._height, self._pixels) != (w, h, p): 

114 self._width = w 

115 self._height = h 

116 self._pixels = p 

117 self._update_layout() 

118 

119 @property 

120 def margins(self): 

121 return tuple( 

122 x * self.config.font_size / self._pixels 

123 for x in self.config.margins_em) 

124 

125 @property 

126 def separator(self): 

127 return self.config.separator_em * self.config.font_size / self._pixels 

128 

129 def rect_to_figure_coords(self, rect): 

130 left, bottom, width, height = rect 

131 return ( 

132 left / self._width, 

133 bottom / self._height, 

134 width / self._width, 

135 height / self._height) 

136 

137 def point_to_axes_coords(self, axes, point): 

138 x, y = point 

139 aleft, abottom, awidth, aheight = axes.get_position().bounds 

140 

141 x_fig = x / self._width 

142 y_fig = y / self._height 

143 

144 x_axes = (x_fig - aleft) / awidth 

145 y_axes = (y_fig - abottom) / aheight 

146 

147 return (x_axes, y_axes) 

148 

149 def get_pixels_factor(self): 

150 try: 

151 r = self._fig.canvas.get_renderer() 

152 return 1.0 / r.points_to_pixels(1.0) 

153 except AttributeError: 

154 return 1.0 

155 

156 def make_limits(self, lims): 

157 a = plot.AutoScaler(space=0.05) 

158 return a.make_scale(lims)[:2] 

159 

160 def get_data_limits(self): 

161 xs = [] 

162 ys = [] 

163 zs = [] 

164 xs.extend(self._axes_xy.get_xaxis().get_data_interval()) 

165 ys.extend(self._axes_xy.get_yaxis().get_data_interval()) 

166 xs.extend(self._axes_xz.get_xaxis().get_data_interval()) 

167 zs.extend(self._axes_xz.get_yaxis().get_data_interval()) 

168 zs.extend(self._axes_zy.get_xaxis().get_data_interval()) 

169 ys.extend(self._axes_zy.get_yaxis().get_data_interval()) 

170 lims = num.zeros((3, 2)) 

171 lims[0, :] = num.nanmin(xs), num.nanmax(xs) 

172 lims[1, :] = num.nanmin(ys), num.nanmax(ys) 

173 lims[2, :] = num.nanmin(zs), num.nanmax(zs) 

174 lims[num.logical_not(num.isfinite(lims))] = 0.0 

175 return lims 

176 

177 def set_xlim(self, xmin, xmax): 

178 self._view_limits[0, :] = xmin, xmax 

179 self._update_layout() 

180 

181 def set_ylim(self, ymin, ymax): 

182 self._view_limits[1, :] = ymin, ymax 

183 self._update_layout() 

184 

185 def set_zlim(self, zmin, zmax): 

186 self._view_limits[2, :] = zmin, zmax 

187 self._update_layout() 

188 

189 def _update_layout(self): 

190 data_limits = self.get_data_limits() 

191 

192 limits = num.zeros((3, 2)) 

193 for i in range(3): 

194 limits[i, :] = self.make_limits(data_limits[i, :]) 

195 

196 mask = num.isfinite(self._view_limits) 

197 limits[mask] = self._view_limits[mask] 

198 

199 deltas = limits[:, 1] - limits[:, 0] 

200 

201 data_w = deltas[0] + deltas[2] 

202 data_h = deltas[1] + deltas[2] 

203 

204 ml, mt, mr, mb = self.margins 

205 ms = self.separator 

206 

207 data_r = data_h / data_w 

208 em = self.config.font_size 

209 w = self._width 

210 h = self._height 

211 fig_w_avail = w - mr - ml - ms 

212 fig_h_avail = h - mt - mb - ms 

213 

214 if fig_w_avail <= 0.0 or fig_h_avail <= 0.0: 

215 raise NotEnoughSpace() 

216 

217 fig_r = fig_h_avail / fig_w_avail 

218 

219 if data_r < fig_r: 

220 data_expanded_h = data_w * fig_r 

221 data_expanded_w = data_w 

222 else: 

223 data_expanded_h = data_h 

224 data_expanded_w = data_h / fig_r 

225 

226 limits[0, 0] -= 0.5 * (data_expanded_w - data_w) 

227 limits[0, 1] += 0.5 * (data_expanded_w - data_w) 

228 limits[1, 0] -= 0.5 * (data_expanded_h - data_h) 

229 limits[1, 1] += 0.5 * (data_expanded_h - data_h) 

230 

231 deltas = limits[:, 1] - limits[:, 0] 

232 

233 w1 = fig_w_avail * deltas[0] / data_expanded_w 

234 w2 = fig_w_avail * deltas[2] / data_expanded_w 

235 

236 h1 = fig_h_avail * deltas[1] / data_expanded_h 

237 h2 = fig_h_avail * deltas[2] / data_expanded_h 

238 

239 rect_xy = [ml, mb+h2+ms, w1, h1] 

240 rect_xz = [ml, mb, w1, h2] 

241 rect_zy = [ml+w1+ms, mb+h2+ms, w2, h1] 

242 

243 axes_xy, axes_xz, axes_zy = self.axes_list 

244 

245 axes_xy.set_position( 

246 self.rect_to_figure_coords(rect_xy), which='both') 

247 axes_xz.set_position( 

248 self.rect_to_figure_coords(rect_xz), which='both') 

249 axes_zy.set_position( 

250 self.rect_to_figure_coords(rect_zy), which='both') 

251 

252 def wcenter(rect): 

253 return rect[0] + rect[2]*0.5 

254 

255 def hcenter(rect): 

256 return rect[1] + rect[3]*0.5 

257 

258 self.set_label_coords( 

259 axes_xy, 'x', [wcenter(rect_xy), h - 1.0*em]) 

260 self.set_label_coords( 

261 axes_xy, 'y', [2.0*em, hcenter(rect_xy)]) 

262 self.set_label_coords( 

263 axes_zy, 'x', [wcenter(rect_zy), h - 1.0*em]) 

264 self.set_label_coords( 

265 axes_xz, 'y', [2.0*em, hcenter(rect_xz)]) 

266 

267 scaler = plot.AutoScaler() 

268 inc = scaler.make_scale( 

269 [0, min(data_expanded_w, data_expanded_h)], override_mode='off')[2] 

270 

271 axes_xy.set_xlim(*limits[0, :]) 

272 axes_xy.set_ylim(*limits[1, :]) 

273 axes_xy.get_xaxis().set_tick_params( 

274 bottom=False, top=True, labelbottom=False, labeltop=True) 

275 axes_xy.get_yaxis().set_tick_params( 

276 left=True, labelleft=True, right=False, labelright=False) 

277 

278 axes_xz.set_xlim(*limits[0, :]) 

279 axes_xz.set_ylim(*limits[2, ::-1]) 

280 axes_xz.get_xaxis().set_tick_params( 

281 bottom=True, top=False, labelbottom=False, labeltop=False) 

282 axes_xz.get_yaxis().set_tick_params( 

283 left=True, labelleft=True, right=True, labelright=False) 

284 

285 axes_zy.set_xlim(*limits[2, :]) 

286 axes_zy.set_ylim(*limits[1, :]) 

287 axes_zy.get_xaxis().set_tick_params( 

288 bottom=True, top=True, labelbottom=False, labeltop=True) 

289 axes_zy.get_yaxis().set_tick_params( 

290 left=False, labelleft=False, right=True, labelright=False) 

291 

292 for axes in self.axes_list: 

293 tl = MultipleLocator(inc) 

294 axes.get_xaxis().set_major_locator(tl) 

295 tl = MultipleLocator(inc) 

296 axes.get_yaxis().set_major_locator(tl) 

297 

298 def set_label_coords(self, axes, which, point): 

299 axis = axes.get_xaxis() if which == 'x' else axes.get_yaxis() 

300 axis.set_label_coords(*self.point_to_axes_coords(axes, point)) 

301 

302 @property 

303 def fig(self): 

304 return self._fig 

305 

306 @property 

307 def axes_xy(self): 

308 return self._axes_xy 

309 

310 @property 

311 def axes_xz(self): 

312 return self._axes_xz 

313 

314 @property 

315 def axes_zy(self): 

316 return self._axes_zy 

317 

318 @property 

319 def axes_list(self): 

320 return [ 

321 self._axes_xy, self._axes_xz, self._axes_zy] 

322 

323 def plot(self, points, *args, **kwargs): 

324 self._axes_xy.plot(points[:, 0], points[:, 1], *args, **kwargs) 

325 self._axes_xz.plot(points[:, 0], points[:, 2], *args, **kwargs) 

326 self._axes_zy.plot(points[:, 2], points[:, 1], *args, **kwargs) 

327 

328 def close(self): 

329 self._disconnect_all() 

330 self._plt.close(self._fig) 

331 

332 def show(self): 

333 self._plt.show() 

334 

335 def set_xlabel(self, s): 

336 self._axes_xy.set_xlabel(s) 

337 

338 def set_ylabel(self, s): 

339 self._axes_xy.set_ylabel(s) 

340 

341 def set_zlabel(self, s): 

342 self._axes_xz.set_ylabel(s) 

343 self._axes_zy.set_xlabel(s)