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

541 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2024-01-02 13:30 +0000

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6''' 

7Matplotlib plotting with some fancy extras. 

8''' 

9 

10from collections import defaultdict 

11import math 

12import logging 

13 

14import numpy as num 

15import matplotlib 

16from matplotlib.axes import Axes 

17# from matplotlib.ticker import MultipleLocator 

18from matplotlib import cm, colors, colorbar, figure 

19 

20from pyrocko.guts import Tuple, Float, Object 

21from pyrocko import plot 

22 

23import scipy.optimize 

24 

25logger = logging.getLogger('pyrocko.plot.smartplot') 

26 

27guts_prefix = 'pf' 

28 

29inch = 2.54 

30 

31 

32def get_callbacks(obj): 

33 try: 

34 return obj.callbacks 

35 except AttributeError: 

36 return obj._callbacks 

37 

38 

39class SmartplotAxes(Axes): 

40 

41 if matplotlib.__version__.split('.') < '3.6'.split('.'): 

42 # Subclassing cla is deprecated on newer mpl but need this fallback for 

43 # older versions. Code is duplicated because mpl behaviour depends 

44 # on the existence of cla in the subclass... 

45 def cla(self): 

46 if hasattr(self, 'callbacks'): 

47 callbacks = self.callbacks 

48 Axes.cla(self) 

49 self.callbacks = callbacks 

50 else: 

51 Axes.cla(self) 

52 

53 else: 

54 def clear(self): 

55 if hasattr(self, 'callbacks'): 

56 callbacks = self.callbacks 

57 Axes.clear(self) 

58 self.callbacks = callbacks 

59 elif hasattr(self, '_callbacks'): 

60 callbacks = self._callbacks 

61 Axes.clear(self) 

62 self._callbacks = callbacks 

63 else: 

64 Axes.clear(self) 

65 

66 

67class SmartplotFigure(figure.Figure): 

68 

69 def set_smartplot(self, plot): 

70 self._smartplot = plot 

71 

72 def draw(self, *args, **kwargs): 

73 if hasattr(self, '_smartplot'): 

74 try: 

75 self._smartplot._update_layout() 

76 except NotEnoughSpace: 

77 logger.error('Figure is too small to show the plot.') 

78 return 

79 

80 return figure.Figure.draw(self, *args, **kwargs) 

81 

82 

83def limits(points): 

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

85 if points.size != 0: 

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

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

88 

89 return lims 

90 

91 

92def wcenter(rect): 

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

94 

95 

96def hcenter(rect): 

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

98 

99 

100def window_min(n, w, ml, mu, s, x): 

101 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + math.floor(x) * s 

102 

103 

104def window_max(n, w, ml, mu, s, x): 

105 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + (math.floor(x)-1) * s 

106 

107 

108def make_smap(cmap, norm=None): 

109 if isinstance(norm, tuple): 

110 norm = colors.Normalize(*norm, clip=False) 

111 smap = cm.ScalarMappable(cmap=cmap, norm=norm) 

112 smap._A = [] # not needed in newer versions of mpl? 

113 return smap 

114 

115 

116def solve_layout_fixed_panels(size, shape, limits, aspects, fracs=None): 

117 

118 weight_aspect = 1000. 

119 

120 sx, sy = size 

121 nx, ny = shape 

122 nvar = nx+ny 

123 vxs, vys = limits 

124 uxs = vxs[:, 1] - vxs[:, 0] 

125 uys = vys[:, 1] - vys[:, 0] 

126 aspects_xx, aspects_yy, aspects_xy = aspects 

127 

128 if fracs is None: 

129 wxs = num.full(nx, sx / nx) 

130 wys = num.full(ny, sy / ny) 

131 else: 

132 frac_x, frac_y = fracs 

133 wxs = sx * frac_x / num.sum(frac_x) 

134 wys = sy * frac_y / num.sum(frac_y) 

135 

136 data = [] 

137 weights = [] 

138 rows = [] 

139 bounds = [] 

140 for ix in range(nx): 

141 u = uxs[ix] 

142 assert u > 0.0 

143 row = num.zeros(nvar) 

144 row[ix] = u 

145 rows.append(row) 

146 data.append(wxs[ix]) 

147 weights.append(1.0 / u) 

148 bounds.append((0, wxs[ix] / u)) 

149 

150 for iy in range(ny): 

151 u = uys[iy] 

152 assert u > 0.0 

153 row = num.zeros(nvar) 

154 row[nx+iy] = u 

155 rows.append(row) 

