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 

13import matplotlib 

14from matplotlib.axes import Axes 

15# from matplotlib.ticker import MultipleLocator 

16from matplotlib import cm, colors, colorbar, figure 

17 

18from pyrocko.guts import Tuple, Float, Object 

19from pyrocko import plot 

20 

21import scipy.optimize 

22 

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

24 

25guts_prefix = 'pf' 

26 

27inch = 2.54 

28 

29 

30class SmartplotAxes(Axes): 

31 

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

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

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

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

36 def cla(self): 

37 if hasattr(self, 'callbacks'): 

38 callbacks = self.callbacks 

39 Axes.cla(self) 

40 self.callbacks = callbacks 

41 else: 

42 Axes.cla(self) 

43 

44 else: 

45 def clear(self): 

46 if hasattr(self, 'callbacks'): 

47 callbacks = self.callbacks 

48 Axes.clear(self) 

49 self.callbacks = callbacks 

50 else: 

51 Axes.clear(self) 

52 

53 

54class SmartplotFigure(figure.Figure): 

55 

56 def set_smartplot(self, plot): 

57 self._smartplot = plot 

58 

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

60 if hasattr(self, '_smartplot'): 

61 try: 

62 self._smartplot._update_layout() 

63 except NotEnoughSpace: 

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

65 return 

66 

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

68 

69 

70def limits(points): 

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

72 if points.size != 0: 

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

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

75 

76 return lims 

77 

78 

79def wcenter(rect): 

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

81 

82 

83def hcenter(rect): 

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

85 

86 

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

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

89 

90 

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

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

93 

94 

95def make_smap(cmap, norm=None): 

96 if isinstance(norm, tuple): 

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

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

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

100 return smap 

101 

102 

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

104 

105 weight_aspect = 1000. 

106 

107 sx, sy = size 

108 nx, ny = shape 

109 nvar = nx+ny 

110 vxs, vys = limits 

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

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

113 aspects_xx, aspects_yy, aspects_xy = aspects 

114 

115 if fracs is None: 

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

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

118 else: 

119 frac_x, frac_y = fracs 

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

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

122 

123 data = [] 

124 weights = [] 

125 rows = [] 

126 bounds = [] 

127 for ix in range(nx): 

128 u = uxs[ix] 

129 assert u > 0.0 

130 row = num.zeros(nvar) 

131 row[ix] = u 

132 rows.append(row) 

133 data.append(wxs[ix]) 

134 weights.append(1.0 / u) 

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

136 

137 for iy in range(ny): 

138 u = uys[iy] 

139 assert u > 0.0 

140 row = num.zeros(nvar) 

141 row[nx+iy] = u 

142 rows.append(row) 

143 data.append(wys[iy]) 

144 weights.append(1.0) 

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

146 

147 for ix1, ix2, aspect in aspects_xx: 

148 row = num.zeros(nvar) 

149 row[ix1] = aspect 

150 row[ix2] = -1.0 

151 weights.append(weight_aspect/aspect) 

152 rows.append(row) 

153 data.append(0.0) 

154 

155 for iy1, iy2, aspect in aspects_yy: 

156 row = num.zeros(nvar) 

157 row[nx+iy1] = aspect 

158 row[nx+iy2] = -1.0 

159 weights.append(weight_aspect/aspect) 

160 rows.append(row) 

161 data.append(0.0) 

162 

163 for ix, iy, aspect in aspects_xy: 

164 row = num.zeros(nvar) 

165 row[ix] = aspect 

166 row[nx+iy] = -1.0 

167 weights.append(weight_aspect/aspect) 

168 rows.append(row) 

169 data.append(0.0) 

170 

171 weights = num.array(weights) 

172 data = num.array(data) 

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

174 data *= weights 

175 

176 bounds = num.array(bounds).T 

177 

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

179 

180 cxs = model[:nx] 

181 cys = model[nx:nx+ny] 

182 

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

184 for ix in range(nx): 

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

186 vmin, vmax = vxs[ix] 

187 udata = vmax - vmin 

188 eps = 1e-7 * u 

189 assert(udata <= u + eps) 

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

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

192 

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

194 for iy in range(ny): 

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

196 vmin, vmax = vys[iy] 

197 udata = vmax - vmin 

198 eps = 1e-7 * u 

199 assert(udata <= u + eps) 

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

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

202 

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

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

205 logger.error( 

206 'Unable to comply with requested aspect ratio ' 

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

208 

209 for ix1, ix2, aspect in aspects_xx: 

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

211 

212 for iy1, iy2, aspect in aspects_yy: 

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

214 

215 for ix, iy, aspect in aspects_xy: 

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

217 

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

219 

220 

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

222 

223 sx, sy = size 

224 nx, ny = shape 

225 vxs, vys = limits 

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

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

228 aspects_xx, aspects_yy, aspects_xy = aspects 

229 

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

231 for i in range(niterations): 

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

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

234 

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

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

237 wxs_used = wxs * uxs / uxs_view 

238 wys_used = wys * uys / uys_view 

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

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

241 

242 fracs_x = wxs_used 

243 fracs_y = wys_used 

244 

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

246 

247 

248class PlotError(Exception): 

249 pass 

250 

251 

252class NotEnoughSpace(PlotError): 

253 pass 

254 

255 

256class PlotConfig(Object): 

257 

258 font_size = Float.T(default=9.0) 

259 

260 size_cm = Tuple.T( 

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

262 

263 margins_em = Tuple.T( 

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

265 

266 separator_em = Float.T(default=1.5) 

267 

268 colorbar_width_em = Float.T(default=2.0) 

269 

270 label_offset_em = Tuple.T( 

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

272 

273 tick_label_offset_em = Tuple.T( 

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

275 

276 @property 

277 def size_inch(self): 

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

279 

280 

281class Plot(object): 

282 

283 def __init__( 

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

285 fig=None, call_mpl_init=True): 

286 

287 if config is None: 

288 config = PlotConfig() 

289 

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

291 

292 dims = [] 

293 for dim in x_dims + y_dims + z_dims: 

294 dim = dim.lstrip('-') 

295 if dim not in dims: 

296 dims.append(dim) 

297 

298 self.config = config 

299 self._disconnect_data = [] 

300 self._width = self._height = self._pixels = None 

301 if call_mpl_init: 

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

303 

304 if fig is None: 

305 fig = self._plt.figure( 

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

307 else: 

308 assert isinstance(fig, SmartplotFigure) 

309 

310 fig.set_smartplot(self) 

311 

312 self._fig = fig 

313 self._colorbar_width = 0.0 

314 self._colorbar_height = 0.0 

315 self._colorbar_axes = [] 

316 

317 self._dims = dims 

318 self._dim_index = self._dims.index 

319 self._ndims = len(dims) 

320 self._labels = {} 

321 self._aspects = {} 

322 

323 self.setup_axes() 

324 

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

326 

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

328 self._last_mpl_view_limits = None 

329 

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

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

332 

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

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

335 

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

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

338 

339 self._mappables = {} 

340 self._updating_layout = False 

341 

342 self._need_update_layout = True 

343 self._update_geometry() 

344 

345 for axes in self.axes_list: 

346 fig.add_axes(axes) 

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

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

349 

350 self._cid_resize = fig.canvas.mpl_connect( 

351 'resize_event', self.resize_handler) 

352 

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

354 

355 self._lim_changed_depth = 0 

356 

357 def reset_size(self): 

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

359 

360 def axes(self, ix, iy): 

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

362 ix = self._x_dims.index(ix) 

363 iy = self._y_dims.index(iy) 

364 

365 return self._axes[iy][ix] 

366 

367 def set_color_dim(self, mappable, dim): 

368 assert dim in self._dims 

369 self._mappables[mappable] = dim 

370 

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

372 self._aspects[ydim, xdim] = aspect 

373 

374 @property 

375 def dims(self): 

376 return self._dims 

377 

378 @property 

379 def fig(self): 

380 return self._fig 

381 

382 @property 

383 def axes_list(self): 

384 axes = [] 

385 for row in self._axes: 

386 axes.extend(row) 

387 return axes 

388 

389 @property 

390 def axes_bottom_list(self): 

391 return self._axes[0] 

392 

393 @property 

394 def axes_left_list(self): 

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

396 

397 def setup_axes(self): 

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

399 nx, ny = self._shape 

400 axes = [] 

401 for iy in range(ny): 

402 axes.append([]) 

403 for ix in range(nx): 

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

405 

406 self._axes = axes 

407 

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

409 axes_.set_autoscale_on(False) 

410 

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

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

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

414 

415 def _disconnect_all(self): 

416 for obj, cid in self._disconnect_data: 

417 obj.callbacks.disconnect(cid) 

418 

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

420 

421 def dpi_changed_handler(self, fig): 

422 if self._updating_layout: 

423 return 

424 

425 self._update_geometry() 

426 

427 def resize_handler(self, event): 

428 if self._updating_layout: 

429 return 

430 

431 self._update_geometry() 

432 

433 def lim_changed_handler(self, axes): 

434 if self._updating_layout: 

435 return 

436 

437 current = self._get_mpl_view_limits() 

438 last = self._last_mpl_view_limits 

439 

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

441 acurrent = current[iy][ix] 

442 alast = last[iy][ix] 

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

444 xdim = self._x_dims[ix] 

445 logger.debug( 

446 'X limits have been changed interactively in subplot ' 

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

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

449 

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

451 ydim = self._y_dims[iy] 

452 logger.debug( 

453 'Y limits have been changed interactively in subplot ' 

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

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

456 

457 self.need_update_layout() 

458 

459 def _update_geometry(self): 

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

461 dp = self.get_device_pixel_ratio() 

462 p = self.get_pixels_factor() * dp 

463 

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

465 logger.debug( 

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

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

468 w, h, p, dp)) 

469 

470 self._width = w # logical pixel 

471 self._height = h # logical pixel 

472 self._pixels = p # logical pixel / point 

473 self._device_pixel_ratio = dp # physical / logical 

474 self.need_update_layout() 

475 

476 @property 

477 def margins(self): 

478 return tuple( 

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

480 for x in self.config.margins_em) 

481 

482 @property 

483 def separator(self): 

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

485 

486 def rect_to_figure_coords(self, rect): 

487 left, bottom, width, height = rect 

488 return ( 

489 left / self._width, 

490 bottom / self._height, 

491 width / self._width, 

492 height / self._height) 

493 

494 def point_to_axes_coords(self, axes, point): 

495 x, y = point 

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

497 

498 x_fig = x / self._width 

499 y_fig = y / self._height 

500 

501 x_axes = (x_fig - aleft) / awidth 

502 y_axes = (y_fig - abottom) / aheight 

503 

504 return (x_axes, y_axes) 

505 

506 def get_pixels_factor(self): 

507 try: 

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

509 return 1.0 / r.points_to_pixels(1.0) 

510 except AttributeError: 

511 return 1.0 

512 

513 def get_device_pixel_ratio(self): 

514 try: 

515 return self._fig.canvas.device_pixel_ratio 

516 except AttributeError: 

517 return 1.0 

518 

519 def make_limits(self, lims): 

520 a = plot.AutoScaler(space=0.05) 

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

522 

523 def iaxes(self): 

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

525 for ix, axes in enumerate(row): 

526 yield iy, ix, axes 

527 

528 def get_data_limits(self): 

529 dim_to_values = defaultdict(list) 

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

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

532 axes.get_yaxis().get_data_interval()) 

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

534 axes.get_xaxis().get_data_interval()) 

535 

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

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

538 

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

540 for idim in range(self._ndims): 

541 dim = self._dims[idim] 

542 if dim in dim_to_values: 

543 vs = num.array( 

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

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

546 if vs.size > 0: 

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

548 else: 

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

550 else: 

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

552 

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

554 return lims 

555 

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

557 assert(vmin <= vmax) 

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

559 

560 def _get_mpl_view_limits(self): 

561 vl = [] 

562 for row in self._axes: 

563 vl_row = [] 

564 for axes in row: 

565 vl_row.append(( 

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

567 axes.get_yaxis().get_view_interval().tolist())) 

568 

569 vl.append(vl_row) 

570 

571 return vl 

572 

573 def _remember_mpl_view_limits(self): 

574 self._last_mpl_view_limits = self._get_mpl_view_limits() 

575 

576 def window_xmin(self, x): 

577 return window_min( 

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

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

580 self.separator, x) 

581 

582 def window_xmax(self, x): 

583 return window_max( 

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

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

586 self.separator, x) 

587 

588 def window_ymin(self, y): 

589 return window_min( 

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

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

592 self.separator, y) 

593 

594 def window_ymax(self, y): 

595 return window_max( 

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

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

598 self.separator, y) 

599 

600 def need_update_layout(self): 

601 self._need_update_layout = True 

602 

603 def _update_layout(self): 

604 assert not self._updating_layout 

605 

606 if not self._need_update_layout: 

607 return 

608 

609 self._updating_layout = True 

610 try: 

611 data_limits = self.get_data_limits() 

612 

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

614 for idim in range(self._ndims): 

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

616 

617 mask = num.isfinite(self._view_limits) 

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

619 

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

621 

622 # data_w = deltas[0] 

623 # data_h = deltas[1] 

624 

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

626 mr += self._colorbar_width 

627 mb += self._colorbar_height 

628 sw = sh = self.separator 

629 

630 nx, ny = self._shape 

631 

632 # data_r = data_h / data_w 

633 em = self.config.font_size / self._pixels 

634 w = self._width 

635 h = self._height 

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

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

638 

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

640 raise NotEnoughSpace() 

641 

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

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

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

645 

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

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

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

649 

650 def get_aspect(dim1, dim2): 

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

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

653 

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

655 

656 aspects_xx = [] 

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

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

659 aspect = get_aspect(xdim2, xdim1) 

660 if aspect: 

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

662 

663 aspects_yy = [] 

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

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

666 aspect = get_aspect(ydim2, ydim1) 

667 if aspect: 

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

669 

670 aspects_xy = [] 

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

672 xdim = self._x_dims[ix] 

673 ydim = self._y_dims[iy] 

674 aspect = get_aspect(ydim, xdim) 

675 if aspect: 

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

677 

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

679 size=(fig_w_avail, fig_h_avail), 

680 shape=(nx, ny), 

681 limits=(x_limits, y_limits), 

682 aspects=( 

683 aspects_xx, 

684 aspects_yy, 

685 aspects_xy)) 

686 

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

688 rect = [ 

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

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

691 aws[ix], ahs[iy]] 

692 

693 axes.set_position( 

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

695 

696 self.set_label_coords( 

697 axes, 'x', [ 

698 wcenter(rect), 

699 self.config.label_offset_em[0]*em 

700 + self._colorbar_height]) 

701 

702 self.set_label_coords( 

703 axes, 'y', [ 

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

705 hcenter(rect)]) 

706 

707 axes.get_xaxis().set_tick_params( 

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

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

710 

711 axes.get_yaxis().set_tick_params( 

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

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

714 

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

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

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

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

719 

720 axes.tick_params( 

721 axis='x', 

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

723 

724 axes.tick_params( 

725 axis='y', 

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

727 

728 self._remember_mpl_view_limits() 

729 

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

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

732 

733 # scaler = plot.AutoScaler() 

734 

735 # aspect tick incs same 

736 # 

737 # inc = scaler.make_scale( 

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

739 # override_mode='off')[2] 

740 # 

741 # for axes in self.axes_list: 

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

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

744 # 

745 # tl = MultipleLocator(inc) 

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

747 # tl = MultipleLocator(inc) 

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

749 

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

751 if orientation == 'horizontal': 

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

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

754 ymin = mb - self._colorbar_height 

755 ymax = mb - self._colorbar_height \ 

756 + self.config.colorbar_width_em * em 

757 else: 

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

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

760 xmin = w - mr + 2 * sw 

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

762 

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

764 axes.set_position( 

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

766 

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

768 dim = self._x_dims[ix] 

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

770 axes.set_xlabel(s) 

771 

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

773 dim = self._y_dims[iy] 

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

775 axes.set_ylabel(s) 

776 

777 finally: 

778 self._updating_layout = False 

779 

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

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

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

783 

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

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

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

787 for ix, axes in enumerate(row): 

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

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

790 

791 def close(self): 

792 self._disconnect_all() 

793 self._plt.close(self._fig) 

794 

795 def show(self): 

796 self._plt.show() 

797 self.reset_size() 

798 

799 def set_label(self, dim, s): 

800 # just set attribute, handle in update_layout 

801 self._labels[dim] = s 

802 

803 def colorbar( 

804 self, dim, 

805 orientation='vertical', 

806 position=None): 

807 

808 if dim not in self._dims: 

809 raise PlotError( 

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

811 

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

813 raise PlotError( 

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

815 

816 mappable = None 

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

818 if dim_ == dim: 

819 if mappable is None: 

820 mappable = mappable_ 

821 else: 

822 mappable_.set_cmap(mappable.get_cmap()) 

823 

824 if mappable is None: 

825 raise PlotError( 

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

827 

828 if position is None: 

829 if orientation == 'vertical': 

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

831 else: 

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

833 

834 em = self.config.font_size / self._pixels 

835 

836 if orientation == 'vertical': 

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

838 self.separator * 2.0 

839 else: 

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

841 self.separator + self.margins[3] 

842 

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

844 self.fig.add_axes(axes) 

845 

846 self._colorbar_axes.append( 

847 (axes, orientation, position)) 

848 

849 self.need_update_layout() 

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

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

852 return colorbar.Colorbar( 

853 axes, mappable, orientation=orientation, label=label) 

854 

855 def __call__(self, *args): 

856 return self.axes(*args) 

857 

858 

859if __name__ == '__main__': 

860 import sys 

861 from pyrocko import util 

862 

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

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

865 

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

867 

868 if 0 in iplots: 

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

870 n = 100 

871 x = num.arange(n) * 2.0 

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

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

874 p.show() 

875 

876 if 1 in iplots: 

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

878 n = 100 

879 x = num.arange(n) * 2.0 

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

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

882 x = num.arange(n) * 2.0 

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

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

885 p.show() 

886 

887 if 11 in iplots: 

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

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

890 n = 100 

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

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

893 p.show() 

894 

895 if 12 in iplots: 

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

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

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

899 n = 100 

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

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

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

903 p.show() 

904 

905 if 13 in iplots: 

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

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

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

909 n = 100 

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

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

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

913 p.show() 

914 

915 if 2 in iplots: 

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

917 

918 n = 100 

919 

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

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

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

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

924 p.show() 

925 

926 if 3 in iplots: 

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

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

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

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

931 

932 n = 100 

933 

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

935 ned[:, 2] *= 0.25 

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

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

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

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

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

941 p.show() 

942 

943 if 5 in iplots: 

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

945 

946 n = 100 

947 

948 t = num.arange(n) 

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

950 xyz[:, 0] *= 0.5 

951 

952 smap = make_smap('summer') 

953 

954 p(0, 0).scatter( 

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

956 p(0, 1).scatter( 

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

958 p(0, 2).scatter( 

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

960 

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

962 

963 p.set_color_dim(smap, 'depth') 

964 

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

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

967 

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

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

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

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

972 

973 p.colorbar('depth') 

974 

975 p.show() 

976 

977 if 6 in iplots: 

978 km = 1000. 

979 p = Plot( 

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

981 

982 nn, ne = 50, 40 

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

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

985 

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

987 g = num.exp( 

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

989 

990 displacement[:, :, 0] = g 

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

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

993 

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

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

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

997 p.set_color_dim(c, 'displacement') 

998 

999 p.colorbar('displacement') 

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

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

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

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

1004 

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

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

1007 p.show()