Perf: 30% to 70% speedup for superimpose

[jberg5]

_get_rotation_matrices was doing np.sum(fixed[:, :, :, np.newaxis] * mobile[:, :, np.newaxis, :], axis=1) which materializes the full m*n*3*3 intermediate (so if m = 100 and n = 10000 you get 9,000,000 floats). This is expensive! _get_rotation_matrices can be made up to 20x faster if you just do this as a single matmul, since the original implementation was just expressing the same contraction as fixed.T @ mobile (ignoring the batch dimension). BLAS runs the computation directly without materializing the intermediate.

This shows up as a 1.3 to 1.7x end to end improvement on superimpose runtime at various dimensions (thanks Amdahl's law).

Tests pass locally. Expect this to not be byte-identical since float reduction order is changing, but the differences are at the level of floating point roundoff.

Claude did some benchmarking on superimpose for me:

x86 — GCP c3-standard-8 (Sapphire Rapids), OpenBLAS 8-thread, best-of-7:

(M, N)	Before (ms)	After (ms)	Speedup
(1, 200)	0.08	0.06	1.29x
(1, 1000)	0.16	0.11	1.49x
(1, 5000)	0.54	0.32	1.69x
(1, 25000)	3.18	2.13	1.49x
(5, 200)	0.17	0.12	1.42x
(5, 1000)	0.55	0.33	1.67x
(5, 5000)	3.17	2.16	1.47x
(5, 25000)	15.33	11.71	1.31x
(25, 200)	0.62	0.40	1.58x
(25, 1000)	2.64	1.52	1.73x
(25, 5000)	16.11	11.33	1.42x
(25, 25000)	75.97	52.06	1.46x
(100, 200)	2.39	1.48	1.62x
(100, 1000)	10.70	6.10	1.75x
(100, 5000)	63.26	41.58	1.52x
(100, 25000)	327.55	206.73	1.58x
(500, 200)	11.87	7.25	1.64x
(500, 1000)	58.11	34.42	1.69x
(500, 5000)	332.39	208.28	1.60x
(500, 25000)	1645.84	1025.11	1.61x

[codspeed-hq]

Merging this PR will regress 0 benchmarks

⚠️ Different runtime environments detected

Some benchmarks with significant performance changes were compared across different runtime environments,
which may affect the accuracy of the results.

Open the report in CodSpeed to investigate

⚡ 4 improved benchmarks
❌ 1 (👁 1) regressed benchmark
✅ 98 untouched benchmarks
🆕 3 new benchmarks
⏩ 14 skipped benchmarks¹

Performance Changes

	Benchmark	BASE	HEAD	Efficiency
⚡	benchmark_superimpose[superimpose_without_outliers]	16.8 ms	13.5 ms	+24.3%
⚡	benchmark_superimpose[superimpose]	2.4 ms	2 ms	+24.08%
⚡	benchmark_match[KmerTable-None]	329.2 µs	295.5 µs	+11.41%
⚡	benchmark_match_kmer_selection[BucketKmerTable(10000)-11*11*1*1***111]	352.7 µs	318.7 µs	+10.69%
🆕	benchmark_superimpose_multi_model[superimpose_homologs]	N/A	36.1 ms	N/A
🆕	benchmark_superimpose_multi_model[superimpose_without_outliers]	N/A	86.5 ms	N/A
🆕	benchmark_superimpose_multi_model[superimpose]	N/A	9.1 ms	N/A
👁	benchmark_cell_list_compute_contacts[MASK]	16.2 ms	18.6 ms	-12.84%

Tip

Curious why this is faster? Comment @codspeedbot explain why this is faster on this PR, or directly use the CodSpeed MCP with your agent.

---

_{Comparing jberg5:faster-superimpose (a90ec2b) with main (2587935)}

Footnotes

1. 14 benchmarks were skipped, so the baseline results were used instead. If they were deleted from the codebase, click here and archive them to remove them from the performance reports. ↩

[padix-key]

Looks good and the codspeed report confirms the improvement 👍. The failing CI test is not related to your change.