perf: vectorize cohort_heterozygosity() for 10-50x speedup

[kunal-10-cloud]

Summary

This PR optimizes cohort_heterozygosity() by loading SNP data once per cohort instead of N times per sample, reducing disk I/O from O(N) → O(1) and enabling vectorized heterozygosity computation across all samples.

Estimated speedup: 10-50× for 100+ sample cohorts
Backward compatible: Yes — external API unchanged

Fixes: #1211

Changes Made

1. New Method: _cohort_count_het_vectorized()

File: malariagen_data/anoph/heterozygosity.py (lines 398-487)

Vectorized heterozygosity computation for multiple samples in a cohort:

• Loads SNP data once for all cohort samples via single snp_calls() call
• Uses GenotypeDaskArray.is_het() for vectorized computation → (variants, samples) shape
• Applies per-sample windowing in manual loop (moving_statistic is 1D-only)
• Returns Dict[sample_id → (windows, counts)] for easy per-sample access

Key optimization: Eliminates O(N) disk I/O operations, replacing with O(1)

2. Refactored: cohort_heterozygosity()

File: malariagen_data/anoph/heterozygosity.py (lines 898-927)

• Replaced sequential per-sample loop with single call to _cohort_count_het_vectorized()
• Maintains identical output format and numerical precision
• Added clear comments documenting vectorization benefit
• Backward compatible: No external API changes

3. Regression Test: test_cohort_count_het_vectorized_regression()

File: tests/anoph/test_heterozygosity.py (lines 265-328)

Validates vectorized method produces identical results as sequential approach:

• Compares outputs element-by-element across all samples
• Verifies numerical identity within floating-point tolerance (rtol=1e-10)
• Tests across all 4 data resource types (ag3_sim, af1_sim, adir1_sim, amin1_sim)
• Uses small cohort (3 samples) to keep test execution fast

Testing

Test Results

• All 28 heterozygosity tests pass (existing + new)
• 4× test_cohort_heterozygosity — PASS
• 4× test_cohort_count_het_vectorized_regression — PASS
• 20× other heterozygosity tests — PASS

Code Quality

• Pre-commit checks: PASS

• trim trailing whitespace

• fix end of files

• ruff linting

• ruff formatting

• Linting: All checks passed

Checklist

• Changes follow project style guidelines
• Pre-commit checks pass
• Ruff linting checks pass
• Ruff formatting checks pass
• All tests pass (28/28)
• Regression test verifies numerical correctness
• New method has comprehensive docstring
• Comments explain optimization benefit
• Backward compatible
• No breaking changes

Additional Notes

For Reviewers

1. The _cohort_count_het_vectorized() method is self-contained and could be reused in other multi-sample methods
2. The same vectorization pattern can be applied to other cohort-level computations
3. Regression test uses small cohort (3 samples) by design — large cohorts are covered by existing integration tests

Future Enhancements

• Optional parallelization of per-sample windowing using concurrent.futures
• Generalization to other multi-sample iterators in the codebase
• Caching layer for frequently accessed cohort parameters

[Copilot]

Pull request overview

This PR targets a major performance improvement for cohort_heterozygosity() by avoiding repeated SNP loading per sample and enabling cohort-wide vectorized heterozygosity computation.

Changes:

• Added a new private vectorized helper _cohort_count_het_vectorized() to compute windowed heterozygosity counts for multiple samples from a single snp_calls() load.
• Refactored cohort_heterozygosity() to use the new vectorized helper and then aggregate per-sample means.
• Added a regression test comparing vectorized vs sequential heterozygosity outputs.

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated 4 comments.

File	Description
malariagen_data/anoph/heterozygosity.py	Introduces vectorized cohort heterozygosity counting and updates cohort_heterozygosity() to use it.
tests/anoph/test_heterozygosity.py	Adds a regression test validating vectorized results match the sequential implementation.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[jonbrenas]

Thanks @kunal-10-cloud, I don't think adding vectorized to the name of the function is needed. It might also be a function that could be useful as a public function.

[kunal-10-cloud]

Hi @jonbrenas i have made the relevant changes and updated the branch as well please look into it once