Extract strata aggregation into it's own function and use in compute_freq_by_strata

Author: jkgoodrich

gnomad/utils/annotations.py

@@ -1861,25 +1861,19 @@ def generate_freq_group_membership_array(
     return ht
 
 
-def compute_freq_by_strata(
+def agg_by_strata(
     mt: hl.MatrixTable,
     entry_agg_funcs: Optional[Dict[str, Tuple[Callable, Callable]]] = None,
     select_fields: Optional[List[str]] = None,
-    group_membership_includes_raw_group: bool = True,
+    group_membership_ht: Optional[hl.Table] = None,
 ) -> hl.Table:
     """
-    Compute call statistics and, when passed, entry aggregation function(s) by strata.
-
-    The computed call statistics are AC, AF, AN, and homozygote_count. The entry
-    aggregation functions are applied to the MatrixTable entries and aggregated. The
-    MatrixTable must contain a 'group_membership' annotation (like the one added by
-    `generate_freq_group_membership_array`) that is a list of bools to aggregate the
-    columns by.
+    Get row expression for annotations of each entry aggregation function(s) by strata.
 
-    .. note::
-        This function is primarily used through `annotate_freq` but can be used
-        independently if desired. Please see the `annotate_freq` function for more
-        complete documentation.
+    The entry aggregation functions are applied to the MatrixTable entries and
+    aggregated. If no `group_membership_ht` (like the one returned by
+    `generate_freq_group_membership_array`) is supplied, `mt` must contain a
+    'group_membership' annotation that is a list of bools to aggregate the columns by.
 
     :param mt: Input MatrixTable.
     :param entry_agg_funcs: Optional dict of entry aggregation functions. When
@@ -1890,15 +1884,9 @@ def compute_freq_by_strata(
         function.
     :param select_fields: Optional list of row fields from `mt` to keep on the output
         Table.
-    :param group_membership_includes_raw_group: Whether the 'group_membership'
-        annotation includes an entry for the 'raw' group, representing all samples. If
-        False, the 'raw' group is inserted as the second element in all added
-        annotations using the same 'group_membership', resulting
-        in array lengths of 'group_membership'+1. If True, the second element of each
-        added annotation is still the 'raw' group, but the group membership is
-        determined by the values in the second element of 'group_membership', and the
-        output annotations will be the same length as 'group_membership'. Default is
-        True.
+    :param group_membership_ht: Optional Table containing group membership annotations
+        to stratify the coverage stats by. If not provided, the 'group_membership'
+        annotation is expected to be present on `mt`.
     :return: Table or MatrixTable with allele frequencies by strata.
     """
     if entry_agg_funcs is None:
@@ -1907,79 +1895,204 @@ def compute_freq_by_strata(
         select_fields = []
 
     n_samples = mt.count_cols()
-    n_groups = len(mt.group_membership.take(1)[0])
-    ht = mt.localize_entries("entries", "cols")
-    ht = ht.annotate_globals(
-        indices_by_group=hl.range(n_groups).map(
-            lambda g_i: hl.range(n_samples).filter(
-                lambda s_i: ht.cols[s_i].group_membership[g_i]
+    global_expr = {}
+    if "adj_group" in mt.index_globals():
+        global_expr["adj_group"] = mt.index_globals().adj_group
+        logger.info("Using the 'adj_group' global annotation found on the input MT.")
+
+    if group_membership_ht is None and "group_membership" not in mt.col:
+        raise ValueError(
+            "The 'group_membership' annotation is not found in the input MatrixTable "
+            "and 'group_membership_ht' is not specified."
+        )
+    elif group_membership_ht is None:
+        logger.info(
+            "'group_membership_ht' is not specified, using sample stratification "
+            "indicated by the 'group_membership' annotation on mt."
+        )
+        n_groups = len(mt.group_membership.take(1)[0])
+    else:
+        logger.info(
+            "'group_membership_ht' is specified, using sample stratification indicated "
+            "by its 'group_membership' annotation."
+        )
+        n_groups = len(group_membership_ht.group_membership.take(1)[0])
+        mt = mt.annotate_cols(
+            group_membership=group_membership_ht[mt.col_key].group_membership
+        )
+        if "adj_group" not in global_expr:
+            if "adj_group" in group_membership_ht.index_globals():
+                global_expr["adj_group"] = mt.index_globals().adj_group
+                logger.info(
+                    "Using the 'adj_group' global annotation on 'group_membership_ht'."
+                )
+            elif "freq_meta" in group_membership_ht.index_globals():
+                logger.info(
+                    "The 'freq_meta' global annotation is found in "
+                    "'group_membership_ht', using it to determine the adj filtered "
+                    "stratification groups."
+                )
+                freq_meta = group_membership_ht.index_globals().freq_meta
+
+            global_expr["adj_group"] = freq_meta.map(
+                lambda x: x.get("group", "NA") == "adj"
             )
+
+    if "adj_group" not in global_expr:
+        global_expr["adj_group"] = hl.range(n_groups).map(lambda x: False)
+
+    n_adj_group = hl.eval(hl.len(global_expr["adj_group"]))
+    if hl.eval(hl.len(global_expr["adj_group"])) != n_groups:
+        raise ValueError(
+            f"The number of elements in the 'adj_group' ({n_adj_group}) global "
+            "annotation does not match the number of elements in the "
+            f"'group_membership' annotation ({n_groups})!",
+        )
+
+    # Keep only the entries needed for the aggregation functions.
+    select_expr = {}
+    if hl.eval(hl.any(global_expr["adj_group"])):
+        select_expr["adj"] = mt.adj
+
+    select_expr.update(**{ann: f[0](mt) for ann, f in entry_agg_funcs.items()})
+    mt = mt.select_entries(**select_expr)
+
+    # Convert MT to HT with a row annotation that is an array of all samples entries
+    # for that variant.
+    ht = mt.localize_entries("entries", "cols")
+
+    # For each stratification group in group_membership, determine the indices of the
+    # samples that belong to that group.
+    global_expr["indices_by_group"] = hl.range(n_groups).map(
+        lambda g_i: hl.range(n_samples).filter(
+            lambda s_i: ht.cols[s_i].group_membership[g_i]
         )
     )
+    ht = ht.annotate_globals(**global_expr)
+
     # Pull out each annotation that will be used in the array aggregation below as its
     # own ArrayExpression. This is important to prevent memory issues when performing
     # the below array aggregations.
     ht = ht.select(
-        *select_fields,
-        adj_array=ht.entries.map(lambda e: e.adj),
-        gt_array=ht.entries.map(lambda e: e.GT),
         **{
-            ann: hl.map(lambda e, s: f[0](e, s), ht.entries, ht.cols)
-            for ann, f in entry_agg_funcs.items()
-        },
+            ann: ht.entries.map(lambda e: e[ann])
+            for ann in select_fields + list(select_expr.keys())
+        }
     )
 
     def _agg_by_group(
-        ht: hl.Table, agg_func: Callable, ann_expr: hl.expr.ArrayExpression, *args
+        ht: hl.Table, agg_func: Callable, ann_expr: hl.expr.ArrayExpression
     ) -> hl.expr.ArrayExpression:
         """
         Aggregate `agg_expr` by group using the `agg_func` function.
 
         :param ht: Input Hail Table.
-        :param agg_func: Aggregation function to apply to `agg_expr`.
-        :param agg_expr: Expression to aggregate by group.
-        :param args: Additional arguments to pass to the `agg_func`.
+        :param agg_func: Aggregation function to apply to `ann_expr`.
+        :param ann_expr: Expression to aggregate by group.
         :return: Aggregated array expression.
         """
-        adj_agg_expr = ht.indices_by_group.map(
-            lambda s_indices: s_indices.aggregate(
-                lambda i: hl.agg.filter(ht.adj_array[i], agg_func(ann_expr[i], *args))
-            )
-        )
-        # Create final agg list by inserting or changing the "raw" group,
-        # representing all samples, in the adj_agg_list.
-        raw_agg_expr = ann_expr.aggregate(lambda x: agg_func(x, *args))
-        if group_membership_includes_raw_group:
-            extend_idx = 2
-        else:
-            extend_idx = 1
-
-        adj_agg_expr = (
-            adj_agg_expr[:1].append(raw_agg_expr).extend(adj_agg_expr[extend_idx:])
+        return hl.map(
+            lambda s_indices, adj: s_indices.aggregate(
+                lambda i: hl.if_else(
+                    adj,
+                    hl.agg.filter(ht.adj[i], agg_func(ann_expr[i])),
+                    agg_func(ann_expr[i]),
+                )
+            ),
+            ht.indices_by_group,
+            ht.adj_group,
         )
 
-        return adj_agg_expr
-
-    freq_expr = _agg_by_group(ht, hl.agg.call_stats, ht.gt_array, ht.alleles)
-
-    # Select non-ref allele (assumes bi-allelic).
-    freq_expr = freq_expr.map(
-        lambda cs: cs.annotate(
-            AC=cs.AC[1],
-            AF=cs.AF[1],
-            homozygote_count=cs.homozygote_count[1],
-        )
-    )
     # Add annotations for any supplied entry transform and aggregation functions.
     ht = ht.select(
         *select_fields,
         **{ann: _agg_by_group(ht, f[1], ht[ann]) for ann, f in entry_agg_funcs.items()},
-        freq=freq_expr,
     )
 
     return ht.drop("cols")
 
 
+def compute_freq_by_strata(
+    mt: hl.MatrixTable,
+    entry_agg_funcs: Optional[Dict[str, Tuple[Callable, Callable]]] = None,
+    select_fields: Optional[List[str]] = None,
+    group_membership_includes_raw_group: bool = True,
+) -> hl.Table:
+    """
+    Compute call statistics and, when passed, entry aggregation function(s) by strata.
+
+    The computed call statistics are AC, AF, AN, and homozygote_count. The entry
+    aggregation functions are applied to the MatrixTable entries and aggregated. The
+    MatrixTable must contain a 'group_membership' annotation (like the one added by
+    `generate_freq_group_membership_array`) that is a list of bools to aggregate the
+    columns by.
+
+    .. note::
+        This function is primarily used through `annotate_freq` but can be used
+        independently if desired. Please see the `annotate_freq` function for more
+        complete documentation.
+
+    :param mt: Input MatrixTable.
+    :param entry_agg_funcs: Optional dict of entry aggregation functions. When
+        specified, additional annotations are added to the output Table/MatrixTable.
+        The keys of the dict are the names of the annotations and the values are tuples
+        of functions. The first function is used to transform the `mt` entries in some
+        way, and the second function is used to aggregate the output from the first
+        function.
+    :param select_fields: Optional list of row fields from `mt` to keep on the output
+        Table.
+    :param group_membership_includes_raw_group: Whether the 'group_membership'
+        annotation includes an entry for the 'raw' group, representing all samples. If
+        False, the 'raw' group is inserted as the second element in all added
+        annotations using the same 'group_membership', resulting
+        in array lengths of 'group_membership'+1. If True, the second element of each
+        added annotation is still the 'raw' group, but the group membership is
+        determined by the values in the second element of 'group_membership', and the
+        output annotations will be the same length as 'group_membership'. Default is
+        True.
+    :return: Table or MatrixTable with allele frequencies by strata.
+    """
+    if not group_membership_includes_raw_group:
+        # Add the 'raw' group to the 'group_membership' annotation.
+        mt = mt.annotate_cols(
+            group_membership=hl.array([mt.group_membership[0]]).extend(
+                mt.group_membership
+            )
+        )
+
+    # Add adj_group global annotation indicating that the second element in
+    # group_membership is 'raw' and all others are 'adj'.
+    mt = mt.annotate_globals(
+        adj_group=hl.range(hl.len(mt.group_membership.take(1)[0])).map(lambda x: x != 1)
+    )
+
+    if entry_agg_funcs is None:
+        entry_agg_funcs = {}
+
+    def _get_freq_expr(gt_expr: hl.expr.CallExpression) -> hl.expr.StructExpression:
+        """
+        Get struct expression with call statistics.
+
+        :param gt_expr: CallExpression to compute call statistics on.
+        :return: StructExpression with call statistics.
+        """
+        # Get the source Table for the CallExpression to grab alleles.
+        ht = gt_expr._indices.source
+        freq_expr = hl.agg.call_stats(gt_expr, ht.alleles)
+        # Select non-ref allele (assumes bi-allelic).
+        freq_expr = freq_expr.annotate(
+            AC=freq_expr.AC[1],
+            AF=freq_expr.AF[1],
+            homozygote_count=freq_expr.homozygote_count[1],
+        )
+
+        return freq_expr
+
+    entry_agg_funcs["freq"] = (lambda x: x.GT, _get_freq_expr)
+
+    return agg_by_strata(mt, entry_agg_funcs, select_fields).drop("adj_group")
+
+
 def update_structured_annotations(
     ht: hl.Table,
     annotation_update_exprs: Dict[str, hl.Expression],