154 lines
4.9 KiB
Python
154 lines
4.9 KiB
Python
|
from matplotlib.patches import _Style, FancyArrowPatch
|
||
|
|
from matplotlib.transforms import IdentityTransform
|
||
|
|
from matplotlib.path import Path
|
||
|
|
import numpy as np
|
||
|
|
|
||
|
|
|
||
|
|
class _FancyAxislineStyle:
|
||
|
|
class SimpleArrow(FancyArrowPatch):
|
||
|
|
"""
|
||
|
|
The artist class that will be returned for SimpleArrow style.
|
||
|
|
"""
|
||
|
|
_ARROW_STYLE = "->"
|
||
|
|
|
||
|
|
def __init__(self, axis_artist, line_path, transform,
|
||
|
|
line_mutation_scale):
|
||
|
|
self._axis_artist = axis_artist
|
||
|
|
self._line_transform = transform
|
||
|
|
self._line_path = line_path
|
||
|
|
self._line_mutation_scale = line_mutation_scale
|
||
|
|
|
||
|
|
FancyArrowPatch.__init__(self,
|
||
|
|
path=self._line_path,
|
||
|
|
arrowstyle=self._ARROW_STYLE,
|
||
|
|
patchA=None,
|
||
|
|
patchB=None,
|
||
|
|
shrinkA=0.,
|
||
|
|
shrinkB=0.,
|
||
|
|
mutation_scale=line_mutation_scale,
|
||
|
|
mutation_aspect=None,
|
||
|
|
transform=IdentityTransform(),
|
||
|
|
)
|
||
|
|
|
||
|
|
def set_line_mutation_scale(self, scale):
|
||
|
|
self.set_mutation_scale(scale*self._line_mutation_scale)
|
||
|
|
|
||
|
|
def _extend_path(self, path, mutation_size=10):
|
||
|
|
"""
|
||
|
|
Extend the path to make a room for drawing arrow.
|
||
|
|
"""
|
||
|
|
from matplotlib.bezier import get_cos_sin
|
||
|
|
|
||
|
|
x0, y0 = path.vertices[-2]
|
||
|
|
x1, y1 = path.vertices[-1]
|
||
|
|
cost, sint = get_cos_sin(x0, y0, x1, y1)
|
||
|
|
|
||
|
|
d = mutation_size * 1.
|
||
|
|
x2, y2 = x1 + cost*d, y1+sint*d
|
||
|
|
|
||
|
|
if path.codes is None:
|
||
|
|
_path = Path(np.concatenate([path.vertices, [[x2, y2]]]))
|
||
|
|
else:
|
||
|
|
_path = Path(np.concatenate([path.vertices, [[x2, y2]]]),
|
||
|
|
np.concatenate([path.codes, [Path.LINETO]]))
|
||
|
|
|
||
|
|
return _path
|
||
|
|
|
||
|
|
def set_path(self, path):
|
||
|
|
self._line_path = path
|
||
|
|
|
||
|
|
def draw(self, renderer):
|
||
|
|
"""
|
||
|
|
Draw the axis line.
|
||
|
|
1) transform the path to the display coordinate.
|
||
|
|
2) extend the path to make a room for arrow
|
||
|
|
3) update the path of the FancyArrowPatch.
|
||
|
|
4) draw
|
||
|
|
"""
|
||
|
|
path_in_disp = self._line_transform.transform_path(self._line_path)
|
||
|
|
mutation_size = self.get_mutation_scale() # line_mutation_scale()
|
||
|
|
extended_path = self._extend_path(path_in_disp,
|
||
|
|
mutation_size=mutation_size)
|
||
|
|
self._path_original = extended_path
|
||
|
|
FancyArrowPatch.draw(self, renderer)
|
||
|
|
|
||
|
|
class FilledArrow(SimpleArrow):
|
||
|
|
"""
|
||
|
|
The artist class that will be returned for SimpleArrow style.
|
||
|
|
"""
|
||
|
|
_ARROW_STYLE = "-|>"
|
||
|
|
|
||
|
|
|
||
|
|
class AxislineStyle(_Style):
|
||
|
|
"""
|
||
|
|
A container class which defines style classes for AxisArtists.
|
||
|
|
|
||
|
|
An instance of any axisline style class is an callable object,
|
||
|
|
whose call signature is ::
|
||
|
|
|
||
|
|
__call__(self, axis_artist, path, transform)
|
||
|
|
|
||
|
|
When called, this should return an `.Artist` with the following methods::
|
||
|
|
|
||
|
|
def set_path(self, path):
|
||
|
|
# set the path for axisline.
|
||
|
|
|
||
|
|
def set_line_mutation_scale(self, scale):
|
||
|
|
# set the scale
|
||
|
|
|
||
|
|
def draw(self, renderer):
|
||
|
|
# draw
|
||
|
|
"""
|
||
|
|
|
||
|
|
_style_list = {}
|
||
|
|
|
||
|
|
class _Base:
|
||
|
|
# The derived classes are required to be able to be initialized
|
||
|
|
# w/o arguments, i.e., all its argument (except self) must have
|
||
|
|
# the default values.
|
||
|
|
|
||
|
|
def __init__(self):
|
||
|
|
"""
|
||
|
|
initialization.
|
||
|
|
"""
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
def __call__(self, axis_artist, transform):
|
||
|
|
"""
|
||
|
|
Given the AxisArtist instance, and transform for the path (set_path
|
||
|
|
method), return the Matplotlib artist for drawing the axis line.
|
||
|
|
"""
|
||
|
|
return self.new_line(axis_artist, transform)
|
||
|
|
|
||
|
|
class SimpleArrow(_Base):
|
||
|
|
"""
|
||
|
|
A simple arrow.
|
||
|
|
"""
|
||
|
|
|
||
|
|
ArrowAxisClass = _FancyAxislineStyle.SimpleArrow
|
||
|
|
|
||
|
|
def __init__(self, size=1):
|
||
|
|
"""
|
||
|
|
Parameters
|
||
|
|
----------
|
||
|
|
size : float
|
||
|
|
Size of the arrow as a fraction of the ticklabel size.
|
||
|
|
"""
|
||
|
|
|
||
|
|
self.size = size
|
||
|
|
super().__init__()
|
||
|
|
|
||
|
|
def new_line(self, axis_artist, transform):
|
||
|
|
|
||
|
|
linepath = Path([(0, 0), (0, 1)])
|
||
|
|
axisline = self.ArrowAxisClass(axis_artist, linepath, transform,
|
||
|
|
line_mutation_scale=self.size)
|
||
|
|
return axisline
|
||
|
|
|
||
|
|
_style_list["->"] = SimpleArrow
|
||
|
|
|
||
|
|
class FilledArrow(SimpleArrow):
|
||
|
|
ArrowAxisClass = _FancyAxislineStyle.FilledArrow
|
||
|
|
|
||
|
|
_style_list["-|>"] = FilledArrow
|