def test_diversity(sample_size, chunks, cohort_allele_count):
ts = simulate_ts(sample_size)
ds = ts_to_dataset(ts, chunks)
ds, subsets = add_cohorts(ds, ts, cohort_key_names=["cohorts"])
if cohort_allele_count is not None:
ds = count_cohort_alleles(ds, merge=False).rename(
{variables.cohort_allele_count: cohort_allele_count})
ds = ds.assign_coords({"cohorts": ["co_0"]})
ds = diversity(ds, cohort_allele_count=cohort_allele_count)
else:
ds = ds.assign_coords({"cohorts": ["co_0"]})
ds = diversity(ds)
div = ds.stat_diversity.sum(axis=0,
skipna=False).sel(cohorts="co_0").values
ts_div = ts.diversity(span_normalise=False)
np.testing.assert_allclose(div, ts_div)
def test_diversity(sample_size, chunks, cohort_allele_count):
ts = msprime.simulate(sample_size,
length=100,
mutation_rate=0.05,
random_seed=42)
ds = ts_to_dataset(ts, chunks) # type: ignore[no-untyped-call]
ds, subsets = add_cohorts(
ds, ts, cohort_key_names=["cohorts"]) # type: ignore[no-untyped-call]
if cohort_allele_count is not None:
ds = count_cohort_alleles(ds, merge=False).rename(
{variables.cohort_allele_count: cohort_allele_count})
ds = ds.assign_coords({"cohorts": ["co_0"]})
ds = diversity(ds, cohort_allele_count=cohort_allele_count)
else:
ds = ds.assign_coords({"cohorts": ["co_0"]})
ds = diversity(ds)
div = ds.stat_diversity.sum(axis=0,
skipna=False).sel(cohorts="co_0").values
ts_div = ts.diversity(span_normalise=False)
np.testing.assert_allclose(div, ts_div)
def time_count_cohort_alleles(self) -> None:
count_cohort_alleles(self.count_cohort_alleles_ds)