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 

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 

29def limits(points): 

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

31 if points.size != 0: 

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

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

34 

35 return lims 

36 

37 

38def wcenter(rect): 

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

40 

41 

42def hcenter(rect): 

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

44 

45 

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

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

48 

49 

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

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

52 

53 

54def make_smap(cmap, norm=None): 

55 if isinstance(norm, tuple): 

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

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

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

59 return smap 

60 

61 

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

63 

64 weight_aspect = 1000. 

65 

66 sx, sy = size 

67 nx, ny = shape 

68 nvar = nx+ny 

69 vxs, vys = limits 

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

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

72 aspects_xx, aspects_yy, aspects_xy = aspects 

73 

74 if fracs is None: 

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

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

77 else: 

78 frac_x, frac_y = fracs 

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

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

81 

82 data = [] 

83 weights = [] 

84 rows = [] 

85 bounds = [] 

86 for ix in range(nx): 

87 u = uxs[ix] 

88 assert u > 0.0 

89 row = num.zeros(nvar) 

90 row[ix] = u 

91 rows.append(row) 

92 data.append(wxs[ix]) 

93 weights.append(1.0 / u) 

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

95 

96 for iy in range(ny): 

97 u = uys[iy] 

98 assert u > 0.0 

99 row = num.zeros(nvar) 

100 row[nx+iy] = u 

101 rows.append(row) 

102 data.append(wys[iy]) 

103 weights.append(1.0) 

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

105 

106 for ix1, ix2, aspect in aspects_xx: 

107 row = num.zeros(nvar) 

108 row[ix1] = aspect 

109 row[ix2] = -1.0 

110 weights.append(weight_aspect/aspect) 

111 rows.append(row) 

112 data.append(0.0) 

113 

114 for iy1, iy2, aspect in aspects_yy: 

115 row = num.zeros(nvar) 

116 row[nx+iy1] = aspect 

117 row[nx+iy2] = -1.0 

118 weights.append(weight_aspect/aspect) 

119 rows.append(row) 

120 data.append(0.0) 

121 

122 for ix, iy, aspect in aspects_xy: 

123 row = num.zeros(nvar) 

124 row[ix] = aspect 

125 row[nx+iy] = -1.0 

126 weights.append(weight_aspect/aspect) 

127 rows.append(row) 

128 data.append(0.0) 

129 

130 weights = num.array(weights) 

131 data = num.array(data) 

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

133 data *= weights 

134 

135 bounds = num.array(bounds).T 

136 

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

138 

139 cxs = model[:nx] 

140 cys = model[nx:nx+ny] 

141 

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

143 for ix in range(nx): 

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

145 vmin, vmax = vxs[ix] 

146 udata = vmax - vmin 

147 eps = 1e-7 * u 

148 assert(udata <= u + eps) 

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

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

151 

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

153 for iy in range(ny): 

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

155 vmin, vmax = vys[iy] 

156 udata = vmax - vmin 

157 eps = 1e-7 * u 

158 assert(udata <= u + eps) 

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

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

161 

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

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

