from __future__ import absolute_import, division, unicode_literals
import numpy as np
from matplotlib.collections import LineCollection, PatchCollection
from matplotlib.patches import Rectangle
from ...core.options import abbreviated_exception
from ..mixins import GeomMixin
from .element import ColorbarPlot
from .path import PathPlot, PolygonPlot
[docs]class SegmentPlot(GeomMixin, ColorbarPlot):
"""
Segments are lines in 2D space where each two key dimensions specify a
(x, y) node of the line.
"""
style_opts = PathPlot.style_opts + ['cmap']
_nonvectorized_styles = ['cmap']
_plot_methods = dict(single='segment')
[docs] def init_artists(self, ax, plot_args, plot_kwargs):
if 'c' in plot_kwargs:
plot_kwargs['array'] = plot_kwargs.pop('c')
if 'vmin' in plot_kwargs and 'vmax' in plot_kwargs:
plot_kwargs['clim'] = plot_kwargs.pop('vmin'), plot_kwargs.pop('vmax')
line_segments = LineCollection(*plot_args, **plot_kwargs)
ax.add_collection(line_segments)
return {'artist': line_segments}
def get_data(self, element, ranges, style):
inds = (1, 0, 3, 2) if self.invert_axes else (0, 1, 2, 3)
dims = element.dimensions()
data = [[(x0, y0), (x1, y1)] for x0, y0, x1, y1
in zip(*(element.dimension_values(d) for d in inds))]
with abbreviated_exception():
style = self._apply_transforms(element, ranges, style)
return (data,), style, {'dimensions': dims}
[docs]class RectanglesPlot(GeomMixin, ColorbarPlot):
"""
Rectanlges are polygons in 2D space where the key dimensions represent
the bottom-left and top-right corner of the rectangle.
"""
style_opts = PolygonPlot.style_opts
_nonvectorized_styles = ['cmap']
[docs] def init_artists(self, ax, plot_args, plot_kwargs):
if 'c' in plot_kwargs:
plot_kwargs['array'] = plot_kwargs.pop('c')
if 'vmin' in plot_kwargs and 'vmax' in plot_kwargs:
plot_kwargs['clim'] = plot_kwargs.pop('vmin'), plot_kwargs.pop('vmax')
line_segments = PatchCollection(*plot_args, **plot_kwargs)
ax.add_collection(line_segments)
return {'artist': line_segments}
def get_data(self, element, ranges, style):
# Get [x0, y0, x1, y1]
inds = (1, 0, 3, 2) if self.invert_axes else (0, 1, 2, 3)
x0s, y0s, x1s, y1s = (element.dimension_values(kd) for kd in inds)
x0s, x1s = np.min([x0s, x1s], axis=0), np.max([x0s, x1s], axis=0)
y0s, y1s = np.min([y0s, y1s], axis=0), np.max([y0s, y1s], axis=0)
dims = element.dimensions()
data = [Rectangle((x0, y0), x1, y1) for (x0, y0, x1, y1)
in zip(x0s, y0s, x1s-x0s, y1s-y0s)]
with abbreviated_exception():
style = self._apply_transforms(element, ranges, style)
return (data,), style, {'dimensions': dims}
