Clip Voronoi regions in cvt_archive_heatmap (#356)

## Description

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

The current cvt_archive_heatmap relies on faraway points to draw the
polygons at the edges of the Voronoi diagram. Thus, if a user zooms out,
they will see that the heatmap actually has polygons going beyond it,
like below. In practice, we rarely notice this because the
cvt_archive_heatmap sets the axis limits to be the bounds of the
archive, so the outer regions are hidden.

![cvt
heatmap](https://github.com/icaros-usc/pyribs/assets/38124174/3b3486e0-cf55-4667-ae74-14f3b8a98d63)

However, there may be times when this behavior is undesirable, e.g.,
drawing two cvt archives next to each other. Thus, this PR makes it
possible to clip the bounds of the Voronoi diagram, like so:

![cvt
heatmap](https://github.com/icaros-usc/pyribs/assets/38124174/217ca661-87eb-43d0-b3f9-3b0ffbae9403)

We accomplish this by taking each polygon and computing its intersection
with the archive’s bounding box. Intersections are computed via the
shapely library: https://shapely.readthedocs.io/ Furthermore, by using
shapely, it is possible to clip to arbitrary shapes:

![cvt
heatmap](https://github.com/icaros-usc/pyribs/assets/38124174/1b316db7-8ad1-4943-8db4-8dbc8d6949be)

There are some tradeoffs to turning this feature on by default, namely:

1. It may be more “intuitive” to have infinite regions at the edge of
the archive, as that is how the archive works in practice — namely,
points outside the archive bounds are inserted into cells at the edges.
2. Computing the intersection of all the polygons with the bounding box
(or bounding shape) can be somewhat time-consuming. However, after
running the script in #355 with `clip=True`, I found that the runtime
for 10,000 cells was still ~2.3 sec (much slower than the 0.6 sec in
that PR, but still very fast).

For these reasons, we will keep this feature turned off by default via
the `clip` parameter to `cvt_archive_heatmap` .

Note that this PR introduces a dependency on shapely, and yet the
feature we use shapely for will be turned off by default since
`clip=False`. However, I believe this is acceptable because shapely is a
well-supported library that is also lightweight (~2MB), so it will not
be a burden to users.

## TODO

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

- [x]  Add shapely to deps
- [x]  Add shapely to pinned reqs
- [x]  Add shapely to intersphinx
- [x]  Implement clipping
- [x]  Add flag to turn clipping on and off
- [x]  Add test for heatmap with and without clipping
- [x]  Enable clipping to an arbitrary polygon
- [x] Add test for clipping to arbitrary polygon, with and without holes

## 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

HISTORY.md

@@ -10,6 +10,7 @@
 - Add visualization of QDax repertoires (#353)
 - Improve cvt_archive_heatmap flexibility (#354)
 - Speed up 2D cvt_archive_heatmap by order of magnitude (#355)
+- Clip Voronoi regions in cvt_archive_heatmap (#356)
 
 #### Documentation
 

docs/conf.py

@@ -256,4 +256,5 @@
     "scipy": ("https://docs.scipy.org/doc/scipy/", None),
     "sklearn": ("https://scikit-learn.org/stable/", None),
     "qdax": ("https://qdax.readthedocs.io/en/latest/", None),
+    "shapely": ("https://shapely.readthedocs.io/en/stable/", None),
 }

pinned_reqs/extras_visualize.txt

@@ -12,3 +12,4 @@ scikit-learn==1.3.0
 scipy==1.10.1
 threadpoolctl==3.0.0
 matplotlib==3.7.2
+shapely==2.0.1

ribs/visualize.py

@@ -23,6 +23,7 @@
 import matplotlib
 import matplotlib.pyplot as plt
 import numpy as np
+import shapely
 from matplotlib.cm import ScalarMappable
 from scipy.spatial import Voronoi  # pylint: disable=no-name-in-module
 
@@ -270,6 +271,7 @@ def cvt_archive_heatmap(archive,
                         vmax=None,
                         cbar="auto",
                         cbar_kwargs=None,
+                        clip=False,
                         plot_centroids=False,
                         plot_samples=False,
                         ms=1):
@@ -343,6 +345,15 @@ def cvt_archive_heatmap(archive,
             the colorbar on the specified Axes.
         cbar_kwargs (dict): Additional kwargs to pass to
             :func:`~matplotlib.pyplot.colorbar`.
+        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
+            polygons are hidden because we set the axis limits to be the archive
+            bounds. Passing `clip=True` will clip the heatmap such that these
+            "outer edge" polygons are within the archive bounds. An arbitrary
+            polygon can also be passed in to clip the heatmap to a custom shape.
+            See `#356 <https://github.com/icaros-usc/pyribs/pull/356>`_ for more
+            info.
         plot_centroids (bool): Whether to plot the cluster centroids.
         plot_samples (bool): Whether to plot the samples used when generating
             the clusters.
@@ -376,6 +387,11 @@ def cvt_archive_heatmap(archive,
         upper_bounds = np.flip(upper_bounds)
         centroids = np.flip(centroids, axis=1)
 
+    # If clip is on, make it default to an archive bounding box.
+    if clip and not isinstance(clip, shapely.Polygon):
+        clip = shapely.box(lower_bounds[0], lower_bounds[1], upper_bounds[0],
+                           upper_bounds[1])
+
     if plot_samples:
         samples = archive.samples
         if transpose_measures:
@@ -443,18 +459,36 @@ def cvt_archive_heatmap(archive,
         if -1 in region or len(region) == 0:
             continue
 
-        if objective is None:
-            # Transparent white (RGBA format) -- this ensures that if a figure
-            # is saved with a transparent background, the empty cells will also
-            # be transparent.
-            facecolors.append(np.array([1.0, 1.0, 1.0, 0.0]))
-            facecolor_cmap_mask.append(False)
+        if clip:
+            # Clip the cell vertices to the polygon. Clipping may cause some
+            # cells to split into two or more polygons, especially if the clip
+            # polygon has holes.
+            polygon = shapely.Polygon(vor.vertices[region])
+            intersection = polygon.intersection(clip)
+            if isinstance(intersection, shapely.MultiPolygon):
+                for polygon in intersection.geoms:
+                    vertices.append(polygon.exterior.coords)
+                n_splits = len(intersection.geoms)
+            else:
+                # The intersection is a single Polygon.
+                vertices.append(intersection.exterior.coords)
+                n_splits = 1
         else:
-            facecolors.append(np.empty(4))
-            facecolor_cmap_mask.append(True)
-            facecolor_objs.append(objective)
-
-        vertices.append(vor.vertices[region])
+            vertices.append(vor.vertices[region])
+            n_splits = 1
+
+        # Repeat values for each split.
+        for _ in range(n_splits):
+            if objective is None:
+                # Transparent white (RGBA format) -- this ensures that if a
+                # figure is saved with a transparent background, the empty cells
+                # will also be transparent.
+                facecolors.append(np.array([1.0, 1.0, 1.0, 0.0]))
+                facecolor_cmap_mask.append(False)
+            else:
+                facecolors.append(np.empty(4))
+                facecolor_cmap_mask.append(True)
+                facecolor_objs.append(objective)
 
     # Compute facecolors from the cmap. We first normalize the objectives and
     # clip them to [0, 1].

setup.py

@@ -3,10 +3,10 @@
 
 from setuptools import find_packages, setup
 
-with open("README.md") as readme_file:
+with open("README.md", encoding="utf-8") as readme_file:
     readme = readme_file.read()
 
-with open("HISTORY.md") as history_file:
+with open("HISTORY.md", encoding="utf-8") as history_file:
     history = history_file.read()
 
 # NOTE: Update pinned_reqs whenever install_requires or extras_require changes.
@@ -24,12 +24,16 @@
 ]
 
 extras_require = {
-    "visualize": ["matplotlib>=3.0.0",],
+    "visualize": [
+        "matplotlib>=3.0.0",
+        "shapely>=2.0.0",
+    ],
     # All dependencies except for dev. Don't worry if there are duplicate
     # dependencies, since setuptools automatically handles duplicates.
     "all": [
         ### visualize ###
         "matplotlib>=3.0.0",
+        "shapely>=2.0.0",
     ],
     "dev": [
         "pip>=20.3",

tests/visualize/visualize_test.py

@@ -16,6 +16,7 @@
 import matplotlib.pyplot as plt
 import numpy as np
 import pytest
+import shapely
 from matplotlib.testing.decorators import image_comparison
 
 from ribs.archives import CVTArchive, GridArchive, SlidingBoundariesArchive
@@ -673,6 +674,99 @@ def test_cvt_archive_heatmap_voronoi_style(cvt_archive):
     cvt_archive_heatmap(cvt_archive, lw=3.0, ec="grey")
 
 
+#
+# cvt_archive_heatmap clip tests
+#
+
+
+@image_comparison(baseline_images=["cvt_archive_heatmap_noclip"],
+                  remove_text=False,
+                  extensions=["png"],
+                  tol=CVT_IMAGE_TOLERANCE)
+def test_cvt_archive_heatmap_noclip(cvt_archive):
+    plt.figure(figsize=(8, 6))
+    cvt_archive_heatmap(cvt_archive, clip=False)
+    plt.xlim(-1.5, 1.5)
+    plt.ylim(-1.5, 1.5)
+
+
+@image_comparison(baseline_images=["cvt_archive_heatmap_clip"],
+                  remove_text=False,
+                  extensions=["png"],
+                  tol=CVT_IMAGE_TOLERANCE)
+def test_cvt_archive_heatmap_clip(cvt_archive):
+    plt.figure(figsize=(8, 6))
+    cvt_archive_heatmap(cvt_archive, clip=True)
+    plt.xlim(-1.5, 1.5)
+    plt.ylim(-1.5, 1.5)
+
+
+@image_comparison(baseline_images=["cvt_archive_heatmap_clip_polygon"],
+                  remove_text=False,
+                  extensions=["png"],
+                  tol=CVT_IMAGE_TOLERANCE)
+def test_cvt_archive_heatmap_clip_polygon(cvt_archive):
+    plt.figure(figsize=(8, 6))
+    cvt_archive_heatmap(
+        cvt_archive,
+        clip=shapely.Polygon(shell=np.array([
+            [-0.75, -0.375],
+            [-0.75, 0.375],
+            [-0.375, 0.75],
+            [0.375, 0.75],
+            [0.75, 0.375],
+            [0.75, -0.375],
+            [0.375, -0.75],
+            [-0.375, -0.75],
+        ]),),
+    )
+    plt.xlim(-1.5, 1.5)
+    plt.ylim(-1.5, 1.5)
+
+
+@image_comparison(
+    baseline_images=["cvt_archive_heatmap_clip_polygon_with_hole"],
+    remove_text=False,
+    extensions=["png"],
+    tol=CVT_IMAGE_TOLERANCE)
+def test_cvt_archive_heatmap_clip_polygon_with_hole(cvt_archive):
+    """This test will force some cells to be split in two."""
+    plt.figure(figsize=(8, 6))
+    cvt_archive_heatmap(
+        cvt_archive,
+        clip=shapely.Polygon(
+            shell=np.array([
+                [-0.75, -0.375],
+                [-0.75, 0.375],
+                [-0.375, 0.75],
+                [0.375, 0.75],
+                [0.75, 0.375],
+                [0.75, -0.375],
+                [0.375, -0.75],
+                [-0.375, -0.75],
+            ]),
+            holes=[
+                # Two holes that split some cells into two parts, and some cells
+                # into three parts.
+                np.array([
+                    [-0.5, 0],
+                    [-0.5, 0.05],
+                    [0.5, 0.05],
+                    [0.5, 0],
+                ]),
+                np.array([
+                    [-0.5, 0.125],
+                    [-0.5, 0.175],
+                    [0.5, 0.175],
+                    [0.5, 0.125],
+                ]),
+            ],
+        ),
+    )
+    plt.xlim(-1.5, 1.5)
+    plt.ylim(-1.5, 1.5)
+
+
 #
 # Parallel coordinate plot test
 #