improve LD pruning with parameter validation and additional tests

Author: adilraza99

malariagen_data/anoph/ld.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+import numpy as np
+
 import allel  # type: ignore
 import xarray as xr
 from numpydoc_decorator import doc  # type: ignore
@@ -31,6 +33,15 @@ def __init__(
             using scikit-allel's `locate_unlinked` function. The resulting dataset
             can be used as input to ADMIXTURE workflows or exported to PLINK format.
 
+            LD pruning is controlled by three parameters:
+
+            - `ld_window_size`: number of SNPs in the sliding window used to
+              compute pairwise r-squared.
+            - `ld_window_step`: number of SNPs to advance the window each
+              iteration.
+            - `ld_threshold`: maximum r-squared value; SNP pairs above this
+              are considered linked and one will be removed.
+
             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.
@@ -68,6 +79,14 @@ def biallelic_snp_calls_ld_pruned(
             sample_query=sample_query, sample_indices=sample_indices
         )
 
+        # Validate LD parameters.
+        if ld_window_size <= 0:
+            raise ValueError(f"ld_window_size must be > 0, got {ld_window_size}")
+        if ld_window_step <= 0:
+            raise ValueError(f"ld_window_step must be > 0, got {ld_window_step}")
+        if not (0 < ld_threshold <= 1):
+            raise ValueError(f"ld_threshold must be in (0, 1], got {ld_threshold}")
+
         # Obtain biallelic SNP calls with thinning applied first.
         ds_snps = self.biallelic_snp_calls(
             region=region,
@@ -99,6 +118,13 @@ def biallelic_snp_calls_ld_pruned(
                 threshold=ld_threshold,
             )
 
+        # Guard against empty result.
+        if not np.any(loc_unlinked):
+            raise ValueError(
+                "LD pruning removed all variants. Consider using a less "
+                "stringent ld_threshold or providing more variants via n_snps."
+            )
+
         # Apply the pruning mask.
         ds_pruned = _dask_compress_dataset(
             ds_snps, indexer=loc_unlinked, dim="variants"

tests/anoph/test_ld.py

@@ -170,3 +170,44 @@ def test_ld_pruned_plink_compatibility(fixture, api: AnophelesLdAnalysis):
     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)
+
+
+@parametrize_with_cases("fixture,api", cases=".")
+def test_ld_pruning_threshold_sensitivity(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)
+
+    try:
+        ds_strict = api.biallelic_snp_calls_ld_pruned(
+            region=region,
+            n_snps=n_snps,
+            site_mask=site_mask,
+            min_minor_ac=1,
+            max_missing_an=0,
+            ld_threshold=0.1,
+        )
+    except ValueError:
+        pytest.skip("Too few variants survive strict LD pruning")
+
+    ds_lenient = api.biallelic_snp_calls_ld_pruned(
+        region=region,
+        n_snps=n_snps,
+        site_mask=site_mask,
+        min_minor_ac=1,
+        max_missing_an=0,
+        ld_threshold=0.5,
+    )
+
+    # A stricter threshold should retain fewer or equal variants.
+    assert ds_strict.sizes["variants"] <= ds_lenient.sizes["variants"]