Add dataset-return options for SNP counts/diplotypes and remove duplicate SNP calls

Author: 31puneet

malariagen_data/anoph/distance.py

@@ -195,7 +195,7 @@ def _biallelic_diplotype_pairwise_distances(
         max_missing_an,
     ):
         # Compute diplotypes.
-        gn, samples = self.biallelic_diplotypes(
+        ds = self.biallelic_diplotypes(
             region=region,
             sample_sets=sample_sets,
             sample_indices=sample_indices,
@@ -211,7 +211,10 @@ def _biallelic_diplotype_pairwise_distances(
             min_minor_ac=min_minor_ac,
             n_snps=n_snps,
             thin_offset=thin_offset,
+            return_dataset=True,
         )
+        gn = ds["call_diplotype"].values
+        samples = ds["sample_id"].values.astype("U")
 
         # Record number of SNPs used.
         n_snps = gn.shape[0]

malariagen_data/anoph/pca.py

@@ -207,7 +207,7 @@ def _pca(
         inline_array,
     ):
         # Load diplotypes.
-        gn, samples = self.biallelic_diplotypes(
+        ds_diplotypes = self.biallelic_diplotypes(
             region=region,
             n_snps=n_snps,
             thin_offset=thin_offset,
@@ -223,7 +223,10 @@ def _pca(
             random_seed=random_seed,
             chunks=chunks,
             inline_array=inline_array,
+            return_dataset=True,
         )
+        gn = ds_diplotypes["call_diplotype"].values
+        samples = ds_diplotypes["sample_id"].values.astype("U")
 
         with self._spinner(desc="Compute PCA"):
             # Exclude any samples prior to computing PCA.

malariagen_data/anoph/snp_data.py

@@ -1452,6 +1452,13 @@ def _snp_allele_counts(
             Compute SNP allele counts. This returns the number of times each
             SNP allele was observed in the selected samples.
         """,
+        parameters=dict(
+            return_dataset="""
+                If True, return an xarray dataset containing SNP calls with
+                `variant_allele_count` added as an extra data variable. If False
+                (default), return a numpy array of allele counts.
+            """
+        ),
         returns="""
             A numpy array of shape (n_variants, 4), where the first column has
             the reference allele (0) counts, the second column has the first
@@ -1481,7 +1488,26 @@ def snp_allele_counts(
         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,
-    ) -> np.ndarray:
+        return_dataset: bool = False,
+    ) -> Any:
+        if return_dataset:
+            ds = self._snp_calls_with_allele_counts(
+                region=region,
+                sample_sets=sample_sets,
+                sample_query=sample_query,
+                sample_query_options=sample_query_options,
+                sample_indices=sample_indices,
+                site_mask=site_mask,
+                site_class=site_class,
+                cohort_size=cohort_size,
+                min_cohort_size=min_cohort_size,
+                max_cohort_size=max_cohort_size,
+                random_seed=random_seed,
+                inline_array=inline_array,
+                chunks=chunks,
+            )
+            return ds
+
         # Change this name if you ever change the behaviour of this function,
         # to invalidate any previously cached data.
         name = "snp_allele_counts_v2"
@@ -1555,6 +1581,52 @@ def snp_allele_counts(
         ac = results["ac"]
         return ac
 
+    def _snp_calls_with_allele_counts(
+        self,
+        *,
+        region,
+        sample_sets,
+        sample_query,
+        sample_query_options,
+        sample_indices,
+        site_mask,
+        site_class,
+        cohort_size,
+        min_cohort_size,
+        max_cohort_size,
+        random_seed,
+        inline_array,
+        chunks,
+    ) -> xr.Dataset:
+        ds = self.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,
+            site_class=site_class,
+            cohort_size=cohort_size,
+            min_cohort_size=min_cohort_size,
+            max_cohort_size=max_cohort_size,
+            random_seed=random_seed,
+            inline_array=inline_array,
+            chunks=chunks,
+        )
+
+        gt = allel.GenotypeDaskArray(ds["call_genotype"].data)
+        ac = gt.count_alleles(max_allele=3)
+        with self._dask_progress(desc="Compute SNP allele counts"):
+            ac = ac.compute().values
+
+        ds = ds.assign(
+            variant_allele_count=(
+                ds["variant_allele"].dims,
+                ac,
+            )
+        )
+        return ds
+
     @_check_types
     @doc(
         summary="""
@@ -1897,8 +1969,8 @@ def biallelic_snp_calls(
             sample_query=sample_query, sample_indices=sample_indices
         )
 
-        # Perform an allele count.
-        ac = self.snp_allele_counts(
+        # Access SNP calls with allele counts in a single pass.
+        ds = self._snp_calls_with_allele_counts(
             region=region,
             sample_sets=sample_sets,
             sample_query=sample_query,
@@ -1913,6 +1985,7 @@ def biallelic_snp_calls(
             inline_array=inline_array,
             chunks=chunks,
         )
+        ac = ds["variant_allele_count"].values
 
         # Locate biallelic SNPs.
         loc_bi = allel.AlleleCountsArray(ac).is_biallelic()
@@ -1921,23 +1994,6 @@ def biallelic_snp_calls(
         ac_bi = ac[loc_bi]
         allele_mapping = _trim_alleles(ac_bi)
 
-        # Set up SNP calls.
-        ds = self.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,
-            site_class=site_class,
-            cohort_size=cohort_size,
-            min_cohort_size=min_cohort_size,
-            max_cohort_size=max_cohort_size,
-            random_seed=random_seed,
-            inline_array=inline_array,
-            chunks=chunks,
-        )
-
         with self._spinner("Prepare biallelic SNP calls"):
             # Subset to biallelic sites.
             ds_bi = _dask_compress_dataset(ds, indexer=loc_bi, dim="variants")
@@ -2032,13 +2088,22 @@ def biallelic_snp_calls(
     @_check_types
     @doc(
         summary="Load biallelic SNP genotypes.",
-        returns=dict(
-            gn="""
-                An array of shape (variants, samples) where each value counts the
-                number of alternate alleles per genotype call.
-            """,
-            samples="Sample identifiers.",
+        parameters=dict(
+            return_dataset="""
+                If True, return an xarray dataset with `call_diplotype` plus
+                `sample_id`, `variant_position`, and `variant_contig`. If False
+                (default), return a tuple `(gn, samples)` for backward
+                compatibility.
+            """
         ),
+        returns="""
+            If `return_dataset` is False (default), return `(gn, samples)`, where
+            `gn` is an array of shape `(variants, samples)` counting alternate
+            alleles per genotype call and `samples` contains sample identifiers.
+            If `return_dataset` is True, return a dataset containing
+            `call_diplotype` with dimensions `(variants, samples)`, plus
+            `sample_id`, `variant_position`, and `variant_contig`.
+        """,
     )
     def biallelic_diplotypes(
         self,
@@ -2059,7 +2124,8 @@ def biallelic_diplotypes(
         thin_offset: base_params.thin_offset = 0,
         inline_array: base_params.inline_array = base_params.inline_array_default,
         chunks: base_params.chunks = base_params.native_chunks,
-    ) -> Tuple[np.ndarray, np.ndarray]:
+        return_dataset: bool = False,
+    ) -> Any:
         # Change this name if you ever change the behaviour of this function, to
         # invalidate any previously cached data.
         name = "biallelic_diplotypes_v2"
@@ -2147,6 +2213,22 @@ def biallelic_diplotypes(
         gn = results["gn"]
         samples = results["samples"]
 
+        if return_dataset:
+            ds = xr.Dataset(
+                coords={
+                    "sample_id": ("samples", samples),
+                    "variant_position": ("variants", results["variant_position"]),
+                    "variant_contig": ("variants", results["variant_contig"]),
+                },
+                data_vars={
+                    "call_diplotype": (
+                        ("variants", "samples"),
+                        gn,
+                    )
+                },
+            )
+            return ds
+
         return gn, samples
 
     def _biallelic_diplotypes(
@@ -2193,6 +2275,8 @@ def _biallelic_diplotypes(
 
         # Load sample IDs
         samples = ds["sample_id"].values.astype("U")
+        variant_position = ds["variant_position"].values
+        variant_contig = ds["variant_contig"].values
 
         # Compute diplotypes as the number of alt alleles per genotype call.
         # with missing calls coded as -127.
@@ -2203,4 +2287,9 @@ def _biallelic_diplotypes(
         missing = np.all(ds["call_genotype"].values == -1, axis=2)
         gn[missing] = -127
 
-        return dict(samples=samples, gn=gn)
+        return dict(
+            samples=samples,
+            gn=gn,
+            variant_position=variant_position,
+            variant_contig=variant_contig,
+        )

tests/anoph/test_snp_data.py

@@ -1146,6 +1146,19 @@ def check_snp_allele_counts(
     )
     assert_array_equal(ac, ac2)
 
+    # Check dataset return mode.
+    ds_ac = api.snp_allele_counts(
+        region=region,
+        sample_sets=sample_sets,
+        sample_query=sample_query,
+        sample_query_options=sample_query_options,
+        site_mask=site_mask,
+        return_dataset=True,
+    )
+    assert isinstance(ds_ac, xr.Dataset)
+    assert "variant_allele_count" in ds_ac
+    assert_array_equal(ds_ac["variant_allele_count"].values, ac)
+
 
 @parametrize_with_cases(
     "fixture,api", cases=".", filter=~ft.has_tag("single-sampleset")
@@ -1453,6 +1466,22 @@ def check_biallelic_snp_calls_and_diplotypes(
     assert samples.shape[0] == gn.shape[1]
     assert samples.tolist() == expected_samples
 
+    # Check dataset return mode.
+    ds_gn = api.biallelic_diplotypes(
+        region=region,
+        sample_sets=sample_sets,
+        site_mask=site_mask,
+        site_class=site_class,
+        min_minor_ac=min_minor_ac,
+        max_missing_an=max_missing_an,
+        n_snps=n_snps,
+        return_dataset=True,
+    )
+    assert isinstance(ds_gn, xr.Dataset)
+    assert "call_diplotype" in ds_gn
+    assert_array_equal(ds_gn["call_diplotype"].values, gn)
+    assert ds_gn["sample_id"].values.tolist() == expected_samples
+
     return ds