""" A module with functions for using LineCollection to create lines.ยดยด."""
import warnings
import numpy as np
from matplotlib import colormaps
from matplotlib import rcParams
from matplotlib.collections import LineCollection
from matplotlib.colors import to_rgba_array
from matplotlib.legend import Legend
from matplotlib.legend_handler import HandlerLineCollection
from mplsoccer.cm import create_transparent_cmap
from mplsoccer.utils import validate_ax
__all__ = ['lines']
[docs]
def lines(xstart, ystart, xend, yend, color=None, n_segments=100, comet=False, transparent=False,
alpha_start=0.01, alpha_end=1, cmap=None, ax=None, vertical=False,
reverse_cmap=False, **kwargs):
""" Plots lines using matplotlib.collections.LineCollection.
This is a fast way to plot multiple lines without loops.
Also enables lines that increase in width or opacity by splitting
the line into n_segments of increasing
width or opacity as the line progresses.
Parameters
----------
xstart, ystart, xend, yend: array-like or scalar.
Commonly, these parameters are 1D arrays.
These should be the start and end coordinates of the lines.
color : A matplotlib color or sequence of colors, defaults to None.
Defaults to None. In that case the marker color is determined
by the value rcParams['lines.color']
n_segments : int, default 100
If comet=True or transparent=True this is used to split the line
into n_segments of increasing width/opacity.
comet : bool default False
Whether to plot the lines increasing in width.
transparent : bool, default False
Whether to plot the lines increasing in opacity.
linewidth or lw : array-like or scalar, default 5.
Multiple linewidths not supported for the comet or transparent lines.
alpha_start: float, default 0.01
The starting alpha value for transparent lines, between 0 (transparent) and 1 (opaque).
If transparent = True the line will be drawn to
linearly increase in opacity between alpha_start and alpha_end.
alpha_end : float, default 1
The ending alpha value for transparent lines, between 0 (transparent) and 1 (opaque).
If transparent = True the line will be drawn to
linearly increase in opacity between alpha_start and alpha_end.
cmap : str, default None
A matplotlib cmap (colormap) name
vertical : bool, default False
If the orientation is vertical (True), then the code switches the x and y coordinates.
reverse_cmap : bool, default False
Whether to reverse the cmap colors.
If the pitch is horizontal and the y-axis is inverted then set this to True.
ax : matplotlib.axes.Axes, default None
The axis to plot on.
**kwargs : All other keyword arguments are passed on to matplotlib.collections.LineCollection.
Returns
-------
LineCollection : matplotlib.collections.LineCollection
Examples
--------
>>> from mplsoccer import Pitch
>>> pitch = Pitch()
>>> fig, ax = pitch.draw()
>>> pitch.lines(20, 20, 45, 80, comet=True, transparent=True, ax=ax)
>>> from mplsoccer.linecollection import lines
>>> import matplotlib.pyplot as plt
>>> fig, ax = plt.subplots()
>>> lines([0.1, 0.4], [0.1, 0.5], [0.9, 0.4], [0.8, 0.8], ax=ax)
"""
validate_ax(ax)
if not isinstance(comet, bool):
raise TypeError("Invalid argument: comet should be bool (True or False).")
if not isinstance(transparent, bool):
raise TypeError("Invalid argument: transparent should be bool (True or False).")
if alpha_start < 0 or alpha_start > 1:
raise TypeError("alpha_start values should be within 0-1 range")
if alpha_end < 0 or alpha_end > 1:
raise TypeError("alpha_end values should be within 0-1 range")
if alpha_start > alpha_end:
msg = "Alpha start > alpha end. The line will increase in transparency nearer to the end"
warnings.warn(msg)
if 'colors' in kwargs:
warnings.warn("lines method takes 'color' as an argument, 'colors' in ignored")
if color is not None and cmap is not None:
raise ValueError("Only use one of color or cmap arguments not both.")
if 'lw' in kwargs and 'linewidth' in kwargs:
raise TypeError("lines got multiple values for 'linewidth' argument (linewidth and lw).")
if 'lw' in kwargs:
lw = kwargs.pop('lw', 5)
elif 'linewidth' in kwargs:
lw = kwargs.pop('linewidth', 5)
else:
lw = 5
xstart = np.ravel(xstart)
ystart = np.ravel(ystart)
xend = np.ravel(xend)
yend = np.ravel(yend)
lw = np.ravel(lw)
if (comet or transparent) and lw.size > 1:
msg = "Multiple linewidths with a comet or transparent line is not implemented."
raise NotImplementedError(msg)
if color is None and cmap is None:
color = rcParams['lines.color']
if (comet or transparent) and cmap is None and to_rgba_array(color).shape[0] > 1:
msg = "Multiple colors with a comet or transparent line is not implemented."
raise NotImplementedError(msg)
if xstart.size != ystart.size:
raise ValueError("xstart and ystart must be the same size")
if xstart.size != xend.size:
raise ValueError("xstart and xend must be the same size")
if ystart.size != yend.size:
raise ValueError("ystart and yend must be the same size")
if lw.size > 1 and lw.size != xstart.size:
raise ValueError("lw and xstart must be the same size")
if lw.size == 1:
lw = lw[0]
if vertical:
ystart, xstart = xstart, ystart
yend, xend = xend, yend
if comet:
lw = np.linspace(1, lw, n_segments)
handler_first_lw = False
else:
handler_first_lw = True
multi_segment = transparent is not False or comet is not False or cmap is not None
if transparent:
cmap = create_transparent_cmap(color, cmap, n_segments, alpha_start, alpha_end)
if isinstance(cmap, str):
cmap = colormaps.get_cmap(cmap)
if cmap is not None:
handler_cmap = True
line_collection = _lines_cmap(xstart, ystart, xend, yend, lw=lw, cmap=cmap, ax=ax,
n_segments=n_segments, multi_segment=multi_segment,
reverse_cmap=reverse_cmap, **kwargs)
else:
handler_cmap = False
line_collection = _lines_no_cmap(xstart, ystart, xend, yend, lw=lw, color=color,
ax=ax, n_segments=n_segments,
multi_segment=multi_segment, **kwargs)
line_collection_handler = HandlerLines(numpoints=n_segments, invert_y=reverse_cmap,
first_lw=handler_first_lw, use_cmap=handler_cmap)
Legend.update_default_handler_map({LineCollection: line_collection_handler})
return line_collection
def _create_segments(xstart, ystart, xend, yend, n_segments=100, multi_segment=False):
if multi_segment:
x = np.linspace(xstart, xend, n_segments + 1)
y = np.linspace(ystart, yend, n_segments + 1)
points = np.array([x, y]).T
points = np.concatenate([points, np.expand_dims(points[:, -1, :], 1)], axis=1)
points = np.expand_dims(points, 1)
segments = np.concatenate([points[:, :, :-2, :],
points[:, :, 1:-1, :],
points[:, :, 2:, :]], axis=1)
segments = np.transpose(segments, (0, 2, 1, 3)).reshape((-1, 3, 2))
else:
segments = np.transpose(np.array([[xstart, ystart], [xend, yend]]), (2, 0, 1))
return segments
def _lines_no_cmap(xstart, ystart, xend, yend, lw=None, color=None, ax=None,
n_segments=100, multi_segment=False, **kwargs):
segments = _create_segments(xstart, ystart, xend, yend,
n_segments=n_segments, multi_segment=multi_segment)
color = to_rgba_array(color)
if (color.shape[0] > 1) and (color.shape[0] != xstart.size):
raise ValueError("xstart and color must be the same size")
line_collection = LineCollection(segments, color=color, linewidth=lw, snap=False, **kwargs)
line_collection = ax.add_collection(line_collection)
return line_collection
def _lines_cmap(xstart, ystart, xend, yend, lw=None, cmap=None, ax=None,
n_segments=100, multi_segment=False, reverse_cmap=False, **kwargs):
segments = _create_segments(xstart, ystart, xend, yend,
n_segments=n_segments, multi_segment=multi_segment)
if reverse_cmap:
cmap = cmap.reversed()
line_collection = LineCollection(segments, cmap=cmap, linewidth=lw, snap=False, **kwargs)
line_collection = ax.add_collection(line_collection)
extent = ax.get_ylim()
pitch_array = np.linspace(extent[0], extent[1], n_segments)
line_collection.set_array(pitch_array)
return line_collection
# Amended from
# https://stackoverflow.com/questions/49223702/adding-a-legend-to-a-matplotlib-plot-with-a-multicolored-line?rq=1
class HandlerLines(HandlerLineCollection):
"""Automatically generated by Pitch.lines() to allow use of linecollection in legend.
"""
def __init__(self, invert_y=False, first_lw=False, use_cmap=False,
marker_pad=0.3, numpoints=None, **kw):
HandlerLineCollection.__init__(self, marker_pad=marker_pad, numpoints=numpoints, **kw)
self.invert_y = invert_y
self.first_lw = first_lw
self.use_cmap = use_cmap
def create_artists(self, legend, artist, xdescent, ydescent,
width, height, fontsize, trans):
x = np.linspace(0, width, self.get_numpoints(legend) + 1)
y = np.zeros(self.get_numpoints(legend) + 1) + height / 2. - ydescent
points = np.array([x, y]).T.reshape(-1, 1, 2)
segments = np.concatenate([points[:-1], points[1:]], axis=1)
lw = artist.get_linewidth()
if self.first_lw:
lw = lw[0]
if self.use_cmap:
cmap = artist.cmap
if self.invert_y:
cmap = cmap.reversed()
line_collection = LineCollection(segments, lw=lw, cmap=cmap,
snap=False, transform=trans)
line_collection.set_array(x)
else:
line_collection = LineCollection(segments, lw=lw, colors=artist.get_colors()[0],
snap=False, transform=trans)
return [line_collection]