Custom functions, that are called each time a figure is drawn

Question

I want to create a matplotlib plot containing arrows, whose head's shape is independent from the data coordinates. This is similar to FancyArrowPatch, but when the arrow length is smaller than the head length is shrank to fit the length of the arrow.

Currently, I solve this by setting the length of the arrow head by transforming the width to display coordinates, calculating the head length in display coordinates and transform it back into data coordinates.

This approach works well as long the axes' dimensions do not change, which can happen due to set_xlim(), set_ylim() or tight_layout() for example. I want to cover these cases, by redrawing the arrow whenever the plot's dimensions do change. At the moment I handle this by registering a function on_draw(event) via

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

but this does only work for interactive backends. I also need a solution for cases, where I save the plot as image file. Is there any other place, where I can register my callback function?

EDIT: Here is the code, I am currently using:

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

Unfortunately the question does not show any of the case that is used, so one cannot help with the actual problem. But did you consider just calling `on_draw` manually before saving? — ImportanceOfBeingErnest, Oct 17 '18 at 15:23
I added the code I am currently using. Calling `on_draw` manually is not an option, since this code should become part of a framework, so it would be inconvenient for the user. — Padix Key, Oct 17 '18 at 15:34
What happens if you connect to the `draw_event`? Would that work? It's sure a bit inefficient. Would you mind extending the code example with a usecase (i.e. how to call that function) and a small recipee (as "run the code, press the L key" or whatever) to see why creating a new arrow is necessary. There might be ways that would not require any callback, but I would need a test case to see if that works. — ImportanceOfBeingErnest, Oct 17 '18 at 16:29
I also tried the `'draw_event'` but I got almost the same result, with the difference, that it requires the plot to be interactively resized for the arrow to be drawn. It does not have any effect for a non-interactive backend. I added a potential use case to the code snippet — Padix Key, Oct 18 '18 at 09:07
Oh yes, there was a problem with `draw_event`. See [this answer](https://stackoverflow.com/a/42972469/4124317), where I create a timer to call the draw manually after the event. — ImportanceOfBeingErnest, Oct 18 '18 at 14:08

ImportanceOfBeingErnest · Accepted Answer · 2018-10-19T00:55:17.940

Here is a solution without callback. I took over mostly the algorithm from the question, because I'm not sure I understand the requirements for the arrow. I'm pretty sure that can be simplified, but that's also beyond the point of the question.

So here we subclass FancyArrow and let it add itself to the axes. We then override the draw method to calculate the needed parameters and then - which is somehow unusual and may in other cases fail - call __init__ again inside the draw method.

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

Padix Key · Answer 2 · 2018-10-18T13:12:51.913

I found a solution to the problem, however, it is not very elegant. The only callback function, I found, that is called in non-interactive backends, is the draw_path() method of AbstractPathEffect subclasses.

I created an AbstractPathEffect subclass that updates the vertices of the arrow head in its draw_path() method.

I am still open for other probably more straight forward solutions to my problem.

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)

Custom functions, that are called each time a figure is drawn

2 Answers2