Merge freq array function and new frequency dictionary builder

gnomad/utils/annotations.py

@@ -125,7 +125,7 @@ def project_max_expr(
                     > 0
                 ),
                 # order the callstats computed by AF in decreasing order
-                lambda x: -x[1].AF[ai]
+                lambda x: -x[1].AF[ai],
                 # take the n_projects projects with largest AF
             )[:n_projects].map(
                 # add the project in the callstats struct
@@ -219,8 +219,10 @@ def faf_expr(
                 ~locus.in_autosome_or_par(),
                 hl.struct(
                     **{
-                        f"faf{str(threshold)[2:]}": hl.experimental.filtering_allele_frequency(
-                            freq[i].AC, freq[i].AN, threshold
+                        f"faf{str(threshold)[2:]}": (
+                            hl.experimental.filtering_allele_frequency(
+                                freq[i].AC, freq[i].AN, threshold
+                            )
                         )
                         for threshold in faf_thresholds
                     }
@@ -1106,7 +1108,7 @@ def region_flag_expr(
     :return: `region_flag` struct row annotation
     """
     prob_flags_expr = (
-        {"non_par": (t.locus.in_x_nonpar() | t.locus.in_y_nonpar())} if non_par else {}
+        {"non_par": t.locus.in_x_nonpar() | t.locus.in_y_nonpar()} if non_par else {}
     )
 
     if prob_regions is not None:
@@ -1188,3 +1190,132 @@ def hemi_expr(
         # mt.GT[0] is alternate allele
         gt.is_haploid() & (sex_expr == male_str) & (gt[0] == 1),
     )
+
+
+def merge_freq_arrays(
+    farrays: List[hl.expr.ArrayExpression],
+    fmeta: List[List[Dict[str, str]]],
+    operation="sum",
+    set_negatives_to_zero: bool = True,
+) -> Tuple[hl.expr.ArrayExpression, Dict[str, int]]:
+    """
+    Merge frequency arrays on the same HT.
+
+    .. note::
+        Arrays do not have to contain the same groupings or order of groupings but
+        the array indices for a freq array in farrays must be the same as its associated
+        frequency metadata index in fmeta i.e., farrays = [freq1, freq2] then fmeta
+        must equal [fmeta1, fmeta2] where fmeta1 contains the metadata information
+        for freq1. If th merge operation is set to "sum", groups in the merged array
+        will be the union of groupings found within the arrays' metadata and all arrays
+        with be summed by grouping. If the operation is set to "diff", the merged array
+        will contain groups only found in the first array of fmeta. Any array containing
+        any of these groups will have thier values subtracted from the values of the first array.
+
+    :param farrays: List of frequency arrays to merge. First entry in the list is the primary array to which other arrays will be added or subtracted. Array must be on the same HT.
+    :param fmeta: List of frequency metadata for arrays being merged.
+    :param operation: Merge operation to perform. Options are "sum" and "diff". If "diff" is passed, the first freq array in the list will have the other arrays subtracted from it.
+    :param set_negatives_to_zero: If True, set negative array values to 0 for AC, AN,
+    AF, and homozygote_count. If False, raise a ValueError. Default is True.
+    :return: Tuple of merged frequency array and its freq index dictionary.
+    """
+    if len(farrays) < 2:
+        raise ValueError("Must provide at least two frequency arrays to merge")
+    if len(farrays) != len(fmeta):
+        raise ValueError("Length of farrays and fmeta must be equal")
+    if operation not in ["sum", "diff"]:
+        raise ValueError("Operation must be either 'sum' or 'diff'")
+
+    # Create a list where each entry is a dictionary whose the key is an aggregation
+    # group and the value is the corresponding index in the freq array. If operation is
+    # set to diff, only use the first entry's meta.
+    fmeta = [hl.dict(hl.enumerate(f).map(lambda x: (x[1], [x[0]]))) for f in fmeta]
+
+    # Merge dictionaries in the list into a single dictionary where key is aggregation
+    # group and the value is a list of the group's index in each of the freq arrays, if
+    # it exists. For sum operation, use keys, aka groups, found in all freq dictionaries.
+    # For "diff" operations, only use key_set from the first entry.
+    fmeta = hl.fold(
+        lambda i, j: hl.dict(
+            (
+                hl.if_else(operation == "sum", (i.key_set() | j.key_set()), i.key_set())
+            ).map(
+                lambda k: (
+                    k,
+                    i.get(k, [hl.missing(hl.tint32)]).extend(
+                        j.get(k, [hl.missing(hl.tint32)])
+                    ),
+                )
+            )
+        ),
+        fmeta[0],
+        fmeta[1:],
+    )
+
+    # Create a list of tuples from the dictionary, sorted by the list of indices for
+    # each aggregation group.
+    fmeta = hl.sorted(fmeta.items(), key=lambda f: f[1])
+
+    # Create a list of the aggregation groups, maintaining the sorted order
+    new_freq_meta = fmeta.map(lambda x: x[0])
+
+    # Create array for each aggregation group of arrays containing the group's freq
+    # values from each freq array
+    freq_meta_idx = fmeta.map(lambda x: hl.zip(farrays, x[1]).map(lambda i: i[0][i[1]]))
+
+    def _sum_or_diff_fields(
+        field_1_expr: str, field_2_expr: str
+    ) -> hl.expr.Int32Expression:
+        """
+        Sum or subtract fields in call statistics struct.
+
+        :param field_1_expr: First field to sum or diff.
+        :param field_2_expr: Second field to sum or diff.
+        :return: Merged field value.
+        """
+        return hl.if_else(
+            operation == "sum",
+            hl.or_else(field_1_expr, 0) + hl.or_else(field_2_expr, 0),
+            hl.or_else(field_1_expr, 0) - hl.or_else(field_2_expr, 0),
+        )
+
+    # Iterate through the groups and their freq lists to merge callstats
+    new_freq = freq_meta_idx.map(
+        lambda x: hl.bind(
+            lambda y: y.annotate(AF=hl.if_else(y.AN > 0, y.AC / y.AN, 0)),
+            hl.fold(
+                lambda i, j: hl.struct(
+                    AC=_sum_or_diff_fields(i.AC, j.AC),
+                    AN=_sum_or_diff_fields(i.AN, j.AN),
+                    homozygote_count=_sum_or_diff_fields(
+                        i.homozygote_count, j.homozygote_count
+                    ),
+                ),
+                x[0].select("AC", "AN", "homozygote_count"),
+                x[1:],
+            ),
+        )
+    )
+    # Check and see if any annotation within the merged array is negative. If so,
+    # raise an error if set_negatives_to_zero is False or set the value to 0 if
+    # set_negatives_to_zero is True.
+    if operation == "diff":
+        negative_value_error_msg = (
+            "Negative values found in merged frequency array. Review data or set"
+            " `set_negatives_to_zero` to True to set negative values to 0."
+        )
+        new_freq = new_freq.map(
+            lambda x: x.annotate(
+                AC=hl.case()
+                .when(set_negatives_to_zero, hl.max(x.AC, 0))
+                .or_error(negative_value_error_msg),
+                AN=hl.case()
+                .when(set_negatives_to_zero, hl.max(x.AN, 0))
+                .or_error(negative_value_error_msg),
+                homozygote_count=hl.case()
+                .when(set_negatives_to_zero, hl.max(x.homozygote_count, 0))
+                .or_error(negative_value_error_msg),
+            )
+        )
+
+    return new_freq, new_freq_meta

gnomad/utils/release.py

@@ -2,14 +2,16 @@
 
 from typing import Dict, List, Optional
 
+import hail as hl
+
 from gnomad.resources.grch38.gnomad import (
     CURRENT_MAJOR_RELEASE,
     GROUPS,
     POPS,
     SEXES,
     SUBSETS,
 )
-from gnomad.utils.vcf import index_globals
+from gnomad.utils.vcf import SORT_ORDER, index_globals
 
 
 def make_faf_index_dict(
@@ -92,3 +94,34 @@ def _get_index(label_groups):
         )
 
     return index_dict
+
+
+def make_freq_index_dict_from_meta(
+    freq_meta: List[Dict[str, str]],
+    label_delimiter: str = "_",
+    sort_order: Optional[List[str]] = SORT_ORDER,
+) -> Dict[str, int]:
+    """
+    Create a dictionary for accessing frequency array.
+
+    :param freq_meta: List of dictionaries containing frequency metadata.
+    :param label_delimiter: Delimiter to use when joining frequency metadata labels.
+    :param sort_order: List of frequency metadata labels to use when sorting the dictionary.
+    :return: Dictionary of frequency metadata.
+    """
+    index_dict = {}
+    for i, f in enumerate(hl.eval(freq_meta)):
+        if sort_order and len(set(f.keys()) - set(sort_order)) < 1:
+            index_dict[
+                label_delimiter.join(
+                    [
+                        f[g]
+                        for g in sorted(
+                            f.keys(),
+                            key=(lambda x: sort_order.index(x)) if sort_order else None,
+                        )
+                    ]
+                )
+            ] = i
+
+    return index_dict