Skip to content

Commit 5a8a5f4

Browse files
apaszkeapaszke
committed
Vectorize patch extraction in Axes3D.plot_surface
When rstride and cstride divide the numbers of rows and columns minus 1, all polygons have the same number of vertices which can be recovered using some stride tricks on NumPy arrays. On a toy benchmark of a 800x800 grid this allows to reduce the time it takes to produce this list of polygons from ~10s to ~60ms.
1 parent 12f2624 commit 5a8a5f4

3 files changed

Lines changed: 135 additions & 20 deletions

File tree

‎lib/matplotlib/cbook/__init__.py‎

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -2002,6 +2002,73 @@ def _array_perimeter(arr):
20022002
))
20032003

20042004

2005+
def _unfold(arr, axis, size, step):
2006+
"""Append an extra dimension containing sliding windows along **axis**.
2007+
2008+
All windows are of size **size** and begin with every **step** elements.
2009+
2010+
Parameters
2011+
----------
2012+
arr : ndarray, shape (N_1, ..., N_k)
2013+
The input array
2014+
2015+
Returns
2016+
-------
2017+
windows : ndarray, shape (N_1, ..., 1 + (N_axis-size)/step, ..., N_k, size)
2018+
"""
2019+
new_shape = list(arr.shape) + [size]
2020+
new_strides = list(arr.strides) + [arr.strides[axis]]
2021+
new_shape[axis] = (new_shape[axis] - size) // step + 1
2022+
new_strides[axis] = new_strides[axis] * step
2023+
return np.lib.stride_tricks.as_strided(arr,
2024+
shape=new_shape,
2025+
strides=new_strides,
2026+
writeable=False)
2027+
2028+
2029+
def _array_patch_perimeters(x, rstride, cstride):
2030+
"""Extract perimeters of patches from **arr**.
2031+
2032+
Extracted patches are of size (**rstride** + 1) x (**cstride** + 1) and
2033+
share perimeters with their neighbors. The ordering of the vertices matches
2034+
that returned by ``_array_perimeter``.
2035+
2036+
Parameters
2037+
----------
2038+
arr : ndarray, shape (N, M)
2039+
2040+
Returns
2041+
-------
2042+
patches : ndarray, shape (N / rstride * M / cstride, 2(rstride + cstride))
2043+
"""
2044+
assert rstride > 0 and cstride > 0
2045+
assert (x.shape[0] - 1) % rstride == 0
2046+
assert (x.shape[1] - 1) % cstride == 0
2047+
# We build up each perimeter from four half-open intervals. Here is an
2048+
# illustrated explanation for rstride == cstride == 3
2049+
#
2050+
# T T T R
2051+
# L R
2052+
# L R
2053+
# L B B B
2054+
#
2055+
# where T means that this element will be in the top array, R for right,
2056+
# B for bottom and L for left. Each of the arrays below has a shape of:
2057+
#
2058+
# (number of perimeters that can be extracted vertically,
2059+
# number of perimeters that can be extracted horizontally,
2060+
# cstride for top and bottom and rstride for left and right)
2061+
#
2062+
# Note that _unfold doesn't incur any memory copies, so the only costly
2063+
# operation here is the np.concatenate.
2064+
top = _unfold(x[:-1:rstride, :-1], 1, cstride, cstride)
2065+
bottom = _unfold(x[rstride::rstride, 1:], 1, cstride, cstride)[..., ::-1]
2066+
right = _unfold(x[:-1, cstride::cstride], 0, rstride, rstride)
2067+
left = _unfold(x[1:, :-1:cstride], 0, rstride, rstride)[..., ::-1]
2068+
return (np.concatenate((top, right, bottom, left), axis=2)
2069+
.reshape(-1, 2 * (rstride + cstride)))
2070+
2071+
20052072
@contextlib.contextmanager
20062073
def _setattr_cm(obj, **kwargs):
20072074
"""Temporarily set some attributes; restore original state at context exit.

‎lib/matplotlib/tests/test_cbook.py‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -591,3 +591,29 @@ def test_warn_external(recwarn):
591591
cbook._warn_external("oops")
592592
assert len(recwarn) == 1
593593
assert recwarn[0].filename == __file__
594+
595+
596+
def test_array_patch_perimeters():
597+
def check(x, rstride, cstride):
598+
rows, cols = x.shape
599+
row_inds = list(range(0, rows-1, rstride)) + [rows-1]
600+
col_inds = list(range(0, cols-1, cstride)) + [cols-1]
601+
polys = []
602+
for rs, rs_next in zip(row_inds[:-1], row_inds[1:]):
603+
for cs, cs_next in zip(col_inds[:-1], col_inds[1:]):
604+
# +1 ensures we share edges between polygons
605+
ps = cbook._array_perimeter(x[rs:rs_next+1, cs:cs_next+1]).T
606+
polys.append(ps)
607+
polys = np.asarray(polys)
608+
assert np.array_equal(polys,
609+
cbook._array_patch_perimeters(
610+
x, rstride=rstride, cstride=cstride))
611+
612+
def divisors(n):
613+
return [i for i in range(1, n + 1) if n % i == 0]
614+
615+
for rows, cols in [(5, 5), (7, 14), (13, 9)]:
616+
x = np.arange(rows * cols).reshape(rows, cols)
617+
for rstride, cstride in itertools.product(divisors(rows - 1),
618+
divisors(cols - 1)):
619+
check(x, rstride=rstride, cstride=cstride)

‎lib/mpl_toolkits/mplot3d/axes3d.py‎

Lines changed: 42 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -1444,6 +1444,17 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
14441444
the input data is larger, it will be downsampled (by slicing) to
14451445
these numbers of points.
14461446
1447+
.. note::
1448+
1449+
To maximize rendering speed consider setting *rstride* and *cstride*
1450+
to divisors of the number of rows minus 1 and columns minus 1
1451+
respectively. For example, given 51 rows rstride can be any of the
1452+
divisors of 50.
1453+
1454+
Similarly, a setting of *rstride* and *cstride* equal to 1 (or
1455+
*rcount* and *ccount* equal the number of rows and columns) can use
1456+
the optimized path.
1457+
14471458
Parameters
14481459
----------
14491460
X, Y, Z : 2d arrays
@@ -1547,25 +1558,33 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
15471558
"semantic or raise an error in matplotlib 3.3. "
15481559
"Please use shade=False instead.")
15491560

1550-
# evenly spaced, and including both endpoints
1551-
row_inds = list(range(0, rows-1, rstride)) + [rows-1]
1552-
col_inds = list(range(0, cols-1, cstride)) + [cols-1]
1553-
15541561
colset = [] # the sampled facecolor
1555-
polys = []
1556-
for rs, rs_next in zip(row_inds[:-1], row_inds[1:]):
1557-
for cs, cs_next in zip(col_inds[:-1], col_inds[1:]):
1558-
ps = [
1559-
# +1 ensures we share edges between polygons
1560-
cbook._array_perimeter(a[rs:rs_next+1, cs:cs_next+1])
1561-
for a in (X, Y, Z)
1562-
]
1563-
# ps = np.stack(ps, axis=-1)
1564-
ps = np.array(ps).T
1565-
polys.append(ps)
1566-
1567-
if fcolors is not None:
1568-
colset.append(fcolors[rs][cs])
1562+
if (rows - 1) % rstride == 0 and \
1563+
(cols - 1) % cstride == 0 and \
1564+
fcolors is None:
1565+
polys = np.stack(
1566+
[cbook._array_patch_perimeters(a, rstride, cstride)
1567+
for a in (X, Y, Z)],
1568+
axis=-1)
1569+
else:
1570+
# evenly spaced, and including both endpoints
1571+
row_inds = list(range(0, rows-1, rstride)) + [rows-1]
1572+
col_inds = list(range(0, cols-1, cstride)) + [cols-1]
1573+
1574+
polys = []
1575+
for rs, rs_next in zip(row_inds[:-1], row_inds[1:]):
1576+
for cs, cs_next in zip(col_inds[:-1], col_inds[1:]):
1577+
ps = [
1578+
# +1 ensures we share edges between polygons
1579+
cbook._array_perimeter(a[rs:rs_next+1, cs:cs_next+1])
1580+
for a in (X, Y, Z)
1581+
]
1582+
# ps = np.stack(ps, axis=-1)
1583+
ps = np.array(ps).T
1584+
polys.append(ps)
1585+
1586+
if fcolors is not None:
1587+
colset.append(fcolors[rs][cs])
15691588

15701589
# note that the striding causes some polygons to have more coordinates
15711590
# than others
@@ -1578,8 +1597,11 @@ def plot_surface(self, X, Y, Z, *args, norm=None, vmin=None,
15781597
polyc.set_facecolors(colset)
15791598
polyc.set_edgecolors(colset)
15801599
elif cmap:
1581-
# doesn't vectorize because polys is jagged
1582-
avg_z = np.array([ps[:, 2].mean() for ps in polys])
1600+
# can't always vectorize, because polys might be jagged
1601+
if isinstance(polys, np.ndarray):
1602+
avg_z = polys[..., 2].mean(axis=-1)
1603+
else:
1604+
avg_z = np.array([ps[:, 2].mean() for ps in polys])
15831605
polyc.set_array(avg_z)
15841606
if vmin is not None or vmax is not None:
15851607
polyc.set_clim(vmin, vmax)

0 commit comments

Comments
 (0)