"""Provides cvt_archive_heatmap."""
from __future__ import annotations
from collections.abc import Sequence
from typing import Literal
import matplotlib.collections
import matplotlib.colors
import matplotlib.pyplot as plt
import numpy as np
import shapely
from matplotlib.axes import Axes
from matplotlib.cm import ScalarMappable
from matplotlib.typing import ColorType
from pandas import DataFrame
from scipy.spatial import Voronoi # pylint: disable=no-name-in-module
from ribs.archives import ArchiveDataFrame, CVTArchive
from ribs.visualize._utils import (
archive_heatmap_1d,
retrieve_cmap,
set_cbar,
validate_df,
validate_heatmap_visual_args,
)
[docs]
def cvt_archive_heatmap(
archive: CVTArchive,
ax: Axes | None = None,
*,
df: DataFrame | ArchiveDataFrame | None = None,
transpose_measures: bool = False,
cmap: str | Sequence[ColorType] | matplotlib.colors.Colormap = "magma",
aspect: Literal["auto", "equal"] | float | None = None,
lw: float = 0.5,
ec: ColorType = "black",
vmin: float | None = None,
vmax: float | None = None,
cbar: Literal["auto"] | None | Axes = "auto",
cbar_kwargs: dict | None = None,
rasterized: bool = False,
clip: bool | shapely.Polygon = False,
plot_centroids: bool = False,
ms: float = 1,
pcm_kwargs: dict | None = None,
plot_samples: None = None,
) -> None:
r"""Plots heatmap of a :class:`~ribs.archives.CVTArchive` with 1D or 2D measure space.
In the 2D case, we create a Voronoi diagram and shade in each cell with a color
corresponding to the objective value of that cell's elite. In the 1D case, we plot a
horizontal series of cells.
Depending on how many cells are in the archive, ``ms`` and ``lw`` may need to be
tuned. If there are too many cells, the Voronoi diagram and centroid markers will
make the entire image appear black. In that case, try turning off the centroids with
``plot_centroids=False`` or even removing the lines completely with ``lw=0``.
Examples:
.. plot::
:context: close-figs
Heatmap of a 2D CVTArchive
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> from ribs.archives import CVTArchive
>>> from ribs.visualize import cvt_archive_heatmap
>>> # Populate the archive with the negative sphere function.
>>> archive = CVTArchive(solution_dim=2,
... centroids=100, ranges=[(-1, 1), (-1, 1)])
>>> x = np.random.uniform(-1, 1, 10000)
>>> y = np.random.uniform(-1, 1, 10000)
>>> archive.add(solution=np.stack((x, y), axis=1),
... objective=-(x**2 + y**2),
... measures=np.stack((x, y), axis=1))
>>> # Plot a heatmap of the archive.
>>> plt.figure(figsize=(8, 6))
>>> cvt_archive_heatmap(archive)
>>> plt.title("Negative sphere function with 2D measures")
>>> plt.xlabel("x coords")
>>> plt.ylabel("y coords")
>>> plt.show()
.. plot::
:context: close-figs
Heatmap of a 1D CVTArchive
>>> import numpy as np
>>> import matplotlib.pyplot as plt
>>> from ribs.archives import CVTArchive
>>> from ribs.visualize import cvt_archive_heatmap
>>> # Populate the archive with the negative sphere function.
>>> archive = CVTArchive(solution_dim=2,
... centroids=20, ranges=[(-1, 1)])
>>> x = np.random.uniform(-1, 1, 1000)
>>> archive.add(solution=np.stack((x, x), axis=1),
... objective=-x**2,
... measures=x[:, None])
>>> # Plot a heatmap of the archive.
>>> plt.figure(figsize=(8, 6))
>>> cvt_archive_heatmap(archive)
>>> plt.title("Negative sphere function with 1D measures")
>>> plt.xlabel("x coords")
>>> plt.show()
Args:
archive: A 1D or 2D :class:`~ribs.archives.CVTArchive`.
ax: Axes on which to plot the heatmap. If ``None``, the current axis will be
used.
df: If provided, we will plot data from this argument instead of the data
currently in the archive. This data can be obtained by, for instance,
calling :meth:`ribs.archives.ArchiveBase.data` with ``return_type="pandas"``
and modifying the resulting :class:`~ribs.archives.ArchiveDataFrame`. Note
that, at a minimum, the data must contain columns for index, objective, and
measures. To display a custom metric, replace the "objective" column.
transpose_measures: By default, the first measure in the archive will appear
along the x-axis, and the second will be along the y-axis. To switch this
behavior (i.e. to transpose the axes), set this to ``True``. Does not apply
for 1D archives.
cmap: The colormap to use when plotting intensity. Either the name of a
:class:`~matplotlib.colors.Colormap`, a list of Matplotlib color
specifications (e.g., an :math:`N \times 3` or :math:`N \times 4` array --
see :class:`~matplotlib.colors.ListedColormap`), or a
:class:`~matplotlib.colors.Colormap` object.
aspect: The aspect ratio of the heatmap (i.e. height/width). Defaults to
``'auto'`` for 2D and ``0.5`` for 1D. ``'equal'`` is the same as
``aspect=1``. See :meth:`matplotlib.axes.Axes.set_aspect` for more info.
lw: Line width when plotting the Voronoi diagram.
ec: Edge color of the cells in the Voronoi diagram. See `here
<https://matplotlib.org/stable/tutorials/colors/colors.html>`_ for more info
on specifying colors in Matplotlib.
vmin: Minimum objective value to use in the plot. If ``None``, the minimum
objective value in the archive is used.
vmax: Maximum objective value to use in the plot. If ``None``, the maximum
objective value in the archive is used.
cbar: By default, this is set to ``'auto'`` which displays the colorbar on the
archive's current :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.
cbar_kwargs: Additional kwargs to pass to :func:`~matplotlib.pyplot.colorbar`.
rasterized: 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: 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
:pr:`356` for more info. Only applies to 2D archives.
plot_centroids: Whether to plot the cluster centroids.
ms: Marker size for centroids.
pcm_kwargs: Additional kwargs to pass to :func:`~matplotlib.pyplot.pcolormesh`.
Only applicable to 1D heatmaps. linewidth and edgecolor are set with the
``lw`` and ``ec`` args.
plot_samples: DEPRECATED.
Raises:
ValueError: The archive's measure dimension must be 1D or 2D.
"""
if plot_samples is not None:
raise ValueError(
"`plot_samples` is deprecated in pyribs 0.9.0, "
"as CVTArchive no longer stores samples."
)
validate_heatmap_visual_args(
aspect,
cbar,
archive.measure_dim,
[1, 2],
"Heatmap can only be plotted for a 1D or 2D CVTArchive",
)
if aspect is None:
# Handles default aspects for different dims.
if archive.measure_dim == 1:
aspect = 0.5
else:
aspect = "auto"
# Try getting the colormap early in case it fails.
cmap = retrieve_cmap(cmap)
# Retrieve archive data.
if df is None:
index_batch = archive.data("index")
objective_batch = archive.data("objective")
else:
df = validate_df(df)
index_batch = df["index"]
objective_batch = df["objective"]
if archive.measure_dim == 1:
# Read in pcm kwargs -- linewidth and edgecolor are overwritten by our args.
pcm_kwargs = {} if pcm_kwargs is None else pcm_kwargs.copy()
pcm_kwargs["linewidth"] = pcm_kwargs["lw"] = lw
pcm_kwargs["edgecolor"] = pcm_kwargs["ec"] = ec
# Sort centroids so they line up left-to-right along the x-axis.
centroids_1d = archive.centroids.squeeze()
centroid_sort_idx = np.argsort(centroids_1d)
sorted_centroids_1d = centroids_1d[centroid_sort_idx]
cell_boundaries = np.concatenate(
(
# Concatenate lower bound.
[archive.lower_bounds[0]],
# The boundaries can be found by taking the midpoints between the
# centroids.
(sorted_centroids_1d[:-1] + sorted_centroids_1d[1:]) / 2.0,
# Concatenate upper bound.
[archive.upper_bounds[0]],
)
)
# centroid_sort_idx tells us which index to place the centroid at such that it
# is sorted, i.e., it maps from the indices in the centroid array to the cell
# indices. This means that if you index with it, e.g., arr[centroid_sort_idx],
# you get a sorted array.
#
# For computing cell_objectives, we need to know the inverse mapping, i.e., the
# mapping from cell indices to centroid indices. This way, when we index with
# it, we get the original order of centroids. This original order then matches
# with the objectives in objective_batch.
inv_idx = np.zeros_like(centroid_sort_idx, dtype=np.int32)
for i, x in enumerate(centroid_sort_idx):
inv_idx[x] = i
# We only want inverse indexes that are actually used in the archive.
selected_inv_idx = inv_idx[index_batch]
cell_objectives = np.full(archive.cells, np.nan)
cell_objectives[selected_inv_idx] = objective_batch
ax = archive_heatmap_1d(
archive,
cell_boundaries=cell_boundaries,
cell_objectives=cell_objectives,
ax=ax,
cmap=cmap,
aspect=aspect,
vmin=vmin,
vmax=vmax,
cbar=cbar,
cbar_kwargs=cbar_kwargs,
rasterized=rasterized,
pcm_kwargs=pcm_kwargs,
)
# Centroids are plotted at y=0.5 so that they appear along the center of the
# diagram.
if plot_centroids:
ax.plot(
archive.centroids[:, 0],
np.full(len(archive.centroids), 0.5),
"o",
c="black",
ms=ms,
)
elif archive.measure_dim == 2:
# Retrieve data from archive.
lower_bounds = archive.lower_bounds
upper_bounds = archive.upper_bounds
centroids = archive.centroids
if transpose_measures:
lower_bounds = np.flip(lower_bounds)
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]
)
# Retrieve and initialize the axis.
ax = plt.gca() if ax is None else ax
ax.set_xlim(lower_bounds[0], upper_bounds[0])
ax.set_ylim(lower_bounds[1], upper_bounds[1])
ax.set_aspect(aspect)
# Add faraway points so that the edge regions of the Voronoi diagram are filled
# in. Refer here for more info:
# https://stackoverflow.com/questions/20515554/colorize-voronoi-diagram
interval = upper_bounds - lower_bounds
scale = 1000
faraway_pts = [
upper_bounds + interval * scale, # Far upper right.
upper_bounds + interval * [-1, 1] * scale, # Far upper left.
lower_bounds + interval * [-1, -1] * scale, # Far bottom left.
lower_bounds + interval * [1, -1] * scale, # Far bottom right.
]
vor = Voronoi(np.append(centroids, faraway_pts, axis=0))
# Calculate objective value for each region. `vor.point_region` contains the
# region index of each point.
region_obj = [None] * len(vor.regions)
min_obj, max_obj = np.inf, -np.inf
pt_to_obj = dict(zip(index_batch, objective_batch, strict=True))
for pt_idx, region_idx in enumerate(
vor.point_region[:-4]
): # Exclude faraway_pts.
if region_idx != -1 and pt_idx in pt_to_obj:
obj = pt_to_obj[pt_idx]
min_obj = min(min_obj, obj)
max_obj = max(max_obj, obj)
region_obj[region_idx] = obj
# Override objective value range.
min_obj = min_obj if vmin is None else vmin
max_obj = max_obj if vmax is None else vmax
# If the min and max are the same, we set a sensible default range.
if min_obj == max_obj:
min_obj, max_obj = min_obj - 0.01, max_obj + 0.01
# Vertices of all cells.
vertices = []
# The facecolor of each cell. Shape (n_regions, 4) for RGBA format, but we do
# not know n_regions in advance.
facecolors = []
# Boolean array indicating which of the facecolors needs to be computed with the
# cmap. The other colors correspond to empty cells. Shape (n_regions,)
facecolor_cmap_mask = []
# The objective corresponding to the regions which must be passed through the
# cmap. Shape (sum(facecolor_cmap_mask),)
facecolor_objs = []
# Cycle through the regions to set up polygon vertices and facecolors.
for region, objective in zip(vor.regions, region_obj, strict=True):
# Checking for -1 is O(n), but n is typically small.
#
# We check length since the first region is an empty list by default:
# https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.Voronoi.html
if -1 in region or len(region) == 0:
continue
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:
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].
normalized_objs = np.clip(
(np.asarray(facecolor_objs) - min_obj) / (max_obj - min_obj), 0.0, 1.0
)
facecolors = np.asarray(facecolors)
facecolors[facecolor_cmap_mask] = cmap(normalized_objs)
# Plot the collection on the axes. Note that this is faster than plotting each
# polygon individually with ax.fill().
ax.add_collection(
matplotlib.collections.PolyCollection(
vertices,
edgecolors=ec,
facecolors=facecolors,
linewidths=lw,
rasterized=rasterized,
)
)
# Create a colorbar.
mappable = ScalarMappable(cmap=cmap)
mappable.set_clim(min_obj, max_obj)
# Plot the sample points and centroids.
if plot_centroids:
ax.plot(centroids[:, 0], centroids[:, 1], "o", c="black", ms=ms)
# Create color bar.
set_cbar(mappable, ax, cbar, cbar_kwargs)