Skip to content

Instantly share code, notes, and snippets.

Embed
What would you like to do?
Fast O(n) sampling of points on a 2D plane.
"""
Copyright (C) 2016 Anton Crombach (anton.crombach@gmail.com)
This program is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.
This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
GNU General Public License for more details.
You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
DESCRIPTION
Python implementation of Daniel Dunbar's C code, with some small
changes. See "A Spatial Data Structure for Fast Poisson-Disk
Sample Generation", in Proc' of SIGGRAPH 2006 for more information.
"""
import numpy as np
import random
import collections
RangeEntry = collections.namedtuple('RangeEntry', ['min', 'max'])
class Point(object):
"""Simple (x,y) point class."""
def __init__(self, x, y):
self.x = x
self.y = y
def length(self):
return np.sqrt(self.x * self.x + self.y * self.y)
def square_dist(self, other):
return (self.x - other.x)**2 + (self.y - other.y)**2
def __add__(self, other):
try:
return Point(self.x + other.x, self.y + other.y)
except AttributeError:
return NotImplemented
def __sub__(self, other):
try:
return Point(self.x - other.x, self.y - other.y)
except AttributeError:
return NotImplemented
def __iadd__(self, other):
try:
self.x += other.x
self.y += other.y
except AttributeError:
return NotImplemented
def __isub__(self, other):
try:
self.x += other.x
self.y += other.y
except AttributeError:
return NotImplemented
def __repr__(self):
return "({0}, {1})".format(self.x, self.y)
class BoundarySampler(object):
"""Sample Poisson-disc style points in O(N)."""
def __init__(self, size, radius):
self.rangelist = []
self.points = []
# Tuple of width and height
self.size = size
self.radius = radius
# Make sure size (x,y) is multiple of radius
assert(self.size[0] % (4 * self.radius) == 0 and
self.size[1] % (4 * self.radius) == 0)
# Grid size is chosen so that 4*radius search only requires searching adjacent cells,
# this also determines max points per cell.
width, height = self.size
self.grid_width = int(0.5 + width / (4.0 * self.radius))
self.grid_height = int(0.5 + height / (4.0 * self.radius))
# Cell size is 4*radius
self.grid_cell_sizex = width / self.grid_width
self.grid_cell_sizey = height / self.grid_height
if options.verbose:
print("Size:", self.size)
print("Radius:", self.radius)
print("Grid:", self.grid_width, self.grid_height)
print("Cell:", self.grid_cell_sizex, self.grid_cell_sizey)
self.grid = np.full(
(self.grid_width, self.grid_height, MAX_POINTS_PER_GRIDCELL), -1, dtype=np.int)
def as_gridxy(self, pt):
width, height = self.size
gx = int(pt.x / self.grid_cell_sizex)
gy = int(pt.y / self.grid_cell_sizey)
assert(gx >= 0 and gy >= 0 and gx < self.grid_width and gy < self.grid_height)
return (gx, gy)
def add_point(self, pt):
self.points.append(pt)
# Update grid for quick lookup
gx, gy = self.as_gridxy(pt)
for k in range(MAX_POINTS_PER_GRIDCELL):
if self.grid[gx, gy, k] == -1:
self.grid[gx, gy, k] = len(self.points)-1
break
assert(k < MAX_POINTS_PER_GRIDCELL)
def add_random_point(self):
"""Generate point in (width, height) rectangle."""
width, height = self.size
self.add_point(Point(width * random.random(), height * random.random()))
def find_neighbour_ranges(self, index):
"""Find nearby neighbours."""
candidate = self.points[index]
range_sqrd = 4 * 4 * self.radius * self.radius
Nx, Ny = 1, 1
gx, gy = self.as_gridxy(candidate)
x_side = int((candidate.x - gx*self.grid_cell_sizex) > self.grid_cell_sizex * .5)
y_side = int((candidate.y - gy*self.grid_cell_sizey) > self.grid_cell_sizey * .5)
iy = 1
for j in range(-Ny, Ny+1):
ix = 1
if j == 0:
iy = y_side
elif j == 1:
iy = 0
for i in range(-Nx, Nx+1):
if i == 0:
ix = x_side
elif i == 1:
ix = 0
# Offset to closest cell point
dx = candidate.x - ((gx + i + ix) * self.grid_cell_sizex)
dy = candidate.y - ((gy + j + iy) * self.grid_cell_sizey)
if dx*dx + dy*dy < range_sqrd:
cx = (gx + i + self.grid_width) % self.grid_width
cy = (gy + j + self.grid_height) % self.grid_height
cell = self.grid[cx, cy, :]
for k in range(MAX_POINTS_PER_GRIDCELL):
if cell[k] == -1:
break
elif cell[k] != index:
v = self.points[cell[k]] - candidate
dist_sqrd = v.x*v.x + v.y*v.y
if dist_sqrd < range_sqrd:
dist = np.sqrt(dist_sqrd)
angle = np.arctan2(v.y, v.x)
theta = np.arccos(0.25 * dist / self.radius)
self.subtract(angle - theta, angle + theta)
def complete(self):
"""Do sampling."""
candidates = []
self.add_random_point()
candidates.append(len(self.points) - 1)
while candidates:
# Pick a random candidate
c = random.randrange(0, len(candidates))
index = candidates[c]
candidate = self.points[index]
# Smart removal of candidate c
candidates[c] = candidates[-1]
candidates.pop()
self.rangelist = [RangeEntry(min=0, max=TAU)]
self.find_neighbour_ranges(index)
while self.rangelist:
re = random.choice(self.rangelist)
angle = re.min + (re.max - re.min) * random.random()
pt = Point(candidate.x + np.cos(angle) * 2 * self.radius,
candidate.y + np.sin(angle) * 2 * self.radius)
if 0 < pt.x < WIDTH and 0 < pt.y < HEIGHT:
self.add_point(pt)
candidates.append(len(self.points)-1)
self.subtract(angle - np.pi/3.0, angle + np.pi/3.0)
def subtract(self, a, b):
"""Subtract range covered by new point from existing ranges."""
if a > TAU:
self.subtract(a - TAU, b - TAU)
elif b < 0.0:
self.subtract(a + TAU, b + TAU)
elif a < 0.0:
self.subtract(0.0, b)
self.subtract(a + TAU, TAU)
elif b > TAU:
self.subtract(a, TAU)
self.subtract(0, b - TAU)
elif not self.rangelist:
pass
else:
# Binary search
if a < self.rangelist[0].min:
pos = -1
else:
lo, mid, hi = 0, 0, len(self.rangelist)
while lo < hi-1:
mid = (lo + hi) // 2
if self.rangelist[mid].min < a:
lo = mid
else:
hi = mid
pos = lo
# Update ranges
if pos == -1:
pos = 0
elif a < self.rangelist[pos].max:
c, d = self.rangelist[pos]
if a - c < SMALLEST_RANGE:
if b < d:
self.rangelist[pos] = self.rangelist[pos]._replace(min=b)
else:
del self.rangelist[pos]
else:
self.rangelist[pos] = self.rangelist[pos]._replace(max=a)
if b < d:
self.rangelist.insert(pos+1, RangeEntry(b, d))
pos += 1
else:
if pos < len(self.rangelist)-1 and b > self.rangelist[pos+1].min:
pos += 1
else:
return
while pos < len(self.rangelist) and b > self.rangelist[pos].min:
if self.rangelist[pos].max - b < SMALLEST_RANGE:
del self.rangelist[pos]
else:
self.rangelist[pos] = self.rangelist[pos]._replace(min=b)
#
# Main function to call
#
def poisson_disc_sampling(size, distance):
"""
Choose random points on a lattice with at least distance between them.
"""
bs = BoundarySampler(size, distance)
bs.complete()
return np.array([[pt.x/size[0], pt.y/size[1]] for pt in bs.points])
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment