Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Significantly improve tight layout performance for cartopy axes #21935

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Dec 21, 2021

Conversation

lukelbd
Copy link
Contributor

@lukelbd lukelbd commented Dec 12, 2021

This PR is adapted from SciTools/cartopy#1956 following our discussion there.

Summary

This PR significantly improves the speed of Axes.get_tightbbox() for non-rectilinear axes (in particular, matplotlib's PolarAxes and cartopy's GeoAxes cartopy GeoAxes with transformed coordinates) by preventing unnecessary and expensive get_window_extent() computations. In the presence of complex, high-resolution artists, it can result in a 2x improvement to the draw time for cartopy GeoAxes when "tight layout" is enabled.

There may be a better way to implement this -- looking forward to everyone's thoughts.

Details

To prevent unnecessary and expensive get_window_extent() computations, Axes.get_tightbbox() skips artists with clip_on set to True and whose clip_box.extents are equivalent to ax.bbox.extents. However, while all axes artists are clipped by TransformedPatchPath(ax.patch) by default, only artists drawn inside rectilinear projections are clipped by ax.bbox (see the below example).

This PR replaces Artist._get_clipping_extent_bbox() with Artist._is_axes_clipped(). Now,Axes.get_tightbbox() includes ax.patch.get_window_extent() in its computation, and skips all artists clipped by either ax.bbox or ax.patch (i.e., artists for which Artist._is_axes_clipped() returns True).

This PR also removes the computation of the intersection of clip_box and clip_path, under the assumption that the independent test of clip_path covers those instances, but perhaps that should be added back.

Example

Here is an example with a cartopy GeoAxes:

import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
fig, ax = plt.subplots(subplot_kw={'projection': ccrs.Robinson()})
N = 5000  # large dataset
lon = np.linspace(-180, 180, N)
lat = np.linspace(-90, 90, N)
data = np.random.rand(N, N)
m = ax.pcolormesh(lon, lat, data, transform=ccrs.PlateCarree())
print(m.get_clip_box() is None)  # returns True
%timeit fig.tight_layout()

Performance before:

638 ms ± 67.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Performance after:

2.37 ms ± 60.2 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

And of course the performance difference is larger the larger the dataset.

Checklist

Tests and Styling

  • Has pytest style unit tests (and pytest passes).
  • Is Flake 8 compliant (install flake8-docstrings and run flake8 --docstring-convention=all).

Documentation

  • (n/a?) New features are documented, with examples if plot related.
  • (n/a?) New features have an entry in doc/users/next_whats_new/ (follow instructions in README.rst there).
  • (n/a?) API changes documented in doc/api/next_api_changes/ (follow instructions in README.rst there).
  • Documentation is sphinx and numpydoc compliant (the docs should build without error).

@lukelbd lukelbd changed the title Improve tightbbox speed for non-rectilinear axes Significantly improve tightbbox speed for non-rectilinear axes Dec 12, 2021
@lukelbd lukelbd changed the title Significantly improve tightbbox speed for non-rectilinear axes Significantly improve tight layout performance for non-rectilinear axes Dec 12, 2021
@lukelbd lukelbd force-pushed the tightbbox-speedup branch 2 times, most recently from 136de27 to c7501dd Compare December 12, 2021 23:01
@tacaswell tacaswell added this to the v3.6.0 milestone Dec 16, 2021
@lukelbd
Copy link
Contributor Author

lukelbd commented Dec 16, 2021

Added a simple test that artist._is_axes_clipped() returns True under default conditions / various projections. Note all non-rectilinear axes (e.g. PolarAxes, MollweideAxes, cartopy GeoAxes) have artist.get_clip_box() set to None.

@lukelbd
Copy link
Contributor Author

lukelbd commented Dec 16, 2021

I can also see the origin of this bug: The Artist._get_clipping_extent_bbox() lines that I removed (previously used inside Axes.get_tightbbox() to skip artists) mirror the below lines in Artist.get_tightbbox().

def get_tightbbox(self, renderer):
"""
Like `.Artist.get_window_extent`, but includes any clipping.
Parameters
----------
renderer : `.RendererBase` subclass
renderer that will be used to draw the figures (i.e.
``fig.canvas.get_renderer()``)
Returns
-------
`.Bbox`
The enclosing bounding box (in figure pixel coordinates).
"""
bbox = self.get_window_extent(renderer)
if self.get_clip_on():
clip_box = self.get_clip_box()
if clip_box is not None:
bbox = Bbox.intersection(bbox, clip_box)
clip_path = self.get_clip_path()
if clip_path is not None and bbox is not None:
clip_path = clip_path.get_fully_transformed_path()
bbox = Bbox.intersection(bbox, clip_path.get_extents())
return bbox

