grpmax vs European AF split nonsynonymous variant analysis

gnomad_qc/v4/analyses/grpmax_comps.py

@@ -8,9 +8,8 @@
 import hail as hl
 from gnomad.utils.vep import (
     CSQ_CODING_HIGH_IMPACT,
-    CSQ_CODING_LOW_IMPACT,
     CSQ_CODING_MEDIUM_IMPACT,
-    CSQ_NON_CODING,
+    filter_vep_to_canonical_transcripts,
 )
 from tabulate import tabulate
 
@@ -82,14 +81,20 @@ def filter_to_threshold(
 
 
 def version_stats(
-    ht: hl.Table, version: str = "v4", can_only: bool = False
+    ht: hl.Table,
+    version: str = "v4",
+    can_only: bool = False,
+    drop_hard_filtered_variants: bool = False,
+    all_variants: bool = False,
 ) -> Dict[str, Dict[str, Dict[str, int]]]:
     """
     Calculate grpmax stats for a given gnomAD version.
 
     :param ht: gnomAD release Table
     :param version: gnomAD version
     :param can_only: Only consider MANE Select and canonical transcripts.
+    :param drop_hard_filtered_variants: Remove only hard filtered variants but keep all other variants regardless of variant QC status
+    :param all_variants: Keep all variants regardless of variant QC status
     :return: Dictionary of grpmax stats
     """
     logger.info("Calculating grpmax stats for %s", version)
@@ -102,30 +107,16 @@ def version_stats(
     t_variants = ht.count()
     if can_only:
         logger.info("Filtering to only MANE Select and canonical transcripts")
-        ht = ht.explode(ht.vep.transcript_consequences)
-        if version == "v2":
-            ht = ht.filter(hl.is_defined(ht.vep.transcript_consequences.canonical))
-        else:
-            ht = ht.filter(
-                hl.is_defined(ht.vep.transcript_consequences.canonical)
-                | hl.is_defined(ht.vep.transcript_consequences.mane_select)
-            )
-        logger.info(
-            "Filtering on mane_select and canonical transcript consequence term to "
-            "keep only non-synonymous variants..."
-        )
+        ht = filter_vep_to_canonical_transcripts(ht, filter_empty_csq=True)
+        # All MANE select transcripts in v4 are also the canonical transcript
+        logger.info("Filtering to keep only non-synonymous variants...")
         ht = ht.filter(
             hl.any(
                 ht.vep.transcript_consequences.consequence_terms.map(
-                    lambda x: ~hl.literal(
-                        CSQ_NON_CODING + CSQ_CODING_LOW_IMPACT
-                    ).contains(x)
+                    lambda x: hl.literal(NS_CONSEQ_TERMS).contains(x)
                 )
             )
         )
-        # Remove duplicate sites after exploding, most_severe_consequence is the same for all
-        # transcript consequences so we can use distinct to remove duplicates
-        ht = ht.distinct()
         ht = ht.checkpoint(
             f"gs://gnomad-tmp-4day/grpmax_comps_{version}_canonical_non_syn.ht",
             overwrite=True,
@@ -147,9 +138,7 @@ def version_stats(
             " variants..."
         )
         ht = ht.filter(
-            ~hl.literal(CSQ_NON_CODING + CSQ_CODING_LOW_IMPACT).contains(
-                ht.vep.most_severe_consequence
-            )
+            hl.literal(NS_CONSEQ_TERMS).contains(ht.vep.most_severe_consequence)
         )
         ht = ht.checkpoint(
             f"gs://gnomad-tmp-4day/grpmax_comps_{version}_non_syn.ht", overwrite=True
@@ -162,17 +151,20 @@ def version_stats(
             ns_variants,
             (ns_variants / t_variants) * 100,
         )
-    if args.drop_hard_filtered_variants:
+    if all_variants:
+        log_message = "of variants, regardless of variant QC status"
+    elif drop_hard_filtered_variants:
         # Filter out AC0 and InbreedingCoeff variants
         ht = ht.filter(~ht.filters.any(lambda x: FILTER_VALUES_TO_DROP.contains(x)))
-        log_message = "ONLY hard filters (RF/VQSR are retained)"
+        log_message = "of variants passing ONLY hard filters (RF/VQSR are retained)"
     else:
+        # Filter to only PASS
         ht = ht.filter(hl.len(ht.filters) == 0)
-        log_message = "all filters"
+        log_message = "of variants passing all filters"
 
     p_ns_variants = ht.count()
     logger.info(
-        "Total number of variants passing %s: %s (%.2f%% of"
+        "Total number %s: %s (%.2f%% of"
         " non-synonymous variants, %.2f%% of total variants)",
         log_message,
         p_ns_variants,
@@ -195,29 +187,26 @@ def version_stats(
             )
             logger.info(
                 "Number of variants after filtering to AF threshold for %s at %s "
-                "threshold in %s: %s (%.2f%% of passing non-synonymous variants, %.2f%%"
+                "threshold in %s: %s ( %.2f%%"
                 " of total variants)",
                 grp_id,
                 threshold,
                 version,
                 t_ht.count(),
-                t_ht.count() / p_ns_variants * 100,
+                # t_ht.count() / p_ns_variants * 100,
                 t_ht.count() / t_variants * 100,
             )
             if version == "v4":
-                grpmax_ga_expr = t_ht.grpmax.gnomad.gen_anc
+                t_ht = t_ht.annotate(grpmax_ga=t_ht.grpmax.gnomad.gen_anc)
             else:
-                grpmax_ga_expr = t_ht.popmax[0].pop
+                t_ht = t_ht.annotate(grpmax_ga=t_ht.popmax[0].pop)
+
+            t_ht = t_ht.filter(DIVERSE_GRPS.contains(t_ht.grpmax_ga))
 
-            t_ht = t_ht.filter(DIVERSE_GRPS.contains(grpmax_ga_expr))
-            if version == "v4":
-                grpmax_ga_expr = t_ht.grpmax.gnomad.gen_anc
-            else:
-                grpmax_ga_expr = t_ht.popmax[0].pop
             # For each diverse genetic ancestry group, aggregate the number of
             # variants where that group is grpmax
             grpmax_by_thresholds[threshold] = t_ht.aggregate(
-                hl.agg.counter(grpmax_ga_expr)
+                hl.agg.counter(t_ht.grpmax_ga)
             )
             # Add a total count for all groups
             grpmax_by_thresholds[threshold]["all"] = t_ht.count()
@@ -268,7 +257,11 @@ def main(args):
         else:
             ht = get_gnomad_public_data("exomes")
         version_dict[version] = version_stats(
-            ht, version=version, can_only=args.canonical_only
+            ht,
+            version=version,
+            can_only=args.canonical_only,
+            drop_hard_filtered_variants=args.drop_hard_filtered_variants,
+            all_variants=args.all_variants,
         )
 
     # Create tables for "all_eur" and "nfe_only" data
@@ -288,5 +281,14 @@ def main(args):
         action="store_true",
         help="Drop variants that have been hard filtered",
     )
+    parser.add_argument(
+        "--all-variants",
+        action="store_true",
+        help="Keep all variants regardless of variant QC status",
+    )
     args = parser.parse_args()
+    if args.drop_hard_filtered_variants and args.all_variants:
+        raise ValueError(
+            "Cannot use --drop-hard-filtered-variants and --all-variants together"
+        )
     main(args)