Clean colors.py

This commit is contained in:
Jensun Ravichandran 2021-06-11 23:07:55 +02:00
parent 56d554ed83
commit 92b8d1785c

View File

@ -1,78 +1,15 @@
"""ProtoFlow color utilities."""
import matplotlib.lines as mlines
from matplotlib import cm
from matplotlib.colors import Normalize, to_hex, to_rgb
def hex_to_rgb(hex_values):
for v in hex_values:
v = v.lstrip('#')
lv = len(v)
c = [int(v[i:i + lv // 3], 16) for i in range(0, lv, lv // 3)]
yield c
def color_scheme(n,
cmap="viridis",
form="hex",
tikz=False,
zero_indexed=False):
"""Return *n* colors from the color scheme.
Arguments:
n (int): number of colors to return
Keyword Arguments:
cmap (str): Name of a matplotlib `colormap\
<https://matplotlib.org/3.1.1/gallery/color/colormap_reference.html>`_.
form (str): Colorformat (supports "hex" and "rgb").
tikz (bool): Output as `TikZ <https://github.com/pgf-tikz/pgf>`_
command.
zero_indexed (bool): Use zero indexing for output array.
Returns:
(list): List of colors
"""
cmap = cm.get_cmap(cmap)
colornorm = Normalize(vmin=1, vmax=n)
hex_map = dict()
rgb_map = dict()
for cl in range(1, n + 1):
if zero_indexed:
hex_map[cl - 1] = to_hex(cmap(colornorm(cl)))
rgb_map[cl - 1] = to_rgb(cmap(colornorm(cl)))
else:
hex_map[cl] = to_hex(cmap(colornorm(cl)))
rgb_map[cl] = to_rgb(cmap(colornorm(cl)))
if tikz:
for k, v in rgb_map.items():
print(f"\\definecolor{{color-{k}}}{{rgb}}{{{v[0]},{v[1]},{v[2]}}}")
if form == "hex":
return hex_map
elif form == "rgb":
return rgb_map
else:
return hex_map
def get_legend_handles(labels, marker="dots", zero_indexed=False):
"""Return matplotlib legend handles and colors."""
handles = list()
n = len(labels)
colors = color_scheme(n,
cmap="viridis",
form="hex",
zero_indexed=zero_indexed)
for label, color in zip(labels, colors.values()):
if marker == "dots":
handle = mlines.Line2D(
[],
[],
color="white",
markerfacecolor=color,
marker="o",
markersize=10,
markeredgecolor="k",
label=label,
)
else:
handle = mlines.Line2D([], [],
color=color,
marker="",
markersize=15,
label=label)
handles.append(handle)
return handles, colors
def rgb_to_hex(rgb_values):
for v in rgb_values:
c = "%02x%02x%02x" % tuple(v)
yield c