Fast O(n) sampling of points on a 2D plane.

poisson_disc_sampling.py

"""
	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])