"""Script containing generic constraint functions that may be used in the constraint pipeline."""
import copy
import functools
import logging
import operator
from typing import Callable, Dict, List, Optional, Tuple, Union
import hail as hl
from hail.utils.misc import divide_null, new_temp_file
from gnomad.assessment.summary_stats import generate_filter_combinations
from gnomad.utils.reference_genome import get_primary_contigs, get_reference_genome
from gnomad.utils.vep import (
add_most_severe_consequence_to_consequence,
explode_by_vep_annotation,
process_consequences,
)
logging.basicConfig(
format="%(asctime)s (%(name)s %(lineno)s): %(message)s",
datefmt="%m/%d/%Y %I:%M:%S %p",
)
logger = logging.getLogger("constraint_utils")
logger.setLevel(logging.INFO)
COVERAGE_CUTOFF = 30
"""
Minimum median exome coverage differentiating high coverage sites from low coverage sites.
Low coverage sites require an extra calibration when computing the proportion of expected variation.
"""
CLASSIC_LOF_ANNOTATIONS = (
"stop_gained",
"splice_donor_variant",
"splice_acceptor_variant",
)
"""Classic loss-of-function VEP annotations."""
DEFAULT_FIELDS_TO_SUM = (
"mu_snp",
"mu",
"observed_variants",
"possible_variants",
"predicted_proportion_observed",
"coverage_correction",
"expected_variants",
)
"""Default fields summed when aggregating expected variants by constraint group."""
DEFAULT_GENCODE_ANNOTATIONS = (
"transcript_id_version",
"gene_id_version",
"level",
"transcript_type",
"start_position",
"end_position",
)
"""Default GENCODE annotations added to a Table by transcript id."""
[docs]def get_mu_annotation_expr(
ht: hl.Table,
mutation_ht: hl.Table,
mu_expr: Union[str, hl.expr.Float64Expression] = "mu_snp",
) -> hl.expr.Float64Expression:
"""
Get mutation rate annotation expression from the mutation rate Table.
.. note::
Function expects that `ht` includes `mutation_ht`'s key fields. Note that these
annotations don't need to be the keys of `ht`. The standard keys used are:
'context', 'ref', 'alt', and 'methylation_level'.
This function converts the mutation rate Table into a dictionary and uses the
mutation rate HT key fields in the input Table to get the mutation rate annotation
expression. The dictionary is used instead of joining the mutation rate Table to
the input Table to avoid unnecessary shuffling given the small size of the
mutation rate Table.
:param ht: Input Table that will be annotated with the returned mutation rate
expression.
:param mutation_ht: Mutation rate Table.
:param mu_expr: Mutation rate expression or annotation name in `mutation_ht`.
Default is 'mu_snp'.
:return: Mutation rate annotation expression.
"""
if isinstance(mu_expr, str):
mu_expr = mutation_ht[mu_expr]
mutation_dict = mutation_ht.aggregate(
hl.agg.group_by(mutation_ht.key, hl.agg.collect(mu_expr)), _localize=False
).map_values(lambda x: x[0])
return mutation_dict.get(hl.struct(**{k: ht[k] for k in mutation_ht.key}))
[docs]def annotate_with_mu(
ht: hl.Table,
mutation_ht: hl.Table,
mu_annotation: str = "mu_snp",
) -> hl.Table:
"""
Annotate SNP mutation rate for the input Table.
.. note::
Function expects that `ht` includes `mutation_ht`'s key fields. Note that these
annotations don't need to be the keys of `ht`. The standard keys used are:
'context', 'ref', 'alt', and 'methylation_level'.
:param ht: Input Table to annotate.
:param mutation_ht: Mutation rate Table.
:param mu_annotation: The name of mutation rate annotation in `mutation_ht`.
Default is 'mu_snp'.
:return: Table with mutational rate annotation added.
"""
mu = get_mu_annotation_expr(ht, mutation_ht, mutation_ht[mu_annotation])
return ht.annotate(
**{mu_annotation: hl.case().when(hl.is_defined(mu), mu).or_error("Missing mu")}
)
def _resolve_row_annotation_expr(
t: Optional[Union[hl.Table, hl.MatrixTable]] = None,
annotation_name: Optional[str] = None,
expr: Optional[hl.expr.Expression] = None,
expr_param_name: Optional[str] = None,
) -> hl.expr.Expression:
"""
Get a row annotation from a Table/MatrixTable, or return a provided expression.
Provides a consistent pattern for functions that accept either an explicit
Hail expression or fall back to a well-known field on a Table. This avoids
duplicating "resolve expr or look it up on ht" logic across callers.
If ``expr`` is provided it is returned directly. Otherwise, ``t`` and
``annotation_name`` must both be supplied and the named row field is looked up
on ``t``.
Example usage inside a public function::
def my_function(
ht: hl.Table,
freq_expr: Optional[hl.expr.StructExpression] = None,
) -> ...:
freq_expr = _resolve_row_annotation_expr(
t=ht,
annotation_name="freq",
expr=freq_expr,
expr_param_name="freq_expr",
)
# freq_expr is now guaranteed to be a valid expression —
# either the caller's explicit value or ht.freq.
:param t: Input Table or MatrixTable to look up ``annotation_name`` on. Required
when ``expr`` is None.
:param annotation_name: Name of the field to retrieve from ``t``. Required when
``expr`` is None.
:param expr: Expression to return directly. When provided, ``t`` and
``annotation_name`` are ignored.
:param expr_param_name: Human-readable name for ``expr``, used in log/error
messages. Defaults to "expr".
:return: The resolved Hail expression.
"""
if expr is None and (t is None or annotation_name is None):
raise ValueError("Either 't' and 'annotation_name' or 'expr' must be provided.")
expr_param_name = expr_param_name or "expr"
if expr is None and annotation_name in t.row:
logger.info(
"%s was not provided, using '%s'.", expr_param_name, annotation_name
)
expr = t[annotation_name]
elif expr is None:
raise ValueError(
f"{expr_param_name} was not provided and '{annotation_name}' is "
f"not present in the input Table or MatrixTable."
)
return expr
[docs]def variant_observed_expr(
freq_expr: Optional[
Union[hl.expr.StructExpression, hl.expr.ArrayExpression]
] = None,
ht: Optional[hl.Table] = None,
singleton: bool = False,
max_af: Optional[float] = None,
count_missing: bool = False,
) -> hl.expr.Int32Expression:
"""
Return 1 if a variant meets frequency criteria, 0 otherwise.
One of ``ht`` or ``freq_expr`` must be provided. If ``freq_expr`` is not provided,
``ht.freq`` is used. When ``freq_expr`` is an ArrayExpression, the first
element is used.
The returned count is 1 when:
- ``singleton`` is True and ``freq_expr.AC == 1``.
- ``max_af`` is not None and ``freq_expr.AC > 0`` and
``freq_expr.AF <= max_af``.
- Neither ``singleton`` nor ``max_af`` is set but ``freq_expr`` is
available: ``freq_expr.AC > 0``.
- Neither ``singleton`` nor ``max_af`` is set and no ``freq_expr`` is
available: unconditionally 1.
:param ht: Input Table. Used to look up ``freq`` when ``freq_expr`` is None.
:param freq_expr: StructExpression (or ArrayExpression of structs) with
``AC`` and ``AF`` fields. When an array, ``freq_expr[0]`` is used.
:param singleton: Count only singletons (AC == 1). When True, ``max_af`` is
ignored. Default is False.
:param max_af: Maximum allele frequency threshold. Default is None (no
cutoff).
:param count_missing: Whether to count a variant whose count expression is
missing (e.g. frequency is None). True counts it as 1, False as 0. Default
is False.
:return: Int32Expression equal to 0 or 1.
"""
if ht is None and freq_expr is None:
raise ValueError("Either ht or freq_expr must be provided.")
if freq_expr is None and (max_af is not None or singleton or "freq" in ht.row):
freq_expr = _resolve_row_annotation_expr(ht, "freq", freq_expr, "freq_expr")
if isinstance(freq_expr, hl.expr.ArrayExpression):
freq_expr = freq_expr[0]
if singleton:
count_var = freq_expr.AC == 1
elif max_af is not None:
count_var = (freq_expr.AC > 0) & (freq_expr.AF <= max_af)
elif freq_expr is not None:
count_var = freq_expr.AC > 0
else:
count_var = True
return hl.int(hl.or_else(count_var, count_missing))
[docs]def variant_observed_and_possible_expr(
freq_expr: hl.ArrayExpression,
max_af: Optional[float] = None,
use_possible_adj: bool = True,
) -> hl.expr.StructExpression:
"""
Return per-variant observed and possible count expressions.
For each element of the frequency array, ``observed_variants`` is 1 when the
variant meets the frequency criteria (AC > 0, and AF <= ``max_af`` if set) and 0
otherwise. ``possible_variants`` uses the same criteria, except that a missing
frequency counts as possible: a variant absent from the dataset could still have
been observed. A variant with a defined frequency that fails the criteria
(AC == 0, or AF > ``max_af``) is therefore 0 in both, so it is excluded from the
denominator as well as the numerator. With ``max_af`` set, this removes common
variants from the possible counts.
When ``freq_expr`` itself is missing (e.g. the variant is not in the dataset),
``observed_variants`` is missing and ``possible_variants`` is 1, or missing when
``use_possible_adj`` is False.
When ``use_possible_adj`` is True (default), the possible count is a scalar
derived from the first (adj) element of the frequency array. When False, it
is an array with one entry per array element, matching the shape of
``observed_variants``.
:param freq_expr: Array of frequency structs with ``AC`` and ``AF`` fields.
:param max_af: Maximum allele frequency threshold. Variants above it are excluded
from both counts. Default is None (no cutoff).
:param use_possible_adj: If True, compute possible count from only the adj
(first) frequency element. If False, compute per-downsampling. Default
is True.
:return: Struct with ``observed_variants`` (array of int) and
``possible_variants`` (int if ``use_possible_adj``, else array of int).
"""
pos_expr = hl.array([freq_expr[0]]) if use_possible_adj else freq_expr
pos_expr = pos_expr.map(
lambda x: variant_observed_expr(freq_expr=x, max_af=max_af, count_missing=True)
)
pos_expr = pos_expr[0] if use_possible_adj else pos_expr
return hl.struct(
observed_variants=freq_expr.map(
lambda x: variant_observed_expr(freq_expr=x, max_af=max_af)
),
possible_variants=pos_expr,
)
[docs]def weighted_sum_agg_expr(
expr: Union[hl.expr.ArrayNumericExpression, hl.expr.NumericExpression],
weight_expr: Union[hl.expr.ArrayNumericExpression, hl.expr.NumericExpression],
) -> Union[hl.expr.ArrayExpression, hl.expr.NumericExpression]:
"""
Return the weighted aggregate sum of ``expr`` weighted by ``weight_expr``.
Both parameters may be scalar or array numeric expressions:
- Both arrays: elements are multiplied pairwise and summed per-element
across rows (``hl.agg.array_sum``).
- Both scalars: standard ``hl.agg.sum(expr * weight_expr)``.
- Mixed: the scalar is broadcast across the array elements and summed
per-element across rows.
:param expr: Numeric expression (scalar or array) to be weighted.
:param weight_expr: Numeric expression (scalar or array) to weight by.
:return: Weighted aggregate sum expression.
"""
expr_is_array = isinstance(expr, hl.expr.ArrayNumericExpression)
weight_is_array = isinstance(weight_expr, hl.expr.ArrayNumericExpression)
if expr_is_array and weight_is_array:
return hl.agg.array_sum(hl.zip(expr, weight_expr).map(lambda x: x[0] * x[1]))
elif not expr_is_array and not weight_is_array:
return hl.agg.sum(expr * weight_expr)
else:
return hl.agg.array_sum(expr * weight_expr)
[docs]def counts_agg_expr(
freq_expr: Optional[
Union[hl.expr.ArrayExpression, hl.expr.StructExpression]
] = None,
ht: Optional[hl.Table] = None,
count_singletons: bool = False,
max_af: Optional[float] = None,
count_missing: bool = False,
) -> hl.expr.StructExpression:
"""
Return an aggregation expression for variant and singleton counts.
Aggregates per-variant counts (via ``variant_observed_expr``) across
rows. Each variant contributes 0 or 1 to the count based on its frequency
metadata (AC, AF). The result is a struct with ``variant_count`` and, when
``count_singletons`` is True, ``singleton_count``.
One of ``freq_expr`` or ``ht`` must be provided. If ``freq_expr`` is not
provided, ``ht.freq`` is used as the fallback.
The shape of ``freq_expr`` controls whether the output counts are scalars or
arrays:
- **StructExpression** (single frequency entry with ``AC`` / ``AF``):
returns scalar ``variant_count`` and ``singleton_count``.
- **ArrayExpression** (array of frequency structs, e.g. one per
downsampling): returns array-valued counts where each element corresponds
to a position in the input array. Internally the array is mapped through
``variant_observed_expr`` and summed element-wise with
``hl.agg.array_sum``.
When ``max_af`` is set, only variants with ``AF <= max_af`` and ``AC > 0``
are counted. Singleton counting (``count_singletons=True``) counts only
variants with ``AC == 1``, independent of the ``max_af`` filter.
:param freq_expr: Frequency expression — an ArrayExpression of structs or a
single StructExpression with ``AC`` and ``AF`` fields. If None, falls
back to ``ht.freq``.
:param ht: Input Table. Used to look up ``freq`` when ``freq_expr`` is None.
:param count_singletons: If True, include a ``singleton_count`` field in
the result. Default is False.
:param max_af: Maximum allele frequency threshold. Variants with
``AF > max_af`` are excluded from ``variant_count``. Does not affect
``singleton_count``. Default is None (no cutoff).
:param count_missing: Whether to count a variant whose frequency is missing.
True counts it as 1, False as 0. Default is False.
:return: Aggregation StructExpression with ``variant_count`` (and
optionally ``singleton_count``). Values are scalars when ``freq_expr`` is
a StructExpression, or arrays when it is an ArrayExpression.
"""
if ht is None and freq_expr is None:
raise ValueError("Either ht or freq_expr must be provided.")
freq_expr = _resolve_row_annotation_expr(ht, "freq", freq_expr, "freq_expr")
is_struct = isinstance(freq_expr, hl.expr.StructExpression)
if is_struct:
freq_expr = hl.array([freq_expr])
params = {"variant_count": {"singleton": False, "max_af": max_af}}
if count_singletons:
params["singleton_count"] = {"singleton": True}
count_expr = {
k: freq_expr.map(
lambda x: variant_observed_expr(
freq_expr=x, **p, count_missing=count_missing
)
)
for k, p in params.items()
}
agg_expr = {k: hl.agg.array_sum(v) for k, v in count_expr.items()}
agg_expr = {k: agg_expr[k][0] for k in agg_expr} if is_struct else agg_expr
return hl.struct(**agg_expr)
[docs]def build_constraint_consequence_groups(
csq_expr: hl.expr.StringExpression,
lof_modifier_expr: hl.expr.StringExpression,
classic_lof_annotations: Tuple[str, ...] = CLASSIC_LOF_ANNOTATIONS,
additional_groupings: Optional[
Dict[str, Dict[str, hl.expr.BooleanExpression]]
] = None,
additional_grouping_combinations: Optional[List[List[str]]] = None,
) -> Tuple[List[hl.expr.BooleanExpression], List[Dict[str, str]]]:
"""
Build constraint consequence groups.
Builds constraint groups based on the consequence expression and LoF
modifier expression. By default, the following groups are built:
- csq_set: synonymous_variant, missense_variant
- lof: classic, hc_lc, hc
The resulting meta and corresponding constraint group filters are:
- {"csq_set": "syn"}: synonymous_variant
- {"csq_set": "mis"}: missense_variant
- {"lof": "classic"}: classic LoF annotations (stop_gained,
splice_donor_variant, splice_acceptor_variant)
- {"lof": "hc_lc"}: LOFTEE HC or LC modifier
- {"lof": "hc"}: LOFTEE HC modifier only
Additional groupings can be added via ``additional_groupings``, and
grouping combinations via ``additional_grouping_combinations``.
.. note::
A combination of more than one grouping builds the **union** of the
combined groups, not their intersection. For example, with an
``additional_groupings`` entry ``{"am": {"am_per_99": ...}}``, the
combination ``["lof", "am"]`` produces groups such as
``{"lof": "hc", "am": "am_per_99"}`` that contain variants that are HC
LoF **or** in the top 1% of AlphaMissense scores. This differs from
frequency metadata (e.g. ``freq_meta``), where multiple keys mean all of
them apply.
:param csq_expr: VEP most severe consequence expression (e.g.,
``ht.most_severe_consequence``).
:param lof_modifier_expr: LOFTEE modifier expression (e.g., ``ht.modifier``).
:param classic_lof_annotations: Classic LoF annotations. Default is
``("stop_gained", "splice_donor_variant", "splice_acceptor_variant")``.
:param additional_groupings: Additional groupings to add to the constraint
groups. Default is None.
:param additional_grouping_combinations: Additional grouping combinations to
add to the constraint groups. Groupings combined in one entry are joined as a
union (see the note above). Default is None.
:return: Tuple of (constraint group filter expressions, meta dicts).
"""
lof_classic_expr = hl.literal(set(classic_lof_annotations)).contains(csq_expr)
lof_hc_expr = lof_modifier_expr == "HC"
lof_hc_lc_expr = lof_hc_expr | (lof_modifier_expr == "LC")
annotation_dict = {
"csq_set": {
"syn": csq_expr == "synonymous_variant",
"mis": csq_expr == "missense_variant",
},
"lof": {
"classic": lof_classic_expr,
"hc_lc": lof_hc_lc_expr,
"hc": lof_hc_expr,
},
}
annotation_dict.update(additional_groupings or {})
grouping_combinations = [["csq_set"], ["lof"]]
grouping_combinations.extend(additional_grouping_combinations or [])
meta = generate_filter_combinations(
grouping_combinations,
{k: list(v.keys()) for k, v in annotation_dict.items()},
)
# Multi-grouping combinations are unions, e.g. LoF or high-scoring missense.
constraint_group_filters = [
functools.reduce(operator.ior, [annotation_dict[k][v] for k, v in m.items()])
for m in meta
]
return constraint_group_filters, meta
[docs]def count_observed_and_possible_by_group(
ht: hl.Table,
possible_expr: hl.expr.Int32Expression,
observed_expr: hl.expr.ArrayExpression,
additional_grouping: Union[List[str], Tuple[str, ...]] = ("methylation_level",),
partition_hint: int = 100,
weight_exprs: Optional[
Union[
List[str],
Dict[str, Union[hl.expr.ArrayExpression, hl.expr.NumericExpression]],
]
] = None,
additional_agg_sum_exprs: Optional[
Union[
List[str],
Dict[str, Union[hl.expr.ArrayExpression, hl.expr.NumericExpression]],
]
] = None,
) -> hl.Table:
"""
Aggregate observed and possible variant counts by substitution context group.
Groups rows by ``context``, ``ref``, ``alt``, and any fields in
``additional_grouping``, then sums observed and possible counts within each
group.
:param ht: Input Table with ``context``, ``ref``, ``alt`` fields and any
fields named in ``additional_grouping``.
:param possible_expr: Per-variant possible count (scalar).
:param observed_expr: Per-variant observed count (array, one element per
array element of the frequency array).
:param additional_grouping: Field names to append to the base
``(context, ref, alt)`` grouping. Default is ``("methylation_level",)``.
:param partition_hint: Target number of partitions for the ``group_by``.
Default is 100.
:param weight_exprs: Weighted sums of ``possible_expr`` to include. Pass
field names (looked up on ``ht``) or a dict mapping output names to
weight expressions. Each produces
``weighted_sum_agg_expr(possible_expr, weight)``.
:param additional_agg_sum_exprs: Extra fields to sum alongside
observed/possible. Pass field names (looked up on ``ht``) or a dict
mapping output names to expressions. Arrays use
``hl.agg.array_sum``; scalars use ``hl.agg.sum``.
:return: Grouped Table with ``observed_variants``, ``possible_variants``,
and any weighted/additional sum fields.
"""
if isinstance(weight_exprs, list):
weight_exprs = {k: ht[k] for k in weight_exprs}
if isinstance(additional_agg_sum_exprs, list):
additional_agg_sum_exprs = {k: ht[k] for k in additional_agg_sum_exprs}
weight_exprs = weight_exprs or {}
additional_agg_sum_exprs = additional_agg_sum_exprs or {}
# Build the grouping struct for the variant count aggregation.
grouping = hl.struct(context=ht.context, ref=ht.ref, alt=ht.alt)
grouping = grouping.annotate(
**{g: ht[g] for g in additional_grouping if g not in grouping}
)
logger.info(
"The following annotations will be used to group the input Table rows when"
" counting variants: %s.",
", ".join(grouping.keys()),
)
agg_expr = {
"observed_variants": hl.agg.array_sum(observed_expr),
"possible_variants": hl.agg.sum(possible_expr),
}
# Update the possible variant count aggregation expression to include weighted sums
# of possible variant counts.
agg_expr.update(
{k: weighted_sum_agg_expr(possible_expr, v) for k, v in weight_exprs.items()}
)
# Get sum aggregation expressions for requested fields.
agg_expr.update(_build_sum_agg_struct(exprs_to_sum=additional_agg_sum_exprs))
# Apply each variant count aggregation in `agg_expr` to get counts for all
# combinations of `grouping`.
ht = ht.group_by(**grouping).partition_hint(partition_hint).aggregate(**agg_expr)
return ht
[docs]def explode_downsamplings_oe(
ht: hl.Table,
downsampling_meta: Dict[str, List[str]],
metrics: List[str] = ["syn", "lof", "mis"],
) -> hl.Table:
"""
Explode observed and expected downsampling counts for each genetic ancestry group and metric.
The input `ht` must contain struct of downsampling information for genetic ancestry
groups under each metric name. For example: 'lof': struct {gen_anc_exp: struct
{global: array<float64>}.
:param ht: Input Table.
:param metrics: List of metrics to explode. Default is '['syn', 'lof', 'mis']'.
:param downsampling_meta: Dictionary containing downsampling metadata. Keys are the
genetic ancestry group names and values are the list of downsamplings for that
genetic ancestry group. Example: {'global': ['5000', '10000'], 'afr': ['5000',
'10000']}.
:return: Table with downsampling counts exploded so that observed and expected
metric counts for each pair of genetic ancestry groups and downsamplings is a
separate row.
"""
ht = ht.select(
_data=[
hl.struct(
gen_anc=gen_anc,
downsampling=downsampling,
**{
f"{metric}.{oe}": ht[metric][f"gen_anc_{oe}"][gen_anc][i]
for oe in ["obs", "exp"]
for metric in metrics
},
)
for gen_anc, downsamplings in downsampling_meta.items()
for i, downsampling in enumerate(downsamplings)
]
)
ht = ht.explode("_data")
ht = ht.transmute(**ht._data)
return ht
[docs]def annotate_mutation_type(
t: Union[hl.MatrixTable, hl.Table],
context_length: Optional[int] = None,
num_scan_context_length: Optional[int] = 100,
) -> Union[hl.MatrixTable, hl.Table]:
"""
Annotate mutation types.
The following annotations are added to the output Table:
- cpg
- transition
- mutation_type - one of "CpG", "non-CpG transition", or "transversion"
- mutation_type_model
..note:
This function uses the term 'mutation_type' because 'variant_type' is already
used in this repo to indicate a variant's multiallelic and SNP/indel status.
:param t: Input Table or MatrixTable.
:param context_length: Length of the 'context' annotation in 't'. If this is not
specified, the value will be determined by examining the first
`num_scan_context_length` values of the 'context' annotation. Default is None.
:param num_scan_context_length: Number of values in the 'context' annotation to use
for determining `context_length` if it is not specified. If set to None, all
values in 'context' will be used. Default is 100.
:return: Table with mutation type annotations added.
"""
if context_length is None:
# Determine the context length by collecting all the context lengths.
if num_scan_context_length is None:
context_lengths = t.aggregate(hl.agg.collect_as_set(hl.len(t.context)))
msg = "all"
else:
context_lengths = hl.len(t.context).take(num_scan_context_length)
msg = f"the first {num_scan_context_length}"
context_lengths = list(filter(None, set(context_lengths)))
if len(context_lengths) > 1:
raise ValueError(
f"More than one length was found among {msg} 'context' values. Length "
"of 'context' should be consistent.",
)
else:
context_length = context_lengths[0]
logger.info(
"Detected a length of %d for context length using %s 'context' values.",
context_length,
msg,
)
# Determine the middle index of the context annotation.
if context_length == 3:
mid_index = 1
elif context_length == 7:
mid_index = 3
else:
raise ValueError(
"The length of context should be either 3 or 7, instead of"
f" {context_length}."
)
transition_expr = hl.is_transition(t.ref, t.alt)
cpg_expr = (
(t.ref == "G") & (t.alt == "A") & (t.context[mid_index - 1 : mid_index] == "C")
) | (
(t.ref == "C")
& (t.alt == "T")
& (t.context[mid_index + 1 : mid_index + 2] == "G")
)
if isinstance(t, hl.MatrixTable):
t = t.annotate_rows(transition=transition_expr, cpg=cpg_expr)
else:
t = t.annotate(transition=transition_expr, cpg=cpg_expr)
mutation_type_expr = (
hl.switch(hl.len(t.context))
.when(
context_length,
hl.case()
.when(t.cpg, "CpG")
.when(t.transition, "non-CpG transition")
.default("transversion"),
)
.or_error("Found 'context' value with unexpected context length!")
)
mutation_type_model_expr = hl.if_else(t.cpg, t.context, "non-CpG")
if isinstance(t, hl.MatrixTable):
return t.annotate_rows(
mutation_type=mutation_type_expr,
mutation_type_model=mutation_type_model_expr,
)
else:
return t.annotate(
mutation_type=mutation_type_expr,
mutation_type_model=mutation_type_model_expr,
)
[docs]def trimer_from_heptamer(
t: Union[hl.MatrixTable, hl.Table],
) -> Union[hl.MatrixTable, hl.Table]:
"""
Trim heptamer context to create trimer context.
:param t: Input MatrixTable or Table with context annotation.
:return: MatrixTable or Table with trimer context annotated.
"""
trimer_expr = hl.if_else(hl.len(t.context) == 7, t.context[2:5], t.context)
return (
t.annotate_rows(context=trimer_expr)
if isinstance(t, hl.MatrixTable)
else t.annotate(context=trimer_expr)
)
[docs]def collapse_strand(
t: Union[hl.Table, hl.MatrixTable],
) -> Union[hl.Table, hl.MatrixTable]:
"""
Return the deduplicated context by collapsing DNA strands.
Function returns the reverse complement for 'ref, 'alt', and 'context' if the
reference allele is either 'G' or 'T'.
The following annotations are added to the output Table:
- was_flipped - whether the 'ref, 'alt', and 'context' were flipped (reverse
complement taken)
:param ht: Input Table.
:return: Table with deduplicated context annotation (ref, alt, context,
was_flipped).
"""
ref_g_or_t_expr = (t.ref == "G") | (t.ref == "T")
collapse_expr = {
"ref": hl.if_else(ref_g_or_t_expr, hl.reverse_complement(t.ref), t.ref),
"alt": hl.if_else(ref_g_or_t_expr, hl.reverse_complement(t.alt), t.alt),
"context": hl.if_else(
ref_g_or_t_expr,
hl.reverse_complement(t.context),
t.context,
),
"was_flipped": ref_g_or_t_expr,
}
return (
t.annotate(**collapse_expr)
if isinstance(t, hl.Table)
else t.annotate_rows(**collapse_expr)
)
[docs]def assemble_constraint_context_ht(
ht: hl.Table,
methylation_ht: hl.Table = None,
gerp_ht: hl.Table = None,
coverage_hts: Dict[str, hl.Table] = None,
an_hts: Dict[str, hl.Table] = None,
freq_hts: Dict[str, hl.Table] = None,
filter_hts: Dict[str, hl.Table] = None,
transformation_funcs: Optional[Dict[str, Callable]] = None,
) -> hl.Table:
"""
Assemble context Table with necessary annotations for constraint calculations.
.. note::
Checks for 'was_split' annotation in Table. If not present, splits
multiallelic sites.
The following annotations are added to the output Table:
- ref: Reference allele.
- alt: Alternate allele.
- context: Trimer context.
- annotations added by `annotate_mutation_type()`, `collapse_strand()`, and
`add_most_severe_csq_to_tc_within_vep_root()`.
Depending on the annotations provided, the following annotations may also be added:
- 'methylation_level': Methylation level annotation will be added if
`methylation_ht` is provided.
- 'gerp': GERP annotation will be added if `gerp_ht` is provided.
- 'coverage': Coverage annotations will be added if `coverage_hts` is provided.
- 'AN': Allele number annotations will be added if `an_hts` is provided.
- 'freq': Frequency annotations will be added if `freq_hts` is provided.
- 'filters': Filter annotations will be added if `filter_hts` is provided.
The `transformation_funcs` parameter can be used to transform the HTs before adding
annotations. The keys should be the annotation names ('coverage', 'gerp', etc.) and
the values should be functions that take the annotation Table and the ht that is
being annotated as input and return the transformed and keyed annotation. If not
provided, the following default transformations are used:
- 'methylation_level': Uses the 'methylation_level' annotation in the
Methylation Table after transforming the methylation level to a 0-2 scale
using `transform_grch37_methylation()` or `transform_grch38_methylation()`.
- 'gerp': Uses the 'S' annotation in the GERP Table. If 'S' is missing, it
defaults to 0.
- 'coverage': If necessary, pulls out the first element of coverage statistics
(which includes all samples). Relevant to v4, where coverage stats include
additional elements to stratify by UK Biobank subset and platforms.
- 'AN': Uses the 'AN' annotation in the allele number Table.
- 'freq': Uses the 'freq' annotation in the frequency Table.
- 'filters': Uses the 'filters' annotation in the filter Table.
The following global annotations may also be added to the output Table:
- 'an_globals': Global allele number annotations 'strata_sample_count' and
'strata_meta' will be added if `an_hts` is provided.
- 'freq_globals': Global frequency annotations 'freq_meta_sample_count' and
'freq_meta' will be added if `freq_hts` is provided.
:param ht: Input context Table with VEP annotation.
:param methylation_ht: Optional Table with methylation annotation. Default is None.
:param gerp_ht: Optional Table with GERP annotation. Default is None.
:param coverage_hts: An optional Dictionary with key as one of 'exomes' or
'genomes' and values as corresponding coverage Tables. Default is None.
:param an_hts: A Dictionary with key as one of 'exomes' or 'genomes' and
values as corresponding allele number Tables. Default is None.
:param transformation_funcs: An optional Dictionary to transform the HTs before
adding annotations. Default is None, which will use some default transformations
as described above.
:return: Table with sites split and necessary annotations.
"""
ref = get_reference_genome(ht.locus)
if ref == hl.get_reference("GRCh37"):
from gnomad.resources.grch37 import transform_grch37_methylation as trans_methyl
else:
from gnomad.resources.grch38 import transform_grch38_methylation as trans_methyl
# Check if context Table is split, and if not, split multiallelic sites.
if "was_split" not in list(ht.row):
ht = hl.split_multi_hts(ht)
# Only need to keep the keys (locus, alleles) and the context and vep annotations.
ht = ht.select("context", "vep")
# Filter Table to only contigs 1-22, X, Y.
ht = hl.filter_intervals(
ht,
[hl.parse_locus_interval(c, ref.name) for c in get_primary_contigs(ref.name)],
)
# Add annotations for 'ref' and 'alt'.
ht = ht.annotate(ref=ht.alleles[0], alt=ht.alleles[1])
# Trim heptamer context to create trimer context.
ht = trimer_from_heptamer(ht)
ht = ht.filter(ht.context.matches(f"[ATCG]{{{3}}}"))
# Annotate mutation type (such as "CpG", "non-CpG transition", "transversion") and
# collapse strands to deduplicate the context.
ht = annotate_mutation_type(collapse_strand(ht))
# Add most_severe_consequence annotation to 'transcript_consequences' within the
# vep root annotation.
csqs = ht.vep.transcript_consequences
csqs = add_most_severe_consequence_to_consequence(csqs)
vep_csq_fields = [
"transcript_id",
"gene_id",
"gene_symbol",
"biotype",
"most_severe_consequence",
"mane_select",
"canonical",
"lof",
"lof_flags",
"sift_score",
"polyphen_score",
"domains",
"uniprot_isoform",
"amino_acids",
"codons",
]
vep_csq_fields = [x for x in vep_csq_fields if x in csqs.dtype.element_type]
ht = ht.annotate(
vep=ht.vep.select(
"most_severe_consequence",
transcript_consequences=csqs.map(lambda x: x.select(*vep_csq_fields)),
)
)
# Add 'methylation_level' and 'gerp' annotations if specified.
transformation_funcs = transformation_funcs or {}
if "methylation_level" not in transformation_funcs:
transformation_funcs["methylation_level"] = lambda x, t: hl.if_else(
ht.cpg, trans_methyl(x)[t.locus].methylation_level, 0
)
if "gerp" not in transformation_funcs:
transformation_funcs["gerp"] = lambda x, t: hl.if_else(
hl.is_missing(x[t.locus].S), 0, x[t.locus].S
)
# If necessary, pull out first element of coverage statistics (which includes all
# samples). Relevant to v4, where coverage stats include additional elements to
# stratify by ukb subset and platforms.
if "coverage" not in transformation_funcs:
transformation_funcs["coverage"] = lambda x, t: (
x[t.locus].coverage_stats[0] if "coverage_stats" in x.row else x[t.locus]
)
if "AN" not in transformation_funcs:
transformation_funcs["AN"] = lambda x, t: x[t.locus].AN
if "freq" not in transformation_funcs:
transformation_funcs["freq"] = lambda x, t: x[t.key].freq
if "filters" not in transformation_funcs:
transformation_funcs["filters"] = lambda x, t: x[t.key].filters
hts = {
"methylation_level": methylation_ht,
"gerp": gerp_ht,
"coverage": coverage_hts,
"AN": an_hts,
"freq": freq_hts,
"filters": filter_hts,
}
hts = {k: v for k, v in hts.items() if v is not None}
exprs = {}
for ann, ann_ht in hts.items():
if isinstance(ann_ht, dict):
ann_expr = hl.struct(
**{k: transformation_funcs[ann](v, ht) for k, v in ann_ht.items()}
)
else:
ann_expr = transformation_funcs[ann](ann_ht, ht)
exprs[ann] = ann_expr
ht = ht.annotate(**exprs)
# Add global annotations for 'AN' and 'freq' if HTs are provided.
global_anns = {
"an_globals": (an_hts, ["strata_sample_count", "strata_meta"]),
"freq_globals": (freq_hts, ["freq_meta_sample_count", "freq_meta"]),
}
global_anns = {k: v for k, v in global_anns.items() if v[0] is not None}
ht = ht.annotate_globals(
**{
k: hl.struct(
**{
data_type: (
ann_ht.select_globals(
*[x for x in g if x in ann_ht.globals]
).index_globals()
)
for data_type, ann_ht in hts.items()
}
)
for k, (hts, g) in global_anns.items()
}
)
return ht
[docs]def calculate_gerp_cutoffs(
ht: hl.Table,
gerp_expr: Optional[hl.expr.Float64Expression] = None,
lower_percentile: float = 0.05,
upper_percentile: float = 0.95,
) -> Tuple[float, float]:
"""
Find GERP cutoffs at the given percentile thresholds.
.. note::
Uses ``hl.agg.approx_quantiles``, so results are approximate.
:param ht: Input Table.
:param gerp_expr: GERP score expression. Default is ``ht.gerp``.
:param lower_percentile: Lower percentile threshold (0-1). Default is 0.05.
:param upper_percentile: Upper percentile threshold (0-1). Default is 0.95.
:return: Tuple of (lower cutoff, upper cutoff) GERP scores.
"""
gerp_expr = _resolve_row_annotation_expr(ht, "gerp", gerp_expr, "gerp_expr")
cutoffs = ht.aggregate(
hl.agg.approx_quantiles(gerp_expr, [lower_percentile, upper_percentile])
)
return cutoffs[0], cutoffs[1]
[docs]def calibration_model_group_expr(
exomes_coverage_expr: hl.expr.Int32Expression,
cpg_expr: hl.expr.BooleanExpression,
low_cov_cutoff: Optional[int] = None,
high_cov_cutoff: int = COVERAGE_CUTOFF,
upper_cov_cutoff: Optional[int] = None,
skip_coverage_model: bool = False,
additional_grouping_exprs: Optional[Dict[str, hl.expr.StringExpression]] = None,
cpg_in_high_only: bool = False,
) -> hl.expr.StructExpression:
"""
Get the calibration model grouping annotation for a variant.
The calibration model expression is a struct with the following fields:
- high_or_low_coverage: Whether the variant belongs to the high or low coverage
calibration model. The variant is assigned to the high coverage model if the
exome coverage is greater than or equal to 'high_cov_cutoff' and less than or
equal to 'upper_cov_cutoff' (if provided). The variant is assigned to the low
coverage model if `skip_coverage_model` is False and the exome coverage is
greater than 'low_cov_cutoff' (if provided) and less than 'high_cov_cutoff'.
- model_group: Struct of the fields the models are grouped by:
- cpg: Whether the variant is a CpG (`cpg_expr`).
- Any fields in `additional_grouping_exprs` (e.g. ``genomic_region``).
The expression is missing for variants that belong to neither model.
:param exomes_coverage_expr: Exome coverage expression.
:param cpg_expr: CpG expression.
:param low_cov_cutoff: Low coverage cutoff. Default is None.
:param high_cov_cutoff: High coverage cutoff. Default is COVERAGE_CUTOFF.
:param upper_cov_cutoff: Upper coverage cutoff. Default is None.
:param skip_coverage_model: Whether to skip the coverage model. Default is False.
:param additional_grouping_exprs: Optional Dictionary of additional expressions to
group by. Default is None.
:param cpg_in_high_only: Whether to only include `cpg_expr` in the model grouping
for the high coverage model. Default is False.
:return: StructExpression with ``high_or_low_coverage`` and ``model_group``
fields.
"""
high_cov_expr = exomes_coverage_expr >= high_cov_cutoff
if upper_cov_cutoff is not None:
high_cov_expr &= exomes_coverage_expr <= upper_cov_cutoff
# Sites above `upper_cov_cutoff` fail the high coverage check, so the upper bound
# on low coverage keeps them out of both models.
low_cov_expr = hl.bool(False) if skip_coverage_model else hl.bool(True)
low_cov_expr &= exomes_coverage_expr < high_cov_cutoff
if low_cov_cutoff is not None:
low_cov_expr &= exomes_coverage_expr > low_cov_cutoff
# Define whether the variant should be included in the high or low coverage model.
model_expr = (
hl.case().when(high_cov_expr, "high").when(low_cov_expr, "low").or_missing()
)
cpg_expr = (
hl.or_missing(model_expr == "high", cpg_expr) if cpg_in_high_only else cpg_expr
)
return hl.or_missing(
hl.is_defined(model_expr),
hl.struct(
high_or_low_coverage=model_expr,
model_group=hl.struct(
cpg=cpg_expr,
**(additional_grouping_exprs or {}),
),
),
)
def _build_sum_agg_struct(
fields_to_sum: Optional[List[str]] = None,
exprs_to_sum: Optional[
Union[hl.expr.StructExpression, Dict[str, hl.expr.NumericExpression]]
] = None,
t: Optional[Union[hl.Table, hl.expr.StructExpression]] = None,
) -> hl.expr.StructExpression:
"""
Return an aggregation expression to sum fields or expressions in a Table/StructExpression.
The aggregation expression is a struct with the sum or array_sum of the fields or
expressions provided in `fields_to_sum` or `exprs_to_sum`.
:param fields_to_sum: List of fields to sum. Must provide ``t`` if using this.
Default is None.
:param exprs_to_sum: Dictionary of expressions to sum. Default is None.
:param t: Optional Table or StructExpression to get `fields_to_sum` from. Default
is None.
:return: Aggregation expression to sum fields or expressions in the Table.
"""
if fields_to_sum is None and exprs_to_sum is None:
raise ValueError("Either 'fields_to_sum' or 'exprs_to_sum' must be provided.")
if fields_to_sum is not None and t is None:
raise ValueError("t must be provided if 'fields_to_sum' is provided.")
exprs_to_sum = exprs_to_sum or {}
exprs_to_sum = hl.struct(**exprs_to_sum, **{f: t[f] for f in fields_to_sum or []})
return hl.struct(
**{
k: (
hl.agg.array_sum(v)
if isinstance(v, hl.ArrayExpression)
else hl.agg.sum(v)
)
for k, v in exprs_to_sum.items()
}
)
[docs]def build_models(
ht: hl.Table,
coverage_expr: hl.expr.Int32Expression,
weighted: bool = False,
keys: Tuple[str, ...] = (
"context",
"ref",
"alt",
"methylation_level",
),
model_group_expr: Optional[hl.expr.StructExpression] = None,
high_cov_definition: int = COVERAGE_CUTOFF,
upper_cov_cutoff: Optional[int] = None,
skip_coverage_model: bool = False,
log10_coverage: bool = True,
additional_grouping: Tuple[str, ...] = (),
) -> Tuple[Optional[hl.expr.ArrayNumericExpression], hl.expr.DictExpression]:
"""
Build coverage and plateau models.
This function builds models (plateau_models) using linear regression to calibrate
mutation rate estimates against the proportion observed of each substitution,
context, and methylation level in `ht`.
Two plateau models are fit, one for CpG transitions, and one for the remainder of
sites (transversions and non CpG transitions).
The plateau models only consider high coverage sites, or sites above a median
coverage of `high_cov_definition` and median coverage below `upper_cov_cutoff`.
Plateau model: adjusts proportion of expected variation based on location in the
genome and CpG status.
The x and y of the plateau models:
- x: `mu_snp` - mutation rate
- y: proportion observed ('observed_variants' / 'possible_variants')
This function also builds models (coverage models) to calibrate the proportion of
expected variation at low coverage sites (sites below `high_cov_definition`).
The coverage models are built by creating a scaling factor across all high coverage
sites, applying this ratio to the low coverage sites, and running a linear
regression.
Coverage model: corrects proportion of expected variation at low coverage sites.
Low coverage sites are defined as sites with median coverage < `high_cov_definition`.
The x and y of the coverage model:
- x: groupings of exome coverage at low coverage sites (log10 transformed if
requested)
- y: sum('observed_variants') / (``high_coverage_scale_factor`` *
sum('possible_variants' * 'mu_snp')) at low coverage sites
``high_coverage_scale_factor`` = sum('observed_variants') /
sum('possible_variants' * 'mu_snp') at high coverage sites
.. note::
This function expects that the input Table (`ht`) contains only high quality
synonymous variants below 0.1% frequency.
The following fields are expected in `ht`:
- context: trinucleotide genomic context.
- ref: the reference allele.
- alt: the alternate allele.
- methylation_level: methylation level.
- mu_snp: mutation rate.
- cpg: whether the variant is a CpG.
- observed_variants: the number of observed variants in the dataset for each
variant. Note that the term "variant" here refers to a specific substitution,
context, methylation level, and coverage combination.
- possible_variants: the number of possible variants in the dataset for each
variant.
:param ht: Input Table.
:param coverage_expr: Expression that defines the coverage metric.
:param weighted: Whether to weight the plateau models (a linear regression
model) by 'possible_variants'. Default is False.
:param keys: Annotations used to group observed and possible variant counts.
Default is ("context", "ref", "alt", "methylation_level").
:param model_group_expr: Expression with ``high_or_low_coverage`` and
``model_group`` annotations to group variants into high or low coverage models.
If ``model_group`` includes a ``genomic_region`` field, the coverage model is
built only on sites where it is "autosome_or_par"; otherwise all sites are
used. If not provided, the ``calibration_model_group_expr`` function is used to
define the grouping, which does not include ``genomic_region``.
:param high_cov_definition: Coverage threshold for high/low classification.
Sites with coverage above this cutoff are considered well covered. Default is
``COVERAGE_CUTOFF``.
:param upper_cov_cutoff: Upper coverage cutoff. Sites with coverage above this
cutoff are excluded from the high coverage Table. Default is None.
:param skip_coverage_model: Whether to skip generating the coverage model. If set
to True, None is returned instead of the coverage model. Default is False.
:param log10_coverage: Whether to convert coverage sites with log10 when building
the coverage model. Default is True.
:param additional_grouping: Additional annotations to group by before
counting the observed and possible variants. Default is ().
:return: Tuple of (coverage model, plateau models). The coverage model is an
``[intercept, slope]`` ArrayNumericExpression, or None when
``skip_coverage_model`` is True.
"""
if model_group_expr is None:
# Define whether the variant should be included in the high or low coverage
# model.
model_group_expr = calibration_model_group_expr(
coverage_expr,
ht.cpg,
low_cov_cutoff=0,
high_cov_cutoff=high_cov_definition,
upper_cov_cutoff=upper_cov_cutoff,
skip_coverage_model=skip_coverage_model,
cpg_in_high_only=True,
)
grouping = keys + additional_grouping
mu_type_fields = ("cpg", "transition", "mutation_type", "mutation_type_model")
# all() accepts a generator expression directly (no list needed) and
# short-circuits on the first False.
has_mu_type = all(x in ht.row for x in mu_type_fields)
grouping += mu_type_fields if has_mu_type else ()
grouping_exprs = {"build_model": model_group_expr}
if not skip_coverage_model:
grouping_exprs["exomes_coverage"] = hl.or_missing(
model_group_expr.high_or_low_coverage == "low", coverage_expr
)
# Collapse per-site rows into one row per (context, ref, alt,
# methylation_level, …) group with summed observed/possible counts.
# model_group_expr and (optionally) exomes_coverage are carried along so
# that plateau and coverage models can be built on the grouped table.
# mu_snp is identical for all rows within a group so grabbing one value suffices.
ht = (
ht.group_by(*grouping, **grouping_exprs)
.aggregate(
mu_snp=hl.agg.take(ht.mu_snp, 1)[0],
**_build_sum_agg_struct(
fields_to_sum=["observed_variants", "possible_variants"], t=ht
),
)
.key_by(*keys)
)
if not has_mu_type:
ht = annotate_mutation_type(ht)
# Build plateau models.
is_high_expr = ht.build_model.high_or_low_coverage == "high"
agg_expr = {
"plateau": hl.agg.filter(
is_high_expr,
build_plateau_models(
ht.mu_snp,
ht.observed_variants,
ht.possible_variants,
model_group_expr=ht.build_model.model_group,
weighted=weighted,
),
)
}
if not skip_coverage_model:
# The coverage model is only built using the full dataset observed variants
# so use the first element if the observed_variants is an array.
obs_is_array = isinstance(ht.observed_variants, hl.expr.ArrayExpression)
obs_expr = ht.observed_variants[0] if obs_is_array else ht.observed_variants
# Restrict the coverage model to autosome/PAR sites when the model grouping
# includes a genomic region, otherwise use all sites.
model_group = ht.build_model.model_group
autosome_or_par_expr = (
model_group.genomic_region == "autosome_or_par"
if "genomic_region" in model_group
else hl.bool(True)
)
# Create a metric that represents the relative mutability of the exome calculated
# on high coverage sites and will be used as scaling factor when building the
# coverage model.
agg_expr["high_coverage_scale_factor"] = hl.agg.filter(
is_high_expr & autosome_or_par_expr,
hl.agg.sum(obs_expr) / hl.agg.sum(ht.possible_variants * ht.mu_snp),
)
# Get the observed variant count and mu_snp for low coverage sites.
agg_expr["coverage"] = hl.agg.filter(
(ht.build_model.high_or_low_coverage == "low") & autosome_or_par_expr,
hl.agg.group_by(
ht.exomes_coverage,
hl.struct(
obs=hl.agg.sum(obs_expr),
mu_snp=hl.agg.sum(ht.possible_variants * ht.mu_snp),
),
),
)
models = ht.aggregate(hl.struct(**agg_expr), _localize=False)
# Build coverage model.
coverage_model = None
if not skip_coverage_model:
coverage_model = models.coverage.map_values(
lambda x: x.annotate(
low_coverage_oe=x.obs / (models.high_coverage_scale_factor * x.mu_snp)
)
)
# TODO: consider weighting here as well.
coverage_model = (
coverage_model.items()
.aggregate(
lambda x: build_coverage_model(
x[1].low_coverage_oe, x[0], log10_coverage=log10_coverage
)
)
.beta
)
return coverage_model, models.plateau
[docs]def build_plateau_models(
mu_snp_expr: hl.expr.Float64Expression,
observed_variants_expr: Union[hl.expr.ArrayExpression, hl.expr.Int64Expression],
possible_variants_expr: Union[hl.expr.ArrayExpression, hl.expr.Int64Expression],
weighted: bool = False,
cpg_expr: Optional[hl.expr.BooleanExpression] = None,
model_group_expr: Optional[hl.expr.StructExpression] = None,
) -> Union[hl.expr.DictExpression, hl.expr.ArrayExpression, hl.expr.StructExpression]:
"""
Build plateau models to calibrate mutation rate against proportion observed.
Fits a linear regression of ``observed_variants_expr / possible_variants_expr``
on ``mu_snp_expr``. When either observed or possible expressions are arrays
(e.g., one model per downsampling), the regression is applied element-wise
via ``hl.agg.array_agg``.
When ``model_group_expr`` or ``cpg_expr`` is provided, the result is a
``DictExpression`` keyed by the grouping struct. Otherwise the result is the
regression beta directly.
:param mu_snp_expr: Mutation rate expression.
:param observed_variants_expr: Observed variant counts (scalar or array).
:param possible_variants_expr: Possible variant counts (scalar or array).
:param weighted: If True, use weighted least squares with
``possible_variants_expr`` as weights. Default is False.
:param cpg_expr: Boolean expression indicating CpG sites. When provided,
adds a ``cpg`` field to the grouping struct.
:param model_group_expr: Struct expression to group by in the aggregation.
:return: Regression betas, optionally grouped by ``model_group_expr``
(and/or ``cpg_expr``).
"""
obs_is_array = isinstance(observed_variants_expr, hl.expr.ArrayExpression)
pos_is_array = isinstance(possible_variants_expr, hl.expr.ArrayExpression)
def _linreg(
o: hl.expr.NumericExpression,
p: hl.expr.NumericExpression,
) -> hl.expr.StructExpression:
"""
Run linear regression of observed/possible on mutation rate.
:param o: Observed variant count expression.
:param p: Possible variant count expression.
:return: Regression beta coefficients.
"""
return hl.agg.linreg(
o / p, [1, mu_snp_expr], weight=p if weighted else None
).beta
if obs_is_array and pos_is_array:
agg_expr = hl.agg.array_agg(
lambda x: _linreg(*x),
hl.zip(observed_variants_expr, possible_variants_expr),
)
elif obs_is_array:
agg_expr = hl.agg.array_agg(
lambda x: _linreg(x, possible_variants_expr), observed_variants_expr
)
elif pos_is_array:
agg_expr = hl.agg.array_agg(
lambda x: _linreg(observed_variants_expr, x), possible_variants_expr
)
else:
agg_expr = _linreg(observed_variants_expr, possible_variants_expr)
if model_group_expr is None and cpg_expr is None:
return agg_expr
model_group_expr = model_group_expr or hl.struct()
if cpg_expr is not None:
model_group_expr = model_group_expr.annotate(cpg=cpg_expr)
return hl.agg.group_by(model_group_expr, agg_expr)
[docs]def build_coverage_model(
low_coverage_oe_expr: hl.expr.Float64Expression,
coverage_expr: hl.expr.Float64Expression,
log10_coverage: bool = False,
) -> hl.expr.StructExpression:
"""
Build coverage model.
This function uses linear regression to build a model of coverage to correct
proportion of expected variation at low coverage sites.
The x and y of the coverage model:
- x: `coverage_expr`
- y: `low_coverage_oe_expr`
:param low_coverage_oe_expr: The Float64Expression of observed:expected ratio
for a given coverage level.
:param coverage_expr: The Float64Expression of the coverage expression.
:param log10_coverage: Whether to convert coverage sites by log10 when building the
coverage model. Default is False.
:return: StructExpression with intercept and slope of the model.
"""
if log10_coverage:
logger.info(
"Converting coverage sites by log10 when building the coverage model."
)
coverage_expr = hl.log10(coverage_expr)
return hl.agg.linreg(low_coverage_oe_expr, [1, coverage_expr])
[docs]def get_all_gen_anc_lengths(
ht: hl.Table,
gen_ancs: Tuple[str, ...],
obs_expr: hl.expr.StructExpression,
) -> List[Tuple[int, str]]:
"""
Get the minimum length of observed variant counts array for each genetic ancestry group downsampling.
The observed variant counts for each genetic ancestry group in `gen_ancs` are specified by
annotations on the `obs_expr` expression.
The function also performs a check that arrays of variant counts within genetic ancestry group
downsamplings all have the same lengths.
:param ht: Input Table containing `obs_expr`.
:param gen_ancs: Genetic ancestry groups used to categorize observed variant counts in downsamplings.
:param obs_expr: Expression for the genetic ancestry group observed variant counts. Should be a
struct containing an array for each gen_anc in `gen_ancs`.
:return: List of ``(length, gen_anc)`` tuples, one per group in ``gen_ancs``,
where ``length`` is the minimum array length for that group.
"""
# TODO: This function will be converted into doing just the length check if there
# is no usage of gen_anc_lengths in the constraint pipeline.
# Get minimum length of downsamplings for each genetic ancestry group.
gen_anc_downsampling_lengths = ht.aggregate(
[hl.agg.min(hl.len(obs_expr[gen_anc])) for gen_anc in gen_ancs]
)
# Zip genetic ancestry group name with their downsampling length.
gen_anc_lengths = list(zip(gen_anc_downsampling_lengths, gen_ancs))
logger.info("Found: %s", "".join(map(str, gen_anc_lengths)))
assert ht.all(
hl.all(
lambda f: f,
[
hl.len(obs_expr[gen_anc]) == length
for length, gen_anc in gen_anc_lengths
],
)
), (
"The arrays of variant counts within genetic ancestry group downsamplings have different"
" lengths!"
)
return gen_anc_lengths
[docs]def get_constraint_grouping_expr(
vep_annotation_expr: hl.expr.StructExpression,
coverage_expr: Optional[hl.expr.Int32Expression] = None,
include_transcript_group: bool = True,
include_canonical_group: bool = True,
include_mane_select_group: bool = False,
) -> Dict[
str,
Union[hl.expr.StringExpression, hl.expr.Int32Expression, hl.expr.BooleanExpression],
]:
"""
Collect annotations used for constraint groupings.
Function collects the following annotations:
- annotation - most_severe_consequence from ``vep_annotation_expr``
- modifier - first non-missing of lof or polyphen_prediction
from ``vep_annotation_expr``, or the literal "None"
- gene - gene_symbol from ``vep_annotation_expr``
- gene_id - gene_id from ``vep_annotation_expr``
- coverage - exome coverage if ``coverage_expr`` is specified
- transcript - transcript_id from ``vep_annotation_expr`` (added when
``include_transcript_group`` is True)
- canonical - from ``vep_annotation_expr`` (added when
``include_canonical_group`` is True)
- mane_select - from ``vep_annotation_expr`` (added when
``include_mane_select_group`` is True)
.. note::
This function expects that the following fields are present in
``vep_annotation_expr``:
- lof
- most_severe_consequence
- gene_symbol
- gene_id
- polyphen_prediction (optional; missing used if absent)
- transcript_id (if ``include_transcript_group`` is True)
- canonical (if ``include_canonical_group`` is True)
- mane_select (if ``include_mane_select_group`` is True)
:param vep_annotation_expr: StructExpression of VEP annotation.
:param coverage_expr: Int32Expression of exome coverage. Default is None.
:param include_transcript_group: Whether to include the transcript annotation in the
groupings. Default is True.
:param include_canonical_group: Whether to include canonical annotation in the
groupings. Default is True.
:param include_mane_select_group: Whether to include mane_select annotation in the
groupings. Default is False.
:return: Dict mapping annotation names to Hail expressions.
"""
lof_expr = vep_annotation_expr.lof
if "polyphen_prediction" in vep_annotation_expr:
polyphen_prediction_expr = vep_annotation_expr.polyphen_prediction
else:
polyphen_prediction_expr = hl.missing(hl.tstr)
# Create constraint annotations to be used for groupings.
groupings = {
"annotation": vep_annotation_expr.most_severe_consequence,
"modifier": hl.coalesce(lof_expr, polyphen_prediction_expr, "None"),
"gene": vep_annotation_expr.gene_symbol,
"gene_id": vep_annotation_expr.gene_id,
}
if coverage_expr is not None:
groupings["coverage"] = coverage_expr
# Add 'transcript' and 'canonical' annotation if requested.
if include_transcript_group:
groupings["transcript"] = vep_annotation_expr.transcript_id
if include_canonical_group:
groupings["canonical"] = hl.or_else(vep_annotation_expr.canonical == 1, False)
if include_mane_select_group:
groupings["mane_select"] = hl.is_defined(vep_annotation_expr.mane_select)
return groupings
[docs]def annotate_exploded_vep_for_constraint_groupings(
ht: hl.Table,
coverage_expr: Optional[hl.expr.Int32Expression] = None,
vep_annotation: str = "transcript_consequences",
include_canonical_group: bool = True,
include_mane_select_group: bool = False,
) -> Tuple[hl.Table, Tuple[str, ...]]:
"""
Explode a VEP annotation and add constraint grouping fields.
Explodes the specified VEP annotation (``vep_annotation``) and adds the
following annotations via ``get_constraint_grouping_expr``:
- annotation - most_severe_consequence from ``vep_annotation``
- modifier - first non-missing of lof or polyphen_prediction
from ``vep_annotation``, or the literal "None"
- gene - gene_symbol from ``vep_annotation``
- gene_id - gene_id from ``vep_annotation``
- coverage - exome coverage if ``coverage_expr`` is specified
- transcript - transcript_id from ``vep_annotation`` (added when
``vep_annotation`` is "transcript_consequences")
- canonical - from ``vep_annotation`` (added when
``include_canonical_group`` is True)
- mane_select - from ``vep_annotation`` (added when
``include_mane_select_group`` is True)
.. note::
This function expects that a ``vep`` annotation is present in ``ht``.
.. note::
Only the grouping fields are kept from the VEP annotation. The returned
Table no longer has the ``vep`` annotation or the exploded
``vep_annotation`` struct, which keeps per-variant tables built from it
small. All other fields of ``ht`` are kept.
:param ht: Input Table.
:param coverage_expr: Expression that defines the coverage metric. Default
is None.
:param vep_annotation: Name of annotation in vep (one of
``"transcript_consequences"`` and ``"worst_csq_by_gene"``) that will be
used for obtaining constraint annotations. Default is
``"transcript_consequences"``.
:param include_canonical_group: Whether to include canonical annotation
in the groupings. Default is True. Ignored unless ``vep_annotation`` is
``"transcript_consequences"``.
:param include_mane_select_group: Whether to include mane_select
annotation in the groupings. Default is False. Ignored unless
``vep_annotation`` is ``"transcript_consequences"``.
:return: Tuple of (Table with one row per element of ``vep_annotation`` and the
grouping fields added in place of the VEP annotation, names of the added
grouping fields).
"""
# Annotate ht with coverage_expr set as a temporary annotation '_coverage_metric'
# before modifying ht.
if coverage_expr is not None:
ht = ht.annotate(_coverage_metric=coverage_expr)
if vep_annotation == "transcript_consequences":
if not include_canonical_group and not include_mane_select_group:
raise ValueError(
"If 'vep_annotation' is 'transcript_consequences', one of either"
" 'include_canonical_group' or 'include_mane_select_group' must be set!"
)
include_transcript_group = True
else:
logger.warning(
"Setting both 'include_canonical_group' and 'include_mane_select_group' to"
" False (options cannot be used unless 'vep_annotation' is"
" 'transcript_consequences')."
)
include_transcript_group = False
include_canonical_group = False
include_mane_select_group = False
# Annotate 'worst_csq_by_gene' to `ht` if it's specified for `vep_annotation`.
if vep_annotation == "worst_csq_by_gene":
ht = process_consequences(ht)
# Explode the specified VEP annotation.
ht = explode_by_vep_annotation(ht, vep_annotation)
# Collect the annotations used for groupings.
groupings = get_constraint_grouping_expr(
ht[vep_annotation],
coverage_expr=None if coverage_expr is None else ht._coverage_metric,
include_transcript_group=include_transcript_group,
include_canonical_group=include_canonical_group,
include_mane_select_group=include_mane_select_group,
)
return ht.transmute(**groupings), tuple(groupings.keys())
[docs]def apply_plateau_models(
mu_expr: hl.Float64Expression,
plateau_models_expr: hl.ArrayExpression,
) -> Union[hl.ArrayExpression, hl.Float64Expression]:
"""
Compute the predicted probability observed.
The predicted probability observed is computed as the mutation rate adjusted by the
plateau model.
:param mu_expr: Mutation rate expression.
:param plateau_models_expr: This can be either a single plateau model, where the
first element is the intercept and the second element is the slope, or an array
of plateau models.
:return: Float64Expression of the predicted probability observed, or an
ArrayExpression of them when ``plateau_models_expr`` is an array of models.
"""
def _apply_model(plateau_model: hl.ArrayExpression) -> hl.Float64Expression:
"""
Apply the plateau model to the mutation rate expression.
:param plateau_model: ArrayExpression of the plateau model.
:return: Float64Expression of the predicted probability observed.
"""
slope = plateau_model[1]
intercept = plateau_model[0]
ppo_expr = mu_expr * slope + intercept
return ppo_expr
if plateau_models_expr.dtype.element_type == hl.tarray(hl.tfloat64):
return plateau_models_expr.map(lambda x: _apply_model(x))
return _apply_model(plateau_models_expr)
[docs]def coverage_correction_expr(
coverage_expr: hl.Float64Expression,
coverage_model: Union[Tuple[float, float], hl.expr.ArrayNumericExpression],
low_coverage_expr: Optional[hl.BooleanExpression] = None,
coverage_cutoff: Optional[int] = None,
log10_coverage: bool = False,
) -> hl.Float64Expression:
"""
Compute the coverage correction expression.
.. note::
One and only one of `low_coverage_expr` or `coverage_cutoff` must be specified.
The coverage correction expression is computed as follows:
- If the coverage is 0, the coverage correction is 0.
- If the low coverage expression (`low_coverage_expr`) is True, or the
coverage (`coverage_expr`) is below the coverage cutoff (`coverage_cutoff`),
the coverage model is applied to the coverage.
- Otherwise, the coverage correction is 1.
:param coverage_expr: Float64Expression of the coverage.
:param coverage_model: Coverage model as ``[intercept, slope]``, either a tuple of
floats or an ArrayNumericExpression (as returned by ``build_models``).
:param low_coverage_expr: Optional BooleanExpression indicating whether the site is
a low coverage site, and the coverage model should be applied. Default is None.
:param coverage_cutoff: Optional coverage cutoff. If specified, the coverage model
is applied to sites with coverage below this cutoff. Default is None.
:param log10_coverage: Whether to convert coverage sites by log10 when applying the
coverage model. Default is False.
:return: Float64Expression of the coverage correction.
"""
if low_coverage_expr is None and coverage_cutoff is None:
raise ValueError(
"Either 'low_coverage_expr' or 'coverage_cutoff' must be specified!"
)
if low_coverage_expr is not None and coverage_cutoff is not None:
raise ValueError(
"Only one of 'low_coverage_expr' or 'coverage_cutoff' can be specified!"
)
if coverage_cutoff is not None:
low_coverage_expr = coverage_expr < coverage_cutoff
if log10_coverage:
cov_corr_expr = hl.log10(coverage_expr)
else:
cov_corr_expr = coverage_expr
return (
hl.case()
.when(coverage_expr == 0, 0)
.when(
low_coverage_expr,
coverage_model[1] * cov_corr_expr + coverage_model[0],
)
.default(1)
)
[docs]def apply_models(
mu_expr: hl.expr.Float64Expression,
plateau_models_expr: hl.expr.ArrayExpression,
possible_variants_expr: hl.expr.Int64Expression,
coverage_model: Optional[
Union[Tuple[float, float], hl.expr.ArrayNumericExpression]
] = None,
coverage_expr: Optional[hl.expr.Int32Expression] = None,
cpg_expr: Optional[hl.expr.BooleanExpression] = None,
model_group_expr: Optional[hl.expr.StructExpression] = None,
high_cov_definition: int = COVERAGE_CUTOFF,
log10_coverage: bool = True,
) -> hl.expr.StructExpression:
"""
Apply calibration models to compute expected variant counts.
Applies plateau and (optionally) coverage models to produce a struct with
``mu``, ``predicted_proportion_observed``, ``expected_variants``, and
(when a coverage model is provided) ``coverage_correction``.
.. note::
When a coverage model is applied, a row whose ``model_group_expr`` is
missing gets missing ``mu``, ``expected_variants``, and
``coverage_correction`` (unless its coverage is 0, which always gives a
correction of 0). ``calibration_model_group_expr`` is missing for sites in
neither model, such as sites above ``upper_cov_cutoff`` or at or below
``low_cov_cutoff``. ``hl.agg.sum`` skips missing values, so filter these rows
out before applying the models if they should not silently drop out of
aggregated expected counts.
:param mu_expr: Mutation rate expression.
:param plateau_models_expr: Single plateau model (array of [intercept,
slope]) or an array of plateau models.
:param possible_variants_expr: Possible variant counts to multiply the
predicted proportion observed by.
:param coverage_model: Coverage model as ``[intercept, slope]``, either a tuple of
floats or an ArrayNumericExpression (as returned by ``build_models``).
Default is None.
:param coverage_expr: Int32Expression of the coverage. Required when
``coverage_model`` is provided or ``model_group_expr`` is None.
:param cpg_expr: BooleanExpression indicating CpG sites. Required when
``model_group_expr`` is None.
:param model_group_expr: Expression with ``high_or_low_coverage``
annotation to group variants into high or low coverage models. If not
provided, ``calibration_model_group_expr`` is used. Default is None.
:param high_cov_definition: Coverage threshold for high/low classification.
Default is ``COVERAGE_CUTOFF``.
:param log10_coverage: Whether to log10-transform coverage when applying
the coverage model. Default is True.
:return: StructExpression with "mu", "predicted_proportion_observed",
"expected_variants", and optionally "coverage_correction".
"""
if coverage_model is not None and coverage_expr is None:
raise ValueError(
"If 'coverage_model' is specified, 'coverage_expr' must also be specified!"
)
if model_group_expr is None and (coverage_expr is None or cpg_expr is None):
raise ValueError(
"If 'model_group_expr' is not specified, 'coverage_expr' and 'cpg_expr' must"
" be specified!"
)
if model_group_expr is None:
# Get the annotations relevant for applying the calibration models.
model_group_expr = calibration_model_group_expr(
coverage_expr,
cpg_expr,
high_cov_cutoff=high_cov_definition,
skip_coverage_model=coverage_model is None,
)
# Apply plateau models.
ppo_expr = apply_plateau_models(mu_expr, plateau_models_expr)
apply_expr = hl.struct(
mu=mu_expr * possible_variants_expr,
predicted_proportion_observed=ppo_expr,
expected_variants=ppo_expr * possible_variants_expr,
)
# Get the coverage correction expression if a coverage model is provided.
if coverage_model is not None:
cov_corr_expr = coverage_correction_expr(
coverage_expr,
coverage_model,
low_coverage_expr=model_group_expr.high_or_low_coverage == "low",
log10_coverage=log10_coverage,
)
apply_expr = apply_expr.annotate(
mu=apply_expr.mu * cov_corr_expr,
expected_variants=apply_expr.expected_variants * cov_corr_expr,
coverage_correction=cov_corr_expr,
)
return apply_expr
[docs]def aggregate_constraint_metrics_expr(
t: Union[hl.Table, hl.StructExpression],
fields_to_sum: Union[List[str], Tuple[str, ...]] = DEFAULT_FIELDS_TO_SUM,
additional_exprs_to_sum: Optional[Dict[str, hl.expr.Expression]] = None,
) -> hl.expr.StructExpression:
"""
Get an aggregation expression for the sum of specified fields and expressions.
An aggregate sum or array sum is created for each field in ``fields_to_sum``.
:param t: Input Table or StructExpression.
:param fields_to_sum: Fields in ``t`` to get an aggregate sum expression for.
Default is ``DEFAULT_FIELDS_TO_SUM``.
:param additional_exprs_to_sum: Dictionary of additional expressions to get
an aggregate sum expression for. Field names are the keys and expressions are
the values. Default is None.
:return: StructExpression with the sum of each field and expression.
"""
return _build_sum_agg_struct(
fields_to_sum=list(fields_to_sum),
exprs_to_sum=additional_exprs_to_sum,
t=t,
)
[docs]def compute_pli(
ht: hl.Table,
obs_expr: hl.expr.Int64Expression,
exp_expr: hl.expr.Float64Expression,
expected_values: Optional[Dict[str, float]] = None,
min_diff_convergence: float = 0.001,
) -> hl.StructExpression:
"""
Compute the pLI score using the observed and expected variant counts.
Full details on pLI can be found in the ExAC paper: Lek, M., Karczewski, K.,
Minikel, E. et al. Analysis of protein-coding genetic variation in 60,706 humans.
Nature 536, 285–291 (2016).
pLI is the probability of being loss-of-function intolerant, and this function
computes that probability using the expectation-maximization (EM) algorithm.
We assume a 3 state model, where each gene fits into one of three categories
with respect loss-of-function variation sensitivity:
- Null: where protein truncating variation is completely tolerated by natural
selection.
- Recessive (Rec): where heterozygous pLoFs are tolerated but homozygous pLoFs
are not.
- Haploinsufficient (LI): where heterozygous pLoFs are not tolerated.
The function requires the expected amount of loss-of-function depletion for each of
these states. The default provided is based on the observed depletion of
protein-truncating variation in the Blekhman autosomal recessive and ClinGen
dosage sensitivity gene sets (Supplementary Information Table 12 of the above
reference):
- Null: 1.0, assume tolerant genes have the expected amount of truncating
variation.
- Rec: 0.463, derived from the empirical mean observed/expected rate of
truncating variation for recessive disease genes (0.463).
- LI: 0.089, derived from the empirical mean observed/expected rate of
truncating variation for severe haploinsufficient genes.
The output StructExpression will include the following annotations:
- pLI: Probability of loss-of-function intolerance; probability that transcript
falls into distribution of haploinsufficient genes.
- pNull: Probability that transcript falls into distribution of unconstrained
genes.
- pRec: Probability that transcript falls into distribution of recessive genes.
:param ht: Input Table containing `obs_expr` and `exp_expr`.
:param obs_expr: Expression for the number of observed variants on each gene or
transcript in `ht`.
:param exp_expr: Expression for the number of expected variants on each gene or
transcript in `ht`.
:param expected_values: Dictionary containing the expected values for 'Null',
'Rec', and 'LI' to use as starting values.
:param min_diff_convergence: Minimum iteration change in LI to consider the EM
model convergence criteria as met. Default is 0.001.
:return: StructExpression for pLI scores.
"""
if expected_values is None:
expected_values = {"Null": 1.0, "Rec": 0.463, "LI": 0.089}
# Set up initial values.
last_pi = {k: 0 for k in expected_values.keys()}
pi = {k: 1 / len(expected_values.keys()) for k in expected_values.keys()}
dpois_expr = {
k: hl.or_missing(
exp_expr > 0, hl.dpois(obs_expr, exp_expr * expected_values[k])
)
for k in pi
}
_ht = ht.select(dpois=dpois_expr)
# Checkpoint the temp HT because it will need to be aggregated several times.
_ht = _ht.checkpoint(new_temp_file(prefix="compute_pli", extension="ht"))
# Calculate pLI scores.
while abs(pi["LI"] - last_pi["LI"]) > min_diff_convergence:
last_pi = copy.deepcopy(pi)
pi_expr = {k: v * _ht.dpois[k] for k, v in pi.items()}
row_sum_expr = hl.sum([pi_expr[k] for k in pi])
pi_expr = {k: pi_expr[k] / row_sum_expr for k, v in pi.items()}
pi = _ht.aggregate({k: hl.agg.mean(pi_expr[k]) for k in pi.keys()})
# Get expression for pLI scores.
pli_expr = {k: v * dpois_expr[k] for k, v in pi.items()}
row_sum_expr = hl.sum([pli_expr[k] for k in pi])
return hl.struct(**{f"p{k}": pli_expr[k] / row_sum_expr for k in pi.keys()})
def _oe_ci_discretized_poisson(
obs_expr: hl.expr.Int64Expression,
exp_expr: hl.expr.Float64Expression,
alpha: float = 0.05,
) -> hl.expr.StructExpression:
"""
Compute OE confidence interval via discretized Poisson CDF.
Sweeps the OE ratio parameter over [0, 2) in steps of 0.001, evaluates
``dpois(obs, exp * x)`` at each point, normalises the cumulative sum, and
reads off the bounds at ``alpha`` and ``1 - alpha``.
:param obs_expr: Observed variant count expression.
:param exp_expr: Expected variant count expression.
:param alpha: Significance level. Default is 0.05.
:return: Struct with ``lower`` and ``upper`` bounds.
"""
# Set up range between 0 and 2.
range_expr = hl.range(0, 2000).map(lambda x: hl.float64(x) / 1000)
range_dpois_expr = range_expr.map(lambda x: hl.dpois(obs_expr, exp_expr * x))
# Compute cumulative density function of the Poisson distribution density.
cumulative_dpois_expr = hl.cumulative_sum(range_dpois_expr)
max_cumulative_dpois_expr = cumulative_dpois_expr[-1]
norm_dpois_expr = cumulative_dpois_expr.map(lambda x: x / max_cumulative_dpois_expr)
# Extract the value of the varying parameter within specified range.
lower_idx_expr = hl.argmax(
norm_dpois_expr.map(lambda x: hl.or_missing(x < alpha, x))
)
upper_idx_expr = hl.argmin(
norm_dpois_expr.map(lambda x: hl.or_missing(x > 1 - alpha, x))
)
return hl.struct(
lower=hl.if_else(obs_expr > 0, range_expr[lower_idx_expr], 0),
upper=range_expr[upper_idx_expr],
)
def _oe_ci_gamma(
obs_expr: hl.expr.Int64Expression,
exp_expr: hl.expr.Float64Expression,
alpha: float = 0.05,
) -> hl.expr.StructExpression:
"""
Compute OE confidence interval using the Gamma distribution.
Uses Hail's ``hl.qgamma`` quantile function.
:param obs_expr: Observed variant count expression.
:param exp_expr: Expected variant count expression.
:param alpha: Significance level. Default is 0.05.
:return: Struct with ``lower`` and ``upper`` bounds.
"""
shape = obs_expr + hl.literal(1.0)
scale = divide_null(hl.literal(1.0), exp_expr)
return hl.struct(
lower=hl.qgamma(hl.literal(alpha), shape, scale),
upper=hl.qgamma(hl.literal(1.0 - alpha), shape, scale),
)
[docs]def oe_confidence_interval(
obs_expr: hl.expr.Int64Expression,
exp_expr: hl.expr.Float64Expression,
alpha: float = 0.05,
method: str = "gamma",
) -> hl.expr.StructExpression:
"""
Compute a confidence interval around the observed/expected ratio.
Two methods are available:
- ``"gamma"`` (default): uses ``hl.qgamma`` to compute exact quantiles of
the Gamma posterior. Fast and precise.
- ``"poisson"``: sweeps a discretized Poisson likelihood over the OE
parameter space [0, 2). Retained for backwards compatibility.
.. note::
The default changed from the discretized Poisson computation to
``"gamma"``, so calls that do not pass ``method`` return different bounds
than before. ``"gamma"`` requires Hail >= 0.2.137, which added
``hl.qgamma``; on older Hail versions pass ``method="poisson"``.
:param obs_expr: Observed variant count expression.
:param exp_expr: Expected variant count expression.
:param alpha: Significance level for the confidence interval. Default is
0.05 (90% CI).
:param method: CI method — ``"gamma"`` or ``"poisson"``. Default is
``"gamma"``.
:return: Struct with ``lower`` and ``upper`` bounds.
"""
if method == "gamma":
if not hasattr(hl, "qgamma"):
raise RuntimeError(
"method='gamma' requires hl.qgamma, available in Hail >= 0.2.137. "
"Use method='poisson' or upgrade Hail."
)
return _oe_ci_gamma(obs_expr, exp_expr, alpha)
elif method == "poisson":
return _oe_ci_discretized_poisson(obs_expr, exp_expr, alpha)
else:
raise ValueError(f"Unknown CI method: {method!r}. Use 'gamma' or 'poisson'.")
[docs]def calculate_raw_z_score(
obs_expr: hl.expr.Int64Expression,
exp_expr: hl.expr.Float64Expression,
) -> hl.expr.Float64Expression:
"""
Compute the signed raw z-score using observed and expected variant counts.
The raw z-scores are positive when the transcript had fewer variants than expected,
and are negative when transcripts had more variants than expected.
:param obs_expr: Observed variant count expression.
:param exp_expr: Expected variant count expression.
:return: Raw z-score expression.
"""
chisq_expr = divide_null((obs_expr - exp_expr) ** 2, exp_expr)
return hl.sqrt(chisq_expr) * hl.if_else(obs_expr > exp_expr, -1, 1)
[docs]def get_constraint_flags(
exp_expr: hl.expr.Float64Expression,
raw_z_expr: hl.expr.Float64Expression,
raw_z_lower_threshold: Optional[Union[float, hl.expr.Float64Expression]] = -5.0,
raw_z_upper_threshold: Optional[Union[float, hl.expr.Float64Expression]] = 5.0,
flag_postfix: str = "",
) -> Dict[str, hl.expr.Expression]:
"""
Determine the constraint flags that define why constraint will not be calculated.
Flags which are added:
- "no_exp_{flag_postfix}" - for genes that have missing or zero expected variants.
- "outlier_{flag_postfix}" - for genes that are raw z-score outliers:
(`raw_z_expr` < `raw_z_lower_threshold`) or (`raw_z_expr` >
`raw_z_upper_threshold`).
:param exp_expr: Expression for the expected variant counts of pLoF, missense, or
synonymous variants.
:param raw_z_expr: Expression for the signed raw z-score of pLoF, missense, or
synonymous variants.
:param raw_z_lower_threshold: Lower threshold for the raw z-score. When `raw_z_expr`
is less than this threshold it is considered an 'outlier'. Default is -5.0.
:param raw_z_upper_threshold: Upper threshold for the raw z-score. When `raw_z_expr`
is greater than this threshold it is considered an 'outlier'. Default is 5.0.
:param flag_postfix: Postfix to add to the end of the constraint flag names.
:return: Dictionary containing expressions for constraint flags.
"""
outlier_expr = False
if raw_z_lower_threshold is not None:
outlier_expr |= hl.or_else(raw_z_expr < raw_z_lower_threshold, False)
if raw_z_upper_threshold is not None:
outlier_expr |= hl.or_else(raw_z_expr > raw_z_upper_threshold, False)
if flag_postfix:
flag_postfix = f"_{flag_postfix}"
constraint_flags = {
f"no_exp{flag_postfix}": hl.or_else(exp_expr <= 0, True),
f"outlier{flag_postfix}": hl.or_else(outlier_expr, False),
}
return constraint_flags
[docs]def calculate_raw_z_score_sd(
raw_z_expr: hl.expr.Float64Expression,
flag_expr: hl.expr.StringExpression,
mirror_neg_raw_z: bool = True,
) -> hl.expr.Expression:
"""
Calculate the standard deviation of the raw z-score.
When using `mirror_neg_raw_z` is True, all the negative raw z-scores (defined by
`raw_z_expr`) are combined with those same z-scores multiplied by -1 (to create a
mirrored distribution).
:param raw_z_expr: Expression for the raw z-score.
:param flag_expr: Expression for the constraint flags. z-score will not be
calculated if flags are present.
:param mirror_neg_raw_z: Whether the standard deviation should be computed using a
mirrored distribution of negative `raw_z_expr`.
:return: StructExpression containing standard deviation of the raw z-score and
the z-score.
"""
filter_expr = (hl.len(flag_expr) == 0) & hl.is_defined(raw_z_expr)
if mirror_neg_raw_z:
filter_expr &= raw_z_expr < 0
sd_expr = hl.agg.explode(
lambda x: hl.agg.stats(x), [raw_z_expr, -raw_z_expr]
).stdev
else:
sd_expr = hl.agg.stats(raw_z_expr).stdev
return hl.agg.filter(filter_expr, sd_expr)
[docs]def add_gencode_transcript_annotations(
ht: hl.Table,
gencode_ht: hl.Table,
annotations: Union[List[str], Tuple[str, ...]] = DEFAULT_GENCODE_ANNOTATIONS,
remove_y_par: bool = True,
) -> hl.Table:
"""
Add GENCODE annotations to Table based on transcript id.
In addition to the annotations specified by ``annotations``, the following
computed annotations are always added:
- chromosome
- cds_length
- num_coding_exons
:param ht: Input Table.
:param gencode_ht: Table with GENCODE annotations.
:param annotations: GENCODE annotations to add. Default is
``DEFAULT_GENCODE_ANNOTATIONS``.
:param remove_y_par: Whether to remove features for the Y chromosome PAR regions.
Default is True because the Y chromosome PAR regions are typically not included
in the constraint calculations and both chrX and chrY will have the same
'transcript_id' field for these regions. This parameter can only be True if
``gencode_ht`` includes a 'transcript_id_version' field because Y_PAR (or
PAR_Y, depending on the GENCODE release) is included in the version of the
transcript, which has been stripped from the 'transcript_id' field.
:return: Table with transcript annotations from GENCODE added.
"""
annotations = list(annotations)
if remove_y_par and "transcript_id_version" not in gencode_ht.row:
raise ValueError(
"remove_y_par is True but 'transcript_id_version' is not in gencode_ht"
)
if remove_y_par:
gencode_ht = gencode_ht.filter(
~gencode_ht.transcript_id_version.endswith("Y_PAR")
& ~gencode_ht.transcript_id_version.endswith("PAR_Y")
)
gencode_ht = gencode_ht.annotate(
length=gencode_ht.interval.end.position
- gencode_ht.interval.start.position
+ 1,
chromosome=gencode_ht.interval.start.contig,
start_position=gencode_ht.interval.start.position,
end_position=gencode_ht.interval.end.position,
)
# Add transcript annotations to input Table.
annotations_to_add = set(annotations + ["chromosome", "transcript_id"])
gencode_transcript_ht = (
gencode_ht.filter(gencode_ht.feature == "transcript")
.select(*annotations_to_add)
.key_by("transcript_id")
.drop("interval")
)
# Obtain CDS annotations from GENCODE file and calculate CDS length and
# number of exons.
annotations_to_add = set(annotations + ["chromosome", "transcript_id", "length"])
gencode_cds_ht = (
gencode_ht.filter(gencode_ht.feature == "CDS")
.select(*annotations_to_add)
.key_by("transcript_id")
.drop("interval")
)
annotations_to_add.remove("length")
gencode_cds_ht = gencode_cds_ht.group_by("transcript_id").aggregate(
cds_length=hl.agg.sum(gencode_cds_ht.length),
num_coding_exons=hl.agg.count(),
)
# Join the transcript and CDS annotations.
gencode_transcript_ht = gencode_transcript_ht.join(gencode_cds_ht).checkpoint(
new_temp_file(prefix="gencode", extension="ht")
)
# Add GENCODE annotations to input Table.
ht = ht.annotate(**gencode_transcript_ht[ht.transcript])
return ht
[docs]def rank_and_assign_bins(
value_expr: hl.expr.Float64Expression,
bin_granularities: Optional[Dict[str, int]] = None,
prefix: str = "",
) -> hl.StructExpression:
"""
Rank rows by a numeric expression and assign bin labels.
**Rank-based binning**: every row receives a unique position in the sorted
order, and bins are derived from that position. This differs from
threshold-based binning (see :func:`annotate_bins_by_threshold`), where
bins are assigned by comparing values against pre-computed boundary values.
Rows are ordered ascending by ``value_expr``. Each row is assigned a
0-based ``{prefix}rank`` and a ``{prefix}bin_{name}`` field for every
entry in ``bin_granularities``, computed as
``hl.int(rank * multiplier / n_rows)``.
Rows where ``value_expr`` is missing are excluded from the ranking and
receive missing rank and bin annotations. ``n_rows`` is therefore the
number of rows with a defined ``value_expr``, so bins are evenly sized
and the full range of bins is populated.
Used by :func:`rank_array_element_metrics` to rank metrics within array
elements.
.. note::
This function evaluates its input when it is called. It needs the number
of rows with a defined ``value_expr`` for the bin denominator, so it counts
them, and it caches the filtered and the ranked rows so the filter and the
sort are not recomputed. In Hail, ``Table.cache()`` ignores the storage
level and writes a checkpoint to a ``persist_Table*`` path under the
temporary directory set by ``hl.init(tmp_dir=...)``, so with a GCS
``tmp_dir`` these Tables are written to GCS. They hold only ``value_expr``
and the key, and are not removed when the function returns.
:param value_expr: Numeric expression to rank by (ascending).
:param bin_granularities: Mapping of bin name to multiplier. Each entry
produces a ``{prefix}bin_{name}`` field. Default is
``{"percentile": 100, "decile": 10, "sextile": 6}``.
:param prefix: String prepended to ``rank`` and ``bin_{name}`` field
names. Default is ``""`` (no prefix).
:return: Struct with ``{prefix}rank`` and ``{prefix}bin_{name}`` fields
for each entry in ``bin_granularities``.
"""
if bin_granularities is None:
bin_granularities = {"percentile": 100, "decile": 10, "sextile": 6}
ht = value_expr._indices.source
source_key = list(ht.key)
ranked_ht = ht.select(_=value_expr)
# Rank only rows with a defined value. `order_by` sorts missing values after
# non-missing ones, so retaining them would assign them the highest ranks and
# inflate the denominator, shifting every row into a lower bin and leaving the
# top bins empty. Rows dropped here get missing rank/bin fields from the join.
# Cache before counting so the filter is not recomputed when ranking below.
ranked_ht = ranked_ht.filter(hl.is_defined(ranked_ht._)).cache()
n_rows = ranked_ht.count()
ranked_ht = ranked_ht.order_by("_").add_index("rank")
ranked_ht = ranked_ht.select(
*source_key,
**{f"{prefix}rank": ranked_ht.rank},
**{
f"{prefix}bin_{name}": hl.int(ranked_ht.rank * multiplier / n_rows)
for name, multiplier in bin_granularities.items()
},
).cache()
return ranked_ht.key_by(*source_key)[ht.key]
[docs]def compute_percentile_thresholds(
ht: hl.Table,
metric_expr: hl.expr.Float64Expression,
outlier_expr: Optional[hl.expr.BooleanExpression] = None,
transcript_filter_expr: Optional[hl.expr.BooleanExpression] = None,
percentiles: Tuple[float, ...] = (1, 5, 10, 15, 25, 50, 75),
quantile_k: int = 1000,
) -> Dict[float, float]:
"""
Compute approximate percentile thresholds for a metric expression.
**Threshold-based binning, step 1**: computes the boundary values that
define bin edges. Its values, keyed by ``(granularity, metric)`` instead of
percentile, form the ``thresholds`` argument of
:func:`annotate_bins_by_threshold` (step 2), which assigns each row to a bin;
see that function for an example.
This two-step approach differs from rank-based binning (see
:func:`rank_and_assign_bins`), where every row receives a unique position
in the sorted order. Threshold-based binning allows thresholds to be
computed on a filtered subset (e.g., representative transcripts) and then
applied to all rows, so multiple rows can share the same bin.
Optionally filters to a subset of rows and excludes outliers, then
computes approximate quantile thresholds at the requested percentiles in
a single aggregation pass.
.. note::
Uses ``hl.agg.approx_quantiles``, so results are approximate. Increase
``quantile_k`` for higher accuracy.
:param ht: Input Table.
:param metric_expr: Float expression to compute thresholds for. Must be
defined on ``ht``.
:param outlier_expr: Optional boolean expression that is ``True`` for rows
to exclude. When None (default), no outlier filtering is applied.
:param transcript_filter_expr: Optional boolean expression that is ``True``
for rows to include. When None (default), all rows are included.
:param percentiles: Percentile values (0-100) at which to compute
thresholds. Default is (1, 5, 10, 15, 25, 50, 75).
:param quantile_k: Accuracy parameter for
:func:`hail.expr.aggregators.approx_quantiles`. Default is 1000.
:return: Dict mapping each percentile to its threshold value.
"""
qs = [p / 100.0 for p in percentiles]
filter_expr = hl.is_defined(metric_expr)
if outlier_expr is not None:
filter_expr = filter_expr & ~outlier_expr
if transcript_filter_expr is not None:
filter_expr = filter_expr & transcript_filter_expr
result = ht.aggregate(
hl.struct(
thresholds=hl.agg.filter(
filter_expr, hl.agg.approx_quantiles(metric_expr, qs, k=quantile_k)
),
n=hl.agg.count_where(filter_expr),
)
)
logger.info(
"Computed percentile thresholds on %d transcripts.",
result.n,
)
return dict(zip(percentiles, result.thresholds))
[docs]def annotate_bins_by_threshold(
ht: hl.Table,
metric_exprs: Dict[str, hl.expr.Float64Expression],
thresholds: Dict[Tuple[str, str], List[float]],
granularities: Union[List[str], Tuple[str, ...]],
field_name: str = "constraint_bins",
) -> hl.Table:
"""
Annotate rows with bin assignments using pre-computed thresholds.
**Threshold-based binning, step 2**: assigns each row to a bin by
comparing its metric value against boundary values produced by
:func:`compute_percentile_thresholds` (step 1). For each
``(granularity, metric)`` pair, the bin equals the number of threshold
boundaries the value exceeds. Bin 0 is the most constrained (below all
thresholds); bin N means the value exceeds all N boundaries.
This differs from rank-based binning (see :func:`rank_and_assign_bins`),
where every row gets a unique position. Here, multiple rows can share a
bin, and the thresholds may have been derived from a different subset of
rows than those being annotated.
:func:`compute_percentile_thresholds` returns a dict of percentile to value,
so ``thresholds`` is built from it by selecting the percentiles for each
granularity. Example use, with decile bins for two metrics:
.. code-block:: python
deciles = [10, 20, 30, 40, 50, 60, 70, 80, 90]
metric_exprs = {"lof": ht.lof_oe_upper, "mis": ht.mis_oe_upper}
thresholds = {}
for metric, expr in metric_exprs.items():
pct_to_value = compute_percentile_thresholds(ht, expr, percentiles=deciles)
thresholds[("decile", metric)] = [pct_to_value[p] for p in deciles]
ht = annotate_bins_by_threshold(ht, metric_exprs, thresholds, ["decile"])
:param ht: Input Table.
:param metric_exprs: Mapping of metric name to the Float64Expression to
bin (e.g. ``{"lof": ht.lof_oe_upper, "mis": ht.mis_oe_upper}``).
:param thresholds: Mapping of ``(granularity, metric)`` to a list of threshold
values in ascending order, built from :func:`compute_percentile_thresholds`
output (see the example above).
:param granularities: Granularity names to iterate over (e.g.
``["decile", "ventile"]``). Each must appear as the first element of
at least one key in ``thresholds``.
:param field_name: Name of the struct field to annotate on ``ht``.
Default is ``"constraint_bins"``.
:return: Table with an added struct field containing per-granularity,
per-metric bin assignments.
"""
miss = hl.missing(hl.tint32)
def _bin_expr(
value_expr: hl.expr.Float64Expression,
threshold_list: List[float],
) -> hl.expr.Int32Expression:
"""
Count how many thresholds a value meets or exceeds.
:param value_expr: Metric value to bin.
:param threshold_list: Boundary values in ascending order.
:return: Number of boundaries met or exceeded (0 = below all thresholds).
"""
arr = hl.literal(threshold_list)
return hl.sum(arr.map(lambda t: hl.int(value_expr >= t)))
granularities_expr = {
gran: hl.struct(
**{
metric: hl.if_else(
hl.is_defined(metric_exprs[metric]),
_bin_expr(metric_exprs[metric], thresholds[(gran, metric)]),
miss,
)
for metric in metric_exprs
}
)
for gran in granularities
}
return ht.annotate(**{field_name: hl.struct(**granularities_expr)})
[docs]def rank_array_element_metrics(
ht: hl.Table,
array_field: str,
element_value_fn: Callable[
[hl.expr.StructExpression], Dict[str, hl.expr.Float64Expression]
],
filter_fn: Optional[Callable[[hl.Table], hl.expr.BooleanExpression]] = None,
bin_granularities: Optional[Dict[str, int]] = None,
rank_field_prefix: str = "",
) -> hl.Table:
"""
Rank metrics within array elements and annotate rank structs back.
**Rank-based binning for array fields**: applies
:func:`rank_and_assign_bins` independently to each element of an array
field. For each element, ``element_value_fn`` extracts named metric
values, which are ranked across rows (optionally on a filtered subset
via ``filter_fn``). Each array element is then annotated with
``{metric_name}_rank`` structs containing rank and bin fields.
Rows not matching ``filter_fn`` get missing rank annotations.
This differs from threshold-based binning (see
:func:`compute_percentile_thresholds` and
:func:`annotate_bins_by_threshold`), where bins are assigned by
comparing values against pre-computed boundaries rather than sorted
position.
.. note::
This function evaluates its input when it is called. It caches the full
input Table (with an added row index) and the extracted metric values,
reads one row to get the array length, and each
:func:`rank_and_assign_bins` call it makes caches its own ranked rows.
Caching avoids recomputing ``ht`` and ``filter_fn`` for every metric, but
in Hail ``Table.cache()`` writes a checkpoint under the temporary directory
(see :func:`rank_and_assign_bins`), so this writes a copy of ``ht`` there.
The copies are not removed when the function returns.
:param ht: Input Table.
:param array_field: Name of the array field on ``ht``.
:param element_value_fn: Function that takes an array element
(StructExpression) and returns a dict mapping metric names to
Float64Expressions to rank. Applied identically to every element.
:param filter_fn: Optional function that takes a Table and returns a
BooleanExpression to filter rows before ranking. When None, all
rows are ranked.
:param bin_granularities: Bin granularities passed to
:func:`rank_and_assign_bins`.
:param rank_field_prefix: Prefix for the ``rank`` and ``bin_{name}``
sub-fields within each rank struct. Passed through to
:func:`rank_and_assign_bins`. Default is ``""`` (no prefix).
:return: Table with ``{metric_name}_rank`` structs added to each array
element. The table is returned with its original key restored.
"""
logger.warning(
"rank_array_element_metrics evaluates its input and checkpoints it, and the"
" ranked values, to the Hail temporary directory (via .cache()); see the"
" function's docstring."
)
original_key = list(ht.key)
ht = ht.add_index("_rank_idx").key_by("_rank_idx").cache()
subset_ht = ht.filter(filter_fn(ht)) if filter_fn is not None else ht
# Extract values to rank using element_value_fn applied via .map().
subset_ht = subset_ht.select(
_rank_values=subset_ht[array_field].map(
lambda elem: hl.struct(**element_value_fn(elem))
)
).cache()
# Metric names come from the type, but the array length is not part of the type,
# so the element count is taken from a sample row.
metric_names = list(subset_ht._rank_values.dtype.element_type.fields)
sample = subset_ht.take(1)
if not sample:
raise ValueError("No rows left to rank after applying 'filter_fn'.")
n_elements = len(sample[0]._rank_values)
# Rank each metric within each array element.
subset_ht = subset_ht.annotate(
_rank_values=[
hl.struct(
**{
name: rank_and_assign_bins(
subset_ht._rank_values[i][name],
bin_granularities,
prefix=rank_field_prefix,
)
for name in metric_names
}
)
for i in range(n_elements)
]
).cache()
# Join ranks back to the original table. Use or_missing so that
# unranked rows get correctly-typed missing rank annotations without
# needing a manually-constructed missing struct for if_else.
ht = ht.annotate(_ranks=subset_ht[ht._rank_idx]._rank_values)
ht = ht.annotate(
**{
array_field: [
ht[array_field][i].annotate(
**{
f"{name}_rank": hl.or_missing(
hl.is_defined(ht._ranks),
ht._ranks[i][name],
)
for name in metric_names
}
)
for i in range(n_elements)
]
}
).drop("_ranks")
return ht.key_by(*original_key).drop("_rank_idx")