与matplotlib的线图上的箭头

时间:2015-12-01 10:17:41

标签: python matplotlib

我想在matplotlib的线条图中添加箭头,如下图所示(使用pgfplots绘制)。

enter image description here

我该怎么办(理想情况下箭头的位置和方向应该是参数)?

以下是一些要试验的代码。

from matplotlib import pyplot
import numpy as np

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

感谢。

4 个答案:

答案 0 :(得分:17)

根据我的经验,使用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词典,直到看起来正确。

答案 1 :(得分:10)

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

编辑:更改了箭头的参数以包含位置&绘制函数的方向。

enter image description here

答案 2 :(得分:4)

不是最好的解决方案,但应该有效:

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

答案 3 :(得分:2)

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

要在图例中包含箭头,您需要制作自定义补丁处理程序并使用 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()

输出如下:

Plot showing the line-arrow combination in the legend.