Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def test_floor_division(self):
a_int32 = hl.array([2, 4, 8, 16, hl.null(tint32)])
a_int64 = a_int32.map(lambda x: hl.int64(x))
a_float32 = a_int32.map(lambda x: hl.float32(x))
a_float64 = a_int32.map(lambda x: hl.float64(x))
int32_4s = hl.array([4, 4, 4, 4, hl.null(tint32)])
int32_3s = hl.array([3, 3, 3, 3, hl.null(tint32)])
int64_3 = hl.int64(3)
int64_3s = int32_3s.map(lambda x: hl.int64(x))
float32_3 = hl.float32(3)
float32_3s = int32_3s.map(lambda x: hl.float32(x))
float64_3 = hl.float64(3)
float64_3s = int32_3s.map(lambda x: hl.float64(x))
expected = [0, 1, 2, 5, None]
expected_inv = [1, 0, 0, 0, None]
self.check_expr(a_int32 // 3, expected, tarray(tint32))
self.check_expr(a_int64 // 3, expected, tarray(tint64))
self.check_expr(a_float32 // 3, expected, tarray(tfloat32))
self.check_expr(a_float64 // 3, expected, tarray(tfloat64))
self.check_expr(3 // a_int32, expected_inv, tarray(tint32))
def field_to_array(ds, field):
return hl.cond(ds[field] != 0, hl.array([field]), hl.empty_array(hl.tstr))
def phase_haploid_proband_x_nonpar(
proband_call: hl.expr.CallExpression,
father_call: hl.expr.CallExpression,
mother_call: hl.expr.CallExpression
) -> hl.expr.ArrayExpression:
"""
Returns phased genotype calls in the case of a haploid proband in the non-PAR region of X
:param CallExpression proband_call: Input proband genotype call
:param CallExpression father_call: Input father genotype call
:param CallExpression mother_call: Input mother genotype call
:return: Array containing: phased proband call, phased father call, phased mother call
:rtype: ArrayExpression
"""
transmitted_allele = hl.zip_with_index(hl.array([mother_call[0], mother_call[1]])).find(lambda m: m[1] == proband_call[0])
return hl.cond(
hl.is_defined(transmitted_allele),
hl.array([
hl.call(proband_call[0], phased=True),
hl.or_missing(father_call.is_haploid(), hl.call(father_call[0], phased=True)),
phase_parent_call(mother_call, transmitted_allele[0])
]),
hl.null(hl.tarray(hl.tcall))
)
Parameters
----------
f : function ( (:class:`.Expression`, :class:`.Expression`) -> :class:`.Expression`)
Function which takes the cumulative value and the next element, and
returns a new value.
zero : :class:`.Expression`
Initial value to pass in as left argument of `f`.
Returns
-------
:class:`.Expression`.
"""
collection = self
if isinstance(collection.dtype, tset):
collection = hl.array(collection)
indices, aggregations = unify_all(collection, zero)
accum_name = Env.get_uid()
elt_name = Env.get_uid()
accum_ref = construct_variable(accum_name, zero.dtype, indices, aggregations)
elt_ref = construct_variable(elt_name, collection.dtype.element_type, collection._indices, collection._aggregations)
body = to_expr(f(accum_ref, elt_ref))
if body.dtype != zero.dtype:
zero_coercer = coercer_from_dtype(zero.dtype)
if zero_coercer.can_coerce(body.dtype):
body = zero_coercer.coerce(body)
else:
body_coercer = coercer_from_dtype(body.dtype)
if body_coercer.can_coerce(zero.dtype):
zero_coerced = body_coercer.coerce(zero)
"g": ["a", "c", "t"],
"t": ["a", "c", "g"],
}
)
context = None
for contig in contigs:
_context = hl.utils.range_table(
ref.contig_length(contig),
n_partitions=int(ref.contig_length(contig) / 5000000),
)
locus_expr = hl.locus(contig=contig, pos=_context.idx + 1, reference_genome=ref)
ref_allele_expr = locus_expr.sequence_context().lower()
if add_all_substitutions:
alleles_expr = hl.array([ref_allele_expr]).extend(
SUBSTITUTIONS_TABLE.get(ref_allele_expr, hl.empty_array(hl.tstr))
)
else:
alleles_expr = [ref_allele_expr]
_context = (
_context.select(locus=locus_expr, alleles=alleles_expr)
.key_by("locus", "alleles")
.drop("idx")
)
if excluded_intervals is not None:
_context = hl.filter_intervals(_context, excluded_intervals, keep=False)
if filter_n:
_context = _context.filter(_context.alleles[0] == "n", keep=False)
call_to_one_hot_alleles_array(father_call, alleles)
)
mother_v = call_to_one_hot_alleles_array(mother_call, alleles)
combinations = hl.flatmap(
lambda f:
hl.zip_with_index(mother_v)
.filter(lambda m: m[1] + f[1] == proband_v)
.map(lambda m: hl.struct(m=m[0], f=f[0])),
hl.zip_with_index(father_v)
)
return (
hl.cond(
hl.is_defined(combinations) & (hl.len(combinations) == 1),
hl.array([
hl.call(father_call[combinations[0].f], mother_call[combinations[0].m], phased=True),
hl.cond(father_call.is_haploid(), hl.call(father_call[0], phased=True), phase_parent_call(father_call, combinations[0].f)),
phase_parent_call(mother_call, combinations[0].m)
]),
hl.null(hl.tarray(hl.tcall))
)
def explode(self, f, array_agg_expr):
if array_agg_expr._aggregations:
raise ExpressionException("'{}.explode' does not support an already-aggregated expression as the argument to 'collection'".format(self.correct_prefix()))
_check_agg_bindings(array_agg_expr, self._agg_bindings)
if isinstance(array_agg_expr.dtype, tset):
array_agg_expr = hl.array(array_agg_expr)
elt = array_agg_expr.dtype.element_type
var = Env.get_uid()
ref = construct_expr(Ref(var), elt, array_agg_expr._indices)
self._agg_bindings.add(var)
aggregated = f(ref)
_check_agg_bindings(aggregated, self._agg_bindings)
self._agg_bindings.remove(var)
if not self._as_scan and not aggregated._aggregations:
raise ExpressionException("'{}.explode' must take mapping that contains aggregation expression.".format(self.correct_prefix()))
indices, _ = unify_all(array_agg_expr, aggregated)
aggregations = hl.utils.LinkedList(Aggregation)
if not self._as_scan:
aggregations = aggregations.push(Aggregation(array_agg_expr, aggregated))
return construct_expr(AggExplode(array_agg_expr._ir, var, aggregated._ir, self._as_scan),
def struct_or_empty(v):
return (hl.case()
.when(cond(v.locus), hl.array([new_struct(v, i)]))
.or_missing())
f : function ( (arg) -> :class:`.Expression`)
Function to transform each element of the collection.
Returns
-------
:class:`.CollectionExpression`.
Collection where each element has been transformed according to `f`.
"""
def transform_ir(array, name, body):
a = ArrayMap(array, name, body)
if isinstance(self.dtype, tset):
a = ToSet(a)
return a
array_map = hl.array(self)._ir_lambda_method(transform_ir, f, self._type.element_type, lambda t: self._type.__class__(t))
if isinstance(self._type, tset):
return hl.set(array_map)
assert isinstance(self._type, tarray)
return array_map