From 613b9119860c4685b7ba88dfcfe0054971321496 Mon Sep 17 00:00:00 2001 From: Eric Larson Date: Thu, 27 Aug 2026 10:34:31 +0200 Subject: [PATCH] Speed up label adjacency computation --- mne/label.py | 268 ++++++++++++++---------------- mne/morph.py | 20 ++- mne/source_space/_source_space.py | 24 ++- mne/tests/test_label.py | 23 ++- 4 files changed, 173 insertions(+), 162 deletions(-) diff --git a/mne/label.py b/mne/label.py index 3eb921ba6be..f0e31e78316 100644 --- a/mne/label.py +++ b/mne/label.py @@ -1710,36 +1710,13 @@ def _verts_within_dist(graph, sources, max_dist): dist : array Distances from source vertex. """ - dist_map = {} - verts_added_last = [] - for source in sources: - dist_map[source] = 0 - verts_added_last.append(source) - - # add neighbors until no more neighbors within max_dist can be found - while len(verts_added_last) > 0: - verts_added = [] - for i in verts_added_last: - v_dist = dist_map[i] - row = graph[[i], :] - neighbor_vert = row.indices - neighbor_dist = row.data - for j, d in zip(neighbor_vert, neighbor_dist): - n_dist = v_dist + d - if j in dist_map: - if n_dist < dist_map[j]: - dist_map[j] = n_dist - else: - if n_dist <= max_dist: - dist_map[j] = n_dist - # we found a new vertex within max_dist - verts_added.append(j) - verts_added_last = verts_added - - verts = np.sort(np.array(list(dist_map.keys()), int)) - dist = np.array([dist_map[v] for v in verts], int) + from scipy.sparse.csgraph import dijkstra - return verts, dist + # ``min_only`` gives the distance to the closest of ``sources`` for every vertex, + # and ``limit`` leaves the ones beyond max_dist at infinity + dist = dijkstra(graph, indices=sources, min_only=True, limit=max_dist) + verts = np.flatnonzero(np.isfinite(dist)) + return verts, dist[verts].astype(int) def _grow_labels(seeds, extents, hemis, names, dist, vert, subject): @@ -1931,6 +1908,8 @@ def _grow_nonoverlapping_labels( subject, seeds_, extents_, hemis, vertices_, graphs, names_ ): """Grow labels while ensuring that they don't overlap.""" + from scipy.sparse.csgraph import dijkstra + labels = [] for hemi in set(hemis): hemi_index = hemis == hemi @@ -1941,56 +1920,34 @@ def _grow_nonoverlapping_labels( n_vertices = len(vertices_[hemi]) n_labels = len(seeds) - # prepare parcellation - parc = np.empty(n_vertices, dtype="int32") - parc[:] = -1 - - # initialize active sources - sources = {} # vert -> (label, dist_from_seed) - edge = [] # queue of vertices to process + # which label each seed vertex belongs to + seed_label = np.full(n_vertices, -1, int) for label, seed in enumerate(seeds): - if np.any(parc[seed] >= 0): + if np.any(seed_label[seed] >= 0): raise ValueError("Overlapping seeds") - parc[seed] = label - for s in np.atleast_1d(seed): - sources[s] = (label, 0.0) - edge.append(s) - - # grow from sources - while edge: - vert_from = edge.pop(0) - label, old_dist = sources[vert_from] - - # add neighbors within allowable distance - row = graph[[vert_from], :] - for vert_to, dist in zip(row.indices, row.data): - # Prevent adding a point that has already been used - # (prevents infinite loop) - if (vert_to == seeds[label]).any(): - continue - new_dist = old_dist + dist - - # abort if outside of extent - if new_dist > extents[label]: - continue - - vert_to_label = parc[vert_to] - if vert_to_label >= 0: - _, vert_to_dist = sources[vert_to] - # abort if the vertex is occupied by a closer seed - if new_dist > vert_to_dist: - continue - elif vert_to in edge: - edge.remove(vert_to) - - # assign label value - parc[vert_to] = label - sources[vert_to] = (label, new_dist) - edge.append(vert_to) + seed_label[seed] = label + + # A multi-source Dijkstra with min_only gives, for every vertex, its distance + # to the closest seed vertex and which seed vertex that was, i.e. exactly the + # non-overlapping assignment we want: each label grows outward until it runs + # into a vertex that some other label reaches sooner. + dist, _, sources = dijkstra( + graph, + indices=np.flatnonzero(seed_label >= 0), + min_only=True, + return_predecessors=True, + ) + parc = np.full(n_vertices, -1, int) + reached = np.isfinite(dist) + parc[reached] = seed_label[sources[reached]] + # then drop vertices that are further away than their own label's extent + parc[reached] = np.where( + dist[reached] <= extents[parc[reached]], parc[reached], -1 + ) # convert parc to labels for i in range(n_labels): - vertices = np.nonzero(parc == i)[0] + vertices = np.flatnonzero(parc == i) name = str(names[i]) label_ = Label(vertices, hemi=hemi, name=name, subject=subject) labels.append(label_) @@ -2118,16 +2075,21 @@ def _cortex_parcellation(subject, n_parcel, hemis, vertices_, graphs, rng): rest -= 1 # merging small labels - # label adjacency matrix + # label adjacency matrix: every vertex belongs to exactly one label at this + # point, so mapping both ends of each graph edge through parc marks all pairs + # of adjacent labels at once. Functionally equivalent to, but much faster than: + # + # for i in range(n_labels): + # vertices = np.nonzero(parc == i)[0] + # label_sizes[i] = len(vertices) + # neighbor_labels = np.unique(parc[graph[vertices, :].indices]) + # label_conn[i, neighbor_labels] = 1 + # n_labels = label_idx + 1 - label_sizes = np.empty(n_labels, dtype=int) + label_sizes = np.bincount(parc, minlength=n_labels) label_conn = np.zeros([n_labels, n_labels], dtype="bool") - for i in range(n_labels): - vertices = np.nonzero(parc == i)[0] - label_sizes[i] = len(vertices) - neighbor_vertices = graph[vertices, :].indices - neighbor_labels = np.unique(np.array(parc[neighbor_vertices])) - label_conn[i, neighbor_labels] = 1 + edges = graph.tocoo() + label_conn[parc[edges.row], parc[edges.col]] = True np.fill_diagonal(label_conn, 0) # merging @@ -2616,30 +2578,53 @@ def _check_values_labels(values, n_labels): ) -def _labels_to_stc_surf(labels, values, tmin, tstep, subject): +def _label_membership(label_indices, n_vertices): + """Get a sparse (n_labels, n_vertices) matrix of which vertices are in each label. + + ``label_indices[li]`` holds the indices, into some length-``n_vertices`` array of + vertices, of the vertices belonging to label ``li``. Multiplying by this matrix (or + its transpose) is how the label-wise operations in this module avoid looping over + labels; see ``_labels_to_stc_surf`` and ``_label_adjacency`` for examples. + """ from scipy import sparse + n_labels = len(label_indices) + rows = np.repeat(np.arange(n_labels), [len(ind) for ind in label_indices]) + cols = np.concatenate([np.empty(0, int)] + list(label_indices)) + return sparse.csr_array( + (np.ones(len(cols)), (rows, cols)), shape=(n_labels, n_vertices) + ) + + +def _labels_to_stc_surf(labels, values, tmin, tstep, subject): subject = _check_labels_subject(labels, subject, "subject") _check_values_labels(values, len(labels)) - vertices = dict(lh=[], rh=[]) - data = dict(lh=[], rh=[]) - for li, label in enumerate(labels): - data[label.hemi].append( - np.repeat(values[li][np.newaxis], len(label.vertices), axis=0) + vertices = list() + data = list() + for hemi in ("lh", "rh"): + idx = np.array( + [li for li, label in enumerate(labels) if label.hemi == hemi], int + ) + label_vertices = [labels[li].vertices for li in idx] + these_vertices = np.unique(np.concatenate([np.empty(0, int)] + label_vertices)) + membership = _label_membership( + [np.searchsorted(these_vertices, verts) for verts in label_vertices], + len(these_vertices), ) - vertices[label.hemi].append(label.vertices) - hemis = ("lh", "rh") - for hemi in hemis: - vertices[hemi] = np.concatenate(vertices[hemi], axis=0) - data[hemi] = np.concatenate(data[hemi], axis=0).astype(float) - cols = np.arange(len(vertices[hemi])) - vertices[hemi], rows = np.unique(vertices[hemi], return_inverse=True) - mat = sparse.coo_array((np.ones(len(rows)), (rows, cols))).tocsr() - mat *= 1.0 / mat.sum(axis=-1) - data[hemi] = mat @ data[hemi] - vertices = [vertices[hemi] for hemi in hemis] - data = np.concatenate([data[hemi] for hemi in hemis], axis=0) - return data, vertices, subject + # ``membership.T @ values[idx]`` sums, for each vertex, the values of every + # label containing it, and ``n_used`` is how many labels that was, so the two + # together give the mean. Functionally equivalent to, but much faster than: + # + # for vi, vertex in enumerate(these_vertices): + # in_label = [ + # ii for ii, verts in enumerate(label_vertices) if vertex in verts + # ] + # this_data[vi] = values[idx[in_label]].mean(axis=0) + # + n_used = membership.sum(axis=0) # number of labels containing each vertex + vertices.append(these_vertices) + data.append((membership.T @ values[idx]) / n_used[:, np.newaxis]) + return np.concatenate(data), vertices, subject _DEFAULT_TABLE_NAME = "MNE-Python Colortable" @@ -3048,6 +3033,37 @@ def select_sources( return new_label +def _label_adjacency(label_src_ind, src_adjacency): + """Turn per-label source space indices plus source adjacency into label adjacency. + + Two labels are adjacent if any vertex of one is adjacent to any vertex of the + other. ``label_src_ind[li]`` holds the indices into the source space of the + vertices belonging to label ``li``. + """ + from scipy import sparse + + n_labels = len(label_src_ind) + membership = _label_membership(label_src_ind, src_adjacency.shape[0]) + # ``counts[i, j]`` is the number of adjacent (vertex in label i, vertex in label j) + # vertex pairs, which is nonzero exactly when the two labels are adjacent. + # Functionally equivalent to, but much faster than: + # + # src_adjacency = src_adjacency.tocsr() + # counts = np.zeros((n_labels, n_labels)) + # for i in range(n_labels): + # for j in range(n_labels): + # counts[i, j] = src_adjacency[label_src_ind[i]][ + # :, label_src_ind[j] + # ].sum() + # + # which needs O(n_labels ** 2) sparse slices, each costing O(nnz(src_adjacency)). + counts = membership @ src_adjacency.tocsr() @ membership.T + row, col = counts.nonzero() # ignores any explicitly stored zeros + return sparse.coo_matrix( + (np.ones(len(row)), (row, col)), shape=(n_labels, n_labels) + ) + + def label_adjacency(labels, src): """Compute adjacency between labels. @@ -3075,9 +3091,7 @@ def label_adjacency(labels, src): ----- .. versionadded:: 1.13 """ - from scipy.sparse import coo_matrix - - src_adjacency = spatial_src_adjacency(src).tocsr() + src_adjacency = spatial_src_adjacency(src) label_src_ind = list() for label in labels: src_hemi = src[0] if label.hemi == "lh" else src[1] @@ -3086,22 +3100,10 @@ def label_adjacency(labels, src): if label.hemi == "rh": src_ind += src[0]["nuse"] label_src_ind.append(src_ind) - - adjacent_label_inds = list() # list of pairs of label indices - for ind1, label1 in enumerate(labels): - for ind2, label2 in enumerate(labels): - # If the labels are on different hemispheres, they are not adjacent. - if label1.hemi != label2.hemi: - continue - - # Get adjacent vertices if any. - adj_verts = src_adjacency[label_src_ind[ind1], :][:, label_src_ind[ind2]] - if adj_verts.data.any(): - adjacent_label_inds.append((ind1, ind2)) - return coo_matrix( - (np.ones(len(adjacent_label_inds)), tuple(zip(*adjacent_label_inds))), - shape=(len(labels), len(labels)), - ) + # Labels in different hemispheres are never adjacent because the source space + # adjacency has no inter-hemispheric edges, so no explicit hemisphere check is + # needed here. + return _label_adjacency(label_src_ind, src_adjacency) @fill_doc @@ -3136,8 +3138,6 @@ def volume_label_adjacency(src, subject, subjects_dir, *, aseg="auto", labels=No ----- .. versionadded:: 1.13 """ - from scipy import sparse - subjects_dir = Path(get_subjects_dir(subjects_dir, raise_error=True)) if aseg == "auto": # use aparc+aseg if auto aseg = _check_fname( @@ -3154,22 +3154,8 @@ def volume_label_adjacency(src, subject, subjects_dir, *, aseg="auto", labels=No labels = get_volume_labels_from_aseg(aseg) vol_labels = _volume_labels(src, (aseg, labels), mri_resolution=False) - src_adjacency = spatial_src_adjacency(src).tocsr() - - label_verts = list() - for label in vol_labels: - label_verts.append(np.searchsorted(src[0]["vertno"], label.vertices)) - - adjacent_label_inds = list() # list of pairs of label indices - for ind1, verts1 in enumerate(label_verts): - for ind2, verts2 in enumerate(label_verts): - # Get adjacent vertices if any. - adj_verts = src_adjacency[verts1, :][:, verts2] - if adj_verts.data.any(): - adjacent_label_inds.append((ind1, ind2)) - - adj = sparse.coo_matrix( - (np.ones(len(adjacent_label_inds)), tuple(zip(*adjacent_label_inds))), - shape=(len(labels), len(labels)), - ) - return adj, labels + src_adjacency = spatial_src_adjacency(src) + label_src_ind = [ + np.searchsorted(src[0]["vertno"], label.vertices) for label in vol_labels + ] + return _label_adjacency(label_src_ind, src_adjacency), labels diff --git a/mne/morph.py b/mne/morph.py index 640ec0b2fa7..aeeaddde35b 100644 --- a/mne/morph.py +++ b/mne/morph.py @@ -1425,12 +1425,20 @@ def _surf_upsampling_mat(idx_from, e, smooth): def _sparse_argmax_nnz_row(csr_mat): """Return index of the maximum non-zero index in each row.""" - n_rows = csr_mat.shape[0] - idx = np.empty(n_rows, dtype=np.int64) - for k in range(n_rows): - row = csr_mat[[k]].tocoo() - idx[k] = row.col[np.argmax(row.data)] - return idx + # Equivalent to, but orders of magnitude faster than, slicing out each row and + # taking ``row.indices[np.argmax(row.data)]``: + # + # for k in range(csr_mat.shape[0]): + # row = csr_mat[[k]].tocoo() + # idx[k] = row.col[np.argmax(row.data)] + # + starts, counts = csr_mat.indptr[:-1], np.diff(csr_mat.indptr) + assert (counts > 0).all() # reduceat below needs every row to be non-empty + row_max = np.maximum.reduceat(csr_mat.data, starts) + # position of the first stored entry in each row that attains the row maximum + is_max = csr_mat.data == np.repeat(row_max, counts) + pos = np.where(is_max, np.arange(len(is_max)), len(is_max)) + return csr_mat.indices[np.minimum.reduceat(pos, starts)] def _get_subject_sphere_tris(subject, subjects_dir): diff --git a/mne/source_space/_source_space.py b/mne/source_space/_source_space.py index 3ef6a9eb962..4d091988b41 100644 --- a/mne/source_space/_source_space.py +++ b/mne/source_space/_source_space.py @@ -2811,15 +2811,25 @@ def add_source_space_distances(src, dist_limit=np.inf, n_jobs=None, *, verbose=N min_idx = min_idx[midx, range_idx] min_dists.append(min_dist) min_idxs.append(min_idx) - # convert to sparse representation + # Convert to sparse representation. Deriving the row/column vertex + # numbers from the flat indices of the entries we keep -- rather than + # from np.meshgrid(vertno, vertno), whose two dense (n_use, n_use) + # index arrays are over 800 MB apiece for an ico-5 source space -- + # and indexing in the narrowest safe dtype roughly halves the peak + # memory of this block. Arrays are freed as soon as they are consumed + # for the same reason. + n_use = len(s["vertno"]) d = np.concatenate([dd[0] for dd in d]).ravel() # already float32 - idx = d > 0 - d = d[idx] - i, j = np.meshgrid(s["vertno"], s["vertno"]) - i = i.ravel()[idx] - j = j.ravel()[idx] + idx_dtype = np.int32 if d.size <= np.iinfo(np.int32).max else np.int64 + vertno = s["vertno"].astype(idx_dtype) + idx = np.flatnonzero(d).astype(idx_dtype, copy=False) # 0 == not computed + data = d[idx] + del d + row = vertno[idx % n_use] + col = vertno[idx // n_use] + del idx, vertno s["dist"] = csr_array( - (d, (i, j)), shape=(s["np"], s["np"]), dtype=np.float32 + (data, (row, col)), shape=(s["np"], s["np"]), dtype=np.float32 ) s["dist_limit"] = np.array([dist_limit], np.float32) diff --git a/mne/tests/test_label.py b/mne/tests/test_label.py index b485b4a476f..a820673804e 100644 --- a/mne/tests/test_label.py +++ b/mne/tests/test_label.py @@ -58,8 +58,8 @@ real_label_fname = data_path / "MEG" / "sample" / "labels" / "Aud-lh.label" v1_label_fname = subjects_dir / "sample" / "label" / "lh.V1.label" +fname_src = data_path / "subjects" / "sample" / "bem" / "sample-oct-4-src.fif" fname_vsrc = data_path / "MEG" / "sample" / "sample_audvis_trunc-meg-vol-7-fwd.fif" -fname_src_fs = data_path / "subjects" / "fsaverage" / "bem" / "fsaverage-ico-5-src.fif" fwd_fname = data_path / "MEG" / "sample" / "sample_audvis_trunc-meg-eeg-oct-6-fwd.fif" src_bad_fname = data_path / "subjects" / "fsaverage" / "bem" / "fsaverage-ico-5-src.fif" @@ -466,6 +466,16 @@ def test_labels_to_stc(): for value, label in zip(values, labels): stc_label = stc.in_label(label) assert (stc_label.data == value).all() + # labels from a single hemisphere, and vertices shared by multiple labels + # (which get averaged, see the docstring) + lh = [ + Label(np.arange(3), hemi="lh", subject="sample"), + Label(np.arange(2, 5), hemi="lh", subject="sample"), + ] + stc = labels_to_stc(lh, np.array([1.0, 3.0])) + assert_array_equal(stc.vertices[0], np.arange(5)) + assert_array_equal(stc.vertices[1], []) + assert_array_equal(stc.data[:, 0], [1.0, 1.0, 2.0, 3.0, 3.0]) stc = read_source_estimate(stc_fname, "sample") @@ -1351,22 +1361,19 @@ def test_volume_label_adjacency(): def test_label_adjacency(): """Test label adjacency.""" pytest.importorskip("nibabel") - src = read_source_spaces(fname_src_fs) - mne.add_source_space_distances(src, dist_limit=0.01, n_jobs=-1) - + src = read_source_spaces(fname_src) labels = mne.read_labels_from_annot( - subject="fsaverage", + subject="sample", subjects_dir=subjects_dir, ) adj = mne.label_adjacency(labels, src) n_neighbors = adj.sum(axis=1) - assert_equal(len(labels), 69) # default number of labels in aseg.mgz + assert_equal(len(labels), 68) # default number of labels in aseg.mgz assert_equal(adj.shape, (len(labels), len(labels))) - assert_equal(n_neighbors.min(), 0) - assert_equal(np.sum(n_neighbors == 0), 1) + assert n_neighbors.min() == 3 input_labels = [ "cuneus-lh",