Voronoi diagram with Jump Flooding algorithm: performance issue

Viewed 210

I am trying to implement the Jump Flooding algorithm to draw a Voronoi diagram from an array of points. The code below works correctly but is extremely slow for large N. The reason it is slow is that there is a loop that adds new information to an array being iterated. This array lists neighbor points. If I do not include this appending, the Voronoi diagram is nonsensical, with most points not being reached by the algorithm. I don't understand why a simple sweep through neighbors identified in the original neighbor search doesn't hit all the points on the grid. Can you suggest what is wrong with my code and if there is a performant alternative to appending all found neighbors to a gigantic array? Or is this perhaps simply a python problem and rewriting in C/C++ etc. will fix the performance issue?

import matplotlib.pyplot as plt
import numpy as np

from time import time

# array = np.zeros((5, 5))
# originalSeeds = {1: [2, 2], 2: [4, 1], 3: [0, 4]}

array = np.zeros((10, 10))
originalSeeds = {1: [2, 2], 2: [4, 1], 3: [0, 4]}
# array = np.zeros((200, 200))
# originalSeeds = {1: [20, 20], 2: [40, 10], 3: [0, 40]}
for key, value in originalSeeds.items():
    array[value[0], value[1]] = key

seeds = {1: [], 2: [], 3: []}
# this is constantly updated list of neighbors

colors = [1, 2, 3]
colorvalues = {1: 1, 2: 2, 3: 3}

# plt.imshow(array)
# plt.show()

start = time()

step = 2**(int(np.log2(array.shape[0])))


def distance(crd1, crd2):
    return np.linalg.norm(crd1 - crd2)


while step >= 1:
    for color in colors:
        originalSeedCoords = originalSeeds[color]
        # neighbor displacments
        disp = step * np.array([[1, 1], [0, 1], [-1, 1], [1, -1], [1, 0], [-1, -1], [0, -1], [-1, 0]])
        for element in disp:
            center = originalSeeds[color]
            idx = center + np.array(element)

            if idx[0] > array.shape[0] - 1 or idx[1] > array.shape[1] - 1 or idx[0] < 0 or idx[1] < 0:
                continue
            else:

                if array[idx[0], idx[1]] == color:
                    continue
                elif array[idx[0], idx[1]] == 0:
                    array[idx[0], idx[1]] = colorvalues[color]
                    seeds[color].append(list(idx))
                else:
                    currColor = array[idx[0], idx[1]]

                    if distance(idx, originalSeeds[color]) < distance(idx, originalSeeds[currColor]):
                        array[idx[0], idx[1]] = colorvalues[color]
                        seeds[color].append(list(idx))
                    else:
                        array[idx[0], idx[1]] = colorvalues[currColor]
                        seeds[currColor].append(list(idx))

    for color in colors:
        for seed in seeds[color]:  # this list may grov during iteration, filling in uncharted points
            for element in disp:
                center = seed

                idx = center + np.array(element)

                if idx[0] > array.shape[0] - 1 or idx[1] > array.shape[1] - 1 or idx[0] < 0 or idx[1] < 0:
                    continue
                else:
                    if array[idx[0], idx[1]] == color:
                        continue
                    elif array[idx[0], idx[1]] == 0:
                        array[idx[0], idx[1]] = colorvalues[color]
                        if list(idx) not in seeds[color]:
                            seeds[color].append(list(idx))
                    else:
                        currColor = array[idx[0], idx[1]]

                        if distance(idx, originalSeeds[color]) < distance(idx, originalSeeds[currColor]):
                            array[idx[0], idx[1]] = colorvalues[color]
                            if list(idx) not in seeds[color]:
                                seeds[color].append(list(idx))
                        else:
                            array[idx[0], idx[1]] = colorvalues[currColor]
                            if list(idx) not in seeds[currColor]:
                                seeds[currColor].append(list(idx))

    step = step // 2
    # step -= 1

2 Answers

