_tools.py 13 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382
  1. # being a bit too dynamic
  2. # pylint: disable=E1101
  3. from __future__ import division
  4. from math import ceil
  5. import warnings
  6. import numpy as np
  7. from pandas.compat import range
  8. from pandas.core.dtypes.common import is_list_like
  9. from pandas.core.dtypes.generic import ABCDataFrame, ABCIndexClass, ABCSeries
  10. def format_date_labels(ax, rot):
  11. # mini version of autofmt_xdate
  12. try:
  13. for label in ax.get_xticklabels():
  14. label.set_ha('right')
  15. label.set_rotation(rot)
  16. fig = ax.get_figure()
  17. fig.subplots_adjust(bottom=0.2)
  18. except Exception: # pragma: no cover
  19. pass
  20. def table(ax, data, rowLabels=None, colLabels=None, **kwargs):
  21. """
  22. Helper function to convert DataFrame and Series to matplotlib.table
  23. Parameters
  24. ----------
  25. ax : Matplotlib axes object
  26. data : DataFrame or Series
  27. data for table contents
  28. kwargs : keywords, optional
  29. keyword arguments which passed to matplotlib.table.table.
  30. If `rowLabels` or `colLabels` is not specified, data index or column
  31. name will be used.
  32. Returns
  33. -------
  34. matplotlib table object
  35. """
  36. if isinstance(data, ABCSeries):
  37. data = data.to_frame()
  38. elif isinstance(data, ABCDataFrame):
  39. pass
  40. else:
  41. raise ValueError('Input data must be DataFrame or Series')
  42. if rowLabels is None:
  43. rowLabels = data.index
  44. if colLabels is None:
  45. colLabels = data.columns
  46. cellText = data.values
  47. import matplotlib.table
  48. table = matplotlib.table.table(ax, cellText=cellText,
  49. rowLabels=rowLabels,
  50. colLabels=colLabels, **kwargs)
  51. return table
  52. def _get_layout(nplots, layout=None, layout_type='box'):
  53. if layout is not None:
  54. if not isinstance(layout, (tuple, list)) or len(layout) != 2:
  55. raise ValueError('Layout must be a tuple of (rows, columns)')
  56. nrows, ncols = layout
  57. # Python 2 compat
  58. ceil_ = lambda x: int(ceil(x))
  59. if nrows == -1 and ncols > 0:
  60. layout = nrows, ncols = (ceil_(float(nplots) / ncols), ncols)
  61. elif ncols == -1 and nrows > 0:
  62. layout = nrows, ncols = (nrows, ceil_(float(nplots) / nrows))
  63. elif ncols <= 0 and nrows <= 0:
  64. msg = "At least one dimension of layout must be positive"
  65. raise ValueError(msg)
  66. if nrows * ncols < nplots:
  67. raise ValueError('Layout of {nrows}x{ncols} must be larger '
  68. 'than required size {nplots}'.format(
  69. nrows=nrows, ncols=ncols, nplots=nplots))
  70. return layout
  71. if layout_type == 'single':
  72. return (1, 1)
  73. elif layout_type == 'horizontal':
  74. return (1, nplots)
  75. elif layout_type == 'vertical':
  76. return (nplots, 1)
  77. layouts = {1: (1, 1), 2: (1, 2), 3: (2, 2), 4: (2, 2)}
  78. try:
  79. return layouts[nplots]
  80. except KeyError:
  81. k = 1
  82. while k ** 2 < nplots:
  83. k += 1
  84. if (k - 1) * k >= nplots:
  85. return k, (k - 1)
  86. else:
  87. return k, k
  88. # copied from matplotlib/pyplot.py and modified for pandas.plotting
  89. def _subplots(naxes=None, sharex=False, sharey=False, squeeze=True,
  90. subplot_kw=None, ax=None, layout=None, layout_type='box',
  91. **fig_kw):
  92. """Create a figure with a set of subplots already made.
  93. This utility wrapper makes it convenient to create common layouts of
  94. subplots, including the enclosing figure object, in a single call.
  95. Keyword arguments:
  96. naxes : int
  97. Number of required axes. Exceeded axes are set invisible. Default is
  98. nrows * ncols.
  99. sharex : bool
  100. If True, the X axis will be shared amongst all subplots.
  101. sharey : bool
  102. If True, the Y axis will be shared amongst all subplots.
  103. squeeze : bool
  104. If True, extra dimensions are squeezed out from the returned axis object:
  105. - if only one subplot is constructed (nrows=ncols=1), the resulting
  106. single Axis object is returned as a scalar.
  107. - for Nx1 or 1xN subplots, the returned object is a 1-d numpy object
  108. array of Axis objects are returned as numpy 1-d arrays.
  109. - for NxM subplots with N>1 and M>1 are returned as a 2d array.
  110. If False, no squeezing is done: the returned axis object is always
  111. a 2-d array containing Axis instances, even if it ends up being 1x1.
  112. subplot_kw : dict
  113. Dict with keywords passed to the add_subplot() call used to create each
  114. subplots.
  115. ax : Matplotlib axis object, optional
  116. layout : tuple
  117. Number of rows and columns of the subplot grid.
  118. If not specified, calculated from naxes and layout_type
  119. layout_type : {'box', 'horziontal', 'vertical'}, default 'box'
  120. Specify how to layout the subplot grid.
  121. fig_kw : Other keyword arguments to be passed to the figure() call.
  122. Note that all keywords not recognized above will be
  123. automatically included here.
  124. Returns:
  125. fig, ax : tuple
  126. - fig is the Matplotlib Figure object
  127. - ax can be either a single axis object or an array of axis objects if
  128. more than one subplot was created. The dimensions of the resulting array
  129. can be controlled with the squeeze keyword, see above.
  130. **Examples:**
  131. x = np.linspace(0, 2*np.pi, 400)
  132. y = np.sin(x**2)
  133. # Just a figure and one subplot
  134. f, ax = plt.subplots()
  135. ax.plot(x, y)
  136. ax.set_title('Simple plot')
  137. # Two subplots, unpack the output array immediately
  138. f, (ax1, ax2) = plt.subplots(1, 2, sharey=True)
  139. ax1.plot(x, y)
  140. ax1.set_title('Sharing Y axis')
  141. ax2.scatter(x, y)
  142. # Four polar axes
  143. plt.subplots(2, 2, subplot_kw=dict(polar=True))
  144. """
  145. import matplotlib.pyplot as plt
  146. if subplot_kw is None:
  147. subplot_kw = {}
  148. if ax is None:
  149. fig = plt.figure(**fig_kw)
  150. else:
  151. if is_list_like(ax):
  152. ax = _flatten(ax)
  153. if layout is not None:
  154. warnings.warn("When passing multiple axes, layout keyword is "
  155. "ignored", UserWarning)
  156. if sharex or sharey:
  157. warnings.warn("When passing multiple axes, sharex and sharey "
  158. "are ignored. These settings must be specified "
  159. "when creating axes", UserWarning,
  160. stacklevel=4)
  161. if len(ax) == naxes:
  162. fig = ax[0].get_figure()
  163. return fig, ax
  164. else:
  165. raise ValueError("The number of passed axes must be {0}, the "
  166. "same as the output plot".format(naxes))
  167. fig = ax.get_figure()
  168. # if ax is passed and a number of subplots is 1, return ax as it is
  169. if naxes == 1:
  170. if squeeze:
  171. return fig, ax
  172. else:
  173. return fig, _flatten(ax)
  174. else:
  175. warnings.warn("To output multiple subplots, the figure containing "
  176. "the passed axes is being cleared", UserWarning,
  177. stacklevel=4)
  178. fig.clear()
  179. nrows, ncols = _get_layout(naxes, layout=layout, layout_type=layout_type)
  180. nplots = nrows * ncols
  181. # Create empty object array to hold all axes. It's easiest to make it 1-d
  182. # so we can just append subplots upon creation, and then
  183. axarr = np.empty(nplots, dtype=object)
  184. # Create first subplot separately, so we can share it if requested
  185. ax0 = fig.add_subplot(nrows, ncols, 1, **subplot_kw)
  186. if sharex:
  187. subplot_kw['sharex'] = ax0
  188. if sharey:
  189. subplot_kw['sharey'] = ax0
  190. axarr[0] = ax0
  191. # Note off-by-one counting because add_subplot uses the MATLAB 1-based
  192. # convention.
  193. for i in range(1, nplots):
  194. kwds = subplot_kw.copy()
  195. # Set sharex and sharey to None for blank/dummy axes, these can
  196. # interfere with proper axis limits on the visible axes if
  197. # they share axes e.g. issue #7528
  198. if i >= naxes:
  199. kwds['sharex'] = None
  200. kwds['sharey'] = None
  201. ax = fig.add_subplot(nrows, ncols, i + 1, **kwds)
  202. axarr[i] = ax
  203. if naxes != nplots:
  204. for ax in axarr[naxes:]:
  205. ax.set_visible(False)
  206. _handle_shared_axes(axarr, nplots, naxes, nrows, ncols, sharex, sharey)
  207. if squeeze:
  208. # Reshape the array to have the final desired dimension (nrow,ncol),
  209. # though discarding unneeded dimensions that equal 1. If we only have
  210. # one subplot, just return it instead of a 1-element array.
  211. if nplots == 1:
  212. axes = axarr[0]
  213. else:
  214. axes = axarr.reshape(nrows, ncols).squeeze()
  215. else:
  216. # returned axis array will be always 2-d, even if nrows=ncols=1
  217. axes = axarr.reshape(nrows, ncols)
  218. return fig, axes
  219. def _remove_labels_from_axis(axis):
  220. for t in axis.get_majorticklabels():
  221. t.set_visible(False)
  222. try:
  223. # set_visible will not be effective if
  224. # minor axis has NullLocator and NullFormattor (default)
  225. import matplotlib.ticker as ticker
  226. if isinstance(axis.get_minor_locator(), ticker.NullLocator):
  227. axis.set_minor_locator(ticker.AutoLocator())
  228. if isinstance(axis.get_minor_formatter(), ticker.NullFormatter):
  229. axis.set_minor_formatter(ticker.FormatStrFormatter(''))
  230. for t in axis.get_minorticklabels():
  231. t.set_visible(False)
  232. except Exception: # pragma no cover
  233. raise
  234. axis.get_label().set_visible(False)
  235. def _handle_shared_axes(axarr, nplots, naxes, nrows, ncols, sharex, sharey):
  236. if nplots > 1:
  237. if nrows > 1:
  238. try:
  239. # first find out the ax layout,
  240. # so that we can correctly handle 'gaps"
  241. layout = np.zeros((nrows + 1, ncols + 1), dtype=np.bool)
  242. for ax in axarr:
  243. layout[ax.rowNum, ax.colNum] = ax.get_visible()
  244. for ax in axarr:
  245. # only the last row of subplots should get x labels -> all
  246. # other off layout handles the case that the subplot is
  247. # the last in the column, because below is no subplot/gap.
  248. if not layout[ax.rowNum + 1, ax.colNum]:
  249. continue
  250. if sharex or len(ax.get_shared_x_axes()
  251. .get_siblings(ax)) > 1:
  252. _remove_labels_from_axis(ax.xaxis)
  253. except IndexError:
  254. # if gridspec is used, ax.rowNum and ax.colNum may different
  255. # from layout shape. in this case, use last_row logic
  256. for ax in axarr:
  257. if ax.is_last_row():
  258. continue
  259. if sharex or len(ax.get_shared_x_axes()
  260. .get_siblings(ax)) > 1:
  261. _remove_labels_from_axis(ax.xaxis)
  262. if ncols > 1:
  263. for ax in axarr:
  264. # only the first column should get y labels -> set all other to
  265. # off as we only have labels in the first column and we always
  266. # have a subplot there, we can skip the layout test
  267. if ax.is_first_col():
  268. continue
  269. if sharey or len(ax.get_shared_y_axes().get_siblings(ax)) > 1:
  270. _remove_labels_from_axis(ax.yaxis)
  271. def _flatten(axes):
  272. if not is_list_like(axes):
  273. return np.array([axes])
  274. elif isinstance(axes, (np.ndarray, ABCIndexClass)):
  275. return axes.ravel()
  276. return np.array(axes)
  277. def _get_all_lines(ax):
  278. lines = ax.get_lines()
  279. if hasattr(ax, 'right_ax'):
  280. lines += ax.right_ax.get_lines()
  281. if hasattr(ax, 'left_ax'):
  282. lines += ax.left_ax.get_lines()
  283. return lines
  284. def _get_xlim(lines):
  285. left, right = np.inf, -np.inf
  286. for l in lines:
  287. x = l.get_xdata(orig=False)
  288. left = min(np.nanmin(x), left)
  289. right = max(np.nanmax(x), right)
  290. return left, right
  291. def _set_ticks_props(axes, xlabelsize=None, xrot=None,
  292. ylabelsize=None, yrot=None):
  293. import matplotlib.pyplot as plt
  294. for ax in _flatten(axes):
  295. if xlabelsize is not None:
  296. plt.setp(ax.get_xticklabels(), fontsize=xlabelsize)
  297. if xrot is not None:
  298. plt.setp(ax.get_xticklabels(), rotation=xrot)
  299. if ylabelsize is not None:
  300. plt.setp(ax.get_yticklabels(), fontsize=ylabelsize)
  301. if yrot is not None:
  302. plt.setp(ax.get_yticklabels(), rotation=yrot)
  303. return axes