fix allele count

Author: daletovar

sgkit/stats/popgen.py

@@ -4,14 +4,14 @@
 import xarray as xr
 from xarray import DataArray, Dataset
 
-from .aggregation import count_alleles
+from .aggregation import count_variant_alleles
 
 
 def diversity(
     ds: Dataset, allele_counts: Hashable = "variant_allele_counts",
 ) -> DataArray:
     if allele_counts not in ds:
-        ds[allele_counts] = count_alleles(ds)
+        ds[allele_counts] = count_variant_alleles(ds)
     ac = ds[allele_counts]
     an = ac.sum(axis=1)
     n_pairs = an * (an - 1) / 2
@@ -34,10 +34,10 @@ def divergence(
     ds1: Dataset, ds2: Dataset, allele_counts: Hashable = "variant_allele_counts",
 ) -> DataArray:
     if allele_counts not in ds1:
-        ds1[allele_counts] = count_alleles(ds1)
+        ds1[allele_counts] = count_variant_alleles(ds1)
     ac1 = ds1[allele_counts]
     if allele_counts not in ds2:
-        ds2[allele_counts] = count_alleles(ds2)
+        ds2[allele_counts] = count_variant_alleles(ds2)
     ac2 = ds2[allele_counts]
 
     an1 = ds1[allele_counts].sum(axis=1)
@@ -65,7 +65,7 @@ def Tajimas_D(
     ds: Dataset, allele_counts: Hashable = "variant_allele_counts",
 ) -> DataArray:
     if allele_counts not in ds:
-        ds[allele_counts] = count_alleles(ds)
+        ds[allele_counts] = count_variant_alleles(ds)
     ac = ds[allele_counts]
 
     # count segregating