Filter training projects (#670)

Author: bpblanken

v03_pipeline/lib/annotations/constants.py

@@ -0,0 +1,5 @@
+PROJECTS_EXCLUDED_FROM_GT_STATS = {
+    # TODO get final list of projects
+    'R0607_gregor_training_project_',
+    'R0610_gregor_training_project_',
+}

v03_pipeline/lib/annotations/mito.py

@@ -2,6 +2,7 @@
 
 import hail as hl
 
+from v03_pipeline.lib.annotations.constants import PROJECTS_EXCLUDED_FROM_GT_STATS
 from v03_pipeline.lib.annotations.enums import MITOTIP_PATHOGENICITIES
 
 MITOTIP_PATHOGENICITIES_LOOKUP = hl.dict(
@@ -14,31 +15,6 @@
 )
 
 
-def _AC_het(row: hl.StructExpression) -> hl.Int32Expression:  # noqa: N802
-    return sum(
-        row.heteroplasmic_samples[project_guid].length()
-        for project_guid in row.heteroplasmic_samples
-    )
-
-
-def _AC_hom(row: hl.StructExpression) -> hl.Int32Expression:  # noqa: N802
-    return sum(
-        row.homoplasmic_samples[project_guid].length()
-        for project_guid in row.homoplasmic_samples
-    )
-
-
-def _AN(row: hl.StructExpression) -> hl.Int32Expression:  # noqa: N802
-    return sum(
-        (
-            row.ref_samples[project_guid].length()
-            + row.heteroplasmic_samples[project_guid].length()
-            + row.homoplasmic_samples[project_guid].length()
-        )
-        for project_guid in row.ref_samples
-    )
-
-
 def common_low_heteroplasmy(ht: hl.Table, **_: Any) -> hl.Expression:
     return ht.common_low_heteroplasmy
 
@@ -90,12 +66,26 @@ def rsid(ht: hl.Table, **_: Any) -> hl.Expression:
     return ht.rsid.find(lambda x: hl.is_defined(x))
 
 
-def gt_stats(ht: hl.Table, sample_lookup_ht: hl.Table, **_: Any) -> hl.Expression:
+def gt_stats(
+    ht: hl.Table,
+    sample_lookup_ht: hl.Table,
+    **_: Any,
+) -> hl.Expression:
     row = sample_lookup_ht[ht.key]
+    AC_het, AC_hom, AN = 0, 0, 0  # noqa: N806
+    for project_guid in row.ref_samples:
+        if project_guid in PROJECTS_EXCLUDED_FROM_GT_STATS:
+            continue
+        ref_samples_length = row.ref_samples[project_guid].length()
+        heteroplasmic_samples_length = row.heteroplasmic_samples[project_guid].length()
+        homoplasmic_samples_length = row.homoplasmic_samples[project_guid].length()
+        AC_het += heteroplasmic_samples_length # noqa: N806
+        AC_hom += homoplasmic_samples_length # noqa: N806
+        AN += (ref_samples_length + heteroplasmic_samples_length + homoplasmic_samples_length) # noqa: N806
     return hl.Struct(
-        AC_het=_AC_het(row),
-        AF_het=hl.float32(_AC_het(row) / _AN(row)),
-        AC_hom=_AC_hom(row),
-        AF_hom=hl.float32(_AC_hom(row) / _AN(row)),
-        AN=_AN(row),
+        AC_het=AC_het,
+        AF_het=hl.float32(AC_het / AN),
+        AC_hom=AC_hom,
+        AF_hom=hl.float32(AC_hom / AN),
+        AN=AN,
     )

v03_pipeline/lib/annotations/snv_indel.py

@@ -3,6 +3,8 @@
 
 import hail as hl
 
+from v03_pipeline.lib.annotations.constants import PROJECTS_EXCLUDED_FROM_GT_STATS
+
 N_ALT_REF = 0
 N_ALT_HET = 1
 N_ALT_HOM = 2
@@ -37,6 +39,8 @@ def gt_stats(
     row = sample_lookup_ht[ht.key]
     AC, AN, hom = 0, 0, 0
     for project_guid in row.ref_samples:
+        if project_guid in PROJECTS_EXCLUDED_FROM_GT_STATS:
+            continue
         ref_samples_length = row.ref_samples[project_guid].length()
         het_samples_length = row.het_samples[project_guid].length()
         hom_samples_length = row.hom_samples[project_guid].length()

v03_pipeline/lib/annotations/snv_indel_test.py

@@ -25,33 +25,37 @@ def test_allele_count_annotations(self) -> None:
             [
                 {
                     'id': 0,
-                    'ref_samples': hl.Struct(project_1={'a', 'c'}, project_2=set()),
-                    'het_samples': hl.Struct(project_1={'b', 'd'}, project_2=set()),
-                    'hom_samples': hl.Struct(project_1={'e', 'f'}, project_2=set()),
+                    'ref_samples': hl.Struct(project_1={'a', 'c'}, project_2=set(), R0607_gregor_training_project_=set()),
+                    'het_samples': hl.Struct(project_1={'b', 'd'}, project_2=set(), R0607_gregor_training_project_=set()),
+                    'hom_samples': hl.Struct(project_1={'e', 'f'}, project_2=set(), R0607_gregor_training_project_={'l', 'm'}),
                 },
                 {
                     'id': 1,
                     'ref_samples': hl.Struct(
                         project_1={'a', 'b', 'c', 'd', 'e', 'f'},
                         project_2=set(),
+                        R0607_gregor_training_project_={'l', 'm'},
                     ),
-                    'het_samples': hl.Struct(project_1=set(), project_2=set()),
-                    'hom_samples': hl.Struct(project_1=set(), project_2=set()),
+                    'het_samples': hl.Struct(project_1=set(), project_2=set(), R0607_gregor_training_project_=set()),
+                    'hom_samples': hl.Struct(project_1=set(), project_2=set(), R0607_gregor_training_project_=set()),
                 },
             ],
             hl.tstruct(
                 id=hl.tint32,
                 ref_samples=hl.tstruct(
                     project_1=hl.tset(hl.tstr),
                     project_2=hl.tset(hl.tstr),
+                    R0607_gregor_training_project_=hl.tset(hl.tstr),
                 ),
                 het_samples=hl.tstruct(
                     project_1=hl.tset(hl.tstr),
                     project_2=hl.tset(hl.tstr),
+                    R0607_gregor_training_project_=hl.tset(hl.tstr),
                 ),
                 hom_samples=hl.tstruct(
                     project_1=hl.tset(hl.tstr),
                     project_2=hl.tset(hl.tstr),
+                    R0607_gregor_training_project_=hl.tset(hl.tstr),
                 ),
             ),
             key='id',