def combine(ts):
# pylint: disable=protected-access
tmp = ts.annotate(
alleles=merge_alleles(ts.data.map(lambda d: d.alleles)),
rsid=hl.find(hl.is_defined, ts.data.map(lambda d: d.rsid)),
filters=hl.set(hl.flatten(ts.data.map(lambda d: hl.array(d.filters)))),
info=hl.struct(
DP=hl.sum(ts.data.map(lambda d: d.info.DP)),
MQ_DP=hl.sum(ts.data.map(lambda d: d.info.MQ_DP)),
QUALapprox=hl.sum(ts.data.map(lambda d: d.info.QUALapprox)),
RAW_MQ=hl.sum(ts.data.map(lambda d: d.info.RAW_MQ)),
VarDP=hl.sum(ts.data.map(lambda d: d.info.VarDP)),
SB=hl.array([
hl.sum(ts.data.map(lambda d: d.info.SB[0])),
hl.sum(ts.data.map(lambda d: d.info.SB[1])),
hl.sum(ts.data.map(lambda d: d.info.SB[2])),
hl.sum(ts.data.map(lambda d: d.info.SB[3]))
])))
tmp = tmp.annotate(
__entries=hl.bind(
lambda combined_allele_index:
hl.range(0, hl.len(tmp.data)).flatmap(
lambda i:
hl.cond(hl.is_missing(tmp.data[i].__entries),
hl.range(0, hl.len(tmp.g[i].__cols))
.map(lambda _: hl.null(tmp.data[i].__entries.dtype.element_type)),
hl.bind(
lambda old_to_new: tmp.data[i].__entries.map(lambda e: renumber_entry(e, old_to_new)),
hl.range(0, hl.len(tmp.data[i].alleles)).map(
lambda j: combined_allele_index[tmp.data[i].alleles[j]])))),
hl.dict(hl.range(0, hl.len(tmp.alleles)).map(
lambda j: hl.tuple([tmp.alleles[j], j])))))
tmp = tmp.annotate_globals(__cols=hl.flatten(tmp.g.map(lambda g: g.__cols)))
return tmp.drop('data', 'g')
def get_tx_expression_expr(
key_expr: hl.expr.StructExpression,
tx_ht: hl.Table,
csq_expr: hl.expr.StructExpression,
gene_field: str = "ensg",
csq_field: str = "csq",
tx_struct: str = "tx_annotation",
) -> hl.expr.Float64Expression:
"""
Pull appropriate transcript expression annotation struct given a specific locus and alleles (provided in `key_expr`).
Assumes that `key_expr` contains a locus and alleles.
Assumes that multi-allelic variants have been split in both `tx_ht` and `key_expr`.
:param row_key_expr: StructExpression containing locus and alleles to search in `tx_ht`.
:param tx_ht: Input Table containing transcript expression information.
:param csq_expr: Input StructExpression that contains VEP consequence information.
:param gene_field: Field in `csq_expr` that contains gene ID.
:param csq_field: Field in `csq_expr` that contains `most_severe_consequence` annotation.
:param tx_struct: StructExpression that contains transcript expression information.
:return: StructExpression that contains transcript expression information for given gene ID in `csq_expr`.
"""
return hl.find(
lambda csq: (csq[gene_field] == csq_expr.gene_id)
& (csq[csq_field] == csq_expr.most_severe_consequence),
tx_ht[key_expr][tx_struct],
)
def combine(ts):
def merge_alleles(alleles):
from hail.expr.functions import _num_allele_type, _allele_ints
return hl.rbind(
alleles.map(lambda a: hl.or_else(a[0], '')).fold(
lambda s, t: hl.cond(hl.len(s) > hl.len(t), s, t), ''),
lambda ref: hl.rbind(
alleles.map(lambda al: hl.rbind(
al[0], lambda r: hl.array([ref]).
extend(al[1:].map(lambda a: hl.rbind(
_num_allele_type(r, a), lambda at: hl.cond(
(_allele_ints['SNP'] == at)
| (_allele_ints['Insertion'] == at)
| (_allele_ints['Deletion'] == at)
| (_allele_ints['MNP'] == at)
| (_allele_ints['Complex'] == at), a + ref[hl.len(
r):], a)))))), lambda lal: hl.
struct(globl=hl.array([ref]).extend(
hl.array(hl.set(hl.flatten(lal)).remove(ref))),
local=lal)))
def renumber_entry(entry, old_to_new) -> StructExpression:
# global index of alternate (non-ref) alleles
return entry.annotate(LA=entry.LA.map(lambda lak: old_to_new[lak]))
if (ts.row.dtype, ts.globals.dtype) not in _merge_function_map:
f = hl.experimental.define_function(
lambda row, gbl: hl.rbind(
merge_alleles(row.data.map(lambda d: d.alleles)), lambda
alleles: hl.struct(
locus=row.locus,
alleles=alleles.globl,
rsid=hl.find(hl.is_defined, row.data.map(lambda d: d.rsid)
),
__entries=hl.bind(
lambda combined_allele_index: hl.
range(0, hl.len(row.data)).flatmap(lambda i: hl.cond(
hl.is_missing(row.data[i].__entries),
hl.range(0, hl.len(gbl.g[i].__cols)).map(
lambda _: hl.null(row.data[i].__entries.dtype.
element_type)),
hl.bind(
lambda old_to_new: row.data[i].__entries.map(
lambda e: renumber_entry(e, old_to_new)),
hl.range(0, hl.len(alleles.local[i])).map(
lambda j: combined_allele_index[
alleles.local[i][j]])))),
hl.dict(
hl.range(0, hl.len(alleles.globl)).map(
lambda j: hl.tuple([alleles.globl[j], j])))))),
ts.row.dtype, ts.globals.dtype)
_merge_function_map[(ts.row.dtype, ts.globals.dtype)] = f
merge_function = _merge_function_map[(ts.row.dtype, ts.globals.dtype)]
ts = Table(
TableMapRows(
ts._tir,
Apply(merge_function._name, merge_function._ret_type,
TopLevelReference('row'), TopLevelReference('global'))))
return ts.transmute_globals(
__cols=hl.flatten(ts.g.map(lambda g: g.__cols)))
def combine(ts):
# pylint: disable=protected-access
tmp = ts.annotate(
alleles=merge_alleles(ts.data.map(lambda d: d.alleles)),
rsid=hl.find(hl.is_defined, ts.data.map(lambda d: d.rsid)),
info=hl.struct(
MQ_DP=hl.sum(ts.data.map(lambda d: d.info.MQ_DP)),
QUALapprox=hl.sum(ts.data.map(lambda d: d.info.QUALapprox)),
RAW_MQ=hl.sum(ts.data.map(lambda d: d.info.RAW_MQ)),
VarDP=hl.sum(ts.data.map(lambda d: d.info.VarDP)),
SB_TABLE=hl.array([
hl.sum(ts.data.map(lambda d: d.info.SB_TABLE[0])),
hl.sum(ts.data.map(lambda d: d.info.SB_TABLE[1])),
hl.sum(ts.data.map(lambda d: d.info.SB_TABLE[2])),
hl.sum(ts.data.map(lambda d: d.info.SB_TABLE[3]))
])))
tmp = tmp.annotate(
__entries=hl.bind(
lambda combined_allele_index:
hl.range(0, hl.len(tmp.data)).flatmap(
lambda i:
hl.cond(hl.is_missing(tmp.data[i].__entries),
hl.range(0, hl.len(tmp.g[i].__cols))
.map(lambda _: hl.null(tmp.data[i].__entries.dtype.element_type)),
hl.bind(
lambda old_to_new: tmp.data[i].__entries.map(lambda e: renumber_entry(e, old_to_new)),
hl.array([0]).extend(
hl.range(0, hl.len(tmp.data[i].alleles)).map(
lambda j: combined_allele_index[tmp.data[i].alleles[j]]))))),
hl.dict(hl.range(1, hl.len(tmp.alleles) + 1).map(
lambda j: hl.tuple([tmp.alleles[j - 1], j])))))
tmp = tmp.annotate_globals(__cols=hl.flatten(tmp.g.map(lambda g: g.__cols)))
return tmp.drop('data', 'g')
def gnomad_coverage_stats_optimized():
mt = hl.read_matrix_table(resource('gnomad_dp_simulation.mt'))
mt = mt.annotate_rows(mean=hl.agg.mean(mt.x),
count_array=hl.rbind(hl.agg.counter(hl.min(100, mt.x)),
lambda c: hl.range(0, 100).map(lambda i: c.get(i, 0))))
mt = mt.annotate_rows(median=hl.rbind(hl.sum(mt.count_array) / 2, lambda s: hl.find(lambda x: x > s,
hl.array_scan(
lambda i, j: i + j,
0,
mt.count_array))),
**{f'above_{x}': hl.sum(mt.count_array[x:]) for x in [1, 5, 10, 15, 20, 25, 30, 50, 100]}
)
mt.rows()._force_count()
def combine_r(ts):
if (ts.row.dtype, ts.globals.dtype) not in _merge_function_map:
f = hl.experimental.define_function(
lambda row, gbl:
hl.struct(
locus=row.locus,
ref_allele=hl.find(hl.is_defined, row.data.map(lambda d: d.ref_allele)),
__entries=hl.range(0, hl.len(row.data)).flatmap(
lambda i:
hl.if_else(hl.is_missing(row.data[i]),
hl.range(0, hl.len(gbl.g[i].__cols))
.map(lambda _: hl.missing(row.data[i].__entries.dtype.element_type)),
row.data[i].__entries))),
ts.row.dtype, ts.globals.dtype)
_merge_function_map[(ts.row.dtype, ts.globals.dtype)] = f
merge_function = _merge_function_map[(ts.row.dtype, ts.globals.dtype)]
ts = Table(TableMapRows(ts._tir, Apply(merge_function._name,
merge_function._ret_type,
TopLevelReference('row'),
TopLevelReference('global'))))
return ts.transmute_globals(__cols=hl.flatten(ts.g.map(lambda g: g.__cols)))
def combine(ts):
def merge_alleles(alleles):
from hail.expr.functions import _num_allele_type, _allele_ints
return hl.rbind(
alleles.map(lambda a: hl.or_else(a[0], ''))
.fold(lambda s, t: hl.cond(hl.len(s) > hl.len(t), s, t), ''),
lambda ref:
hl.rbind(
alleles.map(
lambda al: hl.rbind(
al[0],
lambda r:
hl.array([ref]).extend(
al[1:].map(
lambda a:
hl.rbind(
_num_allele_type(r, a),
lambda at:
hl.cond(
(_allele_ints['SNP'] == at) |
(_allele_ints['Insertion'] == at) |
(_allele_ints['Deletion'] == at) |
(_allele_ints['MNP'] == at) |
(_allele_ints['Complex'] == at),
a + ref[hl.len(r):],
a)))))),
lambda lal:
hl.struct(
globl=hl.array([ref]).extend(hl.array(hl.set(hl.flatten(lal)).remove(ref))),
local=lal)))
def renumber_entry(entry, old_to_new) -> StructExpression:
# global index of alternate (non-ref) alleles
return entry.annotate(LA=entry.LA.map(lambda lak: old_to_new[lak]))
if (ts.row.dtype, ts.globals.dtype) not in _merge_function_map:
f = hl.experimental.define_function(
lambda row, gbl:
hl.rbind(
merge_alleles(row.data.map(lambda d: d.alleles)),
lambda alleles:
hl.struct(
locus=row.locus,
alleles=alleles.globl,
rsid=hl.find(hl.is_defined, row.data.map(lambda d: d.rsid)),
__entries=hl.bind(
lambda combined_allele_index:
hl.range(0, hl.len(row.data)).flatmap(
lambda i:
hl.cond(hl.is_missing(row.data[i].__entries),
hl.range(0, hl.len(gbl.g[i].__cols))
.map(lambda _: hl.null(row.data[i].__entries.dtype.element_type)),
hl.bind(
lambda old_to_new: row.data[i].__entries.map(
lambda e: renumber_entry(e, old_to_new)),
hl.range(0, hl.len(alleles.local[i])).map(
lambda j: combined_allele_index[alleles.local[i][j]])))),
hl.dict(hl.range(0, hl.len(alleles.globl)).map(
lambda j: hl.tuple([alleles.globl[j], j])))))),
ts.row.dtype, ts.globals.dtype)
_merge_function_map[(ts.row.dtype, ts.globals.dtype)] = f
merge_function = _merge_function_map[(ts.row.dtype, ts.globals.dtype)]
ts = Table(TableMapRows(ts._tir, Apply(merge_function._name,
TopLevelReference('row'),
TopLevelReference('global'))))
return ts.transmute_globals(__cols=hl.flatten(ts.g.map(lambda g: g.__cols)))