156 data.append(wys[iy]) 

157 weights.append(1.0) 

158 bounds.append((0, wys[iy] / u)) 

159 

160 for ix1, ix2, aspect in aspects_xx: 

161 row = num.zeros(nvar) 

162 row[ix1] = aspect 

163 row[ix2] = -1.0 

164 weights.append(weight_aspect/aspect) 

165 rows.append(row) 

166 data.append(0.0) 

167 

168 for iy1, iy2, aspect in aspects_yy: 

169 row = num.zeros(nvar) 

170 row[nx+iy1] = aspect 

171 row[nx+iy2] = -1.0 

172 weights.append(weight_aspect/aspect) 

173 rows.append(row) 

174 data.append(0.0) 

175 

176 for ix, iy, aspect in aspects_xy: 

177 row = num.zeros(nvar) 

178 row[ix] = aspect 

179 row[nx+iy] = -1.0 

180 weights.append(weight_aspect/aspect) 

181 rows.append(row) 

182 data.append(0.0) 

183 

184 weights = num.array(weights) 

185 data = num.array(data) 

186 mat = num.vstack(rows) * weights[:, num.newaxis] 

187 data *= weights 

188 

189 bounds = num.array(bounds).T 

190 

191 model = scipy.optimize.lsq_linear(mat, data, bounds).x 

192 

193 cxs = model[:nx] 

194 cys = model[nx:nx+ny] 

195 

196 vlimits_x = num.zeros((nx, 2)) 

197 for ix in range(nx): 

198 u = wxs[ix] / cxs[ix] 

199 vmin, vmax = vxs[ix] 

200 udata = vmax - vmin 

201 eps = 1e-7 * u 

202 assert udata <= u + eps 

203 vlimits_x[ix, 0] = (vmin + vmax) / 2.0 - u / 2.0 

204 vlimits_x[ix, 1] = (vmin + vmax) / 2.0 + u / 2.0 

205 

206 vlimits_y = num.zeros((ny, 2)) 

207 for iy in range(ny): 

208 u = wys[iy] / cys[iy] 

209 vmin, vmax = vys[iy] 

210 udata = vmax - vmin 

211 eps = 1e-7 * u 

212 assert udata <= u + eps 

213 vlimits_y[iy, 0] = (vmin + vmax) / 2.0 - u / 2.0 

214 vlimits_y[iy, 1] = (vmin + vmax) / 2.0 + u / 2.0 

215 

216 def check_aspect(a, awant, eps=1e-2): 

217 if abs(1.0 - (a/awant)) > eps: 

218 logger.error( 

219 'Unable to comply with requested aspect ratio ' 

220 '(wanted: %g, achieved: %g)' % (awant, a)) 

221 

222 for ix1, ix2, aspect in aspects_xx: 

223 check_aspect(cxs[ix2] / cxs[ix1], aspect) 

224 

225 for iy1, iy2, aspect in aspects_yy: 

226 check_aspect(cys[iy2] / cys[iy1], aspect) 

227 

228 for ix, iy, aspect in aspects_xy: 

229 check_aspect(cys[iy] / cxs[ix], aspect) 

230 

231 return (vlimits_x, vlimits_y), (wxs, wys) 

232 

233 

234def solve_layout_iterative(size, shape, limits, aspects, niterations=3): 

235 

236 sx, sy = size 

237 nx, ny = shape 

238 vxs, vys = limits 

239 uxs = vxs[:, 1] - vxs[:, 0] 

240 uys = vys[:, 1] - vys[:, 0] 

241 aspects_xx, aspects_yy, aspects_xy = aspects 

242 

243 fracs_x, fracs_y = num.ones(nx), num.ones(ny) 

244 for i in range(niterations): 

245 (vlimits_x, vlimits_y), (wxs, wys) = solve_layout_fixed_panels( 

246 size, shape, limits, aspects, (fracs_x, fracs_y)) 

247 

248 uxs_view = vlimits_x[:, 1] - vlimits_x[:, 0] 

249 uys_view = vlimits_y[:, 1] - vlimits_y[:, 0] 

250 wxs_used = wxs * uxs / uxs_view 

251 wys_used = wys * uys / uys_view 

252 # wxs_wasted = wxs * (1.0 - uxs / uxs_view) 

253 # wys_wasted = wys * (1.0 - uys / uys_view) 

254 

255 fracs_x = wxs_used 

256 fracs_y = wys_used 

257 

258 return (vlimits_x, vlimits_y), (wxs, wys) 

259 

