1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6from __future__ import print_function
8from collections import defaultdict
9import math
10import logging
12import numpy as num
13import matplotlib
14from matplotlib.axes import Axes
15# from matplotlib.ticker import MultipleLocator
16from matplotlib import cm, colors, colorbar, figure
18from pyrocko.guts import Tuple, Float, Object
19from pyrocko import plot
21import scipy.optimize
23logger = logging.getLogger('pyrocko.plot.smartplot')
25guts_prefix = 'pf'
27inch = 2.54
30class SmartplotAxes(Axes):
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)
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)
54class SmartplotFigure(figure.Figure):
56 def set_smartplot(self, plot):
57 self._smartplot = plot
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
67 return figure.Figure.draw(self, *args, **kwargs)
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)
76 return lims
79def wcenter(rect):
80 return rect[0] + rect[2]*0.5
83def hcenter(rect):
84 return rect[1] + rect[3]*0.5
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
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
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
103def solve_layout_fixed_panels(size, shape, limits, aspects, fracs=None):
105 weight_aspect = 1000.
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
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)
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))
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))
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)
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)
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)
171 weights = num.array(weights)
172 data = num.array(data)
173 mat = num.vstack(rows) * weights[:, num.newaxis]
174 data *= weights
176 bounds = num.array(bounds).T
178 model = scipy.optimize.lsq_linear(mat, data, bounds).x
180 cxs = model[:nx]
181 cys = model[nx:nx+ny]
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
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
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))
209 for ix1, ix2, aspect in aspects_xx:
210 check_aspect(cxs[ix2] / cxs[ix1], aspect)
212 for iy1, iy2, aspect in aspects_yy:
213 check_aspect(cys[iy2] / cys[iy1], aspect)
215 for ix, iy, aspect in aspects_xy:
216 check_aspect(cys[iy] / cxs[ix], aspect)
218 return (vlimits_x, vlimits_y), (wxs, wys)
221def solve_layout_iterative(size, shape, limits, aspects, niterations=3):
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
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))
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)
242 fracs_x = wxs_used
243 fracs_y = wys_used
245 return (vlimits_x, vlimits_y), (wxs, wys)
248class PlotError(Exception):
249 pass
252class NotEnoughSpace(PlotError):
253 pass
256class PlotConfig(Object):
258 font_size = Float.T(default=9.0)
260 size_cm = Tuple.T(
261 2, Float.T(), default=(20., 20.))
263 margins_em = Tuple.T(
264 4, Float.T(), default=(8., 6., 8., 6.))
266 separator_em = Float.T(default=1.5)
268 colorbar_width_em = Float.T(default=2.0)
270 label_offset_em = Tuple.T(
271 2, Float.T(), default=(2., 2.))
273 tick_label_offset_em = Tuple.T(
274 2, Float.T(), default=(0.5, 0.5))
276 @property
277 def size_inch(self):
278 return self.size_cm[0]/inch, self.size_cm[1]/inch
281class Plot(object):
283 def __init__(
284 self, x_dims=['x'], y_dims=['y'], z_dims=[], config=None,
285 fig=None, call_mpl_init=True):
287 if config is None:
288 config = PlotConfig()
290 self._shape = len(x_dims), len(y_dims)
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)
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)
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)
310 fig.set_smartplot(self)
312 self._fig = fig
313 self._colorbar_width = 0.0
314 self._colorbar_height = 0.0
315 self._colorbar_axes = []
317 self._dims = dims
318 self._dim_index = self._dims.index
319 self._ndims = len(dims)
320 self._labels = {}
321 self._aspects = {}
323 self.setup_axes()
325 self._view_limits = num.zeros((self._ndims, 2))
327 self._view_limits[:, :] = num.nan
328 self._last_mpl_view_limits = None
330 self._x_dims = [dim.lstrip('-') for dim in x_dims]
331 self._x_dims_invert = [dim.startswith('-') for dim in x_dims]
333 self._y_dims = [dim.lstrip('-') for dim in y_dims]
334 self._y_dims_invert = [dim.startswith('-') for dim in y_dims]
336 self._z_dims = [dim.lstrip('-') for dim in z_dims]
337 self._z_dims_invert = [dim.startswith('-') for dim in z_dims]
339 self._mappables = {}
340 self._updating_layout = False
342 self._need_update_layout = True
343 self._update_geometry()
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)
350 self._cid_resize = fig.canvas.mpl_connect(
351 'resize_event', self.resize_handler)
353 self._connect(fig, 'dpi_changed', self.dpi_changed_handler)
355 self._lim_changed_depth = 0
357 def reset_size(self):
358 self._fig.set_size_inches(self.config.size_inch)
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)
365 return self._axes[iy][ix]
367 def set_color_dim(self, mappable, dim):
368 assert dim in self._dims
369 self._mappables[mappable] = dim
371 def set_aspect(self, ydim, xdim, aspect=1.0):
372 self._aspects[ydim, xdim] = aspect
374 @property
375 def dims(self):
376 return self._dims
378 @property
379 def fig(self):
380 return self._fig
382 @property
383 def axes_list(self):
384 axes = []
385 for row in self._axes:
386 axes.extend(row)
387 return axes
389 @property
390 def axes_bottom_list(self):
391 return self._axes[0]
393 @property
394 def axes_left_list(self):
395 return [row[0] for row in self._axes]
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))
406 self._axes = axes
408 for _, _, axes_ in self.iaxes():
409 axes_.set_autoscale_on(False)
411 def _connect(self, obj, sig, handler):
412 cid = obj.callbacks.connect(sig, handler)
413 self._disconnect_data.append((obj, cid))
415 def _disconnect_all(self):
416 for obj, cid in self._disconnect_data:
417 obj.callbacks.disconnect(cid)
419 self._fig.canvas.mpl_disconnect(self._cid_resize)
421 def dpi_changed_handler(self, fig):
422 if self._updating_layout:
423 return
425 self._update_geometry()
427 def resize_handler(self, event):
428 if self._updating_layout:
429 return
431 self._update_geometry()
433 def lim_changed_handler(self, axes):
434 if self._updating_layout:
435 return
437 current = self._get_mpl_view_limits()
438 last = self._last_mpl_view_limits
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]))
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]))
457 self.need_update_layout()
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
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))
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()
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)
482 @property
483 def separator(self):
484 return self.config.separator_em * self.config.font_size / self._pixels
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)
494 def point_to_axes_coords(self, axes, point):
495 x, y = point
496 aleft, abottom, awidth, aheight = axes.get_position().bounds
498 x_fig = x / self._width
499 y_fig = y / self._height
501 x_axes = (x_fig - aleft) / awidth
502 y_axes = (y_fig - abottom) / aheight
504 return (x_axes, y_axes)
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
513 def get_device_pixel_ratio(self):
514 try:
515 return self._fig.canvas.device_pixel_ratio
516 except AttributeError:
517 return 1.0
519 def make_limits(self, lims):
520 a = plot.AutoScaler(space=0.05)
521 return a.make_scale(lims)[:2]
523 def iaxes(self):
524 for iy, row in enumerate(self._axes):
525 for ix, axes in enumerate(row):
526 yield iy, ix, axes
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())
536 for mappable, dim in self._mappables.items():
537 dim_to_values[dim].extend(mappable.get_clim())
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
553 lims[num.logical_not(num.isfinite(lims))] = 0.0
554 return lims
556 def set_lim(self, dim, vmin, vmax):
557 assert vmin <= vmax
558 self._view_limits[self._dim_index(dim), :] = vmin, vmax
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()))
569 vl.append(vl_row)
571 return vl
573 def _remember_mpl_view_limits(self):
574 self._last_mpl_view_limits = self._get_mpl_view_limits()
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)
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)
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)
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)
600 def need_update_layout(self):
601 self._need_update_layout = True
603 def _update_layout(self):
604 assert not self._updating_layout
606 if not self._need_update_layout:
607 return
609 self._updating_layout = True
610 try:
611 data_limits = self.get_data_limits()
613 limits = num.zeros((self._ndims, 2))
614 for idim in range(self._ndims):
615 limits[idim, :] = self.make_limits(data_limits[idim, :])
617 mask = num.isfinite(self._view_limits)
618 limits[mask] = self._view_limits[mask]
620 # deltas = limits[:, 1] - limits[:, 0]
622 # data_w = deltas[0]
623 # data_h = deltas[1]
625 ml, mt, mr, mb = self.margins
626 mr += self._colorbar_width
627 mb += self._colorbar_height
628 sw = sh = self.separator
630 nx, ny = self._shape
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
639 if fig_w_avail <= 0.0 or fig_h_avail <= 0.0:
640 raise NotEnoughSpace()
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)]
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)]
650 def get_aspect(dim1, dim2):
651 if (dim2, dim1) in self._aspects:
652 return 1.0/self._aspects[dim2, dim1]
654 return self._aspects.get((dim1, dim2), None)
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))
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))
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))
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))
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]]
693 axes.set_position(
694 self.rect_to_figure_coords(rect), which='both')
696 self.set_label_coords(
697 axes, 'x', [
698 wcenter(rect),
699 self.config.label_offset_em[0]*em
700 + self._colorbar_height])
702 self.set_label_coords(
703 axes, 'y', [
704 self.config.label_offset_em[1]*em,
705 hcenter(rect)])
707 axes.get_xaxis().set_tick_params(
708 bottom=(iy == 0), top=(iy == ny-1),
709 labelbottom=(iy == 0), labeltop=False)
711 axes.get_yaxis().set_tick_params(
712 left=(ix == 0), right=(ix == nx-1),
713 labelleft=(ix == 0), labelright=False)
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])
720 axes.tick_params(
721 axis='x',
722 pad=self.config.tick_label_offset_em[0]*em)
724 axes.tick_params(
725 axis='y',
726 pad=self.config.tick_label_offset_em[0]*em)
728 self._remember_mpl_view_limits()
730 for mappable, dim in self._mappables.items():
731 mappable.set_clim(*limits[self._dim_index(dim)])
733 # scaler = plot.AutoScaler()
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)
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
763 rect = [xmin, ymin, xmax-xmin, ymax-ymin]
764 axes.set_position(
765 self.rect_to_figure_coords(rect), which='both')
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)
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)
777 finally:
778 self._updating_layout = False
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))
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)
791 def close(self):
792 self._disconnect_all()
793 self._plt.close(self._fig)
795 def show(self):
796 self._plt.show()
797 self.reset_size()
799 def set_label(self, dim, s):
800 # just set attribute, handle in update_layout
801 self._labels[dim] = s
803 def colorbar(
804 self, dim,
805 orientation='vertical',
806 position=None):
808 if dim not in self._dims:
809 raise PlotError(
810 'dimension "%s" is not defined')
812 if orientation not in ('vertical', 'horizontal'):
813 raise PlotError(
814 'orientation must be "vertical" or "horizontal"')
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())
824 if mappable is None:
825 raise PlotError(
826 'no mappable registered for dimension "%s"' % dim)
828 if position is None:
829 if orientation == 'vertical':
830 position = (0, self._shape[1])
831 else:
832 position = (0, self._shape[0])
834 em = self.config.font_size / self._pixels
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]
843 axes = SmartplotAxes(self.fig, [0., 0., 1., 1.])
844 self.fig.add_axes(axes)
846 self._colorbar_axes.append(
847 (axes, orientation, position))
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)
855 def __call__(self, *args):
856 return self.axes(*args)
859if __name__ == '__main__':
860 import sys
861 from pyrocko import util
863 logging.getLogger('matplotlib').setLevel(logging.WARNING)
864 util.setup_logging('smartplot', 'debug')
866 iplots = [int(x) for x in sys.argv[1:]]
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()
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()
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()
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()
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()
915 if 2 in iplots:
916 p = Plot(['easting', 'depth'], ['northing', 'depth'])
918 n = 100
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()
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)
932 n = 100
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()
943 if 5 in iplots:
944 p = Plot(['time'], ['northing', 'easting', '-depth'], ['depth'])
946 n = 100
948 t = num.arange(n)
949 xyz = num.random.normal(size=(n, 4))
950 xyz[:, 0] *= 0.5
952 smap = make_smap('summer')
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)
961 p.set_lim('depth', -1., 1.)
963 p.set_color_dim(smap, 'depth')
965 p.set_aspect('northing', 'easting', 1.0)
966 p.set_aspect('northing', 'depth', 1.0)
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]')
973 p.colorbar('depth')
975 p.show()
977 if 6 in iplots:
978 km = 1000.
979 p = Plot(
980 ['easting'], ['northing']*3, ['displacement'])
982 nn, ne = 50, 40
983 n = num.linspace(-5*km, 5*km, nn)
984 e = num.linspace(-10*km, 10*km, ne)
986 displacement = num.zeros((nn, ne, 3))
987 g = num.exp(
988 -(n[:, num.newaxis]**2 + e[num.newaxis, :]**2) / (5*km)**2)
990 displacement[:, :, 0] = g
991 displacement[:, :, 1] = g * 0.5
992 displacement[:, :, 2] = -g * 0.2
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')
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')
1005 p.set_lim('northing', -5.0, 5.0)
1006 p.set_lim('easting', -3.0, 3.0)
1007 p.show()