Avoid slow looping when plotting irregular raster plot using patches.Rectangle

Question

I've written a code to make an irregular raster plot (ie one in which the size of the raster rectangles is variable). Here is a minimum reproducible example below.

The problem is that looping over the blocks is very slow for my example (each plot has a lot of rectangles and there are a lot of plots to make). I tried to convert the coordinates into a list of tuples, but that threw up an error.

Is it possible to get patches.Rectangle to return a list of patches, rather than one, so I can get rid of the loop over the patches and speed up the code?

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

output:

Answer 1

I don't know exactly how to time it, but it seems like that would be the perfect job for a PatchCollection . Does this speed up your graph at all?

EDIT: crude tests seem to show a gain of performance with PatchCollection , especially when N is large. I tested here with 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)

Full code:

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) 

Answer 2

One sentence summary: Use a PolyCollection .

---

Using a collection to draw many shapes is sure more efficient than drawing individual rectangles. The other answer suggests to use a PatchCollection . Even more efficient is the use of a PolyCollection .

The reason is twofold:

1. In a PolyCollection you do not need to create each patch individually
2. It is sufficient to define exactly one shape, and only specify the sizes, colors and offsets.

I made some modifications to the code concerning the color definition (best let the collection do it for you) and the colorbar (use the collection, instead of an independen colorbar)

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

PatchCollection:

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

Single Rectangles:

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

Run all:

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

Timing

For N=1000 I get

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

So in this case using a PatchCollection is a factor 3 more efficient than single rectangles, and using a PolyCollection is a factor 3 more efficient than a PatchCollection .

Overview of the time it takes to create a figure with N rectangles with the 3 different methods from above: