Add rasterized arg for heatmaps (#359)

## Description

<!-- Provide a brief description of the PR's purpose here. -->

A useful feature in heatmap visualizations is to rasterize the heatmap
itself but not the surrounding text of the plot. This is useful when
saving to PDF for instance, as we want to rasterize the heatmap so that
the PDF does not have to render hundreds or thousands of cells. Instead,
it can just render a single image for the heatmap, while the surrounding
text is maintained in vector format.

Previously, we had to rasterize the entire plot, so even the text was
converted to pixels when it can easily be represented in vector format.

The API change is to add a `rasterized` argument to
`grid_archive_heatmap`, `cvt_archive_heatmap`, and
`sliding_boundareis_archive_heatmap`. This is similar to how Matplotlib
methods like
[pcolormesh](https://matplotlib.org/stable/api/_as_gen/matplotlib.pyplot.pcolormesh.html)
handle rasterization.

## TODO

<!-- Notable points that this PR has either accomplished or will
accomplish. -->

- [x] Add rasterized arg in grid_archive_heatmap — rasterized is passed
into pcolormesh, and pcm_kwargs is not allowed to contain rasterized
since that would result in a duplicate kwarg to pcolormesh
- [x] Add rasterized arg in cvt_archive_heatmap — rasterized is passed
into PolyCollection
- [x] Add rasterized arg in sliding_boundaries_archive_heatmap

## Questions

<!-- Any concerns or points of confusion? -->

## Status

- [x] I have read the guidelines in

[CONTRIBUTING.md](https://github.com/icaros-usc/pyribs/blob/master/CONTRIBUTING.md)
- [x] I have formatted my code using `yapf`
- [x] I have tested my code by running `pytest`
- [x] I have linted my code with `pylint`
- [x] I have added a one-line description of my change to the changelog
in
      `HISTORY.md`
- [x] This PR is ready to go

Author: btjanaka

.gitignore

@@ -126,6 +126,7 @@ results/
 
 # Allow matplotlib baseline images.
 !tests/**/*.png
+!tests/**/*.pdf
 
 # Editor, IDE, and OS settings
 .vscode/

HISTORY.md

@@ -14,6 +14,7 @@
 - **Backwards-incompatible:** Allow using kwargs for colorbar in
   parallel_axes_plot (#358)
   - Removes cbar_orientaton and cbar_pad args for parallel_axes_plot
+- Add `rasterized` arg for heatmaps (#359)
 
 #### Documentation
 

ribs/visualize/_cvt_archive_heatmap.py

@@ -25,6 +25,7 @@ def cvt_archive_heatmap(archive,
                         vmax=None,
                         cbar="auto",
                         cbar_kwargs=None,
+                        rasterized=False,
                         clip=False,
                         plot_centroids=False,
                         plot_samples=False,
@@ -99,6 +100,12 @@ def cvt_archive_heatmap(archive,
             the colorbar on the specified Axes.
         cbar_kwargs (dict): Additional kwargs to pass to
             :func:`~matplotlib.pyplot.colorbar`.
+        rasterized (bool): Whether to rasterize the heatmap. This can be useful
+            for saving to a vector format like PDF. Essentially, only the
+            heatmap will be converted to a raster graphic so that the archive
+            cells will not have to be individually rendered. Meanwhile, the
+            surrounding axes, particularly text labels, will remain in vector
+            format.
         clip (bool, shapely.Polygon): Clip the heatmap cells to a given polygon.
             By default, we draw the cells along the outer edges of the heatmap
             as polygons that extend beyond the archive bounds, but these
@@ -259,6 +266,7 @@ def cvt_archive_heatmap(archive,
             edgecolors=ec,
             facecolors=facecolors,
             linewidths=lw,
+            rasterized=rasterized,
         ))
 
     # Create a colorbar.

ribs/visualize/_grid_archive_heatmap.py

@@ -18,8 +18,9 @@ def grid_archive_heatmap(archive,
                          vmin=None,
                          vmax=None,
                          cbar="auto",
-                         pcm_kwargs=None,
-                         cbar_kwargs=None):
+                         cbar_kwargs=None,
+                         rasterized=False,
+                         pcm_kwargs=None):
     """Plots heatmap of a :class:`~ribs.archives.GridArchive` with 1D or 2D
     measure space.
 
@@ -82,10 +83,18 @@ def grid_archive_heatmap(archive,
             :class:`~matplotlib.axes.Axes`. If ``None``, then colorbar is not
             displayed. If this is an :class:`~matplotlib.axes.Axes`, displays
             the colorbar on the specified Axes.
-        pcm_kwargs (dict): Additional kwargs to pass to
-            :func:`~matplotlib.pyplot.pcolormesh`.
         cbar_kwargs (dict): Additional kwargs to pass to
             :func:`~matplotlib.pyplot.colorbar`.
+        rasterized (bool): Whether to rasterize the heatmap. This can be useful
+            for saving to a vector format like PDF. Essentially, only the
+            heatmap will be converted to a raster graphic so that the archive
+            cells will not have to be individually rendered. Meanwhile, the
+            surrounding axes, particularly text labels, will remain in vector
+            format. This is implemented by passing ``rasterized`` to
+            :func:`~matplotlib.pyplot.pcolormesh`, so passing ``"rasterized"``
+            in the ``pcm_kwargs`` below will raise an error.
+        pcm_kwargs (dict): Additional kwargs to pass to
+            :func:`~matplotlib.pyplot.pcolormesh`.
 
     Raises:
         ValueError: The archive's dimension must be 1D or 2D.
@@ -138,6 +147,7 @@ def grid_archive_heatmap(archive,
                           cmap=cmap,
                           vmin=vmin,
                           vmax=vmax,
+                          rasterized=rasterized,
                           **pcm_kwargs)
     elif archive.measure_dim == 2:
         # Retrieve data from archive.
@@ -178,6 +188,7 @@ def grid_archive_heatmap(archive,
                           cmap=cmap,
                           vmin=vmin,
                           vmax=vmax,
+                          rasterized=rasterized,
                           **pcm_kwargs)
 
     # Create color bar.

ribs/visualize/_sliding_boundaries_archive_heatmap.py

@@ -20,7 +20,8 @@ def sliding_boundaries_archive_heatmap(archive,
                                        vmin=None,
                                        vmax=None,
                                        cbar="auto",
-                                       cbar_kwargs=None):
+                                       cbar_kwargs=None,
+                                       rasterized=False):
     """Plots heatmap of a :class:`~ribs.archives.SlidingBoundariesArchive` with
     2D measure space.
 
@@ -90,6 +91,12 @@ def sliding_boundaries_archive_heatmap(archive,
             the colorbar on the specified Axes.
         cbar_kwargs (dict): Additional kwargs to pass to
             :func:`~matplotlib.pyplot.colorbar`.
+        rasterized (bool): Whether to rasterize the heatmap. This can be useful
+            for saving to a vector format like PDF. Essentially, only the
+            heatmap will be converted to a raster graphic so that the archive
+            cells will not have to be individually rendered. Meanwhile, the
+            surrounding axes, particularly text labels, will remain in vector
+            format.
     Raises:
         ValueError: The archive is not 2D.
     """
@@ -138,7 +145,8 @@ def sliding_boundaries_archive_heatmap(archive,
                    c=objective_batch,
                    cmap=cmap,
                    vmin=vmin,
-                   vmax=vmax)
+                   vmax=vmax,
+                   rasterized=rasterized)
     if boundary_lw > 0.0:
         # Careful with bounds here. Lines drawn along the x axis should extend
         # between the y bounds and vice versa -- see
@@ -147,12 +155,14 @@ def sliding_boundaries_archive_heatmap(archive,
                   lower_bounds[1],
                   upper_bounds[1],
                   color='k',
-                  linewidth=boundary_lw)
+                  linewidth=boundary_lw,
+                  rasterized=rasterized)
         ax.hlines(y_boundary,
                   lower_bounds[0],
                   upper_bounds[0],
                   color='k',
-                  linewidth=boundary_lw)
+                  linewidth=boundary_lw,
+                  rasterized=rasterized)
 
     # Create color bar.
     set_cbar(t, ax, cbar, cbar_kwargs)

tests/visualize/baseline_images/visualize_test/cvt_archive_heatmap_rasterized.pdf

@@ -674,6 +674,37 @@ def test_cvt_archive_heatmap_voronoi_style(cvt_archive):
     cvt_archive_heatmap(cvt_archive, lw=3.0, ec="grey")
 
 
+#
+# Rasterization tests
+#
+
+
+@image_comparison(baseline_images=["grid_archive_heatmap_rasterized"],
+                  remove_text=False,
+                  extensions=["pdf"])
+def test_grid_archive_rasterized(grid_archive):
+    plt.figure(figsize=(8, 6))
+    grid_archive_heatmap(grid_archive, rasterized=True)
+
+
+@image_comparison(baseline_images=["cvt_archive_heatmap_rasterized"],
+                  remove_text=False,
+                  extensions=["pdf"])
+def test_cvt_archive_rasterized(cvt_archive):
+    plt.figure(figsize=(8, 6))
+    cvt_archive_heatmap(cvt_archive, rasterized=True)
+
+
+@image_comparison(baseline_images=["sliding_boundaries_heatmap_rasterized"],
+                  remove_text=False,
+                  extensions=["pdf"])
+def test_sliding_boundaries_archive_rasterized(sliding_archive):
+    plt.figure(figsize=(8, 6))
+    sliding_boundaries_archive_heatmap(sliding_archive,
+                                       boundary_lw=1.0,
+                                       rasterized=True)
+
+
 #
 # cvt_archive_heatmap clip tests
 #