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

387 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2024-03-07 11: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 try: 

525 from matplotlib.cm import _cmap_registry 

526 except ImportError: 

527 from matplotlib.cm import cmap_d as _cmap_registry 

528 

529 names = list(_cmap_registry.keys()) 

530 

531 names.sort() 

532 return names 

533 

534 

535def mpl_get_cmap(name): 

536 ''' 

537 Compatibility function to get named MPL colormap. 

538 

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

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

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

542 ''' 

543 

544 try: 

545 from matplotlib import colormaps 

546 return colormaps[name] 

547 except ImportError: 

548 from matplotlib import cm 

549 return cm.get_cmap(name) 

550 

551 

552def mpl_margins( 

553 fig, 

554 left=1.0, top=1.0, right=1.0, bottom=1.0, 

555 wspace=None, hspace=None, 

556 w=None, h=None, 

557 nw=None, nh=None, 

558 all=None, 

559 units='inch'): 

560 

561 ''' 

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

563 

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

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

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

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

568 relative to the subplot width and height. 

569 

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

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

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

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

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

575 :param nw: number of subplots horizontally 

576 :param nh: number of subplots vertically 

577 :param wspace: horizontal spacing between subplots 

578 :param hspace: vertical spacing between subplots 

579 ''' 

580 

581 left, top, right, bottom = map( 

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

583 

584 if w is not None: 

585 left = right = float(w) 

586 

587 if h is not None: 

588 top = bottom = float(h) 

589 

590 if all is not None: 

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

592 

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

594 

595 left *= ufac 

596 right *= ufac 

597 top *= ufac 

598 bottom *= ufac 

599 

600 width, height = fig.get_size_inches() 

601 

602 rel_wspace = None 

603 rel_hspace = None 

604 

605 if wspace is not None: 

606 wspace *= ufac 

607 if nw is None: 

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

609 

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

611 rel_wspace = wspace / wsub 

612 else: 

613 wsub = width - left - right 

614 

615 if hspace is not None: 

616 hspace *= ufac 

617 if nh is None: 

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

619 

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

621 rel_hspace = hspace / hsub 

622 else: 

623 hsub = height - top - bottom 

624 

625 fig.subplots_adjust( 

626 left=left/width, 

627 right=1.0 - right/width, 

628 bottom=bottom/height, 

629 top=1.0 - top/height, 

630 wspace=rel_wspace, 

631 hspace=rel_hspace) 

632 

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

634 xpos *= ufac 

635 ypos *= ufac 

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

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

638 

639 return labelpos 

640 

641 

642mpl_margins.__doc__ %= _doc_units 

643 

644 

645def mpl_labelspace(axes): 

646 ''' 

647 Add some extra padding between label and ax annotations. 

648 ''' 

649 

650 xa = axes.get_xaxis() 

651 ya = axes.get_yaxis() 

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

653 if hasattr(xa, attr): 

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

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

656 break 

657 

658 

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

660 ''' 

661 Get paper size in inch from string. 

662 

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

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

665 

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

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

668 

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

670 ''' 

671 

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

673 

674 

675mpl_papersize.__doc__ %= _doc_papersizes 

676 

677 

678class InvalidColorDef(ValueError): 

679 ''' 

680 Raised for invalid color definitions. 

681 ''' 

682 pass 

683 

684 

685def mpl_graph_color(i): 

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

687 

688 

689def mpl_color(x): 

690 ''' 

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

692 

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

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

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

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

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

698 ''' 

699 

700 import matplotlib.colors 

701 

702 if x in tango_colors: 

703 return to01(tango_colors[x]) 

704 

705 s = x.split('/') 

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

707 try: 

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

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

710 return vals 

711 

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

713 return to01(vals) 

714 

715 except ValueError: 

716 try: 

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

718 except Exception: 

719 pass 

720 

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

722 

723 

724hours = 3600. 

725days = hours*24 

726approx_months = days*30.5 

727approx_years = days*365 

728 

729 

730nice_time_tinc_inc_approx_units = { 

731 'seconds': 1, 

732 'months': approx_months, 

733 'years': approx_years} 

734 

735 

736def nice_time_tick_inc(tinc_approx): 

737 

738 if tinc_approx >= approx_years: 

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

740 

741 elif tinc_approx >= approx_months: 

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

743 for tinc in nice: 

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

745 return tinc, 'months' 

746 

747 elif tinc_approx > days: 

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

749 

750 elif tinc_approx >= 1.0: 

751 nice = [ 

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

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

754 

755 for tinc in nice: 

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

757 return tinc, 'seconds' 

758 

759 else: 

760 return nice_value(tinc_approx), 'seconds' 

761 

762 

763def nice_time_tick_inc_approx_secs(tinc_approx): 

764 v, unit = nice_time_tick_inc(tinc_approx) 

765 return v * nice_time_tinc_inc_approx_units[unit] 

766 

767 

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

769 

770 if tinc_unit == 'years': 

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

772 tmin_year = tt[0] 

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

774 tmin_year += 1 

775 

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

777 

778 tick_times_year = arange2( 

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

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

781 tinc).astype(int) 

782 

783 times = [ 

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

785 for year in tick_times_year] 

786 

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

788 

789 elif tinc_unit == 'months': 

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

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

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

793 tmin_ym += 1 

794 

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

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

797 

798 tick_times_ym = arange2( 

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

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

801 

802 times = [ 

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

804 for ym in tick_times_ym] 

805 

806 labels = [ 

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

808 

809 elif tinc_unit == 'seconds': 

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

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

812 nticks = imax - imin + 1 

813 tmin_ticks = imin * tinc 

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

815 times = times.tolist() 

816 

817 if tinc < 1e-6: 

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

819 elif tinc < 1e-3: 

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

821 elif tinc < 1.0: 

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

823 elif tinc < 60: 

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

825 elif tinc < 3600.*24: 

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

827 else: 

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

829 

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

831 

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

833 labels_weeded = [] 

834 have_ymd = have_hms = False 

835 ymd = hms = '' 

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

837 words = lab.split('.') 

838 if nwords > 2: 

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

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

841 have_hms = True 

842 else: 

843 hms = words[1] 

844 words[1] = '' 

845 else: 

846 have_hms = True 

847 

848 if nwords > 1: 

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

850 have_ymd = True 

851 else: 

852 ymd = words[0] 

853 words[0] = '' 

854 else: 

855 have_ymd = True 

856 

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

858 

859 labels = list(reversed(labels_weeded)) 

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

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

862 hms if not have_hms else '', 

863 ymd if not have_ymd else ''] 

864 

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

866 times[0:0] = [tmin] 

867 

868 return times, labels 

869 

870 

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

872 

873 ''' 

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

875 

876 :param axes: Axes to be configured. 

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

878 

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

880 :type approx_ticks: float 

881 

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

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

884 Snuffler. 

885 ''' 

886 

887 from matplotlib.ticker import Locator, Formatter 

888 

889 class labeled_float(float): 

890 pass 

891 

892 class TimeLocator(Locator): 

893 

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

895 self._approx_ticks = approx_ticks 

896 Locator.__init__(self) 

897 

898 def __call__(self): 

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

900 return self.tick_values(vmin, vmax) 

901 

902 def tick_values(self, vmin, vmax): 

903 if vmax < vmin: 

904 vmin, vmax = vmax, vmin 

905 

906 if vmin == vmax: 

907 return [] 

908 

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

910 tinc, tinc_unit = nice_time_tick_inc(tinc_approx) 

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

912 ftimes = [] 

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

914 ftime = labeled_float(t) 

915 ftime._mpl_label = label 

916 ftimes.append(ftime) 

917 

918 return self.raise_if_exceeds(ftimes) 

919 

920 class TimeFormatter(Formatter): 

921 

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

923 if isinstance(x, labeled_float): 

924 return x._mpl_label 

925 else: 

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

927 

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

929 axes.xaxis.set_major_formatter(TimeFormatter())