使用 matplotlib 的线图上的箭头

Question

我想在带有matplotlib 的线图中添加一个箭头，如下图所示（使用pgfplots 绘制）。

我该怎么做（理想情况下箭头的位置和方向应该是参数）？

这是一些实验代码。

from matplotlib import pyplot
import numpy as np

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

谢谢。

Answer 1

根据我的经验，使用annotate 效果最好。因此，您可以避免使用 ax.arrow 时出现的奇怪变形，这种变形在某种程度上难以控制。

编辑：我已经将它包装成一个小函数。

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

它不是很直观，但很有效。然后，您可以摆弄arrowprops 字典，直到它看起来正确为止。

Answer 2

只需添加一个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()

编辑：更改了箭头的参数以包括要绘制的函数的位置和方向。

Answer 3

不是最好的解决方案，但应该可以：

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

Answer 4

我知道这并不能完全回答所提出的问题，但我认为这可能对其他登陆这里的人有用。我想在我的情节图例中包含箭头，但这里的解决方案没有提到如何。可能有更简单的方法可以做到这一点，但这是我的解决方案：

要在图例中包含箭头，您需要创建一个自定义补丁处理程序并使用 matplotlib.patches.FancyArrow 对象。这是一个最小的工作解决方案。该解决方案借鉴了该线程中的现有解决方案。

首先，导入...

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

现在，我们制作一个自定义图例处理程序。这个处理程序可以为任何线条-补丁组合创建图例艺术家，前提是线条没有标记。

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]

接下来，我们编写一个函数来指定我们想要包含在图例中的补丁类型——在我们的例子中是一个箭头。这是由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

接下来，Thomas 回答中 add_arrow 函数的修改版本使用 FancyArrow 补丁而不是注释。此解决方案可能会导致像Thomas warned against 这样的奇怪包装，但如果箭头是注释，我无法弄清楚如何将箭头放在图例中。

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)

现在，绘制箭头和线的辅助函数。它返回一个 Line2D 对象，这是我们在第一个代码块中编写的图例处理程序所需要的

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

最后，我们制作情节并使用我们的自定义处理程序更新图例的handler_map。

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

这是输出：