From 1fab714e6acc3030f8dd08882359c44cdbe1ca47 Mon Sep 17 00:00:00 2001 From: Kristen Laricchia Date: Fri, 2 Aug 2024 17:11:04 -0400 Subject: [PATCH 1/3] add support for pops --- gnomad/utils/constraint.py | 40 ++++++++++++++++++++++++++++---------- 1 file changed, 30 insertions(+), 10 deletions(-) diff --git a/gnomad/utils/constraint.py b/gnomad/utils/constraint.py index b605474d5..8f4d22c71 100644 --- a/gnomad/utils/constraint.py +++ b/gnomad/utils/constraint.py @@ -62,6 +62,7 @@ def count_variants_by_group( use_table_group_by: bool = False, singleton_expr: Optional[hl.expr.BooleanExpression] = None, max_af: Optional[float] = None, + skip_downsamplings: bool=False, ) -> Union[hl.Table, Any]: """ Count number of observed or possible variants by context, ref, alt, and optionally methylation_level. @@ -145,6 +146,7 @@ def count_variants_by_group( [0].AC == 1`. Default is None. :param max_af: Maximum variant allele frequency to keep. By default, no cutoff is applied. + :param skip_downsamplings: Whether of not to skip pulling the downsampling data. :return: Table including 'variant_count' annotation and if requested, `singleton_count` and downsampling counts. """ @@ -206,12 +208,13 @@ def count_variants_by_group( pop, pop, ) - agg[f"downsampling_counts_{pop}"] = downsampling_counts_expr( + agg[f"downsampling_counts_{pop}"] = pop_counts_expr( freq_expr, freq_meta_expr, pop, max_af=max_af, downsamplings=downsamplings, + skip_downsamplings=skip_downsamplings, ) if count_singletons: logger.info( @@ -220,12 +223,13 @@ def count_variants_by_group( pop, pop, ) - agg[f"singleton_downsampling_counts_{pop}"] = downsampling_counts_expr( + agg[f"singleton_downsampling_counts_{pop}"] = pop_counts_expr( freq_expr, freq_meta_expr, pop, max_af=max_af, downsamplings=downsamplings, + skip_downsamplings=skip_downsamplings, singleton=True, ) # Apply each variant count aggregation in `agg` to get counts for all @@ -238,13 +242,15 @@ def count_variants_by_group( ) -def get_downsampling_freq_indices( + +def get_pop_freq_indices( freq_meta_expr: hl.expr.ArrayExpression, pop: str = "global", variant_quality: str = "adj", genetic_ancestry_label: Optional[str] = None, subset: Optional[str] = None, downsamplings: Optional[List[int]] = None, + skip_downsamplings: bool=False, ) -> hl.expr.ArrayExpression: """ Get indices of dictionaries in meta dictionaries that only have the "downsampling" key with specified `genetic_ancestry_label` and "variant_quality" values. @@ -264,6 +270,7 @@ def get_downsampling_freq_indices( key in `freq_meta_expr`. :param downsamplings: Optional List of integers specifying what downsampling indices to obtain. Default is None, which will return all downsampling indices. + :param skip_downsamplings: Whether of not to skip pulling the downsampling data. :return: ArrayExpression of indices of dictionaries in `freq_meta_expr` that only have the "downsampling" key with specified `genetic_ancestry_label` and "variant_quality" values. @@ -277,12 +284,18 @@ def _get_filter_expr(m: hl.expr.StructExpression) -> hl.expr.BooleanExpression: filter_expr = ( (m.get("group") == variant_quality) & (hl.any([m.get(l, "") == pop for l in gen_anc])) - & m.contains("downsampling") + & ~m.contains("sex") ) - if downsamplings is not None: - filter_expr &= hl.literal(downsamplings).contains( + + if skip_downsamplings: + filter_expr &= ~m.contains("downsampling") + else: + if downsamplings is not None: + filter_expr &= hl.literal(downsamplings).contains( hl.int(m.get("downsampling", "0")) ) + + if subset is None: filter_expr &= ~m.contains("subset") else: @@ -290,12 +303,13 @@ def _get_filter_expr(m: hl.expr.StructExpression) -> hl.expr.BooleanExpression: return filter_expr indices = hl.enumerate(freq_meta_expr).filter(lambda f: _get_filter_expr(f[1])) + + # Get an array of indices and meta dictionaries sorted by "downsampling" key if present. + return hl.sorted(indices, key=lambda f: hl.int(f[1].get("downsampling", "0"))) - # Get an array of indices and meta dictionaries sorted by "downsampling" key. - return hl.sorted(indices, key=lambda f: hl.int(f[1]["downsampling"])) -def downsampling_counts_expr( +def pop_counts_expr( freq_expr: hl.expr.ArrayExpression, freq_meta_expr: hl.expr.ArrayExpression, pop: str = "global", @@ -305,6 +319,7 @@ def downsampling_counts_expr( genetic_ancestry_label: Optional[str] = None, subset: Optional[str] = None, downsamplings: Optional[List[int]] = None, + skip_downsamplings: bool=False, ) -> hl.expr.ArrayExpression: """ Return an aggregation expression to compute an array of counts of all downsamplings found in `freq_expr` where specified criteria is met. @@ -335,17 +350,19 @@ def downsampling_counts_expr( subset will be included. :param downsamplings: Optional List of integers specifying what downsampling indices to obtain. Default is None, which will return all downsampling counts. + :param skip_downsamplings: Whether of not to skip pulling the downsampling data. :return: Aggregation Expression for an array of the variant counts in downsamplings for specified population. """ # Get an array of indices sorted by "downsampling" key. - sorted_indices = get_downsampling_freq_indices( + sorted_indices = get_pop_freq_indices( freq_meta_expr, pop, variant_quality, genetic_ancestry_label, subset, downsamplings, + skip_downsamplings, ).map(lambda x: x[0]) def _get_criteria(i: hl.expr.Int32Expression) -> hl.expr.Int32Expression: @@ -1148,6 +1165,9 @@ def oe_aggregation_expr( agg_expr["gen_anc_obs"] = hl.struct( **{pop: hl.agg.array_sum(ht[f"downsampling_counts_{pop}"]) for pop in pops} ) + agg_expr["gen_anc_oe"] = hl.struct( + **{pop: hl.map(lambda x: divide_null(x[0], x[1]), hl.zip(agg_expr["gen_anc_obs"][pop], agg_expr["gen_anc_exp"][pop])) for pop in pops} + ) agg_expr = hl.struct(**agg_expr) return hl.agg.group_by(filter_expr, agg_expr).get(True, hl.missing(agg_expr.dtype)) From 3dc35c078fa3bb288e700c1aa6a777063a7a268d Mon Sep 17 00:00:00 2001 From: Kristen Laricchia Date: Mon, 5 Aug 2024 10:31:54 -0400 Subject: [PATCH 2/3] small edit --- gnomad/utils/constraint.py | 31 +++++++++++++++++-------------- 1 file changed, 17 insertions(+), 14 deletions(-) diff --git a/gnomad/utils/constraint.py b/gnomad/utils/constraint.py index 8f4d22c71..6ce82c8c7 100644 --- a/gnomad/utils/constraint.py +++ b/gnomad/utils/constraint.py @@ -62,7 +62,7 @@ def count_variants_by_group( use_table_group_by: bool = False, singleton_expr: Optional[hl.expr.BooleanExpression] = None, max_af: Optional[float] = None, - skip_downsamplings: bool=False, + skip_downsamplings: bool = False, ) -> Union[hl.Table, Any]: """ Count number of observed or possible variants by context, ref, alt, and optionally methylation_level. @@ -146,7 +146,7 @@ def count_variants_by_group( [0].AC == 1`. Default is None. :param max_af: Maximum variant allele frequency to keep. By default, no cutoff is applied. - :param skip_downsamplings: Whether of not to skip pulling the downsampling data. + :param skip_downsamplings: Whether or not to skip pulling the downsampling data. :return: Table including 'variant_count' annotation and if requested, `singleton_count` and downsampling counts. """ @@ -242,7 +242,6 @@ def count_variants_by_group( ) - def get_pop_freq_indices( freq_meta_expr: hl.expr.ArrayExpression, pop: str = "global", @@ -250,10 +249,10 @@ def get_pop_freq_indices( genetic_ancestry_label: Optional[str] = None, subset: Optional[str] = None, downsamplings: Optional[List[int]] = None, - skip_downsamplings: bool=False, + skip_downsamplings: bool = False, ) -> hl.expr.ArrayExpression: """ - Get indices of dictionaries in meta dictionaries that only have the "downsampling" key with specified `genetic_ancestry_label` and "variant_quality" values. + Get indices of dictionaries in meta dictionaries with specified `genetic_ancestry_label`, `variant_quality` values, and downsamplings if specified. :param freq_meta_expr: ArrayExpression containing the set of groupings for each element of the `freq_expr` array (e.g., [{'group': 'adj'}, {'group': 'adj', @@ -270,7 +269,7 @@ def get_pop_freq_indices( key in `freq_meta_expr`. :param downsamplings: Optional List of integers specifying what downsampling indices to obtain. Default is None, which will return all downsampling indices. - :param skip_downsamplings: Whether of not to skip pulling the downsampling data. + :param skip_downsamplings: Whether or not to skip pulling the downsampling data. :return: ArrayExpression of indices of dictionaries in `freq_meta_expr` that only have the "downsampling" key with specified `genetic_ancestry_label` and "variant_quality" values. @@ -292,10 +291,9 @@ def _get_filter_expr(m: hl.expr.StructExpression) -> hl.expr.BooleanExpression: else: if downsamplings is not None: filter_expr &= hl.literal(downsamplings).contains( - hl.int(m.get("downsampling", "0")) - ) - - + hl.int(m.get("downsampling", "0")) + ) + if subset is None: filter_expr &= ~m.contains("subset") else: @@ -303,12 +301,11 @@ def _get_filter_expr(m: hl.expr.StructExpression) -> hl.expr.BooleanExpression: return filter_expr indices = hl.enumerate(freq_meta_expr).filter(lambda f: _get_filter_expr(f[1])) - + # Get an array of indices and meta dictionaries sorted by "downsampling" key if present. return hl.sorted(indices, key=lambda f: hl.int(f[1].get("downsampling", "0"))) - def pop_counts_expr( freq_expr: hl.expr.ArrayExpression, freq_meta_expr: hl.expr.ArrayExpression, @@ -319,7 +316,7 @@ def pop_counts_expr( genetic_ancestry_label: Optional[str] = None, subset: Optional[str] = None, downsamplings: Optional[List[int]] = None, - skip_downsamplings: bool=False, + skip_downsamplings: bool = False, ) -> hl.expr.ArrayExpression: """ Return an aggregation expression to compute an array of counts of all downsamplings found in `freq_expr` where specified criteria is met. @@ -1166,7 +1163,13 @@ def oe_aggregation_expr( **{pop: hl.agg.array_sum(ht[f"downsampling_counts_{pop}"]) for pop in pops} ) agg_expr["gen_anc_oe"] = hl.struct( - **{pop: hl.map(lambda x: divide_null(x[0], x[1]), hl.zip(agg_expr["gen_anc_obs"][pop], agg_expr["gen_anc_exp"][pop])) for pop in pops} + **{ + pop: hl.map( + lambda x: divide_null(x[0], x[1]), + hl.zip(agg_expr["gen_anc_obs"][pop], agg_expr["gen_anc_exp"][pop]), + ) + for pop in pops + } ) agg_expr = hl.struct(**agg_expr) From d57e3f3ca5775342b593ab3ad67d282277031aa7 Mon Sep 17 00:00:00 2001 From: Kristen Laricchia Date: Mon, 5 Aug 2024 10:39:27 -0400 Subject: [PATCH 3/3] small edit --- gnomad/utils/constraint.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/gnomad/utils/constraint.py b/gnomad/utils/constraint.py index 6ce82c8c7..90dd06162 100644 --- a/gnomad/utils/constraint.py +++ b/gnomad/utils/constraint.py @@ -302,7 +302,8 @@ def _get_filter_expr(m: hl.expr.StructExpression) -> hl.expr.BooleanExpression: indices = hl.enumerate(freq_meta_expr).filter(lambda f: _get_filter_expr(f[1])) - # Get an array of indices and meta dictionaries sorted by "downsampling" key if present. + # Get an array of indices and meta dictionaries sorted by "downsampling" + # key if present. return hl.sorted(indices, key=lambda f: hl.int(f[1].get("downsampling", "0")))