Add LD pruning support for biallelic SNP calls

adilraza99 · adilraza99 · commit 08903d442464 · 2026-03-07T15:12:04.000+05:30
diff --git a/malariagen_data/anoph/ld.py b/malariagen_data/anoph/ld.py
@@ -0,0 +1,107 @@
+from typing import Optional
+
+import allel  # type: ignore
+import xarray as xr
+from numpydoc_decorator import doc  # type: ignore
+
+from ..util import _check_types, _dask_compress_dataset
+from . import base_params, ld_params, pca_params
+from .snp_data import AnophelesSnpData
+
+
+class AnophelesLdAnalysis(
+    AnophelesSnpData,
+):
+    def __init__(
+        self,
+        **kwargs,
+    ):
+        # N.B., this class is designed to work cooperatively, and
+        # so it's important that any remaining parameters are passed
+        # to the superclass constructor.
+        super().__init__(**kwargs)
+
+    @_check_types
+    @doc(
+        summary="""
+            Access biallelic SNP calls after LD pruning.
+        """,
+        extended_summary="""
+            This function obtains biallelic SNP calls, then performs LD pruning
+            using scikit-allel's `locate_unlinked` function. The resulting dataset
+            can be used as input to ADMIXTURE workflows or exported to PLINK format.
+
+            Note that `n_snps` is required to control memory usage. Without
+            pre-thinning, LD pruning could attempt to materialise millions of
+            variants and run out of memory.
+        """,
+        returns="""
+            A dataset of LD-pruned biallelic SNP calls with the same structure as
+            the output of `biallelic_snp_calls`.
+        """,
+    )
+    def biallelic_snp_calls_ld_pruned(
+        self,
+        region: base_params.regions,
+        n_snps: base_params.n_snps,
+        ld_window_size: ld_params.ld_window_size = ld_params.ld_window_size_default,
+        ld_window_step: ld_params.ld_window_step = ld_params.ld_window_step_default,
+        ld_threshold: ld_params.ld_threshold = ld_params.ld_threshold_default,
+        thin_offset: base_params.thin_offset = 0,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+        sample_query_options: Optional[base_params.sample_query_options] = None,
+        sample_indices: Optional[base_params.sample_indices] = None,
+        site_mask: Optional[base_params.site_mask] = base_params.DEFAULT,
+        min_minor_ac: Optional[
+            base_params.min_minor_ac
+        ] = pca_params.min_minor_ac_default,
+        max_missing_an: Optional[
+            base_params.max_missing_an
+        ] = pca_params.max_missing_an_default,
+        random_seed: base_params.random_seed = 42,
+        inline_array: base_params.inline_array = base_params.inline_array_default,
+        chunks: base_params.chunks = base_params.native_chunks,
+    ) -> xr.Dataset:
+        # Check that either sample_query xor sample_indices are provided.
+        base_params._validate_sample_selection_params(
+            sample_query=sample_query, sample_indices=sample_indices
+        )
+
+        # Obtain biallelic SNP calls with thinning applied first.
+        ds_snps = self.biallelic_snp_calls(
+            region=region,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_query_options=sample_query_options,
+            sample_indices=sample_indices,
+            site_mask=site_mask,
+            min_minor_ac=min_minor_ac,
+            max_missing_an=max_missing_an,
+            n_snps=n_snps,
+            thin_offset=thin_offset,
+            random_seed=random_seed,
+            inline_array=inline_array,
+            chunks=chunks,
+        )
+
+        # Compute genotype reference counts.
+        with self._dask_progress(desc="Computing genotype ref counts"):
+            gt = ds_snps["call_genotype"].data
+            gn = allel.GenotypeDaskArray(gt).to_n_ref(fill=-127).compute()
+
+        # Perform LD pruning.
+        with self._spinner(desc="LD pruning"):
+            loc_unlinked = allel.locate_unlinked(
+                gn,
+                size=ld_window_size,
+                step=ld_window_step,
+                threshold=ld_threshold,
+            )
+
+        # Apply the pruning mask.
+        ds_pruned = _dask_compress_dataset(
+            ds_snps, indexer=loc_unlinked, dim="variants"
+        )
+
+        return ds_pruned
diff --git a/malariagen_data/anoph/ld_params.py b/malariagen_data/anoph/ld_params.py
@@ -0,0 +1,25 @@
+"""Parameters for LD pruning functions."""
+
+from typing_extensions import Annotated, TypeAlias
+
+ld_window_size: TypeAlias = Annotated[
+    int,
+    "Window size in number of SNPs for LD pruning.",
+]
+
+ld_window_size_default: ld_window_size = 500
+
+ld_window_step: TypeAlias = Annotated[
+    int,
+    "Step size in number of SNPs for LD pruning.",
+]
+
+ld_window_step_default: ld_window_step = 200
+
+ld_threshold: TypeAlias = Annotated[
+    float,
+    "r-squared threshold for LD pruning. SNP pairs with r-squared above "
+    "this threshold will be considered linked.",
+]
+
+ld_threshold_default: ld_threshold = 0.1
diff --git a/malariagen_data/anopheles.py b/malariagen_data/anopheles.py
@@ -36,6 +36,7 @@
 from .anoph.sample_metadata import AnophelesSampleMetadata
 from .anoph.snp_data import AnophelesSnpData
 from .anoph.to_plink import PlinkConverter
