Coverage for /usr/local/lib/python3.11/dist-packages/pyrocko/plot/__init__.py: 80%

384 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2023-10-04 08:54 +0000

1# http://pyrocko.org - GPLv3 

2# 

3# The Pyrocko Developers, 21st Century 

4# ---|P------/S----------~Lg---------- 

5 

6''' 

7Utility functions and defintions for a common plot style throughout Pyrocko. 

8 

9Functions with name prefix ``mpl_`` are Matplotlib specific. All others should 

10be toolkit-agnostic. 

11 

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):: 

16 

17 from matplotlib import pyplot as plt 

18 

19 from pyrocko.plot import mpl_init, mpl_margins, mpl_papersize 

20 # from pyrocko.plot import mpl_labelspace 

21 

22 fontsize = 9. # in points 

23 

24 # set some Pyrocko style defaults 

25 mpl_init(fontsize=fontsize) 

26 

27 fig = plt.figure(figsize=mpl_papersize('a4', 'landscape')) 

28 

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) 

32 

33 axes = fig.add_subplot(1, 1, 1) 

34 

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 

38 

39 axes.plot([0, 1], [0, 9]) 

40 

41 axes.set_xlabel('Time [s]') 

42 axes.set_ylabel('Amplitude [m]') 

43 

44 fig.savefig('plot_skeleton.pdf') 

45 

46 plt.show() 

47 

48''' 

49 

50import math 

51import random 

52import time 

53import calendar 

54import numpy as num 

55 

56from pyrocko.util import parse_md, time_to_str, arange2, to_time_float 

57from pyrocko.guts import StringChoice, Float, Int, Bool, Tuple, Object 

58 

59 

60__doc__ += parse_md(__file__) 

61 

62 

63guts_prefix = 'pf' 

64 

65point = 1. 

66inch = 72. 

67cm = 28.3465 

68 

69units_dict = { 

70 'point': point, 

71 'inch': inch, 

72 'cm': cm, 

73} 

74 

75_doc_units = "``'point'``, ``'inch'``, or ``'cm'``" 

76 

77 

78def apply_units(x, units): 

79 if isinstance(units, str): 

80 units = units_dict[units] 

81 

82 if isinstance(x, (int, float)): 

83 return x / units 

84 else: 

85 if isinstance(x, tuple): 

86 return tuple(v / units for v in x) 

87 else: 

88 return list(v / units for v in x) 

89 

90 

91tango_colors = { 

92 'butter1': (252, 233, 79), 

93 'butter2': (237, 212, 0), 

94 'butter3': (196, 160, 0), 

95 'chameleon1': (138, 226, 52), 

96 'chameleon2': (115, 210, 22), 

97 'chameleon3': (78, 154, 6), 

98 'orange1': (252, 175, 62), 

99 'orange2': (245, 121, 0), 

100 'orange3': (206, 92, 0), 

101 'skyblue1': (114, 159, 207), 

102 'skyblue2': (52, 101, 164), 

103 'skyblue3': (32, 74, 135), 

104 'plum1': (173, 127, 168), 

105 'plum2': (117, 80, 123), 

106 'plum3': (92, 53, 102), 

107 'chocolate1': (233, 185, 110), 

108 'chocolate2': (193, 125, 17), 

109 'chocolate3': (143, 89, 2), 

110 'scarletred1': (239, 41, 41), 

111 'scarletred2': (204, 0, 0), 

112 'scarletred3': (164, 0, 0), 

113 'aluminium1': (238, 238, 236), 

114 'aluminium2': (211, 215, 207), 

115 'aluminium3': (186, 189, 182), 

116 'aluminium4': (136, 138, 133), 

117 'aluminium5': (85, 87, 83), 

118 'aluminium6': (46, 52, 54)} 

119 

120 

121graph_colors = [ 

122 tango_colors[_x] for _x in ( 

123 'scarletred2', 

124 'skyblue3', 

125 'chameleon3', 

126 'orange2', 

127 'plum2', 

128 'chocolate2', 

129 'butter2')] 

130 

131 

132def color(x=None): 

133 if x is None: 

134 return tuple([random.randint(0, 255) for _x in 'rgb']) 

135 

136 if isinstance(x, int): 

137 if 0 <= x < len(graph_colors): 

138 return graph_colors[x] 

139 else: 

140 return (0, 0, 0) 

141 

142 elif isinstance(x, str): 

143 if x in tango_colors: 

144 return tango_colors[x] 

145 

146 elif isinstance(x, tuple): 

147 return x 

148 

149 assert False, "Don't know what to do with this color definition: %s" % x 

150 

151 

152def to01(c): 

153 return tuple(x/255. for x in c) 

154 

155 

156def nice_value(x): 

157 ''' 

158 Round x to nice value. 

159 ''' 

160 

161 if x == 0.0: 

162 return 0.0 

163 

164 exp = 1.0 

165 sign = 1 

166 if x < 0.0: 

167 x = -x 

168 sign = -1 

169 while x >= 1.0: 

170 x /= 10.0 

171 exp *= 10.0 

172 while x < 0.1: 

173 x *= 10.0 

174 exp /= 10.0 

175 

176 if x >= 0.75: 

177 return sign * 1.0 * exp 

178 if x >= 0.35: 

179 return sign * 0.5 * exp 

180 if x >= 0.15: 

181 return sign * 0.2 * exp 

182 

183 return sign * 0.1 * exp 

184 

185 

186_papersizes_list = [ 

187 ('a0', (2380., 3368.)), 

188 ('a1', (1684., 2380.)), 

189 ('a2', (1190., 1684.)), 

190 ('a3', (842., 1190.)), 

191 ('a4', (595., 842.)), 

192 ('a5', (421., 595.)), 

193 ('a6', (297., 421.)), 

194 ('a7', (210., 297.)), 

195 ('a8', (148., 210.)), 

196 ('a9', (105., 148.)), 

197 ('a10', (74., 105.)), 

198 ('b0', (2836., 4008.)), 

199 ('b1', (2004., 2836.)), 

200 ('b2', (1418., 2004.)), 

201 ('b3', (1002., 1418.)), 

202 ('b4', (709., 1002.)), 

203 ('b5', (501., 709.)), 

204 ('archa', (648., 864.)), 

205 ('archb', (864., 1296.)), 

206 ('archc', (1296., 1728.)), 

207 ('archd', (1728., 2592.)), 

208 ('arche', (2592., 3456.)), 

209 ('flsa', (612., 936.)), 

210 ('halfletter', (396., 612.)), 

211 ('note', (540., 720.)), 

212 ('letter', (612., 792.)), 

213 ('legal', (612., 1008.)), 

214 ('11x17', (792., 1224.)), 

215 ('ledger', (1224., 792.))] 

216 

217papersizes = dict(_papersizes_list) 

218 

219_doc_papersizes = ', '.join("``'%s'``" % k for (k, _) in _papersizes_list) 

220 

221 

222def papersize(paper, orientation='landscape', units='point'): 

223 

224 ''' 

225 Get paper size from string. 

226 

227 :param paper: string selecting paper size. Choices: %s 

228 :param orientation: ``'landscape'``, or ``'portrait'`` 

229 :param units: Units to be returned. Choices: %s 

230 

231 :returns: ``(width, height)`` 

232 ''' 

233 

234 assert orientation in ('landscape', 'portrait') 

235 

236 w, h = papersizes[paper.lower()] 

237 if orientation == 'landscape': 

238 w, h = h, w 

239 

240 return apply_units((w, h), units) 

241 

242 

243papersize.__doc__ %= (_doc_papersizes, _doc_units) 

244 

245 

246class AutoScaleMode(StringChoice): 

247 ''' 

248 Mode of operation for auto-scaling. 

249 

250 ================ ================================================== 

251 mode description 

252 ================ ================================================== 

253 ``'auto'``: Look at data range and choose one of the choices 

254 below. 

255 ``'min-max'``: Output range is selected to include data range. 

256 ``'0-max'``: Output range shall start at zero and end at data 

257 max. 

258 ``'min-0'``: Output range shall start at data min and end at 

259 zero. 

260 ``'symmetric'``: Output range shall by symmetric by zero. 

261 ``'off'``: Similar to ``'min-max'``, but snap and space are 

262 disabled, such that the output range always 

263 exactly matches the data range. 

264 ================ ================================================== 

265 ''' 

266 choices = ['auto', 'min-max', '0-max', 'min-0', 'symmetric', 'off'] 

267 

268 

269class AutoScaler(Object): 

270 

271 ''' 

272 Tunable 1D autoscaling based on data range. 

273 

274 Instances of this class may be used to determine nice minima, maxima and 

275 increments for ax annotations, as well as suitable common exponents for 

276 notation. 

277 

278 The autoscaling process is guided by the following public attributes: 

279 ''' 

280 

281 approx_ticks = Float.T( 

282 default=7.0, 

283 help='Approximate number of increment steps (tickmarks) to generate.') 

284 

285 mode = AutoScaleMode.T( 

286 default='auto', 

287 help='''Mode of operation for auto-scaling.''') 

288 

289 exp = Int.T( 

290 optional=True, 

291 help='If defined, override automatically determined exponent for ' 

292 'notation by the given value.') 

293 

294 snap = Bool.T( 

295 default=False, 

296 help='If set to True, snap output range to multiples of increment. ' 

297 "This parameter has no effect, if mode is set to ``'off'``.") 

298 

299 inc = Float.T( 

300 optional=True, 

301 help='If defined, override automatically determined tick increment by ' 

302 'the given value.') 

303 

304 space = Float.T( 

305 default=0.0, 

306 help='Add some padding to the range. The value given, is the fraction ' 

307 'by which the output range is increased on each side. If mode is ' 

308 "``'0-max'`` or ``'min-0'``, the end at zero is kept fixed " 

309 'at zero. This parameter has no effect if mode is set to ' 

310 "``'off'``.") 

311 

312 exp_factor = Int.T( 

313 default=3, 

314 help='Exponent of notation is chosen to be a multiple of this value.') 

315 

316 no_exp_interval = Tuple.T( 

317 2, Int.T(), 

318 default=(-3, 5), 

319 help='Range of exponent, for which no exponential notation is a' 

320 'allowed.') 

321 

322 def __init__( 

323 self, 

324 approx_ticks=7.0, 

325 mode='auto', 

326 exp=None, 

327 snap=False, 

328 inc=None, 

329 space=0.0, 

330 exp_factor=3, 

331 no_exp_interval=(-3, 5)): 

332 

333 ''' 

334 Create new AutoScaler instance. 

335 

336 The parameters are described in the AutoScaler documentation. 

337 ''' 

338 

339 Object.__init__( 

340 self, 

341 approx_ticks=approx_ticks, 

342 mode=mode, 

343 exp=exp, 

344 snap=snap, 

345 inc=inc, 

346 space=space, 

347 exp_factor=exp_factor, 

348 no_exp_interval=no_exp_interval) 

349 

350 def make_scale(self, data_range, override_mode=None): 

351 

352 ''' 

353 Get nice minimum, maximum and increment for given data range. 

354 

355 Returns ``(minimum, maximum, increment)`` or ``(maximum, minimum, 

356 -increment)``, depending on whether data_range is ``(data_min, 

357 data_max)`` or ``(data_max, data_min)``. If ``override_mode`` is 

358 defined, the mode attribute is temporarily overridden by the given 

359 value. 

360 ''' 

361 

362 data_min = min(data_range) 

363 data_max = max(data_range) 

364 

365 is_reverse = (data_range[0] > data_range[1]) 

366 

367 a = self.mode 

368 if self.mode == 'auto': 

369 a = self.guess_autoscale_mode(data_min, data_max) 

370 

371 if override_mode is not None: 

372 a = override_mode 

373 

374 mi, ma = 0, 0 

375 if a == 'off': 

376 mi, ma = data_min, data_max 

377 elif a == '0-max': 

378 mi = 0.0 

379 if data_max > 0.0: 

380 ma = data_max 

381 else: 

382 ma = 1.0 

383 elif a == 'min-0': 

384 ma = 0.0 

385 if data_min < 0.0: 

386 mi = data_min 

387 else: 

388 mi = -1.0 

389 elif a == 'min-max': 

390 mi, ma = data_min, data_max 

391 elif a == 'symmetric': 

392 m = max(abs(data_min), abs(data_max)) 

393 mi = -m 

394 ma = m 

395 

396 nmi = mi 

397 if (mi != 0. or a == 'min-max') and a != 'off': 

398 nmi = mi - self.space*(ma-mi) 

399 

400 nma = ma 

401 if (ma != 0. or a == 'min-max') and a != 'off': 

402 nma = ma + self.space*(ma-mi) 

403 

404 mi, ma = nmi, nma 

405 

406 if mi == ma and a != 'off': 

407 mi -= 1.0 

408 ma += 1.0 

409 

410 # make nice tick increment 

411 if self.inc is not None: 

412 inc = self.inc 

413 else: 

414 if self.approx_ticks > 0.: 

415 inc = nice_value((ma-mi) / self.approx_ticks) 

416 else: 

417 inc = nice_value((ma-mi)*10.) 

418 

419 if inc == 0.0: 

420 inc = 1.0 

421 

422 # snap min and max to ticks if this is wanted 

423 if self.snap and a != 'off': 

424 ma = inc * math.ceil(ma/inc) 

425 mi = inc * math.floor(mi/inc) 

426 

427 if is_reverse: 

428 return ma, mi, -inc 

429 else: 

430 return mi, ma, inc 

431 

432 def make_exp(self, x): 

433 ''' 

434 Get nice exponent for notation of ``x``. 

435 

436 For ax annotations, give tick increment as ``x``. 

437 ''' 

438 

439 if self.exp is not None: 

440 return self.exp 

441 

442 x = abs(x) 

443 if x == 0.0: 

444 return 0 

445 

446 if 10**self.no_exp_interval[0] <= x <= 10**self.no_exp_interval[1]: 

447 return 0 

448 

449 return math.floor(math.log10(x)/self.exp_factor)*self.exp_factor 

450 

451 def guess_autoscale_mode(self, data_min, data_max): 

452 ''' 

453 Guess mode of operation, based on data range. 

454 

455 Used to map ``'auto'`` mode to ``'0-max'``, ``'min-0'``, ``'min-max'`` 

456 or ``'symmetric'``. 

457 ''' 

458 

459 a = 'min-max' 

460 if data_min >= 0.0: 

461 if data_min < data_max/2.: 

462 a = '0-max' 

463 else: 

464 a = 'min-max' 

465 if data_max <= 0.0: 

466 if data_max > data_min/2.: 

467 a = 'min-0' 

468 else: 

469 a = 'min-max' 

470 if data_min < 0.0 and data_max > 0.0: 

471 if abs((abs(data_max)-abs(data_min)) / 

472 (abs(data_max)+abs(data_min))) < 0.5: 

473 a = 'symmetric' 

474 else: 

475 a = 'min-max' 

476 return a 

477 

478 

479# below, some convenience functions for matplotlib plotting 

480 

481def mpl_init(fontsize=10): 

482 ''' 

483 Initialize Matplotlib rc parameters Pyrocko style. 

484 

485 Returns the matplotlib.pyplot module for convenience. 

486 ''' 

487 

488 import matplotlib 

489 

490 matplotlib.rcdefaults() 

491 matplotlib.rc('font', size=fontsize) 

