From 93eb092b78bcb4ef3294d445aad8c9a815865813 Mon Sep 17 00:00:00 2001 From: Ben Dichter Date: Thu, 16 Jul 2026 11:17:43 -0400 Subject: [PATCH] Score candidate moves without materializing the permuted matrix The travelling salesman sort dominates rastermap's runtime. On the test dataset in tests/ (569 neurons by 27680 timepoints), a default fit took about 6.5 seconds, of which the PCA accounted for about 0.3 seconds and essentially all of the remainder was the sort. The sort operates on the cluster correlation matrix, so its cost depends on n_clusters rather than on the size of the input, and it stays the bottleneck as datasets grow. The cost was concentrated in new_correlation, which is called once per candidate move. To score a candidate it built a full permuted copy of the correlation matrix with cc[ishift][:, ishift], summed that against BBt, and then built a second full copy to test the reversed ordering. That is four O(n^2) passes and two allocations per candidate, and there are thousands of candidates per iteration. This computes the same quantity, sum over i,j of cc[perm[i], perm[j]] * BBt[i,j], directly from index arithmetic, without ever building the permuted matrix. Two properties make this exact rather than approximate. First, reversing the moved segment's rows and columns is the same operation as reversing that segment of the permutation, so the reversed ordering costs one index reversal instead of a matrix copy. Second, BBt contains exact zeros, roughly half of it on the default path because it is upper triangular, and nearly all of it when locality is high. Skipping those entries is bit-exact, because adding zero to a float32 accumulator initialized to +0.0 is a no-op. The accumulation order over the remaining terms is unchanged. Measured on the test dataset, rounded because repeated runs vary by a few percent to twenty percent: default 6.5s -> 2.0s about 3x locality=0.75, time_lag_window=5 11.8s -> 3.1s about 4x locality=1.0, time_lag_window=10 10.7s -> 1.1s about 10x grid_upsample=0 6.7s -> 2.0s about 3.5x n_splits=3, n_clusters=50, sticky=False 1.6s -> 1.1s about 1.5x n_splits=2, n_clusters=20, nc_splits=10 0.5s -> 0.5s none The gain is largest when locality is high, because BBt is then close to banded and almost the whole sum is skipped. It is negligible for small n_splits configurations, where the sorted matrices are small enough that the removed allocations did not cost much to begin with. Output is bit-identical, not merely close. Verified in 52 comparisons against the previous implementation with no mismatches: six configurations on the test dataset, eleven Rastermap configurations on synthetic data, and 35 cases exercising traveling_salesman and matrix_matching directly, including locality of infinity, the circular basis, a dense user-supplied BBt, a BBt with scattered interior zeros, and an asymmetric cc. In every case the embedding, the sort order, the sorted correlation matrix and the cluster centers match exactly. There is a further factorization available for the n_splits path, since corr_permuted_rows permutes rows but not columns, but it changes the floating point accumulation order and so is left out here. One precondition worth recording: skipping the zeros assumes cc is finite, since 0.0 * inf is nan. That holds for a correlation matrix and for the test data, but it is not checked. --- rastermap/sort.py | 197 +++++++++++++++++++++++++++++++++------------- 1 file changed, 143 insertions(+), 54 deletions(-) diff --git a/rastermap/sort.py b/rastermap/sort.py index 581dac5..3258561 100644 --- a/rastermap/sort.py +++ b/rastermap/sort.py @@ -43,19 +43,103 @@ def elementwise_mult_sum(x, y): return (x * y).sum() -@njit("int64[:] (int64, int64, int64, int64)", nogil=True, cache=True) -def shift_inds(i0, i1, inode, n_nodes): - """ shift segment from i0->i1 to position inode""" - n_seg = i1 - i0 +def bbt_row_bounds(BBt): + """ first and last nonzero column of each row of BBt + + The matching kernels weight every entry of a permuted correlation matrix by + BBt, so any column where BBt is exactly zero contributes nothing to the sum + and can be skipped. BBt is upper triangular in the default sorting path, + so this halves the work; for a dense user-supplied BBt the bounds simply + cover every column and nothing is skipped. + """ + n_rows = BBt.shape[0] + j0 = np.zeros(n_rows, np.int64) + j1 = np.zeros(n_rows, np.int64) + for i in range(n_rows): + nz = np.flatnonzero(BBt[i] != 0) + if len(nz) > 0: + j0[i] = nz[0] + j1[i] = nz[-1] + 1 + return j0, j1 + + +@njit("float32 (float32[:,:], int64[:], float32[:,:], int64[:], int64[:])", nogil=True, + cache=True) +def corr_permuted(cc, perm, BBt, j0, j1): + """ sum over i,j of cc[perm[i], perm[j]] * BBt[i, j] + + Equivalent to elementwise_mult_sum(cc[perm][:, perm], BBt) but computed + without materializing the permuted matrix, which is the dominant cost when + scoring candidate moves. + """ + total = np.float32(0.) + for i in range(perm.shape[0]): + cc_row = cc[perm[i]] + BBt_row = BBt[i] + for j in range(j0[i], j1[i]): + total += cc_row[perm[j]] * BBt_row[j] + return total + + +@njit("float32 (float32[:,:], int64[:], float32[:,:], int64[:], int64[:])", nogil=True, + cache=True) +def corr_permuted_rows(cc, perm, BBt, j0, j1): + """ sum over i,j of cc[perm[i], j] * BBt[i, j] (rows permuted, columns fixed) """ + total = np.float32(0.) + for i in range(perm.shape[0]): + cc_row = cc[perm[i]] + BBt_row = BBt[i] + for j in range(j0[i], j1[i]): + total += cc_row[j] * BBt_row[j] + return total + + +@njit("(int64[:], int64, int64, int64, int64)", nogil=True, cache=True) +def fill_shift_inds(inds, i0, i1, inode, n_nodes): + """ write the shift_inds permutation into a preallocated buffer """ l_seg = inode - i0 + p = 0 if l_seg >= 0: - inds = np.concatenate((np.arange(i0), np.arange(i1, i1 + l_seg), - np.arange(i0, i1), np.arange(i1 + l_seg, n_nodes))) + for k in range(0, i0): + inds[p] = k + p += 1 + for k in range(i1, i1 + l_seg): + inds[p] = k + p += 1 + for k in range(i0, i1): + inds[p] = k + p += 1 + for k in range(i1 + l_seg, n_nodes): + inds[p] = k + p += 1 else: - inds = np.concatenate( - (np.arange(inode), np.arange(i0, i1), np.arange(inode, - i0), np.arange(i1, - n_nodes))) + for k in range(0, inode): + inds[p] = k + p += 1 + for k in range(i0, i1): + inds[p] = k + p += 1 + for k in range(inode, i0): + inds[p] = k + p += 1 + for k in range(i1, n_nodes): + inds[p] = k + p += 1 + + +@njit("(int64[:], int64, int64)", nogil=True, cache=True) +def reverse_segment(inds, start, length): + for k in range(length // 2): + tmp = inds[start + k] + inds[start + k] = inds[start + length - 1 - k] + inds[start + length - 1 - k] = tmp + + +@njit("int64[:] (int64, int64, int64, int64)", nogil=True, cache=True) +def shift_inds(i0, i1, inode, n_nodes): + """ shift segment from i0->i1 to position inode""" + inds = np.empty(n_nodes, np.int64) + fill_shift_inds(inds, i0, i1, inode, n_nodes) return inds @@ -82,38 +166,35 @@ def shift_matrix(cc, i0, i1, inode, n_nodes, ordering): return cc2, ishift -@njit("float32[:] (float32[:,:], int64, int64, int64, int64, float32[:,:])", nogil=True, - cache=True) -def new_correlation(cc, i0, i1, inode, n_nodes, BBt): - """ compute correlation change of moving segment i0:i1+1 to position inode - inode=-1 is at beginning of sequence +@njit("UniTuple(float32, 2) (float32[:,:], int64, int64, int64, int64, float32[:,:], int64[:], int64[:])", + nogil=True, cache=True) +def new_correlation(cc, i0, i1, inode, n_nodes, BBt, j0, j1): + """ compute correlation of moving segment i0:i1 to position inode + + Reversing the rows and columns of the moved segment is the same as reversing + that segment of the permutation, so both orderings are scored straight from + cc without building either permuted matrix. """ - cc_new = shift_matrix_forward(cc, i0, i1, inode, n_nodes)[0] + perm = np.empty(n_nodes, np.int64) + fill_shift_inds(perm, i0, i1, inode, n_nodes) ilength = i1 - i0 - corr_new = elementwise_mult_sum(cc_new, BBt) - ordering = 0 + corr_new = corr_permuted(cc, perm, BBt, j0, j1) + ordering = np.float32(0.) if ilength > 1: - cc0 = cc_new.copy() - cc0[inode:inode + ilength] = cc0[inode:inode + ilength][::-1] - cc0[:, inode:inode + ilength] = cc0[:, inode:inode + ilength][:, ::-1] - - corr_new0 = elementwise_mult_sum(cc0, BBt) + reverse_segment(perm, inode, ilength) + corr_new0 = corr_permuted(cc, perm, BBt, j0, j1) if corr_new0 > corr_new: corr_new = corr_new0 - ordering = 1 - - corr_out = np.zeros(2, np.float32) - corr_out[0] = corr_new - corr_out[1] = ordering + ordering = np.float32(1.) - return corr_out + return corr_new, ordering -@njit("(float32[:,:], int64, int64, int64, float32[:,:], boolean)", nogil=True, - parallel=True, cache=True) -def tsp_fast(cc, n_iter, n_nodes, n_skip, BBt, verbose): +@njit("(float32[:,:], int64, int64, int64, float32[:,:], boolean, int64[:], int64[:])", + nogil=True, parallel=True, cache=True) +def tsp_fast(cc, n_iter, n_nodes, n_skip, BBt, verbose, j0, j1): inds = np.arange(0, n_nodes).astype(np.int32) iinds = np.arange(0, n_nodes).astype(np.int32) seg_len = -1 * np.ones(n_iter) @@ -140,10 +221,11 @@ def tsp_fast(cc, n_iter, n_nodes, n_skip, BBt, verbose): inode = (ix % n_test) * n_skip + k % n_skip i1 = (i0 + ilength) if i1 <= n_nodes and inode + ilength <= n_nodes: - new_corr = new_correlation(cc, i0, i1, inode, n_nodes, BBt) - corr_change = new_corr[0] - corr_orig + corr_val, ordering_val = new_correlation(cc, i0, i1, inode, n_nodes, + BBt, j0, j1) + corr_change = corr_val - corr_orig corr_change_seg[ix] = corr_change - ordering_seg[ix] = new_corr[1] + ordering_seg[ix] = ordering_val #corr_change_seg += np.random.randn(corr_change_seg.shape)*(corr_change_seg**2).mean() ix = corr_change_seg.argmax() corr_change = corr_change_seg[ix] @@ -195,8 +277,11 @@ def traveling_salesman(cc, n_iter=400, locality=0.0, circular=False, n_iter = np.int64(n_iter) + BBt32 = BBt.astype(np.float32) + j0, j1 = bbt_row_bounds(BBt32) + cc, inds, seg_len, start_pos, end_pos, flipped = tsp_fast( - cc, n_iter, n_nodes, n_skip, BBt.astype(np.float32), verbose) + cc, n_iter, n_nodes, n_skip, BBt32, verbose, j0, j1) iter_completed = np.nonzero(seg_len == -1)[0][0] seg_len = seg_len[:iter_completed] start_pos = start_pos[:iter_completed] @@ -204,7 +289,7 @@ def traveling_salesman(cc, n_iter=400, locality=0.0, circular=False, flipped = flipped[:iter_completed] if n_skip > 1: cc, inds2, seg_len2, start_pos2, end_pos2, flipped2 = tsp_fast( - cc, n_iter, n_nodes, 1, BBt.astype(np.float32), verbose) + cc, n_iter, n_nodes, 1, BBt32, verbose, j0, j1) iter_completed = np.nonzero(seg_len2 == -1)[0][0] seg_len = np.append(seg_len, seg_len2[:iter_completed]) start_pos = np.append(start_pos, start_pos2[:iter_completed]) @@ -215,11 +300,11 @@ def traveling_salesman(cc, n_iter=400, locality=0.0, circular=False, return cc, inds, BBt, seg_len, start_pos, end_pos, flipped -@njit("int32[:] (int64, int64, int64, int64, int64)", nogil=True, cache=True) +@njit("int64[:] (int64, int64, int64, int64, int64)", nogil=True, cache=True) def shift_inds_sub(i0, i1, inode, isforward, n_nodes): """ shift segment from i0->i1 to position inode""" n_seg = i1 - i0 + 1 - inds = np.arange(0, n_nodes, 1, np.int32) + inds = np.arange(0, n_nodes, 1, np.int64) inds0 = inds.copy() if inode > i0: inds[i0:i0 + inode - i1] = inds0[i1 + 1:inode + 1] @@ -234,37 +319,37 @@ def shift_inds_sub(i0, i1, inode, isforward, n_nodes): return inds -@njit("(float32[:,:], int32[:])", nogil=True, cache=True) +@njit("(float32[:,:], int64[:])", nogil=True, cache=True) def shift_matrix_sub(cc, ishift): return cc[ishift][:, ishift] -@njit("(float32[:,:], int32[:])", nogil=True, cache=True) +@njit("(float32[:,:], int64[:])", nogil=True, cache=True) def shift_rows(cc, ishift): return cc[ishift] w_add = 4 @njit( - "float32 (float32[:,:], float32[:,:], int64, int64, int64, int64, int64, float32[:,:], float32[:,:])", + "float32 (float32[:,:], float32[:,:], int64, int64, int64, int64, int64, float32[:,:], float32[:,:], int64[:], int64[:], int64[:], int64[:])", nogil=True, cache=True) -def new_correlation_sub(cc, cc_add, i0, i1, inode, n_nodes, isforward, BBt, BBt_add): +def new_correlation_sub(cc, cc_add, i0, i1, inode, n_nodes, isforward, BBt, BBt_add, j0, + j1, j0_add, j1_add): """ compute correlation change of moving segment i0:i1+1 to position inode inode=-1 is at beginning of sequence """ ishift = shift_inds_sub(i0, i1, inode, isforward, n_nodes) - cc2 = shift_matrix_sub(cc, ishift) - cc_add2 = shift_rows(cc_add, ishift) - corr_new = elementwise_mult_sum(cc2, - BBt) + w_add * elementwise_mult_sum(cc_add2, BBt_add) + corr_new = (corr_permuted(cc, ishift, BBt, j0, j1) + + w_add * corr_permuted_rows(cc_add, ishift, BBt_add, j0_add, j1_add)) return corr_new @njit( - "(float32[:,:], float32[:,:], int64, int64, int64, float32[:,:],float32[:,:], boolean)", + "(float32[:,:], float32[:,:], int64, int64, int64, float32[:,:],float32[:,:], boolean, int64[:], int64[:], int64[:], int64[:])", nogil=True, parallel=True, cache=True) -def tsp_sub(cc, cc_add, n_iter, n_nodes, n_skip, BBt, BBt_add, verbose): - inds = np.arange(0, n_nodes).astype(np.int32) +def tsp_sub(cc, cc_add, n_iter, n_nodes, n_skip, BBt, BBt_add, verbose, j0, j1, j0_add, + j1_add): + inds = np.arange(0, n_nodes).astype(np.int64) seg_len = np.ones(n_iter) n_len = 2 n_test = n_nodes // n_skip + 1 @@ -288,10 +373,12 @@ def tsp_sub(cc, cc_add, n_iter, n_nodes, n_skip, BBt, BBt_add, verbose): if i1 < n_nodes and (inode < i0 or inode > i1 + 1) and inode < n_nodes: isforward = 1 new_corr = new_correlation_sub(cc, cc_add, i0, i1, inode, n_nodes, - 1, BBt, BBt_add) + 1, BBt, BBt_add, j0, j1, j0_add, + j1_add) if seg_length > 0: new_corr_backward = new_correlation_sub( - cc, cc_add, i0, i1, inode, n_nodes, 0, BBt, BBt_add) + cc, cc_add, i0, i1, inode, n_nodes, 0, BBt, BBt_add, j0, j1, + j0_add, j1_add) if new_corr < new_corr_backward: isforward = 0 new_corr = new_corr_backward @@ -340,11 +427,13 @@ def matrix_matching(cc, BBt, cc_add, BBt_add, n_iter=400, n_skip=None, verbose=F BBt_add = BBt_add.astype(np.float32) n_iter = np.int64(n_iter) + j0, j1 = bbt_row_bounds(BBt) + j0_add, j1_add = bbt_row_bounds(BBt_add) cc, inds, seg_len = tsp_sub(cc, cc_add, n_iter, n_nodes, n_skip, BBt, BBt_add, - verbose) + verbose, j0, j1, j0_add, j1_add) if n_skip > 1: cc, inds2, seg_len2 = tsp_sub(cc, cc_add, n_iter, n_nodes, 1, BBt, BBt_add, - verbose) + verbose, j0, j1, j0_add, j1_add) inds = inds[inds2] return cc, inds, seg_len