260 

261class PlotError(Exception): 

262 pass 

263 

264 

265class NotEnoughSpace(PlotError): 

266 pass 

267 

268 

269class PlotConfig(Object): 

270 ''' 

271 Configuration for :py:class:`Plot`. 

272 ''' 

273 

274 font_size = Float.T(default=9.0) 

275 

276 size_cm = Tuple.T( 

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

278 

279 margins_em = Tuple.T( 

280 4, Float.T(), default=(8., 6., 8., 6.)) 

281 

282 separator_em = Float.T(default=1.5) 

283 

284 colorbar_width_em = Float.T(default=2.0) 

285 

286 label_offset_em = Tuple.T( 

287 2, Float.T(), default=(2., 2.)) 

288 

289 tick_label_offset_em = Tuple.T( 

290 2, Float.T(), default=(0.5, 0.5)) 

291 

292 @property 

293 def size_inch(self): 

294 return self.size_cm[0]/inch, self.size_cm[1]/inch 

295 

296 

297class Plot(object): 

298 ''' 

299 Matplotlib plotting with some fancy extras. 

300 

301 - Absolute sized figure margins, also for interactive plots. 

302 - Improved label placement for grids of axes. 

303 - Improved shared axis'es across multiple axes, e.g. for cross section 

304 plots. 

305 - Fixed aspect plotting across multiple axis'es on separate axes. 

306 - Automatic subplot sizing based on aspect and data limit constraints. 

307 - Serializable plot configuration. 

308 ''' 

309 

310 def __init__( 

311 self, x_dims=['x'], y_dims=['y'], z_dims=[], config=None, 

312 fig=None, call_mpl_init=True): 

313 

314 if config is None: 

315 config = PlotConfig() 

316 

317 self._shape = len(x_dims), len(y_dims) 

318 

319 dims = [] 

320 for dim in x_dims + y_dims + z_dims: 

321 dim = dim.lstrip('-') 

322 if dim not in dims: 

323 dims.append(dim) 

324 

325 self.config = config 

326 self._disconnect_data = [] 

327 self._width = self._height = self._pixels = None 

328 if call_mpl_init: 

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

330 

331 if fig is None: 

332 fig = self._plt.figure( 

333 figsize=self.config.size_inch, FigureClass=SmartplotFigure) 

334 else: 

335 assert isinstance(fig, SmartplotFigure) 

336 

337 fig.set_smartplot(self) 

338 

339 self._fig = fig 

340 self._colorbar_width = 0.0 

341 self._colorbar_height = 0.0 

342 self._colorbar_axes = [] 

343 

344 self._dims = dims 

345 self._dim_index = self._dims.index 

346 self._ndims = len(dims) 

347 self._labels = {} 

348 self._aspects = {} 

349 

350 self.setup_axes() 

351 

352 self._view_limits = num.zeros((self._ndims, 2)) 

353 

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

355 self._last_mpl_view_limits = None 

356 

357 self._x_dims = [dim.lstrip('-') for dim in x_dims] 

358 self._x_dims_invert = [dim.startswith('-') for dim in x_dims] 

359 

360 self._y_dims = [dim.lstrip('-') for dim in y_dims] 

361 self._y_dims_invert = [dim.startswith('-') for dim in y_dims] 

362 

363 self._z_dims = [dim.lstrip('-') for dim in z_dims] 

364 self._z_dims_invert = [dim.startswith('-') for dim in z_dims] 

365 

366 self._mappables = {} 

367 self._updating_layout = False 

368 

369 self._need_update_layout = True 

370 self._update_geometry() 

371 

372 for axes in self.axes_list: 

373 fig.add_axes(axes) 

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

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

376 

377 self._cid_resize = fig.canvas.mpl_connect( 

378 'resize_event', self.resize_handler) 

379 

380 try: 

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

382 except ValueError: 

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

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

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

386 pass 

387 

388 self._lim_changed_depth = 0 

389 

390 def reset_size(self): 

391 self._fig.set_size_inches(self.config.size_inch) 

392 

393 def axes(self, ix, iy): 

394 if not (isinstance(ix, int) and isinstance(iy, int)): 

395 ix = self._x_dims.index(ix) 

396 iy = self._y_dims.index(iy) 

397 

398 return self._axes[iy][ix] 

399 

400 def set_color_dim(self, mappable, dim): 

401 assert dim in self._dims 

402 self._mappables[mappable] = dim 

403 

404 def set_aspect(self, ydim, xdim, aspect=1.0): 

405 self._aspects[ydim, xdim] = aspect 

406 

407 @property 

408 def dims(self): 

409 return self._dims 

410 

411 @property 

412 def fig(self): 

413 return self._fig 

414 

415 @property 

416 def axes_list(self): 

417 axes = [] 

418 for row in self._axes: 

419 axes.extend(row) 

420 return axes 

421 

422 @property 

423 def axes_bottom_list(self): 

424 return self._axes[0] 

425 

426 @property 

427 def axes_left_list(self): 

428 return [row[0] for row in self._axes] 

429 

430 def setup_axes(self): 

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

432 nx, ny = self._shape 

433 axes = [] 

434 for iy in range(ny): 

435 axes.append([]) 

436 for ix in range(nx): 

437 axes[-1].append(SmartplotAxes(self.fig, rect)) 

438 

439 self._axes = axes 

440 

441 for _, _, axes_ in self.iaxes(): 

442 axes_.set_autoscale_on(False) 

443 

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

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

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

447 

448 def _disconnect_all(self): 

449 for obj, cid in self._disconnect_data: 

450 get_callbacks(obj).disconnect(cid) 

451 

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

453 

454 def dpi_changed_handler(self, fig): 

455 if self._updating_layout: 

456 return 

457 

458 self._update_geometry() 

459 

460 def resize_handler(self, event): 

461 if self._updating_layout: 

462 return 

463 

464 self._update_geometry() 

465 

466 def lim_changed_handler(self, axes): 

467 if self._updating_layout: 

468 return 

469 

470 current = self._get_mpl_view_limits() 

471 last = self._last_mpl_view_limits 

472 if last is None: 

473 return 

474 

475 for iy, ix, axes in self.iaxes(): 

476 acurrent = current[iy][ix] 

477 alast = last[iy][ix] 

478 if acurrent[0] != alast[0]: 

479 xdim = self._x_dims[ix] 

480 logger.debug( 

481 'X limits have been changed interactively in subplot ' 

482 '(%i, %i)' % (ix, iy)) 

483 self.set_lim(xdim, *sorted(acurrent[0])) 

484 

485 if acurrent[1] != alast[1]: 

486 ydim = self._y_dims[iy] 

487 logger.debug( 

488 'Y limits have been changed interactively in subplot ' 

489 '(%i, %i)' % (ix, iy)) 

490 self.set_lim(ydim, *sorted(acurrent[1])) 

491 

492 self.need_update_layout() 

493 

494 def _update_geometry(self): 

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

496 dp = self.get_device_pixel_ratio() 

497 p = self.get_pixels_factor() * dp 

498 

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

500 logger.debug( 

501 'New figure size: %g x %g, ' 

502 'logical-pixel/point: %g, physical-pixel/logical-pixel: %g' % ( 

503 w, h, p, dp)) 

504 

505 self._width = w # logical pixel 

506 self._height = h # logical pixel 

507 self._pixels = p # logical pixel / point 

508 self._device_pixel_ratio = dp # physical / logical 

509 self.need_update_layout() 

510 

511 @property 

512 def margins(self): 

513 return tuple( 

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

515 for x in self.config.margins_em) 

516 

517 @property 

518 def separator(self): 

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

520 

521 def rect_to_figure_coords(self, rect): 

522 left, bottom, width, height = rect 

523 return ( 

524 left / self._width, 

525 bottom / self._height, 

526 width / self._width, 

527 height / self._height) 

528 

529 def point_to_axes_coords(self, axes, point): 

530 x, y = point 

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

532 

533 x_fig = x / self._width 

534 y_fig = y / self._height 

535 

536 x_axes = (x_fig - aleft) / awidth 

537 y_axes = (y_fig - abottom) / aheight 

538 

539 return (x_axes, y_axes) 

540 

541 def get_pixels_factor(self): 

542 try: 

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

544 return 1.0 / r.points_to_pixels(1.0) 

545 except AttributeError: 

546 return 1.0 

547 

548 def get_device_pixel_ratio(self): 

549 try: 

550 return self._fig.canvas.device_pixel_ratio 

551 except AttributeError: 

552 return 1.0 

553 

554 def make_limits(self, lims): 

555 a = plot.AutoScaler(space=0.05) 

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

557 

558 def iaxes(self): 

559 for iy, row in enumerate(self._axes): 

560 for ix, axes in enumerate(row): 

561 yield iy, ix, axes 

562 

563 def get_data_limits(self): 

564 dim_to_values = defaultdict(list) 

565 for iy, ix, axes in self.iaxes(): 

566 dim_to_values[self._y_dims[iy]].extend( 

567 axes.get_yaxis().get_data_interval()) 

568 dim_to_values[self._x_dims[ix]].extend( 

569 axes.get_xaxis().get_data_interval()) 

570 

571 for mappable, dim in self._mappables.items(): 

572 dim_to_values[dim].extend(mappable.get_clim()) 

573 

574 lims = num.zeros((self._ndims, 2)) 

575 for idim in range(self._ndims): 

576 dim = self._dims[idim] 

577 if dim in dim_to_values: 

578 vs = num.array( 

579 dim_to_values[self._dims[idim]], dtype=float) 

580 vs = vs[num.isfinite(vs)] 

581 if vs.size > 0: 

582 lims[idim, :] = num.min(vs), num.max(vs) 

583 else: 

584 lims[idim, :] = num.nan, num.nan 

585 else: 

586 lims[idim, :] = num.nan, num.nan 

587 

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

589 return lims 

590 

591 def set_lim(self, dim, vmin, vmax): 

592 assert vmin <= vmax 

593 self._view_limits[self._dim_index(dim), :] = vmin, vmax 

594 

595 def _get_mpl_view_limits(self): 

596 vl = [] 

597 for row in self._axes: 

598 vl_row = [] 

599 for axes in row: 

600 vl_row.append(( 

601 axes.get_xaxis().get_view_interval().tolist(), 

602 axes.get_yaxis().get_view_interval().tolist())) 

603 

604 vl.append(vl_row) 

605 

606 return vl 

607 

608 def _remember_mpl_view_limits(self): 

609 self._last_mpl_view_limits = self._get_mpl_view_limits() 

610 

611 def window_xmin(self, x): 

612 return window_min( 

613 self._shape[0], self._width, 

614 self.margins[0], self.margins[2] + self._colorbar_width, 

615 self.separator, x) 

616 

617 def window_xmax(self, x): 

618 return window_max( 

619 self._shape[0], self._width, 

620 self.margins[0], self.margins[2] + self._colorbar_width, 

621 self.separator, x) 

622 

623 def window_ymin(self, y): 

624 return window_min( 

625 self._shape[1], self._height, 

626 self.margins[3] + self._colorbar_height, self.margins[1], 

627 self.separator, y) 

628 

629 def window_ymax(self, y): 

630 return window_max( 

631 self._shape[1], self._height, 

632 self.margins[3] + self._colorbar_height, self.margins[1], 

633 self.separator, y) 

634 

635 def need_update_layout(self): 

636 self._need_update_layout = True 

637 

638 def _update_layout(self): 

639 assert not self._updating_layout 

640 

641 if not self._need_update_layout: 

642 return 

643 

644 self._updating_layout = True 

645 try: 

646 data_limits = self.get_data_limits() 

647 

648 limits = num.zeros((self._ndims, 2)) 

649 for idim in range(self._ndims): 

650 limits[idim, :] = self.make_limits(data_limits[idim, :]) 

651 

652 mask = num.isfinite(self._view_limits) 

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

654 

655 # deltas = limits[:, 1] - limits[:, 0] 

656 

657 # data_w = deltas[0] 

658 # data_h = deltas[1] 

659 

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

661 mr += self._colorbar_width 

662 mb += self._colorbar_height 

663 sw = sh = self.separator 

664 

665 nx, ny = self._shape 

666 

667 # data_r = data_h / data_w 

668 em = self.config.font_size 

669 em_pixels = em / self._pixels 

670 w = self._width 

671 h = self._height 

672 fig_w_avail = w - mr - ml - (nx-1) * sw 

673 fig_h_avail = h - mt - mb - (ny-1) * sh 

674 

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

676 raise NotEnoughSpace() 

677 

678 x_limits = num.zeros((nx, 2)) 

679 for ix, xdim in enumerate(self._x_dims): 

680 x_limits[ix, :] = limits[self._dim_index(xdim)] 

681 

682 y_limits = num.zeros((ny, 2)) 

683 for iy, ydim in enumerate(self._y_dims): 

684 y_limits[iy, :] = limits[self._dim_index(ydim)] 

685 

686 def get_aspect(dim1, dim2): 

687 if (dim2, dim1) in self._aspects: 

688 return 1.0/self._aspects[dim2, dim1] 

689 

690 return self._aspects.get((dim1, dim2), None) 

691 

692 aspects_xx = [] 

693 for ix1, xdim1 in enumerate(self._x_dims): 

694 for ix2, xdim2 in enumerate(self._x_dims): 

695 aspect = get_aspect(xdim2, xdim1) 

696 if aspect: 

697 aspects_xx.append((ix1, ix2, aspect)) 

698 

699 aspects_yy = [] 

700 for iy1, ydim1 in enumerate(self._y_dims): 

701 for iy2, ydim2 in enumerate(self._y_dims): 

702 aspect = get_aspect(ydim2, ydim1) 

703 if aspect: 

704 aspects_yy.append((iy1, iy2, aspect)) 

705 

706 aspects_xy = [] 

707 for iy, ix, axes in self.iaxes(): 

708 xdim = self._x_dims[ix] 

709 ydim = self._y_dims[iy] 

710 aspect = get_aspect(ydim, xdim) 

711 if aspect: 

712 aspects_xy.append((ix, iy, aspect)) 

713 

714 (x_limits, y_limits), (aws, ahs) = solve_layout_iterative( 

715 size=(fig_w_avail, fig_h_avail), 

716 shape=(nx, ny), 

717 limits=(x_limits, y_limits), 

718 aspects=( 

719 aspects_xx, 

720 aspects_yy, 

721 aspects_xy)) 

722 

723 for iy, ix, axes in self.iaxes(): 

724 rect = [ 

725 ml + num.sum(aws[:ix])+(ix*sw), 

726 mb + num.sum(ahs[:iy])+(iy*sh), 

727 aws[ix], ahs[iy]] 

728 

729 axes.set_position( 

730 self.rect_to_figure_coords(rect), which='both') 

731 

732 self.set_label_coords( 

733 axes, 'x', [ 

734 wcenter(rect), 

735 self.config.label_offset_em[0]*em_pixels 

736 + self._colorbar_height]) 

737 

738 self.set_label_coords( 

739 axes, 'y', [ 

740 self.config.label_offset_em[1]*em_pixels, 

741 hcenter(rect)]) 

742 

743 axes.get_xaxis().set_tick_params( 

744 bottom=(iy == 0), top=(iy == ny-1), 

745 labelbottom=(iy == 0), labeltop=False) 

746 

747 axes.get_yaxis().set_tick_params( 

748 left=(ix == 0), right=(ix == nx-1), 

749 labelleft=(ix == 0), labelright=False) 

750 

751 istride = -1 if self._x_dims_invert[ix] else 1 

752 axes.set_xlim(*x_limits[ix, ::istride]) 

753 istride = -1 if self._y_dims_invert[iy] else 1 

754 axes.set_ylim(*y_limits[iy, ::istride]) 

755 

756 axes.tick_params( 

757 axis='x', 

758 pad=self.config.tick_label_offset_em[0]*em) 

759 

760 axes.tick_params( 

761 axis='y', 

762 pad=self.config.tick_label_offset_em[0]*em) 

763 

764 self._remember_mpl_view_limits() 

765 

766 for mappable, dim in self._mappables.items(): 

767 mappable.set_clim(*limits[self._dim_index(dim)]) 

768 

769 # scaler = plot.AutoScaler() 

770 

771 # aspect tick incs same 

772 # 

773 # inc = scaler.make_scale( 

774 # [0, min(data_expanded_w, data_expanded_h)], 

775 # override_mode='off')[2] 

776 # 

777 # for axes in self.axes_list: 

778 # axes.set_xlim(*limits[0, :]) 

779 # axes.set_ylim(*limits[1, :]) 

780 # 

781 # tl = MultipleLocator(inc) 

782 # axes.get_xaxis().set_major_locator(tl) 

783 # tl = MultipleLocator(inc) 

784 # axes.get_yaxis().set_major_locator(tl) 

785 

786 for axes, orientation, position in self._colorbar_axes: 

787 if orientation == 'horizontal': 

788 xmin = self.window_xmin(position[0]) 

789 xmax = self.window_xmax(position[1]) 

790 ymin = mb - self._colorbar_height 

791 ymax = mb - self._colorbar_height \ 

792 + self.config.colorbar_width_em * em_pixels 

793 else: 

794 ymin = self.window_ymin(position[0]) 

795 ymax = self.window_ymax(position[1]) 

796 xmin = w - mr + 2 * sw 

797 xmax = w - mr + 2 * sw \ 

798 + self.config.colorbar_width_em * em_pixels 

799 

800 rect = [xmin, ymin, xmax-xmin, ymax-ymin] 

801 axes.set_position( 

802 self.rect_to_figure_coords(rect), which='both') 

803 

804 for ix, axes in enumerate(self.axes_bottom_list): 

805 dim = self._x_dims[ix] 

806 s = self._labels.get(dim, dim) 

807 axes.set_xlabel(s) 

808 

809 for iy, axes in enumerate(self.axes_left_list): 

810 dim = self._y_dims[iy] 

811 s = self._labels.get(dim, dim) 

812 axes.set_ylabel(s) 

813 

814 finally: 

815 self._updating_layout = False 

816 

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

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

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

820 

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

822 for iy, row in enumerate(self._axes): 

823 y = points[:, self._dim_index(self._y_dims[iy])] 

824 for ix, axes in enumerate(row): 

825 x = points[:, self._dim_index(self._x_dims[ix])] 

826 axes.plot(x, y, *args, **kwargs) 

827 

828 def close(self): 

829 self._disconnect_all() 

830 self._plt.close(self._fig) 

831 

832 def show(self): 

833 self._plt.show() 

834 self.reset_size() 

835 

836 def set_label(self, dim, s): 

837 # just set attribute, handle in update_layout 

838 self._labels[dim] = s 

839 

840 def colorbar( 

841 self, dim, 

842 orientation='vertical', 

843 position=None): 

844 

845 if dim not in self._dims: 

846 raise PlotError( 

847 'dimension "%s" is not defined') 

848 

849 if orientation not in ('vertical', 'horizontal'): 

850 raise PlotError( 

851 'orientation must be "vertical" or "horizontal"') 

852 

853 mappable = None 

854 for mappable_, dim_ in self._mappables.items(): 

855 if dim_ == dim: 

856 if mappable is None: 

857 mappable = mappable_ 

858 else: 

859 mappable_.set_cmap(mappable.get_cmap()) 

860 

861 if mappable is None: 

862 raise PlotError( 

863 'no mappable registered for dimension "%s"' % dim) 

864 

865 if position is None: 

866 if orientation == 'vertical': 

867 position = (0, self._shape[1]) 

868 else: 

869 position = (0, self._shape[0]) 

870 

871 em_pixels = self.config.font_size / self._pixels 

872 

873 if orientation == 'vertical': 

874 self._colorbar_width = self.config.colorbar_width_em*em_pixels + \ 

875 self.separator * 2.0 

876 else: 

877 self._colorbar_height = self.config.colorbar_width_em*em_pixels + \ 

878 self.separator + self.margins[3] 

879 

880 axes = SmartplotAxes(self.fig, [0., 0., 1., 1.]) 

881 self.fig.add_axes(axes) 

882 

883 self._colorbar_axes.append( 

884 (axes, orientation, position)) 

885 

886 self.need_update_layout() 

887 # axes.plot([1], [1]) 

888 label = self._labels.get(dim, dim) 

889 return colorbar.Colorbar( 

890 axes, mappable, orientation=orientation, label=label) 

891 

892 def __call__(self, *args): 

893 return self.axes(*args) 

894 

895 

896if __name__ == '__main__': 

897 import sys 

898 from pyrocko import util 

899 

900 logging.getLogger('matplotlib').setLevel(logging.WARNING) 

901 util.setup_logging('smartplot', 'debug') 

902 

903 iplots = [int(x) for x in sys.argv[1:]] 

904 

905 if 0 in iplots: 

906 p = Plot(['x'], ['y']) 

907 n = 100 

908 x = num.arange(n) * 2.0 

909 y = num.random.normal(size=n) 

910 p(0, 0).plot(x, y, 'o') 

911 p.show() 

912 

913 if 1 in iplots: 

914 p = Plot(['x', 'x'], ['y']) 

915 n = 100 

916 x = num.arange(n) * 2.0 

917 y = num.random.normal(size=n) 

918 p(0, 0).plot(x, y, 'o') 

919 x = num.arange(n) * 2.0 

920 y = num.random.normal(size=n) 

921 p(1, 0).plot(x, y, 'o') 

922 p.show() 

923 

924 if 11 in iplots: 

925 p = Plot(['x'], ['y']) 

926 p.set_aspect('y', 'x', 2.0) 

927 n = 100 

928 xy = num.random.normal(size=(n, 2)) 

929 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o') 

930 p.show() 

931 

932 if 12 in iplots: 

933 p = Plot(['x', 'x2'], ['y']) 

934 p.set_aspect('x2', 'x', 2.0) 

935 p.set_aspect('y', 'x', 2.0) 

936 n = 100 

937 xy = num.random.normal(size=(n, 2)) 

938 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o') 

939 p(1, 0).plot(xy[:, 0], xy[:, 1], 'o') 

940 p.show() 

941 

942 if 13 in iplots: 

943 p = Plot(['x'], ['y', 'y2']) 

944 p.set_aspect('y2', 'y', 2.0) 

945 p.set_aspect('y', 'x', 2.0) 

946 n = 100 

947 xy = num.random.normal(size=(n, 2)) 

948 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o') 

949 p(0, 1).plot(xy[:, 0], xy[:, 1], 'o') 

950 p.show() 

951 

952 if 2 in iplots: 

953 p = Plot(['easting', 'depth'], ['northing', 'depth']) 

954 

955 n = 100 

956 

957 ned = num.random.normal(size=(n, 3)) 

958 p(0, 0).plot(ned[:, 1], ned[:, 0], 'o') 

959 p(1, 0).plot(ned[:, 2], ned[:, 0], 'o') 

960 p(0, 1).plot(ned[:, 1], ned[:, 2], 'o') 

961 p.show() 

962 

963 if 3 in iplots: 

964 p = Plot(['easting', 'depth'], ['-depth', 'northing']) 

965 p.set_aspect('easting', 'northing', 1.0) 

966 p.set_aspect('easting', 'depth', 0.5) 

967 p.set_aspect('northing', 'depth', 0.5) 

968 

969 n = 100 

970 

971 ned = num.random.normal(size=(n, 3)) 

972 ned[:, 2] *= 0.25 

973 p(0, 1).plot(ned[:, 1], ned[:, 0], 'o', color='black') 

974 p(0, 0).plot(ned[:, 1], ned[:, 2], 'o') 

975 p(1, 1).plot(ned[:, 2], ned[:, 0], 'o') 

976 p(1, 0).set_visible(False) 

977 p.set_lim('depth', 0., 0.2) 

978 p.show() 

979 

980 if 5 in iplots: 

981 p = Plot(['time'], ['northing', 'easting', '-depth'], ['depth']) 

982 

983 n = 100 

984 

985 t = num.arange(n) 

986 xyz = num.random.normal(size=(n, 4)) 

987 xyz[:, 0] *= 0.5 

988 

989 smap = make_smap('summer') 

990 

991 p(0, 0).scatter( 

992 t, xyz[:, 0], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm) 

993 p(0, 1).scatter( 

994 t, xyz[:, 1], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm) 

995 p(0, 2).scatter( 

996 t, xyz[:, 2], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm) 

997 

998 p.set_lim('depth', -1., 1.) 

999 

1000 p.set_color_dim(smap, 'depth') 

1001 

1002 p.set_aspect('northing', 'easting', 1.0) 

1003 p.set_aspect('northing', 'depth', 1.0) 

1004 

1005 p.set_label('time', 'Time [s]') 

1006 p.set_label('depth', 'Depth [km]') 

1007 p.set_label('easting', 'Easting [km]') 

1008 p.set_label('northing', 'Northing [km]') 

1009 

1010 p.colorbar('depth') 

1011 

1012 p.show() 

1013 

1014 if 6 in iplots: 

1015 km = 1000. 

1016 p = Plot( 

1017 ['easting'], ['northing']*3, ['displacement']) 

1018 

1019 nn, ne = 50, 40 

1020 n = num.linspace(-5*km, 5*km, nn) 

1021 e = num.linspace(-10*km, 10*km, ne) 

1022 

1023 displacement = num.zeros((nn, ne, 3)) 

1024 g = num.exp( 

1025 -(n[:, num.newaxis]**2 + e[num.newaxis, :]**2) / (5*km)**2) 

1026 

1027 displacement[:, :, 0] = g 

1028 displacement[:, :, 1] = g * 0.5 

1029 displacement[:, :, 2] = -g * 0.2 

1030 

1031 for icomp in (0, 1, 2): 

1032 c = p(0, icomp).pcolormesh( 

1033 e/km, n/km, displacement[:, :, icomp], shading='gouraud') 

1034 p.set_color_dim(c, 'displacement') 

1035 

1036 p.colorbar('displacement') 

1037 p.set_lim('displacement', -1.0, 1.0) 

1038 p.set_label('easting', 'Easting [km]') 

1039 p.set_label('northing', 'Northing [km]') 

1040 p.set_aspect('northing', 'easting') 

1041 

1042 p.set_lim('northing', -5.0, 5.0) 

1043 p.set_lim('easting', -3.0, 3.0) 

1044 p.show()