Skip to content

Commit 37a947a

Browse files
authored
Refactor plot utils (#2670)
* Refactor out utility functions. * facetgrid refactor 1. refactor out _easy_facetgrid 2. Combine map_dataarray_line with map_dataarray * flake8 * Refactor out colorbar making to plot.utils._add_colorbar * Refactor out cmap_params, cbar_kwargs processing * Back to map_dataarray_line * lint * small rename * review comment. * Move _infer_line_data back.
1 parent fd2552a commit 37a947a

File tree

3 files changed

+361
-368
lines changed

3 files changed

+361
-368
lines changed

xarray/plot/facetgrid.py

Lines changed: 43 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -7,8 +7,8 @@
77

88
from ..core.formatting import format_item
99
from .utils import (
10-
_determine_cmap_params, _infer_xy_labels, import_matplotlib_pyplot,
11-
label_from_attrs)
10+
_infer_xy_labels, _process_cmap_cbar_kwargs,
11+
import_matplotlib_pyplot, label_from_attrs)
1212

1313
# Overrides axes.labelsize, xtick.major.size, ytick.major.size
1414
# from mpl.rcParams
@@ -219,32 +219,13 @@ def map_dataarray(self, func, x, y, **kwargs):
219219
220220
"""
221221

222-
cmapkw = kwargs.get('cmap')
223-
colorskw = kwargs.get('colors')
224-
cbar_kwargs = kwargs.pop('cbar_kwargs', {})
225-
cbar_kwargs = {} if cbar_kwargs is None else dict(cbar_kwargs)
226-
227222
if kwargs.get('cbar_ax', None) is not None:
228223
raise ValueError('cbar_ax not supported by FacetGrid.')
229224

230-
# colors is mutually exclusive with cmap
231-
if cmapkw and colorskw:
232-
raise ValueError("Can't specify both cmap and colors.")
233-
234-
# These should be consistent with xarray.plot._plot2d
235-
cmap_kwargs = {'plot_data': self.data.values,
236-
# MPL default
237-
'levels': 7 if 'contour' in func.__name__ else None,
238-
'filled': func.__name__ != 'contour',
239-
}
240-
241-
cmap_args = getfullargspec(_determine_cmap_params).args
242-
cmap_kwargs.update((a, kwargs[a]) for a in cmap_args if a in kwargs)
225+
cmap_params, cbar_kwargs = _process_cmap_cbar_kwargs(
226+
func, kwargs, self.data.values)
243227

244-
cmap_params = _determine_cmap_params(**cmap_kwargs)
245-
246-
if colorskw is not None:
247-
cmap_params['cmap'] = None
228+
self._cmap_extend = cmap_params.get('extend')
248229

249230
# Order is important
250231
func_kwargs = kwargs.copy()
@@ -260,7 +241,7 @@ def map_dataarray(self, func, x, y, **kwargs):
260241
# None is the sentinel value
261242
if d is not None:
262243
subset = self.data.loc[d]
263-
mappable = func(subset, x, y, ax=ax, **func_kwargs)
244+
mappable = func(subset, x=x, y=y, ax=ax, **func_kwargs)
264245
self._mappables.append(mappable)
265246

266247
self._cmap_extend = cmap_params.get('extend')
@@ -271,36 +252,24 @@ def map_dataarray(self, func, x, y, **kwargs):
271252

272253
return self
273254

274-
def map_dataarray_line(self, x=None, y=None, hue=None, **kwargs):
275-
"""
276-
Apply a line plot to a 2d facet subset of the data.
277-
278-
Parameters
279-
----------
280-
x, y, hue: string
281-
dimension names for the axes and hues of each facet
282-
283-
Returns
284-
-------
285-
self : FacetGrid object
286-
287-
"""
288-
from .plot import line, _infer_line_data
255+
def map_dataarray_line(self, func, x, y, **kwargs):
256+
from .plot import _infer_line_data
289257

290258
add_legend = kwargs.pop('add_legend', True)
291259
kwargs['add_legend'] = False
260+
func_kwargs = kwargs.copy()
261+
func_kwargs['_labels'] = False
292262

293263
for d, ax in zip(self.name_dicts.flat, self.axes.flat):
294264
# None is the sentinel value
295265
if d is not None:
296266
subset = self.data.loc[d]
297-
mappable = line(subset, x=x, y=y, hue=hue,
298-
ax=ax, _labels=False,
299-
**kwargs)
267+
mappable = func(subset, x=x, y=y, ax=ax, **func_kwargs)
300268
self._mappables.append(mappable)
269+
301270
_, _, hueplt, xlabel, ylabel, huelabel = _infer_line_data(
302271
darray=self.data.loc[self.name_dicts.flat[0]],
303-
x=x, y=y, hue=hue)
272+
x=x, y=y, hue=func_kwargs['hue'])
304273

305274
self._hue_var = hueplt
306275
self._hue_label = huelabel
@@ -520,3 +489,33 @@ def map(self, func, *args, **kwargs):
520489
self._finalize_grid(*args[:2])
521490

522491
return self
492+
493+
494+
def _easy_facetgrid(data, plotfunc, kind, x=None, y=None, row=None,
495+
col=None, col_wrap=None, sharex=True, sharey=True,
496+
aspect=None, size=None, subplot_kws=None, **kwargs):
497+
"""
498+
Convenience method to call xarray.plot.FacetGrid from 2d plotting methods
499+
500+
kwargs are the arguments to 2d plotting method
501+
"""
502+
ax = kwargs.pop('ax', None)
503+
figsize = kwargs.pop('figsize', None)
504+
if ax is not None:
505+
raise ValueError("Can't use axes when making faceted plots.")
506+
if aspect is None:
507+
aspect = 1
508+
if size is None:
509+
size = 3
510+
elif figsize is not None:
511+
raise ValueError('cannot provide both `figsize` and `size` arguments')
512+
513+
g = FacetGrid(data=data, col=col, row=row, col_wrap=col_wrap,
514+
sharex=sharex, sharey=sharey, figsize=figsize,
515+
aspect=aspect, size=size, subplot_kws=subplot_kws)
516+
517+
if kind == 'line':
518+
return g.map_dataarray_line(plotfunc, x, y, **kwargs)
519+
520+
if kind == 'dataarray':
521+
return g.map_dataarray(plotfunc, x, y, **kwargs)

0 commit comments

Comments
 (0)