Merge branch 'benb/lookup_table_refactor' of github.com:broadinstitute/seqr-loading-pipelines into benb/migration2

Author: bpblanken

v03_pipeline/lib/methods/relatedness.py

@@ -30,7 +30,7 @@ def call_relatedness(
     # - brew install lz4
     # - CXXFLAGS='-I/opt/homebrew/include/' HAIL_COMPILE_NATIVES=1 make -C hail install
     # Hail issue here: https://discuss.hail.is/t/noclassdeffounderror-could-not-initialize-class-is-hail-methods-ibsffi/2453
-    kin_ht = hl.identity_by_descent(mt, maf=mt.info.AF[1], min=0.10, max=1.0)
+    kin_ht = hl.identity_by_descent(mt, maf=mt.info.AF[0], min=0.10, max=1.0)
     kin_ht = kin_ht.key_by('i', 'j')
     return kin_ht.select(
         ibd0=kin_ht.ibd.Z0,

v03_pipeline/lib/misc/family_entries.py

@@ -55,16 +55,16 @@ def compute_callset_family_entries_ht(
 
 
 def globalize_ids(ht: hl.Table) -> hl.Table:
-    row = ht.take(1)
-    has_family_entries = row and len(row[0].family_entries) > 0
+    row = ht.take(1)[0] if ht.count() > 0 else None
+    has_family_entries = row and len(row.family_entries) > 0
     ht = ht.annotate_globals(
         family_guids=(
-            [fe[0].family_guid for fe in row[0].family_entries]
+            [fe[0].family_guid for fe in row.family_entries]
             if has_family_entries
             else hl.empty_array(hl.tstr)
         ),
         family_samples=(
-            {fe[0].family_guid: [e.s for e in fe] for fe in row[0].family_entries}
+            {fe[0].family_guid: [e.s for e in fe] for fe in row.family_entries}
             if has_family_entries
             else hl.empty_dict(hl.tstr, hl.tarray(hl.tstr))
         ),
@@ -93,11 +93,10 @@ def deglobalize_ids(ht: hl.Table) -> hl.Table:
     return ht.drop('family_guids', 'family_samples')
 
 
-def remove_new_callset_family_guids(
+def remove_family_guids(
     ht: hl.Table,
-    family_guids: list[str],
+    family_guids: hl.SetExpression,
 ) -> hl.Table:
-    family_guids = hl.set(family_guids)
     # Remove families from the existing project table structure (both the entries arrays and the globals are mutated)
     family_indexes_to_keep = hl.array(
         hl.enumerate(ht.globals.family_guids)
@@ -142,7 +141,7 @@ def join_family_entries_hts(ht: hl.Table, callset_ht: hl.Table) -> hl.Table:
             .default(ht.family_entries.extend(ht.family_entries_1))
         ),
     )
-    # NB: transume because we want to drop the *_1 fields, but preserve other globals
+    # NB: transmute because we want to drop the *_1 fields, but preserve other globals
     return ht.transmute_globals(
         family_guids=ht.family_guids.extend(ht.family_guids_1),
         family_samples=hl.dict(

v03_pipeline/lib/misc/family_entries_test.py

@@ -7,7 +7,7 @@
     deglobalize_ids,
     globalize_ids,
     join_family_entries_hts,
-    remove_new_callset_family_guids,
+    remove_family_guids,
 )
 from v03_pipeline.lib.model import DatasetType
 
@@ -202,7 +202,7 @@ def test_globalize_and_deglobalize(self) -> None:
             ],
         )
 
-    def test_remove_new_callset_family_guids(self) -> None:
+    def test_remove_family_guids(self) -> None:
         family_entries_ht = hl.Table.parallelize(
             [
                 {
@@ -244,7 +244,10 @@ def test_remove_new_callset_family_guids(self) -> None:
                 },
             ),
         )
-        family_entries_ht = remove_new_callset_family_guids(family_entries_ht, ['012'])
+        family_entries_ht = remove_family_guids(
+            family_entries_ht,
+            hl.set(['012']),
+        )
         self.assertCountEqual(
             family_entries_ht.globals.collect(),
             [
@@ -326,7 +329,7 @@ def test_remove_new_callset_family_guids_all_families(self) -> None:
                 },
             ),
         )
-        ht = remove_new_callset_family_guids(family_entries_ht, ['012', '123'])
+        ht = remove_family_guids(family_entries_ht, hl.set(['012', '123']))
         self.assertCountEqual(
             ht.globals.collect(),
             [hl.Struct(family_guids=[], family_samples={})],

v03_pipeline/lib/misc/family_loading_failures.py

@@ -1,3 +1,5 @@
+from collections import defaultdict
+
 import hail as hl
 import numpy as np
 
@@ -10,83 +12,93 @@ def passes_relatedness_check(
     sample_id: str,
     other_id: str,
     relation: Relation,
-) -> bool:
+) -> tuple[bool, str | None]:
     # No relationship to check, return true
     if other_id is None:
-        return True
+        return True, None
     coefficients = relatedness_check_lookup.get(
         (min(sample_id, other_id), max(sample_id, other_id)),
     )
-    if not coefficients:
-        return False
-    return np.allclose(
+    if not coefficients or not np.allclose(
         coefficients,
         relation.coefficients,
         0.1,
-    )
+    ):
+        return (
+            False,
+            f'Sample {sample_id} has expected relation "{relation.value}" to {other_id} but has coefficients {coefficients or []}',
+        )
+    return True, None
 
 
-def passes_all_relatedness_checks(  # noqa: C901
+def all_relatedness_checks(  # noqa: C901
     relatedness_check_lookup: dict[tuple[str, str], list],
     sample: Sample,
-) -> bool:
+) -> list[str]:
+    failure_reasons = []
     for parent_id in [sample.mother, sample.father]:
-        if not passes_relatedness_check(
+        success, reason = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             parent_id,
             Relation.PARENT,
-        ):
-            return False
+        )
+        if not success:
+            failure_reasons.append(reason)
 
     for grandparent_id in [
         sample.maternal_grandmother,
         sample.maternal_grandfather,
         sample.paternal_grandmother,
         sample.paternal_grandfather,
     ]:
-        if not passes_relatedness_check(
+        success, reason = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             grandparent_id,
             Relation.GRANDPARENT,
-        ):
-            return False
+        )
+        if not success:
+            failure_reasons.append(reason)
 
     for sibling_id in sample.siblings:
