Skip to content

feat: add PanUKBB LD pair extraction - #1273

Draft
project-defiant wants to merge 14 commits into
devfrom
ld-pair-extraction-v2
Draft

feat: add PanUKBB LD pair extraction#1273
project-defiant wants to merge 14 commits into
devfrom
ld-pair-extraction-v2

Conversation

@project-defiant

@project-defiant project-defiant commented Jul 25, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds the Gentropy-side PanUKBB LD extraction contract for multi-ancestry fine-mapping.

This PR includes:

  • PanUKBB variant-index preparation for EUR, CSA, and AFR.
  • Population normalization, including the legacy nfeEUR alias.
  • Prepared and filtered LD variant-index helpers.
  • Signed long-format LD extraction from Hail BlockMatrices.
  • Typed MultiAncestryPairwiseLD with ancestry, variantIdI, variantIdJ, and r.
  • Validated JSONL study metadata input using a Pydantic model while preserving the Gentropy dataset schema.
  • FineMappingLocusSetLDAnnotationStep using an explicit ld_registry, final pair deduplication, session-controlled coalescing/write mode, and ancestry-level stats.jsonl.
  • Concrete ancestry-labelled BM paths while retaining legacy template compatibility.
  • Opt-in anonymous S3 session support for public Pan-UKBB BlockMatrices; authenticated S3 remains the default.
  • An opt-in network smoke test for public Pan-UKBB access, excluded from default pytest runs.
  • The portable pre-commit hook fix using python3.
  • Broadcast joins for the small LD-index lookups and persistence of reused locus/pair data.
  • Explicit Gentropy logger-level handling so requested process logging is visible.

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 PairwiseLD matrix.

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), and stats.jsonl reports counts for every requested ancestry, including zeroes.

Study metadata is supplied as one validated JSON object per JSONL line with studyId, ancestry, and sampleSize.

Validation

  • SDKMAN Temurin Java 11.0.23.
  • 36 targeted Spark-backed LD/PanUKBB tests passed.
  • Collector and Nextflow pipeline changes were validated separately; the pipeline collector suite has 80 tests passing.
  • Ruff, mypy, pydoclint, and all pre-commit hooks passed.
  • Default pytest configuration deselects the network smoke test.
  • Latest Gentropy commit: 0f467ec5.

Known limitations

  • The full remote Hail/S3 integration smoke test is intentionally opt-in and was not run in this validation pass.
  • BlockMatrix extraction currently materializes one locus at a time; further performance work is frozen pending hailing_ducks.
  • uv.lock was intentionally left unchanged.

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.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants