def count_genotypes(
ds: Dataset,
dim: Dimension,
call_genotype: Hashable = variables.call_genotype,
call_genotype_mask: Hashable = variables.call_genotype_mask,
merge: bool = True,
) -> Dataset:
variables.validate(
ds,
{
call_genotype_mask: variables.call_genotype_mask_spec,
call_genotype: variables.call_genotype_spec,
},
)
odim = _swap(dim)[:-1]
M, G = ds[call_genotype_mask].any(dim="ploidy"), ds[call_genotype]
n_hom_ref = (G == 0).all(dim="ploidy")
n_hom_alt = ((G > 0) & (G[..., 0] == G)).all(dim="ploidy")
n_non_ref = (G > 0).any(dim="ploidy")
n_het = ~(n_hom_alt | n_hom_ref)
# This would 0 out the `het` case with any missing calls
agg = lambda x: xr.where(M, False, x).sum(dim=dim) # type: ignore[no-untyped-call]
new_ds = create_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]
}
)
return conditional_merge_datasets(ds, new_ds, merge)
def convert_probability_to_call(
ds: Dataset,
call_genotype_probability: str = variables.call_genotype_probability,
threshold: float = 0.9,
merge: bool = True,
) -> Dataset:
"""
Convert genotype probabilities to hard calls.
Parameters
----------
ds
Dataset containing genotype probabilities, such as from :func:`sgkit.io.bgen.read_bgen`.
call_genotype_probability
Genotype probability variable to be converted as defined by
:data:`sgkit.variables.call_genotype_probability_spec`.
threshold
Probability threshold in [0, 1] that must be met or exceeded by at least one genotype
probability in order for any calls to be made -- all values will be -1 (missing)
otherwise. Setting this value to less than or equal to 0 disables any effect it has.
Default value is 0.9.
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:
- `call_genotype` (variants, samples, ploidy): Converted hard calls.
Defined by :data:`sgkit.variables.call_genotype_spec`.
- `call_genotype_mask` (variants, samples, ploidy): Mask for converted hard calls.
Defined by :data:`sgkit.variables.call_genotype_mask_spec`.
"""
if not (0 <= threshold <= 1):
raise ValueError(
f"Threshold must be float in [0, 1], not {threshold}.")
variables.validate(
ds,
{call_genotype_probability: variables.call_genotype_probability_spec})
if ds.dims["genotypes"] != 3:
raise NotImplementedError(
f"Hard call conversion only supported for diploid, biallelic genotypes; "
f"num genotypes in provided probabilities array = {ds.dims['genotypes']}."
)
GP = da.asarray(ds[call_genotype_probability])
# Remove chunking in genotypes dimension, if present
if len(GP.chunks[2]) > 1:
GP = GP.rechunk((None, None, -1))
K = da.empty(2, dtype=np.uint8)
GT = _convert_probability_to_call(GP, K, threshold)
new_ds = create_dataset({
variables.call_genotype: (("variants", "samples", "ploidy"), GT),
variables.call_genotype_mask:
(("variants", "samples", "ploidy"), GT < 0),
})
return conditional_merge_datasets(ds, new_ds, merge)
def count_variant_alleles(
ds: Dataset,
*,
call_allele_count: Hashable = variables.call_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute allele count from per-sample allele counts, or genotype calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_allele_count
Input variable name holding call_allele_count as defined by
:data:`sgkit.variables.call_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_call_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 :data:`sgkit.variables.variant_allele_count_spec`
of allele counts with shape (variants, alleles) and values corresponding to
the number of non-missing occurrences of each allele.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=4, n_sample=2, seed=1)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 1/0 1/0
1 1/0 1/1
2 0/1 1/0
3 0/0 0/0
>>> sg.count_variant_alleles(ds)["variant_allele_count"].values # doctest: +SKIP
array([[2, 2],
[1, 3],
[2, 2],
[4, 0]], dtype=uint64)
"""
ds = define_variable_if_absent(ds, variables.call_allele_count,
call_allele_count, count_call_alleles)
variables.validate(ds,
{call_allele_count: variables.call_allele_count_spec})
new_ds = create_dataset({
variables.variant_allele_count:
ds[call_allele_count].sum(dim="samples")
})
return conditional_merge_datasets(ds, new_ds, merge)
def test_variables__multiple_specs(dummy_ds: xr.Dataset) -> None:
spec = ArrayLikeSpec("baz", "baz doc", kind="i", ndim=1)
invalid_spec = ArrayLikeSpec("baz", "baz doc", kind="i", ndim=2)
variables.validate(dummy_ds, {"foo": spec, "bar": spec})
variables.validate(dummy_ds, {"foo": spec})
variables.validate(dummy_ds, {"bar": spec})
with pytest.raises(ValueError, match="bar does not match the spec"):
variables.validate(dummy_ds, {"bar": invalid_spec})
with pytest.raises(ValueError, match="bar does not match the spec"):
variables.validate(dummy_ds, {"foo": spec}, {"bar": invalid_spec})
def test_variables__validate_by_name(dummy_ds: xr.Dataset) -> None:
spec = ArrayLikeSpec("foo", "foo doc", kind="i", ndim=1)
try:
assert "foo" not in SgkitVariables.registered_variables
name, spec_b = SgkitVariables.register_variable(spec)
assert "foo" in SgkitVariables.registered_variables
assert name == "foo"
assert spec_b == spec
variables.validate(dummy_ds, "foo")
finally:
SgkitVariables.registered_variables.pop("foo", None)
assert "foo" not in SgkitVariables.registered_variables
def test_variables__whole_ds(dummy_ds: xr.Dataset) -> None:
spec_foo = ArrayLikeSpec("foo", "foo doc", kind="i", ndim=1)
spec_bar = ArrayLikeSpec("bar", "bar doc", kind="i", ndim=1)
try:
SgkitVariables.register_variable(spec_foo)
with pytest.raises(ValueError, match="`foo` already registered"):
SgkitVariables.register_variable(spec_foo)
SgkitVariables.register_variable(spec_bar)
variables.validate(dummy_ds)
finally:
SgkitVariables.registered_variables.pop("foo", None)
SgkitVariables.registered_variables.pop("bar", None)
def infer_call_genotype_fill(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
merge: bool = True,
) -> Dataset:
variables.validate(ds, {call_genotype: variables.call_genotype_spec})
mixed_ploidy = ds[variables.call_genotype].attrs.get("mixed_ploidy", False)
if mixed_ploidy:
call_genotype_fill = ds[call_genotype] < -1
else:
call_genotype_fill = xr.full_like(ds[call_genotype], False, "b1")
new_ds = create_dataset({variables.call_genotype_fill: call_genotype_fill})
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def pca_stats(ds: Dataset,
est: BaseEstimator,
*,
merge: bool = True) -> Dataset:
""" Extract attributes from PCA estimator """
new_ds = {
variables.sample_pca_component: (
("variants", "components"),
_get(est, "components_", fn=lambda v: v.T),
),
variables.sample_pca_explained_variance: (
"components",
_get(est, "explained_variance_"),
),
variables.sample_pca_explained_variance_ratio: (
"components",
_get(est, "explained_variance_ratio_"),
),
}
new_ds = Dataset({k: v for k, v in new_ds.items() if v[1] is not None})
if "sample_pca_component" in new_ds and "sample_pca_explained_variance" in new_ds:
new_ds[variables.sample_pca_loading] = new_ds[
variables.sample_pca_component] * np.sqrt(
new_ds[variables.sample_pca_explained_variance].data)
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def allele_frequency(
ds: Dataset,
call_genotype_mask: Hashable,
variant_allele_count: Hashable,
) -> Dataset:
data_vars: Dict[Hashable, Any] = {}
# only compute variant allele count if not already in dataset
if variant_allele_count in ds:
variables.validate(
ds, {variant_allele_count: variables.variant_allele_count_spec}
)
AC = ds[variant_allele_count]
else:
AC = count_variant_alleles(ds, merge=False)[variables.variant_allele_count]
data_vars[variables.variant_allele_count] = AC
M = ds[call_genotype_mask].stack(calls=("samples", "ploidy"))
AN = (~M).sum(dim="calls")
assert AN.shape == (ds.dims["variants"],)
data_vars[variables.variant_allele_total] = AN
data_vars[variables.variant_allele_frequency] = AC / AN
return create_dataset(data_vars)
def infer_call_ploidy(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
call_genotype_non_allele: Hashable = variables.call_genotype_non_allele,
merge: bool = True,
) -> Dataset:
"""Infer the ploidy of each call genotype based on the number of
non-allele values in each call genotype.
Parameters
----------
ds
Dataset containing genotype calls.
call_genotype
Input variable name holding call_genotype as defined by
:data:`sgkit.variables.call_genotype_spec`.
Must be present in ``ds``.
call_genotype_non_allele
Input variable name holding call_genotype_non_allele as defined by
:data:`sgkit.variables.call_genotype_non_allele_spec`.
If the variable is not present in ``ds``, it will be computed
assuming that allele values less than -1 are non-alleles in mixed ploidy
datasets, or that no non-alleles are present in fixed ploidy datasets.
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.call_ploidy_spec`.
"""
ds = define_variable_if_absent(
ds,
variables.call_genotype_non_allele,
call_genotype_non_allele,
infer_non_alleles,
)
mixed_ploidy = ds[variables.call_genotype].attrs.get("mixed_ploidy", False)
if mixed_ploidy:
call_ploidy = (~ds[call_genotype_non_allele]).sum(
axis=-1) # type: ignore[operator]
else:
ploidy = ds[variables.call_genotype].shape[-1]
call_ploidy = xr.full_like(ds[variables.call_genotype][..., 0], ploidy)
new_ds = create_dataset({variables.call_ploidy: call_ploidy})
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def pca_transform(
ds: Dataset,
est: BaseEstimator,
*,
variable: str = "call_alternate_allele_count",
check_missing: bool = True,
merge: bool = True,
) -> Dataset:
""" Apply PCA estimator to new data """
AC = _allele_counts(ds, variable, check_missing=check_missing)
projection = est.transform(da.asarray(AC).T)
new_ds = Dataset(
{variables.sample_pca_projection: (("samples", "components"), projection)}
)
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def infer_sample_ploidy(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
call_ploidy: Hashable = variables.call_ploidy,
merge: bool = True,
) -> Dataset:
"""Infer the ploidy of each sample across all variants based on
the number of non-allele values in call genotypes.
Parameters
----------
ds
Dataset containing genotype calls.
call_genotype
Input variable name holding call_genotype as defined by
:data:`sgkit.variables.call_genotype_spec`.
Must be present in ``ds``.
call_ploidy
Input variable name holding call_ploidy as defined by
:data:`sgkit.variables.call_ploidy_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`infer_call_ploidy`.
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.sample_ploidy_spec`.
"""
ds = define_variable_if_absent(ds, variables.call_ploidy, call_ploidy,
infer_call_ploidy)
# validate against spec
mixed_ploidy = ds[variables.call_genotype].attrs.get("mixed_ploidy", False)
if mixed_ploidy:
sample_ploidy_fixed = (ds[call_ploidy][0, :] == ds[call_ploidy]).all(
axis=-1)
sample_ploidy = xr.where(sample_ploidy_fixed, ds[call_ploidy][0, :],
-1) # type: ignore[no-untyped-call]
else:
ploidy = ds[variables.call_genotype].shape[-1]
sample_ploidy = xr.full_like(ds[call_ploidy][0, ...], ploidy)
new_ds = create_dataset({variables.sample_ploidy: sample_ploidy})
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def individual_heterozygosity(
ds: Dataset,
*,
call_allele_count: Hashable = variables.call_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute per call individual heterozygosity.
Individual heterozygosity is the probability that two alleles
drawn at random without replacement, from an individual at a
given site, are not identical in state. Therefore, individual
heterozygosity is defined for diploid and polyploid calls but
will return nan in the case of haploid calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_allele_count
Input variable name holding call_allele_count as defined by
:data:`sgkit.variables.call_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_call_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 :data:`sgkit.variables.call_heterozygosity_spec`
of per genotype observed heterozygosity with shape (variants, samples)
containing values within the interval [0, 1] or nan if ploidy < 2.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=4, n_sample=2, seed=1)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 1/0 1/0
1 1/0 1/1
2 0/1 1/0
3 0/0 0/0
>>> sg.individual_heterozygosity(ds)["call_heterozygosity"].values # doctest: +NORMALIZE_WHITESPACE
array([[1., 1.],
[1., 0.],
[1., 1.],
[0., 0.]])
"""
ds = define_variable_if_absent(
ds, variables.call_allele_count, call_allele_count, count_call_alleles
)
variables.validate(ds, {call_allele_count: variables.call_allele_count_spec})
AC = da.asarray(ds.call_allele_count)
K = AC.sum(axis=-1)
# use nan denominator to avoid divide by zero with K - 1
K2 = da.where(K > 1, K, np.nan)
AF = AC / K2[..., None]
HI = (1 - da.sum(AF ** 2, axis=-1)) * (K / (K2 - 1))
new_ds = create_dataset(
{variables.call_heterozygosity: (("variants", "samples"), HI)}
)
return conditional_merge_datasets(ds, new_ds, merge)
def sample_stats(
ds: Dataset,
*,
call_genotype_mask: Hashable = variables.call_genotype_mask,
call_genotype: Hashable = variables.call_genotype,
variant_allele_count: Hashable = variables.variant_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute quality control sample statistics from genotype calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_genotype
Input variable name holding call_genotype.
Defined by :data:`sgkit.variables.call_genotype_spec`.
Must be present in ``ds``.
call_genotype_mask
Input variable name holding call_genotype_mask.
Defined by :data:`sgkit.variables.call_genotype_mask_spec`
Must be present in ``ds``.
variant_allele_count
Input variable name holding variant_allele_count,
as 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`.
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:
- :data:`sgkit.variables.sample_n_called_spec` (samples):
The number of variants with called genotypes.
- :data:`sgkit.variables.sample_call_rate_spec` (samples):
The fraction of variants with called genotypes.
- :data:`sgkit.variables.sample_n_het_spec` (samples):
The number of variants with heterozygous calls.
- :data:`sgkit.variables.sample_n_hom_ref_spec` (samples):
The number of variants with homozygous reference calls.
- :data:`sgkit.variables.sample_n_hom_alt_spec` (samples):
The number of variants with homozygous alternate calls.
- :data:`sgkit.variables.sample_n_non_ref_spec` (samples):
The number of variants that are not homozygous reference calls.
"""
variables.validate(
ds,
{
call_genotype: variables.call_genotype_spec,
call_genotype_mask: variables.call_genotype_mask_spec,
},
)
new_ds = xr.merge(
[
call_rate(ds, dim="variants", call_genotype_mask=call_genotype_mask),
count_genotypes(
ds,
dim="variants",
call_genotype=call_genotype,
call_genotype_mask=call_genotype_mask,
merge=False,
),
]
)
return conditional_merge_datasets(ds, variables.validate(new_ds), merge)
def call_allele_frequencies(
ds: Dataset,
*,
call_allele_count: Hashable = variables.call_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute per sample allele frequencies from genotype calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_allele_count
Input variable name holding call_allele_count as defined by
:data:`sgkit.variables.call_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_call_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 :data:`sgkit.variables.call_allele_frequency_spec`
of allele frequencies with shape (variants, samples, alleles) and values
corresponding to the frequency of non-missing occurrences of each allele.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=4, n_sample=2, seed=1)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 1/0 1/0
1 1/0 1/1
2 0/1 1/0
3 0/0 0/0
>>> sg.call_allele_frequencies(ds)["call_allele_frequency"].values # doctest: +NORMALIZE_WHITESPACE
array([[[0.5, 0.5],
[0.5, 0.5]],
<BLANKLINE>
[[0.5, 0.5],
[0. , 1. ]],
<BLANKLINE>
[[0.5, 0.5],
[0.5, 0.5]],
<BLANKLINE>
[[1. , 0. ],
[1. , 0. ]]])
"""
ds = define_variable_if_absent(
ds, variables.call_allele_count, call_allele_count, count_call_alleles
)
variables.validate(ds, {call_allele_count: variables.call_allele_count_spec})
AC = ds[call_allele_count]
K = AC.sum(dim="alleles")
# avoid divide by zero
AF = AC / xr.where(K > 0, K, np.nan) # type: ignore[no-untyped-call]
new_ds = create_dataset({variables.call_allele_frequency: AF})
return conditional_merge_datasets(ds, new_ds, merge)
def count_cohort_alleles(
ds: Dataset,
*,
call_allele_count: Hashable = variables.call_allele_count,
sample_cohort: Hashable = variables.sample_cohort,
merge: bool = True,
) -> Dataset:
"""Compute per cohort allele counts from per-sample allele counts, or genotype calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_allele_count
Input variable name holding call_allele_count as defined by
:data:`sgkit.variables.call_allele_count_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`count_call_alleles`.
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.cohort_allele_count_spec`
of allele counts with shape (variants, cohorts, alleles) and values corresponding to
the number of non-missing occurrences of each allele.
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
>>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.dims["samples"] // 2), dims="samples")
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1 S2 S3
variants
0 0/0 1/0 1/0 0/1
1 1/0 0/1 0/0 1/0
2 1/1 0/0 1/0 0/1
3 1/0 1/1 1/1 1/0
4 1/0 0/0 1/0 1/1
>>> sg.count_cohort_alleles(ds)["cohort_allele_count"].values # doctest: +NORMALIZE_WHITESPACE
array([[[3, 1],
[2, 2]],
<BLANKLINE>
[[2, 2],
[3, 1]],
<BLANKLINE>
[[2, 2],
[2, 2]],
<BLANKLINE>
[[1, 3],
[1, 3]],
<BLANKLINE>
[[3, 1],
[1, 3]]])
"""
ds = define_variable_if_absent(
ds, variables.call_allele_count, call_allele_count, count_call_alleles
)
variables.validate(ds, {call_allele_count: variables.call_allele_count_spec})
n_variants = ds.dims["variants"]
n_alleles = ds.dims["alleles"]
AC, SC = da.asarray(ds[call_allele_count]), da.asarray(ds[sample_cohort])
n_cohorts = SC.max().compute() + 1 # 0-based indexing
C = da.empty(n_cohorts, dtype=np.uint8)
G = da.asarray(ds.call_genotype)
shape = (G.chunks[0], n_cohorts, n_alleles)
AC = da.map_blocks(_count_cohort_alleles, AC, SC, C, chunks=shape, dtype=np.int32)
assert_array_shape(
AC, n_variants, n_cohorts * AC.numblocks[1], n_alleles * AC.numblocks[2]
)
# Stack the blocks and sum across them
# (which will only work because each chunk is guaranteed to have same size)
AC = da.stack([AC.blocks[:, i] for i in range(AC.numblocks[1])]).sum(axis=0)
assert_array_shape(AC, n_variants, n_cohorts, n_alleles)
new_ds = create_dataset(
{variables.cohort_allele_count: (("variants", "cohorts", "alleles"), AC)}
)
return conditional_merge_datasets(ds, new_ds, merge)
def test_variables__no_spec(dummy_ds: xr.Dataset) -> None:
with pytest.raises(ValueError, match="No array spec registered for foo"):
variables.validate(dummy_ds, "foo")
variables.validate(dummy_ds,
"bar") # no spec needed for coordinates or indexes
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 identity_by_state(
ds: Dataset,
*,
call_allele_frequency: Hashable = variables.call_allele_frequency,
merge: bool = True,
) -> Dataset:
"""Compute identity by state (IBS) probabilities between
all pairs of samples.
The IBS probability between a pair of individuals is the
probability that a randomly drawn allele from the first individual
is identical in state with a randomly drawn allele from the second
individual at a single random locus.
Parameters
----------
ds
Dataset containing call genotype alleles.
call_allele_frequency
Input variable name holding call_allele_frequency as defined by
:data:`sgkit.variables.call_allele_frequency_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`call_allele_frequencies`.
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_identity_by_state_spec`
which is a matrix of pairwise IBS probabilities among all samples.
The dimensions are named ``samples_0`` and ``samples_1``.
Raises
------
NotImplementedError
If the variable holding call_allele_frequency is chunked along the
samples dimension.
Warnings
--------
This method does not currently support datasets that are chunked along the
samples dimension.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=2, n_sample=3, seed=2)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1 S2
variants
0 0/0 1/1 1/0
1 1/1 1/1 1/0
>>> sg.identity_by_state(ds)["stat_identity_by_state"].values # doctest: +NORMALIZE_WHITESPACE
array([[1. , 0.5, 0.5],
[0.5, 1. , 0.5],
[0.5, 0.5, 0.5]])
"""
ds = define_variable_if_absent(
ds,
variables.call_allele_frequency,
call_allele_frequency,
call_allele_frequencies,
)
variables.validate(
ds, {call_allele_frequency: variables.call_allele_frequency_spec}
)
af = da.asarray(ds[call_allele_frequency])
if len(af.chunks[1]) > 1:
raise NotImplementedError(
"identity_by_state does not support chunking in the samples dimension"
)
af0 = da.where(da.isnan(af), 0.0, af)
num = da.einsum("ixj,iyj->xy", af0, af0)
called = da.nansum(af, axis=-1)
count = da.einsum("ix,iy->xy", called, called)
denom = da.where(count == 0, np.nan, count)
new_ds = create_dataset(
{
variables.stat_identity_by_state: (
("samples_0", "samples_1"),
num / denom,
)
}
)
return conditional_merge_datasets(ds, new_ds, merge)
def count_call_alleles(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
merge: bool = True,
) -> Dataset:
"""Compute per sample allele counts from genotype calls.
Parameters
----------
ds
Dataset containing genotype calls.
call_genotype
Input variable name holding call_genotype as defined by
:data:`sgkit.variables.call_genotype_spec`.
Must be present in ``ds``.
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.call_allele_count_spec`
of allele counts with shape (variants, samples, alleles) and values corresponding to
the number of non-missing occurrences of each allele.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=4, n_sample=2, seed=1)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 1/0 1/0
1 1/0 1/1
2 0/1 1/0
3 0/0 0/0
>>> sg.count_call_alleles(ds)["call_allele_count"].values # doctest: +NORMALIZE_WHITESPACE
array([[[1, 1],
[1, 1]],
<BLANKLINE>
[[1, 1],
[0, 2]],
<BLANKLINE>
[[1, 1],
[1, 1]],
<BLANKLINE>
[[2, 0],
[2, 0]]], dtype=uint8)
"""
variables.validate(ds, {call_genotype: variables.call_genotype_spec})
n_alleles = ds.dims["alleles"]
G = da.asarray(ds[call_genotype])
shape = (G.chunks[0], G.chunks[1], n_alleles)
N = da.empty(n_alleles, dtype=np.uint8)
new_ds = create_dataset(
{
variables.call_allele_count: (
("variants", "samples", "alleles"),
da.map_blocks(
count_alleles, G, N, chunks=shape, drop_axis=2, new_axis=2
),
)
}
)
return conditional_merge_datasets(ds, new_ds, merge)
def test_variables__validate_by_dummy_spec(dummy_ds: xr.Dataset) -> None:
spec = ArrayLikeSpec("foo", "foo doc", kind="i", ndim=1)
variables.validate(dummy_ds, spec)
def test_variables__no_present_in_ds(dummy_ds: xr.Dataset) -> None:
spec = ArrayLikeSpec("baz", "baz doc", kind="i", ndim=1)
with pytest.raises(ValueError, match="foobarbaz not present in"):
variables.validate(dummy_ds, {"foobarbaz": spec})
def test_variables__alternative_names(dummy_ds: xr.Dataset) -> None:
spec = ArrayLikeSpec("baz", "baz doc", kind="i", ndim=1)
variables.validate(dummy_ds, {"foo": spec, "bar": spec})
def test_variables__invalid_spec_fails(dummy_ds: xr.Dataset) -> None:
invalid_spec = ArrayLikeSpec("foo", "foo doc", kind="i", ndim=2)
with pytest.raises(ValueError, match="foo does not match the spec"):
variables.validate(dummy_ds, invalid_spec)
def hardy_weinberg_test(ds: Dataset,
*,
genotype_counts: Optional[Hashable] = None,
ploidy: Optional[int] = None,
alleles: Optional[int] = None,
merge: bool = True) -> Dataset:
"""Exact test for HWE as described in Wigginton et al. 2005 [1].
Parameters
----------
ds
Dataset containing genotype calls or precomputed genotype counts.
genotype_counts
Name of variable containing precomputed genotype counts, by default
None. If not provided, these counts will be computed automatically
from genotype calls. If present, must correspond to an (`N`, 3) array
where `N` is equal to the number of variants and the 3 columns contain
heterozygous, homozygous reference, and homozygous alternate counts
(in that order) across all samples for a variant.
ploidy
Genotype ploidy, defaults to ``ploidy`` dimension of provided dataset.
If the `ploidy` dimension is not present, then this value must be set explicitly.
Currently HWE calculations are only supported for diploid datasets,
i.e. ``ploidy`` must equal 2.
alleles
Genotype allele count, defaults to ``alleles`` dimension of provided dataset.
If the `alleles` dimension is not present, then this value must be set explicitly.
Currently HWE calculations are only supported for biallelic datasets,
i.e. ``alleles`` must equal 2.
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.
Warnings
--------
This function is only applicable to diploid, biallelic datasets.
Returns
-------
Dataset containing (N = num variants):
variant_hwe_p_value : [array-like, shape: (N, O)]
P values from HWE test for each variant as float in [0, 1].
References
----------
- [1] Wigginton, Janis E., David J. Cutler, and Goncalo R. Abecasis. 2005.
“A Note on Exact Tests of Hardy-Weinberg Equilibrium.” American Journal of
Human Genetics 76 (5): 887–93.
Raises
------
NotImplementedError
If ploidy of provided dataset != 2
NotImplementedError
If maximum number of alleles in provided dataset != 2
"""
ploidy = ploidy or ds.dims.get("ploidy")
if not ploidy:
raise ValueError(
"`ploidy` parameter must be set when not present as dataset dimension."
)
if ploidy != 2:
raise NotImplementedError(
"HWE test only implemented for diploid genotypes")
alleles = alleles or ds.dims.get("alleles")
if not alleles:
raise ValueError(
"`alleles` parameter must be set when not present as dataset dimension."
)
if alleles != 2:
raise NotImplementedError(
"HWE test only implemented for biallelic genotypes")
# Use precomputed genotype counts if provided
if genotype_counts is not None:
variables.validate(ds,
{genotype_counts: variables.genotype_counts_spec})
obs = list(da.asarray(ds[genotype_counts]).T)
# Otherwise compute genotype counts from calls
else:
ds = count_genotypes(ds, dim="samples")
obs = [
da.asarray(ds[v]) for v in
["variant_n_het", "variant_n_hom_ref", "variant_n_hom_alt"]
]
p = da.map_blocks(hardy_weinberg_p_value_vec_jit, *obs)
new_ds = create_dataset({variables.variant_hwe_p_value: ("variants", p)})
return conditional_merge_datasets(ds, new_ds, merge)
def test_variables_in_multi_index(dummy_ds: xr.Dataset) -> None:
# create a multi index
ds = dummy_ds.set_index({"ind": ("foo", "bar")})
spec = ArrayLikeSpec("foo", "foo doc", kind="i", ndim=1)
variables.validate(ds, spec)
def Weir_Goudet_beta(
ds: Dataset,
*,
stat_identity_by_state: Hashable = variables.stat_identity_by_state,
merge: bool = True,
) -> Dataset:
"""Estimate pairwise beta between all pairs of samples as described
in Weir and Goudet 2017 [1].
Beta is the kinship scaled by the average kinship of all pairs of
individuals in the dataset such that the non-diagonal (non-self) values
sum to zero.
Beta may be corrected to more accurately reflect pedigree based kinship
estimates using the formula
:math:`\\hat{\\beta}^c=\\frac{\\hat{\\beta}-\\hat{\\beta}_0}{1-\\hat{\\beta}_0}`
where :math:`\\hat{\\beta}_0` is the estimated beta between samples which are
known to be unrelated [1].
Parameters
----------
ds
Genotype call dataset.
stat_identity_by_state
Input variable name holding stat_identity_by_state as defined
by :data:`sgkit.variables.stat_identity_by_state_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`identity_by_state`.
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_Weir_Goudet_beta_spec`
which is a matrix of estimated pairwise kinship relative to the average
kinship of all pairs of individuals in the dataset.
The dimensions are named ``samples_0`` and ``samples_1``.
Examples
--------
>>> import sgkit as sg
>>> ds = sg.simulate_genotype_call_dataset(n_variant=3, n_sample=3, n_allele=10, seed=3)
>>> # sample 2 "inherits" alleles from samples 0 and 1
>>> ds.call_genotype.data[:, 2, 0] = ds.call_genotype.data[:, 0, 0]
>>> ds.call_genotype.data[:, 2, 1] = ds.call_genotype.data[:, 1, 0]
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1 S2
variants
0 7/1 8/6 7/8
1 9/5 3/6 9/3
2 8/8 8/3 8/8
>>> # estimate beta
>>> ds = sg.Weir_Goudet_beta(ds).compute()
>>> ds.stat_Weir_Goudet_beta.values # doctest: +NORMALIZE_WHITESPACE
array([[ 0.5 , -0.25, 0.25],
[-0.25, 0.25, 0. ],
[ 0.25, 0. , 0.5 ]])
>>> # correct beta assuming least related samples are unrelated
>>> beta = ds.stat_Weir_Goudet_beta
>>> beta0 = beta.min()
>>> beta_corrected = (beta - beta0) / (1 - beta0)
>>> beta_corrected.values # doctest: +NORMALIZE_WHITESPACE
array([[0.6, 0. , 0.4],
[0. , 0.4, 0.2],
[0.4, 0.2, 0.6]])
References
----------
[1] - Bruce, S. Weir, and Jérôme Goudet 2017.
"A Unified Characterization of Population Structure and Relatedness."
Genetics 206 (4): 2085-2103.
"""
ds = define_variable_if_absent(
ds, variables.stat_identity_by_state, stat_identity_by_state, identity_by_state
)
variables.validate(
ds, {stat_identity_by_state: variables.stat_identity_by_state_spec}
)
ibs = ds[stat_identity_by_state].data
# average matching is the mean of non-diagonal elements
num = da.nansum(da.tril(ibs, -1))
denom = da.nansum(da.tril(~da.isnan(ibs), -1))
avg = num / denom
beta = (ibs - avg) / (1 - avg)
new_ds = create_dataset(
{
variables.stat_Weir_Goudet_beta: (
("samples_0", "samples_1"),
beta,
)
}
)
return conditional_merge_datasets(ds, new_ds, merge)
def pc_relate(ds: xr.Dataset,
*,
maf: float = 0.01,
call_genotype: Hashable = variables.call_genotype,
call_genotype_mask: Hashable = variables.call_genotype_mask,
sample_pcs: Hashable = variables.sample_pcs,
merge: bool = True) -> xr.Dataset:
"""Compute PC-Relate as described in Conomos, et al. 2016 [1].
This method computes the kinship coefficient matrix. The kinship coefficient for
a pair of individuals ``i`` and ``j`` is commonly defined to be the probability that
a random allele selected from ``i`` and a random allele selected from ``j`` at
a locus are IBD. Several of the most common family relationships and their
corresponding kinship coefficient:
+--------------------------------------------------+---------------------+
| Relationship | Kinship coefficient |
+==================================================+=====================+
| Individual-self | 1/2 |
+--------------------------------------------------+---------------------+
| full sister/full brother | 1/4 |
+--------------------------------------------------+---------------------+
| mother/father/daughter/son | 1/4 |
+--------------------------------------------------+---------------------+
| grandmother/grandfather/granddaughter/grandson | 1/8 |
+--------------------------------------------------+---------------------+
| aunt/uncle/niece/nephew | 1/8 |
+--------------------------------------------------+---------------------+
| first cousin | 1/16 |
+--------------------------------------------------+---------------------+
| half-sister/half-brother | 1/8 |
+--------------------------------------------------+---------------------+
Parameters
----------
ds
Dataset containing (S = num samples, V = num variants, D = ploidy, PC = num PC)
- genotype calls: (SxVxD)
- genotype calls mask: (SxVxD)
- sample PCs: (PCxS)
maf
individual minor allele frequency filter. If an individual's estimated
individual-specific minor allele frequency at a SNP is less than this value,
that SNP will be excluded from the analysis for that individual.
The default value is 0.01. Must be between (0.0, 0.1).
call_genotype
Input variable name holding call_genotype.
Defined by :data:`sgkit.variables.call_genotype_spec`.
call_genotype_mask
Input variable name holding call_genotype_mask.
Defined by :data:`sgkit.variables.call_genotype_mask_spec`
sample_pcs
Input variable name holding sample_pcs.
Defined by :data:`sgkit.variables.sample_pcs_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.
Warnings
--------
This function is only applicable to diploid, biallelic datasets.
This version is compatible with the R implementation of PC Relate
method from the GENESIS package version 2.18.0.
Returns
-------
Dataset containing (S = num samples):
:data:`sgkit.variables.pc_relate_phi_spec`: (S,S) ArrayLike
pairwise recent kinship coefficient matrix as float in [-0.5, 0.5].
References
----------
[1] - Conomos, Matthew P., Alexander P. Reiner, Bruce S. Weir, and Timothy A. Thornton. 2016.
"Model-Free Estimation of Recent Genetic Relatedness."
American Journal of Human Genetics 98 (1): 127–48.
Raises
------
ValueError
If ploidy of provided dataset != 2
ValueError
If maximum number of alleles in provided dataset != 2
ValueError
Input dataset is missing any of the required variables
ValueError
If maf is not in (0.0, 1.0)
"""
if maf <= 0.0 or maf >= 1.0:
raise ValueError("MAF must be between (0.0, 1.0)")
if "ploidy" in ds.dims and ds.dims["ploidy"] != 2:
raise ValueError("PC Relate only works for diploid genotypes")
if "alleles" in ds.dims and ds.dims["alleles"] != 2:
raise ValueError("PC Relate only works for biallelic genotypes")
variables.validate(
ds,
{
call_genotype: variables.call_genotype_spec,
call_genotype_mask: variables.call_genotype_mask_spec,
sample_pcs: variables.sample_pcs_spec,
},
)
call_g, call_g_mask = _collapse_ploidy(ds, call_genotype,
call_genotype_mask)
imputed_call_g = _impute_genotype_call_with_variant_mean(
call_g, call_g_mask)
# 𝔼[gs|V] = 1β0 + Vβ, where 1 is a length _s_ vector of 1s, and β = (β1,...,βD)^T
# is a length D vector of regression coefficients for each of the PCs
pcs = ds[sample_pcs]
pcsi = da.concatenate([da.ones((1, pcs.shape[1]), dtype=pcs.dtype), pcs],
axis=0)
# Note: dask qr decomp requires no chunking in one dimension, and because number of
# components should be smaller than number of samples in most cases, we disable
# chunking on components
pcsi = pcsi.T.rechunk((None, -1))
q, r = da.linalg.qr(pcsi)
# mu, eq: 3
half_beta = da.linalg.inv(2 * r).dot(q.T).dot(imputed_call_g.T)
mu = pcsi.dot(half_beta).T
# phi, eq: 4
mask = (mu <= maf) | (mu >= 1.0 - maf) | call_g_mask
mu_mask = da.ma.masked_array(mu, mask=mask)
variance = mu_mask * (1.0 - mu_mask)
variance = da.ma.filled(variance, fill_value=0.0)
stddev = da.sqrt(variance)
centered_af = call_g / 2 - mu_mask
centered_af = da.ma.filled(centered_af, fill_value=0.0)
# NOTE: gramian could be a performance bottleneck, and we could explore
# performance improvements like (or maybe sth else):
# * calculating only the pairs we are interested in
# * using an optimized einsum.
assert centered_af.shape == call_g.shape
assert stddev.shape == call_g.shape
phi = gramian(centered_af) / gramian(stddev)
# NOTE: phi is of shape (S x S), S = num samples
assert phi.shape == (call_g.shape[1], ) * 2
new_ds = create_dataset(
{variables.pc_relate_phi: (("sample_x", "sample_y"), phi)})
return conditional_merge_datasets(ds, new_ds, merge)
def filter_partial_calls(
ds: Dataset,
*,
call_genotype: Hashable = variables.call_genotype,
merge: bool = True,
) -> Dataset:
"""Replace partial genotype calls with missing values.
Parameters
----------
ds
Genotype call dataset such as from
:func:`sgkit.create_genotype_call_dataset`.
call_genotype
Input variable name holding call_genotype as defined by
:data:`sgkit.variables.call_genotype_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
-------
Dataset containing `call_genotype_complete` and
`call_genotype_complete_mask` in which partial genotype calls are
replaced with completely missing genotype calls.
Examples
--------
>>> import sgkit as sg
>>> from sgkit.testing import simulate_genotype_call_dataset
>>> ds = simulate_genotype_call_dataset(n_variant=4, n_sample=2, seed=1, missing_pct=0.3)
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 ./0 ./.
1 ./0 1/1
2 0/1 ./0
3 ./0 0/0
>>> ds2 = filter_partial_calls(ds)
>>> ds2['call_genotype'] = ds2['call_genotype_complete']
>>> ds2['call_genotype_mask'] = ds2['call_genotype_complete_mask']
>>> sg.display_genotypes(ds2) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1
variants
0 ./. ./.
1 ./. 1/1
2 0/1 ./.
3 ./. 0/0
Notes
-----
The returned dataset will still contain the initial `call_genotype` and
`call_genotype_mask` variables. Many sgkit functions will default to
using `call_genotype` and/or `call_genotype_mask`, hence it is necessary
to overwrite these variables (see the example) or explicitly pass the new
variables as function arguments in order to remove partial calls from
futher analysis.
"""
variables.validate(ds, {call_genotype: variables.call_genotype_spec})
G = ds[call_genotype]
mixed_ploidy = G.attrs.get("mixed_ploidy", False)
if mixed_ploidy:
P = (G == -1).any(axis=-1) & (G >= -1)
else:
P = (G < 0).any(axis=-1)
F = xr.where(P, -1, G) # type: ignore[no-untyped-call]
new_ds = create_dataset({
variables.call_genotype_complete: F,
variables.call_genotype_complete_mask: F < 0,
})
new_ds[
variables.call_genotype_complete].attrs["mixed_ploidy"] = mixed_ploidy
return conditional_merge_datasets(ds, new_ds, merge)
def cohort_allele_frequencies(
ds: Dataset,
*,
cohort_allele_count: Hashable = variables.cohort_allele_count,
merge: bool = True,
) -> Dataset:
"""Compute allele frequencies for each cohort.
Parameters
----------
ds
Dataset containing genotype calls.
cohort_allele_count
Input variable name holding cohort_allele_count as 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 :data:`sgkit.variables.cohort_allele_frequency_spec`
of allele frequencies with shape (variants, cohorts, alleles) and values
corresponding to the frequency of non-missing occurrences of each allele.
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
>>> ds["sample_cohort"] = xr.DataArray(np.repeat([0, 1], ds.dims["samples"] // 2), dims="samples")
>>> sg.display_genotypes(ds) # doctest: +NORMALIZE_WHITESPACE
samples S0 S1 S2 S3
variants
0 0/0 1/0 1/0 0/1
1 1/0 0/1 0/0 1/0
2 1/1 0/0 1/0 0/1
3 1/0 1/1 1/1 1/0
4 1/0 0/0 1/0 1/1
>>> sg.cohort_allele_frequencies(ds)["cohort_allele_frequency"].values # doctest: +NORMALIZE_WHITESPACE
array([[[0.75, 0.25],
[0.5 , 0.5 ]],
<BLANKLINE>
[[0.5 , 0.5 ],
[0.75, 0.25]],
<BLANKLINE>
[[0.5 , 0.5 ],
[0.5 , 0.5 ]],
<BLANKLINE>
[[0.25, 0.75],
[0.25, 0.75]],
<BLANKLINE>
[[0.75, 0.25],
[0.25, 0.75]]])
"""
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]
AF = AC / AC.sum(dim="alleles")
new_ds = create_dataset({variables.cohort_allele_frequency: AF})
return conditional_merge_datasets(ds, new_ds, merge)
