1# http://pyrocko.org - GPLv3
2#
3# The Pyrocko Developers, 21st Century
4# ---|P------/S----------~Lg----------
6'''
7Utility functions and defintions for a common plot style throughout Pyrocko.
9Functions with name prefix ``mpl_`` are Matplotlib specific. All others should
10be toolkit-agnostic.
12The following skeleton can be used to produce nice PDF figures, with absolute
13sizes derived from paper and font sizes
14(file :file:`/../../examples/plot_skeleton.py`
15in the Pyrocko source directory)::
17 from matplotlib import pyplot as plt
19 from pyrocko.plot import mpl_init, mpl_margins, mpl_papersize
20 # from pyrocko.plot import mpl_labelspace
22 fontsize = 9. # in points
24 # set some Pyrocko style defaults
25 mpl_init(fontsize=fontsize)
27 fig = plt.figure(figsize=mpl_papersize('a4', 'landscape'))
29 # let margins be proportional to selected font size, e.g. top and bottom
30 # margin are set to be 5*fontsize = 45 [points]
31 labelpos = mpl_margins(fig, w=7., h=5., units=fontsize)
33 axes = fig.add_subplot(1, 1, 1)
35 # positioning of axis labels
36 # mpl_labelspace(axes) # either: relative to axis tick labels
37 labelpos(axes, 2., 1.5) # or: relative to left/bottom paper edge
39 axes.plot([0, 1], [0, 9])
41 axes.set_xlabel('Time [s]')
42 axes.set_ylabel('Amplitude [m]')
44 fig.savefig('plot_skeleton.pdf')
46 plt.show()
48'''
49from __future__ import absolute_import
50from pyrocko.util import parse_md
51from pyrocko.guts import StringChoice, Float, Int, Bool, Tuple, Object
53import math
54import random
57try:
58 newstr = unicode
59except NameError:
60 newstr = str
63__doc__ += parse_md(__file__)
66guts_prefix = 'pf'
68point = 1.
69inch = 72.
70cm = 28.3465
72units_dict = {
73 'point': point,
74 'inch': inch,
75 'cm': cm,
76}
78_doc_units = "``'points'``, ``'inch'``, or ``'cm'``"
81def apply_units(x, units):
82 if isinstance(units, (str, newstr)):
83 units = units_dict[units]
85 if isinstance(x, (int, float)):
86 return x / units
87 else:
88 if isinstance(x, tuple):
89 return tuple(v / units for v in x)
90 else:
91 return list(v / units for v in x)
94tango_colors = {
95 'butter1': (252, 233, 79),
96 'butter2': (237, 212, 0),
97 'butter3': (196, 160, 0),
98 'chameleon1': (138, 226, 52),
99 'chameleon2': (115, 210, 22),
100 'chameleon3': (78, 154, 6),
101 'orange1': (252, 175, 62),
102 'orange2': (245, 121, 0),
103 'orange3': (206, 92, 0),
104 'skyblue1': (114, 159, 207),
105 'skyblue2': (52, 101, 164),
106 'skyblue3': (32, 74, 135),
107 'plum1': (173, 127, 168),
108 'plum2': (117, 80, 123),
109 'plum3': (92, 53, 102),
110 'chocolate1': (233, 185, 110),
111 'chocolate2': (193, 125, 17),
112 'chocolate3': (143, 89, 2),
113 'scarletred1': (239, 41, 41),
114 'scarletred2': (204, 0, 0),
115 'scarletred3': (164, 0, 0),
116 'aluminium1': (238, 238, 236),
117 'aluminium2': (211, 215, 207),
118 'aluminium3': (186, 189, 182),
119 'aluminium4': (136, 138, 133),
120 'aluminium5': (85, 87, 83),
121 'aluminium6': (46, 52, 54)}
124graph_colors = [
125 tango_colors[_x] for _x in (
126 'scarletred2',
127 'skyblue3',
128 'chameleon3',
129 'orange2',
130 'plum2',
131 'chocolate2',
132 'butter2')]
135def color(x=None):
136 if x is None:
137 return tuple([random.randint(0, 255) for _x in 'rgb'])
139 if isinstance(x, int):
140 if 0 <= x < len(graph_colors):
141 return graph_colors[x]
142 else:
143 return (0, 0, 0)
145 elif isinstance(x, (str, newstr)):
146 if x in tango_colors:
147 return tango_colors[x]
149 elif isinstance(x, tuple):
150 return x
152 assert False, "Don't know what to do with this color definition: %s" % x
155def to01(c):
156 return tuple(x/255. for x in c)
159def nice_value(x):
160 '''
161 Round x to nice value.
162 '''
164 if x == 0.0:
165 return 0.0
167 exp = 1.0
168 sign = 1
169 if x < 0.0:
170 x = -x
171 sign = -1
172 while x >= 1.0:
173 x /= 10.0
174 exp *= 10.0
175 while x < 0.1:
176 x *= 10.0
177 exp /= 10.0
179 if x >= 0.75:
180 return sign * 1.0 * exp
181 if x >= 0.35:
182 return sign * 0.5 * exp
183 if x >= 0.15:
184 return sign * 0.2 * exp
186 return sign * 0.1 * exp
189_papersizes_list = [
190 ('a0', (2380., 3368.)),
191 ('a1', (1684., 2380.)),
192 ('a2', (1190., 1684.)),
193 ('a3', (842., 1190.)),
194 ('a4', (595., 842.)),
195 ('a5', (421., 595.)),
196 ('a6', (297., 421.)),
197 ('a7', (210., 297.)),
198 ('a8', (148., 210.)),
199 ('a9', (105., 148.)),
200 ('a10', (74., 105.)),
201 ('b0', (2836., 4008.)),
202 ('b1', (2004., 2836.)),
203 ('b2', (1418., 2004.)),
204 ('b3', (1002., 1418.)),
205 ('b4', (709., 1002.)),
206 ('b5', (501., 709.)),
207 ('archa', (648., 864.)),
208 ('archb', (864., 1296.)),
209 ('archc', (1296., 1728.)),
210 ('archd', (1728., 2592.)),
211 ('arche', (2592., 3456.)),
212 ('flsa', (612., 936.)),
213 ('halfletter', (396., 612.)),
214 ('note', (540., 720.)),
215 ('letter', (612., 792.)),
216 ('legal', (612., 1008.)),
217 ('11x17', (792., 1224.)),
218 ('ledger', (1224., 792.))]
220papersizes = dict(_papersizes_list)
222_doc_papersizes = ', '.join("``'%s'``" % k for (k, _) in _papersizes_list)
225def papersize(paper, orientation='landscape', units='point'):
227 '''
228 Get paper size from string.
230 :param paper: string selecting paper size. Choices: %s
231 :param orientation: ``'landscape'``, or ``'portrait'``
232 :param units: Units to be returned. Choices: %s
234 :returns: ``(width, height)``
235 '''
237 assert orientation in ('landscape', 'portrait')
239 w, h = papersizes[paper.lower()]
240 if orientation == 'landscape':
241 w, h = h, w
243 return apply_units((w, h), units)
246papersize.__doc__ %= (_doc_papersizes, _doc_units)
249class AutoScaleMode(StringChoice):
250 '''
251 Mode of operation for auto-scaling.
253 ================ ==================================================
254 mode description
255 ================ ==================================================
256 ``'auto'``: Look at data range and choose one of the choices
257 below.
258 ``'min-max'``: Output range is selected to include data range.
259 ``'0-max'``: Output range shall start at zero and end at data
260 max.
261 ``'min-0'``: Output range shall start at data min and end at
262 zero.
263 ``'symmetric'``: Output range shall by symmetric by zero.
264 ``'off'``: Similar to ``'min-max'``, but snap and space are
265 disabled, such that the output range always
266 exactly matches the data range.
267 ================ ==================================================
268 '''
269 choices = ['auto', 'min-max', '0-max', 'min-0', 'symmetric', 'off']
272class AutoScaler(Object):
274 '''
275 Tunable 1D autoscaling based on data range.
277 Instances of this class may be used to determine nice minima, maxima and
278 increments for ax annotations, as well as suitable common exponents for
279 notation.
281 The autoscaling process is guided by the following public attributes:
282 '''
284 approx_ticks = Float.T(
285 default=7.0,
286 help='Approximate number of increment steps (tickmarks) to generate.')
288 mode = AutoScaleMode.T(
289 default='auto',
290 help='''Mode of operation for auto-scaling.''')
292 exp = Int.T(
293 optional=True,
294 help='If defined, override automatically determined exponent for '
295 'notation by the given value.')
297 snap = Bool.T(
298 default=False,
299 help='If set to True, snap output range to multiples of increment. '
300 'This parameter has no effect, if mode is set to ``\'off\'``.')
302 inc = Float.T(
303 optional=True,
304 help='If defined, override automatically determined tick increment by '
305 'the given value.')
307 space = Float.T(
308 default=0.0,
309 help='Add some padding to the range. The value given, is the fraction '
310 'by which the output range is increased on each side. If mode is '
311 '``\'0-max\'`` or ``\'min-0\'``, the end at zero is kept fixed '
312 'at zero. This parameter has no effect if mode is set to '
313 '``\'off\'``.')
315 exp_factor = Int.T(
316 default=3,
317 help='Exponent of notation is chosen to be a multiple of this value.')
319 no_exp_interval = Tuple.T(
320 2, Int.T(),
321 default=(-3, 5),
322 help='Range of exponent, for which no exponential notation is a'
323 'allowed.')
325 def __init__(
326 self,
327 approx_ticks=7.0,
328 mode='auto',
329 exp=None,
330 snap=False,
331 inc=None,
332 space=0.0,
333 exp_factor=3,
334 no_exp_interval=(-3, 5)):
336 '''
337 Create new AutoScaler instance.
339 The parameters are described in the AutoScaler documentation.
340 '''
342 Object.__init__(
343 self,
344 approx_ticks=approx_ticks,
345 mode=mode,
346 exp=exp,
347 snap=snap,
348 inc=inc,
349 space=space,
350 exp_factor=exp_factor,
351 no_exp_interval=no_exp_interval)
353 def make_scale(self, data_range, override_mode=None):
355 '''
356 Get nice minimum, maximum and increment for given data range.
358 Returns ``(minimum, maximum, increment)`` or ``(maximum, minimum,
359 -increment)``, depending on whether data_range is ``(data_min,
360 data_max)`` or ``(data_max, data_min)``. If ``override_mode`` is
361 defined, the mode attribute is temporarily overridden by the given
362 value.
363 '''
365 data_min = min(data_range)
366 data_max = max(data_range)
368 is_reverse = (data_range[0] > data_range[1])
370 a = self.mode
371 if self.mode == 'auto':
372 a = self.guess_autoscale_mode(data_min, data_max)
374 if override_mode is not None:
375 a = override_mode
377 mi, ma = 0, 0
378 if a == 'off':
379 mi, ma = data_min, data_max
380 elif a == '0-max':
381 mi = 0.0
382 if data_max > 0.0:
383 ma = data_max
384 else:
385 ma = 1.0
386 elif a == 'min-0':
387 ma = 0.0
388 if data_min < 0.0:
389 mi = data_min
390 else:
391 mi = -1.0
392 elif a == 'min-max':
393 mi, ma = data_min, data_max
394 elif a == 'symmetric':
395 m = max(abs(data_min), abs(data_max))
396 mi = -m
397 ma = m
399 nmi = mi
400 if (mi != 0. or a == 'min-max') and a != 'off':
401 nmi = mi - self.space*(ma-mi)
403 nma = ma
404 if (ma != 0. or a == 'min-max') and a != 'off':
405 nma = ma + self.space*(ma-mi)
407 mi, ma = nmi, nma
409 if mi == ma and a != 'off':
410 mi -= 1.0
411 ma += 1.0
413 # make nice tick increment
414 if self.inc is not None:
415 inc = self.inc
416 else:
417 if self.approx_ticks > 0.:
418 inc = nice_value((ma-mi) / self.approx_ticks)
419 else:
420 inc = nice_value((ma-mi)*10.)
422 if inc == 0.0:
423 inc = 1.0
425 # snap min and max to ticks if this is wanted
426 if self.snap and a != 'off':
427 ma = inc * math.ceil(ma/inc)
428 mi = inc * math.floor(mi/inc)
430 if is_reverse:
431 return ma, mi, -inc
432 else:
433 return mi, ma, inc
435 def make_exp(self, x):
436 '''
437 Get nice exponent for notation of ``x``.
439 For ax annotations, give tick increment as ``x``.
440 '''
442 if self.exp is not None:
443 return self.exp
445 x = abs(x)
446 if x == 0.0:
447 return 0
449 if 10**self.no_exp_interval[0] <= x <= 10**self.no_exp_interval[1]:
450 return 0
452 return math.floor(math.log10(x)/self.exp_factor)*self.exp_factor
454 def guess_autoscale_mode(self, data_min, data_max):
455 '''
456 Guess mode of operation, based on data range.
458 Used to map ``'auto'`` mode to ``'0-max'``, ``'min-0'``, ``'min-max'``
459 or ``'symmetric'``.
460 '''
462 a = 'min-max'
463 if data_min >= 0.0:
464 if data_min < data_max/2.:
465 a = '0-max'
466 else:
467 a = 'min-max'
468 if data_max <= 0.0:
469 if data_max > data_min/2.:
470 a = 'min-0'
471 else:
472 a = 'min-max'
473 if data_min < 0.0 and data_max > 0.0:
474 if abs((abs(data_max)-abs(data_min)) /
475 (abs(data_max)+abs(data_min))) < 0.5:
476 a = 'symmetric'
477 else:
478 a = 'min-max'
479 return a
482# below, some convenience functions for matplotlib plotting
484def mpl_init(fontsize=10):
485 '''
486 Initialize Matplotlib rc parameters Pyrocko style.
488 Returns the matplotlib.pyplot module for convenience.
489 '''
491 import matplotlib
493 matplotlib.rcdefaults()
494 matplotlib.rc('font', size=fontsize)
495 matplotlib.rc('axes', linewidth=1.5)
496 matplotlib.rc('xtick', direction='out')
497 matplotlib.rc('ytick', direction='out')
498 ts = fontsize * 0.7071
499 matplotlib.rc('xtick.major', size=ts, width=0.5, pad=ts)
500 matplotlib.rc('ytick.major', size=ts, width=0.5, pad=ts)
501 matplotlib.rc('figure', facecolor='white')
503 try:
504 from cycler import cycler
505 matplotlib.rc(
506 'axes', prop_cycle=cycler(
507 'color', [to01(x) for x in graph_colors]))
508 except (ImportError, KeyError):
509 try:
510 matplotlib.rc('axes', color_cycle=[to01(x) for x in graph_colors])
511 except KeyError:
512 pass
514 from matplotlib import pyplot as plt
515 return plt
518def mpl_margins(
519 fig,
520 left=1.0, top=1.0, right=1.0, bottom=1.0,
521 wspace=None, hspace=None,
522 w=None, h=None,
523 nw=None, nh=None,
524 all=None,
525 units='inch'):
527 '''
528 Adjust Matplotlib subplot params with absolute values in user units.
530 Calls :py:meth:`matplotlib.figure.Figure.subplots_adjust` on ``fig`` with
531 absolute margin widths/heights rather than relative values. If ``wspace``
532 or ``hspace`` are given, the number of subplots must be given in ``nw``
533 and ``nh`` because ``subplots_adjust()`` treats the spacing parameters
534 relative to the subplot width and height.
536 :param units: Unit multiplier or unit as string: %s
537 :param left,right,top,bottom: margin space
538 :param w: set ``left`` and ``right`` at once
539 :param h: set ``top`` and ``bottom`` at once
540 :param all: set ``left``, ``top``, ``right``, and ``bottom`` at once
541 :param nw: number of subplots horizontally
542 :param nh: number of subplots vertically
543 :param wspace: horizontal spacing between subplots
544 :param hspace: vertical spacing between subplots
545 '''
547 left, top, right, bottom = map(
548 float, (left, top, right, bottom))
550 if w is not None:
551 left = right = float(w)
553 if h is not None:
554 top = bottom = float(h)
556 if all is not None:
557 left = right = top = bottom = float(all)
559 ufac = units_dict.get(units, units) / inch
561 left *= ufac
562 right *= ufac
563 top *= ufac
564 bottom *= ufac
566 width, height = fig.get_size_inches()
568 rel_wspace = None
569 rel_hspace = None
571 if wspace is not None:
572 wspace *= ufac
573 if nw is None:
574 raise ValueError('wspace must be given in combination with nw')
576 wsub = (width - left - right - (nw-1) * wspace) / nw
577 rel_wspace = wspace / wsub
578 else:
579 wsub = width - left - right
581 if hspace is not None:
582 hspace *= ufac
583 if nh is None:
584 raise ValueError('hspace must be given in combination with nh')
586 hsub = (height - top - bottom - (nh-1) * hspace) / nh
587 rel_hspace = hspace / hsub
588 else:
589 hsub = height - top - bottom
591 fig.subplots_adjust(
592 left=left/width,
593 right=1.0 - right/width,
594 bottom=bottom/height,
595 top=1.0 - top/height,
596 wspace=rel_wspace,
597 hspace=rel_hspace)
599 def labelpos(axes, xpos=0., ypos=0.):
600 xpos *= ufac
601 ypos *= ufac
602 axes.get_yaxis().set_label_coords(-((left-xpos) / wsub), 0.5)
603 axes.get_xaxis().set_label_coords(0.5, -((bottom-ypos) / hsub))
605 return labelpos
608mpl_margins.__doc__ %= _doc_units
611def mpl_labelspace(axes):
612 '''
613 Add some extra padding between label and ax annotations.
614 '''
616 xa = axes.get_xaxis()
617 ya = axes.get_yaxis()
618 for attr in ('labelpad', 'LABELPAD'):
619 if hasattr(xa, attr):
620 setattr(xa, attr, xa.get_label().get_fontsize())
621 setattr(ya, attr, ya.get_label().get_fontsize())
622 break
625def mpl_papersize(paper, orientation='landscape'):
626 '''
627 Get paper size in inch from string.
629 Returns argument suitable to be passed to the ``figsize`` argument of
630 :py:func:`pyplot.figure`.
632 :param paper: string selecting paper size. Choices: %s
633 :param orientation: ``'landscape'``, or ``'portrait'``
635 :returns: ``(width, height)``
636 '''
638 return papersize(paper, orientation=orientation, units='inch')
641mpl_papersize.__doc__ %= _doc_papersizes
644class InvalidColorDef(ValueError):
645 pass
648def mpl_graph_color(i):
649 return to01(graph_colors[i % len(graph_colors)])
652def mpl_color(x):
653 '''
654 Convert string into color float tuple ranged 0-1 for use with Matplotlib.
656 Accepts tango color names, matplotlib color names, and slash-separated
657 strings. In the latter case, if values are larger than 1., the color
658 is interpreted as 0-255 ranged. Single-valued (grayscale), three-valued
659 (color) and four-valued (color with alpha) are accepted. An
660 :py:exc:`InvalidColorDef` exception is raised when the convertion fails.
661 '''
663 import matplotlib.colors
665 if x in tango_colors:
666 return to01(tango_colors[x])
668 s = x.split('/')
669 if len(s) in (1, 3, 4):
670 try:
671 vals = list(map(float, s))
672 if all(0. <= v <= 1. for v in vals):
673 return vals
675 elif all(0. <= v <= 255. for v in vals):
676 return to01(vals)
678 except ValueError:
679 try:
680 return matplotlib.colors.colorConverter.to_rgba(x)
681 except Exception:
682 pass
684 raise InvalidColorDef('invalid color definition: %s' % x)