How to make a FancyArrow with a round tail?

Viewed 90

Basically, the title already says it all. I would like to create a FancyArrow with a round tail. I bet it's very simple, but somehow I didn't manage to do it, nor find any documentation about it (only this and this). I tried using the capstyle kwarg, but it doesn't do anything:

import matplotlib.pyplot as plt
from matplotlib.patches import FancyArrow


fig, ax = plt.subplots()
for i, cs in enumerate(['butt', 'round', 'projecting']):
    arrow = FancyArrow(0, 0+i, 1, 0, width=0.1, capstyle=cs)
    ax.add_patch(arrow)

ax.relim()
ax.autoscale_view()
fig.show()

enter image description here

What am I missing?

EDIT: I just noticed that there seemed to have been a related bug in the past, however it should be fixed by now. I'm using matplotlib v. 3.5.1.

1 Answers

Ok, so as it turns out, capstyle has nothing to do with the shape of the arrow as a whole, and only affects the contour. So that was a fallacy on my end. Unfortunately, there is currently no built-in option to change the shape of the tail; thus we need to come up with our own solution.

EDIT: Ok, I found a way to make it work by updating the FancyArrow class. Maybe somebody else can improve upon it:

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


class FancierArrow(FancyArrow):
    """
    Like FancyArrow, but lets you choose a round butt.
    """

    def __str__(self):
        return "FancierArrow()"

    def __init__(
        self, x, y, dx, dy, width=0.001, length_includes_head=False,
        head_width=None, head_length=None, shape='full', overhang=0,
        head_starts_at_zero=False, round_tail=False, **kwargs
    ):
        self._round_tail = round_tail
        super().__init__(
            x, y, dx, dy, width=width,
            length_includes_head=length_includes_head,
            head_width=head_width, head_length=head_length, shape=shape,
            overhang=overhang, head_starts_at_zero=head_starts_at_zero,
            **kwargs
        )

    def set_data(
        self, *, x=None, y=None, dx=None, dy=None, width=None,
        head_width=None, head_length=None, round_tail=None,
    ):
        if round_tail is not None:
            self._round_tail = round_tail
        super().__init__(
            x=x, y=y, dx=dx, dy=dy, width=width, head_width=head_width,
            head_length=head_length
        )

    def _make_verts(self):
        super()._make_verts()
        if self._round_tail:
            self._make_tail_round()

    def _make_tail_round(self):
        nvert = 25
        add_rot = np.pi/2
        radius = width/2
        x, y = self._x, self._y
        dx, dy = self._dx, self._dy

        if self._head_starts_at_zero:
            x, y = (self.verts[4] + self.verts[3])/2
            dx, dy = (self.verts[4] - self.verts[3])/2
            add_rot += np.pi/2

        if self._shape in ['left', 'right']:
            radius /= 2
            x, y = (self.verts[4] + self.verts[3])/2
            dx, dy = (self.verts[4] - self.verts[3])/2
            if self._shape == 'left' and not self._head_starts_at_zero:
                add_rot += np.pi/2
            elif self._shape == 'right':
                if self._head_starts_at_zero:
                    add_rot += np.pi
                else:
                    add_rot -= np.pi/2

        angle = np.arctan2(dx, dy) + add_rot
        circ_segs = np.linspace(0, np.pi, nvert+1) + angle
        tail_verts = np.zeros((nvert, 2), float)
        tail_verts[:, 0] = np.sin(circ_segs)[:-1]*radius + x
        tail_verts[:, 1] = np.cos(circ_segs)[:-1]*radius + y
        if self._shape == 'right':
            tail_verts = tail_verts[::-1]

        self.verts = np.array([*self.verts[:4], *tail_verts, *self.verts[4:]])


fig, ax = plt.subplots(dpi=300)
fig.tight_layout()
width = 0.3
xs = np.sin(np.linspace(0, 2*np.pi, 9))[:-1]
ys = np.cos(np.linspace(0, 2*np.pi, 9))[:-1]
for x, y in zip(xs, ys):
    polygon = FancierArrow(
        x, y, x, y, width=width, edgecolor='black',
        round_tail=True,
        head_starts_at_zero=True,
        shape='left'
    )
    ax.add_patch(polygon)

ax.relim()
ax.autoscale_view()
fig.show()

enter image description here

I also made this a feature request on the matplotlib github repository, for those that might want to follow along.

Related