Matplotlib:自定义函数,每次绘制图形时调用

时间:2018-10-17 15:04:26

标签: python matplotlib

我想创建一个包含箭头的 matplotlib 图,其箭头的头部形状与数据坐标无关。这类似于FancyArrowPatch,但是当箭头的长度小于箭头的长度时,将收缩以适合箭头的长度。

目前,我通过设置箭头的长度来解决此问题,方法是将箭头的宽度转换为显示坐标,计算显示坐标中的箭头长度,然后将其转换回数据坐标。

只要轴的尺寸不变,这种方法就可以很好地工作,例如,可能由于set_xlim()set_ylim()tight_layout()而引起。 我想通过在绘制图的尺寸确实发生变化时重新绘制箭头来涵盖这些情况。目前,我通过通过

注册函数on_draw(event)来处理此问题
axes.get_figure().canvas.mpl_connect("resize_event", on_draw)

,但这仅适用于交互式后端。对于将图另存为图像文件的情况,我还需要一种解决方案。还有其他地方可以注册我的回调函数吗?

编辑:这是代码,我当前正在使用:

def draw_adaptive_arrow(axes, x, y, dx, dy,
                        tail_width, head_width, head_ratio, draw_head=True,
                        shape="full", **kwargs):
    from matplotlib.patches import FancyArrow
    from matplotlib.transforms import Bbox

    arrow = None

    def on_draw(event=None):
        """
        Callback function that is called, every time the figure is resized
        Removes the current arrow and replaces it with an arrow with
        recalcualted head
        """
        nonlocal tail_width
        nonlocal head_width
        nonlocal arrow
        if arrow is not None:
            arrow.remove()
        # Create a head that looks equal, independent of the aspect
        # ratio
        # Hence, a transformation into display coordinates has to be
        # performed to fix the head width to length ratio
        # In this transformation only the height and width are
        # interesting, absolute coordinates are not needed
        # -> box origin at (0,0)
        arrow_box = Bbox([(0,0),(0,head_width)])
        arrow_box_display = axes.transData.transform_bbox(arrow_box)
        head_length_display = np.abs(arrow_box_display.height * head_ratio)
        arrow_box_display.x1 = arrow_box_display.x0 + head_length_display
        # Transfrom back to data coordinates for plotting
        arrow_box = axes.transData.inverted().transform_bbox(arrow_box_display)
        head_length = arrow_box.width
        if head_length > np.abs(dx):
            # If the head would be longer than the entire arrow,
            # only draw the arrow head with reduced length
            head_length = np.abs(dx)
        if not draw_head:
            head_length = 0
            head_width = tail_width
        arrow = FancyArrow(
            x, y, dx, dy,
            width=tail_width, head_width=head_width, head_length=head_length,
            length_includes_head=True, **kwargs)
        axes.add_patch(arrow)

    axes.get_figure().canvas.mpl_connect("resize_event", on_draw)



# Some place in the user code...

fig = plt.figure(figsize=(8.0, 3.0))
ax = fig.add_subplot(1,1,1)

