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