[BenjaminDEMAILLE]

Part of #354 (item 24).

Not stacked. #345 and #346 both touch src/solo/count.rs, but neither changes the dedup functions (cell_feature_counts and distribute_multi are what they change), so this is based on main. git merge-tree reports no conflict between this branch and each of #345 and #346.

Problem

Four functions compared every pair of a (cell, gene) UMI set: connected_components, graph_correction_map, cellranger_1mm_map and directional_correction_map. Measured below, 200,000 UMIs in one (cell, gene) take 44 to 109 s per method, and because build_matrix_body calls this inside a parallel per-cell map, one deep gene holds a worker for that long. With CB/UB tags the correction map is computed a second time for the same set.

Change

• OneMmFinder finds the UMIs one mismatch from a given UMI. Above 224 UMIs it looks up the 3 * umi_len single-base variants in the set's own HashMap; up to 224 it compares pair by pair, since most sets are small and a comparison is cheaper than 36 hash misses. There is no extra index: the umi -> count map the caller already built is enough.
• cellranger_1mm_map no longer sorts. "The last 1MM neighbour with a later position in (count, umi) order" is the largest (count, umi) key above the UMI's own, whichever order the neighbours arrive in.
• directional_correction_map still walks the UMIs in (count descending, umi ascending) order, because a UMI takes its hub's corrected value and the hub has to be done first. "The first earlier hub passing the count test" becomes the smallest such key among the neighbours.
• The two union-finds are gone rather than merged. The audit item suggested sharing one. for_each_1mm_component walks each component with a work list and a seen-set, and both connected_components and graph_correction_map use it, which removes more code than a shared union-find would. The representative of a component (highest count, then smallest half-swapped value) depends on the component only, and the smaller raw UMI breaks what is left, so nothing depends on walk order. swap_halves alone is not enough: a UMI longer than --soloUMIlen can share its half-swapped value with another. With {3019, 4042, 976} and --soloUMIlen 5, the first version of this PR corrected 3019 to 4042 by walk order; umis_longer_than_umi_len_match_reference now checks that 4042 goes to 3019. The all_pairs reference applies the same rule. On main such a tie went to the first UMI in hash order, so UB can differ from main in that case only; molecule counts cannot.
• hamming1 is now constant-time bit arithmetic instead of a per-base loop. This is what keeps small sets from getting slower: the scan looks at all pairs, not half of them, and no longer stops early.
• cellranger_1mm and directional count distinct values in an FxHashSet, not a SipHash HashSet. Only .len() is read.

One case stays quadratic on purpose. --soloUMIposition does not have to agree with --soloUMIlen, so a packed UMI can be longer than umi_len. hamming1 only compares the last umi_len bases, so two such UMIs can be neighbours with different high bases, and no lookup of single-base variants would find them. OneMmFinder::new detects this and stays pairwise, to keep results as they were. I did not check whether comparing only part of the UMI is what STAR does in that configuration; it is the existing behaviour, and it is not changed here.

No STAR behaviour is claimed by this PR: the point is to be identical to the code it replaces. The comments that cite STAR's functions were already there.

Measured

cargo test --release --lib umi_dedup_timing -- --ignored --nocapture, on this branch. It runs the old code (kept in the test module) and the new code in one binary, on the same sets: random 12-base UMIs, 1 to 3 reads each. macOS arm64, 16 cores. The machine was busy, with a load average of 65 to 120, so absolute numbers are pessimistic. Up to 1,000 UMIs each number is the best of five rounds; 10,000 and 200,000 are a single round, so the spread between methods at 200,000 on the new code (62 to 208 ms) rests on one round each.

Large sets

UMIs	1MM_All before	after	1MM_Directional before	after	1MM_Directional_UMItools before	after	1MM_CR before	after
1,000	2.09 ms	0.087 ms	1.16 ms	0.078 ms	1.78 ms	0.079 ms	2.11 ms	0.070 ms
10,000	234 ms	2.84 ms	134 ms	1.19 ms	213 ms	1.20 ms	239 ms	1.04 ms
200,000	97.6 s	62 ms	44.2 s	208 ms	87.5 s	65 ms	108.8 s	88 ms

Small sets (ns per call)

UMIs	1MM_All before	after	1MM_Directional before	after	1MM_Directional_UMItools before	after	1MM_CR before	after
1	17	2	66	53	65	49	94	66
2	91	58	81	65	84	70	135	94
3	109	70	123	88	109	89	187	112
5	217	144	148	135	164	137	232	143
10	462	307	254	259	309	262	416	264
20	1,438	999	845	716	1,161	679	1,566	695
50	6,900	3,519	4,215	3,262	6,155	3,343	7,564	3,138
100	25,416	10,720	14,125	11,268	21,573	11,238	27,146	10,971
200	92,343	39,520	51,775	41,058	79,546	40,155	97,320	38,857

The one cell where "after" is above "before" is 1MM_Directional at 10 UMIs, by 2 %. So "not slower" holds everywhere measured but that one, and that one is small enough that I would not claim either way.

Choosing 224

Measured by building with the limit set to 0 and to 1,000,000 and running the same test.

UMIs	Pairwise only, 1MM_All	Lookup only, 1MM_All
10	323	797
50	3,553	10,605
100	11,041	24,544
128	18,065	11,219
200	40,123	53,480
500	236,445	45,906

