def test_window__default_step():
ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, seed=0)
assert not has_windows(ds)
ds = window(ds, 2)
assert has_windows(ds)
np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0])
np.testing.assert_equal(ds[window_start].values, [0, 2, 4, 6, 8])
np.testing.assert_equal(ds[window_stop].values, [2, 4, 6, 8, 10])
def test_window():
ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, seed=0)
assert not has_windows(ds)
ds = window(ds, 2, 2)
assert has_windows(ds)
np.testing.assert_equal(ds[window_contig].values, [0, 0, 0, 0, 0])
np.testing.assert_equal(ds[window_start].values, [0, 2, 4, 6, 8])
np.testing.assert_equal(ds[window_stop].values, [2, 4, 6, 8, 10])
with pytest.raises(MergeWarning):
window(ds, 2, 2)
def test_window_by_position__equal_spaced_windows():
ds = simulate_genotype_call_dataset(n_variant=5, n_sample=3, seed=0)
assert not has_windows(ds)
ds["variant_position"] = (
["variants"],
np.array([1, 4, 6, 8, 12]),
)
ds = window_by_position(ds, size=5, offset=1)
assert has_windows(ds)
np.testing.assert_equal(ds[window_contig].values, [0, 0, 0])
np.testing.assert_equal(ds[window_start].values, [0, 2, 4])
np.testing.assert_equal(ds[window_stop].values, [2, 4, 5])
def test_window_by_position__multiple_contigs():
ds = simulate_genotype_call_dataset(n_variant=10, n_sample=3, n_contig=2)
ds["variant_position"] = (
["variants"],
np.array([1, 4, 6, 8, 12, 1, 21, 25, 40, 55]),
)
ds = window_by_position(ds,
size=10,
window_start_position="variant_position")
assert has_windows(ds)
np.testing.assert_equal(ds[window_contig].values,
[0, 0, 0, 0, 0, 1, 1, 1, 1, 1])
np.testing.assert_equal(ds[window_start].values,
[0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
np.testing.assert_equal(ds[window_stop].values,
[4, 5, 5, 5, 5, 6, 8, 8, 9, 10])
def Garud_H(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
cohorts: Optional[Sequence[Union[int, str]]] = None,
merge: bool = True,
) -> Dataset:
"""Compute the H1, H12, H123 and H2/H1 statistics for detecting signatures
of soft sweeps, as defined in Garud et al. (2015).
By default, values of this statistic are calculated across all variants.
To compute values in windows, call :func:`window` before calling
this function.
Parameters
----------
ds
Genotype call dataset.
call_genotype
Input variable name holding call_genotype as defined by
:data:`sgkit.variables.call_genotype_spec`.
Must be present in ``ds``.
cohorts
The cohorts to compute statistics for, specified as a sequence of
cohort indexes or IDs. None (the default) means compute statistics
for all cohorts.
merge
If True (the default), merge the input dataset and the computed
output variables into a single dataset, otherwise return only
the computed output variables.
See :ref:`dataset_merge` for more details.
Returns
-------
A dataset containing the following variables:
- `stat_Garud_h1` (windows, cohorts): Garud H1 statistic.
Defined by :data:`sgkit.variables.stat_Garud_h1_spec`.
- `stat_Garud_h12` (windows, cohorts): Garud H12 statistic.
Defined by :data:`sgkit.variables.stat_Garud_h12_spec`.
- `stat_Garud_h123` (windows, cohorts): Garud H123 statistic.
Defined by :data:`sgkit.variables.stat_Garud_h123_spec`.
- `stat_Garud_h2_h1` (windows, cohorts): Garud H2/H1 statistic.
Defined by :data:`sgkit.variables.stat_Garud_h2_h1_spec`.
Raises
------
NotImplementedError
If the dataset is not diploid.
Warnings
--------
This function is currently only implemented for diploid datasets.
Examples
--------
>>> import numpy as np
>>> import sgkit as sg
>>> import xarray as xr
>>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
>>> # Divide samples into two cohorts
>>> sample_cohort = np.repeat([0, 1], ds.dims["samples"] // 2)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> # Divide into windows of size three (variants)
>>> ds = sg.window(ds, size=3, step=3)
>>> gh = sg.Garud_H(ds)
>>> gh["stat_Garud_h1"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.25 , 0.375],
[0.375, 0.375]])
>>> gh["stat_Garud_h12"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.375, 0.625],
[0.625, 0.625]])
>>> gh["stat_Garud_h123"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.625, 1. ],
[1. , 1. ]])
>>> gh["stat_Garud_h2_h1"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.75 , 0.33333333],
[0.33333333, 0.33333333]])
"""
if ds.dims["ploidy"] != 2:
raise NotImplementedError(
"Garud H only implemented for diploid genotypes")
if not has_windows(ds):
raise ValueError("Dataset must be windowed for Garud_H")
variables.validate(ds, {call_genotype: variables.call_genotype_spec})
gt = ds[call_genotype]
# convert sample cohorts to haplotype layout
sc = ds.sample_cohort.values
hsc = np.stack((sc, sc), axis=1).ravel() # TODO: assumes diploid
n_cohorts = sc.max() + 1 # 0-based indexing
cohorts = cohorts or range(n_cohorts)
ct = _cohorts_to_array(cohorts, ds.indexes.get("cohorts", None))
gh = window_statistic(
gt,
lambda gt: _Garud_h_cohorts(gt, hsc, n_cohorts, ct),
ds.window_start.values,
ds.window_stop.values,
dtype=np.float64,
# first chunks dimension is windows, computed in window_statistic
chunks=(-1, n_cohorts, N_GARUD_H_STATS),
)
n_windows = ds.window_start.shape[0]
assert_array_shape(gh, n_windows, n_cohorts, N_GARUD_H_STATS)
new_ds = create_dataset({
variables.stat_Garud_h1: (
("windows", "cohorts"),
gh[:, :, 0],
),
variables.stat_Garud_h12: (
("windows", "cohorts"),
gh[:, :, 1],
),
variables.stat_Garud_h123: (
("windows", "cohorts"),
gh[:, :, 2],
),
variables.stat_Garud_h2_h1: (
("windows", "cohorts"),
gh[:, :, 3],
),
})
return conditional_merge_datasets(ds, new_ds, merge)
def diversity(
ds: Dataset,
*,
cohort_allele_count: Hashable = variables.cohort_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute diversity from cohort allele counts.
By default, values of this statistic are calculated per variant.
To compute values in windows, call :func:`window` before calling
this function.
Parameters
----------
ds
Genotype call dataset.
cohort_allele_count
Cohort allele count variable to use or calculate. Defined by
:data:`sgkit.variables.cohort_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_cohort_alleles`.
merge
If True (the default), merge the input dataset and the computed
output variables into a single dataset, otherwise return only
the computed output variables.
See :ref:`dataset_merge` for more details.
Returns
-------
A dataset containing the diversity values, as defined by :data:`sgkit.variables.stat_diversity_spec`.
Shape (variants, cohorts), or (windows, cohorts) if windowing information is available.
Warnings
--------
This method does not currently support datasets that are chunked along the
samples dimension.
Examples
--------
>>> import numpy as np
>>> import sgkit as sg
>>> import xarray as xr
>>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
>>> # Divide samples into two cohorts
>>> sample_cohort = np.repeat([0, 1], ds.dims["samples"] // 2)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> sg.diversity(ds)["stat_diversity"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.5 , 0.66666667],
[0.66666667, 0.5 ],
[0.66666667, 0.66666667],
[0.5 , 0.5 ],
[0.5 , 0.5 ]])
>>> # Divide into windows of size three (variants)
>>> ds = sg.window(ds, size=3)
>>> sg.diversity(ds)["stat_diversity"].values # doctest: +NORMALIZE_WHITESPACE
array([[1.83333333, 1.83333333],
[1. , 1. ]])
"""
ds = define_variable_if_absent(ds, variables.cohort_allele_count,
cohort_allele_count, count_cohort_alleles)
variables.validate(
ds, {cohort_allele_count: variables.cohort_allele_count_spec})
ac = ds[cohort_allele_count]
an = ac.sum(axis=2)
n_pairs = an * (an - 1) / 2
n_same = (ac * (ac - 1) / 2).sum(axis=2)
n_diff = n_pairs - n_same
# replace zeros to avoid divide by zero error
n_pairs_na = n_pairs.where(n_pairs != 0)
pi = n_diff / n_pairs_na
if has_windows(ds):
div = window_statistic(
pi,
np.sum,
ds.window_start.values,
ds.window_stop.values,
dtype=pi.dtype,
axis=0,
)
new_ds = create_dataset(
{variables.stat_diversity: (
("windows", "cohorts"),
div,
)})
else:
new_ds = create_dataset(
{variables.stat_diversity: (
("variants", "cohorts"),
pi,
)})
return conditional_merge_datasets(ds, new_ds, merge)
def divergence(
ds: Dataset,
*,
cohort_allele_count: Hashable = variables.cohort_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute divergence between pairs of cohorts.
The entry at (i, j) is the divergence between for cohort i and cohort j,
except for the case where i and j are the same, in which case the entry
is the diversity for cohort i.
By default, values of this statistic are calculated per variant.
To compute values in windows, call :func:`window` before calling
this function.
Parameters
----------
ds
Genotype call dataset.
cohort_allele_count
Cohort allele count variable to use or calculate. Defined by
:data:`sgkit.variables.cohort_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_cohort_alleles`.
merge
If True (the default), merge the input dataset and the computed
output variables into a single dataset, otherwise return only
the computed output variables.
See :ref:`dataset_merge` for more details.
Returns
-------
A dataset containing the divergence value between pairs of cohorts, as defined by
:data:`sgkit.variables.stat_divergence_spec`.
Shape (variants, cohorts, cohorts), or (windows, cohorts, cohorts) if windowing
information is available.
Warnings
--------
This method does not currently support datasets that are chunked along the
samples dimension.
Examples
--------
>>> import numpy as np
>>> import sgkit as sg
>>> import xarray as xr
>>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
>>> # Divide samples into two cohorts
>>> sample_cohort = np.repeat([0, 1], ds.dims["samples"] // 2)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> sg.divergence(ds)["stat_divergence"].values # doctest: +NORMALIZE_WHITESPACE
array([[[0.5 , 0.5 ],
[0.5 , 0.66666667]],
<BLANKLINE>
[[0.66666667, 0.5 ],
[0.5 , 0.5 ]],
<BLANKLINE>
[[0.66666667, 0.5 ],
[0.5 , 0.66666667]],
<BLANKLINE>
[[0.5 , 0.375 ],
[0.375 , 0.5 ]],
<BLANKLINE>
[[0.5 , 0.625 ],
[0.625 , 0.5 ]]])
>>> # Divide into windows of size three (variants)
>>> ds = sg.window(ds, size=3)
>>> sg.divergence(ds)["stat_divergence"].values # doctest: +NORMALIZE_WHITESPACE
array([[[1.83333333, 1.5 ],
[1.5 , 1.83333333]],
<BLANKLINE>
[[1. , 1. ],
[1. , 1. ]]])
"""
ds = define_variable_if_absent(ds, variables.cohort_allele_count,
cohort_allele_count, count_cohort_alleles)
variables.validate(
ds, {cohort_allele_count: variables.cohort_allele_count_spec})
ac = ds[cohort_allele_count]
n_variants = ds.dims["variants"]
n_cohorts = ds.dims["cohorts"]
ac = da.asarray(ac)
shape = (ac.chunks[0], n_cohorts, n_cohorts)
d = da.map_blocks(_divergence, ac, chunks=shape, dtype=np.float64)
assert_array_shape(d, n_variants, n_cohorts, n_cohorts)
if has_windows(ds):
div = window_statistic(
d,
np.sum,
ds.window_start.values,
ds.window_stop.values,
dtype=d.dtype,
axis=0,
)
new_ds = create_dataset({
variables.stat_divergence: (
("windows", "cohorts_0", "cohorts_1"),
div,
)
})
else:
new_ds = create_dataset({
variables.stat_divergence: (
("variants", "cohorts_0", "cohorts_1"),
d,
)
})
return conditional_merge_datasets(ds, new_ds, merge)
def ld_matrix(
ds: Dataset,
*,
dosage: Hashable = variables.dosage,
threshold: Optional[float] = None,
variant_score: Optional[Hashable] = None,
) -> DataFrame:
"""Compute a sparse linkage disequilibrium (LD) matrix.
This method computes the Rogers Huff R2 value for each pair of variants in
a window, and returns those that exceed the provided threshold, as a sparse
matrix dataframe.
Parameters
----------
ds
Dataset containing genotype dosages. Must already be windowed with :func:`window`.
dosage
Name of genetic dosage variable.
Defined by :data:`sgkit.variables.dosage_spec`.
threshold
R2 threshold below which no variant pairs will be returned. This should almost
always be something at least slightly above 0 to avoid the large density very
near zero LD present in most datasets.
variant_score
Optional name of variable to use to prioritize variant selection
(e.g. minor allele frequency). Defaults to None.
Defined by :data:`sgkit.variables.variant_score_spec`.
Returns
-------
Upper triangle (including diagonal) of LD matrix as COO in dataframe. Fields:
- ``i``: Row (variant) index 1
- ``j``: Row (variant) index 2
- ``value``: R2 value
- ``cmp``: If ``variant_score`` is provided, this is 1, 0, or -1 indicating whether or not ``i > j`` (1), ``i < j`` (-1), or ``i == j`` (0)
Raises
------
ValueError
If the dataset is not windowed.
"""
if not has_windows(ds):
raise ValueError(
"Dataset must be windowed for ld_matrix. See the sgkit.window function."
)
variables.validate(ds, {dosage: variables.dosage_spec})
x = da.asarray(ds[dosage])
threshold = threshold or np.nan
if variant_score is not None:
variables.validate(ds, {variant_score: variables.variant_score_spec})
scores = da.asarray(ds[variant_score])
else:
scores = None
# Find windows in each chunk
window_starts = ds.window_start.values
window_stops = ds.window_stop.values
window_lengths = window_stops - window_starts
chunks = x.chunks[0]
chunk_starts = _sizes_to_start_offsets(chunks)
rel_window_starts, windows_per_chunk = _get_chunked_windows(
chunks, window_starts, window_stops)
rel_window_stops = rel_window_starts + window_lengths
chunk_offset_indexes = _sizes_to_start_offsets(windows_per_chunk)
def to_ld_df(x: ArrayLike, chunk_index: int) -> DataFrame:
chunk_offset_index_start = chunk_offset_indexes[chunk_index]
chunk_offset_index_stop = chunk_offset_indexes[chunk_index + 1]
chunk_window_starts = rel_window_starts[
chunk_offset_index_start:chunk_offset_index_stop]
chunk_window_stops = rel_window_stops[
chunk_offset_index_start:chunk_offset_index_stop]
max_stop = np.max(
chunk_window_stops) if len(chunk_window_stops) > 0 else 0
abs_chunk_start = chunk_starts[chunk_index]
abs_chunk_end = abs_chunk_start + max_stop # this may extend into later chunks
block_x = x[abs_chunk_start:abs_chunk_end]
block_scores = (scores[abs_chunk_start:abs_chunk_end]
if scores is not None else None)
# Look at the next window (not in this chunk) to find out where to stop processing
# windows in this chunk (see _ld_matrix_jit)
if len(window_starts) == chunk_offset_index_stop:
# if there are no more windows, then need to process the all windows in this chunk entirely
chunk_max_window_start = max_stop
else:
# otherwise only process up the start of the next window
chunk_max_window_start = (window_starts[chunk_offset_index_stop] -
chunk_starts[chunk_index])
index_dtype = np.int32
value_dtype = np.float32
f = dask.delayed(_ld_matrix)(
block_x,
chunk_window_starts,
chunk_window_stops,
abs_chunk_start,
chunk_max_window_start,
index_dtype,
value_dtype,
threshold=threshold,
scores=block_scores,
)
meta = [("i", index_dtype), ("j", index_dtype), ("value", value_dtype)]
if scores is not None:
meta.append(("cmp", np.int8))
return dd.from_delayed([f], meta=meta)
return dd.concat([
to_ld_df(x, chunk_index)
for chunk_index in range(len(windows_per_chunk))
])
def observed_heterozygosity(
ds: Dataset,
*,
call_heterozygosity: Hashable = variables.call_heterozygosity,
sample_cohort: Hashable = variables.sample_cohort,
merge: bool = True,
) -> Dataset:
"""Compute per cohort observed heterozygosity.
The observed heterozygosity of a cohort is the mean of individual
heterozygosity values among all samples of that cohort as described
in :func:`individual_heterozygosity`. Calls with a nan value for
individual heterozygosity are ignored when calculating the cohort
mean.
By default, values of this statistic are calculated per variant.
To compute values in windows, call :func:`window_by_position` or :func:`window_by_variant` before calling
this function.
Parameters
----------
ds
Dataset containing genotype calls.
call_heterozygosity
Input variable name holding call_heterozygosity as defined by
:data:`sgkit.variables.call_heterozygosity_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`individual_heterozygosity`.
sample_cohort
Input variable name holding sample_cohort as defined by
:data:`sgkit.variables.sample_cohort_spec`.
merge
If True (the default), merge the input dataset and the computed
output variables into a single dataset, otherwise return only
the computed output variables.
See :ref:`dataset_merge` for more details.
Returns
-------
A dataset containing :data:`sgkit.variables.stat_observed_heterozygosity_spec`
of per cohort observed heterozygosity with shape (variants, cohorts)
containing values within the inteval [0, 1] or nan.
Examples
--------
>>> import numpy as np
>>> import sgkit as sg
>>> import xarray as xr
>>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
>>> # Divide samples into two cohorts
>>> sample_cohort = np.repeat([0, 1], ds.dims["samples"] // 2)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> sg.observed_heterozygosity(ds)["stat_observed_heterozygosity"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.5, 1. ],
[1. , 0.5],
[0. , 1. ],
[0.5, 0.5],
[0.5, 0.5]])
>>> # Divide into windows of size three (variants)
>>> ds = sg.window_by_variant(ds, size=3)
>>> sg.observed_heterozygosity(ds)["stat_observed_heterozygosity"].values # doctest: +NORMALIZE_WHITESPACE
array([[1.5, 2.5],
[1. , 1. ]])
"""
ds = define_variable_if_absent(
ds,
variables.call_heterozygosity,
call_heterozygosity,
individual_heterozygosity,
)
variables.validate(
ds, {call_heterozygosity: variables.call_heterozygosity_spec})
hi = da.asarray(ds[call_heterozygosity])
sc = da.asarray(ds[sample_cohort])
n_cohorts = sc.max().compute() + 1
shape = (hi.chunks[0], n_cohorts)
n = da.zeros(n_cohorts, dtype=np.uint8)
ho = da.map_blocks(
_cohort_observed_heterozygosity,
hi,
sc,
n,
chunks=shape,
drop_axis=1,
new_axis=1,
dtype=np.float64,
)
if has_windows(ds):
ho_sum = window_statistic(
ho,
np.sum,
ds.window_start.values,
ds.window_stop.values,
dtype=ho.dtype,
axis=0,
)
new_ds = create_dataset({
variables.stat_observed_heterozygosity: (
("windows", "cohorts"),
ho_sum,
)
})
else:
new_ds = create_dataset({
variables.stat_observed_heterozygosity: (
("variants", "cohorts"),
ho,
)
})
return conditional_merge_datasets(ds, new_ds, merge)
def Tajimas_D(
ds: Dataset,
*,
variant_allele_count: Hashable = variables.variant_allele_count,
stat_diversity: Hashable = variables.stat_diversity,
merge: bool = True,
) -> Dataset:
"""Compute Tajimas' D for a genotype call dataset.
By default, values of this statistic are calculated per variant.
To compute values in windows, call :func:`window_by_position` or :func:`window_by_variant` before calling
this function.
Parameters
----------
ds
Genotype call dataset.
variant_allele_count
Variant allele count variable to use or calculate. Defined by
:data:`sgkit.variables.variant_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_variant_alleles`.
stat_diversity
Diversity variable to use or calculate. Defined by
:data:`sgkit.variables.stat_diversity_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`diversity`.
merge
If True (the default), merge the input dataset and the computed
output variables into a single dataset, otherwise return only
the computed output variables.
See :ref:`dataset_merge` for more details.
Returns
-------
A dataset containing the Tajimas' D value, as defined by :data:`sgkit.variables.stat_Tajimas_D_spec`.
Shape (variants, cohorts), or (windows, cohorts) if windowing information is available.
Warnings
--------
This method does not currently support datasets that are chunked along the
samples dimension.
Examples
--------
>>> import numpy as np
>>> import sgkit as sg
>>> import xarray as xr
>>> ds = sg.simulate_genotype_call_dataset(n_variant=5, n_sample=4)
>>> # Divide samples into two cohorts
>>> sample_cohort = np.repeat([0, 1], ds.dims["samples"] // 2)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> sg.Tajimas_D(ds)["stat_Tajimas_D"].values # doctest: +NORMALIZE_WHITESPACE
array([[0.88883234, 2.18459998],
[2.18459998, 0.88883234],
[2.18459998, 2.18459998],
[0.88883234, 0.88883234],
[0.88883234, 0.88883234]])
>>> # Divide into windows of size three (variants)
>>> ds = sg.window_by_variant(ds, size=3)
>>> sg.Tajimas_D(ds)["stat_Tajimas_D"].values # doctest: +NORMALIZE_WHITESPACE
array([[2.40517586, 2.40517586],
[1.10393559, 1.10393559]])
"""
ds = define_variable_if_absent(ds, variables.variant_allele_count,
variant_allele_count, count_variant_alleles)
ds = define_variable_if_absent(ds, variables.stat_diversity,
stat_diversity, diversity)
variables.validate(
ds,
{
variant_allele_count: variables.variant_allele_count_spec,
stat_diversity: variables.stat_diversity_spec,
},
)
ac = ds[variant_allele_count]
ac = da.asarray(ac)
# count segregating. Note that this uses the definition in tskit,
# which is the number of alleles - 1. In the biallelic case this
# gives us the number of non-monomorphic sites.
S = (ac > 0).sum(axis=1) - 1
if has_windows(ds):
S = window_statistic(
S,
np.sum,
ds.window_start.values,
ds.window_stop.values,
dtype=S.dtype,
axis=0,
)
# assume number of chromosomes sampled is constant for all variants
# NOTE: even tho ac has dtype uint, we promote the sum to float
# because the computation below requires floats
n = ac.sum(axis=1, dtype="float").max()
# (n-1)th harmonic number
a1 = (1 / da.arange(1, n)).sum()
# calculate Watterson's theta (absolute value)
theta = S / a1
# get diversity
div = ds[stat_diversity]
# N.B., both theta estimates are usually divided by the number of
# (accessible) bases but here we want the absolute difference
d = div - theta[:, np.newaxis]
# calculate the denominator (standard deviation)
a2 = (1 / (da.arange(1, n)**2)).sum()
b1 = (n + 1) / (3 * (n - 1))
b2 = 2 * (n**2 + n + 3) / (9 * n * (n - 1))
c1 = b1 - (1 / a1)
c2 = b2 - ((n + 2) / (a1 * n)) + (a2 / (a1**2))
e1 = c1 / a1
e2 = c2 / (a1**2 + a2)
d_stdev = da.sqrt((e1 * S) + (e2 * S * (S - 1)))
# Let IEEE decide the semantics of division by zero here. The return value
# will be -inf, nan or +inf, depending on the value of the numerator.
# Currently this will raise a RuntimeWarning, if we divide by zero.
D = d / d_stdev[:, np.newaxis]
if has_windows(ds):
new_ds = create_dataset(
{variables.stat_Tajimas_D: (["windows", "cohorts"], D.data)})
else:
new_ds = create_dataset(
{variables.stat_Tajimas_D: (["variants", "cohorts"], D.data)})
return conditional_merge_datasets(ds, new_ds, merge)
