-
Notifications
You must be signed in to change notification settings - Fork 28
/
Copy pathline_label.py
259 lines (218 loc) · 8.55 KB
/
line_label.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
from __future__ import annotations
from typing import TYPE_CHECKING
import matplotlib.dates as mdates
import matplotlib.patheffects as patheffects
import numpy as np
from datetime import timedelta
from matplotlib.text import Text
from .utils import normalize_xydata
if TYPE_CHECKING:
from datetime import datetime
from typing import Any, Literal, Optional, Union
from matplotlib.axes import Axes
from matplotlib.lines import Line2D
Position = Union[float, datetime, np.datetime64]
ColorLike = Any # mpl has no type annotations so this is just a crutch
AutoLiteral = Literal["auto"]
class LineLabel(Text):
"""This artist adds a label onto a preexisting Line2D object"""
_line: Line2D
"""Annotated line"""
_target_x: Position
"""Requested x position of the label, as supplied by the user"""
_ax: Axes
"""Axes containing the line"""
_auto_align: bool
"""Align text with the line (True) or parallel to x axis (False)"""
_xoffset: float
"""An additional x offset for the label"""
_xoffset_logspace: bool
"""Sets whether to treat _xoffset exponentially"""
_yoffset: float
"""An additional y offset for the label"""
_yoffset_logspace: bool
"""Sets whether to treat _yoffset exponentially"""
_label_pos: np.ndarray
"""Position of the label, computed from _target_x and line data"""
_anchor_a: np.ndarray
"""Anchor A for rotation calculation, point of _line neighbouring this label"""
_anchor_b: np.ndarray
"""Anchor B for rotation calculation, point of _line neighbouring this label"""
def __init__(
self,
line: Line2D,
x: Position,
label: Optional[str] = None,
align: Optional[bool] = None,
xoffset: float = 0,
xoffset_logspace: bool = False,
yoffset: float = 0,
yoffset_logspace: bool = False,
outline_color: Optional[Union[AutoLiteral, ColorLike]] = "auto",
outline_width: float = 8,
rotation: Optional[float] = None,
**kwargs,
) -> None:
"""
Parameters
----------
line : Line2D
Line to be decorated.
x : Position
Horizontal target position for the label (in data units).
label : str, optional
Override for line label, by default None.
align : bool, optional
If true, the label is parallel to the line, otherwise horizontal,
by default True.
xoffset : float, optional
An additional x offset for the line label, by default 0.
xoffset_logspace : bool, optional
If true xoffset is applied exponentially to appear linear on a log-axis,
by default False.
yoffset : float, optional
An additional y offset for the line label, by default 0.
yoffset_logspace : bool, optional
If true yoffset is applied exponentially to appear linear on a log-axis,
by default False.
outline_color : None | "auto" | Colorlike
Colour of the outline. If set to "auto", use the background color.
If set to None, do not draw an outline, by default "auto".
outline_width : float
Width of the outline, by default 8.
rotation: float, optional
If set and align = False, controls the angle of the label.
"""
# When rotation is set, align has to be false or None
if rotation is not None and align:
raise ValueError(
f"When rotation is set, align needs to be false or none was {align=}."
)
elif rotation is None:
align = True if (align or align is None) else False
rotation = 0
elif rotation is None and not align:
align = False
rotation = 0
elif not align or align is None:
align = False
rotation = rotation
self._line = line
self._target_x = x
self._ax = line.axes
self._auto_align = align
self._xoffset = xoffset
self._xoffset_logspace = xoffset_logspace
self._yoffset = yoffset
self._yoffset_logspace = yoffset_logspace
label = label or line.get_label()
# Populate self._pos, self._anchor_a, self._anchor_b
self._update_anchors()
self._rotation = rotation
# Set a bunch of default arguments
kwargs.setdefault("color", self._line.get_color())
kwargs.setdefault("clip_on", True)
kwargs.setdefault("zorder", 2.5)
if "ha" not in kwargs:
kwargs.setdefault("horizontalalignment", "center")
if "va" not in kwargs:
kwargs.setdefault("verticalalignment", "center")
# Initialize Text Artist
super().__init__(
*self._label_pos,
label,
rotation=self._rotation,
rotation_mode="anchor",
**kwargs,
)
# Apply outline effect
if outline_color is not None:
if outline_color == "auto":
outline_color = line.axes.get_facecolor()
self.set_path_effects(
[
patheffects.Stroke(
linewidth=outline_width, foreground=outline_color
),
patheffects.Normal(),
]
)
# activate clipping if needed and place on axes
if kwargs["clip_on"]:
self.set_clip_path(self._ax.patch)
self._ax._add_text(self)
def _update_anchors(self):
"""
This helper method computes the position of the textbox and determines
the anchor points needed to adjust the rotation
"""
# Use the mpl-internal float representation (deals with datetime etc)
x = self._line.convert_xunits(self._target_x)
xdata, ydata = normalize_xydata(self._line)
# Convert timedelta to float if needed
if isinstance(self._xoffset, timedelta):
xoffset = mdates.date2num(self._xoffset + self._target_x) - x
else:
xoffset = self._xoffset
# Handle nan values
mask = np.isfinite(ydata)
if mask.sum() == 0:
raise ValueError(f"The line {self._line} only contains nan!")
xdata = xdata[mask]
ydata = ydata[mask]
# Find the first line segment surrounding x
for i, (xa, xb) in enumerate(zip(xdata[:-1], xdata[1:])):
if min(xa, xb) <= x <= max(xa, xb):
ya, yb = ydata[i], ydata[i + 1]
break
else:
raise ValueError("x label location is outside data range!")
# Interpolate y position of label, (interp needs sorted data)
if xa != xb:
dx = np.array((xa, xb))
dy = np.array((ya, yb))
srt = np.argsort(dx)
y = np.interp(x, dx[srt], dy[srt])
else: # Vertical case
y = (ya + yb) / 2
# Apply x offset
if self._xoffset_logspace:
x *= 10**xoffset
else:
x += xoffset
# Apply y offset
if self._yoffset_logspace:
y *= 10**self._yoffset
else:
y += self._yoffset
if not np.isfinite(y):
raise ValueError(
f"{self._line} does not have a well defined value"
f" at x = {self._target_x}. Consider a different position."
)
self._label_pos = np.array((x, y))
self._anchor_a = np.array((xa, ya))
self._anchor_b = np.array((xb, yb))
def __auto_align(self, value=None):
# Helper function to help resize the alignment of
# the label if the window is resized
# Providing the _rotation property
# enables automatic adjustment of the rotation angle
# Adapted from https://stackoverflow.com/a/53111799
if self._auto_align:
# Transform to screen coordinated to make sure the angle is always
# correct regardless of axis scaling etc.
xa, ya = self._ax.transData.transform(self._anchor_a)
xb, yb = self._ax.transData.transform(self._anchor_b)
angle = np.rad2deg(np.arctan2(yb - ya, xb - xa))
# Correct the angle to make sure text is always upright-ish
value = (angle + 90) % 180 - 90
if isinstance(value, (float, int)):
self.__rotation = value
return self.__rotation
@property
def _rotation(self):
return self.__auto_align()
@_rotation.setter
def _rotation(self, rotation):
self.__rotation = self.__auto_align(rotation)