Creates a Voronoi diagram with cell polygons using scipy's Delaunay triangulation (scipy >= 0.9)

voronoi_polygons.py

"""
An adaptation of https://stackoverflow.com/a/15783581/60982
Using ideas from https://stackoverflow.com/a/9471601/60982
"""

import collections
import random

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
import numpy.random.mtrand
from scipy.spatial import Delaunay, KDTree


def voronoi(*, points):
    """
    Return a list of all edges of the voronoi diagram for given input points.
    """
    delaunay: Delaunay = Delaunay(points)
    triangles = delaunay.points[delaunay.vertices]

    circum_centers = np.array(
        [
            triangle_circumscribed_circle(points=triangle)
            for triangle in triangles
        ]
    )

    long_lines_endpoints = []
    line_indices = []
    for triangle_index, triangle in enumerate(triangles):
        circum_center = circum_centers[triangle_index]
        triangle_neighbors = delaunay.neighbors[triangle_index]
        for neighbor_index, neighbor in enumerate(triangle_neighbors):
            if neighbor != -1:
                line_indices.append(
                    (triangle_index, neighbor)
                )
                continue

            ps = \
                triangle[(neighbor_index + 1) % 3] - \
                triangle[(neighbor_index - 1) % 3]
            ps = np.array((ps[1], -ps[0]))

            middle = (
                triangle[(neighbor_index + 1) % 3] +
                triangle[(neighbor_index - 1) % 3]
             ) * 0.5
            di = middle - triangle[neighbor_index]

            ps /= np.linalg.norm(ps)
            di /= np.linalg.norm(di)

            if np.dot(di, ps) < 0.0:
                ps *= -1000.0
            else:
                ps *= 1000.0

            long_lines_endpoints.append(circum_center + ps)
            line_indices.append(
                (
                    triangle_index,
                    len(circum_centers) + len(long_lines_endpoints) - 1
                )
            )

    vertices = np.vstack(
        (circum_centers, long_lines_endpoints)
    )

    # filter out any duplicate lines
    line_indices_sorted = np.sort(line_indices)  # make (1,2) and (2,1) both (1,2)
    line_indices_tuples = [tuple(row) for row in line_indices_sorted]
    line_indices_unique = set(line_indices_tuples)
    line_indices_unique_sorted = sorted(line_indices_unique)

    return vertices, line_indices_unique_sorted


def triangle_circumscribed_circle(*, points):
    rows, columns = points.shape

    A_built_matrix = np.bmat(
        [
            [
                2 * np.dot(points, points.T),
                np.ones((rows, 1))
            ],
            [
                np.ones((1, rows)),
                np.zeros((1, 1))
            ]
        ]
    )

    b = np.hstack(
        (np.sum(points * points, axis=1), np.ones(1))
    )
    x = np.linalg.solve(A_built_matrix, b)
    bary_coordinates = x[:-1]

    circum_center = np.sum(
        points * np.tile(
            bary_coordinates.reshape(
                (points.shape[0], 1)
            ),
            (
                1,
                points.shape[1]
            )
        ),
        axis=0,
    )

    return circum_center


def voronoi_cell_lines(*, points, vertices, line_indices):
    """
    Return a mapping from a voronoi cell to its edges.

    :param points: shape (m,2)
    :param vertices: shape (n,2)
    :param line_indices: shape (o,2)
    :rtype: dict point index -> list of shape (n,2) with vertex indices
    """
    kd_tree: KDTree = KDTree(points)

    cells = collections.defaultdict(list)
    for line_index_0, line_index_1 in line_indices:
        vertex_0, vertex_1 = vertices[line_index_0], vertices[line_index_1]
        middle = (vertex_0 + vertex_1) / 2
        _, (point_0_index, point_1_index) = kd_tree.query(middle, 2)
        cells[point_0_index].append((line_index_0, line_index_1))
        cells[point_1_index].append((line_index_0, line_index_1))

    return cells