164 logger.error( 

165 'Unable to comply with requested aspect ratio ' 

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

167 

168 for ix1, ix2, aspect in aspects_xx: 

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

170 

171 for iy1, iy2, aspect in aspects_yy: 

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

173 

174 for ix, iy, aspect in aspects_xy: 

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

176 

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

178 

179 

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

181 

182 sx, sy = size 

183 nx, ny = shape 

184 vxs, vys = limits 

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

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

187 aspects_xx, aspects_yy, aspects_xy = aspects 

188 

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

190 for i in range(niterations): 

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

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

193 

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

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

196 wxs_used = wxs * uxs / uxs_view 

197 wys_used = wys * uys / uys_view 

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

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

200 

201 fracs_x = wxs_used 

202 fracs_y = wys_used 

203 

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

205 

206 

207class NotEnoughSpace(Exception): 

208 pass 

209 

210 

211class PlotConfig(Object): 

212 

213 font_size = Float.T(default=9.0) 

214 

215 size_cm = Tuple.T( 

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

217 

218 margins_em = Tuple.T( 

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

220 

221 separator_em = Float.T(default=1.5) 

222 

223 colorbar_width_em = Float.T(default=2.0) 

224 

225 label_offset_em = Tuple.T( 

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

227 

228 @property 

229 def size_inch(self): 

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

231 

232 

233class Plot(object): 

234 

235 def __init__( 

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

237 fig=None): 

238 

239 if config is None: 

240 config = PlotConfig() 

241 

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

243 

244 dims = [] 

245 for dim in x_dims + y_dims + z_dims: 

246 dim = dim.lstrip('-') 

247 if dim not in dims: 

248 dims.append(dim) 

249 

250 self.config = config 

251 self._disconnect_data = [] 

252 self._width = self._height = self._pixels = None 

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

254 

255 if fig is None: 

256 fig = self._plt.figure(figsize=self.config.size_inch) 

257 

258 self._fig = fig 

259 self._colorbar_width = 0.0 

260 self._colorbar_height = 0.0 

261 self._colorbar_axes = [] 

262 

263 self._dims = dims 

264 self._dim_index = self._dims.index 

265 self._ndims = len(dims) 

266 self._labels = {} 

267 self._aspects = {} 

268 

269 self.setup_axes() 

270 

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

272 

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

274 self._last_mpl_view_limits = None 

275 

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

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

278 

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

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

281 

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

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

284 

285 self._mappables = {} 

286 self._updating_layout = False 

287 

288 self._update_geometry() 

289 

290 for axes in self.axes_list: 

291 fig.add_axes(axes) 

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

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

294 

295 self._cid_resize = fig.canvas.mpl_connect( 

296 'resize_event', self.resize_handler) 

297 

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

299 

300 self._lim_changed_depth = 0 

301 

302 def axes(self, ix, iy): 

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

304 ix = self._x_dims.index(ix) 

305 iy = self._y_dims.index(iy) 

306 

307 return self._axes[iy][ix] 

308 

309 def set_color_dim(self, mappable, dim): 

310 assert dim in self._dims 

311 self._mappables[mappable] = dim 

312 

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

314 self._aspects[ydim, xdim] = aspect 

315 

316 @property 

317 def dims(self): 

318 return self._dims 

319 

320 @property 

321 def fig(self): 

322 return self._fig 

323 

324 @property 

325 def axes_list(self): 

326 axes = [] 

327 for row in self._axes: 

328 axes.extend(row) 

329 return axes 

330 

331 @property 

332 def axes_bottom_list(self): 

333 return self._axes[0] 

334 

335 @property 

336 def axes_left_list(self): 

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

338 

339 def setup_axes(self): 

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

341 nx, ny = self._shape 

342 axes = [] 

343 for iy in range(ny): 

344 axes.append([]) 

345 for ix in range(nx): 

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

347 

348 self._axes = axes 

349 

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

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

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

353 

354 def _disconnect_all(self): 

355 for obj, cid in self._disconnect_data: 

356 obj.callbacks.disconnect(cid) 

357 

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

359 

360 def dpi_changed_handler(self, fig): 

361 if self._updating_layout: 

362 return 

363 

364 self._update_geometry() 

365 

366 def resize_handler(self, event): 

367 if self._updating_layout: 

368 return 

369 

370 self._update_geometry() 

371 

372 def lim_changed_handler(self, axes): 

373 if self._updating_layout: 

374 return 

375 

376 current = self._get_mpl_view_limits() 

377 last = self._last_mpl_view_limits 

378 

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

380 acurrent = current[iy][ix] 

381 alast = last[iy][ix] 

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

383 xdim = self._x_dims[ix] 

384 logger.debug( 

385 'X limits have been changed interactively in subplot ' 

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

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

388 

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

390 ydim = self._y_dims[iy] 

391 logger.debug( 

392 'Y limits have been changed interactively in subplot ' 

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

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

395 

396 self._update_layout() 

397 

398 def _update_geometry(self): 

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

400 p = self.get_pixels_factor() 

401 

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

403 self._width = w 

404 self._height = h 

405 self._pixels = p 

406 self._update_layout() 

407 

408 @property 

409 def margins(self): 

410 return tuple( 

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

412 for x in self.config.margins_em) 

413 

414 @property 

415 def separator(self): 

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

417 

418 def rect_to_figure_coords(self, rect): 

419 left, bottom, width, height = rect 

420 return ( 

421 left / self._width, 

422 bottom / self._height, 

423 width / self._width, 

424 height / self._height) 

425 

426 def point_to_axes_coords(self, axes, point): 

427 x, y = point 

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

429 

430 x_fig = x / self._width 

431 y_fig = y / self._height 

432 

433 x_axes = (x_fig - aleft) / awidth 

434 y_axes = (y_fig - abottom) / aheight 

435 

436 return (x_axes, y_axes) 

437 

438 def get_pixels_factor(self): 

439 try: 

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

441 return 1.0 / r.points_to_pixels(1.0) 

442 except AttributeError: 

443 return 1.0 

444 

445 def make_limits(self, lims): 

446 a = plot.AutoScaler(space=0.05) 

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

448 

449 def iaxes(self): 

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

451 for ix, axes in enumerate(row): 

452 yield iy, ix, axes 

453 

454 def get_data_limits(self): 

455 dim_to_values = defaultdict(list) 

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

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

458 axes.get_yaxis().get_data_interval()) 

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

460 axes.get_xaxis().get_data_interval()) 

461 

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

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

464 

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

466 for idim in range(self._ndims): 

467 dim = self._dims[idim] 

468 if dim in dim_to_values: 

469 vs = num.array( 

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

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

472 if vs.size > 0: 

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

474 else: 

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

476 else: 

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

478 

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

480 return lims 

481 

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

483 assert(vmin <= vmax) 

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

485 

486 def _get_mpl_view_limits(self): 

487 vl = [] 

488 for row in self._axes: 

489 vl_row = [] 

490 for axes in row: 

491 vl_row.append(( 

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

493 axes.get_yaxis().get_view_interval().tolist())) 

494 

495 vl.append(vl_row) 

496 

497 return vl 

498 

499 def _remember_mpl_view_limits(self): 

500 self._last_mpl_view_limits = self._get_mpl_view_limits() 

501 

502 def window_xmin(self, x): 

503 return window_min( 

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

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

506 self.separator, x) 

507 

508 def window_xmax(self, x): 

509 return window_max( 

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

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

512 self.separator, x) 

513 

514 def window_ymin(self, y): 

515 return window_min( 

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

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

518 self.separator, y) 

519 

520 def window_ymax(self, y): 

521 return window_max( 

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

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

524 self.separator, y) 

525 

526 def _update_layout(self): 

527 assert not self._updating_layout 

528 self._updating_layout = True 

529 data_limits = self.get_data_limits() 

530 

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

532 for idim in range(self._ndims): 

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

534 

535 mask = num.isfinite(self._view_limits) 

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

537 

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

539 

540 # data_w = deltas[0] 

541 # data_h = deltas[1] 

542 

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

544 mr += self._colorbar_width 

545 mb += self._colorbar_height 

546 sw = sh = self.separator 

547 

548 nx, ny = self._shape 

549 

550 # data_r = data_h / data_w 

551 em = self.config.font_size 

552 w = self._width 

553 h = self._height 

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

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

556 

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

558 raise NotEnoughSpace() 

559 

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

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

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

563 

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

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

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

567 

568 def get_aspect(dim1, dim2): 

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

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

571 

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

573 

574 aspects_xx = [] 

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

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

577 aspect = get_aspect(xdim2, xdim1) 

578 if aspect: 

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

580 

581 aspects_yy = [] 

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

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

584 aspect = get_aspect(ydim2, ydim1) 

585 if aspect: 

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

587 

588 aspects_xy = [] 

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

590 xdim = self._x_dims[ix] 

591 ydim = self._y_dims[iy] 

592 aspect = get_aspect(ydim, xdim) 

593 if aspect: 

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

595 

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

597 size=(fig_w_avail, fig_h_avail), 

598 shape=(nx, ny), 

599 limits=(x_limits, y_limits), 

600 aspects=( 

601 aspects_xx, 

602 aspects_yy, 

603 aspects_xy)) 

604 

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

606 rect = [ 

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

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

609 aws[ix], ahs[iy]] 

610 

611 axes.set_position( 

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

613 

614 self.set_label_coords( 

615 axes, 'x', [ 

616 wcenter(rect), 

617 self.config.label_offset_em[0]*em + self._colorbar_height]) 

618 

619 self.set_label_coords( 

620 axes, 'y', [ 

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

622 hcenter(rect)]) 

623 

624 axes.get_xaxis().set_tick_params( 

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

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

627 

628 axes.get_yaxis().set_tick_params( 

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

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

631 

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

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

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

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

636 

637 self._remember_mpl_view_limits() 

638 

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

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

641 

642 # scaler = plot.AutoScaler() 

643 

644 # aspect tick incs same 

645 # 

646 # inc = scaler.make_scale( 

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

648 # override_mode='off')[2] 

649 # 

650 # for axes in self.axes_list: 

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

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

653 # 

654 # tl = MultipleLocator(inc) 

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

656 # tl = MultipleLocator(inc) 

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

658 

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

660 if orientation == 'horizontal': 

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

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

663 ymin = mb - self._colorbar_height 

664 ymax = mb - self._colorbar_height \ 

665 + self.config.colorbar_width_em * em 

666 else: 

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

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

669 xmin = w - mr + 2 * sw 

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

671 

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

673 axes.set_position( 

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

675 

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

677 dim = self._x_dims[ix] 

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

679 axes.set_xlabel(s) 

680 

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

682 dim = self._y_dims[iy] 

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

684 axes.set_ylabel(s) 

685 

686 self._updating_layout = False 

687 

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

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

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

691 

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

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

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

695 for ix, axes in enumerate(row): 

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

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

698 

699 def close(self): 

700 self._disconnect_all() 

701 self._plt.close(self._fig) 

702 

703 def show(self): 

704 self._plt.show() 

705 

706 def set_label(self, dim, s): 

707 # just set attrbitute handle in update_layout 

708 self._labels[dim] = s 

709 

710 def colorbar( 

711 self, dim, 

712 orientation='vertical', 

713 position=None): 

714 

715 if dim not in self._dims: 

716 raise Exception( 

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

718 

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

720 raise Exception( 

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

722 

723 mappable = None 

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

725 if dim_ == dim: 

726 if mappable is None: 

727 mappable = mappable_ 

728 else: 

729 mappable_.set_cmap(mappable.get_cmap()) 

730 

731 if mappable is None: 

732 raise Exception( 

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

734 

735 if position is None: 

736 if orientation == 'vertical': 

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

738 else: 

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

740 

741 em = self.config.font_size 

742 

743 if orientation == 'vertical': 

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

745 self.separator * 2.0 

746 else: 

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

748 self.separator + self.margins[3] 

749 

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

751 self.fig.add_axes(axes) 

752 

753 self._colorbar_axes.append( 

754 (axes, orientation, position)) 

755 

756 self._update_layout() 

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

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

759 return colorbar.Colorbar( 

760 axes, mappable, orientation=orientation, label=label) 

761 

762 def __call__(self, *args): 

763 return self.axes(*args) 

764 

765 

766if __name__ == '__main__': 

767 import sys 

768 from pyrocko import util 

769 

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

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

772 

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

774 

775 if 0 in iplots: 

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

777 n = 100 

778 x = num.arange(n) * 2.0 

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

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

781 p.show() 

782 

783 if 1 in iplots: 

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

785 n = 100 

786 x = num.arange(n) * 2.0 

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

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

789 x = num.arange(n) * 2.0 

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

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

792 p.show() 

793 

794 if 11 in iplots: 

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

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

797 n = 100 

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

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

800 p.show() 

801 

802 if 12 in iplots: 

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

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

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

806 n = 100 

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

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

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

810 p.show() 

811 

812 if 13 in iplots: 

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

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

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

816 n = 100 

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

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

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

820 p.show() 

821 

822 if 2 in iplots: 

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

824 

825 n = 100 

826 

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

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

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

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

831 p.show() 

832 

833 if 3 in iplots: 

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

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

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

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

838 

839 n = 100 

840 

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

842 ned[:, 2] *= 0.25 

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

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

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

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

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

848 p.show() 

849 

850 if 5 in iplots: 

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

852 

853 n = 100 

854 

855 t = num.arange(n) 

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

857 xyz[:, 0] *= 0.5 

858 

859 smap = make_smap('summer') 

860 

861 p(0, 0).scatter( 

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

863 p(0, 1).scatter( 

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

865 p(0, 2).scatter( 

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

867 

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

869 

870 p.set_color_dim(smap, 'depth') 

871 

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

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

874 

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

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

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

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

879 

880 p.colorbar('depth') 

881 

882 p.show() 

883 

884 if 6 in iplots: 

885 km = 1000. 

886 p = Plot( 

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

888 

889 nn, ne = 50, 40 

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

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

892 

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

894 g = num.exp( 

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

896 

897 displacement[:, :, 0] = g 

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

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

900 

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

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

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

904 p.set_color_dim(c, 'displacement') 

905 

906 p.colorbar('displacement') 

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

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

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

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

911 

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

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

914 p.show()