diff --git a/pybop/plot/voronoi.py b/pybop/plot/voronoi.py index 4dad41223..9283251dc 100644 --- a/pybop/plot/voronoi.py +++ b/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]] diff --git a/tests/unit/test_optimisation.py b/tests/unit/test_optimisation.py index b3c094d0c..ba22828bc 100644 --- a/tests/unit/test_optimisation.py +++ b/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) diff --git a/tests/unit/test_plots.py b/tests/unit/test_plots.py index 733b81d42..869f89dfb 100644 --- a/tests/unit/test_plots.py +++ b/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(