使用 patch.Rectangle 绘制不规则光栅图时避免缓慢循环

Question

我编写了一个代码来制作不规则的光栅图（即光栅矩形的大小是可变的）。这是下面的最小可重现示例。

问题在于，对于我的示例，循环块非常慢（每个图都有很多矩形，并且有很多图要制作）。我试图将坐标转换为元组列表，但这会引发错误。

是否有可能让 patch.Rectangle 返回一个补丁列表，而不是一个，这样我就可以摆脱补丁上的循环并加快代码速度？

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import matplotlib.colorbar as cbar

fig,ax=plt.subplots(1)

rng=6
plt.ylim(0,rng)
plt.xlim(0,rng)

N = 30
x = np.random.rand(N)*rng
y = np.random.rand(N)*rng
s = np.random.rand(N)
colors=np.random.rand(N)
normal = plt.Normalize(0,1) # my numbers from 0-1

cmap=plt.cm.RdYlBu_r
c=cmap(colors)

for i in range(N):
    val=0.5
    rect=patches.Rectangle((x[i],y[i]),s[i],s[i],
                            edgecolor='black',
                            linewidth = 1,
                            facecolor = c[i],
                            ) 
    ax.add_patch(rect)

cax, _ = cbar.make_axes(ax) 
cb2 = cbar.ColorbarBase(cax, cmap=cmap,norm=normal) 

plt.savefig("test.png")

输出：

Answer 1

一句话总结：使用PolyCollection。

---

使用集合绘制许多形状肯定比绘制单个矩形更有效。 other answer 建议使用 PatchCollection。更有效的是使用 PolyCollection。

原因有两个：

1. 在 PolyCollection 中，您无需单独创建每个补丁
2. 只定义一个形状就足够了，并且只指定大小、颜色和偏移量。

我对有关颜色定义（最好让集合为您完成）和颜色条（使用集合，而不是独立的颜色条）的代码进行了一些修改

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from matplotlib.collections import PatchCollection, PolyCollection
import matplotlib.transforms as mtrans

补丁集合：

def patchcoll(N, show=False):
    fig,ax=plt.subplots()

    rng=6
    plt.ylim(0,rng+1)
    plt.xlim(0,rng+1)

    x = np.random.rand(N)*rng
    y = np.random.rand(N)*rng
    s = np.random.rand(N)
    c = np.random.rand(N)
    norm = plt.Normalize(0,1) # my numbers from 0-1
    cmap=plt.cm.RdYlBu_r

    pat = []

    for i in range(N):
        rect=patches.Rectangle((x[i],y[i]),s[i],s[i])
        pat.append(rect)

    col = PatchCollection(pat, cmap=cmap, norm=norm)
    col.set_array(c)
    col.set_edgecolor('k')
    col.set_linewidth(1.)
    ax.add_collection(col)


    fig.colorbar(col)
    if show:
        plt.show()
    else:
        fig.canvas.draw() 

    plt.close()

PolyCollection：

def polycoll(N, show=False):
    fig,ax=plt.subplots()

    rng=6
    plt.ylim(0,rng)
    plt.xlim(0,rng)

    x = np.random.rand(N)*rng
    y = np.random.rand(N)*rng
    s = np.random.rand(N)
    c = np.random.rand(N)
    norm = plt.Normalize(0,1) # my numbers from 0-1
    cmap=plt.cm.RdYlBu_r

    offsets = np.c_[x,y]
    verts = list(zip([0,1,1,0,0], [0,0,1,1,0]))

    col = PolyCollection([verts], sizes=s, offsets=offsets, 
                         transOffset=mtrans.IdentityTransform(),
                         offset_position="data", cmap=cmap, norm=norm)

    col.set_array(c)
    col.set_edgecolor('k')
    col.set_linewidth(1.)
    ax.add_collection(col)

    fig.colorbar(col)

    if show:
        plt.show()
    else:
        fig.canvas.draw() 

    plt.close()

单个矩形：

def rectangles(N, show=False):
    fig,ax=plt.subplots()

    rng=6
    plt.ylim(0,rng)
    plt.xlim(0,rng)

    x = np.random.rand(N)*rng
    y = np.random.rand(N)*rng
    s = np.random.rand(N)
    c = np.random.rand(N)
    norm = plt.Normalize(0,1) # my numbers from 0-1

    cmap=plt.cm.RdYlBu_r

    for i in range(N):
        rect=patches.Rectangle((x[i],y[i]),s[i],s[i], 
                               facecolor=cmap(norm(c[i])), edgecolor="k", linewidth=1)
        ax.add_patch(rect)


    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    fig.colorbar(sm)

    if show:
        plt.show()
    else:
        fig.canvas.draw() 

    plt.close()

全部运行：

patchcoll(30, show=True)
polycoll(30,show=True)
rectangles(30,show=True)

时间

对于N=1000 我明白了

%timeit(rectangles(1000))
757 ms ± 4.26 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit(patchcoll(1000))
184 ms ± 462 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit(polycoll(1000))
58.3 ms ± 146 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

因此，在这种情况下，使用 PatchCollection 的效率比单个矩形高 3 倍，使用 PolyCollection 的效率比 PatchCollection 高 3 倍。

概述使用上述 3 种不同方法创建具有 N 矩形的图形所需的时间：

Answer 2

我不知道具体如何计时，但对于PatchCollection 来说，这似乎是一项完美的工作。这会加速你的图表吗？

编辑： 粗略的测试似乎显示PatchCollection 的性能有所提高，尤其是当 N 很大时。我在这里用 N=1000 进行了测试：

%timeit withCollection()
316 ms ± 5.41 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

%timeit withoutCollection()
772 ms ± 30.5 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

完整代码：

import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
import matplotlib.colorbar as cbar
from matplotlib.collections import PatchCollection

fig,ax=plt.subplots()

rng=6
plt.ylim(0,rng)
plt.xlim(0,rng)

N = 30
x = np.random.rand(N)*rng
y = np.random.rand(N)*rng
s = np.random.rand(N)
colors=np.random.rand(N)
normal = plt.Normalize(0,1) # my numbers from 0-1

cmap=plt.cm.RdYlBu_r
c=cmap(colors)
pat = []

for i in range(N):
    rect=patches.Rectangle((x[i],y[i]),s[i],s[i])
    pat.append(rect)

col = PatchCollection(pat)
col.set_facecolor(c)
col.set_edgecolor('k')
col.set_linewidth(1.)
ax.add_collection(col)


cax, _ = cbar.make_axes(ax) 
cb2 = cbar.ColorbarBase(cax, cmap=cmap,norm=normal)