Scale points for Voronoi plots

pybop/plot/voronoi.py

@@ -194,9 +194,41 @@ def interpolate_point(p, q, axis, boundary_val):
     return np.array([boundary_val, s]) if axis == 0 else np.array([s, boundary_val])
 
 
+def assign_nearest_value(x, y, f, xi, yi):
+    """
+    Computes an array of values given by the score of the nearest point.
+
+    Parameters
+    ----------
+    x : array-like
+        The x coordinates of points with known scores.
+    y : array-like
+        The y coordinates of points with known scores.
+    f : array-like
+        The score function at the given x and y coordinates.
+    xi : array-like
+        The x coordinates of grid points.
+    yi : array-like
+        The y coordinates of grid points.
+
+    Returns
+    -------
+        A numpy array containing the scores corresponding to the grid points.
+    """
+    # Create a KD-tree for efficient nearest neighbor search
+    tree = cKDTree(np.column_stack((x, y)))
+
+    # Find the nearest point for each grid point
+    _, indices = tree.query(np.column_stack((xi.ravel(), yi.ravel())))
+    zi = f[indices].reshape(xi.shape)
+
+    return zi
+
+
 def surface(
     optim: Union[BaseOptimiser, Optimisation],
     bounds=None,
+    normalise=True,
     resolution=250,
     show=True,
     **layout_kwargs,
@@ -211,6 +243,9 @@ def surface(
     bounds : numpy.ndarray, optional
         A 2x2 array specifying the [min, max] bounds for each parameter. If None, uses
         `cost.parameters.get_bounds_for_plotly`.
+    normalise : bool, optional
+        If True, the voronoi regions are computed using the Euclidean distance between
+        points normalised with respect to the bounds (default: True).
     resolution : int, optional
         Resolution of the plot. Default is 500.
     show : bool, optional
@@ -235,24 +270,51 @@ def surface(
         bounds if bounds is not None else [param.bounds for param in optim.parameters]
     )[:2]
 
-    # Compute regions
-    x, y, f, regions = _voronoi_regions(x_optim, y_optim, f, xlim, ylim)
-
     # Create a grid for plot
     xi = np.linspace(xlim[0], xlim[1], resolution)
     yi = np.linspace(ylim[0], ylim[1], resolution)
     xi, yi = np.meshgrid(xi, yi)
 
-    # Create a KD-tree for efficient nearest neighbor search
-    tree = cKDTree(np.column_stack((x, y)))
+    if normalise:
+        if xlim[1] <= xlim[0] or ylim[1] <= ylim[0]:
+            raise ValueError("Lower bounds must be strictly less than upper bounds.")
 
-    # Find the nearest point for each grid point
-    _, indices = tree.query(np.column_stack((xi.ravel(), yi.ravel())))
-    zi = f[indices].reshape(xi.shape)
+        # Normalise the region
+        x_range = xlim[1] - xlim[0]
+        y_range = ylim[1] - ylim[0]
+        norm_x_optim = (np.asarray(x_optim) - xlim[0]) / x_range
+        norm_y_optim = (np.asarray(y_optim) - ylim[0]) / y_range
+
+        # Compute regions
+        norm_x, norm_y, f, norm_regions = _voronoi_regions(
+            norm_x_optim, norm_y_optim, f, (0, 1), (0, 1)
+        )
+
+        # Create a normalised grid
+        norm_xi = np.linspace(0, 1, resolution)
+        norm_xi, norm_yi = np.meshgrid(norm_xi, norm_xi)
+
+        # Assign a value to each point in the grid
+        zi = assign_nearest_value(norm_x, norm_y, f, norm_xi, norm_yi)
+
+        # Rescale for plotting
+        regions = []
+        for norm_region in norm_regions:
+            region = np.empty_like(norm_region)
+            region[:, 0] = norm_region[:, 0] * x_range + xlim[0]
+            region[:, 1] = norm_region[:, 1] * y_range + ylim[0]
+            regions.append(region)
+
+    else:
+        # Compute regions
+        x, y, f, regions = _voronoi_regions(x_optim, y_optim, f, xlim, ylim)
+
+        # Assign a value to each point in the grid
+        zi = assign_nearest_value(x, y, f, xi, yi)
 
     # Calculate the size of each Voronoi region
     region_sizes = np.array([len(region) for region in regions])
-    normalized_sizes = (region_sizes - region_sizes.min()) / (
+    relative_sizes = (region_sizes - region_sizes.min()) / (
         region_sizes.max() - region_sizes.min()
     )
 
@@ -272,7 +334,7 @@ def surface(
     )
 
     # Add Voronoi edges
-    for region, size in zip(regions, normalized_sizes):
+    for region, size in zip(regions, relative_sizes):
         x_region = region[:, 0].tolist() + [region[0, 0]]
         y_region = region[:, 1].tolist() + [region[0, 1]]
 

tests/unit/test_optimisation.py

@@ -163,7 +163,7 @@ def check_bounds_handling(optim, expected_bounds, should_raise=False):
                 with pytest.raises(
                     ValueError, match="Either all bounds or no bounds must be set"
                 ):
-                    optim = optimiser(cost=cost, bounds=expected_bounds)
+                    optimiser(cost=cost, bounds=expected_bounds)
             else:
                 assert optim.bounds == expected_bounds
 
@@ -180,8 +180,9 @@ def check_multistart(optim, n_iters, multistarts):
         assert_log_update(optim)
         check_incorrect_update(optim)
 
-        multistart_optim = optimiser(cost, multistart=2, max_iterations=2)
-        check_multistart(multistart_optim, 2, 2)
+        # Test multistart
+        multistart_optim = optimiser(cost, multistart=2, max_iterations=6)
+        check_multistart(multistart_optim, 6, 2)
 
         if optimiser in [pybop.GradientDescent, pybop.Adam, pybop.NelderMead]:
             optim = optimiser(cost=cost, bounds=cost_bounds)

tests/unit/test_plots.py

@@ -147,11 +147,16 @@ def test_optim_plots(self, optim):
         pybop.plot.contour(optim, gradient=True, steps=5)
 
         # Plot voronoi
-        pybop.plot.surface(optim)
+        pybop.plot.surface(optim, normalise=False)
 
         # Plot voronoi w/ bounds
         pybop.plot.surface(optim, bounds=bounds)
 
+        with pytest.raises(
+            ValueError, match="Lower bounds must be strictly less than upper bounds."
+        ):
+            pybop.plot.surface(optim, bounds=[[0.5, 0.8], [0.7, 0.4]])
+
     @pytest.fixture
     def posterior_summary(self, fitting_problem):
         posterior = pybop.LogPosterior(