[bugfix] Fix trio matrix (hail-is#5525)

* [bugfix] Fix trio matrix It was totally broken in the default case. * Fix
akotlar · Mar 6, 2019 · 52f3c12 · 52f3c12
1 parent 45df77b
commit 52f3c12
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 3 deletions.
diff --git a/hail/python/hail/methods/family_methods.py b/hail/python/hail/methods/family_methods.py
@@ -67,7 +67,7 @@ def trio_matrix(dataset, pedigree, complete_trios=False) -> MatrixTable:
     samples = mt[k].collect()
 
     pedigree = pedigree.filter_to(samples)
-    trios = pedigree.complete_trios() if complete_trios else pedigree.trios()
+    trios = pedigree.complete_trios() if complete_trios else pedigree.trios
     n_trios = len(trios)
 
     sample_idx = {}
@@ -76,8 +76,8 @@ def trio_matrix(dataset, pedigree, complete_trios=False) -> MatrixTable:
 
     trios = [hl.Struct(
         id=sample_idx[t.s],
-        pat_id=sample_idx[t.pat_id],
-        mat_id=sample_idx[t.mat_id],
+        pat_id=None if t.pat_id is None else sample_idx[t.pat_id],
+        mat_id=None if t.mat_id is None else sample_idx[t.mat_id],
         is_female=t.is_female,
         fam_id=t.fam_id) for t in trios]
     trios_type = hl.dtype('array<struct{id:int,pat_id:int,mat_id:int,is_female:bool,fam_id:str}>')

diff --git a/hail/python/test/hail/methods/test_family_methods.py b/hail/python/test/hail/methods/test_family_methods.py
@@ -93,6 +93,12 @@ def test_trio_matrix_null_keys(self):
         tt = hl.trio_matrix(mt, ped, complete_trios=True)
         self.assertEqual(tt.count_cols(), 0)
 
+    def test_trio_matrix_incomplete_trios(self):
+        ped = hl.Pedigree.read(resource('triomatrix.fam'))
+        mt = hl.import_vcf(resource('triomatrix.vcf'))
+        hl.trio_matrix(mt, ped, complete_trios=False)
+
+
     def test_mendel_errors(self):
         mt = hl.import_vcf(resource('mendel.vcf'))
         ped = hl.Pedigree.read(resource('mendel.fam'))