How can I place an arrow at the end of the curve in matplotlib?

Question

I want arrows next to a curve. For example:

import numpy as np
import matplotlib.pyplot as plt

X = np.linspace(0,4*np.pi,10000)
Y = np.sin(X)

shift = 0.1
seg_size = 300

i = 0
plt.plot(X,Y,color='blue')
while i +seg_size < len(X):
    x = X[i:i+seg_size]
    y = Y[i:i+seg_size]+shift
    plt.plot(x,y,color='black')
    #input here command for arrow head
    i += seg_size*2

plt.show()

I tried to calculate the angle next of the line at the end of the curve and plot the arrow head lines, but I'm doing something wrong and the arrow head are deformed. Any hints?

Answer 1

The FancyArrowPatch class takes a path as an argument, so I thought that you could use that.

1) For each line segment, create a matplotlib.path.Path instance.

2) Use path instance to draw arrow.

import numpy as np
import matplotlib.pyplot as plt; plt.ion()
from matplotlib.patches import FancyArrowPatch, PathPatch
from matplotlib.path import Path

def create_path(x,y):
    vertices = zip(x,y)
    codes = [Path.MOVETO] + (len(vertices)-1) * [Path.CURVE3]
    return Path(vertices, codes)

X = np.linspace(0,4*np.pi,10000)
Y = np.sin(X)

fig, ax = plt.subplots(1,1)
ax.plot(X,Y,color='blue')

shift = 0.1
seg_size = 300

i = 0
while i +seg_size < len(X):
    x = X[i:i+seg_size]
    y = Y[i:i+seg_size]+shift

    path = create_path(x,y)

    # for testing path 
    # patch = PathPatch(path, facecolor='none', lw=2)
    # ax.add_patch(patch)

    arrow = FancyArrowPatch(path=path, color='r')
    ax.add_artist(arrow)

    i += seg_size*2

Unfortunately, that does not work, as the path that is passed to FancyArrowPatch cannot have more than 2 segments (not documented, but there is a check in ensure_quadratic_bezier).

So you have to cheat. Below I use the last 2 points of each segment to draw the arrow.

import numpy as np
import matplotlib.pyplot as plt; plt.ion()
from matplotlib.patches import FancyArrowPatch

X = np.linspace(0,4*np.pi,10000)
Y = np.sin(X)

fig, ax = plt.subplots(1,1)
ax.plot(X,Y,color='blue')

shift = 0.1
seg_size = 300

i = 0
while i +seg_size < len(X):
    x = X[i:i+seg_size]
    y = Y[i:i+seg_size]+shift

    ax.plot(x, y, 'k')

    posA, posB = zip(x[-2:], y[-2:])
    edge_width = 2.
    arrowstyle = "fancy,head_length={},head_width={},tail_width={}".format(2*edge_width, 3*edge_width, edge_width)
    arrow = FancyArrowPatch(posA=posA, posB=posB, arrowstyle=arrowstyle, color='k')
    ax.add_artist(arrow)

    i += seg_size*2