Skip to content

Commit

Permalink
add heuristics for determining colorbar extend and nice ticks for dis…
Browse files Browse the repository at this point in the history 
…crete colormaps/colorbars
  • Loading branch information
Joe Hamman committed Aug 4, 2015
1 parent 1a01795 commit 119c1c8
Show file tree
Hide file tree
Showing 2 changed files with 81 additions and 31 deletions.
68 changes: 51 additions & 17 deletions xray/plot/plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -218,10 +218,13 @@ def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
import matplotlib as mpl

calc_data = plot_data[~pd.isnull(plot_data)]
bounds_set = True
if vmin is None:
vmin = np.percentile(calc_data, 2) if robust else calc_data.min()
bounds_set = False
if vmax is None:
vmax = np.percentile(calc_data, 98) if robust else calc_data.max()
bounds_set = False

# Simple heuristics for whether these data should have a divergent map
divergent = ((vmin < 0) and (vmax > 0)) or center is not None
Expand Down Expand Up @@ -250,6 +253,7 @@ def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
cmap = _load_default_cmap()

if extend is None:
extend_set = False
extend_min = calc_data.min() < vmin
extend_max = calc_data.max() > vmax
if extend_min and extend_max:
Expand All @@ -260,39 +264,67 @@ def _determine_cmap_params(plot_data, vmin=None, vmax=None, cmap=None,
extend = 'max'
else:
extend = 'neither'
else:
extend_set = True

if levels is not None:
cmap, cnorm = _determine_discrete_cmap_params(cmap, levels,
vmin, vmax,
extend)
cmap, cnorm, extend = _determine_discrete_cmap_params(cmap, levels,
vmin, vmax,
extend,
bounds_set,
extend_set)
else:
cnorm = mpl.colors.Normalize(vmin=vmin, vmax=vmax)

return vmin, vmax, cmap, extend, cnorm


def _determine_discrete_cmap_params(cmap, levels, vmin, vmax, extend):
def _determine_discrete_cmap_params(cmap, levels, vmin, vmax, extend,
bounds_set, extend_set):
"""
Build a discrete colormap and normalization of the data.
"""
import matplotlib as mpl
import matplotlib.pyplot as plt

if extend == 'both':
ext_n = 2
elif extend in ['min', 'max']:
ext_n = 1
else:
ext_n = 0
def extension_colors(extend):
if extend == 'both':
ext_n = 2
elif extend in ['min', 'max']:
ext_n = 1
else:
ext_n = 0
return ext_n

if isinstance(levels, int):
vmax += 10 * np.finfo(float).eps # Add small epison to include vmax
cticks = np.linspace(vmin, vmax, num=levels + 1, endpoint=True)
n_colors = levels + ext_n
if not bounds_set:
# if there were not user provided bounds, use MaxNLocator to pick
# a nice set of ticks
ticker = mpl.ticker.MaxNLocator(levels)
cticks = ticker.tick_values(vmin, vmax)
print(cticks)
else:
# otherwise, use the user provided vmin/vmax
cticks = np.linspace(vmin, vmax, num=levels + 1, endpoint=True)
ext_n = extension_colors(extend)
n_colors = len(cticks) + ext_n - 1
else:
try:
n_colors = len(levels) + ext_n - 1
cticks = np.asarray(levels)
if not extend_set:
extend_min = cticks[0] > vmin
extend_max = cticks[-1] < vmax
if extend_min and extend_max:
extend = 'both'
elif extend_min:
extend = 'min'
elif extend_max:
extend = 'max'
else:
extend = 'neither'
ext_n = extension_colors(extend)
n_colors = len(levels) + ext_n - 1
except TypeError as e:
print('Unexpected type (%s) given for levels' % type(levels))
raise e
Expand All @@ -311,7 +343,7 @@ def _determine_discrete_cmap_params(cmap, levels, vmin, vmax, extend):

cmap, cnorm = mpl.colors.from_levels_and_colors(cticks, pal, extend=extend)

return cmap, cnorm
return cmap, cnorm, extend


# MUST run before any 2d plotting functions are defined since
Expand Down Expand Up @@ -425,10 +457,12 @@ def newplotfunc(darray, ax=None, xincrease=None, yincrease=None,
# passing it to the colorbar is sufficient for imshow and
# pcolormesh
kwargs['extend'] = extend
kwargs['levels'] = levels

ax, primitive = plotfunc(x, y, z, ax=ax, cmap=cmap, norm=cnorm,
**kwargs)
if 'norm' not in kwargs:
# This allows the user to pass in a custom norm coming via kwargs
kwargs['norm'] = cnorm

ax, primitive = plotfunc(x, y, z, ax=ax, cmap=cmap, **kwargs)

ax.set_xlabel(xlab)
ax.set_ylabel(ylab)
Expand Down
44 changes: 30 additions & 14 deletions xray/test/test_plot.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -188,48 +188,57 @@ def test_integer_levels(self):
levels = 8
vmin = -5
vmax = 5
cmap, cnorm = _determine_discrete_cmap_params('Spectral', levels,
vmin, vmax, 'neither')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Spectral', levels, vmin, vmax, 'neither', True, True)
self.assertEqual(cmap.N, levels)
self.assertEqual(cnorm.N, levels + 1)
self.assertEqual(cnorm.vmin, vmin)
self.assertEqual(cnorm.vmax, vmax + 10 * np.finfo(float).eps)

cmap, cnorm = _determine_discrete_cmap_params('Blues', levels,
vmin, vmax, 'both')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Blues', levels, vmin, vmax, 'both', True, True)
# extension colors are not included here
self.assertEqual(cmap.N, levels)
self.assertEqual(cnorm.N, levels + 1)
self.assertEqual(cnorm.vmin, vmin)
self.assertEqual(cnorm.vmax, vmax + 10 * np.finfo(float).eps)

# heuristics for picking nice ticks
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Spectral', levels, vmin, vmax, 'neither', False, True)
self.assertGreaterEqual(cnorm.vmax, vmax)
self.assertLessEqual(cnorm.vmin, vmin)

def test_list_levels(self):
levels = [-4, -2, 0, 2, 4]
vmin = -5
vmax = 5

cmap, cnorm = _determine_discrete_cmap_params('Spectral', levels,
vmin, vmax, 'neither')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Spectral', levels, vmin, vmax, 'neither', True, True)
self.assertEqual(cmap.N, len(levels) - 1)
self.assertEqual(cnorm.N, len(levels))
self.assertEqual(cnorm.vmin, min(levels))
self.assertEqual(cnorm.vmax, max(levels))

cmap, cnorm = _determine_discrete_cmap_params('Greens_r', levels,
vmin, vmax, 'both')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Greens_r', levels, vmin, vmax, 'both', True, True)
self.assertEqual(cmap.N, len(levels) - 1)
self.assertEqual(cnorm.N, len(levels))
self.assertEqual(cnorm.vmin, min(levels))
self.assertEqual(cnorm.vmax, max(levels))

# levels as an array
cmap, cnorm = _determine_discrete_cmap_params('Greens_r',
np.array(levels),
vmin, vmax, 'both')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Greens_r', np.array(levels), vmin, vmax, 'both', True, True)
# levels as a DataArray
cmap, cnorm = _determine_discrete_cmap_params('Greens_r',
DataArray(levels),
vmin, vmax, 'both')
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Greens_r', DataArray(levels), vmin, vmax, 'both', True, True)

# heuristics for picking extend when using list of levels
cmap, cnorm, extend = _determine_discrete_cmap_params(
'Greens_r', DataArray(levels), -5, 3, 'both', True, False)
self.assertEqual(extend, 'min')


class Common2dMixin:
Expand Down Expand Up @@ -337,6 +346,13 @@ def test_extend(self):
artist = self.plotmethod(vmin=-10, vmax=0)
self.assertEqual(artist.extend, 'max')

def test_levels(self):
artist = self.plotmethod(levels=[-0.5, -0.4, 0.1])
self.assertEqual(artist.extend, 'both')

artist = self.plotmethod(levels=3)
self.assertEqual(artist.extend, 'neither')


class TestContour(Common2dMixin, PlotTestCase):

Expand Down

0 comments on commit 119c1c8

Please sign in to comment.