1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

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

5 

6from __future__ import print_function 

7 

8from collections import defaultdict 

9import math 

10import logging 

11 

12import numpy as num 

13from matplotlib.axes import Axes 

14# from matplotlib.ticker import MultipleLocator 

15from matplotlib import cm, colors, colorbar, figure 

16 

17from pyrocko.guts import Tuple, Float, Object 

18from pyrocko import plot 

19 

20import scipy.optimize 

21 

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

23 

24guts_prefix = 'pf' 

25 

26inch = 2.54 

27 

28 

29class SmartPlotFigure(figure.Figure): 

30 

31 def set_smartplot(self, plot): 

32 self._smartplot = plot 

33 

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

35 if hasattr(self, '_smartplot'): 

36 try: 

37 self._smartplot._update_layout() 

38 except NotEnoughSpace: 

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

40 return 

41 

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

43 

44 

45def limits(points): 

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

47 if points.size != 0: 

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

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

50 

51 return lims 

52 

53 

54def wcenter(rect): 

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

56 

57 

58def hcenter(rect): 

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

60 

61 

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

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

64 

65 

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

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

68 

69 

70def make_smap(cmap, norm=None): 

71 if isinstance(norm, tuple): 

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

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

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

75 return smap 

76 

77 

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

79 

80 weight_aspect = 1000. 

81 

82 sx, sy = size 

83 nx, ny = shape 

84 nvar = nx+ny 

85 vxs, vys = limits 

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

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

88 aspects_xx, aspects_yy, aspects_xy = aspects 

89 

90 if fracs is None: 

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

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

93 else: 

94 frac_x, frac_y = fracs 

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

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

97 

98 data = [] 

99 weights = [] 

100 rows = [] 

101 bounds = [] 

102 for ix in range(nx): 

103 u = uxs[ix] 

104 assert u > 0.0 

105 row = num.zeros(nvar) 

106 row[ix] = u 

107 rows.append(row) 

108 data.append(wxs[ix]) 

109 weights.append(1.0 / u) 

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

111 

112 for iy in range(ny): 

113 u = uys[iy] 

114 assert u > 0.0 

115 row = num.zeros(nvar) 

116 row[nx+iy] = u 

117 rows.append(row) 

118 data.append(wys[iy]) 

119 weights.append(1.0) 

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

121 

122 for ix1, ix2, aspect in aspects_xx: 

123 row = num.zeros(nvar) 

124 row[ix1] = aspect 

125 row[ix2] = -1.0 

126 weights.append(weight_aspect/aspect) 

127 rows.append(row) 

128 data.append(0.0) 

129 

130 for iy1, iy2, aspect in aspects_yy: 

131 row = num.zeros(nvar) 

132 row[nx+iy1] = aspect 

133 row[nx+iy2] = -1.0 

134 weights.append(weight_aspect/aspect) 

135 rows.append(row) 

136 data.append(0.0) 

137 

138 for ix, iy, aspect in aspects_xy: 

139 row = num.zeros(nvar) 

140 row[ix] = aspect 

141 row[nx+iy] = -1.0 

142 weights.append(weight_aspect/aspect) 

143 rows.append(row) 

144 data.append(0.0) 

145 

146 weights = num.array(weights) 

147 data = num.array(data) 

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

149 data *= weights 

150 

151 bounds = num.array(bounds).T 

152 

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

154 

155 cxs = model[:nx] 

156 cys = model[nx:nx+ny] 

157 

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

159 for ix in range(nx): 

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

161 vmin, vmax = vxs[ix] 

162 udata = vmax - vmin 

163 eps = 1e-7 * u 

164 assert(udata <= u + eps) 

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

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

167 

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

169 for iy in range(ny): 

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

171 vmin, vmax = vys[iy] 

172 udata = vmax - vmin 

173 eps = 1e-7 * u 

174 assert(udata <= u + eps) 

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

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

177 

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

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

