Variant summary stats

Author: tomwhite

requirements.txt

@@ -2,5 +2,6 @@ numpy
 xarray
 dask[array]
 scipy
+typing-extensions
 numba
 zarr

setup.cfg

@@ -59,7 +59,7 @@ ignore =
 profile = black
 default_section = THIRDPARTY
 known_first_party = sgkit
-known_third_party = dask,fire,glow,hail,hypothesis,invoke,numba,numpy,pandas,pkg_resources,pyspark,pytest,setuptools,sgkit_plink,xarray,yaml,zarr
+known_third_party = dask,fire,glow,hail,hypothesis,invoke,numba,numpy,pandas,pkg_resources,pyspark,pytest,setuptools,sgkit_plink,typing_extensions,xarray,yaml,zarr
 multi_line_output = 3
 include_trailing_comma = True
 force_grid_wrap = 0

sgkit/stats/aggregation.py

@@ -1,10 +1,13 @@
 import dask.array as da
 import numpy as np
 from numba import guvectorize
+from typing_extensions import Literal
 from xarray import Dataset
 
 from ..typing import ArrayLike
 
+Dimension = Literal["samples", "variants"]
+
 
 @guvectorize(  # type: ignore
     [
@@ -155,3 +158,59 @@ def count_variant_alleles(ds: Dataset, merge: bool = True) -> Dataset:
         }
     )
     return ds.merge(new_ds) if merge else new_ds
+
+
+def _swap(dim: Dimension) -> Dimension:
+    return "samples" if dim == "variants" else "variants"
+
+
+def call_rate(ds: Dataset, dim: Dimension) -> Dataset:
+    odim = _swap(dim)[:-1]
+    n_called = (~ds["call_genotype_mask"].any(dim="ploidy")).sum(dim=dim)
+    return xr.Dataset(
+        {f"{odim}_n_called": n_called, f"{odim}_call_rate": n_called / ds.dims[dim]}
+    )
+
+
+def genotype_count(ds: Dataset, dim: Dimension) -> Dataset:
+    odim = _swap(dim)[:-1]
+    M, G = ds["call_genotype_mask"].any(dim="ploidy"), ds["call_genotype"]
+    n_het = (G > 0).any(dim="ploidy") & (G == 0).any(dim="ploidy")
+    n_hom_ref = (G == 0).all(dim="ploidy")
+    n_hom_alt = (G > 0).all(dim="ploidy")
+    n_non_ref = (G > 0).any(dim="ploidy")
+    agg = lambda x: xr.where(M, False, x).sum(dim=dim)  # type: ignore[no-untyped-call]
+    return xr.Dataset(
+        {
+            f"{odim}_n_het": agg(n_het),  # type: ignore[no-untyped-call]
+            f"{odim}_n_hom_ref": agg(n_hom_ref),  # type: ignore[no-untyped-call]
+            f"{odim}_n_hom_alt": agg(n_hom_alt),  # type: ignore[no-untyped-call]
+            f"{odim}_n_non_ref": agg(n_non_ref),  # type: ignore[no-untyped-call]
+        }
+    )
+
+
+def allele_frequency(ds: Dataset) -> Dataset:
+    AC = count_alleles(ds)
+
+    M = ds["call_genotype_mask"].stack(calls=("samples", "ploidy"))
+    AN = (~M).sum(dim="calls")  # type: ignore
+    assert AN.shape == (ds.dims["variants"],)
+
+    return xr.Dataset(
+        {
+            "variant_allele_count": AC,
+            "variant_allele_total": AN,
+            "variant_allele_frequency": AC / AN,
+        }
+    )
+
+
+def variant_stats(ds: Dataset) -> Dataset:
+    return xr.merge(
+        [
+            call_rate(ds, dim="samples"),
+            genotype_count(ds, dim="samples"),
+            allele_frequency(ds),
+        ]
+    )

sgkit/tests/test_aggregation.py

@@ -4,7 +4,7 @@
 import xarray as xr
 from xarray import Dataset
 
-from sgkit.stats.aggregation import count_call_alleles, count_variant_alleles
+from sgkit.stats.aggregation import count_call_alleles, count_variant_alleles, variant_stats
 from sgkit.testing import simulate_genotype_call_dataset
 from sgkit.typing import ArrayLike
 
@@ -202,3 +202,26 @@ def test_count_call_alleles__chunked():
     ds["call_genotype"] = ds["call_genotype"].chunk(chunks=(5, 5, 1))  # type: ignore[arg-type]
     ac2 = count_call_alleles(ds)
     xr.testing.assert_equal(ac1, ac2)  # type: ignore[no-untyped-call]
+
+
+def test_variant_stats():
+    ds = get_dataset(
+        [[[1, 0], [-1, -1]], [[1, 0], [1, 1]], [[0, 1], [1, 0]], [[-1, -1], [0, 0]]]
+    )
+    vs = variant_stats(ds)
+
+    np.testing.assert_equal(vs["variant_n_called"], np.array([1, 2, 2, 1]))
+    np.testing.assert_equal(vs["variant_call_rate"], np.array([0.5, 1.0, 1.0, 0.5]))
+    np.testing.assert_equal(vs["variant_n_hom_ref"], np.array([0, 0, 0, 1]))
+    np.testing.assert_equal(vs["variant_n_hom_alt"], np.array([0, 1, 0, 0]))
+    np.testing.assert_equal(vs["variant_n_het"], np.array([1, 1, 2, 0]))
+    np.testing.assert_equal(vs["variant_n_non_ref"], np.array([1, 2, 2, 0]))
+    np.testing.assert_equal(vs["variant_n_non_ref"], np.array([1, 2, 2, 0]))
+    np.testing.assert_equal(
+        vs["variant_allele_count"], np.array([[1, 1], [1, 3], [2, 2], [2, 0]])
+    )
+    np.testing.assert_equal(vs["variant_allele_total"], np.array([2, 4, 4, 2]))
+    np.testing.assert_equal(
+        vs["variant_allele_frequency"],
+        np.array([[0.5, 0.5], [0.25, 0.75], [0.5, 0.5], [1, 0]]),
+    )