Skip to content

Speed up the travelling salesman sort with bit-identical output #44

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bendichter wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Speed up the travelling salesman sort with bit-identical output #44

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
197 changes: 143 additions & 54 deletions rastermap/sort.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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


Expand All @@ -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)
Expand All @@ -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]
Expand Down Expand Up @@ -195,16 +277,19 @@ 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]
end_pos = end_pos[:iter_completed]
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])
Expand All @@ -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]
Expand All @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down