180 logger.error( 

181 'Unable to comply with requested aspect ratio ' 

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

183 

184 for ix1, ix2, aspect in aspects_xx: 

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

186 

187 for iy1, iy2, aspect in aspects_yy: 

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

189 

190 for ix, iy, aspect in aspects_xy: 

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

192 

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

194 

195 

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

197 

198 sx, sy = size 

199 nx, ny = shape 

200 vxs, vys = limits 

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

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

203 aspects_xx, aspects_yy, aspects_xy = aspects 

204 

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

206 for i in range(niterations): 

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

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

209 

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

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

212 wxs_used = wxs * uxs / uxs_view 

213 wys_used = wys * uys / uys_view 

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

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

216 

217 fracs_x = wxs_used 

218 fracs_y = wys_used 

219 

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

221 

222 

223class PlotError(Exception): 

224 pass 

225 

226 

227class NotEnoughSpace(PlotError): 

228 pass 

229 

230 

231class PlotConfig(Object): 

232 

233 font_size = Float.T(default=9.0) 

234 

235 size_cm = Tuple.T( 

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

237 

238 margins_em = Tuple.T( 

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

240 

241 separator_em = Float.T(default=1.5) 

242 

243 colorbar_width_em = Float.T(default=2.0) 

244 

245 label_offset_em = Tuple.T( 

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

247 

248 @property 

249 def size_inch(self): 

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

251 

252 

253class Plot(object): 

254 

255 def __init__( 

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

257 fig=None, call_mpl_init=True): 

258 

259 if config is None: 

260 config = PlotConfig() 

261 

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

263 

264 dims = [] 

265 for dim in x_dims + y_dims + z_dims: 

266 dim = dim.lstrip('-') 

267 if dim not in dims: 

268 dims.append(dim) 

269 

270 self.config = config 

271 self._disconnect_data = [] 

272 self._width = self._height = self._pixels = None 

273 if call_mpl_init: 

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

275 

276 if fig is None: 

277 fig = self._plt.figure( 

278 figsize=self.config.size_inch, Figure=SmartPlotFigure) 

279 else: 

280 assert isinstance(fig, SmartPlotFigure) 

281 

282 fig.set_smartplot(self) 

283 

284 self._fig = fig 

285 self._colorbar_width = 0.0 

286 self._colorbar_height = 0.0 

287 self._colorbar_axes = [] 

288 

289 self._dims = dims 

290 self._dim_index = self._dims.index 

291 self._ndims = len(dims) 

292 self._labels = {} 

293 self._aspects = {} 

294 

295 self.setup_axes() 

296 

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

298 

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

300 self._last_mpl_view_limits = None 

301 

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

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

304 

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

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

307 

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

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

310 

311 self._mappables = {} 

312 self._updating_layout = False 

313 

314 self._need_update_layout = True 

315 self._update_geometry() 

316 

317 for axes in self.axes_list: 

318 fig.add_axes(axes) 

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

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

321 

322 self._cid_resize = fig.canvas.mpl_connect( 

323 'resize_event', self.resize_handler) 

324 

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

326 

327 self._lim_changed_depth = 0 

328 

329 def axes(self, ix, iy): 

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

331 ix = self._x_dims.index(ix) 

332 iy = self._y_dims.index(iy) 

333 

334 return self._axes[iy][ix] 

335 

336 def set_color_dim(self, mappable, dim): 

337 assert dim in self._dims 

338 self._mappables[mappable] = dim 

339 

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

341 self._aspects[ydim, xdim] = aspect 

342 

343 @property 

344 def dims(self): 

345 return self._dims 

346 

347 @property 

348 def fig(self): 

349 return self._fig 

350 

351 @property 

352 def axes_list(self): 

353 axes = [] 

354 for row in self._axes: 

355 axes.extend(row) 

356 return axes 

357 

358 @property 

359 def axes_bottom_list(self): 

360 return self._axes[0] 

361 

362 @property 

363 def axes_left_list(self): 

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

365 

366 def setup_axes(self): 

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

368 nx, ny = self._shape 

369 axes = [] 

370 for iy in range(ny): 

371 axes.append([]) 

372 for ix in range(nx): 

373 axes[-1].append(Axes(self.fig, rect)) 

374 

375 self._axes = axes 

376 

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

378 axes_.set_autoscale_on(False) 

379 

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

381 cid = obj.callbacks.connect(sig, handler) 

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

383 

384 def _disconnect_all(self): 

385 for obj, cid in self._disconnect_data: 

386 obj.callbacks.disconnect(cid) 

387 

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

389 

390 def dpi_changed_handler(self, fig): 

391 if self._updating_layout: 

392 return 

393 

394 self._update_geometry() 

395 

396 def resize_handler(self, event): 

397 if self._updating_layout: 

398 return 

399 

400 self._update_geometry() 

401 

402 def lim_changed_handler(self, axes): 

403 if self._updating_layout: 

404 return 

405 

406 current = self._get_mpl_view_limits() 

407 last = self._last_mpl_view_limits 

408 

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

410 acurrent = current[iy][ix] 

411 alast = last[iy][ix] 

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

413 xdim = self._x_dims[ix] 

414 logger.debug( 

415 'X limits have been changed interactively in subplot ' 

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

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

418 

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

420 ydim = self._y_dims[iy] 

421 logger.debug( 

422 'Y limits have been changed interactively in subplot ' 

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

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

425 

426 self.need_update_layout() 

427 

428 def _update_geometry(self): 

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

430 p = self.get_pixels_factor() 

431 

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

433 self._width = w 

434 self._height = h 

435 self._pixels = p 

436 self.need_update_layout() 

437 

438 @property 

439 def margins(self): 

440 return tuple( 

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

442 for x in self.config.margins_em) 

443 

444 @property 

445 def separator(self): 

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

447 

448 def rect_to_figure_coords(self, rect): 

449 left, bottom, width, height = rect 

450 return ( 

451 left / self._width, 

452 bottom / self._height, 

453 width / self._width, 

454 height / self._height) 

455 

456 def point_to_axes_coords(self, axes, point): 

457 x, y = point 

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

459 

460 x_fig = x / self._width 

461 y_fig = y / self._height 

462 

463 x_axes = (x_fig - aleft) / awidth 

464 y_axes = (y_fig - abottom) / aheight 

465 

466 return (x_axes, y_axes) 

467 

468 def get_pixels_factor(self): 

469 try: 

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

471 return 1.0 / r.points_to_pixels(1.0) 

472 except AttributeError: 

473 return 1.0 

474 

475 def make_limits(self, lims): 

476 a = plot.AutoScaler(space=0.05) 

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

478 

479 def iaxes(self): 

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

481 for ix, axes in enumerate(row): 

482 yield iy, ix, axes 

483 

484 def get_data_limits(self): 

485 dim_to_values = defaultdict(list) 

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

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

488 axes.get_yaxis().get_data_interval()) 

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

490 axes.get_xaxis().get_data_interval()) 

491 

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

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

494 

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

496 for idim in range(self._ndims): 

497 dim = self._dims[idim] 

498 if dim in dim_to_values: 

499 vs = num.array( 

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

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

502 if vs.size > 0: 

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

504 else: 

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

506 else: 

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

508 

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

510 return lims 

511 

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

513 assert(vmin <= vmax) 

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

515 

516 def _get_mpl_view_limits(self): 

517 vl = [] 

518 for row in self._axes: 

519 vl_row = [] 

520 for axes in row: 

521 vl_row.append(( 

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

523 axes.get_yaxis().get_view_interval().tolist())) 

524 

525 vl.append(vl_row) 

526 

527 return vl 

528 

529 def _remember_mpl_view_limits(self): 

530 self._last_mpl_view_limits = self._get_mpl_view_limits() 

531 

532 def window_xmin(self, x): 

533 return window_min( 

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

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

536 self.separator, x) 

537 

538 def window_xmax(self, x): 

539 return window_max( 

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

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

542 self.separator, x) 

543 

544 def window_ymin(self, y): 

545 return window_min( 

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

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

548 self.separator, y) 

549 

550 def window_ymax(self, y): 

551 return window_max( 

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

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

554 self.separator, y) 

555 

556 def need_update_layout(self): 

557 self._need_update_layout = True 

558 

559 def _update_layout(self): 

560 assert not self._updating_layout 

561 

562 if not self._need_update_layout: 

563 return 

564 

565 self._updating_layout = True 

566 try: 

567 data_limits = self.get_data_limits() 

568 

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

570 for idim in range(self._ndims): 

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

572 

573 mask = num.isfinite(self._view_limits) 

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

575 

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

577 

578 # data_w = deltas[0] 

579 # data_h = deltas[1] 

580 

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

582 mr += self._colorbar_width 

583 mb += self._colorbar_height 

584 sw = sh = self.separator 

585 

586 nx, ny = self._shape 

587 

588 # data_r = data_h / data_w 

589 em = self.config.font_size 

590 w = self._width 

591 h = self._height 

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

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

594 

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

596 raise NotEnoughSpace() 

597 

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

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

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

601 

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

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

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

605 

606 def get_aspect(dim1, dim2): 

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

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

609 

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

611 

612 aspects_xx = [] 

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

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

615 aspect = get_aspect(xdim2, xdim1) 

616 if aspect: 

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

618 

619 aspects_yy = [] 

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

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

622 aspect = get_aspect(ydim2, ydim1) 

623 if aspect: 

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

625 

626 aspects_xy = [] 

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

628 xdim = self._x_dims[ix] 

629 ydim = self._y_dims[iy] 

630 aspect = get_aspect(ydim, xdim) 

631 if aspect: 

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

633 

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

635 size=(fig_w_avail, fig_h_avail), 

636 shape=(nx, ny), 

637 limits=(x_limits, y_limits), 

638 aspects=( 

639 aspects_xx, 

640 aspects_yy, 

641 aspects_xy)) 

642 

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

644 rect = [ 

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

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

647 aws[ix], ahs[iy]] 

648 

649 axes.set_position( 

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

651 

652 self.set_label_coords( 

653 axes, 'x', [ 

654 wcenter(rect), 

655 self.config.label_offset_em[0]*em 

656 + self._colorbar_height]) 

657 

658 self.set_label_coords( 

659 axes, 'y', [ 

660 self.config.label_offset_em[1]*em, 

661 hcenter(rect)]) 

662 

663 axes.get_xaxis().set_tick_params( 

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

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

666 

667 axes.get_yaxis().set_tick_params( 

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

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

670 

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

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

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

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

675 

676 self._remember_mpl_view_limits() 

677 

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

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

680 

681 # scaler = plot.AutoScaler() 

682 

683 # aspect tick incs same 

684 # 

685 # inc = scaler.make_scale( 

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

687 # override_mode='off')[2] 

688 # 

689 # for axes in self.axes_list: 

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

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

692 # 

693 # tl = MultipleLocator(inc) 

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

695 # tl = MultipleLocator(inc) 

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

697 

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

699 if orientation == 'horizontal': 

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

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

702 ymin = mb - self._colorbar_height 

703 ymax = mb - self._colorbar_height \ 

704 + self.config.colorbar_width_em * em 

705 else: 

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

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

708 xmin = w - mr + 2 * sw 

709 xmax = w - mr + 2 * sw + self.config.colorbar_width_em * em 

710 

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

712 axes.set_position( 

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

714 

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

716 dim = self._x_dims[ix] 

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

718 axes.set_xlabel(s) 

719 

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

721 dim = self._y_dims[iy] 

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

723 axes.set_ylabel(s) 

724 

725 finally: 

726 self._updating_layout = False 

727 

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

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

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

731 

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

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

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

735 for ix, axes in enumerate(row): 

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

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

738 

739 def close(self): 

740 self._disconnect_all() 

741 self._plt.close(self._fig) 

742 

743 def show(self): 

744 self._plt.show() 

745 

746 def set_label(self, dim, s): 

747 # just set attribute, handle in update_layout 

748 self._labels[dim] = s 

749 

750 def colorbar( 

751 self, dim, 

752 orientation='vertical', 

753 position=None): 

754 

755 if dim not in self._dims: 

756 raise PlotError( 

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

758 

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

760 raise PlotError( 

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

762 

763 mappable = None 

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

765 if dim_ == dim: 

766 if mappable is None: 

767 mappable = mappable_ 

768 else: 

769 mappable_.set_cmap(mappable.get_cmap()) 

770 

771 if mappable is None: 

772 raise PlotError( 

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

774 

775 if position is None: 

776 if orientation == 'vertical': 

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

778 else: 

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

780 

781 em = self.config.font_size 

782 

783 if orientation == 'vertical': 

784 self._colorbar_width = self.config.colorbar_width_em*em + \ 

785 self.separator * 2.0 

786 else: 

787 self._colorbar_height = self.config.colorbar_width_em*em + \ 

788 self.separator + self.margins[3] 

789 

790 axes = Axes(self.fig, [0., 0., 1., 1.]) 

791 self.fig.add_axes(axes) 

792 

793 self._colorbar_axes.append( 

794 (axes, orientation, position)) 

795 

796 self.need_update_layout() 

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

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

799 return colorbar.Colorbar( 

800 axes, mappable, orientation=orientation, label=label) 

801 

802 def __call__(self, *args): 

803 return self.axes(*args) 

804 

805 

806if __name__ == '__main__': 

807 import sys 

808 from pyrocko import util 

809 

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

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

812 

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

814 

815 if 0 in iplots: 

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

817 n = 100 

818 x = num.arange(n) * 2.0 

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

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

821 p.show() 

822 

823 if 1 in iplots: 

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

825 n = 100 

826 x = num.arange(n) * 2.0 

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

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

829 x = num.arange(n) * 2.0 

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

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

832 p.show() 

833 

834 if 11 in iplots: 

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

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

837 n = 100 

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

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

840 p.show() 

841 

842 if 12 in iplots: 

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

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

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

846 n = 100 

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

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

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

850 p.show() 

851 

852 if 13 in iplots: 

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

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

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

856 n = 100 

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

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

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

860 p.show() 

861 

862 if 2 in iplots: 

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

864 

865 n = 100 

866 

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

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

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

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

871 p.show() 

872 

873 if 3 in iplots: 

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

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

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

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

878 

879 n = 100 

880 

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

882 ned[:, 2] *= 0.25 

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

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

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

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

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

888 p.show() 

889 

890 if 5 in iplots: 

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

892 

893 n = 100 

894 

895 t = num.arange(n) 

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

897 xyz[:, 0] *= 0.5 

898 

899 smap = make_smap('summer') 

900 

901 p(0, 0).scatter( 

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

903 p(0, 1).scatter( 

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

905 p(0, 2).scatter( 

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

907 

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

909 

910 p.set_color_dim(smap, 'depth') 

911 

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

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

914 

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

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

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

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

919 

920 p.colorbar('depth') 

921 

922 p.show() 

923 

924 if 6 in iplots: 

925 km = 1000. 

926 p = Plot( 

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

928 

929 nn, ne = 50, 40 

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

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

932 

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

934 g = num.exp( 

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

936 

937 displacement[:, :, 0] = g 

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

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

940 

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

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

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

944 p.set_color_dim(c, 'displacement') 

945 

946 p.colorbar('displacement') 

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

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

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

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

951 

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

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

954 p.show()