feat: add PanUKBB LD pair extraction - #1273
Draft
project-defiant wants to merge 14 commits into
Draft
Conversation
The multi-ancestry fine-mapping pipeline needs local PanUKBB LD reference assets before it can query per-locus LD. This exposes the previously configured PanUKBB variant-index step, factors the allele-aligned index preparation into testable Spark helpers, and adds bounded BlockMatrix filtering for test subsets without copying the full matrix. Constraint: PanUKBB source indexes are Hail tables and matrices; tests must avoid remote reads and full matrix copies. Rejected: Implement the full CSA/AFR generalisation in this slice | issue 3 owns multi-population expansion after the EUR tracer is stable. Rejected: Materialise BlockMatrix to numpy for reference extraction | test assets should preserve Hail BlockMatrix shape and avoid unnecessary IO. Confidence: medium Scope-risk: moderate Directive: Keep alleleOrder in the prepared index; downstream LD sign correction depends on it. Tested: env JAVA_HOME=/usr/lib/jvm/java-8-openjdk-arm64 PATH=/usr/lib/jvm/java-8-openjdk-arm64/bin:/home/misio-pysio/.nvm/versions/node/v25.7.0/lib/node_modules/@openai/codex/node_modules/@openai/codex-linux-arm64/vendor/aarch64-unknown-linux-musl/codex-path:/home/misio-pysio/.codex/tmp/arg0/codex-arg0LU7AXk:/home/linuxbrew/.linuxbrew/bin:/home/linuxbrew/.linuxbrew/sbin:/home/misio-pysio/.bun/bin:/home/misio-pysio/.opencode/bin:/home/misio-pysio/.local/share/pnpm/bin:/home/linuxbrew/.linuxbrew/bin:/home/linuxbrew/.linuxbrew/sbin:/home/misio-pysio/.sdkman/candidates/sbt/current/bin:/home/misio-pysio/.sdkman/candidates/java/current/bin:/home/misio-pysio/.local/bin:/home/misio-pysio/.bun/bin:/home/misio-pysio/.nvm/versions/node/v25.7.0/bin:/home/misio-pysio/.opencode/bin:/home/misio-pysio/.cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/misio-pysio/.lmstudio/bin:/user/localgo/bin:/home/misio-pysio/.lmstudio/bin:/user/localgo/bin uv run --frozen --group test pytest tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py tests/gentropy/test_config.py tests/gentropy/test_pan_ukb_ingestion.py --no-cov -n 0 Tested: uv run --frozen --group dev ruff check src/gentropy/datasource/pan_ukbb_ld/ld.py src/gentropy/pan_ukb_ingestion.py src/gentropy/config.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py tests/gentropy/test_config.py tests/gentropy/test_pan_ukb_ingestion.py Tested: uv run --frozen --group dev ruff format --check src/gentropy/datasource/pan_ukbb_ld/ld.py src/gentropy/pan_ukb_ingestion.py src/gentropy/config.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py tests/gentropy/test_config.py tests/gentropy/test_pan_ukb_ingestion.py Not-tested: Real PanUKBB S3/Hail reference extraction; local Java 17 is not installed, so Spark tests were run with installed Java 8 instead of the target Java 17 runtime.
The LD reference preparation step must accept both pipeline ancestry labels and PanUKBB population labels before the pipeline can prepare CSA, EUR, and AFR assets from manifest-derived metadata. This centralizes PanUKBB population normalization, routes explicit PanUKBB labels through the PanUKBB interface, and verifies that every requested population writes full and filtered reference assets. Constraint: Pipeline metadata uses nfe/csa/afr while PanUKBB paths use EUR/CSA/AFR. Rejected: Keep a private mapping in LDMatrixInterface | that would leave the reference-preparation step and matrix query path with divergent label semantics. Rejected: Silently pass unsupported labels into PanUKBB path templates | invalid POP paths would fail later and be harder to diagnose. Confidence: high Scope-risk: moderate Directive: Use normalize_pan_ukbb_population for every PanUKBB path or ancestry argument; do not add another mapping table. Tested: env JAVA_HOME=/usr/lib/jvm/java-8-openjdk-arm64 PATH=/usr/lib/jvm/java-8-openjdk-arm64/bin:/home/misio-pysio/.nvm/versions/node/v25.7.0/lib/node_modules/@openai/codex/node_modules/@openai/codex-linux-arm64/vendor/aarch64-unknown-linux-musl/codex-path:/home/misio-pysio/.codex/tmp/arg0/codex-arg0LU7AXk:/home/linuxbrew/.linuxbrew/bin:/home/linuxbrew/.linuxbrew/sbin:/home/misio-pysio/.bun/bin:/home/misio-pysio/.opencode/bin:/home/misio-pysio/.local/share/pnpm/bin:/home/linuxbrew/.linuxbrew/bin:/home/linuxbrew/.linuxbrew/sbin:/home/misio-pysio/.sdkman/candidates/sbt/current/bin:/home/misio-pysio/.sdkman/candidates/java/current/bin:/home/misio-pysio/.local/bin:/home/misio-pysio/.bun/bin:/home/misio-pysio/.nvm/versions/node/v25.7.0/bin:/home/misio-pysio/.opencode/bin:/home/misio-pysio/.cargo/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin:/usr/games:/usr/local/games:/snap/bin:/home/misio-pysio/.lmstudio/bin:/user/localgo/bin:/home/misio-pysio/.lmstudio/bin:/user/localgo/bin UV_CACHE_DIR=/tmp/uv-cache TMPDIR=/tmp uv run --frozen --group test pytest tests/gentropy/test_pan_ukb_ingestion.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py tests/gentropy/method/test_ld_matrix_interface.py --no-cov -n 0 Tested: env UV_CACHE_DIR=/tmp/uv-cache TMPDIR=/tmp uv run --frozen --group dev ruff check src/gentropy/datasource/pan_ukbb_ld/ld.py src/gentropy/pan_ukb_ingestion.py src/gentropy/method/ld_matrix_interface.py tests/gentropy/test_pan_ukb_ingestion.py tests/gentropy/method/test_ld_matrix_interface.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py Tested: env UV_CACHE_DIR=/tmp/uv-cache TMPDIR=/tmp uv run --frozen --group dev ruff format --check src/gentropy/datasource/pan_ukbb_ld/ld.py src/gentropy/pan_ukb_ingestion.py src/gentropy/method/ld_matrix_interface.py tests/gentropy/test_pan_ukb_ingestion.py tests/gentropy/method/test_ld_matrix_interface.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py Not-tested: Real PanUKBB S3/Hail extraction for CSA/EUR/AFR; only mocked local path/filter behavior was tested.
MultiSuSiE consumes signed LD correlations in long form, keyed by ancestry and variant pair, not thresholded ldSet/r2 values. This adds a PanUKBB extraction helper that reuses the existing allele-order sign correction and converts a prepared filtered matrix slice into deterministic pair rows. A small writer wrapper persists the resulting rows as parquet for downstream annotation steps. Constraint: Issue 4 should not read full remote PanUKBB matrices in tests. Rejected: Extend PairwiseLD schema with ancestry | existing PairwiseLD is genome-agnostic and validates square row counts, while this annotation contract needs multi-ancestry long rows. Confidence: medium Scope-risk: narrow Directive: Keep r signed; do not square values or rename to r2. Tested: uv run --frozen --group test pytest tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py --no-cov -n 0 Tested: uv run --frozen --group dev ruff check src/gentropy/datasource/pan_ukbb_ld/ld.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py Tested: uv run --frozen --group dev ruff format --check src/gentropy/datasource/pan_ukbb_ld/ld.py tests/gentropy/datasource/ukbiobank/test_pan_ukbb_ld.py Tested: uv run --frozen --group dev mypy src/gentropy/datasource/pan_ukbb_ld/ld.py Not-tested: pre-commit run --all-files; local pre-commit Python 3.11 hook env fails before hooks with missing encodings. Not-tested: Real PanUKBB S3/Hail extraction and collector/Nextflow wiring; this branch implements the core extraction helper only.
setup-dev installs Python pre-commit correctly. The uv-managed python3.11 symlink makes virtualenv hook environments resolve to /install. Those hook environments then fail before importing encodings. Using python3 lets pre-commit resolve a working interpreter without hardcoding a path. Constraint: Gentropy still uses Python pre-commit rather than prek in this repository. Constraint: uv-lock required refreshing uv.lock to the current lockfile revision. Rejected: Hardcode the canonical uv Python path | not portable. Rejected: Switch Gentropy setup to prek | user clarified pre-commit should work. Confidence: high Scope-risk: narrow Directive: Do not restore python3.11 until uv symlink virtualenv creation is verified. Tested: make setup-dev Tested: uv run --no-sync pre-commit run --all-files Tested: targeted PanUKBB LD pytest passed
The hook-fix commit refreshed uv.lock only for the uv-lock hook. This PR should not include repository-wide lockfile metadata churn. Constraint: Keep the pre-commit fix and PanUKBB LD changes intact. Rejected: Keep refreshed uv.lock | thousands of unrelated PR lines. Confidence: high Scope-risk: narrow Tested: SKIP=uv-lock uv run --no-sync pre-commit run --all-files Not-tested: uv-lock hook, intentionally skipped to keep uv.lock out of the PR.
Reference-data bootstrap should prepare PanUKBB LD variant indexes only. It must not expose a path that writes or syncs multi-terabyte BlockMatrices. The ingestion step now writes filtered index datasets only. LD matrix access keeps its own default BlockMatrix path for query-time use. Allele alignment now joins oriented allele keys before unmatched positions. This avoids positional expansion at multi-allelic sites. Constraint: Local reference prep must not write PanUKBB BlockMatrices. Rejected: Optional filtered BlockMatrix writes | unsafe large writes are too easy. Confidence: high Scope-risk: moderate Directive: Do not add BlockMatrix writes back to PanUKBBVariantIndexStep. Tested: uv run ruff check changed files Tested: uv run mypy with repo-equivalent flags on changed implementation files Tested: targeted pytest for PanUKBB config, ingestion, and LD reference tests Tested: uv run pre-commit run --all-files passed before restoring uv.lock churn Not-tested: Full Gentropy test suite
Add a typed multi-ancestry container around independent PairwiseLD datasets. Constraint: Existing PairwiseLD schema and consumers remain unchanged Rejected: Add ancestry to PairwiseLD | one matrix cannot represent multiple ancestry dimensions Confidence: high Scope-risk: moderate Directive: Use MultiAncestryPairwiseLD for combined outputs Tested: pre-commit hooks; Ruff; diff check Not-tested: Spark tests because local Java is incompatible with Spark startup
Keep ancestry-aware LD pairs global. Reconstruct a square matrix only after selecting one locus. Add one typed metadata row per study. Constraint: PairwiseLD remains the locus-specific square matrix contract Rejected: Emit FineMappingLoci | reuse the original locus-set input Confidence: high Scope-risk: moderate Directive: Filter global pairs by locus before constructing PairwiseLD Tested: Ruff and pre-commit hooks Not-tested: Spark runtime tests; Java 24 blocks the local Spark gateway
Resolve ancestry-specific indexes and BlockMatrices for each locus. Emit one global typed LD-pair dataset. Constraint: Reuse FineMappingLocusSet without locus keys in global pairs Rejected: Emit a separate fine-mapping locus dataset | redundant materialization Confidence: medium Scope-risk: moderate Directive: Filter global pairs by locus before constructing PairwiseLD Tested: Ruff, focused helper tests, pre-commit hooks Not-tested: Hail integration; local Spark is incompatible with Java 24
Keep inherited LD validation and type the variant helper explicitly. Constraint: Preserve the global LD-pair and locus annotation contracts Rejected: Performance refactor | separate benchmarked scope Confidence: high Scope-risk: narrow Directive: Optimize BlockMatrix reads separately from cleanup Tested: Ruff, mypy, focused tests, pre-commit hooks Not-tested: Spark and Hail integration under Java 24
The fine-mapping LD annotation contract now consumes validated JSONL metadata, writes deterministic ancestry pair statistics, and supports concrete ancestry-labelled registries. Public Pan-UKBB matrices can be accessed through an explicit anonymous S3 session mode while authenticated S3 remains unchanged.
Constraint: Pan-UKBB public BlockMatrices require Hadoop anonymous credentials and legacy templated paths must retain population normalization.
Rejected: Make anonymous S3 the default | authenticated deployments must remain unchanged.
Confidence: high
Scope-risk: moderate
Reversibility: clean
Directive: Keep concrete registry ancestry labels exact; normalize only legacy {POP} matrix templates.
Tested: 33 Gentropy LD/Pan-UKBB tests, 4 anonymous S3 tests, Ruff, mypy, pydoclint, and default pytest deselection of network smoke test.
Not-tested: Opt-in public S3/Hail network smoke test; uv-lock hook requires unrelated lockfile regeneration.
The image ships Hadoop 3.3.4. The S3 connector requested 3.3.6 and caused runtime class-loading failures. Public Pan-UKBB access also needs Hadoop's anonymous provider. Constraint: Spark bundles Hadoop 3.3.4. Constraint: Pan-UKBB uses regional AWS S3 endpoints. Rejected: Keep hadoop-aws 3.3.6 | incompatible with the image runtime. Confidence: high Scope-risk: moderate Directive: Keep hadoop-aws aligned with Spark's bundled Hadoop version. Tested: Six session tests; image rebuild; S3A metadata read. Tested: 3x3 EUR Pan-UKBB matrix read. Not-tested: Full multi-run Nextflow integration.
Filter Hail entries before Spark conversion and preserve sparse pair ordering. Constraint: Large PanUKBB loci make dense P^2 materialization unsafe Rejected: Convert the full BlockMatrix before filtering | unnecessary transfer Confidence: high Scope-risk: moderate Directive: Preserve the sparse-pair contract downstream Tested: PanUKBB and PairwiseLD tests; pre-commit all-files Not-tested: Remote BlockMatrix smoke test
Avoid repeated work in sparse LD extraction. Materialize reused data and expose requested log levels. Constraint: Public Pan-UKBB BlockMatrices remain remote. Rejected: More concurrency; profiling points to remote I/O. Confidence: medium Scope-risk: moderate Directive: Revisit dense materialisation with hailing_ducks. Tested: 36 targeted tests; all pre-commit hooks. Not-tested: Full suite and production-scale run.
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
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
Adds the Gentropy-side PanUKBB LD extraction contract for multi-ancestry fine-mapping.
This PR includes:
nfe→EURalias.MultiAncestryPairwiseLDwithancestry,variantIdI,variantIdJ, andr.FineMappingLocusSetLDAnnotationStepusing an explicitld_registry, final pair deduplication, session-controlled coalescing/write mode, and ancestry-levelstats.jsonl.python3.Output contract
The annotation step does not emit a redundant fine-mapping-locus dataset. Downstream code filters the global LD pairs by ancestry and locus variants before constructing the existing square
PairwiseLDmatrix.Variants absent from an ancestry-specific LD index are omitted from that ancestry's extracted pairs. The final pairwise output is deduplicated by
(ancestry, variantIdI, variantIdJ), andstats.jsonlreports counts for every requested ancestry, including zeroes.Study metadata is supplied as one validated JSON object per JSONL line with
studyId,ancestry, andsampleSize.Validation
0f467ec5.Known limitations
hailing_ducks.uv.lockwas intentionally left unchanged.