diff --git a/src/spikeinterface/core/generate.py b/src/spikeinterface/core/generate.py index 2c0ca2cfd5..c7b484f890 100644 --- a/src/spikeinterface/core/generate.py +++ b/src/spikeinterface/core/generate.py @@ -2167,21 +2167,207 @@ def generate_channel_locations(num_channels, num_columns, contact_spacing_um): return channel_locations -def _generate_multimodal(rng, size, num_modes, lim0, lim1): - bins = np.linspace(lim0, lim1, 10000) - bin_step = bins[1] - bins[0] - prob = np.zeros(bins.size) +def _multimodal_density(values, num_modes, lim0, lim1): + """ + Unnormalized density of the "multimodal" ("by layer") distribution, evaluated at `values`. + """ + values = np.asarray(values, dtype="float64") + density = np.zeros(values.shape, dtype="float64") mode_step = (lim1 - lim0) / (num_modes + 1) + sigma = mode_step / 5.0 for i in range(num_modes): center = mode_step * (i + 1) - sigma = mode_step / 5.0 - prob += np.exp(-((bins - center) ** 2) / (2 * sigma**2)) + density += np.exp(-((values - center) ** 2) / (2 * sigma**2)) + return density + + +def _generate_multimodal(rng, size, num_modes, lim0, lim1): + bins = np.linspace(lim0, lim1, 10000) + bin_step = bins[1] - bins[0] + prob = _multimodal_density(bins, num_modes, lim0, lim1) prob /= np.sum(prob) choices = rng.choice(np.arange(bins.size), size, p=prob) values = bins[choices] + rng.uniform(low=-bin_step / 2, high=bin_step / 2, size=size) return values +def _indices_closer_than(points, minimum_distance): + """ + Return the sorted indices of the points that lie within `minimum_distance` of another point. + + This answers the same question as taking the full pairwise distance matrix and collecting the + rows that contain a violation, but it only ever compares points that share a grid cell or one + of its neighbours. Cells are `minimum_distance` wide, so a violating pair can never be more + than one cell apart on any axis and nothing outside that block has to be looked at. The cost + then follows the number of points rather than its square. + + Parameters + ---------- + points : numpy.ndarray + A (num_points, 3) array of positions. + minimum_distance : float + Positions strictly closer together than this are reported. + + Returns + ------- + indices : numpy.ndarray + The sorted, unique indices of the offending points. + """ + num_points = points.shape[0] + if num_points < 2: + return np.zeros(0, dtype="int64") + + coordinates = np.asarray(points, dtype="float64") + origin = coordinates.min(axis=0) + cells = np.floor((coordinates - origin) / minimum_distance).astype("int64") + + # group the point indices by cell, so a cell can be looked up without scanning every point + buckets = {} + for index, cell in enumerate(map(tuple, cells)): + buckets.setdefault(cell, []).append(index) + + offsets = [(dx, dy, dz) for dx in (-1, 0, 1) for dy in (-1, 0, 1) for dz in (-1, 0, 1)] + squared_minimum = minimum_distance * minimum_distance + offending = set() + for cell, members in buckets.items(): + neighbourhood = [] + for dx, dy, dz in offsets: + neighbourhood.extend(buckets.get((cell[0] + dx, cell[1] + dy, cell[2] + dz), ())) + neighbour_indices = np.asarray(neighbourhood, dtype="int64") + deltas = coordinates[members][:, np.newaxis, :] - coordinates[neighbour_indices][np.newaxis, :, :] + squared = np.sum(deltas * deltas, axis=2) + # a point always finds itself at distance zero, so drop the self pairing before testing + squared[np.asarray(members)[:, np.newaxis] == neighbour_indices[np.newaxis, :]] = np.inf + offending.update(np.asarray(members)[np.any(squared < squared_minimum, axis=1)].tolist()) + + return np.array(sorted(offending), dtype="int64") + + +def _poisson_disk_sampling_3d(rng, lower, upper, radius, min_points=0, num_candidates=30, max_sweeps=20): + """ + Fill an axis aligned 3D box with points that are all at least `radius` apart, using Bridson's + algorithm. + + Accepted points are bucketed into a uniform grid whose cell size is `radius / sqrt(3)`, so the + diagonal of a cell is exactly `radius` and a cell can therefore hold at most one point. Nothing + further than `radius` from a candidate can rule it out, and a candidate is never further than + `2 * radius` from the point it was drawn around, so only a small fixed block of cells around + that point has to be looked at instead of every point accepted so far. That is what keeps the + whole fill linear in the number of points rather than quadratic. + + Parameters + ---------- + rng : numpy.random.Generator + The random generator to draw from. + lower : array-like + The (x, y, z) corner of the box with the smallest coordinates. + upper : array-like + The (x, y, z) corner of the box with the largest coordinates. + radius : float + The minimum distance between any two returned points. Must be strictly positive. + min_points : int, default: 0 + Stop as soon as this many points have been placed. Only whole sweeps of the box are ever + run, so the points returned always cover the whole box rather than one corner of it. + num_candidates : int, default: 30 + How many candidates are drawn around an active point before it is retired. This is the + "k" of Bridson's algorithm, higher values pack the box tighter at a linear extra cost. + max_sweeps : int, default: 20 + How many times every placed point is given another go at spawning neighbours. Sweeping + stops early as soon as a sweep adds nothing, so this is only an upper bound. + + Returns + ------- + points : numpy.ndarray + A (num_points, 3) float32 array of points inside the box, pairwise further apart than + `radius`. The number of points depends on the box and on the radius, it is not chosen. + """ + lower = np.asarray(lower, dtype="float64") + upper = np.asarray(upper, dtype="float64") + extent = np.maximum(upper - lower, 0.0) + + # a flat box is still usable, but candidates must not be pushed off the box along a flat axis + flat_axes = extent <= 0.0 + if np.all(flat_axes): + return lower.astype("float32").reshape(1, 3) + + cell_size = radius / np.sqrt(3.0) + grid_shape = np.maximum(np.ceil(extent / cell_size), 1.0).astype("int64") + # each cell holds the index of the single point that falls in it, -1 means the cell is empty + grid = np.full(tuple(grid_shape), -1, dtype="int64") + + points = np.empty((256, 3), dtype="float32") + num_points = 0 + + def store(point): + nonlocal points, num_points + if num_points == points.shape[0]: + points = np.concatenate((points, np.empty_like(points))) + points[num_points] = point + cell = np.clip(((np.asarray(point, dtype="float64") - lower) / cell_size).astype("int64"), 0, grid_shape - 1) + grid[cell[0], cell[1], cell[2]] = num_points + num_points += 1 + return num_points - 1 + + # points are stored as float32, so the check is done on the float32 values with a hair of head + # room, that way the constraint still holds for the array that is handed back to the caller + squared_radius = (radius * (1.0 + 1e-6)) ** 2 + # a candidate sits within 2 * radius of its parent and can only clash with a point within + # radius of itself, so nothing outside this many cells around the parent can matter + block_halfwidth = int(np.ceil(3.0 * np.sqrt(3.0))) + + # in a box that is thin compared to the radius, most of the shell around a point lands outside + # the box, so draw proportionally more candidates to keep about num_candidates of them usable + thinness = np.prod(np.clip(extent[~flat_axes] / (2.0 * radius), 0.0, 1.0)) + num_draws = num_candidates * int(np.clip(round(1.0 / max(thinness, 1e-3)), 1, 32)) + + store((lower + rng.random(3) * extent).astype("float32")) + + # One Bridson pass leaves gaps, because a point is retired as soon as num_candidates tries in a + # row miss. Sweeping again over every point already placed fills most of those gaps, which is + # what brings the packing up to what the old rejection sampler reached at high iteration + # counts. Extra sweeps are only worth their cost when the caller still needs more points, so + # stop as soon as there are enough or a sweep stops adding any. + for _ in range(max_sweeps): + before_sweep = num_points + active = list(range(num_points)) + + while len(active) > 0: + # Bridson picks a random active point, which grows the front in every direction at once + which = int(rng.integers(len(active))) + parent = points[active[which]].astype("float64") + + directions = rng.normal(size=(num_draws, 3)) + directions[:, flat_axes] = 0.0 + directions /= np.linalg.norm(directions, axis=1, keepdims=True) + # uniform over the volume of the spherical shell [radius, 2 * radius) + offsets = radius * np.cbrt(1.0 + 7.0 * rng.random((num_draws, 1))) + candidates = (parent + directions * offsets).astype("float32") + + parent_cell = np.clip(((parent - lower) / cell_size).astype("int64"), 0, grid_shape - 1) + low = np.maximum(parent_cell - block_halfwidth, 0) + high = np.minimum(parent_cell + block_halfwidth + 1, grid_shape) + block = grid[low[0] : high[0], low[1] : high[1], low[2] : high[2]] + neighbours = points[block[block >= 0]].astype("float64") + + candidates_64 = candidates.astype("float64") + inside = np.all((candidates_64 >= lower) & (candidates_64 <= upper), axis=1) + deltas = candidates_64[:, np.newaxis, :] - neighbours[np.newaxis, :, :] + clear = np.all(np.sum(deltas * deltas, axis=2) > squared_radius, axis=1) + valid = np.flatnonzero(inside & clear) + + if valid.size > 0: + active.append(store(candidates[valid[0]])) + else: + # nothing else fits around this parent, retire it + active[which] = active[-1] + active.pop() + + if num_points >= min_points or num_points == before_sweep: + break + + return points[:num_points] + + def generate_unit_locations( num_units, channel_locations, @@ -2207,8 +2393,11 @@ def generate_unit_locations( * z coordinates are within a specified range `(minimum_z, maximum_z)` 2) the distance between any two units is greater than a specified minimum value - If the minimum distance constraint cannot be met within the allowed number of iterations, - the function can either raise an exception or issue a warning based on the `distance_strict` flag. + The locations are drawn with Bridson's Poisson disk sampling, accelerated by a uniform grid, so + the cost grows with the number of units rather than with its square. + + If the box is too small to hold `num_units` locations that far apart, the function can either + raise an exception or issue a warning based on the `distance_strict` flag. Parameters ---------- @@ -2227,18 +2416,18 @@ def generate_unit_locations( minimum_distance : float, default: 20.0 The minimum allowable distance in micrometers between any two units max_iteration : int, default: 100 - The maximum number of iterations to attempt generating unit locations that meet - the minimum distance constraint. + The maximum number of times the sampling radius is reduced while looking for enough + locations. The sampler no longer needs many attempts, so this is mostly kept for backward + compatibility and the default is far above what is used in practice. distance_strict : bool, default: False If True, the function will raise an exception if a solution meeting the distance - constraint cannot be found within the maximum number of iterations. If False, a warning - will be issued. + constraint cannot be found. If False, a warning will be issued. distribution : "uniform" | "multimodal", default: "uniform" How units are spread. "uniform" is units everywhere "multimodal" mimic the distribution of units 'by layer' on the 'y' axis (dim=1) Important note, when using multimodal in conjonction of minimum_distance not None, there is not garanty - of a true multimodal because units that do not respect the distance of move again and are most chance to be in between layers. + of a true multimodal because the distance constraint caps how many units a layer can hold. num_modes : int, default 2 In case of distribution="multimodal", this is the number of modes (layers) seed : int or None, optional @@ -2250,61 +2439,83 @@ def generate_unit_locations( A 2D array of shape (num_units, 3), where each row represents the (x, y, z) coordinates of a generated unit location. """ + if distribution not in ("uniform", "multimodal"): + raise ValueError("generate_unit_locations has wrong distribution must be 'uniform' or 'multimodal'") + rng = np.random.default_rng(seed=seed) - units_locations = np.zeros((num_units, 3), dtype="float32") minimum_x, maximum_x = np.min(channel_locations[:, 0]) - margin_um, np.max(channel_locations[:, 0]) + margin_um minimum_y, maximum_y = np.min(channel_locations[:, 1]) - margin_um, np.max(channel_locations[:, 1]) + margin_um - units_locations[:, 0] = rng.uniform(minimum_x, maximum_x, size=num_units) - if distribution == "uniform": - units_locations[:, 1] = rng.uniform(minimum_y, maximum_y, size=num_units) - elif distribution == "multimodal": - units_locations[:, 1] = _generate_multimodal(rng, num_units, num_modes, minimum_y, maximum_y) - else: - raise ValueError("generate_unit_locations has wrong distribution must be 'uniform' or ") - units_locations[:, 2] = rng.uniform(minimum_z, maximum_z, size=num_units) - - if minimum_distance is not None: - solution_found = False - renew_inds = None - for i in range(max_iteration): - distances = np.linalg.norm(units_locations[:, np.newaxis] - units_locations[np.newaxis, :], axis=2) - inds0, inds1 = np.nonzero(distances < minimum_distance) - mask = inds0 != inds1 - inds0 = inds0[mask] - inds1 = inds1[mask] - - if inds0.size > 0: - if renew_inds is None: - renew_inds = np.unique(inds0) - else: - # random only bad ones in the previous set - renew_inds = renew_inds[np.isin(renew_inds, np.unique(inds0))] - - units_locations[:, 0][renew_inds] = rng.uniform(minimum_x, maximum_x, size=renew_inds.size) - if distribution == "uniform": - units_locations[:, 1][renew_inds] = rng.uniform(minimum_y, maximum_y, size=renew_inds.size) - - elif distribution == "multimodal": - units_locations[:, 1][renew_inds] = _generate_multimodal( - rng, renew_inds.size, num_modes, minimum_y, maximum_y - ) - units_locations[:, 2][renew_inds] = rng.uniform(minimum_z, maximum_z, size=renew_inds.size) - - else: - solution_found = True - break - - if not solution_found: - if distance_strict: - raise ValueError( - f"generate_unit_locations(): no solution for {minimum_distance=} and {max_iteration=} " - "You can use distance_strict=False or reduce minimum distance" - ) - else: - warnings.warn(f"generate_unit_locations(): no solution for {minimum_distance=} and {max_iteration=}") + def draw_unconstrained(size): + locations = np.zeros((size, 3), dtype="float32") + locations[:, 0] = rng.uniform(minimum_x, maximum_x, size=size) + if distribution == "uniform": + locations[:, 1] = rng.uniform(minimum_y, maximum_y, size=size) + else: + locations[:, 1] = _generate_multimodal(rng, size, num_modes, minimum_y, maximum_y) + locations[:, 2] = rng.uniform(minimum_z, maximum_z, size=size) + return locations + + if minimum_distance is None or minimum_distance <= 0 or num_units == 0: + return draw_unconstrained(num_units) + + # First keep drawing positions uniformly and redrawing the ones that sit too close, which is + # what this function has always done. That leaves the units spread the way callers already + # depend on: they are only ever pushed apart as far as minimum_distance actually requires, + # rather than as far as the box would allow. The one thing that changes is how the offending + # units are found, a uniform grid instead of the full pairwise distance matrix, which is what + # made the old loop cost grow with the square of num_units. + units_locations = draw_unconstrained(num_units) + renew_inds = None + for _ in range(max_iteration): + too_close = _indices_closer_than(units_locations, minimum_distance) + if too_close.size == 0: + return units_locations + # narrow to the ones that were already bad last round, matching the previous behaviour + renew_inds = too_close if renew_inds is None else renew_inds[np.isin(renew_inds, too_close)] + units_locations[:, 0][renew_inds] = rng.uniform(minimum_x, maximum_x, size=renew_inds.size) + if distribution == "uniform": + units_locations[:, 1][renew_inds] = rng.uniform(minimum_y, maximum_y, size=renew_inds.size) + else: + units_locations[:, 1][renew_inds] = _generate_multimodal( + rng, renew_inds.size, num_modes, minimum_y, maximum_y + ) + units_locations[:, 2][renew_inds] = rng.uniform(minimum_z, maximum_z, size=renew_inds.size) + + # Redrawing never converged. That is the case the old implementation gave up on, warning and + # handing back locations that still broke the constraint it promises. A box packed close to + # what minimum_distance allows is exactly where independent redraws keep colliding, so fall + # back to placing the units directly with Poisson disk sampling, which fills such a box in one + # pass instead of hoping a redraw happens to land clear. + lower = np.array([minimum_x, minimum_y, minimum_z], dtype="float64") + upper = np.array([maximum_x, maximum_y, maximum_z], dtype="float64") + points = _poisson_disk_sampling_3d(rng, lower, upper, minimum_distance, min_points=num_units) + + if points.shape[0] >= num_units: + if distribution == "uniform": + keep = rng.choice(points.shape[0], size=num_units, replace=False) + else: + # every point already respects the distance, so the layering is applied by picking + # which of them to keep rather than by moving any of them + weights = _multimodal_density(points[:, 1], num_modes, minimum_y, maximum_y) + # drawing without replacement needs at least num_units strictly positive weights, and a + # far away enough point can have its density underflow to zero + weights = np.maximum(weights, np.finfo("float64").tiny) + keep = rng.choice(points.shape[0], size=num_units, replace=False, p=weights / np.sum(weights)) + return points[keep] + + # the box cannot hold num_units points that far apart, report it the way the previous sampler + # did and pad the shortfall with unconstrained draws so the returned shape is still honoured + if distance_strict: + raise ValueError( + f"generate_unit_locations(): no solution for {minimum_distance=} and {max_iteration=} " + "You can use distance_strict=False or reduce minimum distance" + ) + warnings.warn(f"generate_unit_locations(): no solution for {minimum_distance=} and {max_iteration=}") + units_locations = draw_unconstrained(num_units) + units_locations[: points.shape[0]] = points return units_locations diff --git a/src/spikeinterface/core/tests/test_generate.py b/src/spikeinterface/core/tests/test_generate.py index 66f84a2914..fa9eb08c9b 100644 --- a/src/spikeinterface/core/tests/test_generate.py +++ b/src/spikeinterface/core/tests/test_generate.py @@ -423,6 +423,111 @@ def test_generate_unit_locations(): # plt.show() +def _pairwise_distances(unit_locations): + distances = np.linalg.norm(unit_locations[:, np.newaxis] - unit_locations[np.newaxis, :], axis=2) + dist_flat = np.triu(distances, k=1).flatten() + return dist_flat[dist_flat > 0] + + +def test_generate_unit_locations_dense(): + # a 384 contact probe packed close to what minimum_distance allows. The rejection sampler that + # used to back this function could not solve this within max_iteration, so it warned and handed + # back locations that broke the very constraint minimum_distance promises. + seed = 0 + + probe = generate_multi_columns_probe(num_columns=4, num_contact_per_column=96, xpitch=20, ypitch=20) + channel_locations = probe.contact_positions + assert channel_locations.shape[0] == 384 + + num_units = 550 + minimum_distance = 20.0 + margin_um = 20.0 + minimum_z, maximum_z = 5.0, 40.0 + + # distance_strict raises rather than warns, so this fails loudly if the constraint is not met + unit_locations = generate_unit_locations( + num_units, + channel_locations, + margin_um=margin_um, + minimum_z=minimum_z, + maximum_z=maximum_z, + minimum_distance=minimum_distance, + distance_strict=True, + seed=seed, + ) + + assert unit_locations.shape == (num_units, 3) + assert np.all(_pairwise_distances(unit_locations) > minimum_distance) + + # and they must still sit inside the requested box + assert np.all(unit_locations[:, 0] >= np.min(channel_locations[:, 0]) - margin_um) + assert np.all(unit_locations[:, 0] <= np.max(channel_locations[:, 0]) + margin_um) + assert np.all(unit_locations[:, 1] >= np.min(channel_locations[:, 1]) - margin_um) + assert np.all(unit_locations[:, 1] <= np.max(channel_locations[:, 1]) + margin_um) + assert np.all(unit_locations[:, 2] >= minimum_z) + assert np.all(unit_locations[:, 2] <= maximum_z) + + +def test_generate_unit_locations_multimodal(): + seed = 0 + + probe = generate_multi_columns_probe(num_columns=4, num_contact_per_column=96, xpitch=20, ypitch=20) + channel_locations = probe.contact_positions + + num_units = 200 + num_modes = 2 + minimum_distance = 20.0 + margin_um = 20.0 + + unit_locations = generate_unit_locations( + num_units, + channel_locations, + margin_um=margin_um, + minimum_distance=minimum_distance, + distribution="multimodal", + num_modes=num_modes, + distance_strict=True, + seed=seed, + ) + + assert unit_locations.shape == (num_units, 3) + assert np.all(_pairwise_distances(unit_locations) > minimum_distance) + + # the layers must still be visible: count how many units sit near a mode center against what a + # flat distribution over the same windows would give + minimum_y = np.min(channel_locations[:, 1]) - margin_um + maximum_y = np.max(channel_locations[:, 1]) + margin_um + mode_step = (maximum_y - minimum_y) / (num_modes + 1) + half_window = mode_step / 4 + near_a_mode = np.zeros(num_units, dtype=bool) + for i in range(num_modes): + near_a_mode |= np.abs(unit_locations[:, 1] - mode_step * (i + 1)) < half_window + flat_expectation = 2 * num_modes * half_window / (maximum_y - minimum_y) + assert np.sum(near_a_mode) / num_units > 1.5 * flat_expectation + + +def test_generate_unit_locations_no_solution(): + seed = 0 + + probe = generate_multi_columns_probe(num_columns=2, num_contact_per_column=20, xpitch=20, ypitch=20) + channel_locations = probe.contact_positions + + # far more units than the box can hold with that spacing + num_units = 500 + minimum_distance = 30.0 + + with pytest.raises(ValueError): + generate_unit_locations( + num_units, channel_locations, minimum_distance=minimum_distance, distance_strict=True, seed=seed + ) + + with pytest.warns(UserWarning): + unit_locations = generate_unit_locations( + num_units, channel_locations, minimum_distance=minimum_distance, distance_strict=False, seed=seed + ) + assert unit_locations.shape == (num_units, 3) + + def test_generate_templates(): seed = 0