Poisson disk sampling in n-dimensions

poisson_disk_sampling.py

"""Poisson disk sampling in n-dimensions.

This is part of an effort to add Poisson disk sampling into SciPy:

https://github.com/scipy/scipy/pull/13918

---------------------------

MIT License

Copyright (c) 2021 Pamphile Tupui ROY

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.

"""
from typing import List, Tuple

import numpy as np


def kernel_factory(dim: int = 2, radius: float = 1) -> np.ndarray:
    """Create a kernel.

    It will be a donut with ``r1=radius`` and ``r2=2*radius``.
    The center point's value is 2 while the donut has 1s and
    the rest is 0.

    array([[0, 0, 1, 0, 0],
           [0, 1, 1, 1, 0],
           [1, 1, 2, 1, 1],
           [0, 1, 1, 1, 0],
           [0, 0, 1, 0, 0]])

    """
    xx = np.meshgrid(*[np.arange(radius * 4+1) for _ in range(dim)])

    diam = 2 * radius

    # hypersphere L2 distance
    dists = np.sqrt(np.sum((np.array(xx)-diam)**2, axis=0))

    donut = (dists <= diam) & (dists >= radius).astype(int)
    kernel = donut+(dists < radius).astype(int) * 2
    return kernel


def candidates_around_point(
        point: np.ndarray, kernel: np.ndarray, sample_grid: np.ndarray
) -> Tuple[np.ndarray, List[float]]:
    """Sample candidates with respect to a kernel and existing samples."""
    kernel_len = len(kernel)-1
    half_len = kernel_len // 2
    size = sample_grid.shape[0]

    # subspace: hypercube around point
    sub_space = [(idx_-kernel_len//2, idx_ + kernel_len//2 + 1) for idx_ in point]

    # mask kernel for feasible points and close to boundaries (half kernel)
    bounds = [np.arange(*idx_) for idx_ in sub_space]
    space_idx = np.meshgrid(*bounds, indexing='ij')
    kernel_mask = np.logical_and.reduce(
        [np.logical_and(0 <= idx_, idx_ < size) for idx_ in space_idx]
    )

    sub_space = np.clip(sub_space, 0, size)
    slices = [slice(*idx_) for idx_ in sub_space]
    hypercube = sample_grid[tuple(np.flip(slices))]

    # mask point's kernel on the sample_grid
    kernel_ = kernel[kernel_mask].reshape(hypercube.shape).astype(float)
    kernel_[kernel_ == 2] = np.inf
    hypercube += kernel_

    # select new candidates from possible
    candidates = np.argwhere(hypercube == 1)
    candidates = (candidates+point) - half_len

    return candidates


def poisson_sampling(
        n, dim: int = 2, radius: float = 1, grid_size: int = 10, seed=None,
) -> np.ndarray:
    """Poisson disk sampling.

    Sample in ``[0, grid_size-1]**dim``.
    """
    rng = np.random.default_rng(seed)

    # hold the samples: a cell can contain a single sample
    sample_grid = np.zeros((grid_size,) * dim)

    kernel = kernel_factory(dim=dim, radius=radius)
    half_len = len(kernel) // 2

    # add a point in the inner domain: at least 1 kernel away from borders
    point = rng.integers(half_len, high=grid_size-half_len-1, size=dim)
    points = [point]
    candidates = [point]

    for _ in range(n-1):
        idx_new_point = rng.integers(len(candidates))
        point = candidates[idx_new_point]
        points.append(point)

        candidates_ = candidates_around_point(point, kernel, sample_grid)

        # TODO make this smarter
        candidates = np.argwhere(np.logical_and(0<sample_grid, sample_grid<np.inf))
        candidates = np.flip(candidates)

        # filter with bounds
        candidates_idx = np.where(
            np.all(half_len <= candidates, axis=1)
            & np.all(candidates <= grid_size-half_len-1, axis=1)
        )
        candidates = candidates[candidates_idx]

        if len(candidates) == 0:
            break

    points = np.array(points)

    return points

Plotting is just a matter of adding the following in the main function poisson_sampling :

import matplotlib as mpl
import matplotlib.pyplot as plt

fig, ax = plt.subplots()
cmap = mpl.cm.get_cmap("gray_r").copy()
cmap.set_bad(color='b', alpha=.2)
ax.imshow(sample_grid.T,cmap=cmap)
ax.scatter(points[:-1, 1], points[:-1, 0], c='r')
plt.show()