poisson_disk_sampling.py

# Experimental Python implementation of Poisson disk sampling with O(n) space and time complexity.
# Reference: https://www.cs.ubc.ca/~rbridson/docs/bridson-siggraph07-poissondisk.pdf

import random
import numpy as np
import matplotlib.pyplot as plt
    
S2 = np.sqrt(2)

def runit(r, shape, k=10):
    
    m = np.ceil(np.asarray(shape) * S2 / r).astype(int) + 1
    
    def is_valid_point(x):        
        if not (0 <= x[0] < shape[0] and 0 <= x[1] < shape[1]):
            return False
        
        gridpos = x_to_gridpos(x)
        
        for p0 in range(gridpos[0]-2, gridpos[0]+3):
            for p1 in range(gridpos[1]-2, gridpos[1]+3):
                if 0 <= p0 < m[0] and  0 <= p1 < m[1] and grid[p0, p1] > -1:
                    if np.sum((samples[grid[p0, p1]] - x)**2) < r**2:
                        return False
        return True
        
        
    def x_to_gridpos(x):
        x = np.asarray(x) * S2 / r
        return tuple(np.round(x).astype(int))
        
    
    grid = -np.ones(m, dtype=np.int32)
    
    x0 = np.random.rand(2) * shape
    
    grid[x_to_gridpos(x0)] = 0
    
    active = {0}
    samples = [x0]
    
    while len(active) > 0:
        i = random.sample(active, 1)[0]
        x0 = samples[i]
    
        for ki in range(k):
            sr = np.random.rand() * r + r
            sa = np.random.rand() * 2 * np.pi
            x = x0 + np.array([np.cos(sa) * sr, np.sin(sa) * sr])
            
            if is_valid_point(x):
                j = len(samples)
                grid[x_to_gridpos(x)] = j
                active.add(j)
                samples.append(x)                
                plt.plot([x0[0], x[0]], [x0[1], x[1]], 'g', alpha=0.6)                
                break
            
        if ki == k - 1:                
            active.remove(i)
            
    print(len(samples), np.prod(m))
            
    return samples, grid
    
r = 2.0
shape = [100, 80]

samples, grid = runit(r, shape)

plt.plot(*np.transpose(samples), 'k.')

plt.figure()
plt.imshow(grid)