492 matplotlib.rc('axes', linewidth=1.5) 

493 matplotlib.rc('xtick', direction='out') 

494 matplotlib.rc('ytick', direction='out') 

495 ts = fontsize * 0.7071 

496 matplotlib.rc('xtick.major', size=ts, width=0.5, pad=ts) 

497 matplotlib.rc('ytick.major', size=ts, width=0.5, pad=ts) 

498 matplotlib.rc('figure', facecolor='white') 

499 

500 try: 

501 from cycler import cycler 

502 matplotlib.rc( 

503 'axes', prop_cycle=cycler( 

504 'color', [to01(x) for x in graph_colors])) 

505 except (ImportError, KeyError): 

506 try: 

507 matplotlib.rc('axes', color_cycle=[to01(x) for x in graph_colors]) 

508 except KeyError: 

509 pass 

510 

511 from matplotlib import pyplot as plt 

512 return plt 

513 

514 

515def mpl_get_cmap_names(): 

516 ''' 

517 Compatibility function to get named MPL colormap names. 

518 ''' 

519 

520 try: 

521 from matplotlib import colormaps 

522 names = list(colormaps.keys()) 

523 except ImportError: 

524 from matplotlib.cm import _cmap_registry 

525 names = list(_cmap_registry.keys()) 

526 

527 names.sort() 

528 return names 

529 

530 

531def mpl_get_cmap(name): 

532 ''' 

533 Compatibility function to get named MPL colormap. 

534 

535 The function matplotlib.cm.get_cmap has been removed in MPL 3.8 but the 

536 suggested replacement is not available in slightly older versions of MPL, 

537 e.g. 3.3 (default on Debian 11). 

538 ''' 

539 

540 try: 

541 from matplotlib import colormaps 

542 return colormaps[name] 

543 except ImportError: 

544 from matplotlib import cm 

545 return cm.get_cmap(name) 

546 

547 

548def mpl_margins( 

549 fig, 

550 left=1.0, top=1.0, right=1.0, bottom=1.0, 

551 wspace=None, hspace=None, 

552 w=None, h=None, 

553 nw=None, nh=None, 

554 all=None, 

555 units='inch'): 

556 

557 ''' 

558 Adjust Matplotlib subplot params with absolute values in user units. 

559 

560 Calls :py:meth:`matplotlib.figure.Figure.subplots_adjust` on ``fig`` with 

561 absolute margin widths/heights rather than relative values. If ``wspace`` 

562 or ``hspace`` are given, the number of subplots must be given in ``nw`` 

563 and ``nh`` because ``subplots_adjust()`` treats the spacing parameters 

564 relative to the subplot width and height. 

565 

566 :param units: Unit multiplier or unit as string: %s 

567 :param left,right,top,bottom: margin space 

568 :param w: set ``left`` and ``right`` at once 

569 :param h: set ``top`` and ``bottom`` at once 

570 :param all: set ``left``, ``top``, ``right``, and ``bottom`` at once 

571 :param nw: number of subplots horizontally 

572 :param nh: number of subplots vertically 

573 :param wspace: horizontal spacing between subplots 

574 :param hspace: vertical spacing between subplots 

575 ''' 

576 

577 left, top, right, bottom = map( 

578 float, (left, top, right, bottom)) 

579 

580 if w is not None: 

581 left = right = float(w) 

582 

583 if h is not None: 

584 top = bottom = float(h) 

585 

586 if all is not None: 

587 left = right = top = bottom = float(all) 

588 

589 ufac = units_dict.get(units, units) / inch 

590 

591 left *= ufac 

592 right *= ufac 

593 top *= ufac 

594 bottom *= ufac 

595 

596 width, height = fig.get_size_inches() 

597 

598 rel_wspace = None 

599 rel_hspace = None 

600 

601 if wspace is not None: 

602 wspace *= ufac 

603 if nw is None: 

604 raise ValueError('wspace must be given in combination with nw') 

605 

606 wsub = (width - left - right - (nw-1) * wspace) / nw 

607 rel_wspace = wspace / wsub 

608 else: 

609 wsub = width - left - right 

610 

611 if hspace is not None: 

612 hspace *= ufac 

613 if nh is None: 

614 raise ValueError('hspace must be given in combination with nh') 

615 

616 hsub = (height - top - bottom - (nh-1) * hspace) / nh 

617 rel_hspace = hspace / hsub 

618 else: 

619 hsub = height - top - bottom 

620 

621 fig.subplots_adjust( 

622 left=left/width, 

623 right=1.0 - right/width, 

624 bottom=bottom/height, 

625 top=1.0 - top/height, 

626 wspace=rel_wspace, 

627 hspace=rel_hspace) 

628 

629 def labelpos(axes, xpos=0., ypos=0.): 

630 xpos *= ufac 

631 ypos *= ufac 

632 axes.get_yaxis().set_label_coords(-((left-xpos) / wsub), 0.5) 

633 axes.get_xaxis().set_label_coords(0.5, -((bottom-ypos) / hsub)) 

634 

635 return labelpos 

636 

637 

638mpl_margins.__doc__ %= _doc_units 

639 

640 

641def mpl_labelspace(axes): 

642 ''' 

643 Add some extra padding between label and ax annotations. 

644 ''' 

645 

646 xa = axes.get_xaxis() 

647 ya = axes.get_yaxis() 

648 for attr in ('labelpad', 'LABELPAD'): 

649 if hasattr(xa, attr): 

650 setattr(xa, attr, xa.get_label().get_fontsize()) 

651 setattr(ya, attr, ya.get_label().get_fontsize()) 

652 break 

653 

654 

655def mpl_papersize(paper, orientation='landscape'): 

656 ''' 

657 Get paper size in inch from string. 

658 

659 Returns argument suitable to be passed to the ``figsize`` argument of 

660 :py:func:`matplotlib.pyplot.figure`. 

661 

662 :param paper: string selecting paper size. Choices: %s 

663 :param orientation: ``'landscape'``, or ``'portrait'`` 

664 

665 :returns: ``(width, height)`` 

666 ''' 

667 

668 return papersize(paper, orientation=orientation, units='inch') 

669 

670 

671mpl_papersize.__doc__ %= _doc_papersizes 

672 

673 

674class InvalidColorDef(ValueError): 

675 ''' 

676 Raised for invalid color definitions. 

677 ''' 

678 pass 

679 

680 

681def mpl_graph_color(i): 

682 return to01(graph_colors[i % len(graph_colors)]) 

683 

684 

685def mpl_color(x): 

686 ''' 

687 Convert string into color float tuple ranged 0-1 for use with Matplotlib. 

688 

689 Accepts tango color names, matplotlib color names, and slash-separated 

690 strings. In the latter case, if values are larger than 1., the color 

691 is interpreted as 0-255 ranged. Single-valued (grayscale), three-valued 

692 (color) and four-valued (color with alpha) are accepted. An 

693 :py:exc:`InvalidColorDef` exception is raised when the convertion fails. 

694 ''' 

695 

696 import matplotlib.colors 

697 

698 if x in tango_colors: 

699 return to01(tango_colors[x]) 

700 

701 s = x.split('/') 

702 if len(s) in (1, 3, 4): 

703 try: 

704 vals = list(map(float, s)) 

705 if all(0. <= v <= 1. for v in vals): 

706 return vals 

707 

708 elif all(0. <= v <= 255. for v in vals): 

709 return to01(vals) 

710 

711 except ValueError: 

712 try: 

713 return matplotlib.colors.colorConverter.to_rgba(x) 

714 except Exception: 

715 pass 

716 

717 raise InvalidColorDef('invalid color definition: %s' % x) 

718 

719 

720hours = 3600. 

721days = hours*24 

722approx_months = days*30.5 

723approx_years = days*365 

