Use grid accelerated Poisson disk sampling in generate_unit_locations - #4730
Use grid accelerated Poisson disk sampling in generate_unit_locations#4730adityasingh2400 wants to merge 2 commits into
Conversation
The old sampler drew every unit at once and then enforced minimum_distance by rejection, rebuilding the full dense (N, N, 3) broadcast difference and reducing it to an (N, N) distance matrix on every pass. With no spatial index that is O(N^2) in both time and peak memory per pass, and the loop only warns when it runs out of iterations, so past a moderate packing density callers silently got unit locations that broke the constraint minimum_distance promises. Bridson's algorithm with a uniform grid over the box only ever compares a candidate against the points in the cells around it, so filling the box is linear in the number of points placed. The box is filled and num_units points are then drawn from that pool, which keeps the signature, the distance_strict raise and the warning text unchanged. Fixes SpikeInterface#3419
The Poisson disk packing changed how far apart units end up, not just the minimum they respect. On the small probe used by the postprocessing tests it put the first two units 36um apart in xy where the previous sampler put them 17um apart, so their sparsity no longer overlapped and merging them produced a unit with no channels. Restore the original draw-and-redraw loop, which only pushes units as far apart as minimum_distance requires, and replace just the pairwise distance matrix with a uniform grid lookup. That removes the quadratic cost the issue was about while leaving the output identical to before. Poisson disk sampling is kept as the fallback for the dense case the redraw loop cannot solve, which is the case the old code silently returned invalid locations for.
|
Fixed the postprocessing failures. The cause was mine, not a flaky test. Switching to a Poisson disk packing did not just enforce So I put the original draw-and-redraw loop back, since it only ever pushes units as far apart as the constraint actually requires, and replaced only the part the issue was really about: the full pairwise distance matrix is now a uniform grid lookup in Poisson disk sampling is still here, as the fallback for when redrawing does not converge. That is the dense case Checks I ran:
|
1 participant
Fixes #3419.
generate_unit_locationsdrew every unit at once and then enforcedminimum_distanceby rejection. Each of up tomax_iterationpasses rebuilt the full dense(N, N, 3)broadcast difference, reduced it to an(N, N)distance matrix, and re-randomised the offenders. There is no spatial index, so every pass is O(N^2) in time and in peak memory, and the loop only warns when it runs out of iterations. Past a moderate packing density that means callers silently receive unit locations that violate the constraintminimum_distancepromises.This replaces that block with Bridson's Poisson disk sampling over the
(minimum_x, maximum_x) x (minimum_y, maximum_y) x (minimum_z, maximum_z)box. Accepted points go into a uniform grid whose cell size isminimum_distance / sqrt(3), so a cell holds at most one point and a candidate only has to be compared against the points in the small block of cells around it.num_unitsof the resulting points are then drawn withrng.choice(..., replace=False).Everything is pure numpy. No new dependency, in particular no scipy.
On a 384 contact Neuropixels shaped probe (4 columns, 20 um pitch) at the default
minimum_distance=20.0and the defaultmax_iteration=100, medians over 5 seeds:The worst pair column is the point of the issue. At 500 and 600 units the old sampler gave up and returned locations 3.5 um apart while
minimum_distancewas 20.The scaling is the other half. On a 6000 contact probe where nothing saturates:
Below a few hundred units the new sampler costs tens of milliseconds more, because a per point loop with a fixed numpy overhead cannot beat one small dense matrix multiply. Above that the quadratic term takes over and the new sampler wins, and the 200 MB
(N, N)matrix at N=5000 is gone entirely.Notes on behaviour:
distribution="multimodal"is preserved by weighting which pooled points to keep by the layer density rather than by moving points, so the distance constraint and the layering hold at the same time. Measured against the old sampler on the same probe, the fraction of units landing near a mode centre went from 0.62 to 0.74 at 200 units and from 0.57 to 0.67 at 300 units, against a flat baseline of 0.33.distance_strictraise and the warning text are unchanged, and when the box genuinely cannot holdnum_unitspoints that far apart the returned array is still(num_units, 3).minimum_distance=Nonekeeps the old plain uniform draw path untouched.max_iterationstays in the signature and now bounds how many times the sampling radius is reduced while looking for enough locations.ValueErrormessage for a baddistribution("must be 'uniform' or ") is completed, and it is now raised before any sampling work.Tested with a 1536 configuration sweep over 4 probe geometries, 6 unit counts, 4 z ranges, 4 minimum distances, 2 margins and both distributions, comparing against the old sampler given its most favourable setting of
max_iteration=500. The new sampler solves 41 configurations the old one could not, loses 1, and takes 73 s against 449 s in total. The single loss is a saturation boundary coin flip: on that configuration the old sampler solves 10 of 40 seeds and the new one solves 13 of 40, it just happened to lose on seed 0.test_generate_unit_locationsis untouched. Addedtest_generate_unit_locations_dense(550 units on a 384 contact probe withdistance_strict=True, plus box bounds),test_generate_unit_locations_multimodal(distance plus layer concentration against a flat baseline), andtest_generate_unit_locations_no_solution(raise and warn paths keep the shape). 550 rather than 500 because at the defaultmax_iteration=100the old sampler fails 500 on only 2 of 8 seeds but fails 550 and above on 8 of 8, so 550 is the reliable threshold. Box capacity there is 646 to 673 points across seeds, leaving headroom for CI variation.src/spikeinterface/core/tests/is green at 326 passed and 5 skipped. Ingeneration/tests/, threetest_template_database.pycases fail withModuleNotFoundError: No module named 's3fs', which I confirmed fail identically withgenerate.pyrestored frommain, so they are a missing optional dependency and unrelated.Disclosure: this change was prepared with AI assistance. I have reviewed and tested it.