While Artist.get_tightbbox() starts the if block with a valid bbox, Artist._get_clipping_extent_bbox() starts with bbox = None... so when it gets to Line 361, the clip_path is ignored, because there is no bbox to intersect with.

@lukelbd
Copy link
Contributor Author

lukelbd commented Dec 17, 2021

I've improved this PR in a force push:

  • Renamed _is_axes_clipped() to _fully_clipped_to_axes (thanks @tacaswell).
  • Added test lines that assert _fully_clipped_to_axes() is False for non-default clipping situations (re: @jklymak).
  • Removed the redundant addition of ax.patch to the bboxes (it's included in get_default_bbox_extra_artists()).
  • Condensed multiple return statements into a single boolean return statement (possibly cleaner?)
  • Moved the line that skips artists from Axes.get_tightbbox() to Axes.get_default_bbox_extra_artists(). Note this means the public get_default_bbox_extra_artists() will return different values compared to previous versions.

I also made the following additional changes:

  • Removed an unnecessary bbox is not None check in Artist.get_tightbbox() (this is always True).
  • Removed axis instances from get_default_bbox_extra_artists() (was redundant -- see Axes.get_tightbbox()).
  • Replaced Axes.get_tightbbox() reference to [xy]axis with ax._get_axis_list() (consistent with removed lines).
  • Always include artists that don't internally implement clipping (should fix obscure tight layout "bugs").

The latter points also resolve a second inefficiency: Currently, Axes.get_tightbbox() both 1) computes the xaxis and yaxis tight bbox using for_layout_only=True on these lines:

if self.axison:
if self.xaxis.get_visible():
try:
bb_xaxis = self.xaxis.get_tightbbox(
renderer, for_layout_only=for_layout_only)
except TypeError:
# in case downstream library has redefined axis:
bb_xaxis = self.xaxis.get_tightbbox(renderer)
if bb_xaxis:
bb.append(bb_xaxis)
if self.yaxis.get_visible():
try:
bb_yaxis = self.yaxis.get_tightbbox(
renderer, for_layout_only=for_layout_only)
except TypeError:
# in case downstream library has redefined axis:
bb_yaxis = self.yaxis.get_tightbbox(renderer)
if bb_yaxis:
bb.append(bb_yaxis)

and 2) includes xaxis and yaxis as "extra artists" in get_default_bbox_extra_artists(). This meant that their bounding boxes would be effectively calculated twice if they were not "skipped" by the _get_clipping_extent_bbox check. For some reason, this check was successfully skipping xaxis and yaxis in rectilinear axes, because they have a clip_on set to True and clip_box set to ax.bbox (despite the fact that they are not clipped when drawn..... weird). However, this failed to skip non-rectilinear axis instances, because they only have a clip_path set to ax.patch. After this PR, axis instances are excluded from Axes.get_default_bbox_extra_artists(), and their extents are always calculated once.

I've also verified that, as before, ax.patch and ax.spines are still included as default "extra artists". Here's a test I used to determine the default clipping settings for various axes components:

import matplotlib.pyplot as plt
import cartopy.crs as ccrs
fig = plt.figure(figsize=(4, 2))
ax1 = fig.add_subplot(131)
ax2 = fig.add_subplot(132, projection='polar')
ax3 = fig.add_subplot(133, projection='hammer')
for ax in (ax1, ax2, ax3):
    print(ax)
    for a in (ax.patch, tuple(ax.spines.values())[0]):
        clip_box = a.get_clip_box()
        b1 = clip_box is not None and np.all(clip_box.extents == ax.bbox.extents)
        clip_path = a.get_clip_path()
        b2 = clip_path is not None and clip_path._patch is ax.patch
        print(a, '\nclip on?', a.get_clip_on(), '\naxes clip box?', b1, '\naxes clip path?', b2, '\nFULLY CLIPPED?', a._fully_clipped_to_axes())
    print()

Results:

AxesSubplot(0.125,0.11;0.227941x0.77)
Rectangle(xy=(0, 0), width=1, height=1, angle=0)
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False
Spine
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False

PolarAxesSubplot(0.398529,0.11;0.227941x0.77)
Wedge(center=(0.5, 0.5), r=0.5, theta1=0, theta2=360, width=None)
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False
Spine
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False

HammerAxesSubplot(0.672059,0.11;0.227941x0.77)
Circle(xy=(0.5, 0.5), radius=0.5)
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False
Spine
clip on? True
axes clip box? False
axes clip path? False
FULLY CLIPPED? False

@jklymak
Copy link
Member

jklymak commented Dec 17, 2021

This all looks good, but are we sure it has similar performance to before? https://github.com/matplotlib/mpl-bench has benchmarking mechanisms, in particular you could use the constrained_layout and tight_layout tests?

@lukelbd
Copy link
Contributor Author

