Matplotlib: Как получить цветовой градиент в виде стрелки рядом с сюжетом? - PullRequest
Новая
       20

Matplotlib: Как получить цветовой градиент в виде стрелки рядом с сюжетом?

1 голос
/ 17 мая 2019

Я пытаюсь создать график с matplotlib, который включает несколько различных линий, искусственно смещенных, которые окрашены в соответствии с цветовой картой RedBlue mpl.cm.RdBu. Теперь я хочу, чтобы стрелка рядом с графиком действовала как эффективная цветовая шкала, означая, что у нее должен быть градиент цвета.

До сих пор мне удалось создать саму стрелку, используя annotate с помощью этого ответа и нарисовав "Радужную стрелу" внутри графика, используя этот блестящий ответ (Примечание. Вам понадобится matplotlib 2.2.4 или старше для запуска этой части кода, см. Комментарии.).

Это MWE, который я могу произвести до сих пор: 

import numpy as np
import matplotlib as mpl
import matplotlib.pyplot as plt
import matplotlib.transforms
import matplotlib.path
from matplotlib.collections import LineCollection

# from https://stackoverflow.com/questions/47163796/using-colormap-with-annotate-arrow-in-matplotlib
def rainbowarrow(ax, start, end, cmap="viridis", n=50,lw=3):
    cmap = plt.get_cmap(cmap,n)
    # Arrow shaft: LineCollection
    x = np.linspace(start[0],end[0],n)
    y = np.linspace(start[1],end[1],n)
    points = np.array([x,y]).T.reshape(-1,1,2)
    segments = np.concatenate([points[:-1],points[1:]], axis=1)
    lc = LineCollection(segments, cmap=cmap, linewidth=lw)
    lc.set_array(np.linspace(0,1,n))
    ax.add_collection(lc)
    # Arrow head: Triangle
    tricoords = [(0,-0.4),(0.5,0),(0,0.4),(0,-0.4)]
    angle = np.arctan2(end[1]-start[1],end[0]-start[0])
    rot = matplotlib.transforms.Affine2D().rotate(angle)
    tricoords2 = rot.transform(tricoords)
    tri = matplotlib.path.Path(tricoords2, closed=True)
    ax.scatter(end[0],end[1], c=1, s=(2*lw)**2, marker=tri, cmap=cmap,vmin=0)
    ax.autoscale_view()

def plot_arrow(data,n):
    fig, subfig = plt.subplots(1,1,figsize=(6.28,10)) # plotsize, etc
    colorP=mpl.cm.RdBu(0.2)

    i = 0
    while i<=n-1:
        subfig.plot(data[i,0], (data[i,1])+i, lw=2, color=mpl.cm.RdBu(1-i/20)) # plot of data
        i=i+1

    subfig.annotate('', xy=(1.1,0), xycoords='axes fraction', xytext=(1.1,1), 
            arrowprops=dict(arrowstyle="<-", lw = 3))
    subfig.annotate('A', xy=(1.1,0), xycoords='axes fraction', xytext=(1.1,1))
    subfig.annotate('B', xy=(1.1,0), xycoords='axes fraction', xytext=(1.1,0))

    rainbowarrow(subfig, (1.1,3), (1.1,5), cmap='RdBu_r', n=100,lw=3)

    plt.show(fig)
    plt.close(fig)


# things to plot
np.random.seed(19680802)
n = 20
i = 0
data = np.empty([n,2,10])
while i<=n-1:
    data[i]=np.sin(np.random.rand(10))
    i = i+1

# actual plot
plot_arrow(data,n)

Вот сгенерированный график:

Resulting matplotlib plot

В двух словах: я хочу, чтобы стрелка аннотации вне графика имела цвет карты цветов, как маленькая радужная стрелка внутри графика.

...