def isin(element, test_elements, assume_unique=False, invert=False):
element = da.asarray(element)
test_elements = da.asarray(test_elements)
element_axes = tuple(range(element.ndim))
test_axes = tuple(i + element.ndim for i in range(test_elements.ndim))
mapped = da.atop(_isin_kernel, element_axes + test_axes,
element, element_axes,
test_elements, test_axes,
adjust_chunks={axis: lambda _: 1
for axis in test_axes},
dtype=bool,
assume_unique=assume_unique)
result = mapped.any(axis=test_axes)
if invert:
result = ~result
return result
def test_boolean_numpy_array_slicing():
with pytest.raises(IndexError):
da.asarray(range(2))[np.array([True])]
with pytest.raises(IndexError):
da.asarray(range(2))[np.array([False, False, False])]
x = np.arange(5)
ind = np.array([True, False, False, False, True])
assert_eq(da.asarray(x)[ind], x[ind])
# https://github.com/dask/dask/issues/3706
ind = np.array([True])
assert_eq(da.asarray([0])[ind], np.arange(1)[ind])
def test_tile_basic(reps):
a = da.asarray([0, 1, 2])
b = [[1, 2], [3, 4]]
assert_eq(np.tile(a.compute(), reps), da.tile(a, reps))
assert_eq(np.tile(b, reps), da.tile(b, reps))
def regenie_transform(
G: ArrayLike,
X: ArrayLike,
Y: ArrayLike,
contigs: ArrayLike,
*,
variant_block_size: Optional[Union[int, Tuple[int, ...]]] = None,
sample_block_size: Optional[Union[int, Tuple[int, ...]]] = None,
alphas: Optional[ArrayLike] = None,
add_intercept: bool = True,
orthogonalize: bool = False,
normalize: bool = False,
_glow_adj_dof: bool = False,
_glow_adj_alpha: bool = False,
_glow_adj_scaling: bool = False,
) -> Dataset:
"""Regenie trait transformation.
Parameters
----------
G
[array-like, shape: (M, N)]
Genotype data array, `M` samples by `N` variants.
X
[array-like, shape: (M, C)]
Covariate array, `M` samples by `C` covariates.
Y
[array-like, shape: (M, O)]
Outcome array, `M` samples by `O` outcomes.
contigs
[array-like, shape: (N,)]
Variant contigs as monotonic increasting integer contig index.
See the `regenie` function for documentation on remaining fields.
Returns
-------
A dataset containing the following variables:
- `base_prediction` (blocks, alphas, samples, outcomes): Stage 1
predictions from ridge regression reduction .
- `meta_prediction` (samples, outcomes): Stage 2 predictions from
the best meta estimator trained on the out-of-sample Stage 1
predictions.
- `loco_prediction` (contigs, samples, outcomes): LOCO predictions
resulting from Stage 2 predictions ignoring effects for variant
blocks on held out contigs. This will be absent if the
data provided does not contain at least 2 contigs.
Raises
------
ValueError
If `G`, `X`, and `Y` do not have the same size along
the first (samples) dimension.
"""
if not G.shape[0] == X.shape[0] == Y.shape[0]:
raise ValueError(
"All data arrays must have same size along first (samples) dimension "
f"(shapes provided: G={G.shape}, X={X.shape}, Y={Y.shape})")
n_sample = Y.shape[0]
n_variant = G.shape[1]
if alphas is not None:
alphas = np.asarray(alphas)
G, X, Y = da.asarray(G), da.asarray(X), da.asarray(Y)
contigs = da.asarray(contigs)
# Set default block sizes if not provided
if variant_block_size is None:
# Block in groups of 1000, unless dataset is small
# enough to default to 2 blocks (typically for tests)
variant_block_size = min(1000, n_variant // 2)
if sample_block_size is None:
# Break into 10 chunks of approximately equal size
sample_block_size = tuple(
split_array_chunks(n_sample, min(10, n_sample)))
assert sum(sample_block_size) == n_sample
if normalize:
# See: https://github.com/projectglow/glow/issues/255
dof = 1 if _glow_adj_dof else 0
G = (G - G.mean(axis=0)) / G.std(axis=0, ddof=dof)
Y = (Y - Y.mean(axis=0)) / Y.std(axis=0)
X = (X - X.mean(axis=0)) / X.std(axis=0)
if add_intercept:
X = da.concatenate([da.ones((X.shape[0], 1), dtype=X.dtype), X],
axis=1)
# TODO: Test this after finding out whether or not there was a good reason
# it was precluded in glow by unit covariate regularization:
# https://github.com/projectglow/glow/issues/266
if orthogonalize: # pragma: no cover
G = G - X @ da.linalg.lstsq(X, G)[0]
Y = Y - X @ da.linalg.lstsq(X, Y)[0]
G = G / G.std(axis=0)
Y = Y / Y.std(axis=0)
X = da.zeros(shape=(n_sample, 0), dtype=G.dtype)
variant_chunk_start, variant_chunk_size = _variant_block_indexes(
variant_block_size, contigs)
G = G.rechunk(chunks=(sample_block_size, tuple(variant_chunk_size)))
X = X.rechunk(chunks=(sample_block_size, -1))
Y = Y.rechunk(chunks=(sample_block_size, -1))
YP1 = _stage_1(G, X, Y, alphas=alphas)
B2, YP2 = _stage_2(
YP1,
X,
Y,
alphas=alphas,
_glow_adj_alpha=_glow_adj_alpha,
_glow_adj_scaling=_glow_adj_scaling,
)
YP3 = _stage_3(B2, YP1, X, Y, contigs, variant_chunk_start)
data_vars: Dict[Hashable, Any] = {}
data_vars[variables.base_prediction] = xr.DataArray(
YP1,
dims=("blocks", "alphas", "samples", "outcomes"),
attrs={"description": DESC_BASE_PRED},
)
data_vars[variables.meta_prediction] = xr.DataArray(
YP2,
dims=("samples", "outcomes"),
attrs={"description": DESC_META_PRED})
if YP3 is not None:
data_vars[variables.loco_prediction] = xr.DataArray(
YP3,
dims=("contigs", "samples", "outcomes"),
attrs={"description": DESC_LOCO_PRED},
)
return create_dataset(data_vars)
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 gradient(f, *varargs, **kwargs):
f = da.asarray(f)
kwargs["edge_order"] = math.ceil(kwargs.get("edge_order", 1))
if kwargs["edge_order"] > 2:
raise ValueError("edge_order must be less than or equal to 2.")
drop_result_list = False
axis = kwargs.pop("axis", None)
if axis is None:
axis = tuple(range(f.ndim))
elif isinstance(axis, Integral):
drop_result_list = True
axis = (axis,)
axis = validate_axis(axis, f.ndim)
if len(axis) != len(set(axis)):
raise ValueError("duplicate axes not allowed")
axis = tuple(ax % f.ndim for ax in axis)
if varargs == ():
varargs = (1,)
if len(varargs) == 1:
varargs = len(axis) * varargs
if len(varargs) != len(axis):
raise TypeError(
"Spacing must either be a single scalar, or a scalar / "
"1d-array per axis"
)
if issubclass(f.dtype.type, (np.bool8, Integral)):
f = f.astype(float)
elif issubclass(f.dtype.type, Real) and f.dtype.itemsize < 4:
f = f.astype(float)
results = []
for i, ax in enumerate(axis):
for c in f.chunks[ax]:
if np.min(c) < kwargs["edge_order"] + 1:
raise ValueError(
'Chunk size must be larger than edge_order + 1. '
'Minimum chunk for aixs {} is {}. Rechunk to '
'proceed.'.format(np.min(c), ax))
if np.isscalar(varargs[i]):
array_locs = None
else:
if isinstance(varargs[i], da.Array):
raise NotImplementedError(
'dask array coordinated is not supported.')
# coordinate position for each block taking overlap into
# account
chunk = np.array(f.chunks[ax])
array_loc_stop = np.cumsum(chunk) + 1
array_loc_start = array_loc_stop - chunk - 2
array_loc_stop[-1] -= 1
array_loc_start[0] = 0
array_locs = (array_loc_start, array_loc_stop)
results.append(f.map_overlap(
_gradient_kernel,
dtype=f.dtype,
depth={j: 1 if j == ax else 0 for j in range(f.ndim)},
boundary="none",
coord=varargs[i],
axis=ax,
array_locs=array_locs,
grad_kwargs=kwargs,
))
if drop_result_list:
results = results[0]
return results
ax.set_aspect('equal')
plt.show()
# +
fig, ax = plt.subplots(figsize=[10, 10], constrained_layout=True)
base_extent = np.array(
[-dims[1] // 2, dims[1] // 2, -dims[2] // 2, dims[2] // 2])
ax.scatter(*cpc,
c=cropindices,
cmap='nipy_spectral',
zorder=5,
linewidths=1,
edgecolors='black')
cfac = 4
coarse_mask = da.coarsen(np.all, da.asarray(mask), {0: cfac, 1: cfac})
cropdata = da.coarsen(np.mean, data[cropindices], {1: cfac, 2: cfac}).persist()
xlim, ylim = np.array([ax.get_xlim(), ax.get_ylim()])
vmin, vmax = da.nanmin(cropdata).compute(), da.nanmax(cropdata).compute()
for i in range(len(cropdata)):
plt.imshow(
np.where(coarse_mask, cropdata[i], np.nan).T,
extent=base_extent +
np.array([cpc[0, i], cpc[0, i], cpc[1, i], cpc[1, i]]),
origin='lower',
#alpha=0.5,
cmap='gray',
vmin=vmin,
vmax=vmax,
)
aligned_prior_fluxes = aligned_prior_fluxes.transpose("flux_time", "dim_y",
"dim_x", "realization")
aligned_influences = aligned_influences.transpose("observation", "flux_time",
"dim_y", "dim_x")
write_progress_message("Rechunked to square")
aligned_influences = aligned_influences.fillna(0)
aligned_true_fluxes.astype(np.float32).load()
aligned_prior_fluxes.astype(np.float32).load()
aligned_influences.astype(np.float32).load()
write_progress_message("Loaded data")
# 23 min for seven towers over a month
# 11 min next run
# This includes the realignment time deferred with dask above
sparse_influences = sparse.COO(aligned_influences.values)
aligned_influences.data = da.asarray(sparse_influences)
print(datetime.datetime.now(UTC).strftime("%c"), "Converted to COO")
flush_output_streams()
posterior_var_atts = aligned_prior_fluxes.attrs.copy()
posterior_var_atts.update(
dict(
long_name="posterior_fluxes",
units=PRIOR_FLUXES_MATCHED[PRIOR_FLUX_NAME].attrs["units"],
description="posterior fluxes using dask for a month",
origin="OI using dask for a month",
prior_flux_name=PRIOR_FLUX_NAME,
flux_window=FLUX_WINDOW,
observation_window=OBS_WINDOW,
ancillary_variables="reduced_posterior_covariance",
def test_unique():
a = [1, 2, 2]
assert_allclose(unique(a), [1, 2])
a = da.asarray(a)
assert_allclose(unique(a), [1, 2])
def test_inverse_transform(self, array):
a = dpp.LabelEncoder()
assert_eq_ar(a.inverse_transform(a.fit_transform(array)),
da.asarray(array))
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 test_asarray():
y = da.asarray(xr.DataArray([1, 2, 3.0]))
assert isinstance(y, da.Array)
assert_eq(y, y)
def segment(
image,
channels,
model_type,
diameter,
fast_mode=False,
use_anisotropy=True,
iou_depth=2,
iou_threshold=0.7,
):
"""Use cellpose to segment nuclei in fluorescence data.
Parameters
----------
image : array of shape (z, y, x, channel)
Image used for detection of objects
channels : array of int with size 2
See cellpose
model_type : str
"cyto" or "nuclei"
diameter : tuple of size 3
Approximate diameter (in pixels) of a segmented region, i.e. cell width
fast_mode : bool
In fast mode, network averaging, tiling, and augmentation are turned off.
use_anisotropy : bool
If true, use anisotropy parameter of cellpose
iou_depth: dask depth parameter
Number of pixels of overlap to use in intersection-over-union calculation when
linking segments across neighboring, overlapping dask chunk regions.
iou_threshold: float
Minimum intersection-over-union in neighboring, overlapping dask chunk regions
to be considered the same segment. The region for calculating IOU is given by the
iou_depth parameter.
Returns:
segments : array of int32 with same shape as input
Each segmented cell is assigned a number and all its pixels contain that value (0 is background)
"""
assert image.ndim == 4, image.ndim
assert image.shape[-1] in {1, 2}, image.shape
assert diameter[1] == diameter[2], diameter
diameter_yx = diameter[1]
anisotropy = diameter[0] / diameter[1] if use_anisotropy else None
image = da.asarray(image)
image = image.rechunk({-1: -1}) # color channel is chunked together
depth = tuple(np.ceil(diameter).astype(np.int64))
boundary = "reflect"
# No chunking in channel direction
image = da.overlap.overlap(image, depth + (0, ), boundary)
block_iter = zip(
np.ndindex(*image.numblocks),
map(
functools.partial(operator.getitem, image),
da.core.slices_from_chunks(image.chunks),
),
)
labeled_blocks = np.empty(image.numblocks[:-1], dtype=object)
total = None
for index, input_block in block_iter:
labeled_block, n = dask.delayed(segment_chunk, nout=2)(
input_block,
channels,
model_type,
diameter_yx,
anisotropy,
fast_mode,
index,
)
shape = input_block.shape[:-1]
labeled_block = da.from_delayed(labeled_block,
shape=shape,
dtype=np.int32)
n = dask.delayed(np.int32)(n)
n = da.from_delayed(n, shape=(), dtype=np.int32)
total = n if total is None else total + n
block_label_offset = da.where(labeled_block > 0, total, np.int32(0))
labeled_block += block_label_offset
labeled_blocks[index[:-1]] = labeled_block
total += n
# Put all the blocks together
block_labeled = da.block(labeled_blocks.tolist())
depth = da.overlap.coerce_depth(len(depth), depth)
if np.prod(block_labeled.numblocks) > 1:
iou_depth = da.overlap.coerce_depth(len(depth), iou_depth)
if any(iou_depth[ax] > depth[ax] for ax in depth.keys()):
raise DistSegError("iou_depth (%s) > depth (%s)" %
(iou_depth, depth))
trim_depth = {k: depth[k] - iou_depth[k] for k in depth.keys()}
block_labeled = da.overlap.trim_internal(block_labeled,
trim_depth,
boundary=boundary)
block_labeled = link_labels(
block_labeled,
total,
iou_depth,
iou_threshold=iou_threshold,
)
block_labeled = da.overlap.trim_internal(block_labeled,
iou_depth,
boundary=boundary)
else:
block_labeled = da.overlap.trim_internal(block_labeled,
depth,
boundary=boundary)
return block_labeled
def dask_hist1d(
a: Array, bins=None, range=None, normed=False, weights=None, density=None
):
"""
Blocked variant of :func:`numpy.histogram`, but using the fast-histogram module.
Parameters
----------
a : array_like
Input data. The histogram is computed over the flattened array.
bins : int or sequence of scalars, optional
Either an iterable specifying the ``bins`` or the number of ``bins``
and a ``range`` argument is required as computing ``min`` and ``max``
over blocked arrays is an expensive operation that must be performed
explicitly.
If `bins` is an int, it defines the number of equal-width
bins in the given range (10, by default). If `bins` is a
sequence, it defines a monotonically increasing array of bin edges,
including the rightmost edge, allowing for non-uniform bin widths.
range : (float, float), optional
The lower and upper range of the bins. If not provided, range
is simply ``(a.min(), a.max())``. Values outside the range are
ignored. The first element of the range must be less than or
equal to the second. `range` affects the automatic bin
computation as well. While bin width is computed to be optimal
based on the actual data within `range`, the bin count will fill
the entire range including portions containing no data.
normed : bool, optional
This is equivalent to the ``density`` argument, but produces incorrect
results for unequal bin widths. It should not be used.
weights : array_like, optional
A dask.array.Array of weights, of the same block structure as ``a``. Each value in
``a`` only contributes its associated weight towards the bin count
(instead of 1). If ``density`` is True, the weights are
normalized, so that the integral of the density over the range
remains 1.
density : bool, optional
If ``False``, the result will contain the number of samples in
each bin. If ``True``, the result is the value of the
probability *density* function at the bin, normalized such that
the *integral* over the range is 1. Note that the sum of the
histogram values will not be equal to 1 unless bins of unity
width are chosen; it is not a probability *mass* function.
Overrides the ``normed`` keyword if given.
If ``density`` is True, ``bins`` cannot be a single-number delayed
value. It must be a concrete number, or a (possibly-delayed)
array/sequence of the bin edges.
Returns
-------
hist : dask Array
The values of the histogram. See `density` and `weights` for a
description of the possible semantics.
bin_edges : dask Array of dtype float
Return the bin edges ``(length(hist)+1)``.
Examples
--------
Using number of bins and range:
>>> import dask.array as da
>>> import numpy as np
>>> x = da.from_array(np.arange(10000), chunks=10)
>>> h, bins = da.histogram(x, bins=10, range=[0, 10000])
>>> bins
array([ 0., 1000., 2000., 3000., 4000., 5000., 6000., 7000.,
8000., 9000., 10000.])
>>> h.compute()
array([1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000, 1000])
Explicitly specifying the bins:
>>> h, bins = da.histogram(x, bins=np.array([0, 5000, 10000]))
>>> bins
array([ 0, 5000, 10000])
>>> h.compute()
array([5000, 5000])
"""
if isinstance(bins, Array):
scalar_bins = bins.ndim == 0
# ^ `np.ndim` is not implemented by Dask array.
elif isinstance(bins, Delayed):
scalar_bins = bins._length is None or bins._length == 1
else:
scalar_bins = np.ndim(bins) == 0
if bins is None or (scalar_bins and range is None):
raise ValueError(
"dask.array.histogram requires either specifying "
"bins as an iterable or specifying both a range and "
"the number of bins"
)
if weights is not None and weights.chunks != a.chunks:
raise ValueError("Input array and weights must have the same chunked structure")
if normed is not False:
raise ValueError(
"The normed= keyword argument has been deprecated. "
"Please use density instead. "
"See the numpy.histogram docstring for more information."
)
if density and scalar_bins and isinstance(bins, (Array, Delayed)):
raise NotImplementedError(
"When `density` is True, `bins` cannot be a scalar Dask object. "
"It must be a concrete number or a (possibly-delayed) array/sequence of bin edges."
)
for argname, val in [("bins", bins), ("range", range), ("weights", weights)]:
if not isinstance(bins, (Array, Delayed)) and is_dask_collection(bins):
raise TypeError(
"Dask types besides Array and Delayed are not supported "
"for `histogram`. For argument `{}`, got: {!r}".format(argname, val)
)
if range is not None:
try:
if len(range) != 2:
raise ValueError(
f"range must be a sequence or array of length 2, but got {len(range)} items"
)
if isinstance(range, (Array, np.ndarray)) and range.shape != (2,):
raise ValueError(
f"range must be a 1-dimensional array of two items, but got an array of shape {range.shape}"
)
except TypeError:
raise TypeError(
f"Expected a sequence or array for range, not {range}"
) from None
token = tokenize(a, bins, range, weights, density)
name = "histogram-sum-" + token
if scalar_bins:
bins = _linspace_from_delayed(range[0], range[1], bins + 1)
# ^ NOTE `range[1]` is safe because of the above check, and the initial check
# that range must not be None if `scalar_bins`
else:
if not isinstance(bins, (Array, np.ndarray)):
bins = asarray(bins)
if bins.ndim != 1:
raise ValueError(
f"bins must be a 1-dimensional array or sequence, got shape {bins.shape}"
)
(bins_ref, range_ref), deps = unpack_collections([bins, range])
# Map the histogram to all bins, forming a 2D array of histograms, stacked for each chunk
if weights is None:
dsk = {
(name, i, 0): (_block_fast_hist1d, k, bins_ref, range_ref)
for i, k in enumerate(flatten(a.__dask_keys__()))
}
dtype = np.histogram([])[0].dtype
else:
a_keys = flatten(a.__dask_keys__())
w_keys = flatten(weights.__dask_keys__())
dsk = {
(name, i, 0): (_block_fast_hist1d, k, bins_ref, range_ref, w)
for i, (k, w) in enumerate(zip(a_keys, w_keys))
}
dtype = weights.dtype
deps = (a,) + deps
if weights is not None:
deps += (weights,)
graph = HighLevelGraph.from_collections(name, dsk, dependencies=deps)
# Turn graph into a 2D Array of shape (nchunks, nbins)
nchunks = len(list(flatten(a.__dask_keys__())))
nbins = bins.size - 1 # since `bins` is 1D
chunks = ((1,) * nchunks, (nbins,))
mapped = Array(graph, name, chunks, dtype=dtype)
# Sum over chunks to get the final histogram
n = mapped.sum(axis=0)
# We need to replicate normed and density options from numpy
if density is not None:
if density:
db = asarray(np.diff(bins).astype(float), chunks=n.chunks)
return n / db / n.sum(), bins
else:
return n, bins
else:
return n, bins
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 pbs(
ds: Dataset,
*,
stat_Fst: Hashable = variables.stat_Fst,
cohorts: Optional[Sequence[Union[Tuple[int, int, int],
Tuple[str, str, str]]]] = None,
merge: bool = True,
) -> Dataset:
"""Compute the population branching statistic (PBS) between cohort triples.
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.
stat_Fst
Fst variable to use or calculate. Defined by
:data:`sgkit.variables.stat_Fst_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`Fst`.
cohorts
The cohort triples to compute statistics for, specified as a sequence of
tuples 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 PBS value between cohort triples, as defined by
:data:`sgkit.variables.stat_pbs_spec`.
Shape (variants, cohorts, cohorts, cohorts), or
(windows, cohorts, 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=6)
>>> # Divide samples into three named cohorts
>>> n_cohorts = 3
>>> sample_cohort = np.repeat(range(n_cohorts), ds.dims["samples"] // n_cohorts)
>>> ds["sample_cohort"] = xr.DataArray(sample_cohort, dims="samples")
>>> cohort_names = [f"co_{i}" for i in range(n_cohorts)]
>>> ds = ds.assign_coords({"cohorts_0": cohort_names, "cohorts_1": cohort_names, "cohorts_2": cohort_names})
>>> # Divide into two windows of size three (variants)
>>> ds = sg.window_by_variant(ds, size=3)
>>> sg.pbs(ds)["stat_pbs"].sel(cohorts_0="co_0", cohorts_1="co_1", cohorts_2="co_2").values # doctest: +NORMALIZE_WHITESPACE
array([ 0. , -0.160898])
"""
ds = define_variable_if_absent(ds, variables.stat_Fst, stat_Fst, Fst)
variables.validate(ds, {stat_Fst: variables.stat_Fst_spec})
fst = ds[variables.stat_Fst]
fst = fst.clip(min=0, max=(1 - np.finfo(float).epsneg))
t = -np.log(1 - fst)
n_cohorts = ds.dims["cohorts"]
n_windows = ds.dims["windows"]
assert_array_shape(t, n_windows, n_cohorts, n_cohorts)
# calculate PBS triples
t = da.asarray(t)
shape = (t.chunks[0], n_cohorts, n_cohorts, n_cohorts)
cohorts = cohorts or list(itertools.combinations(range(n_cohorts),
3)) # type: ignore
ct = _cohorts_to_array(cohorts, ds.indexes.get("cohorts_0", None))
p = da.map_blocks(lambda t: _pbs_cohorts(t, ct),
t,
chunks=shape,
new_axis=3,
dtype=np.float64)
assert_array_shape(p, n_windows, n_cohorts, n_cohorts, n_cohorts)
new_ds = create_dataset({
variables.stat_pbs:
(["windows", "cohorts_0", "cohorts_1", "cohorts_2"], p)
})
return conditional_merge_datasets(ds, new_ds, merge)
def test_lazy_nop(self):
src = self.realistic_cube[:2, :3, :10, :10]
src.data = da.asarray(src.data, chunks=((1, 1), (2, 1), (10,), (10,)))
res = regrid_area_weighted(src, src)
self.assertTrue(res.has_lazy_data())
self.assertEqual(res, src)
results = []
for variable_name, variable in ds.data_vars.items():
y = delayed(convert_variable)(variable_name, variable)
results.append(y)
total = delayed(sum)(results)
total.visualize()
# +
# %%time
print(total.compute(), "data variables converted")
# -
for variable_name in ["lat", "lon", "prob"]:
variable = ds[variable_name]
zarr_path = target_path.joinpath(variable_name)
if zarr_path.exists():
shutil.rmtree(zarr_path)
da.asarray(variable.data).to_zarr(str(zarr_path))
variable_name = "time"
time = ds[variable_name]
time_in_days = np.array(time[0], dtype="datetime64[D]") + np.arange(
(len(time) - 1) * days_per_month + 1)
time_in_days.shape
zarr_path = target_path.joinpath(variable_name)
if zarr_path.exists():
shutil.rmtree(zarr_path)
da.from_array(time_in_days).to_zarr(str(zarr_path))
def gwas_linear_regression(
ds: Dataset,
*,
dosage: Hashable,
covariates: Union[Hashable, Sequence[Hashable]],
traits: Union[Hashable, Sequence[Hashable]],
add_intercept: bool = True,
call_genotype: Hashable = variables.call_genotype,
merge: bool = True,
) -> Dataset:
"""Run linear regression to identify continuous trait associations with genetic variants.
This method solves OLS regressions for each variant simultaneously and reports
effect statistics as defined in [1]. This is facilitated by the removal of
sample (i.e. person/individual) covariates through orthogonal projection
of both the genetic variant and phenotype data [2]. A consequence of this
rotation is that effect sizes and significances cannot be reported for
covariates, only variants.
Parameters
----------
ds
Dataset containing necessary dependent and independent variables.
dosage
Name of genetic dosage variable.
Defined by :data:`sgkit.variables.dosage_spec`.
covariates
Names of covariate variables (1D or 2D).
Defined by :data:`sgkit.variables.covariates_spec`.
traits
Names of trait variables (1D or 2D).
Defined by :data:`sgkit.variables.traits_spec`.
add_intercept
Add intercept term to covariate set, by default True.
call_genotype
Input variable name holding call_genotype.
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.
Warnings
--------
Regression statistics from this implementation are only valid when an
intercept is present. The `add_intercept` flag is a convenience for adding one
when not already present, but there is currently no parameterization for
intercept-free regression.
Additionally, both covariate and trait arrays will be rechunked to have blocks
along the sample (row) dimension but not the column dimension (i.e.
they must be tall and skinny).
Returns
-------
Dataset containing (N = num variants, O = num traits):
variant_linreg_beta : [array-like, shape: (N, O)]
Beta values associated with each variant and trait
variant_linreg_t_value : [array-like, shape: (N, O)]
T statistics for each beta
variant_linreg_p_value : [array-like, shape: (N, O)]
P values as float in [0, 1]
References
----------
- [1] Hastie, Trevor, Robert Tibshirani, and Jerome Friedman. 2009. The Elements
of Statistical Learning: Data Mining, Inference, and Prediction, Second Edition.
Springer Science & Business Media.
- [2] Loh, Po-Ru, George Tucker, Brendan K. Bulik-Sullivan, Bjarni J. Vilhjálmsson,
Hilary K. Finucane, Rany M. Salem, Daniel I. Chasman, et al. 2015. “Efficient
Bayesian Mixed-Model Analysis Increases Association Power in Large Cohorts.”
Nature Genetics 47 (3): 284–90.
"""
if isinstance(covariates, Hashable):
covariates = [covariates]
if isinstance(traits, Hashable):
traits = [traits]
variables.validate(
ds,
{dosage: variables.dosage_spec},
{c: variables.covariates_spec for c in covariates},
{t: variables.traits_spec for t in traits},
)
G = _get_loop_covariates(ds, dosage=dosage, call_genotype=call_genotype)
if len(covariates) == 0:
if add_intercept:
X = da.ones((ds.dims["samples"], 1), dtype=np.float32)
else:
raise ValueError("add_intercept must be True if no covariates specified")
else:
X = da.asarray(concat_2d(ds[list(covariates)], dims=("samples", "covariates")))
if add_intercept:
X = da.concatenate([da.ones((X.shape[0], 1), dtype=X.dtype), X], axis=1)
# Note: dask qr decomp (used by lstsq) requires no chunking in one
# dimension, and because dim 0 will be far greater than the number
# of covariates for the large majority of use cases, chunking
# should be removed from dim 1. Also, dim 0 should have the same chunking
# as G dim 1, so that when XLP is computed in linear_regression() the
# two arrays have the same chunking.
X = X.rechunk((G.chunksize[1], -1))
Y = da.asarray(concat_2d(ds[list(traits)], dims=("samples", "traits")))
# Like covariates, traits must also be tall-skinny arrays
Y = Y.rechunk((None, -1))
res = linear_regression(G.T, X, Y)
new_ds = create_dataset(
{
variables.variant_linreg_beta: (("variants", "traits"), res.beta),
variables.variant_linreg_t_value: (("variants", "traits"), res.t_value),
variables.variant_linreg_p_value: (("variants", "traits"), res.p_value),
}
)
return conditional_merge_datasets(ds, new_ds, merge)
def hardy_weinberg_test(ds: Dataset,
*,
genotype_count: 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_count
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_count is not None:
variables.validate(ds, {genotype_count: variables.genotype_count_spec})
obs = list(da.asarray(ds[genotype_count]).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 linear_regression(
XL: ArrayLike, XC: ArrayLike, Y: ArrayLike
) -> LinearRegressionResult:
"""Efficient linear regression estimation for multiple covariate sets
Parameters
----------
XL
[array-like, shape: (M, N)]
"Loop" covariates for which N separate regressions will be run
XC
[array-like, shape: (M, P)]
"Core" covariates included in the regressions for each loop
covariate. All P core covariates are used in each of the N
loop covariate regressions.
Y
[array-like, shape: (M, O)]
Continuous outcomes
Returns
-------
Dataclass containing:
beta : [array-like, shape: (N, O)]
Beta values associated with each loop covariate and outcome
t_value : [array-like, shape: (N, O)]
T statistics for each beta
p_value : [array-like, shape: (N, O)]
P values as float in [0, 1]
"""
XL, XC = da.asarray(XL), da.asarray(XC) # Coerce for `lstsq`
if set([x.ndim for x in [XL, XC, Y]]) != {2}:
raise ValueError("All arguments must be 2D")
n_core_covar, n_loop_covar, n_obs, n_outcome = (
XC.shape[1],
XL.shape[1],
Y.shape[0],
Y.shape[1],
)
dof = n_obs - n_core_covar - 1
if dof < 1:
raise ValueError(
"Number of observations (N) too small to calculate sampling statistics. "
"N must be greater than number of core covariates (C) plus one. "
f"Arguments provided: N={n_obs}, C={n_core_covar}."
)
# Apply orthogonal projection to eliminate core covariates
# Note: QR factorization or SVD should be used here to find
# what are effectively OLS residuals rather than matrix inverse
# to avoid need for MxM array; additionally, dask.lstsq fails
# with numpy arrays
LS = XC @ da.linalg.lstsq(XC, XL)[0]
assert XL.chunksize == LS.chunksize
XLP = XL - LS
assert XLP.shape == (n_obs, n_loop_covar)
YP = Y - XC @ da.linalg.lstsq(XC, Y)[0]
assert YP.shape == (n_obs, n_outcome)
# Estimate coefficients for each loop covariate
# Note: A key assumption here is that 0-mean residuals
# from projection require no extra terms in variance
# estimate for loop covariates (columns of G), which is
# only true when an intercept is present.
XLPS = (XLP ** 2).sum(axis=0, keepdims=True).T
assert XLPS.shape == (n_loop_covar, 1)
B = (XLP.T @ YP) / XLPS
assert B.shape == (n_loop_covar, n_outcome)
# Compute residuals for each loop covariate and outcome separately
YR = YP[:, np.newaxis, :] - XLP[..., np.newaxis] * B[np.newaxis, ...]
assert YR.shape == (n_obs, n_loop_covar, n_outcome)
RSS = (YR ** 2).sum(axis=0)
assert RSS.shape == (n_loop_covar, n_outcome)
# Get t-statistics for coefficient estimates
T = B / np.sqrt(RSS / dof / XLPS)
assert T.shape == (n_loop_covar, n_outcome)
# Match to p-values
# Note: t dist not implemented in Dask so this must be delayed,
# see https://github.com/dask/dask/issues/6857
P = da.map_blocks(
lambda t: 2 * stats.distributions.t.sf(np.abs(t), dof), T, dtype="float64"
)
assert P.shape == (n_loop_covar, n_outcome)
return LinearRegressionResult(beta=B, t_value=T, p_value=P)
def test_ms_create(tmp_path, chunks, num_chans, corr_types, sources):
# Set up
rs = np.random.RandomState(42)
ms_path = tmp_path / "create.ms"
ms_table_name = str(ms_path)
ant_table_name = "::".join((ms_table_name, "ANTENNA"))
ddid_table_name = "::".join((ms_table_name, "DATA_DESCRIPTION"))
pol_table_name = "::".join((ms_table_name, "POLARIZATION"))
spw_table_name = "::".join((ms_table_name, "SPECTRAL_WINDOW"))
# SOURCE is an optional MS sub-table
src_table_name = "::".join((ms_table_name, "SOURCE"))
ms_datasets = []
ant_datasets = []
ddid_datasets = []
pol_datasets = []
spw_datasets = []
src_datasets = []
# For comparison
all_data_desc_id = []
all_data = []
# Create ANTENNA dataset of 64 antennas
# Each column in the ANTENNA has a fixed shape so we
# can represent all rows with one dataset
na = 64
position = da.random.random((na, 3)) * 10000
offset = da.random.random((na, 3))
names = np.array(['ANTENNA-%d' % i for i in range(na)], dtype=np.object)
ds = Dataset({
'POSITION': (("row", "xyz"), position),
'OFFSET': (("row", "xyz"), offset),
'NAME': (("row", ), da.from_array(names, chunks=na)),
})
ant_datasets.append(ds)
# Create SOURCE datasets
for s, (name, direction, rest_freq) in enumerate(sources):
dask_num_lines = da.full((1, ), len(rest_freq), dtype=np.int32)
dask_direction = da.asarray(direction)[None, :]
dask_rest_freq = da.asarray(rest_freq)[None, :]
dask_name = da.asarray(np.asarray([name], dtype=np.object), chunks=1)
ds = Dataset({
"NUM_LINES": (("row", ), dask_num_lines),
"NAME": (("row", ), dask_name),
"REST_FREQUENCY": (("row", "line"), dask_rest_freq),
"DIRECTION": (("row", "dir"), dask_direction),
})
src_datasets.append(ds)
# Create POLARISATION datasets.
# Dataset per output row required because column shapes are variable
for r, corr_type in enumerate(corr_types):
dask_num_corr = da.full((1, ), len(corr_type), dtype=np.int32)
dask_corr_type = da.from_array(corr_type,
chunks=len(corr_type))[None, :]
ds = Dataset({
"NUM_CORR": (("row", ), dask_num_corr),
"CORR_TYPE": (("row", "corr"), dask_corr_type),
})
pol_datasets.append(ds)
# Create multiple MeerKAT L-band SPECTRAL_WINDOW datasets
# Dataset per output row required because column shapes are variable
for num_chan in num_chans:
dask_num_chan = da.full((1, ), num_chan, dtype=np.int32)
dask_chan_freq = da.linspace(.856e9,
2 * .856e9,
num_chan,
chunks=num_chan)[None, :]
dask_chan_width = da.full((1, num_chan), .856e9 / num_chan)
ds = Dataset({
"NUM_CHAN": (("row", ), dask_num_chan),
"CHAN_FREQ": (("row", "chan"), dask_chan_freq),
"CHAN_WIDTH": (("row", "chan"), dask_chan_width),
})
spw_datasets.append(ds)
# For each cartesian product of SPECTRAL_WINDOW and POLARIZATION
# create a corresponding DATA_DESCRIPTION.
# Each column has fixed shape so we handle all rows at once
spw_ids, pol_ids = zip(
*product(range(len(num_chans)), range(len(corr_types))))
dask_spw_ids = da.asarray(np.asarray(spw_ids, dtype=np.int32))
dask_pol_ids = da.asarray(np.asarray(pol_ids, dtype=np.int32))
ddid_datasets.append(
Dataset({
"SPECTRAL_WINDOW_ID": (("row", ), dask_spw_ids),
"POLARIZATION_ID": (("row", ), dask_pol_ids),
}))
# Now create the associated MS dataset
for ddid, (spw_id, pol_id) in enumerate(zip(spw_ids, pol_ids)):
# Infer row, chan and correlation shape
row = sum(chunks['row'])
chan = spw_datasets[spw_id].CHAN_FREQ.shape[1]
corr = pol_datasets[pol_id].CORR_TYPE.shape[1]
# Create some dask vis data
dims = ("row", "chan", "corr")
np_data = (rs.normal(size=(row, chan, corr)) +
1j * rs.normal(size=(row, chan, corr))).astype(np.complex64)
data_chunks = tuple((chunks['row'], chan, corr))
dask_data = da.from_array(np_data, chunks=data_chunks)
# Create dask ddid column
dask_ddid = da.full(row, ddid, chunks=chunks['row'], dtype=np.int32)
dataset = Dataset({
'DATA': (dims, dask_data),
'DATA_DESC_ID': (("row", ), dask_ddid)
})
ms_datasets.append(dataset)
all_data.append(dask_data)
all_data_desc_id.append(dask_ddid)
ms_writes = xds_to_table(ms_datasets, ms_table_name, columns="ALL")
ant_writes = xds_to_table(ant_datasets, ant_table_name, columns="ALL")
pol_writes = xds_to_table(pol_datasets, pol_table_name, columns="ALL")
spw_writes = xds_to_table(spw_datasets, spw_table_name, columns="ALL")
ddid_writes = xds_to_table(ddid_datasets, ddid_table_name, columns="ALL")
source_writes = xds_to_table(src_datasets, src_table_name, columns="ALL")
dask.compute(ms_writes, ant_writes, pol_writes, spw_writes, ddid_writes,
source_writes)
# Check ANTENNA table correctly created
with pt.table(ant_table_name, ack=False) as A:
assert_array_equal(A.getcol("NAME"), names)
assert_array_equal(A.getcol("POSITION"), position)
assert_array_equal(A.getcol("OFFSET"), offset)
required_desc = pt.required_ms_desc("ANTENNA")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(A.colnames()) == set(required_columns)
# Check POLARIZATION table correctly created
with pt.table(pol_table_name, ack=False) as P:
for r, corr_type in enumerate(corr_types):
assert_array_equal(P.getcol("CORR_TYPE", startrow=r, nrow=1),
[corr_type])
assert_array_equal(P.getcol("NUM_CORR", startrow=r, nrow=1),
[len(corr_type)])
required_desc = pt.required_ms_desc("POLARIZATION")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(P.colnames()) == set(required_columns)
# Check SPECTRAL_WINDOW table correctly created
with pt.table(spw_table_name, ack=False) as S:
for r, num_chan in enumerate(num_chans):
assert_array_equal(
S.getcol("NUM_CHAN", startrow=r, nrow=1)[0], num_chan)
assert_array_equal(
S.getcol("CHAN_FREQ", startrow=r, nrow=1)[0],
np.linspace(.856e9, 2 * .856e9, num_chan))
assert_array_equal(
S.getcol("CHAN_WIDTH", startrow=r, nrow=1)[0],
np.full(num_chan, .856e9 / num_chan))
required_desc = pt.required_ms_desc("SPECTRAL_WINDOW")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(S.colnames()) == set(required_columns)
# We should get a cartesian product out
with pt.table(ddid_table_name, ack=False) as D:
spw_id, pol_id = zip(
*product(range(len(num_chans)), range(len(corr_types))))
assert_array_equal(pol_id, D.getcol("POLARIZATION_ID"))
assert_array_equal(spw_id, D.getcol("SPECTRAL_WINDOW_ID"))
required_desc = pt.required_ms_desc("DATA_DESCRIPTION")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(D.colnames()) == set(required_columns)
with pt.table(src_table_name, ack=False) as S:
for r, (name, direction, rest_freq) in enumerate(sources):
assert_array_equal(S.getcol("NAME", startrow=r, nrow=1)[0], [name])
assert_array_equal(S.getcol("REST_FREQUENCY", startrow=r, nrow=1),
[rest_freq])
assert_array_equal(S.getcol("DIRECTION", startrow=r, nrow=1),
[direction])
with pt.table(ms_table_name, ack=False) as T:
# DATA_DESC_ID's are all the same shape
assert_array_equal(T.getcol("DATA_DESC_ID"),
da.concatenate(all_data_desc_id))
# DATA is variably shaped (on DATA_DESC_ID) so we
# compared each one separately.
for ddid, data in enumerate(all_data):
ms_data = T.getcol("DATA", startrow=ddid * row, nrow=row)
assert_array_equal(ms_data, data)
required_desc = pt.required_ms_desc()
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
# Check we have the required columns
assert set(T.colnames()) == required_columns.union(
["DATA", "DATA_DESC_ID"])
def gradient(f, *varargs, **kwargs):
f = da.asarray(f)
kwargs["edge_order"] = math.ceil(kwargs.get("edge_order", 1))
if kwargs["edge_order"] > 2:
raise ValueError("edge_order must be less than or equal to 2.")
drop_result_list = False
axis = kwargs.pop("axis", None)
if axis is None:
axis = tuple(range(f.ndim))
elif isinstance(axis, Integral):
drop_result_list = True
axis = (axis, )
axis = validate_axis(axis, f.ndim)
if len(axis) != len(set(axis)):
raise ValueError("duplicate axes not allowed")
axis = tuple(ax % f.ndim for ax in axis)
if varargs == ():
varargs = (1, )
if len(varargs) == 1:
varargs = len(axis) * varargs
if len(varargs) != len(axis):
raise TypeError(
"Spacing must either be a single scalar, or a scalar / "
"1d-array per axis")
if issubclass(f.dtype.type, (np.bool8, Integral)):
f = f.astype(float)
elif issubclass(f.dtype.type, Real) and f.dtype.itemsize < 4:
f = f.astype(float)
results = []
for i, ax in enumerate(axis):
for c in f.chunks[ax]:
if np.min(c) < kwargs["edge_order"] + 1:
raise ValueError(
'Chunk size must be larger than edge_order + 1. '
'Minimum chunk for aixs {} is {}. Rechunk to '
'proceed.'.format(np.min(c), ax))
if np.isscalar(varargs[i]):
array_locs = None
else:
if isinstance(varargs[i], da.Array):
raise NotImplementedError(
'dask array coordinated is not supported.')
# coordinate position for each block taking overlap into
# account
chunk = np.array(f.chunks[ax])
array_loc_stop = np.cumsum(chunk) + 1
array_loc_start = array_loc_stop - chunk - 2
array_loc_stop[-1] -= 1
array_loc_start[0] = 0
array_locs = (array_loc_start, array_loc_stop)
results.append(
f.map_overlap(
_gradient_kernel,
dtype=f.dtype,
depth={j: 1 if j == ax else 0
for j in range(f.ndim)},
boundary="none",
coord=varargs[i],
axis=ax,
array_locs=array_locs,
grad_kwargs=kwargs,
))
if drop_result_list:
results = results[0]
return results
def count_cohort_alleles(
ds: Dataset,
*,
call_allele_count: Hashable = variables.call_allele_count,
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`.
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)
ds['call_dosage_mask'] = ds['call_genotype_probability_mask']
return ds
def rechunk_to_zarr(
ds: Dataset,
store: Union[PathType, MutableMapping], # type: ignore[type-arg]
*,
mode: str = "w",
chunk_length: int = 10_000,
chunk_width: int = 10_000,
compressor: Any = zarr.Blosc(cname="zstd", clevel=7, shuffle=2),
compute: bool = True) -> ZarrStore:
ds = pack_variables(ds)
for v in ['call_genotype_probability', 'call_genotype_probability_mask']:
chunk_size = da.asarray(ds[v]).chunksize[0]
if chunk_length % chunk_size != 0:
raise ValueError(
f"Chunk size in variant dimension for variable '{v}' ({chunk_size}) "
f"must evenly divide target chunk size {chunk_length}")
ds[v] = ds[v].chunk(chunks=dict(samples=chunk_width))
encoding = encode_variables(ds, compressor=compressor)
return ds.to_zarr(store, mode=mode, encoding=encoding, compute=compute)
def rechunk_from_zarr(store: Union[PathType, MutableMapping],
chunk_length: int = 10_000,
chunk_width: int = 10_000,
mask_and_scale: bool = True):
ds = xr.open_zarr(store, mask_and_scale=mask_and_scale)
for v in ['call_genotype_probability', 'call_genotype_probability_mask']:
def threshold_local(image, block_size, method='gaussian', offset=0,
mode='reflect', param=None, cval=0):
"""Compute a threshold mask image based on local pixel neighborhood.
Also known as adaptive or dynamic thresholding[1]_. The threshold value is
the weighted mean for the local neighborhood of a pixel subtracted by a
constant. Alternatively the threshold can be determined dynamically by a
given function, using the 'generic' method.
Parameters
----------
image : (N, M) dask ndarray
Input image.
block_size : int or list/tuple/array
Size of pixel neighborhood which is used to calculate the
threshold value.
(1) A single value for use in all dimensions or
(2) A tuple, list, or array with length equal to image.ndim
method : {'generic', 'gaussian', 'mean', 'median'}, optional
Method used to determine adaptive threshold for local neighbourhood in
weighted mean image.
* 'generic': use custom function (see `param` parameter)
* 'gaussian': apply gaussian filter (see `param` parameter for custom\
sigma value)
* 'mean': apply arithmetic mean filter
* 'median': apply median rank filter
By default the 'gaussian' method is used.
offset : float, optional
Constant subtracted from weighted mean of neighborhood to calculate
the local threshold value. Default offset is 0.
mode : {'reflect', 'constant', 'nearest', 'mirror', 'wrap'}, optional
The mode parameter determines how the array borders are handled, where
cval is the value when mode is equal to 'constant'.
Default is 'reflect'.
param : {int, function}, optional
Either specify sigma for 'gaussian' method or function object for
'generic' method. This functions takes the flat array of local
neighbourhood as a single argument and returns the calculated
threshold for the centre pixel.
cval : float, optional
Value to fill past edges of input if mode is 'constant'.
Returns
-------
threshold : (N, M) dask ndarray
Threshold image. All pixels in the input image higher than the
corresponding pixel in the threshold image are considered foreground.
References
----------
.. [1] https://docs.opencv.org/modules/imgproc/doc/miscellaneous_transformations.html?highlight=threshold
Examples
--------
>>> import dask.array as da
>>> image = da.random.random((1000, 1000), chunks=(100, 100))
>>> result = threshold_local(image, 15, 'gaussian')
""" # noqa
image = image.astype(np.float64)
if method == 'generic':
if not callable(param):
raise ValueError("Must include a valid function to use as the "
"'param' keyword argument.")
thresh_image = _generic.generic_filter(image, param, block_size,
mode=mode, cval=cval)
elif method == 'gaussian':
if param is None:
sigma = (da.asarray(block_size) - 1) / 6.0
else:
sigma = param
thresh_image = _gaussian.gaussian_filter(image, sigma, mode=mode,
cval=cval)
elif method == 'mean':
thresh_image = _generic.generic_filter(
image, dispatch_threshold_local_mean(image), block_size, mode=mode,
cval=cval)
elif method == 'median':
thresh_image = _order.median_filter(image, block_size, mode=mode,
cval=cval)
else:
raise ValueError("Invalid method specified. Please use `generic`, "
"`gaussian`, `mean`, or `median`.")
return thresh_image - offset
raise(ValueError("data can only be rechunked to monthly and yearly zarr archives by now"))
ds
# -
chunk_dict = {"lat": 1, "lon": 1, "prob": 1}
ds_rechunked = ds.chunk(chunk_dict)
ds_rechunked
# +
# overwite potentially existing zarr files?
overwrite = False # if False, raises zarr.errors.ContainsArrayError if zarr archive already exists
# +
# stupid nanny and worker messages: just wait, computations are running (see dashboard for status)
results = []
for variable_name, variable in tqdm(ds_rechunked.variables.items()):
zarr_dir_path = target_path.joinpath(variable_name)
zarr_dir_name = str(zarr_dir_path)
print(zarr_dir_name)
if overwrite & zarr_dir_path.exists():
print('overwrtiting...')
shutil.rmtree(zarr_dir_path)
da.asarray(variable.data).to_zarr(zarr_dir_name)
# -
def window_statistic(
values: ArrayLike,
statistic: Callable[..., ArrayLike],
window_starts: ArrayLike,
window_stops: ArrayLike,
dtype: DType,
chunks: Any = None,
new_axis: Union[None, int, Iterable[int]] = None,
**kwargs: Any,
) -> da.Array:
values = da.asarray(values)
desired_chunks = chunks or values.chunks
window_lengths = window_stops - window_starts
depth = np.max(window_lengths) # type: ignore[no-untyped-call]
# Dask will raise an error if the last chunk size is smaller than the depth
# Workaround by rechunking to combine the last two chunks in first axis
# See https://github.com/dask/dask/issues/6597
if depth > values.chunks[0][-1]:
chunk0 = values.chunks[0]
new_chunk0 = tuple(list(chunk0[:-2]) + [chunk0[-2] + chunk0[-1]])
values = values.rechunk({0: new_chunk0})
chunks = values.chunks[0]
rel_window_starts, windows_per_chunk = _get_chunked_windows(
chunks, window_starts, window_stops)
# Add depth for map_overlap
rel_window_starts = rel_window_starts + depth
rel_window_stops = rel_window_starts + window_lengths
chunk_offsets = _sizes_to_start_offsets(windows_per_chunk)
def blockwise_moving_stat(x: ArrayLike,
block_info: Any = None) -> ArrayLike:
if block_info is None or len(block_info) == 0:
return np.array([])
chunk_number = block_info[0]["chunk-location"][0]
chunk_offset_start = chunk_offsets[chunk_number]
chunk_offset_stop = chunk_offsets[chunk_number + 1]
chunk_window_starts = rel_window_starts[
chunk_offset_start:chunk_offset_stop]
chunk_window_stops = rel_window_stops[
chunk_offset_start:chunk_offset_stop]
out = np.array([
statistic(x[i:j], **kwargs)
for i, j in zip(chunk_window_starts, chunk_window_stops)
])
return out
if values.ndim == 1:
new_chunks = (tuple(windows_per_chunk), )
else:
# depth is 0 except in first axis
depth = {0: depth}
# new chunks are same except in first axis
new_chunks = tuple([tuple(windows_per_chunk)] +
list(desired_chunks[1:])) # type: ignore
return values.map_overlap(
blockwise_moving_stat,
dtype=dtype,
chunks=new_chunks,
depth=depth,
boundary=0,
trim=False,
new_axis=new_axis,
)
def Fst(
ds: Dataset,
*,
estimator: Optional[str] = None,
stat_divergence: Hashable = variables.stat_divergence,
merge: bool = True,
) -> Dataset:
"""Compute Fst between pairs of cohorts.
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.
estimator
Determines the formula to use for computing Fst.
If None (the default), or ``Hudson``, Fst is calculated
using the method of Hudson (1992) elaborated by Bhatia et al. (2013),
(the same estimator as scikit-allel).
Other supported estimators include ``Nei`` (1986), (the same estimator
as tskit).
stat_divergence
Divergence variable to use or calculate. Defined by
:data:`sgkit.variables.stat_divergence_spec`.
If the variable is not present in ``ds``, it will be computed
using :func:`divergence`.
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 Fst value between pairs of cohorts, as defined by
:data:`sgkit.variables.stat_Fst_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.Fst(ds)["stat_Fst"].values # doctest: +NORMALIZE_WHITESPACE
array([[[ nan, -0.16666667],
[-0.16666667, nan]],
<BLANKLINE>
[[ nan, -0.16666667],
[-0.16666667, nan]],
<BLANKLINE>
[[ nan, -0.33333333],
[-0.33333333, nan]],
<BLANKLINE>
[[ nan, -0.33333333],
[-0.33333333, nan]],
<BLANKLINE>
[[ nan, 0.2 ],
[ 0.2 , nan]]])
>>> # Divide into windows of size three (variants)
>>> ds = sg.window_by_variant(ds, size=3)
>>> sg.Fst(ds)["stat_Fst"].values # doctest: +NORMALIZE_WHITESPACE
array([[[ nan, -0.22222222],
[-0.22222222, nan]],
<BLANKLINE>
[[ nan, 0. ],
[ 0. , nan]]])
"""
known_estimators = {"Hudson": _Fst_Hudson, "Nei": _Fst_Nei}
if estimator is not None and estimator not in known_estimators:
raise ValueError(
f"Estimator '{estimator}' is not a known estimator: {known_estimators.keys()}"
)
estimator = estimator or "Hudson"
ds = define_variable_if_absent(ds, variables.stat_divergence,
stat_divergence, divergence)
variables.validate(ds, {stat_divergence: variables.stat_divergence_spec})
n_cohorts = ds.dims["cohorts"]
gs = da.asarray(ds.stat_divergence)
shape = (gs.chunks[0], n_cohorts, n_cohorts)
fst = da.map_blocks(known_estimators[estimator],
gs,
chunks=shape,
dtype=np.float64)
# TODO: reinstate assert (first dim could be either variants or windows)
# assert_array_shape(fst, n_windows, n_cohorts, n_cohorts)
new_ds = create_dataset(
{variables.stat_Fst: (("windows", "cohorts_0", "cohorts_1"), fst)})
return conditional_merge_datasets(ds, new_ds, merge)
def test_tile_np_kroncompare_examples(shape, reps):
x = np.random.random(shape)
d = da.asarray(x)
assert_eq(np.tile(x, reps), da.tile(d, reps))
def pairwise_distance(
x: ArrayLike,
metric: str = "euclidean",
) -> np.ndarray:
"""Calculates the pairwise distance between all pairs of row vectors in the
given two dimensional array x.
To illustrate the algorithm consider the following (4, 5) two dimensional array:
[e.00, e.01, e.02, e.03, e.04]
[e.10, e.11, e.12, e.13, e.14]
[e.20, e.21, e.22, e.23, e.24]
[e.30, e.31, e.32, e.33, e.34]
The rows of the above matrix are the set of vectors. Now let's label all
the vectors as v0, v1, v2, v3.
The result will be a two dimensional symmetric matrix which will contain
the distance between all pairs. Since there are 4 vectors, calculating the
distance between each vector and every other vector, will result in 16
distances and the resultant array will be of size (4, 4) as follows:
[v0.v0, v0.v1, v0.v2, v0.v3]
[v1.v0, v1.v1, v1.v2, v1.v3]
[v2.v0, v2.v1, v2.v2, v2.v3]
[v3.v0, v3.v1, v3.v2, v3.v3]
The (i, j) position in the resulting array (matrix) denotes the distance
between vi and vj vectors.
Negative and nan values are considered as missing values. They are ignored
for all distance metric calculations.
Parameters
----------
x
[array-like, shape: (M, N)]
An array like two dimensional matrix. The rows are the
vectors used for comparison, i.e. for pairwise distance.
metric
The distance metric to use. The distance function can be
'euclidean' or 'correlation'.
Returns
-------
[array-like, shape: (M, M)]
A two dimensional distance matrix, which will be symmetric. The dimension
will be (M, M). The (i, j) position in the resulting array
(matrix) denotes the distance between ith and jth row vectors
in the input array.
Examples
--------
>>> from sgkit.distance.api import pairwise_distance
>>> import dask.array as da
>>> x = da.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]]).rechunk(2, 2)
>>> pairwise_distance(x, metric='euclidean')
array([[0. , 2.44948974, 4.69041576],
[2.44948974, 0. , 5.47722558],
[4.69041576, 5.47722558, 0. ]])
>>> import numpy as np
>>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]])
>>> pairwise_distance(x, metric='euclidean')
array([[0. , 2.44948974, 4.69041576],
[2.44948974, 0. , 5.47722558],
[4.69041576, 5.47722558, 0. ]])
>>> x = np.array([[6, 4, 1,], [4, 5, 2], [9, 7, 3]])
>>> pairwise_distance(x, metric='correlation')
array([[1.11022302e-16, 2.62956526e-01, 2.82353505e-03],
[2.62956526e-01, 0.00000000e+00, 2.14285714e-01],
[2.82353505e-03, 2.14285714e-01, 0.00000000e+00]])
"""
try:
metric_ufunc = getattr(metrics, metric)
except AttributeError:
raise NotImplementedError(f"Given metric: {metric} is not implemented.")
x = da.asarray(x)
x_distance = da.blockwise(
# Lambda wraps reshape for broadcast
lambda _x, _y: metric_ufunc(_x[:, None, :], _y),
"jk",
x,
"ji",
x,
"ki",
dtype="float64",
concatenate=True,
)
x_distance = da.triu(x_distance, 1) + da.triu(x_distance).T
return x_distance.compute()
def potential_energy(self):
"""
Specific potential energy calculation.
Calculates the specific potencial energy
of dark matter, star and gas particles.
Returns
-------
gx : `galaxy object`
New instanced galaxy specific potencial energy calculated for
stars, dark matter and gas particles.
Examples
--------
This returns the specific potential energy of stars, dark matter and
gas particles.
>>> import galaxychop as gc
>>> galaxy = gc.Galaxy(...)
>>> gpot = galaxy.potential_energy()
>>> pot_s, pot_dm, pot_g = gpot.pot_s, gpot.pot_dm, gpot.pot_g
Note
----
If the potentials are entered when the `galaxy` object is instanced,
then, the calculation of `potential_energy` will raise a `ValueError`.
"""
m_s = self.arr_.m_s
x_s = self.arr_.x_s
y_s = self.arr_.y_s
z_s = self.arr_.z_s
m_dm = self.arr_.m_dm
x_dm = self.arr_.x_dm
y_dm = self.arr_.y_dm
z_dm = self.arr_.z_dm
m_g = self.arr_.m_g
x_g = self.arr_.x_g
y_g = self.arr_.y_g
z_g = self.arr_.z_g
pot_s = self.arr_.pot_s
pot_dm = self.arr_.pot_dm
pot_g = self.arr_.pot_g
pot_s = self.arr_.pot_s
pot_dm = self.arr_.pot_dm
pot_g = self.arr_.pot_g
eps_s = self.arr_.eps_s
eps_dm = self.arr_.eps_dm
eps_g = self.arr_.eps_g
potential = np.concatenate([pot_s, pot_dm, pot_g])
if np.all(potential == 0.0):
x = np.hstack((x_s, x_dm, x_g))
y = np.hstack((y_s, y_dm, y_g))
z = np.hstack((z_s, z_dm, z_g))
m = np.hstack((m_s, m_dm, m_g))
eps = np.max([eps_s, eps_dm, eps_g])
pot = utils.potential(
da.asarray(x, chunks=100),
da.asarray(y, chunks=100),
da.asarray(z, chunks=100),
da.asarray(m, chunks=100),
da.asarray(eps),
)
num_s = len(m_s)
num = len(m_s) + len(m_dm)
pot_s = pot[:num_s]
pot_dm = pot[num_s:num]
pot_g = pot[num:]
new = attr.asdict(self, recurse=False)
del new["arr_"]
new.update(
pot_s=-pot_s * (u.km / u.s) ** 2,
pot_dm=-pot_dm * (u.km / u.s) ** 2,
pot_g=-pot_g * (u.km / u.s) ** 2,
)
return Galaxy(**new)
else:
raise ValueError("Potentials are already calculated")
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)