-        if not passes_relatedness_check(
+        success, reason = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             sibling_id,
             Relation.SIBLING,
-        ):
-            return False
+        )
+        if not success:
+            failure_reasons.append(reason)
 
     for half_sibling_id in sample.half_siblings:
         # NB: A "half sibling" parsed from the pedigree may actually be a sibling, so we allow those
         # through as well.
-        if not passes_relatedness_check(
+        success1, _ = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             half_sibling_id,
-            Relation.HALF_SIBLING,
-        ) and not passes_relatedness_check(
+            Relation.SIBLING,
+        )
+        success2, reason = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             half_sibling_id,
-            Relation.SIBLING,
-        ):
-            return False
+            Relation.HALF_SIBLING,
+        )
+        if not success1 and not success2:
+            failure_reasons.append(reason)
 
     for aunt_nephew_id in sample.aunt_nephews:
-        if not passes_relatedness_check(
+        success, reason = passes_relatedness_check(
             relatedness_check_lookup,
             sample.sample_id,
             aunt_nephew_id,
             Relation.AUNT_NEPHEW,
-        ):
-            return False
-    return True
+        )
+        if not success:
+            failure_reasons.append(reason)
+    return failure_reasons
 
 
 def build_relatedness_check_lookup(
@@ -99,7 +111,9 @@ def build_relatedness_check_lookup(
         j=remap_lookup.get(relatedness_check_ht.j, relatedness_check_ht.j),
     )
     return {
-        (r.i, r.j): list(r.drop('i', 'j').values())
+        # NB: samples are sorted in the original ibd but not necessarily
+        # sorted after remapping
+        (min(r.i, r.j), max(r.i, r.j)): list(r.drop('i', 'j').values())
         for r in relatedness_check_ht.collect()
     }
 
@@ -119,43 +133,50 @@ def build_sex_check_lookup(
 def get_families_failed_missing_samples(
     mt: hl.MatrixTable,
     families: set[Family],
-) -> set[Family]:
+) -> dict[Family, list[str]]:
     callset_samples = set(mt.cols().s.collect())
-    failed_families = set()
+    failed_families = {}
     for family in families:
-        if len(family.samples.keys() - callset_samples) > 0:
-            failed_families.add(family)
+        missing_samples = family.samples.keys() - callset_samples
+        if len(missing_samples) > 0:
+            # NB: This is an array of a single element for consistency with
+            # the other checks.
+            failed_families[family] = [f'Missing samples: {missing_samples}']
     return failed_families
 
 
 def get_families_failed_relatedness_check(
     families: set[Family],
     relatedness_check_ht: hl.Table,
     remap_lookup: hl.dict,
-) -> set[Family]:
+) -> dict[Family, list[str]]:
     relatedness_check_lookup = build_relatedness_check_lookup(
         relatedness_check_ht,
         remap_lookup,
     )
-    failed_families = set()
+    failed_families = defaultdict(list)
     for family in families:
         for sample in family.samples.values():
-            if not passes_all_relatedness_checks(relatedness_check_lookup, sample):
-                failed_families.add(family)
-                break
-    return failed_families
+            failure_reasons = all_relatedness_checks(
+                relatedness_check_lookup,
+                sample,
+            )
+            if failure_reasons:
+                failed_families[family].extend(failure_reasons)
+    return dict(failed_families)
 
 
 def get_families_failed_sex_check(
     families: set[Family],
     sex_check_ht: hl.Table,
     remap_lookup: hl.dict,
-) -> set[Family]:
+) -> dict[Family, list[str]]:
     sex_check_lookup = build_sex_check_lookup(sex_check_ht, remap_lookup)
-    failed_families = set()
+    failed_families = defaultdict(list)
     for family in families:
         for sample_id in family.samples:
             if family.samples[sample_id].sex != sex_check_lookup[sample_id]:
-                failed_families.add(family)
-                break
-    return failed_families
+                failed_families[family].append(
+                    f'Sample {sample_id} has pedigree sex {family.samples[sample_id].sex.value} but imputed sex {sex_check_lookup[sample_id].value}',
+                )
+    return dict(failed_families)