from __future__ import annotations
import hashlib
import typing
from warnings import warn
import numpy as np
import pandas as pd
from mizani.bounds import rescale
from ..exceptions import PlotnineWarning
from ..mapping.aes import rename_aesthetics
from ..scales.scale_continuous import scale_continuous
from .guide import guide
if typing.TYPE_CHECKING:
from plotnine.typing import ScaleContinuous
[docs]class guide_colorbar(guide):
"""
Guide colorbar
Parameters
----------
barwidth : float
Width (in pixels) of the colorbar.
barheight : float
Height (in pixels) of the colorbar. The height is multiplied by
a factor of 5.
nbin : int
Number of bins for drawing a colorbar. A larger value yields
a smoother colorbar. Default is 20.
raster : bool
Whether to render the colorbar as a raster object.
ticks : bool
Whether tick marks on colorbar should be visible.
draw_ulim : bool
Whether to show the upper limit tick marks.
draw_llim : bool
Whether to show the lower limit tick marks.
direction : str in ``['horizontal', 'vertical']``
Direction of the guide.
kwargs : dict
Parameters passed on to :class:`.guide`
"""
# bar
barwidth = None
barheight = None
nbin = 20 # maximum number of bins
raster = True
# ticks
ticks = True
draw_ulim = True
draw_llim = True
# parameter
available_aes = {"colour", "color", "fill"}
def train(self, scale: ScaleContinuous, aesthetic=None):
if aesthetic is None:
aesthetic = scale.aesthetics[0]
# Do nothing if scales are inappropriate
if set(scale.aesthetics) & self.available_aes == 0:
warn("colorbar guide needs appropriate scales.", PlotnineWarning)
return None
if not isinstance(scale, scale_continuous):
warn("colorbar guide needs continuous scales", PlotnineWarning)
return None
# value = breaks (numeric) is used for determining the
# position of ticks
limits = scale.limits
breaks = scale.get_bounded_breaks()
if not len(breaks):
return None
self.key = pd.DataFrame(
{
aesthetic: scale.map(breaks),
"label": scale.get_labels(breaks),
"value": breaks,
}
)
bar = np.linspace(limits[0], limits[1], self.nbin)
self.bar = pd.DataFrame({"color": scale.map(bar), "value": bar})
labels = " ".join(str(x) for x in self.key["label"])
info = "\n".join(
[
self.title,
labels,
" ".join(self.bar["color"].tolist()),
self.__class__.__name__,
]
)
self.hash = hashlib.md5(info.encode("utf-8")).hexdigest()
return self
def merge(self, other):
"""
Simply discards the other guide
"""
return self
def create_geoms(self, plot):
"""
Return self if colorbar will be drawn and None if not
This guide is not geom based
"""
for l in plot.layers:
exclude = set()
if isinstance(l.show_legend, dict):
l.show_legend = rename_aesthetics(l.show_legend)
exclude = {ae for ae, val in l.show_legend.items() if not val}
elif l.show_legend not in (None, True):
continue
matched = self.legend_aesthetics(l, plot)
# layer uses guide
if set(matched) - exclude:
break
# no break, no layer uses this guide
else:
return None
return self
def draw(self):
"""
Draw guide
Returns
-------
out : matplotlib.offsetbox.Offsetbox
A drawing of this legend
"""
from matplotlib.offsetbox import (
DrawingArea,
HPacker,
TextArea,
VPacker,
)
obverse = slice(0, None)
reverse = slice(None, None, -1)
nbars = len(self.bar)
direction = self.direction
colors = self.bar["color"].tolist()
labels = self.key["label"].tolist()
_targets = self.theme._targets
_property = self.theme.themeables.property
# 1.45 makes the default colourbar wider than the
# legend entry boxes.
width = (self.barwidth or _property("legend_key_width")) * 1.45
height = (self.barheight or _property("legend_key_height")) * 1.45
height *= 5
length = height
# When there is more than one guide, we keep
# record of all of them using lists
if "legend_title" not in _targets:
_targets["legend_title"] = []
if "legend_text_colorbar" not in _targets:
_targets["legend_text_colorbar"] = []
# .5 puts the ticks in the middle of the bars when
# raster=False. So when raster=True the ticks are
# in between interpolation points and the matching is
# close though not exactly right.
_from = self.bar["value"].min(), self.bar["value"].max()
tick_locations = (
rescale(self.key["value"], (0.5, nbars - 0.5), _from)
* length
/ nbars
)
if direction == "horizontal":
width, height = height, width
length = width
if self.reverse:
colors = colors[::-1]
labels = labels[::-1]
tick_locations = length - tick_locations[::-1]
# title #
title_box = TextArea(self.title, textprops={"color": "black"})
_targets["legend_title"].append(title_box)
# colorbar and ticks #
da = DrawingArea(width, height, 0, 0)
if self.raster:
add_interpolated_colorbar(da, colors, direction)
else:
add_segmented_colorbar(da, colors, direction)
if self.ticks:
_locations = tick_locations
if not self.draw_ulim:
_locations = _locations[:-1]
if not self.draw_llim:
_locations = _locations[1:]
add_ticks(da, _locations, direction)
# labels #
if self.label:
labels_da, legend_text = create_labels(
da, labels, tick_locations, direction
)
_targets["legend_text_colorbar"].extend(legend_text)
else:
labels_da = DrawingArea(0, 0)
# colorbar + labels #
if direction == "vertical":
packer, align = HPacker, "bottom"
align = "center"
else:
packer, align = VPacker, "right"
align = "center"
slc = obverse if self.label_position == "right" else reverse
if self.label_position in ("right", "bottom"):
slc = obverse
else:
slc = reverse
main_box = packer(
children=[da, labels_da][slc],
sep=self._label_margin,
align=align,
pad=0,
)
# title + colorbar(with labels) #
lookup = {
"right": (HPacker, reverse),
"left": (HPacker, obverse),
"bottom": (VPacker, reverse),
"top": (VPacker, obverse),
}
packer, slc = lookup[self.title_position]
children = [title_box, main_box][slc]
box = packer(
children=children,
sep=self._title_margin,
align=self._title_align,
pad=0,
)
return box
def add_interpolated_colorbar(da, colors, direction):
"""
Add 'rastered' colorbar to DrawingArea
"""
from matplotlib.collections import QuadMesh
from matplotlib.colors import ListedColormap
# Special case that arises due to not so useful
# aesthetic mapping.
if len(colors) == 1:
colors = [colors[0], colors[0]]
# Number of horizontal egdes(breaks) in the grid
# No need to create more nbreak than colors, provided
# no. of colors = no. of breaks = no. of cmap colors
# the shading does a perfect interpolation
nbreak = len(colors)
if direction == "vertical":
mesh_width = 1
mesh_height = nbreak - 1
linewidth = da.height / mesh_height
# Construct rectangular meshgrid
# The values(Z) at each vertex are just the
# normalized (onto [0, 1]) vertical distance
x = np.array([0, da.width])
y = np.arange(0, nbreak) * linewidth
X, Y = np.meshgrid(x, y)
Z = Y / y.max()
else:
mesh_width = nbreak - 1
mesh_height = 1
linewidth = da.width / mesh_width
x = np.arange(0, nbreak) * linewidth
y = np.array([0, da.height])
X, Y = np.meshgrid(x, y)
Z = X / x.max()
# As a 3D (mesh_width x mesh_height x 2) coordinates array
coordinates = np.stack([X, Y], axis=-1)
cmap = ListedColormap(colors)
coll = QuadMesh(
coordinates,
antialiased=False,
shading="gouraud",
linewidth=0,
cmap=cmap,
array=Z.ravel(),
)
da.add_artist(coll)
def add_segmented_colorbar(da, colors, direction):
"""
Add 'non-rastered' colorbar to DrawingArea
"""
from matplotlib.collections import PolyCollection
nbreak = len(colors)
if direction == "vertical":
linewidth = da.height / nbreak
verts = []
x1, x2 = 0, da.width
for i in range(nbreak):
y1 = i * linewidth
y2 = y1 + linewidth
verts.append(((x1, y1), (x1, y2), (x2, y2), (x2, y1)))
else:
linewidth = da.width / nbreak
verts = []
y1, y2 = 0, da.height
for i in range(nbreak):
x1 = i * linewidth
x2 = x1 + linewidth
verts.append(((x1, y1), (x1, y2), (x2, y2), (x2, y1)))
coll = PolyCollection(
verts, facecolors=colors, linewidth=0, antialiased=False
)
da.add_artist(coll)
def add_ticks(da, locations, direction):
from matplotlib.collections import LineCollection
segments = []
if direction == "vertical":
x1, x2, x3, x4 = np.array([0.0, 1 / 5, 4 / 5, 1.0]) * da.width
for y in locations:
segments.extend(
[
((x1, y), (x2, y)),
((x3, y), (x4, y)),
]
)
else:
y1, y2, y3, y4 = np.array([0.0, 1 / 5, 4 / 5, 1.0]) * da.height
for x in locations:
segments.extend(
[
((x, y1), (x, y2)),
((x, y3), (x, y4)),
]
)
coll = LineCollection(
segments, color="#CCCCCC", linewidth=1, antialiased=False
)
da.add_artist(coll)
def create_labels(da, labels, locations, direction):
"""
Return an OffsetBox with label texts
"""
# The box dimensions are determined by the size of
# the text objects. We put two dummy children at
# either end to gaurantee that when center packed
# the labels in the labels_box matchup with the ticks.
from matplotlib.text import Text
from .._mpl.offsetbox import MyAuxTransformBox
fontsize = 9
labels_box = MyAuxTransformBox()
xs, ys = [0] * len(labels), locations
ha, va = "left", "center"
x1, y1 = 0, 0
x2, y2 = 0, da.height
if direction == "horizontal":
xs, ys = ys, xs
ha, va = "center", "top"
x2, y2 = da.width, 0
txt1 = Text(x1, y1, "", horizontalalignment=ha, verticalalignment=va)
txt2 = Text(x2, y2, "", horizontalalignment=ha, verticalalignment=va)
labels_box.add_artist(txt1)
labels_box.add_artist(txt2)
legend_text = []
for i, (x, y, text) in enumerate(zip(xs, ys, labels)):
txt = Text(
x,
y,
text,
size=fontsize,
horizontalalignment=ha,
verticalalignment=va,
)
labels_box.add_artist(txt)
legend_text.append(txt)
return labels_box, legend_text
guide_colourbar = guide_colorbar