lukelbd commented Dec 17, 2021

Here are some very quick benchmarks with 2-column figures of map projections (thanks for pointing me to mpl-benchmark -- don't have time to look into that now, but will later). Repeating the tests several times gives qualitatively similar results. If I have time later I can post more rigorous benchmarks.

In principle this fits as a matplotlib PR, but the speedup does seem to be cartopy-specific. I get essentially zero speedup for 1) non-rectilinear matplotlib axes and 2) by omitting the transform=ccrs.PlateCarree() line in the cartopy example. It seems that for high-res datasets, the limitation in cartopy is translating coordinates (triggered in both get_window_extent and draw). In those cases, this PR can deliver a roughly 2x speedup.

Cartopy JPG test

Test code:

import time
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import proplot as pplt
fig, axs = plt.subplots(ncols=2, figsize=(8, 4), subplot_kw={'projection': ccrs.Stereographic()})
t1 = time.time()
lon = np.linspace(0, 360, 3000)
lat = np.linspace(-60, 60, 500)
state = np.random.RandomState(51423)
data = state.rand(len(lat) - 1, len(lon) - 1)
for ax in axs:
    ax.coastlines()
    ax.pcolormesh(lon, lat, data, transform=ccrs.PlateCarree())
t2 = time.time()
print('Plot time:', t2 - t1)
fig.savefig('test.jpg', dpi=300, bbox_inches=None)
# fig.savefig('test.jpg', dpi=300, bbox_inches='tight')
t3 = time.time()
print('Save time:', t3 - t2)

Results with bbox_inches=None:

Plot time: 5.106376886367798
Save time: 4.24917197227478

Results with bbox_inches='tight', after this PR (save time ~20% slower):

Plot time: 5.4417970180511475
Save time: 5.905539035797119

Results with bbox_inches='tight', before this PR (save time ~120% slower):

Plot time: 5.133382081985474
Save time: 10.721824884414673

Cartopy PDF test

Test code:

import time
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import proplot as pplt
fig, axs = plt.subplots(ncols=2, figsize=(8, 4), subplot_kw={'projection': ccrs.Stereographic()})
t1 = time.time()
lon = np.linspace(0, 360, 100)
lat = np.linspace(-60, 60, 100)
state = np.random.RandomState(51423)
data = state.rand(len(lat) - 1, len(lon) - 1)
for ax in axs:
    ax.coastlines()
    ax.pcolormesh(lon, lat, data, transform=ccrs.PlateCarree())
t2 = time.time()
print('Plot time:', t2 - t1)
fig.savefig('test.pdf')
t3 = time.time()
print('Save time:', t3 - t2)

Results with bbox_inches=None:

Plot time: 0.28498291969299316
Save time: 2.3959097862243652

Results with bbox_inches='tight', after this PR (save time ~10% slower):

Plot time: 0.24462389945983887
Save time: 2.699641227722168

Results with bbox_inches='tight', before this PR (save time ~60% slower):

Plot time: 0.2555239200592041
Save time: 3.705993890762329

Matplotlib JPG test

Test code:

import time
import numpy as np
import matplotlib.pyplot as plt
import cartopy.crs as ccrs
import proplot as pplt
fig, axs = plt.subplots(ncols=2, figsize=(8, 4), subplot_kw={'projection': 'hammer'})
t1 = time.time()
lon = np.linspace(0, 360, 500)
lat = np.linspace(-60, 60, 500)
state = np.random.RandomState(51423)
data = state.rand(len(lat) - 1, len(lon) - 1)
for ax in axs:
    ax.pcolormesh(lon, lat, data)
t2 = time.time()
print('Plot time:', t2 - t1)
fig.savefig('test.jpg', dpi=300, bbox_inches='tight')
t3 = time.time()
print('Save time:', t3 - t2)

Results with bbox_inches=None:

0.04705810546875
Save time: 12.776510953903198

Results with bbox_inches='tight', after this PR (save time about the same):

Plot time: 0.05183219909667969
Save time: 11.54852819442749

Results with bbox_inches='tight', before this PR (save time about the same):

Plot time: 0.053359031677246094
Save time: 11.422415018081665

@lukelbd lukelbd changed the title Significantly improve tight layout performance for non-rectilinear axes Significantly improve tight layout performance for cartopy axes Dec 17, 2021
if (artist.get_visible() and artist.get_in_layout())]
# always include types that do not internally implement clipping
# to axes. may have clip_on set to True and clip_box equivalent
# to ax.bbox but then ignore these properties during draws.
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

That is .... a fun fact.

@dstansby dstansby merged commit af76ddf into matplotlib:main Dec 21, 2021
@lukelbd lukelbd deleted the tightbbox-speedup branch January 12, 2022 21:52
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

4 participants