def index_genotype_ac_array(g, item, array_cls, vector_cls):
# apply indexing operation to underlying values
out = g.values[item]
# decide whether to wrap the output, if so how
wrap_array = (
hasattr(out, 'ndim') and out.ndim == 3 and # dimensionality preserved
out.shape[2] == g.shape[2] and # alleles preserved
not contains_newaxis(item))
wrap_vector = (
# single row selection
isinstance(item, int) or (
# other way to make a single row selection
isinstance(item, tuple) and len(item) == 2
and isinstance(item[0], int)
and isinstance(item[1], (slice, list, np.ndarray, type(Ellipsis))))
or (
# single column selection
isinstance(item, tuple) and len(item) == 2
and isinstance(item[0], (slice, list, np.ndarray))
and isinstance(item[1], int)))
if wrap_array:
out = array_cls(out)
if wrap_vector:
out = vector_cls(out)
return out
def index_haplotype_array(h, item, cls):
# apply indexing operation on underlying values
out = h.values[item]
# decide whether to wrap the result as HaplotypeArray
wrap_array = (
hasattr(out, 'ndim') and out.ndim == 2 and # dimensionality preserved
not contains_newaxis(item))
if wrap_array:
out = cls(out)
return out
def index_genotype_ac_vector(g, item, cls):
# apply indexing operation on underlying values
out = g.values[item]
# decide whether to wrap the result
wrap = (
hasattr(out, 'ndim') and out.ndim == 2 and # dimensionality preserved
out.shape[1] == g.shape[1] and # alleles preserved
not contains_newaxis(item))
if wrap:
out = cls(out)
return out
def index_allele_counts_array(ac, item, cls):
# apply indexing operation on underlying values
out = ac.values[item]
# decide whether to wrap the result as HaplotypeArray
wrap_array = (
hasattr(out, 'ndim') and out.ndim == 2 and # dimensionality preserved
ac.shape[1] == out.shape[1] and # number of alleles preserved
not contains_newaxis(item))
if wrap_array:
out = cls(out)
return out
def index_genotype_vector(g, item, cls):
# apply indexing operation on underlying values
out = g.values[item]
# decide whether to wrap the result
wrap = (
hasattr(out, 'ndim') and out.ndim == 2 and # dimensionality preserved
out.shape[1] == g.shape[1] and # ploidy preserved
not contains_newaxis(item))
if wrap:
out = cls(out)
if g.mask is not None:
out.mask = g.mask[item]
if g.is_phased is not None:
out.is_phased = g.is_phased[item]
return out
