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
51import math
52import random
53import time
54import calendar
55import numpy as num
57from pyrocko.util import parse_md, time_to_str, arange2, to_time_float
58from pyrocko.guts import StringChoice, Float, Int, Bool, Tuple, Object
61try:
62 newstr = unicode
63except NameError:
64 newstr = str
67__doc__ += parse_md(__file__)
70guts_prefix = 'pf'
72point = 1.
73inch = 72.
74cm = 28.3465
76units_dict = {
77 'point': point,
78 'inch': inch,
79 'cm': cm,
80}
82_doc_units = "``'points'``, ``'inch'``, or ``'cm'``"
85def apply_units(x, units):
86 if isinstance(units, (str, newstr)):
87 units = units_dict[units]
89 if isinstance(x, (int, float)):
90 return x / units
91 else:
92 if isinstance(x, tuple):
93 return tuple(v / units for v in x)
94 else:
95 return list(v / units for v in x)
98tango_colors = {
99 'butter1': (252, 233, 79),
100 'butter2': (237, 212, 0),
101 'butter3': (196, 160, 0),
102 'chameleon1': (138, 226, 52),
103 'chameleon2': (115, 210, 22),
104 'chameleon3': (78, 154, 6),
105 'orange1': (252, 175, 62),
106 'orange2': (245, 121, 0),
107 'orange3': (206, 92, 0),
108 'skyblue1': (114, 159, 207),
109 'skyblue2': (52, 101, 164),
110 'skyblue3': (32, 74, 135),
111 'plum1': (173, 127, 168),
112 'plum2': (117, 80, 123),
113 'plum3': (92, 53, 102),
114 'chocolate1': (233, 185, 110),
115 'chocolate2': (193, 125, 17),
116 'chocolate3': (143, 89, 2),
117 'scarletred1': (239, 41, 41),
118 'scarletred2': (204, 0, 0),
119 'scarletred3': (164, 0, 0),
120 'aluminium1': (238, 238, 236),
121 'aluminium2': (211, 215, 207),
122 'aluminium3': (186, 189, 182),
123 'aluminium4': (136, 138, 133),
124 'aluminium5': (85, 87, 83),
125 'aluminium6': (46, 52, 54)}
128graph_colors = [
129 tango_colors[_x] for _x in (
130 'scarletred2',
131 'skyblue3',
132 'chameleon3',
133 'orange2',
134 'plum2',
135 'chocolate2',
136 'butter2')]
139def color(x=None):
140 if x is None:
141 return tuple([random.randint(0, 255) for _x in 'rgb'])
143 if isinstance(x, int):
144 if 0 <= x < len(graph_colors):
145 return graph_colors[x]
146 else:
147 return (0, 0, 0)
149 elif isinstance(x, (str, newstr)):
150 if x in tango_colors:
151 return tango_colors[x]
153 elif isinstance(x, tuple):
154 return x
156 assert False, "Don't know what to do with this color definition: %s" % x
159def to01(c):
160 return tuple(x/255. for x in c)
163def nice_value(x):
164 '''
165 Round x to nice value.
166 '''
168 if x == 0.0:
169 return 0.0
171 exp = 1.0
172 sign = 1
173 if x < 0.0:
174 x = -x
175 sign = -1
176 while x >= 1.0:
177 x /= 10.0
178 exp *= 10.0
179 while x < 0.1:
180 x *= 10.0
181 exp /= 10.0
183 if x >= 0.75:
184 return sign * 1.0 * exp
185 if x >= 0.35:
186 return sign * 0.5 * exp
187 if x >= 0.15:
188 return sign * 0.2 * exp
190 return sign * 0.1 * exp
193_papersizes_list = [
194 ('a0', (2380., 3368.)),
195 ('a1', (1684., 2380.)),
196 ('a2', (1190., 1684.)),
197 ('a3', (842., 1190.)),
198 ('a4', (595., 842.)),
199 ('a5', (421., 595.)),
200 ('a6', (297., 421.)),
201 ('a7', (210., 297.)),
202 ('a8', (148., 210.)),
203 ('a9', (105., 148.)),
204 ('a10', (74., 105.)),
205 ('b0', (2836., 4008.)),
206 ('b1', (2004., 2836.)),
207 ('b2', (1418., 2004.)),
208 ('b3', (1002., 1418.)),
209 ('b4', (709., 1002.)),
210 ('b5', (501., 709.)),
211 ('archa', (648., 864.)),
212 ('archb', (864., 1296.)),
213 ('archc', (1296., 1728.)),
214 ('archd', (1728., 2592.)),
215 ('arche', (2592., 3456.)),
216 ('flsa', (612., 936.)),
217 ('halfletter', (396., 612.)),
218 ('note', (540., 720.)),
219 ('letter', (612., 792.)),
220 ('legal', (612., 1008.)),
221 ('11x17', (792., 1224.)),
222 ('ledger', (1224., 792.))]
224papersizes = dict(_papersizes_list)
226_doc_papersizes = ', '.join("``'%s'``" % k for (k, _) in _papersizes_list)
229def papersize(paper, orientation='landscape', units='point'):
231 '''
232 Get paper size from string.
234 :param paper: string selecting paper size. Choices: %s
235 :param orientation: ``'landscape'``, or ``'portrait'``
236 :param units: Units to be returned. Choices: %s
238 :returns: ``(width, height)``
239 '''
241 assert orientation in ('landscape', 'portrait')
243 w, h = papersizes[paper.lower()]
244 if orientation == 'landscape':
245 w, h = h, w
247 return apply_units((w, h), units)
250papersize.__doc__ %= (_doc_papersizes, _doc_units)
253class AutoScaleMode(StringChoice):
254 '''
255 Mode of operation for auto-scaling.
257 ================ ==================================================
258 mode description
259 ================ ==================================================
260 ``'auto'``: Look at data range and choose one of the choices
261 below.
262 ``'min-max'``: Output range is selected to include data range.
263 ``'0-max'``: Output range shall start at zero and end at data
264 max.
265 ``'min-0'``: Output range shall start at data min and end at
266 zero.
267 ``'symmetric'``: Output range shall by symmetric by zero.
268 ``'off'``: Similar to ``'min-max'``, but snap and space are
269 disabled, such that the output range always
270 exactly matches the data range.
271 ================ ==================================================
272 '''
273 choices = ['auto', 'min-max', '0-max', 'min-0', 'symmetric', 'off']
276class AutoScaler(Object):
278 '''
279 Tunable 1D autoscaling based on data range.
281 Instances of this class may be used to determine nice minima, maxima and
282 increments for ax annotations, as well as suitable common exponents for
283 notation.
285 The autoscaling process is guided by the following public attributes:
286 '''
288 approx_ticks = Float.T(
289 default=7.0,
290 help='Approximate number of increment steps (tickmarks) to generate.')
292 mode = AutoScaleMode.T(
293 default='auto',
294 help='''Mode of operation for auto-scaling.''')
296 exp = Int.T(
297 optional=True,
298 help='If defined, override automatically determined exponent for '
299 'notation by the given value.')
301 snap = Bool.T(
302 default=False,
303 help='If set to True, snap output range to multiples of increment. '
304 'This parameter has no effect, if mode is set to ``\'off\'``.')
306 inc = Float.T(
307 optional=True,
308 help='If defined, override automatically determined tick increment by '
309 'the given value.')
311 space = Float.T(
312 default=0.0,
313 help='Add some padding to the range. The value given, is the fraction '
314 'by which the output range is increased on each side. If mode is '
315 '``\'0-max\'`` or ``\'min-0\'``, the end at zero is kept fixed '
316 'at zero. This parameter has no effect if mode is set to '
317 '``\'off\'``.')
319 exp_factor = Int.T(
320 default=3,
321 help='Exponent of notation is chosen to be a multiple of this value.')
323 no_exp_interval = Tuple.T(
324 2, Int.T(),
325 default=(-3, 5),
326 help='Range of exponent, for which no exponential notation is a'
327 'allowed.')
329 def __init__(
330 self,
331 approx_ticks=7.0,
332 mode='auto',
333 exp=None,
334 snap=False,
335 inc=None,
336 space=0.0,
337 exp_factor=3,
338 no_exp_interval=(-3, 5)):
340 '''
341 Create new AutoScaler instance.
343 The parameters are described in the AutoScaler documentation.
344 '''
346 Object.__init__(
347 self,
348 approx_ticks=approx_ticks,
349 mode=mode,
350 exp=exp,
351 snap=snap,
352 inc=inc,
353 space=space,
354 exp_factor=exp_factor,
355 no_exp_interval=no_exp_interval)
357 def make_scale(self, data_range, override_mode=None):
359 '''
360 Get nice minimum, maximum and increment for given data range.
362 Returns ``(minimum, maximum, increment)`` or ``(maximum, minimum,
363 -increment)``, depending on whether data_range is ``(data_min,
364 data_max)`` or ``(data_max, data_min)``. If ``override_mode`` is
365 defined, the mode attribute is temporarily overridden by the given
366 value.
367 '''
369 data_min = min(data_range)
370 data_max = max(data_range)
372 is_reverse = (data_range[0] > data_range[1])
374 a = self.mode
375 if self.mode == 'auto':
376 a = self.guess_autoscale_mode(data_min, data_max)
378 if override_mode is not None:
379 a = override_mode
381 mi, ma = 0, 0
382 if a == 'off':
383 mi, ma = data_min, data_max
384 elif a == '0-max':
385 mi = 0.0
386 if data_max > 0.0:
387 ma = data_max
388 else:
389 ma = 1.0
390 elif a == 'min-0':
391 ma = 0.0
392 if data_min < 0.0:
393 mi = data_min
394 else:
395 mi = -1.0
396 elif a == 'min-max':
397 mi, ma = data_min, data_max
398 elif a == 'symmetric':
399 m = max(abs(data_min), abs(data_max))
400 mi = -m
401 ma = m
403 nmi = mi
404 if (mi != 0. or a == 'min-max') and a != 'off':
405 nmi = mi - self.space*(ma-mi)
407 nma = ma
408 if (ma != 0. or a == 'min-max') and a != 'off':
409 nma = ma + self.space*(ma-mi)
411 mi, ma = nmi, nma
413 if mi == ma and a != 'off':
414 mi -= 1.0
415 ma += 1.0
417 # make nice tick increment
418 if self.inc is not None:
419 inc = self.inc
420 else:
421 if self.approx_ticks > 0.:
422 inc = nice_value((ma-mi) / self.approx_ticks)
423 else:
424 inc = nice_value((ma-mi)*10.)
426 if inc == 0.0:
427 inc = 1.0
429 # snap min and max to ticks if this is wanted
430 if self.snap and a != 'off':
431 ma = inc * math.ceil(ma/inc)
432 mi = inc * math.floor(mi/inc)
434 if is_reverse:
435 return ma, mi, -inc
436 else:
437 return mi, ma, inc
439 def make_exp(self, x):
440 '''
441 Get nice exponent for notation of ``x``.
443 For ax annotations, give tick increment as ``x``.
444 '''
446 if self.exp is not None:
447 return self.exp
449 x = abs(x)
450 if x == 0.0:
451 return 0
453 if 10**self.no_exp_interval[0] <= x <= 10**self.no_exp_interval[1]:
454 return 0
456 return math.floor(math.log10(x)/self.exp_factor)*self.exp_factor
458 def guess_autoscale_mode(self, data_min, data_max):
459 '''
460 Guess mode of operation, based on data range.
462 Used to map ``'auto'`` mode to ``'0-max'``, ``'min-0'``, ``'min-max'``
463 or ``'symmetric'``.
464 '''
466 a = 'min-max'
467 if data_min >= 0.0:
468 if data_min < data_max/2.:
469 a = '0-max'
470 else:
471 a = 'min-max'
472 if data_max <= 0.0:
473 if data_max > data_min/2.:
474 a = 'min-0'
475 else:
476 a = 'min-max'
477 if data_min < 0.0 and data_max > 0.0:
478 if abs((abs(data_max)-abs(data_min)) /
479 (abs(data_max)+abs(data_min))) < 0.5:
480 a = 'symmetric'
481 else:
482 a = 'min-max'
483 return a
486# below, some convenience functions for matplotlib plotting
488def mpl_init(fontsize=10):
489 '''
490 Initialize Matplotlib rc parameters Pyrocko style.
492 Returns the matplotlib.pyplot module for convenience.
493 '''
495 import matplotlib
497 matplotlib.rcdefaults()
498 matplotlib.rc('font', size=fontsize)
499 matplotlib.rc('axes', linewidth=1.5)
500 matplotlib.rc('xtick', direction='out')
501 matplotlib.rc('ytick', direction='out')
502 ts = fontsize * 0.7071
503 matplotlib.rc('xtick.major', size=ts, width=0.5, pad=ts)
504 matplotlib.rc('ytick.major', size=ts, width=0.5, pad=ts)
505 matplotlib.rc('figure', facecolor='white')
507 try:
508 from cycler import cycler
509 matplotlib.rc(
510 'axes', prop_cycle=cycler(
511 'color', [to01(x) for x in graph_colors]))
512 except (ImportError, KeyError):
513 try:
514 matplotlib.rc('axes', color_cycle=[to01(x) for x in graph_colors])
515 except KeyError:
516 pass
518 from matplotlib import pyplot as plt
519 return plt
522def mpl_margins(
523 fig,
524 left=1.0, top=1.0, right=1.0, bottom=1.0,
525 wspace=None, hspace=None,
526 w=None, h=None,
527 nw=None, nh=None,
528 all=None,
529 units='inch'):
531 '''
532 Adjust Matplotlib subplot params with absolute values in user units.
534 Calls :py:meth:`matplotlib.figure.Figure.subplots_adjust` on ``fig`` with
535 absolute margin widths/heights rather than relative values. If ``wspace``
536 or ``hspace`` are given, the number of subplots must be given in ``nw``
537 and ``nh`` because ``subplots_adjust()`` treats the spacing parameters
538 relative to the subplot width and height.
540 :param units: Unit multiplier or unit as string: %s
541 :param left,right,top,bottom: margin space
542 :param w: set ``left`` and ``right`` at once
543 :param h: set ``top`` and ``bottom`` at once
544 :param all: set ``left``, ``top``, ``right``, and ``bottom`` at once
545 :param nw: number of subplots horizontally
546 :param nh: number of subplots vertically
547 :param wspace: horizontal spacing between subplots
548 :param hspace: vertical spacing between subplots
549 '''
551 left, top, right, bottom = map(
552 float, (left, top, right, bottom))
554 if w is not None:
555 left = right = float(w)
557 if h is not None:
558 top = bottom = float(h)
560 if all is not None:
561 left = right = top = bottom = float(all)
563 ufac = units_dict.get(units, units) / inch
565 left *= ufac
566 right *= ufac
567 top *= ufac
568 bottom *= ufac
570 width, height = fig.get_size_inches()
572 rel_wspace = None
573 rel_hspace = None
575 if wspace is not None:
576 wspace *= ufac
577 if nw is None:
578 raise ValueError('wspace must be given in combination with nw')
580 wsub = (width - left - right - (nw-1) * wspace) / nw
581 rel_wspace = wspace / wsub
582 else:
583 wsub = width - left - right
585 if hspace is not None:
586 hspace *= ufac
587 if nh is None:
588 raise ValueError('hspace must be given in combination with nh')
590 hsub = (height - top - bottom - (nh-1) * hspace) / nh
591 rel_hspace = hspace / hsub
592 else:
593 hsub = height - top - bottom
595 fig.subplots_adjust(
596 left=left/width,
597 right=1.0 - right/width,
598 bottom=bottom/height,
599 top=1.0 - top/height,
600 wspace=rel_wspace,
601 hspace=rel_hspace)
603 def labelpos(axes, xpos=0., ypos=0.):
604 xpos *= ufac
605 ypos *= ufac
606 axes.get_yaxis().set_label_coords(-((left-xpos) / wsub), 0.5)
607 axes.get_xaxis().set_label_coords(0.5, -((bottom-ypos) / hsub))
609 return labelpos
612mpl_margins.__doc__ %= _doc_units
615def mpl_labelspace(axes):
616 '''
617 Add some extra padding between label and ax annotations.
618 '''
620 xa = axes.get_xaxis()
621 ya = axes.get_yaxis()
622 for attr in ('labelpad', 'LABELPAD'):
623 if hasattr(xa, attr):
624 setattr(xa, attr, xa.get_label().get_fontsize())
625 setattr(ya, attr, ya.get_label().get_fontsize())
626 break
629def mpl_papersize(paper, orientation='landscape'):
630 '''
631 Get paper size in inch from string.
633 Returns argument suitable to be passed to the ``figsize`` argument of
634 :py:func:`pyplot.figure`.
636 :param paper: string selecting paper size. Choices: %s
637 :param orientation: ``'landscape'``, or ``'portrait'``
639 :returns: ``(width, height)``
640 '''
642 return papersize(paper, orientation=orientation, units='inch')
645mpl_papersize.__doc__ %= _doc_papersizes
648class InvalidColorDef(ValueError):
649 pass
652def mpl_graph_color(i):
653 return to01(graph_colors[i % len(graph_colors)])
656def mpl_color(x):
657 '''
658 Convert string into color float tuple ranged 0-1 for use with Matplotlib.
660 Accepts tango color names, matplotlib color names, and slash-separated
661 strings. In the latter case, if values are larger than 1., the color
662 is interpreted as 0-255 ranged. Single-valued (grayscale), three-valued
663 (color) and four-valued (color with alpha) are accepted. An
664 :py:exc:`InvalidColorDef` exception is raised when the convertion fails.
665 '''
667 import matplotlib.colors
669 if x in tango_colors:
670 return to01(tango_colors[x])
672 s = x.split('/')
673 if len(s) in (1, 3, 4):
674 try:
675 vals = list(map(float, s))
676 if all(0. <= v <= 1. for v in vals):
677 return vals
679 elif all(0. <= v <= 255. for v in vals):
680 return to01(vals)
682 except ValueError:
683 try:
684 return matplotlib.colors.colorConverter.to_rgba(x)
685 except Exception:
686 pass
688 raise InvalidColorDef('invalid color definition: %s' % x)
691def nice_time_tick_inc(tinc_approx):
692 hours = 3600.
693 days = hours*24
694 approx_months = days*30.5
695 approx_years = days*365
697 if tinc_approx >= approx_years:
698 return max(1.0, nice_value(tinc_approx / approx_years)), 'years'
700 elif tinc_approx >= approx_months:
701 nice = [1, 2, 3, 6]
702 for tinc in nice:
703 if tinc*approx_months >= tinc_approx or tinc == nice[-1]:
704 return tinc, 'months'
706 elif tinc_approx > days:
707 return nice_value(tinc_approx / days) * days, 'seconds'
709 elif tinc_approx >= 1.0:
710 nice = [
711 1., 2., 5., 10., 20, 30, 60., 120., 300., 600., 1200., 1800.,
712 1*hours, 2*hours, 3*hours, 6*hours, 12*hours, days, 2*days]
714 for tinc in nice:
715 if tinc >= tinc_approx or tinc == nice[-1]:
716 return tinc, 'seconds'
718 else:
719 return nice_value(tinc_approx), 'seconds'
722def time_tick_labels(tmin, tmax, tinc, tinc_unit):
724 if tinc_unit == 'years':
725 tt = time.gmtime(int(tmin))
726 tmin_year = tt[0]
727 if tt[1:6] != (1, 1, 0, 0, 0):
728 tmin_year += 1
730 tmax_year = time.gmtime(int(tmax))[0]
732 tick_times_year = arange2(
733 math.ceil(tmin_year/tinc)*tinc,
734 math.floor(tmax_year/tinc)*tinc,
735 tinc).astype(int)
737 times = [
738 to_time_float(calendar.timegm((year, 1, 1, 0, 0, 0)))
739 for year in tick_times_year]
741 labels = ['%04i' % year for year in tick_times_year]
743 elif tinc_unit == 'months':
744 tt = time.gmtime(int(tmin))
745 tmin_ym = tt[0] * 12 + (tt[1] - 1)
746 if tt[2:6] != (1, 0, 0, 0):
747 tmin_ym += 1
749 tt = time.gmtime(int(tmax))
750 tmax_ym = tt[0] * 12 + (tt[1] - 1)
752 tick_times_ym = arange2(
753 math.ceil(tmin_ym/tinc)*tinc,
754 math.floor(tmax_ym/tinc)*tinc, tinc).astype(int)
756 times = [
757 to_time_float(calendar.timegm((ym // 12, ym % 12 + 1, 1, 0, 0, 0)))
758 for ym in tick_times_ym]
760 labels = [
761 '%04i-%02i' % (ym // 12, ym % 12 + 1) for ym in tick_times_ym]
763 elif tinc_unit == 'seconds':
764 imin = int(num.ceil(tmin/tinc))
765 imax = int(num.floor(tmax/tinc))
766 nticks = imax - imin + 1
767 tmin_ticks = imin * tinc
768 times = tmin_ticks + num.arange(nticks) * tinc
769 times = times.tolist()
771 if tinc < 1e-6:
772 fmt = '%Y-%m-%d.%H:%M:%S.9FRAC'
773 elif tinc < 1e-3:
774 fmt = '%Y-%m-%d.%H:%M:%S.6FRAC'
775 elif tinc < 1.0:
776 fmt = '%Y-%m-%d.%H:%M:%S.3FRAC'
777 elif tinc < 60:
778 fmt = '%Y-%m-%d.%H:%M:%S'
779 elif tinc < 3600.*24:
780 fmt = '%Y-%m-%d.%H:%M'
781 else:
782 fmt = '%Y-%m-%d'
784 nwords = len(fmt.split('.'))
786 labels = [time_to_str(t, format=fmt) for t in times]
787 labels_weeded = []
788 have_ymd = have_hms = False
789 ymd = hms = ''
790 for ilab, lab in reversed(list(enumerate(labels))):
791 words = lab.split('.')
792 if nwords > 2:
793 words[2] = '.' + words[2]
794 if float(words[2]) == 0.0: # or (ilab == 0 and not have_hms):
795 have_hms = True
796 else:
797 hms = words[1]
798 words[1] = ''
799 else:
800 have_hms = True
802 if nwords > 1:
803 if words[1] in ('00:00', '00:00:00'): # or (ilab == 0 and not have_ymd): # noqa
804 have_ymd = True
805 else:
806 ymd = words[0]
807 words[0] = ''
808 else:
809 have_ymd = True
811 labels_weeded.append('\n'.join(reversed(words)))
813 labels = list(reversed(labels_weeded))
814 if (not have_ymd or not have_hms) and (hms or ymd):
815 words = ([''] if nwords > 2 else []) + [
816 hms if not have_hms else '',
817 ymd if not have_ymd else '']
819 labels[0:0] = ['\n'.join(words)]
820 times[0:0] = [tmin]
822 return times, labels