Lookup is not monotonic. On random 12-base UMIs almost every lookup is a miss, and the rows line up with how full the map is: 100 and 200 keys fill 78 % of 128 and 256 buckets, and 128 and 500 keys about 50 % of 256 and 1,024. hashbrown grows at 7/8, so 224 keys are the most that 256 buckets hold and 225 move to 512 buckets. That is why the limit is 224 and not, say, 200. I did not run 224 or 225 themselves; 200 and 500 are the nearest rows.

This is not optimal everywhere: at 128 UMIs lookup alone would take 11,219 ns where this takes 18,022 ns, against 40,696 ns before.

End to end, main against this branch

300,000 reads on one barcode and one gene of a small synthetic genome, 150,000 random 12-base parent UMIs, 10 % of reads with one changed base: 148,070 distinct UMIs. --runThreadN 4, --outSAMtype BAM SortedByCoordinate --outSAMattributes NH HI CB UB. "Same records" is samtools view | sort compared in full, so it includes UB on all 300,000 records. The genome is not in the repository; the timing table above is what reproduces from the branch alone.

--soloUMIdedup	--soloUMIfiltering	matrix.mtx	BAM records	Molecules	main	branch
Exact	-	identical	identical	148,070	3.3 s	3.1 s
NoDedup	-	identical	identical	300,000	2.8 s	3.4 s
1MM_All	-	identical	identical	106,557	216.3 s	1.7 s
1MM_Directional	-	identical	identical	125,476	62.3 s	1.5 s
1MM_Directional_UMItools	-	identical	identical	115,493	119.3 s	2.1 s
1MM_CR	-	identical	identical	116,180	140.1 s	3.0 s
1MM_CR	MultiGeneUMI_CR	identical	identical	116,180	146.8 s	0.9 s
1MM_All	MultiGeneUMI	identical	identical	106,557	153.9 s	2.9 s

With a single gene, the two filtering rows run the code paths but no UMI is actually multi-gene. Multi-gene behaviour on the changed functions is covered by the existing multi_gene_umi_cr_* tests only.

test/solo_diff_docker.sh

Run from a copy of this branch: PASS on the 3' run (rustar-aligner matrix matches real STARsolo exactly) and on both 5' PE matrices (Gene and GeneFull, byte-identical). Two limits on what that shows: the container's STAR is Debian's 2.7.10b, not 2.7.11b, and the fixtures have 3 entries each, so they do not reach the lookup. It only says small sets still agree with a real STARsolo.

Tests

In src/solo/count.rs:

• umi_dedup_matches_all_pairs_reference: 3,000 random sets, against the old code kept in mod all_pairs, for all six methods, comparing dedup_count, umi_correction_map and the full cellranger_1mm_map. UMI lengths 1 to 12, so anywhere from nearly full to sparse; counts capped at 1, 2, 3, 10 or 1000 to force ties and chains. One set in twenty is over the limit. It asserts that more than 500 sets had corrections and more than 100 reached the lookup.
• one_mm_finder_paths_agree: the two paths of OneMmFinder return the same neighbours on the same set, including sets under the limit.
• hamming1_matches_the_base_by_base_reference: 100,000 pairs against the old loop, lengths 1 to 32, including bits past len.
• umis_longer_than_umi_len_match_reference: the longer-than-umi_len case above, on 896 draws of 8-base UMIs at umi_len 5.
• umi_dedup_timing: #[ignore]d, the table above.

I checked that these tests can fail by changing the code by hand and running them, then putting it back:

Change	Caught by
CR takes its lowest neighbour	umi_dedup_matches_all_pairs_reference, umis_longer_than_umi_len_match_reference
CR ignores the "later" test	those two, cellranger_1mm_collapses_neighbor, exact_dedup_corrects_nothing
Directional takes the last hub	umi_dedup_matches_all_pairs_reference, umis_longer_than_umi_len_match_reference
Directional drops the "earlier" test	umi_dedup_matches_all_pairs_reference
Directional stops following chains	umi_dedup_matches_all_pairs_reference
Graph tie-break flipped	umi_dedup_matches_all_pairs_reference, umis_longer_than_umi_len_match_reference
Lookup skips base 0	one_mm_finder_paths_agree, umi_dedup_matches_all_pairs_reference
Longer UMIs go to the lookup	umis_longer_than_umi_len_match_reference

The reference adds about 190 lines to the test module. The whole solo::count test module took 6 to 16 s in a debug build, depending on the load, and 0.5 s in release.

cargo test --release (817 passed, 1 ignored), cargo clippy --all-targets and cargo fmt --check are clean on macOS.

Not checked

• test/parity_suite.sh was not run.
• No real 10x data: every timing above is on random UMIs, where neighbours are sparse at 12 bases. Dense sets are in the randomized test, which checks results, not time, so time on a dense set of 200,000 is not measured.
• UMIs over 32 bases (--soloUMIlen 33 or more) are not compared with the old code. A packed u64 cannot hold them, and the old hamming1 wrapped its shifts in that case, so there is no meaningful result to keep. Such sets stay on the pairwise path.
• The CB/UB path computes the correction twice, as before. I left that alone.

🤖 Generated with Claude Code