from __future__ import annotations
import typing
from collections import Counter
from contextlib import suppress
from warnings import warn
import numpy as np
from ..doctools import document
from ..exceptions import PlotnineWarning
from ..utils import SIZE_FACTOR, make_line_segments, match, to_rgba
from .geom import geom
if typing.TYPE_CHECKING:
from typing import Any, Literal, Sequence
import numpy.typing as npt
import pandas as pd
from matplotlib.path import Path
from plotnine.iapi import panel_view
from plotnine.typing import Axes, Coord, DrawingArea, Layer, TupleFloat2
[docs]@document
class geom_path(geom):
"""
Connected points
{usage}
Parameters
----------
{common_parameters}
lineend : str (default: butt)
Line end style, of of *butt*, *round* or *projecting.*
This option is applied for solid linetypes.
linejoin : str (default: round)
Line join style, one of *round*, *miter* or *bevel*.
This option is applied for solid linetypes.
arrow : plotnine.geoms.geom_path.arrow (default: None)
Arrow specification. Default is no arrow.
See Also
--------
plotnine.geoms.arrow : for adding arrowhead(s) to paths.
"""
DEFAULT_AES = {
"alpha": 1,
"color": "black",
"linetype": "solid",
"size": 0.5,
}
REQUIRED_AES = {"x", "y"}
DEFAULT_PARAMS = {
"stat": "identity",
"position": "identity",
"na_rm": False,
"lineend": "butt",
"linejoin": "round",
"arrow": None,
}
def handle_na(self, data: pd.DataFrame) -> pd.DataFrame:
def keep(x: Sequence[float]) -> npt.NDArray[np.bool_]:
# first non-missing to last non-missing
first = match([False], x, nomatch=1, start=0)[0]
last = len(x) - match([False], x[::-1], nomatch=1, start=0)[0]
bool_idx = np.hstack(
[
np.repeat(False, first),
np.repeat(True, last - first),
np.repeat(False, len(x) - last),
]
)
return bool_idx
# Get indices where any row for the select aesthetics has
# NaNs at the beginning or the end. Those we drop
bool_idx = (
data[["x", "y", "size", "color", "linetype"]]
.isna() # Missing
.apply(keep, axis=0)
) # Beginning or the End
bool_idx = np.all(bool_idx, axis=1) # Across the aesthetics
# return data
n1 = len(data)
data = data[bool_idx]
data.reset_index(drop=True, inplace=True)
n2 = len(data)
if n2 != n1 and not self.params["na_rm"]:
msg = "geom_path: Removed {} rows containing missing values."
warn(msg.format(n1 - n2), PlotnineWarning)
return data
def draw_panel(
self,
data: pd.DataFrame,
panel_params: panel_view,
coord: Coord,
ax: Axes,
**params: Any,
):
if not any(data["group"].duplicated()):
warn(
"geom_path: Each group consist of only one "
"observation. Do you need to adjust the "
"group aesthetic?",
PlotnineWarning,
)
# drop lines with less than two points
c = Counter(data["group"])
counts = np.array([c[v] for v in data["group"]])
data = data[counts >= 2]
if len(data) < 2:
return
# dataframe mergesort is stable, we rely on that here
data = data.sort_values("group", kind="mergesort")
data.reset_index(drop=True, inplace=True)
# When the parameters of the path are not constant
# with in the group, then the lines that make the paths
# can be drawn as separate segments
cols = {"color", "size", "linetype", "alpha", "group"}
cols = cols & set(data.columns)
num_unique_rows = len(data.drop_duplicates(cols))
ngroup = len(np.unique(data["group"].to_numpy()))
constant = num_unique_rows == ngroup
params["constant"] = constant
if not constant:
self.draw_group(data, panel_params, coord, ax, **params)
else:
for _, gdata in data.groupby("group"):
gdata.reset_index(inplace=True, drop=True)
self.draw_group(gdata, panel_params, coord, ax, **params)
@staticmethod
def draw_group(
data: pd.DataFrame,
panel_params: panel_view,
coord: Coord,
ax: Axes,
**params: Any,
):
data = coord.transform(data, panel_params, munch=True)
data["size"] *= SIZE_FACTOR
if "constant" in params:
constant: bool = params.pop("constant")
else:
constant = len(np.unique(data["group"].to_numpy())) == 1
if not constant:
_draw_segments(data, ax, **params)
else:
_draw_lines(data, ax, **params)
if "arrow" in params and params["arrow"]:
params["arrow"].draw(
data, panel_params, coord, ax, constant=constant, **params
)
@staticmethod
def draw_legend(
data: pd.Series[Any], da: DrawingArea, lyr: Layer
) -> DrawingArea:
"""
Draw a horizontal line in the box
Parameters
----------
data : Series
Data Row
da : DrawingArea
Canvas
lyr : layer
Layer
Returns
-------
out : DrawingArea
"""
from matplotlib.lines import Line2D
data["size"] *= SIZE_FACTOR
x = [0, da.width]
y = [0.5 * da.height] * 2
color = to_rgba(data["color"], data["alpha"])
key = Line2D(
x,
y,
linestyle=data["linetype"],
linewidth=data["size"],
color=color,
solid_capstyle="butt",
antialiased=False,
)
da.add_artist(key)
return da
[docs]class arrow:
"""
Define arrow (actually an arrowhead)
This is used to define arrow heads for
:class:`.geom_path`.
Parameters
----------
angle : int | float
angle in degrees between the tail a
single edge.
length : int | float
of the edge in "inches"
ends : str in ``['last', 'first', 'both']``
At which end of the line to draw the
arrowhead
type : str in ``['open', 'closed']``
When it is closed, it is also filled
"""
def __init__(
self,
angle: float = 30,
length: float = 0.2,
ends: Literal["first", "last", "both"] = "last",
type: Literal["open", "closed"] = "open",
):
self.angle = angle
self.length = length
self.ends = ends
self.type = type
def draw(
self,
data: pd.DataFrame,
panel_params: panel_view,
coord: Coord,
ax: Axes,
constant: bool = True,
**params: Any,
):
"""
Draw arrows at the end(s) of the lines
Parameters
----------
data : dataframe
Data to be plotted by this geom. This is the
dataframe created in the plot_build pipeline.
panel_params : panel_view
The scale information as may be required by the
axes. At this point, that information is about
ranges, ticks and labels. Attributes are of interest
to the geom are::
'panel_params.x.range' # tuple
'panel_params.y.range' # tuple
coord : coord
Coordinate (e.g. coord_cartesian) system of the
geom.
ax : axes
Axes on which to plot.
constant: bool
If the path attributes vary along the way. If false,
the arrows are per segment of the path
params : dict
Combined parameters for the geom and stat. Also
includes the 'zorder'.
"""
first = self.ends in ("first", "both")
last = self.ends in ("last", "both")
data = data.sort_values("group", kind="mergesort")
data["color"] = to_rgba(data["color"], data["alpha"])
if self.type == "open":
data["facecolor"] = "none"
else:
data["facecolor"] = data["color"]
if not constant:
from matplotlib.collections import PathCollection
# Get segments/points (x1, y1) -> (x2, y2)
# for which to calculate the arrow heads
idx1: list[int] = []
idx2: list[int] = []
for _, df in data.groupby("group"):
idx1.extend(df.index[:-1].to_list())
idx2.extend(df.index[1:].to_list())
d = {
"zorder": params["zorder"],
"rasterized": params["raster"],
"edgecolor": data.loc[idx1, "color"],
"facecolor": data.loc[idx1, "facecolor"],
"linewidth": data.loc[idx1, "size"],
"linestyle": data.loc[idx1, "linetype"],
}
x1 = data.loc[idx1, "x"].to_numpy()
y1 = data.loc[idx1, "y"].to_numpy()
x2 = data.loc[idx2, "x"].to_numpy()
y2 = data.loc[idx2, "y"].to_numpy()
if first:
paths = self.get_paths(x1, y1, x2, y2, panel_params, coord, ax)
coll = PathCollection(paths, **d)
ax.add_collection(coll)
if last:
x1, y1, x2, y2 = x2, y2, x1, y1
paths = self.get_paths(x1, y1, x2, y2, panel_params, coord, ax)
coll = PathCollection(paths, **d)
ax.add_collection(coll)
else:
from matplotlib.patches import PathPatch
d = {
"zorder": params["zorder"],
"rasterized": params["raster"],
"edgecolor": data["color"].iloc[0],
"facecolor": data["facecolor"].iloc[0],
"linewidth": data["size"].iloc[0],
"linestyle": data["linetype"].iloc[0],
"joinstyle": "round",
"capstyle": "butt",
}
if first:
x1, x2 = data["x"].iloc[0:2]
y1, y2 = data["y"].iloc[0:2]
x1, y1, x2, y2 = (np.array([i]) for i in (x1, y1, x2, y2))
paths = self.get_paths(x1, y1, x2, y2, panel_params, coord, ax)
patch = PathPatch(paths[0], **d)
ax.add_artist(patch)
if last:
x1, x2 = data["x"].iloc[-2:]
y1, y2 = data["y"].iloc[-2:]
x1, y1, x2, y2 = x2, y2, x1, y1
x1, y1, x2, y2 = (np.array([i]) for i in (x1, y1, x2, y2))
paths = self.get_paths(x1, y1, x2, y2, panel_params, coord, ax)
patch = PathPatch(paths[0], **d)
ax.add_artist(patch)
def get_paths(
self,
x1: npt.ArrayLike,
y1: npt.ArrayLike,
x2: npt.ArrayLike,
y2: npt.ArrayLike,
panel_params: panel_view,
coord: Coord,
ax: Axes,
) -> list[Path]:
"""
Compute paths that create the arrow heads
Parameters
----------
x1, y1, x2, y2 : array_like
List of points that define the tails of the arrows.
The arrow heads will be at x1, y1. If you need them
at x2, y2 reverse the input.
panel_params : panel_view
The scale information as may be required by the
axes. At this point, that information is about
ranges, ticks and labels. Attributes are of interest
to the geom are::
'panel_params.x.range' # tuple
'panel_params.y.range' # tuple
coord : coord
Coordinate (e.g. coord_cartesian) system of the
geom.
ax : axes
Axes on which to plot.
Returns
-------
out : list of Path
Paths that create arrow heads
"""
from matplotlib.path import Path
# The arrowhead path has 3 vertices,
# plus a dummy vertex for the STOP code
dummy = (0, 0)
# codes list remains the same after initialization
codes = [Path.MOVETO, Path.LINETO, Path.LINETO, Path.STOP]
# We need the axes dimensions so that we can
# compute scaling factors
width, height = _axes_get_size_inches(ax)
ranges = coord.range(panel_params)
width_ = np.ptp(ranges.x)
height_ = np.ptp(ranges.y)
# scaling factors to prevent skewed arrowheads
lx = self.length * width_ / width
ly = self.length * height_ / height
# angle in radians
a = self.angle * np.pi / 180
# direction of arrow head
xdiff, ydiff = x2 - x1, y2 - y1 # type: ignore
rotations = np.arctan2(ydiff / ly, xdiff / lx)
# Arrow head vertices
v1x = x1 + lx * np.cos(rotations + a)
v1y = y1 + ly * np.sin(rotations + a)
v2x = x1 + lx * np.cos(rotations - a)
v2y = y1 + ly * np.sin(rotations - a)
# create a path for each arrow head
paths = []
for t in zip(v1x, v1y, x1, y1, v2x, v2y): # type: ignore
verts = [t[:2], t[2:4], t[4:], dummy]
paths.append(Path(verts, codes))
return paths
def _draw_segments(data: pd.DataFrame, ax: Axes, **params: Any):
"""
Draw independent line segments between all the
points
"""
from matplotlib.collections import LineCollection
color = to_rgba(data["color"], data["alpha"])
# All we do is line-up all the points in a group
# into segments, all in a single list.
# Along the way the other parameters are put in
# sequences accordingly
indices: list[int] = [] # for attributes of starting point of each segment
_segments = []
for _, df in data.groupby("group"):
idx = df.index
indices.extend(idx[:-1].to_list()) # One line from two points
x = data["x"].iloc[idx]
y = data["y"].iloc[idx]
_segments.append(make_line_segments(x, y, ispath=True))
segments = np.vstack(_segments)
if color is None:
edgecolor = color
else:
edgecolor = [color[i] for i in indices]
linewidth = data.loc[indices, "size"]
linestyle = data.loc[indices, "linetype"]
coll = LineCollection(
segments,
edgecolor=edgecolor,
linewidth=linewidth,
linestyle=linestyle,
zorder=params["zorder"],
rasterized=params["raster"],
)
ax.add_collection(coll)
def _draw_lines(data: pd.DataFrame, ax: Axes, **params: Any):
"""
Draw a path with the same characteristics from the
first point to the last point
"""
from matplotlib.lines import Line2D
color = to_rgba(data["color"].iloc[0], data["alpha"].iloc[0])
join_style = _get_joinstyle(data, params)
lines = Line2D(
data["x"],
data["y"],
color=color,
linewidth=data["size"].iloc[0],
linestyle=data["linetype"].iloc[0],
zorder=params["zorder"],
rasterized=params["raster"],
**join_style,
)
ax.add_artist(lines)
def _get_joinstyle(
data: pd.DataFrame, params: dict[str, Any]
) -> dict[str, Any]:
with suppress(KeyError):
if params["linejoin"] == "mitre":
params["linejoin"] = "miter"
with suppress(KeyError):
if params["lineend"] == "square":
params["lineend"] = "projecting"
joinstyle = params.get("linejoin", "miter")
capstyle = params.get("lineend", "butt")
d = {}
if data["linetype"].iloc[0] == "solid":
d["solid_joinstyle"] = joinstyle
d["solid_capstyle"] = capstyle
elif data["linetype"].iloc[0] == "dashed":
d["dash_joinstyle"] = joinstyle
d["dash_capstyle"] = capstyle
return d
def _axes_get_size_inches(ax: Axes) -> TupleFloat2:
"""
Size of axes in inches
Parameters
----------
ax : axes
Axes
Returns
-------
out : tuple[float, float]
(width, height) of ax in inches
"""
fig = ax.get_figure()
bbox = ax.get_window_extent().transformed(
fig.dpi_scale_trans.inverted() # pyright: ignore
)
return bbox.width, bbox.height