ENH: Poisson disk sampling for arbitrary bounds

[mplough-kobold]

Is your feature request related to a problem? Please describe.

#13918 added Poisson disk sampling of the unit hypercube to scipy.stats.qmc. Poisson disk sampling is widely used in image processing and image generation; see scikit-image/scikit-image#2380 for example use cases. Sampling of the unit hypercube is not sufficient for image processing applications because an image can be any aspect ratio. The sampling algorithm itself must be aware of the bounds from which to draw samples; scaling samples from the unit hypercube using scipy.stats.qmc.scale results in loss of the radius distance property through unequal scaling. In the two dimensional case for example, a circular radius gets squashed into an ellipse as shown below.

Original sampling of unit hypercube

After scaling

Describe the solution you'd like.

It will be best to modify scipy.stats.qmc.PoissonDisk to optionally accept, along with the existing dimension parameter d, a d-dimensional u_bounds parameter similar to scipy.stats.qmc.scale. If u_bounds is left unspecified the Poisson disk sampling will proceed on the unit hypercube. If it is specified the samples will be scaled prior to radius consideration.

Describe alternatives you've considered.

Sample scaling via scipiy.stats.qmc.scale is not possible due to the distortion it introduces.

Subclassing scipy.stats.qmc.PoissonDisk to allow pre-scaling is not easy due to the tight coupling of the radius parameter and the initialization of the cell grid.

Currently it's necessary to either create a parallel PoissonDisk implementation or to not use SciPy for Poisson disk sampling from arbitrary bounds.

It's possible to calculate the image aspect ratio and do rejection sampling on the unit hypercube, rejecting samples that lie outside a "hyperrectangle". For 2D that looks like this:

import numpy as np
from scipy.stats import qmc


class AspectRatioRejection:
    def __init__(self, width: int, height: int):
        """Rejection sampling to get around scipy.stats.qmc.PoissonDisk scaling issues.

        See https://github.com/scipy/scipy/issues/20288.
        """
        if width >= height:
            self.xmax = 1.0
            self.ymax = height / width
        else:
            self.xmax = width / height
            self.ymax = 1.0

    def __call__(self, sample: tuple[float, float]) -> bool:
        """Reject sample if it's outside the aspect ratio."""
        return sample[0] > self.xmax or sample[1] > self.ymax


def fill_space_blue_noise_samples(
    width: int,
    height: int,
    radius_fraction: float = 0.05,
    rng: np.random.Generator | None = None
):
    """Fill image area with blue noise samples.

    Samples are spaced at least `radius_fraction` apart, where radius_fraction is a fraction of the image width.
    """
    aspect_reject = AspectRatioRejection(width, height)

    rng = np.random.default_rng() if rng is None else rng
    engine = qmc.PoissonDisk(d=2, radius=radius_fraction, seed=rng)
    samples = [s * max(width, height) for s in engine.fill_space() if not aspect_reject(s)]

    return samples

However, the performance gets worse and worse as the space's aspect ratio gets larger. For the case of a line, all samples will be almost surely rejected. However, it's very easy (and performant) to generate samples in a lower dimensional space so this performance hit can be avoided entirely through prescaling.

Additional context (e.g. screenshots, GIFs)

No response

[tupui]

Hi @mplough-kobold thank you for the suggestion. Also you detailed report is highly appreciated 🙏

That's a legitimate ask 👍 would you be interested in making a PR? I am happy to help with reviewing it or helping you implement the solution.

[mplough-kobold]

@tupui - glad the report was helpful. I'd be interested in making a PR but unfortunately don't have time right now to dig in. So this is definitely up for grabs if someone else wants to pick this up.

[tupui]

Thanks for letting us know. I also don't have time at the moment. Feel free to come back when you do have time 😃