724 

725 

726nice_time_tinc_inc_approx_units = { 

727 'seconds': 1, 

728 'months': approx_months, 

729 'years': approx_years} 

730 

731 

732def nice_time_tick_inc(tinc_approx): 

733 

734 if tinc_approx >= approx_years: 

735 return max(1.0, nice_value(tinc_approx / approx_years)), 'years' 

736 

737 elif tinc_approx >= approx_months: 

738 nice = [1, 2, 3, 6] 

739 for tinc in nice: 

740 if tinc*approx_months >= tinc_approx or tinc == nice[-1]: 

741 return tinc, 'months' 

742 

743 elif tinc_approx > days: 

744 return nice_value(tinc_approx / days) * days, 'seconds' 

745 

746 elif tinc_approx >= 1.0: 

747 nice = [ 

748 1., 2., 5., 10., 20., 30., 60., 120., 300., 600., 1200., 1800., 

749 1*hours, 2*hours, 3*hours, 6*hours, 12*hours, days, 2*days] 

750 

751 for tinc in nice: 

752 if tinc >= tinc_approx or tinc == nice[-1]: 

753 return tinc, 'seconds' 

754 

755 else: 

756 return nice_value(tinc_approx), 'seconds' 

757 

758 

759def nice_time_tick_inc_approx_secs(tinc_approx): 

760 v, unit = nice_time_tick_inc(tinc_approx) 

761 return v * nice_time_tinc_inc_approx_units[unit] 

762 

763 

764def time_tick_labels(tmin, tmax, tinc, tinc_unit): 

765 

766 if tinc_unit == 'years': 

767 tt = time.gmtime(int(tmin)) 

768 tmin_year = tt[0] 

769 if tt[1:6] != (1, 1, 0, 0, 0): 

770 tmin_year += 1 

771 

772 tmax_year = time.gmtime(int(tmax))[0] 

773 

774 tick_times_year = arange2( 

775 math.ceil(tmin_year/tinc)*tinc, 

776 math.floor(tmax_year/tinc)*tinc, 

777 tinc).astype(int) 

778 

779 times = [ 

780 to_time_float(calendar.timegm((year, 1, 1, 0, 0, 0))) 

781 for year in tick_times_year] 

782 

783 labels = ['%04i' % year for year in tick_times_year] 

784 

785 elif tinc_unit == 'months': 

786 tt = time.gmtime(int(tmin)) 

787 tmin_ym = tt[0] * 12 + (tt[1] - 1) 

788 if tt[2:6] != (1, 0, 0, 0): 

789 tmin_ym += 1 

790 

791 tt = time.gmtime(int(tmax)) 

792 tmax_ym = tt[0] * 12 + (tt[1] - 1) 

793 

794 tick_times_ym = arange2( 

795 math.ceil(tmin_ym/tinc)*tinc, 

796 math.floor(tmax_ym/tinc)*tinc, tinc).astype(int) 

797 