Jump flooding algorithm is designed for GPU where the function is executed on all pixels in parallel and a ping-pong buffer is used to store the result of last pass. It should not be used in such as sequential way, not to say implemented in Python.

You can perform several micro-optimizations like avoiding the re-computation of expression (eg. array indexing), avoiding Numpy calls on very small array (since they are more costly than doing the actual array computation due to many checks and allocations), etc. Moreover, you can use a set to speed up the not in operator and actually even remove them. Note that you cannot add new values to an iterating set. One solution to overcome this is to add the value in a new set and iterate over the new set later, and so on until the new set is empty. Here is the resulting code:

import matplotlib.pyplot as plt
import numpy as np
import math


# array = np.zeros((5, 5))
# originalSeeds = {1: [2, 2], 2: [4, 1], 3: [0, 4]}

array = np.zeros((100, 100))
originalSeeds = {1: [2, 2], 2: [4, 1], 3: [0, 4]}
# array = np.zeros((200, 200))
# originalSeeds = {1: [20, 20], 2: [40, 10], 3: [0, 40]}
for key, value in originalSeeds.items():
    array[value[0], value[1]] = key

seeds = {1: set(), 2: set(), 3: set()}
# this is constantly updated list of neighbors

colors = [1, 2, 3]
colorvalues = {1: 1, 2: 2, 3: 3}

# plt.imshow(array)
# plt.show()


step = 2**(int(np.log2(array.shape[0])))


def distance(crd1, crd2):
    return math.sqrt((crd1[0] - crd2[0])**2 + (crd1[1] - crd2[1])**2)


n, m = array.shape
while step >= 1:
    disp = [tuple(e) for e in step * np.array([[1, 1], [0, 1], [-1, 1], [1, -1], [1, 0], [-1, -1], [0, -1], [-1, 0]])]
    for color in colors:
        originalSeedCoords = originalSeeds[color]
        # neighbor displacments
        for element in disp:
            center = originalSeeds[color]

            x = int(center[0] + element[0])
            y = int(center[1] + element[1])
            idx = (x, y)

            if x > n - 1 or y > m - 1 or x < 0 or y < 0:
                continue
            else:
                if array[x, y] == color:
                    continue
                elif array[x, y] == 0:
                    array[x, y] = colorvalues[color]
                    seeds[color].add(idx)
                else:
                    currColor = array[x, y]

                    if distance(idx, originalSeeds[color]) < distance(idx, originalSeeds[currColor]):
                        array[x, y] = colorvalues[color]
                        seeds[color].add(idx)
                    else:
                        array[x, y] = colorvalues[currColor]
                        seeds[currColor].add(idx)

    for color in colors:
        curSeeds = seeds[color]
        while len(curSeeds) > 0:
            newSeeds = set()
            for seed in curSeeds:  # this list may grow during iteration, filling in uncharted points
                for element in disp:
                    center = seed

                    x = int(center[0] + element[0])
                    y = int(center[1] + element[1])
                    idx = (x, y)

                    if x > n - 1 or y > m - 1 or x < 0 or y < 0:
                        continue
                    else:
                        if array[x, y] == color:
                            continue
                        elif array[x, y] == 0:
                            array[x, y] = colorvalues[color]
                            newSeeds.add(idx)
                        else:
                            currColor = array[x, y]

                            if distance(idx, originalSeeds[color]) < distance(idx, originalSeeds[currColor]):
                                array[x, y] = colorvalues[color]
                                newSeeds.add(idx)
                            else:
                                array[x, y] = colorvalues[currColor]
                                if currColor != color:
                                    seeds[currColor].add(idx)
                                else:
                                    newSeeds.add(idx)
            seeds[color].update(newSeeds)
            curSeeds = newSeeds

    step = step // 2
    # step -= 1

This unchecked code is about 12 times faster on my machine.

Using a plain Python code with loops and basic operations executed by the CPython interpreter is not great performance wise. The Numba's JIT or Cython could be used to speed up a lot the code. Rewriting the algorithm in C++ could also help a lot. I think it is also possible to improve even more the algorithm by using a better data structure than a set.

Related