merge

Author: bpblanken

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)

v03_pipeline/lib/misc/family_loading_failures_test.py

@@ -3,13 +3,17 @@
 import hail as hl
 
 from v03_pipeline.lib.misc.family_loading_failures import (
+    all_relatedness_checks,
     build_relatedness_check_lookup,
     build_sex_check_lookup,
-    passes_all_relatedness_checks,
+    get_families_failed_sex_check,
 )
-from v03_pipeline.lib.misc.pedigree import Sample
+from v03_pipeline.lib.misc.io import import_pedigree
+from v03_pipeline.lib.misc.pedigree import Sample, parse_pedigree_ht_to_families
 from v03_pipeline.lib.model import Ploidy
 
+TEST_PEDIGREE_6 = 'v03_pipeline/var/test/pedigrees/test_pedigree_6.tsv'
+
 
 class FamilyLoadingFailuresTest(unittest.TestCase):
     def test_build_relatedness_check_lookup(self):
@@ -40,7 +44,7 @@ def test_build_relatedness_check_lookup(self):
                 hl.dict({'ROS_006_18Y03226_D1': 'remapped_id'}),
             ),
             {
-                ('remapped_id', 'ROS_007_19Y05939_D1'): [
+                ('ROS_007_19Y05939_D1', 'remapped_id'): [
                     0.0,
                     1.0,
                     0.0,
@@ -77,7 +81,7 @@ def test_build_sex_check_lookup(self):
             },
         )
 
-    def test_passes_all_relatedness_checks(self):
+    def test_all_relatedness_checks(self):
         relatedness_check_lookup = {
             # Parent
             ('sample_1', 'sample_2'): [
@@ -98,9 +102,8 @@ def test_passes_all_relatedness_checks(self):
             paternal_grandfather='sample_3',
             half_siblings=['sample_4'],
         )
-        self.assertTrue(
-            passes_all_relatedness_checks(relatedness_check_lookup, sample),
-        )
+        failure_reasons = all_relatedness_checks(relatedness_check_lookup, sample)
+        self.assertListEqual(failure_reasons, [])
 
         # Defined grandparent missing in relatedness table
         sample = Sample(
@@ -110,8 +113,15 @@ def test_passes_all_relatedness_checks(self):
             paternal_grandfather='sample_3',
             paternal_grandmother='sample_5',
         )
-        self.assertFalse(
-            passes_all_relatedness_checks(relatedness_check_lookup, sample),
+        failure_reasons = all_relatedness_checks(
+            relatedness_check_lookup,
+            sample,
+        )
+        self.assertListEqual(
+            failure_reasons,
+            [
+                'Sample sample_1 has expected relation "grandparent" to sample_5 but has coefficients []',
+            ],
         )
 
         # Sibling is actually a half sibling.
@@ -126,6 +136,73 @@ def test_passes_all_relatedness_checks(self):
             paternal_grandfather='sample_3',
             siblings=['sample_4'],
         )
-        self.assertFalse(
-            passes_all_relatedness_checks(relatedness_check_lookup, sample),
+        failure_reasons = all_relatedness_checks(
+            relatedness_check_lookup,
+            sample,
+        )
+        self.assertListEqual(
+            failure_reasons,
+            [
+                'Sample sample_1 has expected relation "sibling" to sample_4 but has coefficients [0.5, 0.5, 0, 0.25]',
+            ],
+        )
+
+        relatedness_check_lookup = {
+            **relatedness_check_lookup,
+            ('sample_1', 'sample_2'): [
+                0.5,
+                0.5,
+                0.5,
+                0.5,
+            ],
+        }
+        sample = Sample(
+            sex=Ploidy.FEMALE,
+            sample_id='sample_1',
+            mother='sample_2',
+            paternal_grandfather='sample_3',
+            siblings=['sample_4'],
+        )
+        failure_reasons = all_relatedness_checks(
+            relatedness_check_lookup,
+            sample,
+        )
+        self.assertListEqual(
+            failure_reasons,
+            [
+                'Sample sample_1 has expected relation "parent" to sample_2 but has coefficients [0.5, 0.5, 0.5, 0.5]',
+                'Sample sample_1 has expected relation "sibling" to sample_4 but has coefficients [0.5, 0.5, 0, 0.25]',
+            ],
+        )
+
+    def test_get_families_failed_sex_check(self):
+        sex_check_ht = hl.Table.parallelize(
+            [
+                {'s': 'ROS_006_18Y03226_D1', 'sex': 'M'},
+                {'s': 'ROS_006_18Y03227_D1', 'sex': 'F'},
+                {'s': 'ROS_006_18Y03228_D1', 'sex': 'F'},
+                {'s': 'ROS_007_19Y05919_D1', 'sex': 'F'},
+                {'s': 'ROS_007_19Y05939_D1', 'sex': 'F'},
+                {'s': 'ROS_007_19Y05987_D1', 'sex': 'F'},
+            ],
+            hl.tstruct(
+                s=hl.tstr,
+                sex=hl.tstr,
+            ),
+            key='s',
+        )
+        pedigree_ht = import_pedigree(TEST_PEDIGREE_6)
+        failed_families = get_families_failed_sex_check(
+            parse_pedigree_ht_to_families(pedigree_ht),
+            sex_check_ht,
+            {},
+        )
+        self.assertCountEqual(
+            failed_families.values(),
+            [
+                [
+                    'Sample ROS_006_18Y03226_D1 has pedigree sex F but imputed sex M',
+                    'Sample ROS_006_18Y03227_D1 has pedigree sex M but imputed sex F',
+                ],
+            ],
         )

v03_pipeline/lib/misc/io.py

@@ -146,7 +146,12 @@ def select_relevant_fields(
     dataset_type: DatasetType,
 ) -> hl.MatrixTable:
     mt = mt.select_globals()
-    mt = mt.select_rows(*dataset_type.row_fields)
+    optional_row_fields = [
+        row_field
+        for row_field in dataset_type.optional_row_fields
+        if hasattr(mt, row_field)
+    ]
+    mt = mt.select_rows(*dataset_type.row_fields, *optional_row_fields)
     mt = mt.select_cols(*dataset_type.col_fields)
     return mt.select_entries(*dataset_type.entries_fields)