CODEWITHSUNDEEP
pythonmatplotlib
Yotam
Yotam

Reputation: 10685

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

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?

Upvotes: 5

Views: 3791

Answers (2)

Paul Brodersen
Paul Brodersen

Reputation: 13031

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

enter image description here

Upvotes: 2

Gustav Rasmussen
Gustav Rasmussen

Reputation: 3961

How about using the 

numpy.annotate()

function?

See examples of arrows created by the pyplot.annotate() function here:

https://matplotlib.org/examples/pylab_examples/annotation_demo.html

Upvotes: 0

Related Questions