def form_polygons(*, cells):
    """
    Form polygons by storing vertex indices in (counter-)clockwise order.
    """
    polygons = {}
    for polygon_index, line_indices in cells.items():
        # get a directed graph which contains both directions and
        # arbitrarily follow one of both
        directed_graph = \
            line_indices + \
            [
                (line_index_1, line_index_0)
                for (line_index_0, line_index_1) in line_indices
            ]

        directed_graph_map = collections.defaultdict(list)
        for (line_index_0, line_index_1) in directed_graph:
            directed_graph_map[line_index_0].append(line_index_1)

        ordered_edges = []
        current_edge = directed_graph[0]
        while len(ordered_edges) < len(line_indices):
            line_index_0 = current_edge[1]
            line_index_1 = \
                directed_graph_map[line_index_0][0] \
                if directed_graph_map[line_index_0][0] != current_edge[0] \
                else directed_graph_map[line_index_0][1]

            next_edge = (line_index_0, line_index_1)
            ordered_edges.append(next_edge)
            current_edge = next_edge

        polygons[polygon_index] = [
            line_index_0 for (line_index_0, line_index_1) in ordered_edges
        ]

    return polygons


def close_outer_cells(*, cells):
    """
    Close the outer cells.
    """
    for polygon_index, line_indices in cells.items():
        dangling_lines = []
        for line_index_0, line_index_1 in line_indices:
            connections = list(
                filter(
                    lambda i12_:
                    (line_index_0, line_index_1) != (i12_[0], i12_[1])
                    and
                    (
                            line_index_0 == i12_[0] or
                            line_index_0 == i12_[1] or
                            line_index_1 == i12_[0] or
                            line_index_1 == i12_[1]
                    ),
                    line_indices
                )
            )
            assert 1 <= len(connections) <= 2
            if len(connections) == 1:
                dangling_lines.append((line_index_0, line_index_1))
        assert len(dangling_lines) in {0, 2}
        if len(dangling_lines) == 2:
            (i11, i12), (i21, i22) = dangling_lines

            # determine which line ends are unconnected
            connected = list(
                filter(
                    lambda i12_:
                        (i12_[0], i12_[1]) != (i11, i12) and
                        (i12_[0] == i11 or i12_[1] == i11),
                    line_indices
                )
            )
            i11_unconnected = not connected

            connected = list(
                filter(
                    lambda i12_:
                        (i12_[0], i12_[1]) != (i21, i22) and
                        (i12_[0] == i21 or i12_[1] == i21),
                    line_indices
                )
            )
            i21_unconnected = not connected

            start_index = i11 if i11_unconnected else i12
            end_index = i21 if i21_unconnected else i22

            cells[polygon_index].append((start_index, end_index))

    return cells


def voronoi_polygons_for_points(*, points):
    """
    Return the voronoi polygon for each input point.

    :param points: shape (n, 2)
    :rtype: list of n polygons where each polygon is an array of vertices
    """
    vertices, line_indices = voronoi(points=points)
    cells = voronoi_cell_lines(
        points=points,
        vertices=vertices,
        line_indices=line_indices,
    )
    cells = close_outer_cells(cells=cells)
    voronoi_polygons = form_polygons(cells=cells)

    voronoi_polygons = [
        vertices[np.asarray(voronoi_polygons[point_index])]
        for point_index in range(len(points))
    ]

    return voronoi_polygons


def main():
    random.seed(1)
    numpy.random.mtrand.seed(1)

    plt.figure(figsize=(4.5, 4.5))
    axes = plt.subplot(1, 1, 1)
    plt.axis(
        [-0.05, 1.05, -0.05, 1.05]
    )

    points = np.random.random((100, 2))
    polygons = voronoi_polygons_for_points(points=points)
    for polygon in polygons:
        polygon_object: matplotlib.patches.Polygon = \
            matplotlib.patches.Polygon(
                polygon,
                facecolor=[
                    0.3 + 0.7 * random.random(),  # red channel
                    0.3 + 0.7 * random.random(),  # green channel
                    0.3 + 0.7 * random.random(),  # blue channel
                    1,                            # alpha channel
                ],
            )
        axes.add_patch(polygon_object)

    x, y = points[:, 0], points[:, 1]
    plt.scatter(
        x=x,
        y=y,
        c=[(0, 0, 0, 1)] * len(points),  # Opaque black
        marker='.',                      # Dot
        zorder=2,                        # Foreground
    )

    plt.show()


if __name__ == '__main__':
    main()