Unify logic for computing vmin and vmax in ribs.visualize

HISTORY.md

@@ -10,11 +10,12 @@
   archive.data() ({pr}`721`)
 - Add `archive_ecdf` to visualize ECDF and ECCDF of archives ({pr}`719`)
 - Add `archive_histogram` to visualize objective values in archives ({pr}`714`,
-  {pr}`723`)
+  {pr}`724`)
 - Implement `novelty_threshold` decay in `ProximityArchive` ({pr}`709`)
 
 #### Improvements
 
+- Unify logic for computing vmin and vmax in ribs.visualize ({pr}`725`)
 - Remove torch from dev-array-api extra since it is in the all extra ({pr}`715`)
 - Regenerate ribs.visualize baseline images due to recent update with matplotlib
   ({pr}`711`, {pr}`713`)

ribs/visualize/_archive_histogram.py

@@ -6,13 +6,12 @@
 
 import matplotlib.colors
 import matplotlib.pyplot as plt
-import numpy as np
 from matplotlib.axes import Axes
 from matplotlib.typing import ColorType
 from pandas import DataFrame
 
 from ribs.archives import ArchiveBase, ArchiveDataFrame
-from ribs.visualize._utils import retrieve_cmap, validate_df
+from ribs.visualize._utils import compute_vmin_vmax, retrieve_cmap, validate_df
 
 
 def archive_histogram(
@@ -182,16 +181,7 @@ def archive_histogram(
         df = validate_df(df)
         objectives = df["objective"]
 
-    # Compute vmin and vmax.
-    if len(objectives) == 0:
-        # Sensible defaults when there is no elite in the archive. The colorbar for the
-        # heatmap functions usually defaults to -0.1 and 0.1 when no objectives exist in
-        # the archive.
-        vmin = -0.1 if vmin is None else vmin
-        vmax = 0.1 if vmax is None else vmax
-    else:
-        vmin = np.min(objectives) if vmin is None else vmin
-        vmax = np.max(objectives) if vmax is None else vmax
+    vmin, vmax = compute_vmin_vmax(vmin, vmax, objectives)
 
     # Initialize axis.
     ax = plt.gca() if ax is None else ax

ribs/visualize/_cvt_archive_3d_plot.py

@@ -18,6 +18,7 @@
 
 from ribs.archives import ArchiveDataFrame, CVTArchive
 from ribs.visualize._utils import (
+    compute_vmin_vmax,
     retrieve_cmap,
     set_cbar,
     validate_df,
@@ -260,22 +261,7 @@ def cvt_archive_3d_plot(
         upper_bounds = upper_bounds[measure_order]
         centroids = centroids[:, measure_order]
 
-    # Compute objective value range.
-    if vmin is None:
-        # Defaulting to -inf (and inf for max_obj) allows the computations after this to
-        # proceed smoothly.
-        min_obj = np.min(objective_batch) if len(objective_batch) > 0 else -np.inf
-    else:
-        min_obj = vmin
-
-    if vmax is None:
-        max_obj = np.max(objective_batch) if len(objective_batch) > 0 else np.inf
-    else:
-        max_obj = 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
+    vmin, vmax = compute_vmin_vmax(vmin, vmax, objective_batch)
 
     # Default ax behavior.
     if ax is None:
@@ -368,9 +354,7 @@ def cvt_archive_3d_plot(
     objs = np.asarray(objs)
     cmap_idx = ~np.isnan(objs)
     cmap_objs = objs[cmap_idx]
-    normalized_objs = np.clip(
-        (np.asarray(cmap_objs) - min_obj) / (max_obj - min_obj), 0.0, 1.0
-    )
+    normalized_objs = np.clip((np.asarray(cmap_objs) - vmin) / (vmax - vmin), 0.0, 1.0)
 
     # Create an array of facecolors in RGBA format that defaults to transparent white.
     facecolors = np.full((len(objs), 4), [1.0, 1.0, 1.0, 0.0])
@@ -397,8 +381,8 @@ def cvt_archive_3d_plot(
             s=elite_ms,
             c=objective_batch,
             cmap=cmap,
-            vmin=min_obj,
-            vmax=max_obj,
+            vmin=vmin,
+            vmax=vmax,
             lw=0.0,
             alpha=elite_alpha,
         )
@@ -409,5 +393,5 @@ def cvt_archive_3d_plot(
 
     # Create color bar.
     mappable = ScalarMappable(cmap=cmap)
-    mappable.set_clim(min_obj, max_obj)
+    mappable.set_clim(vmin, vmax)
     set_cbar(mappable, ax, cbar, cbar_kwargs)

ribs/visualize/_cvt_archive_heatmap.py

@@ -19,6 +19,7 @@
 from ribs.archives import ArchiveDataFrame, CVTArchive
 from ribs.visualize._utils import (
     archive_heatmap_1d,
+    compute_vmin_vmax,
     retrieve_cmap,
     set_cbar,
     validate_df,
@@ -307,24 +308,14 @@ def cvt_archive_heatmap(
         # 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
+                region_obj[region_idx] = pt_to_obj[pt_idx]
 
-        # 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
+        vmin, vmax = compute_vmin_vmax(vmin, vmax, objective_batch)
 
         # Vertices of all cells.
         vertices = []
@@ -381,7 +372,7 @@ def cvt_archive_heatmap(
         # 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
+            (np.asarray(facecolor_objs) - vmin) / (vmax - vmin), 0.0, 1.0
         )
         facecolors = np.asarray(facecolors)
         facecolors[facecolor_cmap_mask] = cmap(normalized_objs)
@@ -400,7 +391,7 @@ def cvt_archive_heatmap(
 
         # Create a colorbar.
         mappable = ScalarMappable(cmap=cmap)
-        mappable.set_clim(min_obj, max_obj)
+        mappable.set_clim(vmin, vmax)
 
         # Plot the sample points and centroids.
         if plot_centroids:

ribs/visualize/_grid_archive_heatmap.py

@@ -15,6 +15,7 @@
 from ribs.archives import ArchiveDataFrame, GridArchive
 from ribs.visualize._utils import (
     archive_heatmap_1d,
+    compute_vmin_vmax,
     retrieve_cmap,
     set_cbar,
     validate_df,
@@ -218,16 +219,7 @@ def grid_archive_heatmap(
 
         # Create the plot.
         pcm_kwargs = {} if pcm_kwargs is None else pcm_kwargs
-        vmin = (
-            np.min(objective_batch)
-            if vmin is None and len(objective_batch) > 0
-            else vmin
-        )
-        vmax = (
-            np.max(objective_batch)
-            if vmax is None and len(objective_batch) > 0
-            else vmax
-        )
+        vmin, vmax = compute_vmin_vmax(vmin, vmax, objective_batch)
         t = ax.pcolormesh(
             x_bounds,
             y_bounds,

ribs/visualize/_parallel_axes_plot.py

@@ -20,7 +20,12 @@
     ProximityArchive,
     SlidingBoundariesArchive,
 )
-from ribs.visualize._utils import retrieve_cmap, set_cbar, validate_df
+from ribs.visualize._utils import (
+    compute_vmin_vmax,
+    retrieve_cmap,
+    set_cbar,
+    validate_df,
+)
 
 
 def parallel_axes_plot(
@@ -200,8 +205,7 @@ def parallel_axes_plot(
     upper_bounds = upper_bounds[cols]
 
     host_ax = plt.gca() if ax is None else ax  # Try to get current axis.
-    vmin = df["objective"].min() if vmin is None else vmin
-    vmax = df["objective"].max() if vmax is None else vmax
+    vmin, vmax = compute_vmin_vmax(vmin, vmax, df["objective"])
     norm = matplotlib.colors.Normalize(vmin=vmin, vmax=vmax, clip=True)
     if sort_archive:
         df = df.sort_values("objective")

ribs/visualize/_proximity_archive_plot.py

@@ -14,6 +14,7 @@
 
 from ribs.archives import ArchiveDataFrame, ProximityArchive
 from ribs.visualize._utils import (
+    compute_vmin_vmax,
     retrieve_cmap,
     set_cbar,
     validate_df,
@@ -201,12 +202,7 @@ def proximity_archive_plot(
     ax.set_aspect(aspect)
 
     # Create the plot.
-    vmin = (
-        np.min(objective_batch) if vmin is None and len(objective_batch) > 0 else vmin
-    )
-    vmax = (
-        np.max(objective_batch) if vmax is None and len(objective_batch) > 0 else vmax
-    )
+    vmin, vmax = compute_vmin_vmax(vmin, vmax, objective_batch)
     t = ax.scatter(
         x,
         y,

ribs/visualize/_sliding_boundaries_archive_heatmap.py

@@ -14,6 +14,7 @@
 
 from ribs.archives import ArchiveDataFrame, SlidingBoundariesArchive
 from ribs.visualize._utils import (
+    compute_vmin_vmax,
     retrieve_cmap,
     set_cbar,
     validate_df,
@@ -161,12 +162,7 @@ def sliding_boundaries_archive_heatmap(
     ax.set_aspect(aspect)
 
     # Create the plot.
-    vmin = (
-        np.min(objective_batch) if vmin is None and len(objective_batch) > 0 else vmin
-    )
-    vmax = (
-        np.max(objective_batch) if vmax is None and len(objective_batch) > 0 else vmax
-    )
+    vmin, vmax = compute_vmin_vmax(vmin, vmax, objective_batch)
     t = ax.scatter(
         x,
         y,

ribs/visualize/_utils.py

@@ -87,6 +87,99 @@ def set_cbar(
         cbar.figure.colorbar(t, ax=cbar, **cbar_kwargs)
 
 
+# Use this offset to prevent vmin and vmax from being too close to each other.
+OBJECTIVE_OFFSET = 0.1
+
+
+def compute_vmin_vmax(  # pylint: disable = too-many-return-statements
+    vmin: float | None,
+    vmax: float | None,
+    objectives: np.ndarray,
+) -> tuple[float, float]:
+    """Computes vmin and vmax based on the user's args and objectives in the archive.
+
+    Args:
+        vmin: User-supplied value for vmin.
+        vmax: User-supplied value for vmax.
+        objectives: Array of objective values.
+
+    Returns:
+        Tuple containing the new vmin and vmax.
+
+    Raises:
+        ValueError: vmin and vmax were both passed in, but vmin is greater than vmax (it
+            must be less than or equal to vmax).
+    """
+    has_objectives = len(objectives) > 0
+
+    # Cache min and max objectives.
+    if has_objectives:
+        min_obj = np.min(objectives)
+        max_obj = np.max(objectives)
+    else:
+        min_obj = None
+        max_obj = None
+
+    # Determine new_vmin and new_vmax. This depends on three conditions:
+    # 1. What is the value of vmin?
+    # 2. What is the value of vmax? (This is combined with (1) in the branches.)
+    # 3. Are there any objectives present?
+    # The guiding principle is that we should always return reasonable and valid values
+    # for vmin and vmax. The values are valid if vmin < vmax (strictly less than; equal
+    # is not okay).
+    if vmin is None and vmax is None:
+        if has_objectives:
+            # Neither vmin nor vmax were passed in, and there are objectives in the
+            # archive, so we use min_obj and max_obj.
+            if min_obj == max_obj:
+                # We use strict equality here rather than isclose. isclose checks for a tiny
+                # difference, and we are okay with tiny differences.
+                #
+                # Move the objectives apart since they are equal.
+                return (min_obj - OBJECTIVE_OFFSET, max_obj + OBJECTIVE_OFFSET)
+            else:
+                # Here, the objectives are far enough away, so set them directly.
+                return (min_obj, max_obj)
+        else:
+            # Neither vmin nor vmax were passed in, and there are no objectives, so we
+            # can choose any default value.
+            return (-OBJECTIVE_OFFSET, OBJECTIVE_OFFSET)
+    elif vmin is not None and vmax is None:
+        # vmin is passed in, but we need to decide how to set vmax.
+        if has_objectives:
+            if vmin < max_obj:
+                # Ideally, we can just use max_obj as vmax.
+                return (vmin, max_obj)
+            else:
+                # However, if vmin >= max_obj, we choose our own default.
+                return (vmin, vmin + 2.0 * OBJECTIVE_OFFSET)
+        else:
+            # If there are no objectives, we choose our own default.
+            return (vmin, vmin + 2.0 * OBJECTIVE_OFFSET)
+    elif vmin is None and vmax is not None:
+        # vmax is passed in, but we need to decide how to set vmin.
+        if has_objectives:
+            if min_obj < vmax:
+                # Ideally, we can just use min_obj as vmin.
+                return (min_obj, vmax)
+            else:
+                # However, if min_obj is >= vmax, we choose our own default.
+                return (vmax - 2.0 * OBJECTIVE_OFFSET, vmax)
+        else:
+            # If there are no objectives, we choose our own default.
+            return (vmax - 2.0 * OBJECTIVE_OFFSET, vmax)
+    else:  # vmin is not None and vmax is not None
+        # Both vmin and vmax are passed in. Take them as is, subject to verification.
+        if vmax < vmin:
+            raise ValueError(
+                f"vmax ({vmax}) must be greater than or equal to vmin ({vmin})"
+            )
+        elif vmin == vmax:
+            # If they're equal, set a sensible default range.
+            return (vmin - OBJECTIVE_OFFSET, vmax + OBJECTIVE_OFFSET)
+        return (vmin, vmax)
+
+
 def archive_heatmap_1d(
     archive: GridArchive | CVTArchive,
     *,
@@ -142,8 +235,9 @@ def archive_heatmap_1d(
 
     # Create the plot.
     pcm_kwargs = {} if pcm_kwargs is None else pcm_kwargs
-    vmin = np.nanmin(cell_objectives) if vmin is None and not archive.empty else vmin
-    vmax = np.nanmax(cell_objectives) if vmax is None and not archive.empty else vmax
+    vmin, vmax = compute_vmin_vmax(
+        vmin, vmax, cell_objectives[~np.isnan(cell_objectives)]
+    )
     t = ax.pcolormesh(
         cell_boundaries,
         # y-bounds; needs a sensible default so that aspect ratio is consistent.