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
28class SmartplotAxes(Axes):
30 if matplotlib.__version__.split('.') < '3.6'.split('.'):
31 # Subclassing cla is deprecated on newer mpl but need this fallback for
32 # older versions. Code is duplicated because mpl behaviour depends
33 # on the existence of cla in the subclass...
34 def cla(self):
35 if hasattr(self, 'callbacks'):
36 callbacks = self.callbacks
37 Axes.cla(self)
38 self.callbacks = callbacks
39 else:
40 Axes.cla(self)
42 else:
43 def clear(self):
44 if hasattr(self, 'callbacks'):
45 callbacks = self.callbacks
46 Axes.clear(self)
47 self.callbacks = callbacks
48 else:
49 Axes.clear(self)
52class SmartplotFigure(figure.Figure):
54 def set_smartplot(self, plot):
55 self._smartplot = plot
57 def draw(self, *args, **kwargs):
58 if hasattr(self, '_smartplot'):
59 try:
60 self._smartplot._update_layout()
61 except NotEnoughSpace:
62 logger.error('Figure is too small to show the plot.')
63 return
65 return figure.Figure.draw(self, *args, **kwargs)
68def limits(points):
69 lims = num.zeros((3, 2))
70 if points.size != 0:
71 lims[:, 0] = num.min(points, axis=0)
72 lims[:, 1] = num.max(points, axis=0)
74 return lims
77def wcenter(rect):
78 return rect[0] + rect[2]*0.5
81def hcenter(rect):
82 return rect[1] + rect[3]*0.5
85def window_min(n, w, ml, mu, s, x):
86 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + math.floor(x) * s
89def window_max(n, w, ml, mu, s, x):
90 return ml + x/float(n) * (w - (ml + mu + (n-1)*s)) + (math.floor(x)-1) * s
93def make_smap(cmap, norm=None):
94 if isinstance(norm, tuple):
95 norm = colors.Normalize(*norm, clip=False)
96 smap = cm.ScalarMappable(cmap=cmap, norm=norm)
97 smap._A = [] # not needed in newer versions of mpl?
98 return smap
101def solve_layout_fixed_panels(size, shape, limits, aspects, fracs=None):
103 weight_aspect = 1000.
105 sx, sy = size
106 nx, ny = shape
107 nvar = nx+ny
108 vxs, vys = limits
109 uxs = vxs[:, 1] - vxs[:, 0]
110 uys = vys[:, 1] - vys[:, 0]
111 aspects_xx, aspects_yy, aspects_xy = aspects
113 if fracs is None:
114 wxs = num.full(nx, sx / nx)
115 wys = num.full(ny, sy / ny)
116 else:
117 frac_x, frac_y = fracs
118 wxs = sx * frac_x / num.sum(frac_x)
119 wys = sy * frac_y / num.sum(frac_y)
121 data = []
122 weights = []
123 rows = []
124 bounds = []
125 for ix in range(nx):
126 u = uxs[ix]
127 assert u > 0.0
128 row = num.zeros(nvar)
129 row[ix] = u
130 rows.append(row)
131 data.append(wxs[ix])
132 weights.append(1.0 / u)
133 bounds.append((0, wxs[ix] / u))
135 for iy in range(ny):
136 u = uys[iy]
137 assert u > 0.0
138 row = num.zeros(nvar)
139 row[nx+iy] = u
140 rows.append(row)
141 data.append(wys[iy])
142 weights.append(1.0)
143 bounds.append((0, wys[iy] / u))
145 for ix1, ix2, aspect in aspects_xx:
146 row = num.zeros(nvar)
147 row[ix1] = aspect
148 row[ix2] = -1.0
149 weights.append(weight_aspect/aspect)
150 rows.append(row)
151 data.append(0.0)
153 for iy1, iy2, aspect in aspects_yy:
154 row = num.zeros(nvar)
155 row[nx+iy1] = aspect
156 row[nx+iy2] = -1.0
157 weights.append(weight_aspect/aspect)
158 rows.append(row)
159 data.append(0.0)
161 for ix, iy, aspect in aspects_xy:
162 row = num.zeros(nvar)
163 row[ix] = aspect
164 row[nx+iy] = -1.0
165 weights.append(weight_aspect/aspect)
166 rows.append(row)
167 data.append(0.0)
169 weights = num.array(weights)
170 data = num.array(data)
171 mat = num.vstack(rows) * weights[:, num.newaxis]
172 data *= weights
174 bounds = num.array(bounds).T
176 model = scipy.optimize.lsq_linear(mat, data, bounds).x
178 cxs = model[:nx]
179 cys = model[nx:nx+ny]
181 vlimits_x = num.zeros((nx, 2))
182 for ix in range(nx):
183 u = wxs[ix] / cxs[ix]
184 vmin, vmax = vxs[ix]
185 udata = vmax - vmin
186 eps = 1e-7 * u
187 assert udata <= u + eps
188 vlimits_x[ix, 0] = (vmin + vmax) / 2.0 - u / 2.0
189 vlimits_x[ix, 1] = (vmin + vmax) / 2.0 + u / 2.0
191 vlimits_y = num.zeros((ny, 2))
192 for iy in range(ny):
193 u = wys[iy] / cys[iy]
194 vmin, vmax = vys[iy]
195 udata = vmax - vmin
196 eps = 1e-7 * u
197 assert udata <= u + eps
198 vlimits_y[iy, 0] = (vmin + vmax) / 2.0 - u / 2.0
199 vlimits_y[iy, 1] = (vmin + vmax) / 2.0 + u / 2.0
201 def check_aspect(a, awant, eps=1e-2):
202 if abs(1.0 - (a/awant)) > eps:
203 logger.error(
204 'Unable to comply with requested aspect ratio '
205 '(wanted: %g, achieved: %g)' % (awant, a))
207 for ix1, ix2, aspect in aspects_xx:
208 check_aspect(cxs[ix2] / cxs[ix1], aspect)
210 for iy1, iy2, aspect in aspects_yy:
211 check_aspect(cys[iy2] / cys[iy1], aspect)
213 for ix, iy, aspect in aspects_xy:
214 check_aspect(cys[iy] / cxs[ix], aspect)
216 return (vlimits_x, vlimits_y), (wxs, wys)
219def solve_layout_iterative(size, shape, limits, aspects, niterations=3):
221 sx, sy = size
222 nx, ny = shape
223 vxs, vys = limits
224 uxs = vxs[:, 1] - vxs[:, 0]
225 uys = vys[:, 1] - vys[:, 0]
226 aspects_xx, aspects_yy, aspects_xy = aspects
228 fracs_x, fracs_y = num.ones(nx), num.ones(ny)
229 for i in range(niterations):
230 (vlimits_x, vlimits_y), (wxs, wys) = solve_layout_fixed_panels(
231 size, shape, limits, aspects, (fracs_x, fracs_y))
233 uxs_view = vlimits_x[:, 1] - vlimits_x[:, 0]
234 uys_view = vlimits_y[:, 1] - vlimits_y[:, 0]
235 wxs_used = wxs * uxs / uxs_view
236 wys_used = wys * uys / uys_view
237 # wxs_wasted = wxs * (1.0 - uxs / uxs_view)
238 # wys_wasted = wys * (1.0 - uys / uys_view)
240 fracs_x = wxs_used
241 fracs_y = wys_used
243 return (vlimits_x, vlimits_y), (wxs, wys)
246class PlotError(Exception):
247 pass
250class NotEnoughSpace(PlotError):
251 pass
254class PlotConfig(Object):
256 font_size = Float.T(default=9.0)
258 size_cm = Tuple.T(
259 2, Float.T(), default=(20., 20.))
261 margins_em = Tuple.T(
262 4, Float.T(), default=(8., 6., 8., 6.))
264 separator_em = Float.T(default=1.5)
266 colorbar_width_em = Float.T(default=2.0)
268 label_offset_em = Tuple.T(
269 2, Float.T(), default=(2., 2.))
271 tick_label_offset_em = Tuple.T(
272 2, Float.T(), default=(0.5, 0.5))
274 @property
275 def size_inch(self):
276 return self.size_cm[0]/inch, self.size_cm[1]/inch
279class Plot(object):
281 def __init__(
282 self, x_dims=['x'], y_dims=['y'], z_dims=[], config=None,
283 fig=None, call_mpl_init=True):
285 if config is None:
286 config = PlotConfig()
288 self._shape = len(x_dims), len(y_dims)
290 dims = []
291 for dim in x_dims + y_dims + z_dims:
292 dim = dim.lstrip('-')
293 if dim not in dims:
294 dims.append(dim)
296 self.config = config
297 self._disconnect_data = []
298 self._width = self._height = self._pixels = None
299 if call_mpl_init:
300 self._plt = plot.mpl_init(self.config.font_size)
302 if fig is None:
303 fig = self._plt.figure(
304 figsize=self.config.size_inch, FigureClass=SmartplotFigure)
305 else:
306 assert isinstance(fig, SmartplotFigure)
308 fig.set_smartplot(self)
310 self._fig = fig
311 self._colorbar_width = 0.0
312 self._colorbar_height = 0.0
313 self._colorbar_axes = []
315 self._dims = dims
316 self._dim_index = self._dims.index
317 self._ndims = len(dims)
318 self._labels = {}
319 self._aspects = {}
321 self.setup_axes()
323 self._view_limits = num.zeros((self._ndims, 2))
325 self._view_limits[:, :] = num.nan
326 self._last_mpl_view_limits = None
328 self._x_dims = [dim.lstrip('-') for dim in x_dims]
329 self._x_dims_invert = [dim.startswith('-') for dim in x_dims]
331 self._y_dims = [dim.lstrip('-') for dim in y_dims]
332 self._y_dims_invert = [dim.startswith('-') for dim in y_dims]
334 self._z_dims = [dim.lstrip('-') for dim in z_dims]
335 self._z_dims_invert = [dim.startswith('-') for dim in z_dims]
337 self._mappables = {}
338 self._updating_layout = False
340 self._need_update_layout = True
341 self._update_geometry()
343 for axes in self.axes_list:
344 fig.add_axes(axes)
345 self._connect(axes, 'xlim_changed', self.lim_changed_handler)
346 self._connect(axes, 'ylim_changed', self.lim_changed_handler)
348 self._cid_resize = fig.canvas.mpl_connect(
349 'resize_event', self.resize_handler)
351 self._connect(fig, 'dpi_changed', self.dpi_changed_handler)
353 self._lim_changed_depth = 0
355 def reset_size(self):
356 self._fig.set_size_inches(self.config.size_inch)
358 def axes(self, ix, iy):
359 if not (isinstance(ix, int) and isinstance(iy, int)):
360 ix = self._x_dims.index(ix)
361 iy = self._y_dims.index(iy)
363 return self._axes[iy][ix]
365 def set_color_dim(self, mappable, dim):
366 assert dim in self._dims
367 self._mappables[mappable] = dim
369 def set_aspect(self, ydim, xdim, aspect=1.0):
370 self._aspects[ydim, xdim] = aspect
372 @property
373 def dims(self):
374 return self._dims
376 @property
377 def fig(self):
378 return self._fig
380 @property
381 def axes_list(self):
382 axes = []
383 for row in self._axes:
384 axes.extend(row)
385 return axes
387 @property
388 def axes_bottom_list(self):
389 return self._axes[0]
391 @property
392 def axes_left_list(self):
393 return [row[0] for row in self._axes]
395 def setup_axes(self):
396 rect = [0., 0., 1., 1.]
397 nx, ny = self._shape
398 axes = []
399 for iy in range(ny):
400 axes.append([])
401 for ix in range(nx):
402 axes[-1].append(SmartplotAxes(self.fig, rect))
404 self._axes = axes
406 for _, _, axes_ in self.iaxes():
407 axes_.set_autoscale_on(False)
409 def _connect(self, obj, sig, handler):
410 cid = obj.callbacks.connect(sig, handler)
411 self._disconnect_data.append((obj, cid))
413 def _disconnect_all(self):
414 for obj, cid in self._disconnect_data:
415 obj.callbacks.disconnect(cid)
417 self._fig.canvas.mpl_disconnect(self._cid_resize)
419 def dpi_changed_handler(self, fig):
420 if self._updating_layout:
421 return
423 self._update_geometry()
425 def resize_handler(self, event):
426 if self._updating_layout:
427 return
429 self._update_geometry()
431 def lim_changed_handler(self, axes):
432 if self._updating_layout:
433 return
435 current = self._get_mpl_view_limits()
436 last = self._last_mpl_view_limits
438 for iy, ix, axes in self.iaxes():
439 acurrent = current[iy][ix]
440 alast = last[iy][ix]
441 if acurrent[0] != alast[0]:
442 xdim = self._x_dims[ix]
443 logger.debug(
444 'X limits have been changed interactively in subplot '
445 '(%i, %i)' % (ix, iy))
446 self.set_lim(xdim, *sorted(acurrent[0]))
448 if acurrent[1] != alast[1]:
449 ydim = self._y_dims[iy]
450 logger.debug(
451 'Y limits have been changed interactively in subplot '
452 '(%i, %i)' % (ix, iy))
453 self.set_lim(ydim, *sorted(acurrent[1]))
455 self.need_update_layout()
457 def _update_geometry(self):
458 w, h = self._fig.canvas.get_width_height()
459 dp = self.get_device_pixel_ratio()
460 p = self.get_pixels_factor() * dp
462 if (self._width, self._height, self._pixels) != (w, h, p, dp):
463 logger.debug(
464 'New figure size: %g x %g, '
465 'logical-pixel/point: %g, physical-pixel/logical-pixel: %g' % (
466 w, h, p, dp))
468 self._width = w # logical pixel
469 self._height = h # logical pixel
470 self._pixels = p # logical pixel / point
471 self._device_pixel_ratio = dp # physical / logical
472 self.need_update_layout()
474 @property
475 def margins(self):
476 return tuple(
477 x * self.config.font_size / self._pixels
478 for x in self.config.margins_em)
480 @property
481 def separator(self):
482 return self.config.separator_em * self.config.font_size / self._pixels
484 def rect_to_figure_coords(self, rect):
485 left, bottom, width, height = rect
486 return (
487 left / self._width,
488 bottom / self._height,
489 width / self._width,
490 height / self._height)
492 def point_to_axes_coords(self, axes, point):
493 x, y = point
494 aleft, abottom, awidth, aheight = axes.get_position().bounds
496 x_fig = x / self._width
497 y_fig = y / self._height
499 x_axes = (x_fig - aleft) / awidth
500 y_axes = (y_fig - abottom) / aheight
502 return (x_axes, y_axes)
504 def get_pixels_factor(self):
505 try:
506 r = self._fig.canvas.get_renderer()
507 return 1.0 / r.points_to_pixels(1.0)
508 except AttributeError:
509 return 1.0
511 def get_device_pixel_ratio(self):
512 try:
513 return self._fig.canvas.device_pixel_ratio
514 except AttributeError:
515 return 1.0
517 def make_limits(self, lims):
518 a = plot.AutoScaler(space=0.05)
519 return a.make_scale(lims)[:2]
521 def iaxes(self):
522 for iy, row in enumerate(self._axes):
523 for ix, axes in enumerate(row):
524 yield iy, ix, axes
526 def get_data_limits(self):
527 dim_to_values = defaultdict(list)
528 for iy, ix, axes in self.iaxes():
529 dim_to_values[self._y_dims[iy]].extend(
530 axes.get_yaxis().get_data_interval())
531 dim_to_values[self._x_dims[ix]].extend(
532 axes.get_xaxis().get_data_interval())
534 for mappable, dim in self._mappables.items():
535 dim_to_values[dim].extend(mappable.get_clim())
537 lims = num.zeros((self._ndims, 2))
538 for idim in range(self._ndims):
539 dim = self._dims[idim]
540 if dim in dim_to_values:
541 vs = num.array(
542 dim_to_values[self._dims[idim]], dtype=float)
543 vs = vs[num.isfinite(vs)]
544 if vs.size > 0:
545 lims[idim, :] = num.min(vs), num.max(vs)
546 else:
547 lims[idim, :] = num.nan, num.nan
548 else:
549 lims[idim, :] = num.nan, num.nan
551 lims[num.logical_not(num.isfinite(lims))] = 0.0
552 return lims
554 def set_lim(self, dim, vmin, vmax):
555 assert vmin <= vmax
556 self._view_limits[self._dim_index(dim), :] = vmin, vmax
558 def _get_mpl_view_limits(self):
559 vl = []
560 for row in self._axes:
561 vl_row = []
562 for axes in row:
563 vl_row.append((
564 axes.get_xaxis().get_view_interval().tolist(),
565 axes.get_yaxis().get_view_interval().tolist()))
567 vl.append(vl_row)
569 return vl
571 def _remember_mpl_view_limits(self):
572 self._last_mpl_view_limits = self._get_mpl_view_limits()
574 def window_xmin(self, x):
575 return window_min(
576 self._shape[0], self._width,
577 self.margins[0], self.margins[2] + self._colorbar_width,
578 self.separator, x)
580 def window_xmax(self, x):
581 return window_max(
582 self._shape[0], self._width,
583 self.margins[0], self.margins[2] + self._colorbar_width,
584 self.separator, x)
586 def window_ymin(self, y):
587 return window_min(
588 self._shape[1], self._height,
589 self.margins[3] + self._colorbar_height, self.margins[1],
590 self.separator, y)
592 def window_ymax(self, y):
593 return window_max(
594 self._shape[1], self._height,
595 self.margins[3] + self._colorbar_height, self.margins[1],
596 self.separator, y)
598 def need_update_layout(self):
599 self._need_update_layout = True
601 def _update_layout(self):
602 assert not self._updating_layout
604 if not self._need_update_layout:
605 return
607 self._updating_layout = True
608 try:
609 data_limits = self.get_data_limits()
611 limits = num.zeros((self._ndims, 2))
612 for idim in range(self._ndims):
613 limits[idim, :] = self.make_limits(data_limits[idim, :])
615 mask = num.isfinite(self._view_limits)
616 limits[mask] = self._view_limits[mask]
618 # deltas = limits[:, 1] - limits[:, 0]
620 # data_w = deltas[0]
621 # data_h = deltas[1]
623 ml, mt, mr, mb = self.margins
624 mr += self._colorbar_width
625 mb += self._colorbar_height
626 sw = sh = self.separator
628 nx, ny = self._shape
630 # data_r = data_h / data_w
631 em = self.config.font_size / self._pixels
632 w = self._width
633 h = self._height
634 fig_w_avail = w - mr - ml - (nx-1) * sw
635 fig_h_avail = h - mt - mb - (ny-1) * sh
637 if fig_w_avail <= 0.0 or fig_h_avail <= 0.0:
638 raise NotEnoughSpace()
640 x_limits = num.zeros((nx, 2))
641 for ix, xdim in enumerate(self._x_dims):
642 x_limits[ix, :] = limits[self._dim_index(xdim)]
644 y_limits = num.zeros((ny, 2))
645 for iy, ydim in enumerate(self._y_dims):
646 y_limits[iy, :] = limits[self._dim_index(ydim)]
648 def get_aspect(dim1, dim2):
649 if (dim2, dim1) in self._aspects:
650 return 1.0/self._aspects[dim2, dim1]
652 return self._aspects.get((dim1, dim2), None)
654 aspects_xx = []
655 for ix1, xdim1 in enumerate(self._x_dims):
656 for ix2, xdim2 in enumerate(self._x_dims):
657 aspect = get_aspect(xdim2, xdim1)
658 if aspect:
659 aspects_xx.append((ix1, ix2, aspect))
661 aspects_yy = []
662 for iy1, ydim1 in enumerate(self._y_dims):
663 for iy2, ydim2 in enumerate(self._y_dims):
664 aspect = get_aspect(ydim2, ydim1)
665 if aspect:
666 aspects_yy.append((iy1, iy2, aspect))
668 aspects_xy = []
669 for iy, ix, axes in self.iaxes():
670 xdim = self._x_dims[ix]
671 ydim = self._y_dims[iy]
672 aspect = get_aspect(ydim, xdim)
673 if aspect:
674 aspects_xy.append((ix, iy, aspect))
676 (x_limits, y_limits), (aws, ahs) = solve_layout_iterative(
677 size=(fig_w_avail, fig_h_avail),
678 shape=(nx, ny),
679 limits=(x_limits, y_limits),
680 aspects=(
681 aspects_xx,
682 aspects_yy,
683 aspects_xy))
685 for iy, ix, axes in self.iaxes():
686 rect = [
687 ml + num.sum(aws[:ix])+(ix*sw),
688 mb + num.sum(ahs[:iy])+(iy*sh),
689 aws[ix], ahs[iy]]
691 axes.set_position(
692 self.rect_to_figure_coords(rect), which='both')
694 self.set_label_coords(
695 axes, 'x', [
696 wcenter(rect),
697 self.config.label_offset_em[0]*em
698 + self._colorbar_height])
700 self.set_label_coords(
701 axes, 'y', [
702 self.config.label_offset_em[1]*em,
703 hcenter(rect)])
705 axes.get_xaxis().set_tick_params(
706 bottom=(iy == 0), top=(iy == ny-1),
707 labelbottom=(iy == 0), labeltop=False)
709 axes.get_yaxis().set_tick_params(
710 left=(ix == 0), right=(ix == nx-1),
711 labelleft=(ix == 0), labelright=False)
713 istride = -1 if self._x_dims_invert[ix] else 1
714 axes.set_xlim(*x_limits[ix, ::istride])
715 istride = -1 if self._y_dims_invert[iy] else 1
716 axes.set_ylim(*y_limits[iy, ::istride])
718 axes.tick_params(
719 axis='x',
720 pad=self.config.tick_label_offset_em[0]*em)
722 axes.tick_params(
723 axis='y',
724 pad=self.config.tick_label_offset_em[0]*em)
726 self._remember_mpl_view_limits()
728 for mappable, dim in self._mappables.items():
729 mappable.set_clim(*limits[self._dim_index(dim)])
731 # scaler = plot.AutoScaler()
733 # aspect tick incs same
734 #
735 # inc = scaler.make_scale(
736 # [0, min(data_expanded_w, data_expanded_h)],
737 # override_mode='off')[2]
738 #
739 # for axes in self.axes_list:
740 # axes.set_xlim(*limits[0, :])
741 # axes.set_ylim(*limits[1, :])
742 #
743 # tl = MultipleLocator(inc)
744 # axes.get_xaxis().set_major_locator(tl)
745 # tl = MultipleLocator(inc)
746 # axes.get_yaxis().set_major_locator(tl)
748 for axes, orientation, position in self._colorbar_axes:
749 if orientation == 'horizontal':
750 xmin = self.window_xmin(position[0])
751 xmax = self.window_xmax(position[1])
752 ymin = mb - self._colorbar_height
753 ymax = mb - self._colorbar_height \
754 + self.config.colorbar_width_em * em
755 else:
756 ymin = self.window_ymin(position[0])
757 ymax = self.window_ymax(position[1])
758 xmin = w - mr + 2 * sw
759 xmax = w - mr + 2 * sw + self.config.colorbar_width_em * em
761 rect = [xmin, ymin, xmax-xmin, ymax-ymin]
762 axes.set_position(
763 self.rect_to_figure_coords(rect), which='both')
765 for ix, axes in enumerate(self.axes_bottom_list):
766 dim = self._x_dims[ix]
767 s = self._labels.get(dim, dim)
768 axes.set_xlabel(s)
770 for iy, axes in enumerate(self.axes_left_list):
771 dim = self._y_dims[iy]
772 s = self._labels.get(dim, dim)
773 axes.set_ylabel(s)
775 finally:
776 self._updating_layout = False
778 def set_label_coords(self, axes, which, point):
779 axis = axes.get_xaxis() if which == 'x' else axes.get_yaxis()
780 axis.set_label_coords(*self.point_to_axes_coords(axes, point))
782 def plot(self, points, *args, **kwargs):
783 for iy, row in enumerate(self._axes):
784 y = points[:, self._dim_index(self._y_dims[iy])]
785 for ix, axes in enumerate(row):
786 x = points[:, self._dim_index(self._x_dims[ix])]
787 axes.plot(x, y, *args, **kwargs)
789 def close(self):
790 self._disconnect_all()
791 self._plt.close(self._fig)
793 def show(self):
794 self._plt.show()
795 self.reset_size()
797 def set_label(self, dim, s):
798 # just set attribute, handle in update_layout
799 self._labels[dim] = s
801 def colorbar(
802 self, dim,
803 orientation='vertical',
804 position=None):
806 if dim not in self._dims:
807 raise PlotError(
808 'dimension "%s" is not defined')
810 if orientation not in ('vertical', 'horizontal'):
811 raise PlotError(
812 'orientation must be "vertical" or "horizontal"')
814 mappable = None
815 for mappable_, dim_ in self._mappables.items():
816 if dim_ == dim:
817 if mappable is None:
818 mappable = mappable_
819 else:
820 mappable_.set_cmap(mappable.get_cmap())
822 if mappable is None:
823 raise PlotError(
824 'no mappable registered for dimension "%s"' % dim)
826 if position is None:
827 if orientation == 'vertical':
828 position = (0, self._shape[1])
829 else:
830 position = (0, self._shape[0])
832 em = self.config.font_size / self._pixels
834 if orientation == 'vertical':
835 self._colorbar_width = self.config.colorbar_width_em*em + \
836 self.separator * 2.0
837 else:
838 self._colorbar_height = self.config.colorbar_width_em*em + \
839 self.separator + self.margins[3]
841 axes = SmartplotAxes(self.fig, [0., 0., 1., 1.])
842 self.fig.add_axes(axes)
844 self._colorbar_axes.append(
845 (axes, orientation, position))
847 self.need_update_layout()
848 # axes.plot([1], [1])
849 label = self._labels.get(dim, dim)
850 return colorbar.Colorbar(
851 axes, mappable, orientation=orientation, label=label)
853 def __call__(self, *args):
854 return self.axes(*args)
857if __name__ == '__main__':
858 import sys
859 from pyrocko import util
861 logging.getLogger('matplotlib').setLevel(logging.WARNING)
862 util.setup_logging('smartplot', 'debug')
864 iplots = [int(x) for x in sys.argv[1:]]
866 if 0 in iplots:
867 p = Plot(['x'], ['y'])
868 n = 100
869 x = num.arange(n) * 2.0
870 y = num.random.normal(size=n)
871 p(0, 0).plot(x, y, 'o')
872 p.show()
874 if 1 in iplots:
875 p = Plot(['x', 'x'], ['y'])
876 n = 100
877 x = num.arange(n) * 2.0
878 y = num.random.normal(size=n)
879 p(0, 0).plot(x, y, 'o')
880 x = num.arange(n) * 2.0
881 y = num.random.normal(size=n)
882 p(1, 0).plot(x, y, 'o')
883 p.show()
885 if 11 in iplots:
886 p = Plot(['x'], ['y'])
887 p.set_aspect('y', 'x', 2.0)
888 n = 100
889 xy = num.random.normal(size=(n, 2))
890 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o')
891 p.show()
893 if 12 in iplots:
894 p = Plot(['x', 'x2'], ['y'])
895 p.set_aspect('x2', 'x', 2.0)
896 p.set_aspect('y', 'x', 2.0)
897 n = 100
898 xy = num.random.normal(size=(n, 2))
899 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o')
900 p(1, 0).plot(xy[:, 0], xy[:, 1], 'o')
901 p.show()
903 if 13 in iplots:
904 p = Plot(['x'], ['y', 'y2'])
905 p.set_aspect('y2', 'y', 2.0)
906 p.set_aspect('y', 'x', 2.0)
907 n = 100
908 xy = num.random.normal(size=(n, 2))
909 p(0, 0).plot(xy[:, 0], xy[:, 1], 'o')
910 p(0, 1).plot(xy[:, 0], xy[:, 1], 'o')
911 p.show()
913 if 2 in iplots:
914 p = Plot(['easting', 'depth'], ['northing', 'depth'])
916 n = 100
918 ned = num.random.normal(size=(n, 3))
919 p(0, 0).plot(ned[:, 1], ned[:, 0], 'o')
920 p(1, 0).plot(ned[:, 2], ned[:, 0], 'o')
921 p(0, 1).plot(ned[:, 1], ned[:, 2], 'o')
922 p.show()
924 if 3 in iplots:
925 p = Plot(['easting', 'depth'], ['-depth', 'northing'])
926 p.set_aspect('easting', 'northing', 1.0)
927 p.set_aspect('easting', 'depth', 0.5)
928 p.set_aspect('northing', 'depth', 0.5)
930 n = 100
932 ned = num.random.normal(size=(n, 3))
933 ned[:, 2] *= 0.25
934 p(0, 1).plot(ned[:, 1], ned[:, 0], 'o', color='black')
935 p(0, 0).plot(ned[:, 1], ned[:, 2], 'o')
936 p(1, 1).plot(ned[:, 2], ned[:, 0], 'o')
937 p(1, 0).set_visible(False)
938 p.set_lim('depth', 0., 0.2)
939 p.show()
941 if 5 in iplots:
942 p = Plot(['time'], ['northing', 'easting', '-depth'], ['depth'])
944 n = 100
946 t = num.arange(n)
947 xyz = num.random.normal(size=(n, 4))
948 xyz[:, 0] *= 0.5
950 smap = make_smap('summer')
952 p(0, 0).scatter(
953 t, xyz[:, 0], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
954 p(0, 1).scatter(
955 t, xyz[:, 1], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
956 p(0, 2).scatter(
957 t, xyz[:, 2], c=xyz[:, 2], cmap=smap.cmap, norm=smap.norm)
959 p.set_lim('depth', -1., 1.)
961 p.set_color_dim(smap, 'depth')
963 p.set_aspect('northing', 'easting', 1.0)
964 p.set_aspect('northing', 'depth', 1.0)
966 p.set_label('time', 'Time [s]')
967 p.set_label('depth', 'Depth [km]')
968 p.set_label('easting', 'Easting [km]')
969 p.set_label('northing', 'Northing [km]')
971 p.colorbar('depth')
973 p.show()
975 if 6 in iplots:
976 km = 1000.
977 p = Plot(
978 ['easting'], ['northing']*3, ['displacement'])
980 nn, ne = 50, 40
981 n = num.linspace(-5*km, 5*km, nn)
982 e = num.linspace(-10*km, 10*km, ne)
984 displacement = num.zeros((nn, ne, 3))
985 g = num.exp(
986 -(n[:, num.newaxis]**2 + e[num.newaxis, :]**2) / (5*km)**2)
988 displacement[:, :, 0] = g
989 displacement[:, :, 1] = g * 0.5
990 displacement[:, :, 2] = -g * 0.2
992 for icomp in (0, 1, 2):
993 c = p(0, icomp).pcolormesh(
994 e/km, n/km, displacement[:, :, icomp], shading='gouraud')
995 p.set_color_dim(c, 'displacement')
997 p.colorbar('displacement')
998 p.set_lim('displacement', -1.0, 1.0)
999 p.set_label('easting', 'Easting [km]')
1000 p.set_label('northing', 'Northing [km]')
1001 p.set_aspect('northing', 'easting')
1003 p.set_lim('northing', -5.0, 5.0)
1004 p.set_lim('easting', -3.0, 3.0)
1005 p.show()