Speed up the travelling salesman sort with bit-identical output#44
Open
bendichter wants to merge 1 commit into
Open
Speed up the travelling salesman sort with bit-identical output#44bendichter wants to merge 1 commit into
bendichter wants to merge 1 commit into
Conversation
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
This makes the travelling salesman sort roughly three times faster on the default settings, and close to ten times faster on the high locality settings used for sequence finding, while producing bit-identical output. The change is confined to
rastermap/sort.pyand does not alter the algorithm, only how each candidate move is scored.What Changed
new_correlationis called once per candidate move, and for each candidate it materialized a full permuted copy of the correlation matrix, summed it againstBBt, then materialized a second copy to test the reversed ordering. The replacement computes the same sum directly from index arithmetic without building either matrix.Two observations keep this exact. Reversing the moved segment's rows and columns is equivalent to reversing that segment of the permutation, which is an index relabeling rather than arithmetic. And
BBtcontains exact zeros, about half of it on the default path because it is upper triangular, so those terms can be skipped without changing the result.Results
Measured on the dataset in
tests/, on one 10 core machine:locality=0.75, time_lag_window=5locality=1.0, time_lag_window=10grid_upsample=0n_splits=3, n_clusters=50, sticky=Falsen_splits=2, n_clusters=20, nc_splits=10Repeated runs on the same machine vary by a few percent to twenty percent, and the high locality case landed anywhere between nine and eleven times faster depending on the run.
The high locality case improves the most because
BBtis then close to banded, so almost the entire sum is skipped rather than merely half of it. That is the regime the sequence finding settings target, so the optimization happens to help most where the sorting is slowest.The small
n_splitscase is unchanged. It sorts 20 node matrices, where the per-candidate work is small and the allocations this removes did not cost much in the first place. That is the configuration intest_rastermap_splits, so anyone benchmarking that test specifically will not see a difference.Verification
Output is bit-identical rather than approximately equal. The scripts are in this gist if you want to run the comparison yourself: record the sorting output on
main, record it again on this branch, and diff the two recordings. It reports a difference of any size in any case, so it fails loudly rather than quietly passing on a tolerance.I compared against the previous implementation in 52 cases with no mismatches:
tests/, comparingembedding,isort,cc,U_nodesandX_embedding.Rastermapconfigurations on synthetic data, coveringlocality,time_lag_window,n_splits,sticky,grid_upsample,mean_time,run_scaled_kmeansandtime_bin.traveling_salesmanandmatrix_matchingdirectly, including locality values up to infinity, the circular basis, a dense user-suppliedBBt, aBBtwith scattered interior zeros, an asymmetriccc, and bothn_skipsettings.The bit-exactness rests on numba's
.sum()using a sequential float32 accumulator in memory order rather than numpy's pairwise summation. I confirmed this empirically on numba 0.60, where the two disagree on 187 of 200 random trials, so the hand-written loop reproduces numba's existing behavior and not numpy's.What I Left Out
The
n_splitspath improves the least because itsBBt_addis dense and no zeros can be skipped, so only the allocations are saved. A further factorization is available there, sincecorr_permuted_rowspermutes rows but not columns and can therefore be precomputed as a single matrix product, reducing the per-candidate cost from O(n*m) to O(n). I left it out because it changes the floating point accumulation order, and with a greedy threshold of 1e-3 a small numerical difference can change which move is selected and therefore the final sort. That seemed worth discussing separately rather than folding into a change whose main claim is that nothing moves.Note on Tooling
I used Claude Code to profile this and to write the patch. The reason I am comfortable proposing it is that the result is bit-identical, so the claim does not rest on trusting the tool: the comparison reproduces the previous implementation's output exactly on your own test data, and the diff is a mechanical restructuring of one function. Happy to walk through it, and happy to drop it if you would rather not carry a change of this shape.