# 90 degree tip
draw_adaptive_arrow(
    ax, 0, 0, 4, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
# Still 90 degree tip
draw_adaptive_arrow(
    ax, 5, 0, 2, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
# Smaller head, since otherwise head would be longer than entire arrow
draw_adaptive_arrow(
    ax, 8, 0, 0.5, 0, tail_width=0.4, head_width=0.8, head_ratio=0.5
)
ax.set_xlim(0,10)
ax.set_ylim(-1,1)

# Does not work in non-interactive backend
plt.savefig("test.pdf")
# But works in interactive backend
plt.show()

2 个答案:

答案 0 :(得分:1)

这是没有回调的解决方案。我主要从问题中接管了算法,因为我不确定我是否了解箭头的要求。我很确定可以简化,但是这也不是问题的重点。

因此,在这里我们将FancyArrow子类化,并将其添加到轴中。然后,我们重写draw方法以计算所需的参数,然后-在某种程度上是异常的,在其他情况下可能会失败-在draw方法中再次调用__init__

import matplotlib.pyplot as plt
import numpy as np
from matplotlib.patches import FancyArrow
from matplotlib.transforms import Bbox

class MyArrow(FancyArrow):

    def __init__(self,  *args, **kwargs):
        self.ax = args[0]
        self.args = args[1:]
        self.kw = kwargs
        self.head_ratio = self.kw.pop("head_ratio", 1)
        self.draw_head = self.kw.pop("draw_head", True)
        self.kw.update(length_includes_head=True)
        super().__init__(*self.args,**self.kw)
        self.ax.add_patch(self)
        self.trans = self.get_transform()

    def draw(self, renderer):
        self.kw.update(transform = self.trans)

        arrow_box = Bbox([(0,0),(0,self.kw["head_width"])])
        arrow_box_display = self.ax.transData.transform_bbox(arrow_box)
        head_length_display = np.abs(arrow_box_display.height * self.head_ratio)
        arrow_box_display.x1 = arrow_box_display.x0 + head_length_display
        # Transfrom back to data coordinates for plotting
        arrow_box = self.ax.transData.inverted().transform_bbox(arrow_box_display)
        self.kw["head_length"] = arrow_box.width
        if self.kw["head_length"] > np.abs(self.args[2]):
            # If the head would be longer than the entire arrow,
            # only draw the arrow head with reduced length
            self.kw["head_length"] = np.abs(self.args[2])
        if not self.draw_head:
            self.kw["head_length"] = 0
            self.kw["head_width"] = self.kw["width"]    

        super().__init__(*self.args,**self.kw)
        self.set_clip_path(self.ax.patch)
        self.ax._update_patch_limits(self)
        super().draw(renderer)



fig = plt.figure(figsize=(8.0, 3.0))
ax = fig.add_subplot(1,1,1)

# 90 degree tip
MyArrow( ax, 0, 0, 4, 0, width=0.4, head_width=0.8, head_ratio=0.5 )

MyArrow( ax, 5, 0, 2, 0, width=0.4, head_width=0.8, head_ratio=0.5 )
# Smaller head, since otherwise head would be longer than entire arrow
MyArrow( ax, 8, 0, 0.5, 0, width=0.4, head_width=0.8, head_ratio=0.5 )
ax.set_xlim(0,10)
ax.set_ylim(-1,1)

# Does not work in non-interactive backend
plt.savefig("test.pdf")
# But works in interactive backend
plt.show()

答案 1 :(得分:0)

我找到了解决问题的方法,但是,它不是很优雅。 我发现,在非交互式后端中调用的唯一回调函数是draw_path()子类的AbstractPathEffect方法。

我创建了一个AbstractPathEffect子类,该子类更新了箭头的顶点 在其draw_path()方法中。

我仍然愿意寻求其他可能更直接的解决方案。

import numpy as np
from numpy.linalg import norm
from matplotlib.patches import FancyArrow
from matplotlib.patheffects import AbstractPathEffect

class AdaptiveFancyArrow(FancyArrow):
    """
    A `FancyArrow` with fixed head shape.
    The length of the head is proportional to the width the head
    in display coordinates.
    If the head length is longer than the length of the entire
    arrow, the head length is limited to the arrow length.
    """

    def __init__(self, x, y, dx, dy,
                 tail_width, head_width, head_ratio, draw_head=True,
                 shape="full", **kwargs):
        if not draw_head:
            head_width = tail_width
        super().__init__(
            x, y, dx, dy,
            width=tail_width, head_width=head_width,
            overhang=0, shape=shape,
            length_includes_head=True, **kwargs
        )
        self.set_path_effects(
            [_ArrowHeadCorrect(self, head_ratio, draw_head)]
        )


class _ArrowHeadCorrect(AbstractPathEffect):
    """
    Updates the arrow head length every time the arrow is rendered
    """

    def __init__(self, arrow, head_ratio, draw_head):
        self._arrow = arrow
        self._head_ratio = head_ratio
        self._draw_head = draw_head

    def draw_path(self, renderer, gc, tpath, affine, rgbFace=None):
        # Indices to certain vertices in the arrow
        TIP = 0
        HEAD_OUTER_1 = 1
        HEAD_INNER_1 = 2
        TAIL_1 = 3
        TAIL_2 = 4
        HEAD_INNER_2 = 5
        HEAD_OUTER_2 = 6

        transform = self._arrow.axes.transData

        vert = tpath.vertices
        # Transform data coordiantes to display coordinates
        vert = transform.transform(vert)
        # The direction vector alnog the arrow
        arrow_vec = vert[TIP] - (vert[TAIL_1] + vert[TAIL_2]) / 2
        tail_width = norm(vert[TAIL_2] - vert[TAIL_1])
        # Calculate head length from head width
        head_width = norm(vert[HEAD_OUTER_2] - vert[HEAD_OUTER_1])
        head_length = head_width * self._head_ratio
        if head_length > norm(arrow_vec):
            # If the head would be longer than the entire arrow,
            # only draw the arrow head with reduced length
            head_length = norm(arrow_vec)
        # The new head start vector; is on the arrow vector
        if self._draw_head:
            head_start = \
            vert[TIP] - head_length * arrow_vec/norm(arrow_vec)
        else:
            head_start = vert[TIP]
        # vector that is orthogonal to the arrow vector
        arrow_vec_ortho = vert[TAIL_2] - vert[TAIL_1]
        # Make unit vector
        arrow_vec_ortho = arrow_vec_ortho / norm(arrow_vec_ortho)
        # Adjust vertices of the arrow head
        vert[HEAD_OUTER_1] = head_start - arrow_vec_ortho * head_width/2
        vert[HEAD_OUTER_2] = head_start + arrow_vec_ortho * head_width/2
        vert[HEAD_INNER_1] = head_start - arrow_vec_ortho * tail_width/2
        vert[HEAD_INNER_2] = head_start + arrow_vec_ortho * tail_width/2
        # Transform back to data coordinates
        # and modify path with manipulated vertices
        tpath.vertices = transform.inverted().transform(vert)
        renderer.draw_path(gc, tpath, affine, rgbFace)
相关问题