plotmpl.py 41 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262
  1. #!/usr/bin/env python3
  2. #
  3. # Plot CSV files with matplotlib.
  4. #
  5. # Example:
  6. # ./scripts/plotmpl.py bench.csv -xSIZE -ybench_read -obench.svg
  7. #
  8. # Copyright (c) 2022, The littlefs authors.
  9. # SPDX-License-Identifier: BSD-3-Clause
  10. #
  11. import codecs
  12. import collections as co
  13. import csv
  14. import io
  15. import itertools as it
  16. import logging
  17. import math as m
  18. import numpy as np
  19. import os
  20. import shlex
  21. import shutil
  22. import time
  23. import matplotlib as mpl
  24. import matplotlib.pyplot as plt
  25. # some nicer colors borrowed from Seaborn
  26. # note these include a non-opaque alpha
  27. COLORS = [
  28. '#4c72b0bf', # blue
  29. '#dd8452bf', # orange
  30. '#55a868bf', # green
  31. '#c44e52bf', # red
  32. '#8172b3bf', # purple
  33. '#937860bf', # brown
  34. '#da8bc3bf', # pink
  35. '#8c8c8cbf', # gray
  36. '#ccb974bf', # yellow
  37. '#64b5cdbf', # cyan
  38. ]
  39. COLORS_DARK = [
  40. '#a1c9f4bf', # blue
  41. '#ffb482bf', # orange
  42. '#8de5a1bf', # green
  43. '#ff9f9bbf', # red
  44. '#d0bbffbf', # purple
  45. '#debb9bbf', # brown
  46. '#fab0e4bf', # pink
  47. '#cfcfcfbf', # gray
  48. '#fffea3bf', # yellow
  49. '#b9f2f0bf', # cyan
  50. ]
  51. ALPHAS = [0.75]
  52. FORMATS = ['-']
  53. FORMATS_POINTS = ['.']
  54. FORMATS_POINTS_AND_LINES = ['.-']
  55. WIDTH = 750
  56. HEIGHT = 350
  57. FONT_SIZE = 11
  58. SI_PREFIXES = {
  59. 18: 'E',
  60. 15: 'P',
  61. 12: 'T',
  62. 9: 'G',
  63. 6: 'M',
  64. 3: 'K',
  65. 0: '',
  66. -3: 'm',
  67. -6: 'u',
  68. -9: 'n',
  69. -12: 'p',
  70. -15: 'f',
  71. -18: 'a',
  72. }
  73. SI2_PREFIXES = {
  74. 60: 'Ei',
  75. 50: 'Pi',
  76. 40: 'Ti',
  77. 30: 'Gi',
  78. 20: 'Mi',
  79. 10: 'Ki',
  80. 0: '',
  81. -10: 'mi',
  82. -20: 'ui',
  83. -30: 'ni',
  84. -40: 'pi',
  85. -50: 'fi',
  86. -60: 'ai',
  87. }
  88. # formatter for matplotlib
  89. def si(x):
  90. if x == 0:
  91. return '0'
  92. # figure out prefix and scale
  93. p = 3*int(m.log(abs(x), 10**3))
  94. p = min(18, max(-18, p))
  95. # format with 3 digits of precision
  96. s = '%.3f' % (abs(x) / (10.0**p))
  97. s = s[:3+1]
  98. # truncate but only digits that follow the dot
  99. if '.' in s:
  100. s = s.rstrip('0')
  101. s = s.rstrip('.')
  102. return '%s%s%s' % ('-' if x < 0 else '', s, SI_PREFIXES[p])
  103. # formatter for matplotlib
  104. def si2(x):
  105. if x == 0:
  106. return '0'
  107. # figure out prefix and scale
  108. p = 10*int(m.log(abs(x), 2**10))
  109. p = min(30, max(-30, p))
  110. # format with 3 digits of precision
  111. s = '%.3f' % (abs(x) / (2.0**p))
  112. s = s[:3+1]
  113. # truncate but only digits that follow the dot
  114. if '.' in s:
  115. s = s.rstrip('0')
  116. s = s.rstrip('.')
  117. return '%s%s%s' % ('-' if x < 0 else '', s, SI2_PREFIXES[p])
  118. # parse escape strings
  119. def escape(s):
  120. return codecs.escape_decode(s.encode('utf8'))[0].decode('utf8')
  121. # we want to use MaxNLocator, but since MaxNLocator forces multiples of 10
  122. # to be an option, we can't really...
  123. class AutoMultipleLocator(mpl.ticker.MultipleLocator):
  124. def __init__(self, base, nbins=None):
  125. # note base needs to be floats to avoid integer pow issues
  126. self.base = float(base)
  127. self.nbins = nbins
  128. super().__init__(self.base)
  129. def __call__(self):
  130. # find best tick count, conveniently matplotlib has a function for this
  131. vmin, vmax = self.axis.get_view_interval()
  132. vmin, vmax = mpl.transforms.nonsingular(vmin, vmax, 1e-12, 1e-13)
  133. if self.nbins is not None:
  134. nbins = self.nbins
  135. else:
  136. nbins = np.clip(self.axis.get_tick_space(), 1, 9)
  137. # find the best power, use this as our locator's actual base
  138. scale = self.base ** (m.ceil(m.log((vmax-vmin) / (nbins+1), self.base)))
  139. self.set_params(scale)
  140. return super().__call__()
  141. def openio(path, mode='r', buffering=-1):
  142. # allow '-' for stdin/stdout
  143. if path == '-':
  144. if mode == 'r':
  145. return os.fdopen(os.dup(sys.stdin.fileno()), mode, buffering)
  146. else:
  147. return os.fdopen(os.dup(sys.stdout.fileno()), mode, buffering)
  148. else:
  149. return open(path, mode, buffering)
  150. # parse different data representations
  151. def dat(x):
  152. # allow the first part of an a/b fraction
  153. if '/' in x:
  154. x, _ = x.split('/', 1)
  155. # first try as int
  156. try:
  157. return int(x, 0)
  158. except ValueError:
  159. pass
  160. # then try as float
  161. try:
  162. return float(x)
  163. # just don't allow infinity or nan
  164. if m.isinf(x) or m.isnan(x):
  165. raise ValueError("invalid dat %r" % x)
  166. except ValueError:
  167. pass
  168. # else give up
  169. raise ValueError("invalid dat %r" % x)
  170. def collect(csv_paths, renames=[]):
  171. # collect results from CSV files
  172. results = []
  173. for path in csv_paths:
  174. try:
  175. with openio(path) as f:
  176. reader = csv.DictReader(f, restval='')
  177. for r in reader:
  178. results.append(r)
  179. except FileNotFoundError:
  180. pass
  181. if renames:
  182. for r in results:
  183. # make a copy so renames can overlap
  184. r_ = {}
  185. for new_k, old_k in renames:
  186. if old_k in r:
  187. r_[new_k] = r[old_k]
  188. r.update(r_)
  189. return results
  190. def dataset(results, x=None, y=None, define=[]):
  191. # organize by 'by', x, and y
  192. dataset = {}
  193. i = 0
  194. for r in results:
  195. # filter results by matching defines
  196. if not all(k in r and r[k] in vs for k, vs in define):
  197. continue
  198. # find xs
  199. if x is not None:
  200. if x not in r:
  201. continue
  202. try:
  203. x_ = dat(r[x])
  204. except ValueError:
  205. continue
  206. else:
  207. x_ = i
  208. i += 1
  209. # find ys
  210. if y is not None:
  211. if y not in r:
  212. continue
  213. try:
  214. y_ = dat(r[y])
  215. except ValueError:
  216. continue
  217. else:
  218. y_ = None
  219. if y_ is not None:
  220. dataset[x_] = y_ + dataset.get(x_, 0)
  221. else:
  222. dataset[x_] = y_ or dataset.get(x_, None)
  223. return dataset
  224. def datasets(results, by=None, x=None, y=None, define=[]):
  225. # filter results by matching defines
  226. results_ = []
  227. for r in results:
  228. if all(k in r and r[k] in vs for k, vs in define):
  229. results_.append(r)
  230. results = results_
  231. # if y not specified, try to guess from data
  232. if y is None:
  233. y = co.OrderedDict()
  234. for r in results:
  235. for k, v in r.items():
  236. if (by is None or k not in by) and v.strip():
  237. try:
  238. dat(v)
  239. y[k] = True
  240. except ValueError:
  241. y[k] = False
  242. y = list(k for k,v in y.items() if v)
  243. if by is not None:
  244. # find all 'by' values
  245. ks = set()
  246. for r in results:
  247. ks.add(tuple(r.get(k, '') for k in by))
  248. ks = sorted(ks)
  249. # collect all datasets
  250. datasets = co.OrderedDict()
  251. for ks_ in (ks if by is not None else [()]):
  252. for x_ in (x if x is not None else [None]):
  253. for y_ in y:
  254. # hide x/y if there is only one field
  255. k_x = x_ if len(x or []) > 1 else ''
  256. k_y = y_ if len(y or []) > 1 or (not ks_ and not k_x) else ''
  257. datasets[ks_ + (k_x, k_y)] = dataset(
  258. results,
  259. x_,
  260. y_,
  261. [(by_, {k_}) for by_, k_ in zip(by, ks_)]
  262. if by is not None else [])
  263. return datasets
  264. # some classes for organizing subplots into a grid
  265. class Subplot:
  266. def __init__(self, **args):
  267. self.x = 0
  268. self.y = 0
  269. self.xspan = 1
  270. self.yspan = 1
  271. self.args = args
  272. class Grid:
  273. def __init__(self, subplot, width=1.0, height=1.0):
  274. self.xweights = [width]
  275. self.yweights = [height]
  276. self.map = {(0,0): subplot}
  277. self.subplots = [subplot]
  278. def __repr__(self):
  279. return 'Grid(%r, %r)' % (self.xweights, self.yweights)
  280. @property
  281. def width(self):
  282. return len(self.xweights)
  283. @property
  284. def height(self):
  285. return len(self.yweights)
  286. def __iter__(self):
  287. return iter(self.subplots)
  288. def __getitem__(self, i):
  289. x, y = i
  290. if x < 0:
  291. x += len(self.xweights)
  292. if y < 0:
  293. y += len(self.yweights)
  294. return self.map[(x,y)]
  295. def merge(self, other, dir):
  296. if dir in ['above', 'below']:
  297. # first scale the two grids so they line up
  298. self_xweights = self.xweights
  299. other_xweights = other.xweights
  300. self_w = sum(self_xweights)
  301. other_w = sum(other_xweights)
  302. ratio = self_w / other_w
  303. other_xweights = [s*ratio for s in other_xweights]
  304. # now interleave xweights as needed
  305. new_xweights = []
  306. self_map = {}
  307. other_map = {}
  308. self_i = 0
  309. other_i = 0
  310. self_xweight = (self_xweights[self_i]
  311. if self_i < len(self_xweights) else m.inf)
  312. other_xweight = (other_xweights[other_i]
  313. if other_i < len(other_xweights) else m.inf)
  314. while self_i < len(self_xweights) and other_i < len(other_xweights):
  315. if other_xweight - self_xweight > 0.0000001:
  316. new_xweights.append(self_xweight)
  317. other_xweight -= self_xweight
  318. new_i = len(new_xweights)-1
  319. for j in range(len(self.yweights)):
  320. self_map[(new_i, j)] = self.map[(self_i, j)]
  321. for j in range(len(other.yweights)):
  322. other_map[(new_i, j)] = other.map[(other_i, j)]
  323. for s in other.subplots:
  324. if s.x+s.xspan-1 == new_i:
  325. s.xspan += 1
  326. elif s.x > new_i:
  327. s.x += 1
  328. self_i += 1
  329. self_xweight = (self_xweights[self_i]
  330. if self_i < len(self_xweights) else m.inf)
  331. elif self_xweight - other_xweight > 0.0000001:
  332. new_xweights.append(other_xweight)
  333. self_xweight -= other_xweight
  334. new_i = len(new_xweights)-1
  335. for j in range(len(other.yweights)):
  336. other_map[(new_i, j)] = other.map[(other_i, j)]
  337. for j in range(len(self.yweights)):
  338. self_map[(new_i, j)] = self.map[(self_i, j)]
  339. for s in self.subplots:
  340. if s.x+s.xspan-1 == new_i:
  341. s.xspan += 1
  342. elif s.x > new_i:
  343. s.x += 1
  344. other_i += 1
  345. other_xweight = (other_xweights[other_i]
  346. if other_i < len(other_xweights) else m.inf)
  347. else:
  348. new_xweights.append(self_xweight)
  349. new_i = len(new_xweights)-1
  350. for j in range(len(self.yweights)):
  351. self_map[(new_i, j)] = self.map[(self_i, j)]
  352. for j in range(len(other.yweights)):
  353. other_map[(new_i, j)] = other.map[(other_i, j)]
  354. self_i += 1
  355. self_xweight = (self_xweights[self_i]
  356. if self_i < len(self_xweights) else m.inf)
  357. other_i += 1
  358. other_xweight = (other_xweights[other_i]
  359. if other_i < len(other_xweights) else m.inf)
  360. # squish so ratios are preserved
  361. self_h = sum(self.yweights)
  362. other_h = sum(other.yweights)
  363. ratio = (self_h-other_h) / self_h
  364. self_yweights = [s*ratio for s in self.yweights]
  365. # finally concatenate the two grids
  366. if dir == 'above':
  367. for s in other.subplots:
  368. s.y += len(self_yweights)
  369. self.subplots.extend(other.subplots)
  370. self.xweights = new_xweights
  371. self.yweights = self_yweights + other.yweights
  372. self.map = self_map | {(x, y+len(self_yweights)): s
  373. for (x, y), s in other_map.items()}
  374. else:
  375. for s in self.subplots:
  376. s.y += len(other.yweights)
  377. self.subplots.extend(other.subplots)
  378. self.xweights = new_xweights
  379. self.yweights = other.yweights + self_yweights
  380. self.map = other_map | {(x, y+len(other.yweights)): s
  381. for (x, y), s in self_map.items()}
  382. if dir in ['right', 'left']:
  383. # first scale the two grids so they line up
  384. self_yweights = self.yweights
  385. other_yweights = other.yweights
  386. self_h = sum(self_yweights)
  387. other_h = sum(other_yweights)
  388. ratio = self_h / other_h
  389. other_yweights = [s*ratio for s in other_yweights]
  390. # now interleave yweights as needed
  391. new_yweights = []
  392. self_map = {}
  393. other_map = {}
  394. self_i = 0
  395. other_i = 0
  396. self_yweight = (self_yweights[self_i]
  397. if self_i < len(self_yweights) else m.inf)
  398. other_yweight = (other_yweights[other_i]
  399. if other_i < len(other_yweights) else m.inf)
  400. while self_i < len(self_yweights) and other_i < len(other_yweights):
  401. if other_yweight - self_yweight > 0.0000001:
  402. new_yweights.append(self_yweight)
  403. other_yweight -= self_yweight
  404. new_i = len(new_yweights)-1
  405. for j in range(len(self.xweights)):
  406. self_map[(j, new_i)] = self.map[(j, self_i)]
  407. for j in range(len(other.xweights)):
  408. other_map[(j, new_i)] = other.map[(j, other_i)]
  409. for s in other.subplots:
  410. if s.y+s.yspan-1 == new_i:
  411. s.yspan += 1
  412. elif s.y > new_i:
  413. s.y += 1
  414. self_i += 1
  415. self_yweight = (self_yweights[self_i]
  416. if self_i < len(self_yweights) else m.inf)
  417. elif self_yweight - other_yweight > 0.0000001:
  418. new_yweights.append(other_yweight)
  419. self_yweight -= other_yweight
  420. new_i = len(new_yweights)-1
  421. for j in range(len(other.xweights)):
  422. other_map[(j, new_i)] = other.map[(j, other_i)]
  423. for j in range(len(self.xweights)):
  424. self_map[(j, new_i)] = self.map[(j, self_i)]
  425. for s in self.subplots:
  426. if s.y+s.yspan-1 == new_i:
  427. s.yspan += 1
  428. elif s.y > new_i:
  429. s.y += 1
  430. other_i += 1
  431. other_yweight = (other_yweights[other_i]
  432. if other_i < len(other_yweights) else m.inf)
  433. else:
  434. new_yweights.append(self_yweight)
  435. new_i = len(new_yweights)-1
  436. for j in range(len(self.xweights)):
  437. self_map[(j, new_i)] = self.map[(j, self_i)]
  438. for j in range(len(other.xweights)):
  439. other_map[(j, new_i)] = other.map[(j, other_i)]
  440. self_i += 1
  441. self_yweight = (self_yweights[self_i]
  442. if self_i < len(self_yweights) else m.inf)
  443. other_i += 1
  444. other_yweight = (other_yweights[other_i]
  445. if other_i < len(other_yweights) else m.inf)
  446. # squish so ratios are preserved
  447. self_w = sum(self.xweights)
  448. other_w = sum(other.xweights)
  449. ratio = (self_w-other_w) / self_w
  450. self_xweights = [s*ratio for s in self.xweights]
  451. # finally concatenate the two grids
  452. if dir == 'right':
  453. for s in other.subplots:
  454. s.x += len(self_xweights)
  455. self.subplots.extend(other.subplots)
  456. self.xweights = self_xweights + other.xweights
  457. self.yweights = new_yweights
  458. self.map = self_map | {(x+len(self_xweights), y): s
  459. for (x, y), s in other_map.items()}
  460. else:
  461. for s in self.subplots:
  462. s.x += len(other.xweights)
  463. self.subplots.extend(other.subplots)
  464. self.xweights = other.xweights + self_xweights
  465. self.yweights = new_yweights
  466. self.map = other_map | {(x+len(other.xweights), y): s
  467. for (x, y), s in self_map.items()}
  468. def scale(self, width, height):
  469. self.xweights = [s*width for s in self.xweights]
  470. self.yweights = [s*height for s in self.yweights]
  471. @classmethod
  472. def fromargs(cls, width=1.0, height=1.0, *,
  473. subplots=[],
  474. **args):
  475. grid = cls(Subplot(**args))
  476. for dir, subargs in subplots:
  477. subgrid = cls.fromargs(
  478. width=subargs.pop('width',
  479. 0.5 if dir in ['right', 'left'] else width),
  480. height=subargs.pop('height',
  481. 0.5 if dir in ['above', 'below'] else height),
  482. **subargs)
  483. grid.merge(subgrid, dir)
  484. grid.scale(width, height)
  485. return grid
  486. def main(csv_paths, output, *,
  487. svg=False,
  488. png=False,
  489. quiet=False,
  490. by=None,
  491. x=None,
  492. y=None,
  493. define=[],
  494. points=False,
  495. points_and_lines=False,
  496. colors=None,
  497. formats=None,
  498. width=WIDTH,
  499. height=HEIGHT,
  500. xlim=(None,None),
  501. ylim=(None,None),
  502. xlog=False,
  503. ylog=False,
  504. x2=False,
  505. y2=False,
  506. xticks=None,
  507. yticks=None,
  508. xunits=None,
  509. yunits=None,
  510. xlabel=None,
  511. ylabel=None,
  512. xticklabels=None,
  513. yticklabels=None,
  514. title=None,
  515. legend_right=False,
  516. legend_above=False,
  517. legend_below=False,
  518. dark=False,
  519. ggplot=False,
  520. xkcd=False,
  521. github=False,
  522. font=None,
  523. font_size=FONT_SIZE,
  524. font_color=None,
  525. foreground=None,
  526. background=None,
  527. subplot={},
  528. subplots=[],
  529. **args):
  530. # guess the output format
  531. if not png and not svg:
  532. if output.endswith('.png'):
  533. png = True
  534. else:
  535. svg = True
  536. # some shortcuts for color schemes
  537. if github:
  538. ggplot = True
  539. if font_color is None:
  540. if dark:
  541. font_color = '#c9d1d9'
  542. else:
  543. font_color = '#24292f'
  544. if foreground is None:
  545. if dark:
  546. foreground = '#343942'
  547. else:
  548. foreground = '#eff1f3'
  549. if background is None:
  550. if dark:
  551. background = '#0d1117'
  552. else:
  553. background = '#ffffff'
  554. # what colors/alphas/formats to use?
  555. if colors is not None:
  556. colors_ = colors
  557. elif dark:
  558. colors_ = COLORS_DARK
  559. else:
  560. colors_ = COLORS
  561. if formats is not None:
  562. formats_ = formats
  563. elif points_and_lines:
  564. formats_ = FORMATS_POINTS_AND_LINES
  565. elif points:
  566. formats_ = FORMATS_POINTS
  567. else:
  568. formats_ = FORMATS
  569. if font_color is not None:
  570. font_color_ = font_color
  571. elif dark:
  572. font_color_ = '#ffffff'
  573. else:
  574. font_color_ = '#000000'
  575. if foreground is not None:
  576. foreground_ = foreground
  577. elif dark:
  578. foreground_ = '#333333'
  579. else:
  580. foreground_ = '#e5e5e5'
  581. if background is not None:
  582. background_ = background
  583. elif dark:
  584. background_ = '#000000'
  585. else:
  586. background_ = '#ffffff'
  587. # configure some matplotlib settings
  588. if xkcd:
  589. # the font search here prints a bunch of unhelpful warnings
  590. logging.getLogger('matplotlib.font_manager').setLevel(logging.ERROR)
  591. plt.xkcd()
  592. # turn off the white outline, this breaks some things
  593. plt.rc('path', effects=[])
  594. if ggplot:
  595. plt.style.use('ggplot')
  596. plt.rc('patch', linewidth=0)
  597. plt.rc('axes', facecolor=foreground_, edgecolor=background_)
  598. plt.rc('grid', color=background_)
  599. # fix the the gridlines when ggplot+xkcd
  600. if xkcd:
  601. plt.rc('grid', linewidth=1)
  602. plt.rc('axes.spines', bottom=False, left=False)
  603. if dark:
  604. plt.style.use('dark_background')
  605. plt.rc('savefig', facecolor='auto', edgecolor='auto')
  606. # fix ggplot when dark
  607. if ggplot:
  608. plt.rc('axes',
  609. facecolor=foreground_,
  610. edgecolor=background_)
  611. plt.rc('grid', color=background_)
  612. if font is not None:
  613. plt.rc('font', family=font)
  614. plt.rc('font', size=font_size)
  615. plt.rc('text', color=font_color_)
  616. plt.rc('figure',
  617. titlesize='medium',
  618. labelsize='small')
  619. plt.rc('axes',
  620. titlesize='small',
  621. labelsize='small',
  622. labelcolor=font_color_)
  623. if not ggplot:
  624. plt.rc('axes', edgecolor=font_color_)
  625. plt.rc('xtick', labelsize='small', color=font_color_)
  626. plt.rc('ytick', labelsize='small', color=font_color_)
  627. plt.rc('legend',
  628. fontsize='small',
  629. fancybox=False,
  630. framealpha=None,
  631. edgecolor=foreground_,
  632. borderaxespad=0)
  633. plt.rc('axes.spines', top=False, right=False)
  634. plt.rc('figure', facecolor=background_, edgecolor=background_)
  635. if not ggplot:
  636. plt.rc('axes', facecolor='#00000000')
  637. # I think the svg backend just ignores DPI, but seems to use something
  638. # equivalent to 96, maybe this is the default for SVG rendering?
  639. plt.rc('figure', dpi=96)
  640. # separate out renames
  641. renames = list(it.chain.from_iterable(
  642. ((k, v) for v in vs)
  643. for k, vs in it.chain(by or [], x or [], y or [])))
  644. if by is not None:
  645. by = [k for k, _ in by]
  646. if x is not None:
  647. x = [k for k, _ in x]
  648. if y is not None:
  649. y = [k for k, _ in y]
  650. # first collect results from CSV files
  651. results = collect(csv_paths, renames)
  652. # then extract the requested datasets
  653. datasets_ = datasets(results, by, x, y, define)
  654. # figure out formats/colors here so that subplot defines
  655. # don't change them later, that'd be bad
  656. dataformats_ = {
  657. name: formats_[i % len(formats_)]
  658. for i, name in enumerate(datasets_.keys())}
  659. datacolors_ = {
  660. name: colors_[i % len(colors_)]
  661. for i, name in enumerate(datasets_.keys())}
  662. # create a grid of subplots
  663. grid = Grid.fromargs(
  664. subplots=subplots + subplot.pop('subplots', []),
  665. **subplot)
  666. # create a matplotlib plot
  667. fig = plt.figure(figsize=(
  668. width/plt.rcParams['figure.dpi'],
  669. height/plt.rcParams['figure.dpi']),
  670. layout='constrained',
  671. # we need a linewidth to keep xkcd mode happy
  672. linewidth=8 if xkcd else 0)
  673. gs = fig.add_gridspec(
  674. grid.height
  675. + (1 if legend_above else 0)
  676. + (1 if legend_below else 0),
  677. grid.width
  678. + (1 if legend_right else 0),
  679. height_ratios=([0.001] if legend_above else [])
  680. + [max(s, 0.01) for s in reversed(grid.yweights)]
  681. + ([0.001] if legend_below else []),
  682. width_ratios=[max(s, 0.01) for s in grid.xweights]
  683. + ([0.001] if legend_right else []))
  684. # first create axes so that plots can interact with each other
  685. for s in grid:
  686. s.ax = fig.add_subplot(gs[
  687. grid.height-(s.y+s.yspan) + (1 if legend_above else 0)
  688. : grid.height-s.y + (1 if legend_above else 0),
  689. s.x
  690. : s.x+s.xspan])
  691. # now plot each subplot
  692. for s in grid:
  693. # allow subplot params to override global params
  694. define_ = define + s.args.get('define', [])
  695. xlim_ = s.args.get('xlim', xlim)
  696. ylim_ = s.args.get('ylim', ylim)
  697. xlog_ = s.args.get('xlog', False) or xlog
  698. ylog_ = s.args.get('ylog', False) or ylog
  699. x2_ = s.args.get('x2', False) or x2
  700. y2_ = s.args.get('y2', False) or y2
  701. xticks_ = s.args.get('xticks', xticks)
  702. yticks_ = s.args.get('yticks', yticks)
  703. xunits_ = s.args.get('xunits', xunits)
  704. yunits_ = s.args.get('yunits', yunits)
  705. xticklabels_ = s.args.get('xticklabels', xticklabels)
  706. yticklabels_ = s.args.get('yticklabels', yticklabels)
  707. # label/titles are handled a bit differently in subplots
  708. subtitle = s.args.get('title')
  709. xsublabel = s.args.get('xlabel')
  710. ysublabel = s.args.get('ylabel')
  711. # allow shortened ranges
  712. if len(xlim_) == 1:
  713. xlim_ = (0, xlim_[0])
  714. if len(ylim_) == 1:
  715. ylim_ = (0, ylim_[0])
  716. # data can be constrained by subplot-specific defines,
  717. # so re-extract for each plot
  718. subdatasets = datasets(results, by, x, y, define_)
  719. # plot!
  720. ax = s.ax
  721. for name, dataset in subdatasets.items():
  722. dats = sorted((x,y) for x,y in dataset.items())
  723. ax.plot([x for x,_ in dats], [y for _,y in dats],
  724. dataformats_[name],
  725. color=datacolors_[name],
  726. label=','.join(k for k in name if k))
  727. # axes scaling
  728. if xlog_:
  729. ax.set_xscale('symlog')
  730. ax.xaxis.set_minor_locator(mpl.ticker.NullLocator())
  731. if ylog_:
  732. ax.set_yscale('symlog')
  733. ax.yaxis.set_minor_locator(mpl.ticker.NullLocator())
  734. # axes limits
  735. ax.set_xlim(
  736. xlim_[0] if xlim_[0] is not None
  737. else min(it.chain([0], (k
  738. for r in subdatasets.values()
  739. for k, v in r.items()
  740. if v is not None))),
  741. xlim_[1] if xlim_[1] is not None
  742. else max(it.chain([0], (k
  743. for r in subdatasets.values()
  744. for k, v in r.items()
  745. if v is not None))))
  746. ax.set_ylim(
  747. ylim_[0] if ylim_[0] is not None
  748. else min(it.chain([0], (v
  749. for r in subdatasets.values()
  750. for _, v in r.items()
  751. if v is not None))),
  752. ylim_[1] if ylim_[1] is not None
  753. else max(it.chain([0], (v
  754. for r in subdatasets.values()
  755. for _, v in r.items()
  756. if v is not None))))
  757. # axes ticks
  758. if x2_:
  759. ax.xaxis.set_major_formatter(lambda x, pos:
  760. si2(x)+(xunits_ if xunits_ else ''))
  761. if xticklabels_ is not None:
  762. ax.xaxis.set_ticklabels(xticklabels_)
  763. if xticks_ is None:
  764. ax.xaxis.set_major_locator(AutoMultipleLocator(2))
  765. elif isinstance(xticks_, list):
  766. ax.xaxis.set_major_locator(mpl.ticker.FixedLocator(xticks_))
  767. elif xticks_ != 0:
  768. ax.xaxis.set_major_locator(AutoMultipleLocator(2, xticks_-1))
  769. else:
  770. ax.xaxis.set_major_locator(mpl.ticker.NullLocator())
  771. else:
  772. ax.xaxis.set_major_formatter(lambda x, pos:
  773. si(x)+(xunits_ if xunits_ else ''))
  774. if xticklabels_ is not None:
  775. ax.xaxis.set_ticklabels(xticklabels_)
  776. if xticks_ is None:
  777. ax.xaxis.set_major_locator(mpl.ticker.AutoLocator())
  778. elif isinstance(xticks_, list):
  779. ax.xaxis.set_major_locator(mpl.ticker.FixedLocator(xticks_))
  780. elif xticks_ != 0:
  781. ax.xaxis.set_major_locator(mpl.ticker.MaxNLocator(xticks_-1))
  782. else:
  783. ax.xaxis.set_major_locator(mpl.ticker.NullLocator())
  784. if y2_:
  785. ax.yaxis.set_major_formatter(lambda x, pos:
  786. si2(x)+(yunits_ if yunits_ else ''))
  787. if yticklabels_ is not None:
  788. ax.yaxis.set_ticklabels(yticklabels_)
  789. if yticks_ is None:
  790. ax.yaxis.set_major_locator(AutoMultipleLocator(2))
  791. elif isinstance(yticks_, list):
  792. ax.yaxis.set_major_locator(mpl.ticker.FixedLocator(yticks_))
  793. elif yticks_ != 0:
  794. ax.yaxis.set_major_locator(AutoMultipleLocator(2, yticks_-1))
  795. else:
  796. ax.yaxis.set_major_locator(mpl.ticker.NullLocator())
  797. else:
  798. ax.yaxis.set_major_formatter(lambda x, pos:
  799. si(x)+(yunits_ if yunits_ else ''))
  800. if yticklabels_ is not None:
  801. ax.yaxis.set_ticklabels(yticklabels_)
  802. if yticks_ is None:
  803. ax.yaxis.set_major_locator(mpl.ticker.AutoLocator())
  804. elif isinstance(yticks_, list):
  805. ax.yaxis.set_major_locator(mpl.ticker.FixedLocator(yticks_))
  806. elif yticks_ != 0:
  807. ax.yaxis.set_major_locator(mpl.ticker.MaxNLocator(yticks_-1))
  808. else:
  809. ax.yaxis.set_major_locator(mpl.ticker.NullLocator())
  810. if ggplot:
  811. ax.grid(sketch_params=None)
  812. # axes subplot labels
  813. if xsublabel is not None:
  814. ax.set_xlabel(escape(xsublabel))
  815. if ysublabel is not None:
  816. ax.set_ylabel(escape(ysublabel))
  817. if subtitle is not None:
  818. ax.set_title(escape(subtitle))
  819. # add a legend? a bit tricky with matplotlib
  820. #
  821. # the best solution I've found is a dedicated, invisible axes for the
  822. # legend, hacky, but it works.
  823. #
  824. # note this was written before constrained_layout supported legend
  825. # collisions, hopefully this is added in the future
  826. labels = co.OrderedDict()
  827. for s in grid:
  828. for h, l in zip(*s.ax.get_legend_handles_labels()):
  829. labels[l] = h
  830. if legend_right:
  831. ax = fig.add_subplot(gs[(1 if legend_above else 0):,-1])
  832. ax.set_axis_off()
  833. ax.legend(
  834. labels.values(),
  835. labels.keys(),
  836. loc='upper left',
  837. fancybox=False,
  838. borderaxespad=0)
  839. if legend_above:
  840. ax = fig.add_subplot(gs[0, :grid.width])
  841. ax.set_axis_off()
  842. # try different column counts until we fit in the axes
  843. for ncol in reversed(range(1, len(labels)+1)):
  844. legend_ = ax.legend(
  845. labels.values(),
  846. labels.keys(),
  847. loc='upper center',
  848. ncol=ncol,
  849. fancybox=False,
  850. borderaxespad=0)
  851. if (legend_.get_window_extent().width
  852. <= ax.get_window_extent().width):
  853. break
  854. if legend_below:
  855. ax = fig.add_subplot(gs[-1, :grid.width])
  856. ax.set_axis_off()
  857. # big hack to get xlabel above the legend! but hey this
  858. # works really well actually
  859. if xlabel:
  860. ax.set_title(escape(xlabel),
  861. size=plt.rcParams['axes.labelsize'],
  862. weight=plt.rcParams['axes.labelweight'])
  863. # try different column counts until we fit in the axes
  864. for ncol in reversed(range(1, len(labels)+1)):
  865. legend_ = ax.legend(
  866. labels.values(),
  867. labels.keys(),
  868. loc='upper center',
  869. ncol=ncol,
  870. fancybox=False,
  871. borderaxespad=0)
  872. if (legend_.get_window_extent().width
  873. <= ax.get_window_extent().width):
  874. break
  875. # axes labels, NOTE we reposition these below
  876. if xlabel is not None and not legend_below:
  877. fig.supxlabel(escape(xlabel))
  878. if ylabel is not None:
  879. fig.supylabel(escape(ylabel))
  880. if title is not None:
  881. fig.suptitle(escape(title))
  882. # precompute constrained layout and find midpoints to adjust things
  883. # that should be centered so they are actually centered
  884. fig.canvas.draw()
  885. xmid = (grid[0,0].ax.get_position().x0 + grid[-1,0].ax.get_position().x1)/2
  886. ymid = (grid[0,0].ax.get_position().y0 + grid[0,-1].ax.get_position().y1)/2
  887. if xlabel is not None and not legend_below:
  888. fig.supxlabel(escape(xlabel), x=xmid)
  889. if ylabel is not None:
  890. fig.supylabel(escape(ylabel), y=ymid)
  891. if title is not None:
  892. fig.suptitle(escape(title), x=xmid)
  893. # write the figure!
  894. plt.savefig(output, format='png' if png else 'svg')
  895. # some stats
  896. if not quiet:
  897. print('updated %s, %s datasets, %s points' % (
  898. output,
  899. len(datasets_),
  900. sum(len(dataset) for dataset in datasets_.values())))
  901. if __name__ == "__main__":
  902. import sys
  903. import argparse
  904. parser = argparse.ArgumentParser(
  905. description="Plot CSV files with matplotlib.",
  906. allow_abbrev=False)
  907. parser.add_argument(
  908. 'csv_paths',
  909. nargs='*',
  910. help="Input *.csv files.")
  911. output_rule = parser.add_argument(
  912. '-o', '--output',
  913. required=True,
  914. help="Output *.svg/*.png file.")
  915. parser.add_argument(
  916. '--svg',
  917. action='store_true',
  918. help="Output an svg file. By default this is infered.")
  919. parser.add_argument(
  920. '--png',
  921. action='store_true',
  922. help="Output a png file. By default this is infered.")
  923. parser.add_argument(
  924. '-q', '--quiet',
  925. action='store_true',
  926. help="Don't print info.")
  927. parser.add_argument(
  928. '-b', '--by',
  929. action='append',
  930. type=lambda x: (
  931. lambda k,v=None: (k, v.split(',') if v is not None else ())
  932. )(*x.split('=', 1)),
  933. help="Group by this field. Can rename fields with new_name=old_name.")
  934. parser.add_argument(
  935. '-x',
  936. action='append',
  937. type=lambda x: (
  938. lambda k,v=None: (k, v.split(',') if v is not None else ())
  939. )(*x.split('=', 1)),
  940. help="Field to use for the x-axis. Can rename fields with "
  941. "new_name=old_name.")
  942. parser.add_argument(
  943. '-y',
  944. action='append',
  945. type=lambda x: (
  946. lambda k,v=None: (k, v.split(',') if v is not None else ())
  947. )(*x.split('=', 1)),
  948. help="Field to use for the y-axis. Can rename fields with "
  949. "new_name=old_name.")
  950. parser.add_argument(
  951. '-D', '--define',
  952. type=lambda x: (lambda k,v: (k, set(v.split(','))))(*x.split('=', 1)),
  953. action='append',
  954. help="Only include results where this field is this value. May include "
  955. "comma-separated options.")
  956. parser.add_argument(
  957. '-.', '--points',
  958. action='store_true',
  959. help="Only draw data points.")
  960. parser.add_argument(
  961. '-!', '--points-and-lines',
  962. action='store_true',
  963. help="Draw data points and lines.")
  964. parser.add_argument(
  965. '--colors',
  966. type=lambda x: [x.strip() for x in x.split(',')],
  967. help="Comma-separated hex colors to use.")
  968. parser.add_argument(
  969. '--formats',
  970. type=lambda x: [x.strip().replace('0',',') for x in x.split(',')],
  971. help="Comma-separated matplotlib formats to use. Allows '0' as an "
  972. "alternative for ','.")
  973. parser.add_argument(
  974. '-W', '--width',
  975. type=lambda x: int(x, 0),
  976. help="Width in pixels. Defaults to %r." % WIDTH)
  977. parser.add_argument(
  978. '-H', '--height',
  979. type=lambda x: int(x, 0),
  980. help="Height in pixels. Defaults to %r." % HEIGHT)
  981. parser.add_argument(
  982. '-X', '--xlim',
  983. type=lambda x: tuple(
  984. dat(x) if x.strip() else None
  985. for x in x.split(',')),
  986. help="Range for the x-axis.")
  987. parser.add_argument(
  988. '-Y', '--ylim',
  989. type=lambda x: tuple(
  990. dat(x) if x.strip() else None
  991. for x in x.split(',')),
  992. help="Range for the y-axis.")
  993. parser.add_argument(
  994. '--xlog',
  995. action='store_true',
  996. help="Use a logarithmic x-axis.")
  997. parser.add_argument(
  998. '--ylog',
  999. action='store_true',
  1000. help="Use a logarithmic y-axis.")
  1001. parser.add_argument(
  1002. '--x2',
  1003. action='store_true',
  1004. help="Use base-2 prefixes for the x-axis.")
  1005. parser.add_argument(
  1006. '--y2',
  1007. action='store_true',
  1008. help="Use base-2 prefixes for the y-axis.")
  1009. parser.add_argument(
  1010. '--xticks',
  1011. type=lambda x: int(x, 0) if ',' not in x
  1012. else [dat(x) for x in x.split(',')],
  1013. help="Ticks for the x-axis. This can be explicit comma-separated "
  1014. "ticks, the number of ticks, or 0 to disable.")
  1015. parser.add_argument(
  1016. '--yticks',
  1017. type=lambda x: int(x, 0) if ',' not in x
  1018. else [dat(x) for x in x.split(',')],
  1019. help="Ticks for the y-axis. This can be explicit comma-separated "
  1020. "ticks, the number of ticks, or 0 to disable.")
  1021. parser.add_argument(
  1022. '--xunits',
  1023. help="Units for the x-axis.")
  1024. parser.add_argument(
  1025. '--yunits',
  1026. help="Units for the y-axis.")
  1027. parser.add_argument(
  1028. '--xlabel',
  1029. help="Add a label to the x-axis.")
  1030. parser.add_argument(
  1031. '--ylabel',
  1032. help="Add a label to the y-axis.")
  1033. parser.add_argument(
  1034. '--xticklabels',
  1035. type=lambda x:
  1036. [x.strip() for x in x.split(',')]
  1037. if x.strip() else [],
  1038. help="Comma separated xticklabels.")
  1039. parser.add_argument(
  1040. '--yticklabels',
  1041. type=lambda x:
  1042. [x.strip() for x in x.split(',')]
  1043. if x.strip() else [],
  1044. help="Comma separated yticklabels.")
  1045. parser.add_argument(
  1046. '-t', '--title',
  1047. help="Add a title.")
  1048. parser.add_argument(
  1049. '-l', '--legend-right',
  1050. action='store_true',
  1051. help="Place a legend to the right.")
  1052. parser.add_argument(
  1053. '--legend-above',
  1054. action='store_true',
  1055. help="Place a legend above.")
  1056. parser.add_argument(
  1057. '--legend-below',
  1058. action='store_true',
  1059. help="Place a legend below.")
  1060. parser.add_argument(
  1061. '--dark',
  1062. action='store_true',
  1063. help="Use the dark style.")
  1064. parser.add_argument(
  1065. '--ggplot',
  1066. action='store_true',
  1067. help="Use the ggplot style.")
  1068. parser.add_argument(
  1069. '--xkcd',
  1070. action='store_true',
  1071. help="Use the xkcd style.")
  1072. parser.add_argument(
  1073. '--github',
  1074. action='store_true',
  1075. help="Use the ggplot style with GitHub colors.")
  1076. parser.add_argument(
  1077. '--font',
  1078. type=lambda x: [x.strip() for x in x.split(',')],
  1079. help="Font family for matplotlib.")
  1080. parser.add_argument(
  1081. '--font-size',
  1082. help="Font size for matplotlib. Defaults to %r." % FONT_SIZE)
  1083. parser.add_argument(
  1084. '--font-color',
  1085. help="Color for the font and other line elements.")
  1086. parser.add_argument(
  1087. '--foreground',
  1088. help="Foreground color to use.")
  1089. parser.add_argument(
  1090. '--background',
  1091. help="Background color to use.")
  1092. class AppendSubplot(argparse.Action):
  1093. @staticmethod
  1094. def parse(value):
  1095. import copy
  1096. subparser = copy.deepcopy(parser)
  1097. next(a for a in subparser._actions
  1098. if '--output' in a.option_strings).required = False
  1099. next(a for a in subparser._actions
  1100. if '--width' in a.option_strings).type = float
  1101. next(a for a in subparser._actions
  1102. if '--height' in a.option_strings).type = float
  1103. return subparser.parse_intermixed_args(shlex.split(value or ""))
  1104. def __call__(self, parser, namespace, value, option):
  1105. if not hasattr(namespace, 'subplots'):
  1106. namespace.subplots = []
  1107. namespace.subplots.append((
  1108. option.split('-')[-1],
  1109. self.__class__.parse(value)))
  1110. parser.add_argument(
  1111. '--subplot-above',
  1112. action=AppendSubplot,
  1113. help="Add subplot above with the same dataset. Takes an arg string to "
  1114. "control the subplot which supports most (but not all) of the "
  1115. "parameters listed here. The relative dimensions of the subplot "
  1116. "can be controlled with -W/-H which now take a percentage.")
  1117. parser.add_argument(
  1118. '--subplot-below',
  1119. action=AppendSubplot,
  1120. help="Add subplot below with the same dataset.")
  1121. parser.add_argument(
  1122. '--subplot-left',
  1123. action=AppendSubplot,
  1124. help="Add subplot left with the same dataset.")
  1125. parser.add_argument(
  1126. '--subplot-right',
  1127. action=AppendSubplot,
  1128. help="Add subplot right with the same dataset.")
  1129. parser.add_argument(
  1130. '--subplot',
  1131. type=AppendSubplot.parse,
  1132. help="Add subplot-specific arguments to the main plot.")
  1133. def dictify(ns):
  1134. if hasattr(ns, 'subplots'):
  1135. ns.subplots = [(dir, dictify(subplot_ns))
  1136. for dir, subplot_ns in ns.subplots]
  1137. if ns.subplot is not None:
  1138. ns.subplot = dictify(ns.subplot)
  1139. return {k: v
  1140. for k, v in vars(ns).items()
  1141. if v is not None}
  1142. sys.exit(main(**dictify(parser.parse_intermixed_args())))