WNN: fix segfault if multiple cells have identical coordinates

ilia-kats · ilia-kats · commit 31e472ecc3db · 2025-08-14T11:30:03.000+02:00
closes #158
diff --git a/muon/_core/preproc.py b/muon/_core/preproc.py
@@ -524,7 +524,7 @@ def neighbors(
             (
                 distances[:, 1:].reshape(-1),
                 nn_indices[:, 1:].reshape(-1),
-                np.concatenate((nn_indices[:, 0] * n_multineighbors, (nn_indices[:, 1:].size,))),
+                np.arange((nn_indices.shape[0] + 1) * n_multineighbors, step=n_multineighbors),
             ),
             shape=(rep.shape[0], rep.shape[0]),
         )

Original file line number	Diff line number	Diff line change
`@@ -524,7 +524,7 @@ def neighbors(`
`524`	`524`	`(`
`525`	`525`	`distances[:, 1:].reshape(-1),`
`526`	`526`	`nn_indices[:, 1:].reshape(-1),`
`527`		`- np.concatenate((nn_indices[:, 0] * n_multineighbors, (nn_indices[:, 1:].size,))),`
	`527`	`+ np.arange((nn_indices.shape[0] + 1) * n_multineighbors, step=n_multineighbors),`
`528`	`528`	`),`
`529`	`529`	`shape=(rep.shape[0], rep.shape[0]),`
`530`	`530`	`)`