1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from collections import defaultdict
7import math
8import logging
10import numpy as num
11import matplotlib
12from matplotlib.axes import Axes
13# from matplotlib.ticker import MultipleLocator
14from matplotlib import cm, colors, colorbar, figure
16from pyrocko.guts import Tuple, Float, Object
17from pyrocko import plot
19import scipy.optimize
21logger = logging.getLogger('pyrocko.plot.smartplot')
23guts_prefix = 'pf'
25inch = 2.54
28def get_callbacks(obj):
29 try:
30 return obj.callbacks
31 except AttributeError:
32 return obj._callbacks
35class SmartplotAxes(Axes):
37 if matplotlib.__version__.split('.') < '3.6'.split('.'):
38 # Subclassing cla is deprecated on newer mpl but need this fallback for
39 # older versions. Code is duplicated because mpl behaviour depends
40 # on the existence of cla in the subclass...
41 def cla(self):
42 if hasattr(self, 'callbacks'):
43 callbacks = self.callbacks
44 Axes.cla(self)
45 self.callbacks = callbacks
46 else:
47 Axes.cla(self)
49 else:
50 def clear(self):
51 if hasattr(self, 'callbacks'):
52 callbacks = self.callbacks
53 Axes.clear(self)
54 self.callbacks = callbacks
55 elif hasattr(self, '_callbacks'):
56 callbacks = self._callbacks
57 Axes.clear(self)
58 self._callbacks = callbacks
59 else:
60 Axes.clear(self)
63class SmartplotFigure(figure.Figure):
65 def set_smartplot(self, plot):
66 self._smartplot = plot
68 def draw(self, *args, **kwargs):
69 if hasattr(self, '_smartplot'):
70 try:
71 self._smartplot._update_layout()
72 except NotEnoughSpace:
73 logger.error('Figure is too small to show the plot.')
74 return
76 return figure.Figure.draw(self, *args, **kwargs)
79def limits(points):
80 lims = num.zeros((3, 2))
81 if points.size != 0:
82 lims[:, 0] = num.min(points, axis=0)
83 lims[:, 1] = num.max(points, axis=0)
85 return lims
88def wcenter(rect):
89 return rect[0] + rect[2]*0.5
92def hcenter(rect):
93 return rect[1] + rect[3]*0.5
96def window_min(n, w, ml, mu, s, x):
97 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + math.floor(x) * s
100def window_max(n, w, ml, mu, s, x):
101 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + (math.floor(x)-1) * s
104def make_smap(cmap, norm=None):
105 if isinstance(norm, tuple):
106 norm = colors.Normalize(*norm, clip=False)
107 smap = cm.ScalarMappable(cmap=cmap, norm=norm)
108 smap._A = [] # not needed in newer versions of mpl?
109 return smap
112def solve_layout_fixed_panels(size, shape, limits, aspects, fracs=None):
114 weight_aspect = 1000.
116 sx, sy = size
117 nx, ny = shape
118 nvar = nx+ny
119 vxs, vys = limits
120 uxs = vxs[:, 1] - vxs[:, 0]
121 uys = vys[:, 1] - vys[:, 0]
122 aspects_xx, aspects_yy, aspects_xy = aspects
124 if fracs is None:
125 wxs = num.full(nx, sx / nx)
126 wys = num.full(ny, sy / ny)
127 else:
128 frac_x, frac_y = fracs
129 wxs = sx * frac_x / num.sum(frac_x)
130 wys = sy * frac_y / num.sum(frac_y)
132 data = []
133 weights = []
134 rows = []
135 bounds = []
136 for ix in range(nx):
137 u = uxs[ix]
138 assert u > 0.0
139 row = num.zeros(nvar)
140 row[ix] = u
141 rows.append(row)
142 data.append(wxs[ix])
143 weights.append(1.0 / u)
144 bounds.append((0, wxs[ix] / u))
146 for iy in range(ny):
147 u = uys[iy]
148 assert u > 0.0
149 row = num.zeros(nvar)
150 row[nx+iy] = u
151 rows.append(row)
152 data.append(wys[iy])
153 weights.append(1.0)
154 bounds.append((0, wys[iy] / u))
156 for ix1, ix2, aspect in aspects_xx:
157 row = num.zeros(nvar)
158 row[ix1] = aspect
159 row[ix2] = -1.0
160 weights.append(weight_aspect/aspect)
161 rows.append(row)
162 data.append(0.0)
164 for iy1, iy2, aspect in aspects_yy:
165 row = num.zeros(nvar)
166 row[nx+iy1] = aspect
167 row[nx+iy2] = -1.0
168 weights.append(weight_aspect/aspect)
169 rows.append(row)
170 data.append(0.0)
172 for ix, iy, aspect in aspects_xy:
173 row = num.zeros(nvar)
174 row[ix] = aspect
175 row[nx+iy] = -1.0
176 weights.append(weight_aspect/aspect)
177 rows.append(row)
178 data.append(0.0)
180 weights = num.array(weights)
181 data = num.array(data)
182 mat = num.vstack(rows) * weights[:, num.newaxis]
183 data *= weights
185 bounds = num.array(bounds).T
187 model = scipy.optimize.lsq_linear(mat, data, bounds).x
189 cxs = model[:nx]
190 cys = model[nx:nx+ny]
192 vlimits_x = num.zeros((nx, 2))
193 for ix in range(nx):
194 u = wxs[ix] / cxs[ix]
195 vmin, vmax = vxs[ix]
196 udata = vmax - vmin
197 eps = 1e-7 * u
198 assert udata <= u + eps
199 vlimits_x[ix, 0] = (vmin + vmax) / 2.0 - u / 2.0
200 vlimits_x[ix, 1] = (vmin + vmax) / 2.0 + u / 2.0
202 vlimits_y = num.zeros((ny, 2))
203 for iy in range(ny):
204 u = wys[iy] / cys[iy]
205 vmin, vmax = vys[iy]
206 udata = vmax - vmin
207 eps = 1e-7 * u
208 assert udata <= u + eps
209 vlimits_y[iy, 0] = (vmin + vmax) / 2.0 - u / 2.0
210 vlimits_y[iy, 1] = (vmin + vmax) / 2.0 + u / 2.0
212 def check_aspect(a, awant, eps=1e-2):
213 if abs(1.0 - (a/awant)) > eps:
214 logger.error(
215 'Unable to comply with requested aspect ratio '
216 '(wanted: %g, achieved: %g)' % (awant, a))
218 for ix1, ix2, aspect in aspects_xx:
219 check_aspect(cxs[ix2] / cxs[ix1], aspect)
221 for iy1, iy2, aspect in aspects_yy:
222 check_aspect(cys[iy2] / cys[iy1], aspect)
224 for ix, iy, aspect in aspects_xy:
225 check_aspect(cys[iy] / cxs[ix], aspect)
227 return (vlimits_x, vlimits_y), (wxs, wys)
230def solve_layout_iterative(size, shape, limits, aspects, niterations=3):
232 sx, sy = size
233 nx, ny = shape
234 vxs, vys = limits
235 uxs = vxs[:, 1] - vxs[:, 0]
236 uys = vys[:, 1] - vys[:, 0]
237 aspects_xx, aspects_yy, aspects_xy = aspects
239 fracs_x, fracs_y = num.ones(nx), num.ones(ny)
240 for i in range(niterations):
241 (vlimits_x, vlimits_y), (wxs, wys) = solve_layout_fixed_panels(
242 size, shape, limits, aspects, (fracs_x, fracs_y))
244 uxs_view = vlimits_x[:, 1] - vlimits_x[:, 0]
245 uys_view = vlimits_y[:, 1] - vlimits_y[:, 0]
246 wxs_used = wxs * uxs / uxs_view
247 wys_used = wys * uys / uys_view
248 # wxs_wasted = wxs * (1.0 - uxs / uxs_view)
249 # wys_wasted = wys * (1.0 - uys / uys_view)
251 fracs_x = wxs_used
252 fracs_y = wys_used
254 return (vlimits_x, vlimits_y), (wxs, wys)
257class PlotError(Exception):
258 pass
261class NotEnoughSpace(PlotError):
262 pass
265class PlotConfig(Object):
267 font_size = Float.T(default=9.0)
269 size_cm = Tuple.T(
270 2, Float.T(), default=(20., 20.))
272 margins_em = Tuple.T(
273 4, Float.T(), default=(8., 6., 8., 6.))
275 separator_em = Float.T(default=1.5)
277 colorbar_width_em = Float.T(default=2.0)
279 label_offset_em = Tuple.T(
280 2, Float.T(), default=(2., 2.))
282 tick_label_offset_em = Tuple.T(
283 2, Float.T(), default=(0.5, 0.5))
285 @property
286 def size_inch(self):
287 return self.size_cm[0]/inch, self.size_cm[1]/inch
290class Plot(object):
292 def __init__(
293 self, x_dims=['x'], y_dims=['y'], z_dims=[], config=None,
294 fig=None, call_mpl_init=True):
296 if config is None:
297 config = PlotConfig()
299 self._shape = len(x_dims), len(y_dims)
301 dims = []
302 for dim in x_dims + y_dims + z_dims:
303 dim = dim.lstrip('-')
304 if dim not in dims:
305 dims.append(dim)
307 self.config = config
308 self._disconnect_data = []
309 self._width = self._height = self._pixels = None
310 if call_mpl_init:
311 self._plt = plot.mpl_init(self.config.font_size)
313 if fig is None:
314 fig = self._plt.figure(
315 figsize=self.config.size_inch, FigureClass=SmartplotFigure)
316 else:
317 assert isinstance(fig, SmartplotFigure)
319 fig.set_smartplot(self)
321 self._fig = fig
322 self._colorbar_width = 0.0
323 self._colorbar_height = 0.0
324 self._colorbar_axes = []
326 self._dims = dims
327 self._dim_index = self._dims.index
328 self._ndims = len(dims)
329 self._labels = {}
330 self._aspects = {}
332 self.setup_axes()
334 self._view_limits = num.zeros((self._ndims, 2))
336 self._view_limits[:, :] = num.nan
337 self._last_mpl_view_limits = None
339 self._x_dims = [dim.lstrip('-') for dim in x_dims]
340 self._x_dims_invert = [dim.startswith('-') for dim in x_dims]
342 self._y_dims = [dim.lstrip('-') for dim in y_dims]
343 self._y_dims_invert = [dim.startswith('-') for dim in y_dims]
345 self._z_dims = [dim.lstrip('-') for dim in z_dims]
346 self._z_dims_invert = [dim.startswith('-') for dim in z_dims]
348 self._mappables = {}
349 self._updating_layout = False
351 self._need_update_layout = True
352 self._update_geometry()
354 for axes in self.axes_list:
355 fig.add_axes(axes)
356 self._connect(axes, 'xlim_changed', self.lim_changed_handler)
357 self._connect(axes, 'ylim_changed', self.lim_changed_handler)
359 self._cid_resize = fig.canvas.mpl_connect(
360 'resize_event', self.resize_handler)
362 try:
363 self._connect(fig, 'dpi_changed', self.dpi_changed_handler)
364 except ValueError:
365 # 'dpi_changed' event has been removed in MPL 3.8.
366 # canvas 'resize_event' may be sufficient but needs to be checked.
367 # https://matplotlib.org/stable/api/prev_api_changes/api_changes_3.8.0.html#text-get-rotation
368 pass
370 self._lim_changed_depth = 0
372 def reset_size(self):
373 self._fig.set_size_inches(self.config.size_inch)
375 def axes(self, ix, iy):
376 if not (isinstance(ix, int) and isinstance(iy, int)):
377 ix = self._x_dims.index(ix)
378 iy = self._y_dims.index(iy)
380 return self._axes[iy][ix]
382 def set_color_dim(self, mappable, dim):
383 assert dim in self._dims
384 self._mappables[mappable] = dim
386 def set_aspect(self, ydim, xdim, aspect=1.0):
387 self._aspects[ydim, xdim] = aspect
389 @property
390 def dims(self):
391 return self._dims
393 @property
394 def fig(self):
395 return self._fig
397 @property
398 def axes_list(self):
399 axes = []
400 for row in self._axes:
401 axes.extend(row)
402 return axes
404 @property
405 def axes_bottom_list(self):
406 return self._axes[0]
408 @property
409 def axes_left_list(self):
410 return [row[0] for row in self._axes]
412 def setup_axes(self):
413 rect = [0., 0., 1., 1.]
414 nx, ny = self._shape
415 axes = []
416 for iy in range(ny):
417 axes.append([])
418 for ix in range(nx):
419 axes[-1].append(SmartplotAxes(self.fig, rect))
421 self._axes = axes
423 for _, _, axes_ in self.iaxes():
424 axes_.set_autoscale_on(False)
426 def _connect(self, obj, sig, handler):
427 cid = get_callbacks(obj).connect(sig, handler)
428 self._disconnect_data.append((obj, cid))
430 def _disconnect_all(self):
431 for obj, cid in self._disconnect_data:
432 get_callbacks(obj).disconnect(cid)
434 self._fig.canvas.mpl_disconnect(self._cid_resize)
436 def dpi_changed_handler(self, fig):
437 if self._updating_layout:
438 return
440 self._update_geometry()
442 def resize_handler(self, event):
443 if self._updating_layout:
444 return
446 self._update_geometry()
448 def lim_changed_handler(self, axes):
449 if self._updating_layout:
450 return
452 current = self._get_mpl_view_limits()
453 last = self._last_mpl_view_limits
454 if last is None:
455 return
457 for iy, ix, axes in self.iaxes():
458 acurrent = current[iy][ix]
459 alast = last[iy][ix]
460 if acurrent[0] != alast[0]:
461 xdim = self._x_dims[ix]
462 logger.debug(
463 'X limits have been changed interactively in subplot '
464 '(%i, %i)' % (ix, iy))
465 self.set_lim(xdim, *sorted(acurrent[0]))
467 if acurrent[1] != alast[1]:
468 ydim = self._y_dims[iy]
469 logger.debug(
470 'Y limits have been changed interactively in subplot '
471 '(%i, %i)' % (ix, iy))
472 self.set_lim(ydim, *sorted(acurrent[1]))
474 self.need_update_layout()
476 def _update_geometry(self):
477 w, h = self._fig.canvas.get_width_height()
478 dp = self.get_device_pixel_ratio()
479 p = self.get_pixels_factor() * dp
481 if (self._width, self._height, self._pixels) != (w, h, p, dp):
482 logger.debug(
483 'New figure size: %g x %g, '
484 'logical-pixel/point: %g, physical-pixel/logical-pixel: %g' % (
485 w, h, p, dp))
487 self._width = w # logical pixel
488 self._height = h # logical pixel
489 self._pixels = p # logical pixel / point
490 self._device_pixel_ratio = dp # physical / logical
491 self.need_update_layout()
493 @property
494 def margins(self):
495 return tuple(
496 x * self.config.font_size / self._pixels
497 for x in self.config.margins_em)
499 @property
500 def separator(self):
501 return self.config.separator_em * self.config.font_size / self._pixels
503 def rect_to_figure_coords(self, rect):
504 left, bottom, width, height = rect
505 return (
506 left / self._width,
507 bottom / self._height,
508 width / self._width,
509 height / self._height)
511 def point_to_axes_coords(self, axes, point):
512 x, y = point
513 aleft, abottom, awidth, aheight = axes.get_position().bounds
515 x_fig = x / self._width
516 y_fig = y / self._height
518 x_axes = (x_fig - aleft) / awidth
519 y_axes = (y_fig - abottom) / aheight
521 return (x_axes, y_axes)
523 def get_pixels_factor(self):
524 try:
525 r = self._fig.canvas.get_renderer()
526 return 1.0 / r.points_to_pixels(1.0)
527 except AttributeError:
528 return 1.0
530 def get_device_pixel_ratio(self):
531 try:
532 return self._fig.canvas.device_pixel_ratio
533 except AttributeError:
534 return 1.0
536 def make_limits(self, lims):
537 a = plot.AutoScaler(space=0.05)
538 return a.make_scale(lims)[:2]
540 def iaxes(self):
541 for iy, row in enumerate(self._axes):
542 for ix, axes in enumerate(row):
543 yield iy, ix, axes
545 def get_data_limits(self):
546 dim_to_values = defaultdict(list)
547 for iy, ix, axes in self.iaxes():
548 dim_to_values[self._y_dims[iy]].extend(
549 axes.get_yaxis().get_data_interval())
550 dim_to_values[self._x_dims[ix]].extend(
551 axes.get_xaxis().get_data_interval())
553 for mappable, dim in self._mappables.items():
554 dim_to_values[dim].extend(mappable.get_clim())
556 lims = num.zeros((self._ndims, 2))
557 for idim in range(self._ndims):
558 dim = self._dims[idim]
559 if dim in dim_to_values:
560 vs = num.array(
561 dim_to_values[self._dims[idim]], dtype=float)
562 vs = vs[num.isfinite(vs)]
563 if vs.size > 0:
564 lims[idim, :] = num.min(vs), num.max(vs)
565 else:
566 lims[idim, :] = num.nan, num.nan
567 else:
568 lims[idim, :] = num.nan, num.nan
570 lims[num.logical_not(num.isfinite(lims))] = 0.0
571 return lims
573 def set_lim(self, dim, vmin, vmax):
574 assert vmin <= vmax
575 self._view_limits[self._dim_index(dim), :] = vmin, vmax
577 def _get_mpl_view_limits(self):
578 vl = []
579 for row in self._axes:
580 vl_row = []
581 for axes in row:
582 vl_row.append((
583 axes.get_xaxis().get_view_interval().tolist(),
584 axes.get_yaxis().get_view_interval().tolist()))
586 vl.append(vl_row)
588 return vl
590 def _remember_mpl_view_limits(self):
591 self._last_mpl_view_limits = self._get_mpl_view_limits()
593 def window_xmin(self, x):
594 return window_min(
595 self._shape[0], self._width,
596 self.margins[0], self.margins[2] + self._colorbar_width,
597 self.separator, x)
599 def window_xmax(self, x):
600 return window_max(
601 self._shape[0], self._width,
602 self.margins[0], self.margins[2] + self._colorbar_width,
603 self.separator, x)
605 def window_ymin(self, y):
606 return window_min(
607 self._shape[1], self._height,
608 self.margins[3] + self._colorbar_height, self.margins[1],
609 self.separator, y)
611 def window_ymax(self, y):
612 return window_max(
613 self._shape[1], self._height,
614 self.margins[3] + self._colorbar_height, self.margins[1],
615 self.separator, y)
617 def need_update_layout(self):
618 self._need_update_layout = True
620 def _update_layout(self):
621 assert not self._updating_layout
623 if not self._need_update_layout:
624 return
626 self._updating_layout = True
627 try:
628 data_limits = self.get_data_limits()
630 limits = num.zeros((self._ndims, 2))
631 for idim in range(self._ndims):
632 limits[idim, :] = self.make_limits(data_limits[idim, :])
634 mask = num.isfinite(self._view_limits)
635 limits[mask] = self._view_limits[mask]
637 # deltas = limits[:, 1] - limits[:, 0]
639 # data_w = deltas[0]
640 # data_h = deltas[1]
642 ml, mt, mr, mb = self.margins
643 mr += self._colorbar_width
644 mb += self._colorbar_height
645 sw = sh = self.separator
647 nx, ny = self._shape
649 # data_r = data_h / data_w
650 em = self.config.font_size
651 em_pixels = em / self._pixels
652 w = self._width
653 h = self._height
654 fig_w_avail = w - mr - ml - (nx-1) * sw
655 fig_h_avail = h - mt - mb - (ny-1) * sh
657 if fig_w_avail <= 0.0 or fig_h_avail <= 0.0:
658 raise NotEnoughSpace()
660 x_limits = num.zeros((nx, 2))
661 for ix, xdim in enumerate(self._x_dims):
662 x_limits[ix, :] = limits[self._dim_index(xdim)]
664 y_limits = num.zeros((ny, 2))
665 for iy, ydim in enumerate(self._y_dims):
666 y_limits[iy, :] = limits[self._dim_index(ydim)]
668 def get_aspect(dim1, dim2):
669 if (dim2, dim1) in self._aspects:
670 return 1.0/self._aspects[dim2, dim1]
672 return self._aspects.get((dim1, dim2), None)
674 aspects_xx = []
675 for ix1, xdim1 in enumerate(self._x_dims):
676 for ix2, xdim2 in enumerate(self._x_dims):
677 aspect = get_aspect(xdim2, xdim1)
678 if aspect:
679 aspects_xx.append((ix1, ix2, aspect))
681 aspects_yy = []
682 for iy1, ydim1 in enumerate(self._y_dims):
683 for iy2, ydim2 in enumerate(self._y_dims):
684 aspect = get_aspect(ydim2, ydim1)
685 if aspect:
686 aspects_yy.append((iy1, iy2, aspect))
688 aspects_xy = []
689 for iy, ix, axes in self.iaxes():
690 xdim = self._x_dims[ix]
691 ydim = self._y_dims[iy]
692 aspect = get_aspect(ydim, xdim)
693 if aspect:
694 aspects_xy.append((ix, iy, aspect))
696 (x_limits, y_limits), (aws, ahs) = solve_layout_iterative(
697 size=(fig_w_avail, fig_h_avail),
698 shape=(nx, ny),
699 limits=(x_limits, y_limits),
700 aspects=(
701 aspects_xx,
702 aspects_yy,
703 aspects_xy))
705 for iy, ix, axes in self.iaxes():
706 rect = [
707 ml + num.sum(aws[:ix])+(ix*sw),
708 mb + num.sum(ahs[:iy])+(iy*sh),
709 aws[ix], ahs[iy]]
711 axes.set_position(
712 self.rect_to_figure_coords(rect), which='both')
714 self.set_label_coords(
715 axes, 'x', [
716 wcenter(rect),
717 self.config.label_offset_em[0]*em_pixels
718 + self._colorbar_height])
720 self.set_label_coords(
721 axes, 'y', [
722 self.config.label_offset_em[1]*em_pixels,
723 hcenter(rect)])
725 axes.get_xaxis().set_tick_params(
726 bottom=(iy == 0), top=(iy == ny-1),
727 labelbottom=(iy == 0), labeltop=False)
729 axes.get_yaxis().set_tick_params(
730 left=(ix == 0), right=(ix == nx-1),
731 labelleft=(ix == 0), labelright=False)
733 istride = -1 if self._x_dims_invert[ix] else 1
734 axes.set_xlim(*x_limits[ix, ::istride])
735 istride = -1 if self._y_dims_invert[iy] else 1
736 axes.set_ylim(*y_limits[iy, ::istride])
738 axes.tick_params(
739 axis='x',
740 pad=self.config.tick_label_offset_em[0]*em)
742 axes.tick_params(
743 axis='y',
744 pad=self.config.tick_label_offset_em[0]*em)
746 self._remember_mpl_view_limits()
748 for mappable, dim in self._mappables.items():
749 mappable.set_clim(*limits[self._dim_index(dim)])
751 # scaler = plot.AutoScaler()
753 # aspect tick incs same
754 #
755 # inc = scaler.make_scale(
756 # [0, min(data_expanded_w, data_expanded_h)],
757 # override_mode='off')[2]
758 #
759 # for axes in self.axes_list:
760 # axes.set_xlim(*limits[0, :])
761 # axes.set_ylim(*limits[1, :])
762 #
763 # tl = MultipleLocator(inc)
764 # axes.get_xaxis().set_major_locator(tl)
765 # tl = MultipleLocator(inc)
766 # axes.get_yaxis().set_major_locator(tl)
768 for axes, orientation, position in self._colorbar_axes:
769 if orientation == 'horizontal':
770 xmin = self.window_xmin(position[0])
771 xmax = self.window_xmax(position[1])
772 ymin = mb - self._colorbar_height
773 ymax = mb - self._colorbar_height \
774 + self.config.colorbar_width_em * em_pixels
775 else:
776 ymin = self.window_ymin(position[0])
777 ymax = self.window_ymax(position[1])
778 xmin = w - mr + 2 * sw
779 xmax = w - mr + 2 * sw \
780 + self.config.colorbar_width_em * em_pixels
782 rect = [xmin, ymin, xmax-xmin, ymax-ymin]
783 axes.set_position(
784 self.rect_to_figure_coords(rect), which='both')
786 for ix, axes in enumerate(self.axes_bottom_list):
787 dim = self._x_dims[ix]
788 s = self._labels.get(dim, dim)
789 axes.set_xlabel(s)
791 for iy, axes in enumerate(self.axes_left_list):
792 dim = self._y_dims[iy]
793 s = self._labels.get(dim, dim)
794 axes.set_ylabel(s)
796 finally:
797 self._updating_layout = False
799 def set_label_coords(self, axes, which, point):
800 axis = axes.get_xaxis() if which == 'x' else axes.get_yaxis()
801 axis.set_label_coords(*self.point_to_axes_coords(axes, point))
803 def plot(self, points, *args, **kwargs):
804 for iy, row in enumerate(self._axes):
805 y = points[:, self._dim_index(self._y_dims[iy])]
806 for ix, axes in enumerate(row):
807 x = points[:, self._dim_index(self._x_dims[ix])]
808 axes.plot(x, y, *args, **kwargs)
810 def close(self):
811 self._disconnect_all()
812 self._plt.close(self._fig)
814 def show(self):
815 self._plt.show()
816 self.reset_size()
818 def set_label(self, dim, s):
819 # just set attribute, handle in update_layout
820 self._labels[dim] = s
822 def colorbar(
823 self, dim,
824 orientation='vertical',
825 position=None):
827 if dim not in self._dims:
828 raise PlotError(
829 'dimension "%s" is not defined')
831 if orientation not in ('vertical', 'horizontal'):
832 raise PlotError(
833 'orientation must be "vertical" or "horizontal"')
835 mappable = None
836 for mappable_, dim_ in self._mappables.items():
837 if dim_ == dim:
838 if mappable is None:
839 mappable = mappable_
840 else:
841 mappable_.set_cmap(mappable.get_cmap())
843 if mappable is None:
844 raise PlotError(
845 'no mappable registered for dimension "%s"' % dim)
847 if position is None:
848 if orientation == 'vertical':
849 position = (0, self._shape[1])
850 else:
851 position = (0, self._shape[0])
853 em_pixels = self.config.font_size / self._pixels
855 if orientation == 'vertical':
856 self._colorbar_width = self.config.colorbar_width_em*em_pixels + \
857 self.separator * 2.0
858 else:
859 self._colorbar_height = self.config.colorbar_width_em*em_pixels + \
860 self.separator + self.margins[3]
862 axes = SmartplotAxes(self.fig, [0., 0., 1., 1.])
863 self.fig.add_axes(axes)
865 self._colorbar_axes.append(
866 (axes, orientation, position))
868 self.need_update_layout()
869 # axes.plot([1], [1])
870 label = self._labels.get(dim, dim)
871 return colorbar.Colorbar(
872 axes, mappable, orientation=orientation, label=label)
874 def __call__(self, *args):
875 return self.axes(*args)
878if __name__ == '__main__':
879 import sys
880 from pyrocko import util
882 logging.getLogger('matplotlib').setLevel(logging.WARNING)
883 util.setup_logging('smartplot', 'debug')
885 iplots = [int(x) for x in sys.argv[1:]]
887 if 0 in iplots:
888 p = Plot(['x'], ['y'])
889 n = 100
890 x = num.arange(n) * 2.0
891 y = num.random.normal(size=n)
892 p(0, 0).plot(x, y, 'o')
893 p.show()
895 if 1 in iplots:
896 p = Plot(['x', 'x'], ['y'])
897 n = 100
898 x = num.arange(n) * 2.0
899 y = num.random.normal(size=n)
900 p(0, 0).plot(x, y, 'o')
901 x = num.arange(n) * 2.0
902 y = num.random.normal(size=n)
903 p(1, 0).plot(x, y, 'o')
904 p.show()
906 if 11 in iplots:
907 p = Plot(['x'], ['y'])
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.show()
914 if 12 in iplots:
915 p = Plot(['x', 'x2'], ['y'])
916 p.set_aspect('x2', 'x', 2.0)
917 p.set_aspect('y', 'x', 2.0)
918 n = 100
919 xy = num.random.normal(size=(n, 2))
920 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o')
921 p(1, 0).plot(xy[:, 0], xy[:, 1], 'o')
922 p.show()
924 if 13 in iplots:
925 p = Plot(['x'], ['y', 'y2'])
926 p.set_aspect('y2', 'y', 2.0)
927 p.set_aspect('y', 'x', 2.0)
928 n = 100
929 xy = num.random.normal(size=(n, 2))
930 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o')
931 p(0, 1).plot(xy[:, 0], xy[:, 1], 'o')
932 p.show()
934 if 2 in iplots:
935 p = Plot(['easting', 'depth'], ['northing', 'depth'])
937 n = 100
939 ned = num.random.normal(size=(n, 3))
940 p(0, 0).plot(ned[:, 1], ned[:, 0], 'o')
941 p(1, 0).plot(ned[:, 2], ned[:, 0], 'o')
942 p(0, 1).plot(ned[:, 1], ned[:, 2], 'o')
943 p.show()
945 if 3 in iplots:
946 p = Plot(['easting', 'depth'], ['-depth', 'northing'])
947 p.set_aspect('easting', 'northing', 1.0)
948 p.set_aspect('easting', 'depth', 0.5)
949 p.set_aspect('northing', 'depth', 0.5)
951 n = 100
953 ned = num.random.normal(size=(n, 3))
954 ned[:, 2] *= 0.25
955 p(0, 1).plot(ned[:, 1], ned[:, 0], 'o', color='black')
956 p(0, 0).plot(ned[:, 1], ned[:, 2], 'o')
957 p(1, 1).plot(ned[:, 2], ned[:, 0], 'o')
958 p(1, 0).set_visible(False)
959 p.set_lim('depth', 0., 0.2)
960 p.show()
962 if 5 in iplots:
963 p = Plot(['time'], ['northing', 'easting', '-depth'], ['depth'])
965 n = 100
967 t = num.arange(n)
968 xyz = num.random.normal(size=(n, 4))
969 xyz[:, 0] *= 0.5
971 smap = make_smap('summer')
973 p(0, 0).scatter(
974 t, xyz[:, 0], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
975 p(0, 1).scatter(
976 t, xyz[:, 1], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
977 p(0, 2).scatter(
978 t, xyz[:, 2], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
980 p.set_lim('depth', -1., 1.)
982 p.set_color_dim(smap, 'depth')
984 p.set_aspect('northing', 'easting', 1.0)
985 p.set_aspect('northing', 'depth', 1.0)
987 p.set_label('time', 'Time [s]')
988 p.set_label('depth', 'Depth [km]')
989 p.set_label('easting', 'Easting [km]')
990 p.set_label('northing', 'Northing [km]')
992 p.colorbar('depth')
994 p.show()
996 if 6 in iplots:
997 km = 1000.
998 p = Plot(
999 ['easting'], ['northing']*3, ['displacement'])
1001 nn, ne = 50, 40
1002 n = num.linspace(-5*km, 5*km, nn)
1003 e = num.linspace(-10*km, 10*km, ne)
1005 displacement = num.zeros((nn, ne, 3))
1006 g = num.exp(
1007 -(n[:, num.newaxis]**2 + e[num.newaxis, :]**2) / (5*km)**2)
1009 displacement[:, :, 0] = g
1010 displacement[:, :, 1] = g * 0.5
1011 displacement[:, :, 2] = -g * 0.2
1013 for icomp in (0, 1, 2):
1014 c = p(0, icomp).pcolormesh(
1015 e/km, n/km, displacement[:, :, icomp], shading='gouraud')
1016 p.set_color_dim(c, 'displacement')
1018 p.colorbar('displacement')
1019 p.set_lim('displacement', -1.0, 1.0)
1020 p.set_label('easting', 'Easting [km]')
1021 p.set_label('northing', 'Northing [km]')
1022 p.set_aspect('northing', 'easting')
1024 p.set_lim('northing', -5.0, 5.0)
1025 p.set_lim('easting', -3.0, 3.0)
1026 p.show()