+from .anoph.ld import AnophelesLdAnalysis
 from .anoph.g123 import AnophelesG123Analysis
 from .anoph.fst import AnophelesFstAnalysis
 from .anoph.h12 import AnophelesH12Analysis
@@ -88,6 +89,7 @@ class AnophelesDataResource(
     AnophelesDistanceAnalysis,
     AnophelesPca,
     PlinkConverter,
+    AnophelesLdAnalysis,
     AnophelesIgv,
     AnophelesKaryotypeAnalysis,
     AnophelesAimData,
diff --git a/tests/anoph/test_ld.py b/tests/anoph/test_ld.py
@@ -0,0 +1,172 @@
+import random
+
+import pytest
+from pytest_cases import parametrize_with_cases
+
+from malariagen_data import af1 as _af1
+from malariagen_data import ag3 as _ag3
+
+from malariagen_data.anoph.ld import AnophelesLdAnalysis
+
+
+@pytest.fixture
+def ag3_sim_api(ag3_sim_fixture):
+    return AnophelesLdAnalysis(
+        url=ag3_sim_fixture.url,
+        public_url=ag3_sim_fixture.url,
+        config_path=_ag3.CONFIG_PATH,
+        major_version_number=_ag3.MAJOR_VERSION_NUMBER,
+        major_version_path=_ag3.MAJOR_VERSION_PATH,
+        pre=True,
+        aim_metadata_dtype={
+            "aim_species_fraction_arab": "float64",
+            "aim_species_fraction_colu": "float64",
+            "aim_species_fraction_colu_no2l": "float64",
+            "aim_species_gambcolu_arabiensis": object,
+            "aim_species_gambiae_coluzzii": object,
+            "aim_species": object,
+        },
+        gff_gene_type="gene",
+        gff_gene_name_attribute="Name",
+        gff_default_attributes=("ID", "Parent", "Name", "description"),
+        default_site_mask="gamb_colu_arab",
+        results_cache=ag3_sim_fixture.results_cache_path.as_posix(),
+        taxon_colors=_ag3.TAXON_COLORS,
+        virtual_contigs=_ag3.VIRTUAL_CONTIGS,
+    )
+
+
+@pytest.fixture
+def af1_sim_api(af1_sim_fixture):
+    return AnophelesLdAnalysis(
+        url=af1_sim_fixture.url,
+        public_url=af1_sim_fixture.url,
+        config_path=_af1.CONFIG_PATH,
+        major_version_number=_af1.MAJOR_VERSION_NUMBER,
+        major_version_path=_af1.MAJOR_VERSION_PATH,
+        pre=False,
+        gff_gene_type="protein_coding_gene",
+        gff_gene_name_attribute="Note",
+        gff_default_attributes=("ID", "Parent", "Note", "description"),
+        default_site_mask="funestus",
+        results_cache=af1_sim_fixture.results_cache_path.as_posix(),
+        taxon_colors=_af1.TAXON_COLORS,
+    )
+
+
+def case_ag3_sim(ag3_sim_fixture, ag3_sim_api):
+    return ag3_sim_fixture, ag3_sim_api
+
+
+def case_af1_sim(af1_sim_fixture, af1_sim_api):
+    return af1_sim_fixture, af1_sim_api
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_ld_pruning_returns_fewer_snps(fixture, api: AnophelesLdAnalysis):
+    region = random.choice(api.contigs)
+    site_mask = random.choice(api.site_mask_ids)
+    ds_full = api.biallelic_snp_calls(
+        region=region,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+    n_available = ds_full.sizes["variants"]
+    if n_available < 10:
+        pytest.skip("Not enough variants for LD pruning test")
+
+    n_snps = min(n_available, 200)
+
+    ds_pruned = api.biallelic_snp_calls_ld_pruned(
+        region=region,
+        n_snps=n_snps,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+
+    # Pruned dataset should have fewer or equal variants.
+    assert ds_pruned.sizes["variants"] <= n_snps
+    assert ds_pruned.sizes["variants"] > 0
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_ld_pruned_dataset_structure(fixture, api: AnophelesLdAnalysis):
+    region = random.choice(api.contigs)
+    site_mask = random.choice(api.site_mask_ids)
+    ds_full = api.biallelic_snp_calls(
+        region=region,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+    n_available = ds_full.sizes["variants"]
+    if n_available < 10:
+        pytest.skip("Not enough variants for LD pruning test")
+
+    n_snps = min(n_available, 200)
+
+    ds_pruned = api.biallelic_snp_calls_ld_pruned(
+        region=region,
+        n_snps=n_snps,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+
+    # Check expected coordinates.
+    assert "sample_id" in ds_pruned.coords
+    assert "variant_position" in ds_pruned.coords
+    assert "variant_contig" in ds_pruned.coords
+
+    # Check expected data variables.
+    assert "variant_allele" in ds_pruned.data_vars
+    assert "call_genotype" in ds_pruned.data_vars
+
+    # Check dimensions.
+    assert "variants" in ds_pruned.dims
+    assert "samples" in ds_pruned.dims
+    assert "ploidy" in ds_pruned.dims
+    assert "alleles" in ds_pruned.dims
+
+    # Check alleles are biallelic.
+    assert ds_pruned.sizes["alleles"] == 2
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_ld_pruned_plink_compatibility(fixture, api: AnophelesLdAnalysis):
+    region = random.choice(api.contigs)
+    site_mask = random.choice(api.site_mask_ids)
+    ds_full = api.biallelic_snp_calls(
+        region=region,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+    n_available = ds_full.sizes["variants"]
+    if n_available < 10:
+        pytest.skip("Not enough variants for LD pruning test")
+
+    n_snps = min(n_available, 200)
+
+    ds_pruned = api.biallelic_snp_calls_ld_pruned(
+        region=region,
+        n_snps=n_snps,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+    )
+
+    # Verify the pruned dataset has all variables required by PlinkConverter.
+    assert "call_genotype" in ds_pruned
+    assert "variant_allele" in ds_pruned
+    assert "variant_contig" in ds_pruned.coords
+    assert "variant_position" in ds_pruned.coords
+    assert "sample_id" in ds_pruned.coords
+
+    # Verify shapes are internally consistent.
+    n_variants = ds_pruned.sizes["variants"]
+    n_samples = ds_pruned.sizes["samples"]
+    assert ds_pruned["call_genotype"].shape == (n_variants, n_samples, 2)
+    assert ds_pruned["variant_allele"].shape == (n_variants, 2)
