fimama.voronoi

  1# from icecream import ic
  2
  3import numpy as np
  4from matplotlib.axes import Axes
  5from matplotlib.collections import PatchCollection
  6from matplotlib.colors import LinearSegmentedColormap
  7from matplotlib.patches import Polygon
  8
  9from scipy.spatial import Voronoi, voronoi_plot_2d
 10
 11from .configuration import VoronoiConfiguration
 12
 13
 14def _generate_voronoi_grid(heightmap: np.ndarray | None = None):
 15    np.random.seed(1234)
 16    if heightmap is not None:
 17        grid_shape = (heightmap.shape[0] - 1, heightmap.shape[1] - 1, )
 18    else:
 19        grid_shape = (200, 200)
 20
 21    # make up data points
 22    base_points = np.mgrid[
 23        0:(grid_shape[0]):(grid_shape[0]+1)*1j,
 24        0:(grid_shape[1]):(grid_shape[1]+1)*1j
 25    ]
 26    base_points = base_points.reshape(2, -1).T
 27    points = base_points + np.random.random_sample(base_points.shape) - .5
 28
 29    # TODO 0.6: better dummy points or better yet exclude outermost ring of
 30    # points while keeping them saved as the basis of extending the map.
 31    # add 4 distant dummy points
 32    dummy_points = []
 33    dummy_points = [[999, 999], [-999, 999], [999, -999], [-999, -999]]
 34    base_points = np.append(base_points, dummy_points, axis=0)
 35    points = np.append(points, dummy_points, axis=0)
 36
 37    # compute Voronoi tessellation
 38    voronoi = Voronoi(points)
 39    return voronoi, points, dummy_points, grid_shape
 40
 41
 42def voronoi_map(
 43    axes: Axes,
 44    fig,
 45    heightmap: np.ndarray = None,
 46    colormap: LinearSegmentedColormap = None,
 47    config: VoronoiConfiguration = None,
 48):
 49    """
 50    Plot a voronoi heightmap.
 51
 52    Parameters
 53    ----------
 54    axes : Axes
 55        axes to plot on
 56    fig : _type_
 57        Figure we are plotting into. Used for adding a colorbar.
 58    heightmap : np.ndarray, optional
 59        The heightmap, by default None
 60    colormap : LinearSegmentedColormap, optional
 61        _description_, by default None
 62    config : VoronoiConfiguration, optional
 63        _description_, by default None
 64    """
 65    if colormap is None:
 66        colormap = "terrain"
 67    if config is None:
 68        config = VoronoiConfiguration()
 69
 70    voronoi, points, dummy_points, grid_shape = _generate_voronoi_grid(
 71        heightmap=heightmap)
 72
 73    if config.plot_voronoi_grid:
 74        voronoi_plot_2d(
 75            vor=voronoi,
 76            ax=axes,
 77            show_vertices=config.show_vertices,
 78            show_points=config.show_points
 79        )
 80
 81    # map regions to points
 82    regions_to_points = np.argsort(voronoi.point_region)
 83
 84    # Polygon patches for the voronoi regions
 85    polygons = []
 86    polygon_index = 0
 87    heights = np.zeros(len(points) - len(dummy_points))
 88    heightmap = heightmap.flatten()
 89    for i, region in enumerate(voronoi.regions):
 90        if -1 not in region and len(region) > 0:
 91            vertices = [voronoi.vertices[i] for i in region]
 92            polygon = Polygon(vertices, closed=True)
 93            polygons.append(polygon)
 94            # center = np.mean(polygon.get_xy(), axis=0)
 95            # axes.text(center[0], center[1], f"{i}")
 96            # index, = np.where(voronoi.point_region == i)
 97            # index = index[0]
 98            # ic(index, regions_to_points[i-1])
 99            heights[polygon_index] = heightmap[regions_to_points[i-1]]
100            polygon_index += 1
101            # axes.text(
102            #     point[0], point[1],
103            #     f"{j}, {i}",
104            #     horizontalalignment='center',
105            #     verticalalignment='center'
106            # )
107    # ic(points_in_regions[1:])
108    # ic(region_to_polygon[1:])
109    # ic(voronoi.point_region)
110    # ic(regions_to_points)
111    poly_collection = PatchCollection(patches=polygons, cmap=colormap)
112
113    poly_collection.set_array(heights)
114    axes.add_collection(poly_collection)
115    fig.colorbar(poly_collection, ax=axes)
116
117    # fix the range of axes
118    axes.set_xlim([-1, grid_shape[0]+1])
119    axes.set_ylim([-1, grid_shape[1]+1])