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

wnvonmuf  于 2023-10-24  发布在  其他
关注(0)|答案(2)|浏览(113)

我想创建一个包含箭头的 matplotlib 图,其头部形状与数据坐标无关。这类似于FancyArrowPatch,但当箭头长度小于头部长度时,头部长度会缩小以适应箭头的长度。
目前,我通过将宽度转换为显示坐标来设置箭头的长度,计算显示坐标中的头部长度并将其转换回数据坐标来解决这个问题。
只要轴的尺寸不改变,这种方法就可以很好地工作,例如,由于set_xlim()set_ylim()tight_layout(),可能会发生这种情况。我想涵盖这些情况,只要图的尺寸改变,就通过重新绘制箭头来解决。目前,我通过注册函数on_draw(event)来处理这个问题,通过

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

但这只适用于交互式后端。2我还需要一个解决方案的情况下,我保存的情节作为图像文件。3有没有其他地方,我可以注册我的回调函数?
编辑:以下是我目前使用的代码:

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()
kpbwa7wx

kpbwa7wx1#

这是一个没有回调的解决方案。我从问题中接管了大部分算法,因为我不确定我是否理解了箭头的要求。我很确定这可以简化,但这也超出了问题的重点。
所以这里我们子类化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()
r6l8ljro

r6l8ljro2#

我找到了一个解决这个问题的方法,但是,它不是很优雅。我发现,在非交互式后端中调用的唯一回调函数是AbstractPathEffect子类的draw_path()方法。
我创建了一个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)

相关问题