Skip to content

add Gram matrix fast path for pairwise_distance - #27

Open
kevinkorfmann wants to merge 1 commit into
kr-colab:mainfrom
kevinkorfmann:feat/fast-pairwise-distance
Open

add Gram matrix fast path for pairwise_distance#27
kevinkorfmann wants to merge 1 commit into
kr-colab:mainfrom
kevinkorfmann:feat/fast-pairwise-distance

Conversation

@kevinkorfmann

Copy link
Copy Markdown
Contributor

Summary

  • Add Gram matrix fast path for pairwise_distance() with euclidean/sqeuclidean metrics
  • Uses d²(i,j) = ||x_i||² + ||x_j||² - 2*x_i·x_j via a single X @ X.T (n x n) instead of materializing all
    (n_pairs, n_variants) difference vectors
  • Falls back to batched pairwise approach for cityblock metric or when missing data is present
  • ~13x speedup, exact correctness (0.0 error vs scipy reference)

Benchmarks (NVIDIA A100, 200 haplotypes x 50k variants)

n_haps n_snps Before (ms) After (ms) Speedup
200 50,000 ~75 5.8 ~13x
200 100,000 ~150 17.1 ~9x

Test plan

  • pixi run pytest tests/ -k 'decomposition or pca or pcoa or distance' (26 passed)
  • pixi run python debug/bench_pairwise_distance.py

@kevinkorfmann
kevinkorfmann force-pushed the feat/fast-pairwise-distance branch from 6b12cea to a540f61 Compare August 17, 2026 16:24
@kevinkorfmann

Copy link
Copy Markdown
Contributor Author

Rebased onto the latest main and simplified the implementation. Euclidean and squared-Euclidean distances now use the memory-safe chunked Gram path, while missing-data and cityblock cases retain the existing fallback.

@nspope

nspope commented Aug 17, 2026

Copy link
Copy Markdown
Collaborator

Sorry, this was rewritten to be multiallelic safe in #142 which is on the big multiallelic fix branch that's going to get merged in soon by #184. Can you wait until this is in, as it'll likely have merge conflicts?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants