Arrow on a line plot

Question

I'd like to add an arrow to a line plot with matplotlib like in the plot below (drawn with pgfplots).

How can I do (position and direction of the arrow should be parameters ideally)?

Here is some code to experiment.

from matplotlib import pyplot
import numpy as np

t = np.linspace(-2, 2, 100)
plt.plot(t, np.sin(t))
plt.show()

Thanks.

Answer 1

The easiest way is to use "marker" in plot:

from matplotlib import pyplot
import numpy as np

t = np.linspace(-2, 2, 100)
plt.plot(t, np.sin(t), marker='<',[enter image description here][1] markevery=[50])
plt.show()

Answer 2

I've found that quiver() works better than arrow() or annotate() when the x and y axes have very different scales. Here's my helper function for plotting a line with arrows:

def plot_with_arrows(ax, x, y, color="g", label="", n_arrows=2):
    ax.plot(x, y, rasterized=True, color=color, label=label)
    x_range = x.max() - x.min()
    y_range = y.max() - y.min()
    for i in np.linspace(x.keys().min(), x.keys().max(), n_arrows * 2 + 1).astype(np.int32)[1::2]:
        direction = np.array([(x[i+5] - x[i]), (y[i+5] - y[i])])
        direction = direction / (np.sqrt(np.sum(np.power(direction, 2)))) * 0.05
        direction[0] /= x_range
        direction[1] /= y_range
        ax.quiver(x[i], y[i], direction[0], direction[1], color=color)

Answer 3

I know this doesn't exactly answer the question as asked, but I thought this could be useful to other people landing here. I wanted to include the arrow in my plot's legend, but the solutions here don't mention how. There may be an easier way to do this, but here is my solution:

To include the arrow in your legend, you need to make a custom patch handler and use the matplotlib.patches.FancyArrow object. Here is a minimal working solution. This solution piggybacks off of the existing solutions in this thread.

First, the imports...

import matplotlib.pyplot as plt
from matplotlib.legend_handler import HandlerPatch
import matplotlib.patches as patches
from matplotlib.lines import Line2D

import numpy as np

Now, we make a custom legend handler. This handler can create legend artists for any line-patch combination, granted that the line has no markers.

class HandlerLinePatch(HandlerPatch):
    def __init__(self, linehandle=None, **kw):
        HandlerPatch.__init__(self, **kw)
        self.linehandle=linehandle
    
    def create_artists(self, legend, orig_handle, 
                       xdescent, ydescent, width, 
                       height, fontsize, trans):
        p = super().create_artists(legend, orig_handle, 
                                   xdescent, descent, 
                                   width, height, fontsize, 
                                   trans)
        line = Line2D([0,width],[height/2.,height/2.])
        if self.linehandle is None:
            line.set_linestyle('-')
            line._color = orig_handle._edgecolor
        else:
            self.update_prop(line, self.linehandle, legend)
            line.set_drawstyle('default')
            line.set_marker('')
            line.set_transform(trans)
        return [p[0],line]

Next, we write a function that specifies the type of patch we want to include in the legend - an arrow in our case. This is courtesy of Javier's answer here.

def make_legend_arrow(legend, orig_handle,
                      xdescent, ydescent,
                      width, height, fontsize):
    p = patches.FancyArrow(width/2., height/2., width/5., 0, 
                           length_includes_head=True, width=0, 
                           head_width=height, head_length=height, 
                           overhang=0.2)
    return p

Next, a modified version of the add_arrow function from Thomas' answer that uses the FancyArrow patch rather than annotations. This solution might cause weird wrapping like Thomas warned against, but I couldn't figure out how to put the arrow in the legend if the arrow is an annotation.

def add_arrow(line, ax, position=None, direction='right', color=None, label=''):
    """
    add an arrow to a line.

    line:       Line2D object
    position:   x-position of the arrow. If None, mean of xdata is taken
    direction:  'left' or 'right'
    color:      if None, line color is taken.
    label:      label for arrow
    """
    if color is None:
        color = line.get_color()

    xdata = line.get_xdata()
    ydata = line.get_ydata()

    if position is None:
        position = xdata.mean()
    # find closest index
    start_ind = np.argmin(np.absolute(xdata - position))
    if direction == 'right':
        end_ind = start_ind + 1
    else:
        end_ind = start_ind - 1
    
    dx = xdata[end_ind] - xdata[start_ind]
    dy = ydata[end_ind] - ydata[start_ind]
    size = abs(dx) * 5.
    x = xdata[start_ind] + (np.sign(dx) * size/2.)
    y = ydata[start_ind] + (np.sign(dy) * size/2.)

    arrow = patches.FancyArrow(x, y, dx, dy, color=color, width=0, 
                               head_width=size, head_length=size, 
                               label=label,length_includes_head=True, 
                               overhang=0.3, zorder=10)
    ax.add_patch(arrow)

Now, a helper function to plot both the arrow and the line. It returns a Line2D object, which is needed for the legend handler we wrote in the first code block

def plot_line_with_arrow(x,y,ax=None,label='',**kw):
    if ax is None:
        ax = plt.gca()
    line = ax.plot(x,y,**kw)[0]
    add_arrow(line, ax, label=label)
    return line

Finally, we make the plot and update the legend's handler_map with our custom handler.

t = np.linspace(-2, 2, 100)
y = np.sin(t)

line = plot_line_with_arrow(t,y,label='Path', linestyle=':')
plt.gca().set_aspect('equal')

plt.legend(handler_map={patches.FancyArrow : 
                        HandlerLinePatch(patch_func=make_legend_arrow, 
                                         linehandle=line)})
plt.show()

Here is the output:

Answer 4

In my experience this works best by using annotate. Thereby you avoid the weird warping you get with ax.arrow which is somehow hard to control.

EDIT: I've wrapped it into a little function.

from matplotlib import pyplot as plt
import numpy as np


def add_arrow(line, position=None, direction='right', size=15, color=None):
    """
    add an arrow to a line.

    line:       Line2D object
    position:   x-position of the arrow. If None, mean of xdata is taken
    direction:  'left' or 'right'
    size:       size of the arrow in fontsize points
    color:      if None, line color is taken.
    """
    if color is None:
        color = line.get_color()

    xdata = line.get_xdata()
    ydata = line.get_ydata()

    if position is None:
        position = xdata.mean()
    # find closest index
    start_ind = np.argmin(np.absolute(xdata - position))
    if direction == 'right':
        end_ind = start_ind + 1
    else:
        end_ind = start_ind - 1

    line.axes.annotate('',
        xytext=(xdata[start_ind], ydata[start_ind]),
        xy=(xdata[end_ind], ydata[end_ind]),
        arrowprops=dict(arrowstyle="->", color=color),
        size=size
    )


t = np.linspace(-2, 2, 100)
y = np.sin(t)
# return the handle of the line
line = plt.plot(t, y)[0]

add_arrow(line)

plt.show()

It's not very intuitive but it works. You can then fiddle with the arrowprops dictionary until it looks right.

Answer 5

Just add a plt.arrow():

from matplotlib import pyplot as plt
import numpy as np

# your function
def f(t): return np.sin(t)

t = np.linspace(-2, 2, 100)
plt.plot(t, f(t))
plt.arrow(0, f(0), 0.01, f(0.01)-f(0), shape='full', lw=0, length_includes_head=True, head_width=.05)
plt.show()

EDIT: Changed parameters of arrow to include position & direction of function to draw.

Answer 6

Not the nicest solution, but should work:

import matplotlib.pyplot as plt
import numpy as np


def makeArrow(ax,pos,function,direction):
    delta = 0.0001 if direction >= 0 else -0.0001
    ax.arrow(pos,function(pos),pos+delta,function(pos+delta),head_width=0.05,head_length=0.1)

fun = np.sin
t = np.linspace(-2, 2, 100)
ax = plt.axes()
ax.plot(t, fun(t))
makeArrow(ax,0,fun,+1)

plt.show()