798 times = [ 

799 to_time_float(calendar.timegm((ym // 12, ym % 12 + 1, 1, 0, 0, 0))) 

800 for ym in tick_times_ym] 

801 

802 labels = [ 

803 '%04i-%02i' % (ym // 12, ym % 12 + 1) for ym in tick_times_ym] 

804 

805 elif tinc_unit == 'seconds': 

806 imin = int(num.ceil(tmin/tinc)) 

807 imax = int(num.floor(tmax/tinc)) 

808 nticks = imax - imin + 1 

809 tmin_ticks = imin * tinc 

810 times = tmin_ticks + num.arange(nticks) * tinc 

811 times = times.tolist() 

812 

813 if tinc < 1e-6: 

814 fmt = '%Y-%m-%d.%H:%M:%S.9FRAC' 

815 elif tinc < 1e-3: 

816 fmt = '%Y-%m-%d.%H:%M:%S.6FRAC' 

817 elif tinc < 1.0: 

818 fmt = '%Y-%m-%d.%H:%M:%S.3FRAC' 

819 elif tinc < 60: 

820 fmt = '%Y-%m-%d.%H:%M:%S' 

821 elif tinc < 3600.*24: 

822 fmt = '%Y-%m-%d.%H:%M' 

823 else: 

824 fmt = '%Y-%m-%d' 

825 

826 nwords = len(fmt.split('.')) 

827 

828 labels = [time_to_str(t, format=fmt) for t in times] 

829 labels_weeded = [] 

830 have_ymd = have_hms = False 

831 ymd = hms = '' 

832 for ilab, lab in reversed(list(enumerate(labels))): 

833 words = lab.split('.') 

834 if nwords > 2: 

835 words[2] = '.' + words[2] 

836 if float(words[2]) == 0.0: # or (ilab == 0 and not have_hms): 

837 have_hms = True 

838 else: 

839 hms = words[1] 

840 words[1] = '' 

841 else: 

842 have_hms = True 

843 

844 if nwords > 1: 

845 if words[1] in ('00:00', '00:00:00'): # or (ilab == 0 and not have_ymd): # noqa 

846 have_ymd = True 

847 else: 

848 ymd = words[0] 

849 words[0] = '' 

850 else: 

851 have_ymd = True 

852 

853 labels_weeded.append('\n'.join(reversed(words))) 

854 

855 labels = list(reversed(labels_weeded)) 

856 if (not have_ymd or not have_hms) and (hms or ymd): 

857 words = ([''] if nwords > 2 else []) + [ 

858 hms if not have_hms else '', 

859 ymd if not have_ymd else ''] 

860 

861 labels[0:0] = ['\n'.join(words)] 

862 times[0:0] = [tmin] 

863 

864 return times, labels 

865 

866 

867def mpl_time_axis(axes, approx_ticks=5.): 

868 

869 ''' 

870 Configure x axis of a matplotlib axes object for interactive time display. 

871 

872 :param axes: Axes to be configured. 

873 :type axes: :py:class:`matplotlib.axes.Axes` 

874 

875 :param approx_ticks: Approximate number of ticks to create. 

876 :type approx_ticks: float 

877 

878 This function tries to use nice tick increments and tick labels for time 

879 ranges from microseconds to years, similar to how this is handled in 

880 Snuffler. 

881 ''' 

882 

883 from matplotlib.ticker import Locator, Formatter 

884 

885 class labeled_float(float): 

886 pass 

887 

888 class TimeLocator(Locator): 

889 

890 def __init__(self, approx_ticks=5.): 

891 self._approx_ticks = approx_ticks 

892 Locator.__init__(self) 

893 

894 def __call__(self): 

895 vmin, vmax = self.axis.get_view_interval() 

896 return self.tick_values(vmin, vmax) 

897 

898 def tick_values(self, vmin, vmax): 

899 if vmax < vmin: 

900 vmin, vmax = vmax, vmin 

901 

902 if vmin == vmax: 

903 return [] 

904 

905 tinc_approx = (vmax - vmin) / self._approx_ticks 

906 tinc, tinc_unit = nice_time_tick_inc(tinc_approx) 

907 times, labels = time_tick_labels(vmin, vmax, tinc, tinc_unit) 

908 ftimes = [] 

909 for t, label in zip(times, labels): 

910 ftime = labeled_float(t) 

911 ftime._mpl_label = label 

912 ftimes.append(ftime) 

913 

914 return self.raise_if_exceeds(ftimes) 

915 

916 class TimeFormatter(Formatter): 

917 

918 def __call__(self, x, pos=None): 

919 if isinstance(x, labeled_float): 

920 return x._mpl_label 

921 else: 

922 return time_to_str(x, format='%Y-%m-%d %H:%M:%S.6FRAC') 

923 

924 axes.xaxis.set_major_locator(TimeLocator(approx_ticks=approx_ticks)) 

925 axes.xaxis.set_major_formatter(TimeFormatter())