Add support for writing Anopheles SNP data to the plink binary file format

.gitignore

@@ -1,3 +1,4 @@
+debugging_biallelic_snp_calls.ipynb
 .idea
 .vscode
 __pycache__

malariagen_data/anoph/snp_data.py

@@ -1627,12 +1627,23 @@ def biallelic_snp_calls(
             # Store alleles, transformed.
             variant_allele = ds_bi["variant_allele"].data
             variant_allele = variant_allele.rechunk((variant_allele.chunks[0], -1))
+
+            # Chunk allele mapping according to same variant_allele.
+            allele_mapping_chunked = da.from_array(
+                allele_mapping, chunks=variant_allele.chunks
+            )
+
+            # Apply allele mapping blockwise to variant_allele.
             variant_allele_out = da.map_blocks(
-                lambda block: apply_allele_mapping(block, allele_mapping, max_allele=1),
+                lambda allele, map: apply_allele_mapping(allele, map, max_allele=1),
                 variant_allele,
+                allele_mapping_chunked,
                 dtype=variant_allele.dtype,
                 chunks=(variant_allele.chunks[0], [2]),
             )
+            # variant_allele_out = apply_allele_mapping(
+            #    variant_allele.compute(), allele_mapping, max_allele=1
+            # )
             data_vars["variant_allele"] = ("variants", "alleles"), variant_allele_out
 
             # Store allele counts, transformed, so we don't have to recompute.

malariagen_data/anoph/to_plink.py

@@ -0,0 +1,182 @@
+from typing import Optional
+
+import allel  # type: ignore
+import numpy as np
+import os
+import bed_reader
+from dask.diagnostics import ProgressBar
+
+from .snp_data import AnophelesSnpData
+from . import base_params
+
+
+class PlinkConverter(
+    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)
+
+    def _create_plink_outfile(
+        self,
+        *,
+        results_dir,
+        region,
+        n_snps,
+        min_minor_ac,
+        thin_offset,
+        max_missing_an,
+    ):
+        return f"{results_dir}/{region}.{n_snps}.{min_minor_ac}.{thin_offset}.{max_missing_an}"
+
+    def _biallelic_snps_to_plink(
+        self,
+        *,
+        results_dir,
+        region,
+        n_snps,
+        thin_offset,
+        sample_sets,
+        sample_query,
+        sample_indices,
+        site_mask,
+        min_minor_ac,
+        max_missing_an,
+        random_seed,
+        inline_array,
+        chunks,
+    ):
+        # Define output file
+        plink_file_path = self._create_plink_outfile(
+            results_dir=results_dir,
+            region=region,
+            n_snps=n_snps,
+            thin_offset=thin_offset,
+            min_minor_ac=min_minor_ac,
+            max_missing_an=max_missing_an,
+        )
+
+        bed_file_path = f"{plink_file_path}.bed"
+        if os.path.exists(bed_file_path):
+            return plink_file_path
+
+        # Get snps
+        ds_snps = self.biallelic_snp_calls(
+            region=region,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_indices=sample_indices,
+            site_mask=site_mask,
+            random_seed=random_seed,
+            inline_array=inline_array,
+            chunks=chunks,
+            n_snps=n_snps,
+        )
+
+        # Filter SNPs for segregating sites only
+        with self._spinner("Subsetting to segregating sites"):
+            gt = ds_snps["call_genotype"].data.compute()
+            print("count alleles")
+            with ProgressBar():
+                ac = allel.GenotypeArray(gt).count_alleles(max_allele=3)
+                print("ascertain segregating  sites")
+                n_chroms = ds_snps.dims["samples"] * 2
+                an_called = ac.sum(axis=1)
+                an_missing = n_chroms - an_called
+                min_ref_ac = min_minor_ac
+                max_ref_ac = n_chroms - min_minor_ac
+                loc_sites = (
+                    (ac[:, 0] >= min_ref_ac)
+                    & (ac[:, 0] <= max_ref_ac)
+                    & (an_missing <= max_missing_an)
+                )
+                print(f"ascertained {np.count_nonzero(loc_sites):,} sites")
+
+        # Set up dataset with required vars for plink conversion
+        print("Set up dataset")
+        ds_snps_asc = ds_snps[
+            [
+                "variant_contig",
+                "variant_position",
+                "variant_allele",
+                "sample_id",
+                "call_genotype",
+            ]
+        ].isel(alleles=slice(0, 2))
+
+        # Compute gt ref counts
+        with self._spinner("Computing genotype ref counts"):
+            gt_asc = ds_snps_asc["call_genotype"].data.compute()
+            gn_ref = allel.GenotypeDaskArray(gt_asc).to_n_ref(fill=-127)
+            with ProgressBar():
+                gn_ref = gn_ref.compute()
+
+        loc_var = np.any(gn_ref != gn_ref[:, 0, np.newaxis], axis=1)
+
+        with ProgressBar():
+            ds_snps_final = ds_snps_asc[
+                ["variant_contig", "variant_position", "variant_allele", "sample_id"]
+            ].isel(variants=loc_var)
+
+        # Init vars for input to bed reader
+        gn_ref_final = gn_ref[loc_var]
+        val = gn_ref_final.T
+        alleles = ds_snps_final["variant_allele"].values
+        properties = {
+            "iid": ds_snps_final["sample_id"].values,
+            "chromosome": ds_snps_final["variant_contig"].values,
+            "bp_position": ds_snps_final["variant_position"].values,
+            "allele_1": alleles[:, 0],
+            "allele_2": alleles[:, 1],
+        }
+
+        bed_reader.to_bed(
+            filepath=bed_file_path,
+            val=val,
+            properties=properties,
+            count_A1=True,
+        )
+
+        print(f"PLINK files written to to: {plink_file_path}")
+
+        return plink_file_path
+
+    def biallelic_snps_to_plink(
+        self,
+        results_dir,
+        region: base_params.regions,
+        n_snps: base_params.n_snps,
+        thin_offset: base_params.thin_offset = 0,
+        sample_sets: Optional[base_params.sample_sets] = None,
+        sample_query: Optional[base_params.sample_query] = None,
+        sample_indices: Optional[base_params.sample_indices] = None,
+        site_mask: Optional[base_params.site_mask] = None,
+        min_minor_ac: Optional[base_params.min_minor_ac] = 0,
+        max_missing_an: Optional[base_params.max_missing_an] = 0,
+        random_seed: base_params.random_seed = 42,
+        inline_array: base_params.inline_array = base_params.inline_array_default,
+        chunks: base_params.chunks = base_params.chunks_default,
+    ):
+        params = dict(
+            results_dir=results_dir,
+            region=region,
+            n_snps=n_snps,
+            thin_offset=thin_offset,
+            sample_sets=sample_sets,
+            sample_query=sample_query,
+            sample_indices=sample_indices,
+            site_mask=site_mask,
+            min_minor_ac=min_minor_ac,
+            max_missing_an=max_missing_an,
+            random_seed=random_seed,
+        )
+
+        filepath = self._biallelic_snps_to_plink(
+            inline_array=inline_array, chunks=chunks, **params
+        )
+        return filepath

malariagen_data/anopheles.py

@@ -49,6 +49,7 @@
 from .anoph.pca import AnophelesPca
 from .anoph.sample_metadata import AnophelesSampleMetadata, locate_cohorts
 from .anoph.snp_data import AnophelesSnpData
+from .anoph.to_plink import PlinkConverter
 from .anoph.g123 import AnophelesG123Analysis
 from .anoph.fst import AnophelesFstAnalysis
 from .anoph.h12 import AnophelesH12Analysis
@@ -105,6 +106,7 @@ class AnophelesDataResource(
     AnophelesG123Analysis,
     AnophelesFstAnalysis,
     AnophelesSnpFrequencyAnalysis,
+    PlinkConverter,
     AnophelesPca,
     AnophelesIgv,
     AnophelesAimData,