Get pLOFs for each gene

gnomad_toolbox/filtering/vep.py

@@ -1,9 +1,21 @@
 """Functions to filter gnomAD sites HT by VEP annotations."""
 
+from typing import List
+
 import hail as hl
-from gnomad.utils.vep import LOF_CSQ_SET, filter_vep_transcript_csqs_expr
+from gnomad.resources.grch37.reference_data import gencode as grch37_gencode
+from gnomad.resources.grch38.reference_data import gencode as grch38_gencode
+from gnomad.utils.vep import (
+    LOF_CSQ_SET,
+    filter_vep_transcript_csqs,
+    filter_vep_transcript_csqs_expr,
+)
 
-from gnomad_toolbox.load_data import _get_gnomad_release
+from gnomad_toolbox.load_data import (
+    _get_gnomad_release,
+    get_coverage_for_variant,
+    gnomad_session,
+)
 
 
 # TODO: Check these csq sets, the ones in the code don't match what is listed on the
@@ -116,23 +128,102 @@ def filter_by_consequence_category(
     return ht.filter(filter_expr)
 
 
-# TODO: The following was in one of the notebooks, and I think we should add a wrapper
-#  around this function to make it much simpler instead of using it in the notebook.
+def get_gene_intervals(gene_symbol: str, version: str) -> List[hl.utils.Interval]:
+    """
+    Get the genomic intervals for a given gene symbol.
+
+    :param gene_symbol: Gene symbol.
+    :param version: Dataset version.
+    :return: List of intervals for the specified gene.
+    """
+    gen_ht = grch37_gencode.ht() if version.startswith("2.") else grch38_gencode.ht()
+    intervals = (
+        gen_ht.filter(
+            (gen_ht.feature == "gene")
+            & (gen_ht.gene_name.lower() == gene_symbol.lower())
+        )
+        .select()
+        .collect()
+    )
+    if not intervals:
+        raise ValueError(f"No interval found for gene: {gene_symbol}")
+    return [row["interval"] for row in intervals]
 
-# Filter to LOFTEE high-confidence variants for certain genes
 
-# In this example, we are filtering to variants in ASH1L that are LOFTEE high-confidence
-# (with no flags) in the MANE select transcript.
+def filter_hc_variants(var_ht: hl.Table, gene_symbol: str, version: str) -> hl.Table:
+    """
+    Filter variants to high-confidence LOFTEE variants with optional transcript selection.
+
+    :param var_ht: Variants Table.
+    :param gene_symbol: Gene symbol.
+    :param version: Dataset version.
+    :return: Filtered variants Hail Table.
+    """
+    return filter_vep_transcript_csqs(
+        var_ht,
+        synonymous=False,
+        mane_select=version.startswith("4."),
+        genes=[gene_symbol.upper()],
+        match_by_gene_symbol=True,
+        additional_filtering_criteria=[
+            lambda x: (x.lof == "HC")
+            & (hl.is_missing(x.lof_flags) | (x.lof_flags == ""))
+        ],
+    )
+
+
+def filter_to_plofs(
+    gene_symbol: str, select_fields: bool = False, **kwargs
+) -> hl.Table:
+    """
+    Filter to observed pLoF variants used for gene constraint metrics.
+
+    .. note::
+
+                pLOF variants meets the following requirements:
+                - High-confidence LOFTEE variants (without any flags),
+                - Only variants in the MANE Select transcript,
+                - PASS variants that are SNVs with MAF ≤ 0.1%,
+                - Exome median depth ≥ 30 (# TODO: This is changing in v4 constraint?)
+
+    :param gene_symbol: Gene symbol.
+    :param select_fields: Whether to limit the output to specific fields.
+    :return: Table with pLoF variants.
+    """
+    var_version = kwargs.pop("version", gnomad_session.version)
+    var_ht = _get_gnomad_release(dataset="variant", version=var_version, **kwargs)
+    cov_ht = get_coverage_for_variant(var_version, **kwargs)
+
+    # Get gene intervals and filter tables
+    intervals = get_gene_intervals(gene_symbol, var_version)
+    var_ht = hl.filter_intervals(var_ht, intervals)
+    cov_ht = hl.filter_intervals(cov_ht, intervals)
+
+    # Filter to high-confidence LOFTEE variants
+    var_ht = filter_hc_variants(var_ht, gene_symbol, var_version)
+
+    # Version-specific expressions
+    if var_version.startswith("2."):
+        allele_type_expr = var_ht.allele_type
+        cov_cut_expr = cov_ht[var_ht.locus].median
+    else:
+        allele_type_expr = var_ht.allele_info.allele_type
+        cov_cut_expr = cov_ht[var_ht.locus].median_approx
+
+    # Apply final filters
+    var_ht = var_ht.filter(
+        (hl.len(var_ht.filters) == 0)
+        & (allele_type_expr == "snv")
+        & (var_ht.freq[0].AF <= 0.001)
+        & (cov_cut_expr >= 30)
+    )
+
+    # Select specific fields if requested
+    if select_fields:
+        var_ht = var_ht.select(
+            freq=var_ht.freq[0],
+            csq=var_ht.vep.transcript_consequences[0].consequence_terms,
+            coverage=cov_ht[var_ht.locus],
+        )
 
-# from gnomad.utils.vep import filter_vep_transcript_csqs
-# ht = get_gnomad_release(data_type='exomes', version='4.1')
-# ht = filter_vep_transcript_csqs(
-#    ht,
-#    synonymous=False,
-#    mane_select=True,
-#    genes=["ASH1L"],
-#    match_by_gene_symbol=True,
-#    additional_filtering_criteria=[lambda x: (x.lof == "HC") & hl.is_missing(x.lof_flags)],
-# )
-# ht.show()
-# ht.count()
+    return var_ht

gnomad_toolbox/load_data.py

@@ -200,3 +200,22 @@ def get_gnomad_release(
     :return: Hail Table for requested dataset, data type, and version.
     """
     return _get_gnomad_release(dataset=dataset, data_type=data_type, version=version)
+
+
+def get_coverage_for_variant(variant_version: str, **kwargs) -> hl.Table:
+    """
+    Get the appropriate coverage table based on the provided version.
+
+    :param variant_version: Version of the variant dataset.
+    :return: Hail Table for the corresponding coverage dataset.
+    """
+    if variant_version.startswith("4."):
+        return _get_gnomad_release(dataset="coverage", version="4.0", **kwargs)
+    elif variant_version.startswith("3."):
+        return _get_gnomad_release(dataset="coverage", version="3.0.1", **kwargs)
+    elif variant_version.startswith("2."):
+        return _get_gnomad_release(dataset="coverage", version="2.1", **kwargs)
+    else:
+        raise ValueError(
+            f"Unrecognized version: '{variant_version}'. Please specify a valid version."
+        )