handle wes/wgs inheritance edge case

hail_search/definitions.py

@@ -13,10 +13,10 @@ def family_entries_field(self) -> str:
         }[self]
 
     @property
-    def passes_inheritance_field(self) -> str:
+    def failed_family_sample_field(self) -> str:
         return {
-            SampleType.WES: 'wes_passes_inheritance',
-            SampleType.WGS: 'wgs_passes_inheritance',
+            SampleType.WES: 'wes_failed_family_sample_guids',
+            SampleType.WGS: 'wgs_failed_family_sample_guids',
         }[self]
 
     @property

hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/README.txt

@@ -1,3 +1,3 @@
 This folder comprises a Hail (www.hail.is) native Table or MatrixTable.
-  Written with version 0.2.130-bea04d9c79b5
-  Created at 2024/10/02 14:46:35
+  Written with version 0.2.132-678e1f52b999
+  Created at 2024/10/28 16:21:30

hail_search/queries/base.py

@@ -473,7 +473,7 @@ def _filter_single_entries_table(self, ht, project_families, inheritance_filter=
         ht, sorted_family_sample_data = self._add_entry_sample_families(ht, project_families, is_merged_ht)
         ht = self._filter_quality(ht, quality_filter, **kwargs)
         ht, ch_ht = self._filter_inheritance(
-            ht, None, inheritance_filter, sorted_family_sample_data,
+            ht, None, inheritance_filter, sorted_family_sample_data, self._annotate_families_inheritance
         )
         ht = self._apply_entry_filters(ht)
         ch_ht = self._apply_entry_filters(ch_ht)
@@ -573,7 +573,7 @@ def _get_sample_type(cls, family_index, ht_globals):
         return ht_globals.sample_type
 
     def _filter_inheritance(
-        self, ht, comp_het_ht, inheritance_filter, sorted_family_sample_data,
+        self, ht, comp_het_ht, inheritance_filter, sorted_family_sample_data, annotate_func,
         annotation='family_entries', entries_ht_field='family_entries'
     ):
         any_valid_entry = lambda x: self.GENOTYPE_QUERY_MAP[HAS_ALT](x.GT)
@@ -584,12 +584,12 @@ def _filter_inheritance(
             any_valid_entry = lambda x: prev_any_valid_entry(x) & (x.affected_id == AFFECTED_ID)
 
         ht = ht.annotate(**{
-            annotation: ht[entries_ht_field].map(
+            entries_ht_field: ht[entries_ht_field].map(
                 lambda entries: hl.or_missing(entries.any(any_valid_entry), entries)
             )})
 
         if self._has_comp_het_search:
-            comp_het_ht = self._annotate_families_inheritance(
+            comp_het_ht = annotate_func(
                 comp_het_ht if comp_het_ht is not None else ht, COMPOUND_HET, inheritance_filter,
                 sorted_family_sample_data, annotation, entries_ht_field
             )
@@ -598,7 +598,7 @@ def _filter_inheritance(
             # No sample-specific inheritance filtering needed
             sorted_family_sample_data = []
 
-        ht = None if self._inheritance_mode == COMPOUND_HET else self._annotate_families_inheritance(
+        ht = None if self._inheritance_mode == COMPOUND_HET else annotate_func(
             ht, self._inheritance_mode, inheritance_filter, sorted_family_sample_data,
             annotation, entries_ht_field
         )
@@ -609,6 +609,27 @@ def _annotate_families_inheritance(
         self, ht, inheritance_mode, inheritance_filter, sorted_family_sample_data,
         annotation, entries_ht_field,
     ):
+        entry_indices_by_gt = self._get_entry_indices_by_gt_map(
+            inheritance_filter, inheritance_mode, sorted_family_sample_data
+        )
+
+        for genotype, entry_indices in entry_indices_by_gt.items():
+            if not entry_indices:
+                continue
+            entry_indices = hl.dict(entry_indices)
+            ht = ht.annotate(**{
+                annotation: hl.enumerate(ht[entries_ht_field]).starmap(
+                    lambda family_i, family_samples: hl.or_missing(
+                        ~entry_indices.contains(family_i) | entry_indices[family_i].all(
+                            lambda sample_i: self.GENOTYPE_QUERY_MAP[genotype](family_samples[sample_i].GT)
+                        ), family_samples,
+                    ),
+                )
+            })
+
+        return ht
+
+    def _get_entry_indices_by_gt_map(self, inheritance_filter, inheritance_mode, sorted_family_sample_data):
         individual_genotype_filter = (inheritance_filter or {}).get('genotype')
 
         # Create a mapping of genotypes to check against a list of samples for a family
@@ -630,21 +651,7 @@ def _annotate_families_inheritance(
             ]
             self.max_unaffected_samples = max(family_unaffected_counts) if family_unaffected_counts else 0
 
-        for genotype, entry_indices in entry_indices_by_gt.items():
-            if not entry_indices:
-                continue
-            entry_indices = hl.dict(entry_indices)
-            ht = ht.annotate(**{
-                annotation: hl.enumerate(ht[entries_ht_field]).starmap(
-                    lambda family_i, family_samples: hl.or_missing(
-                        ~entry_indices.contains(family_i) | entry_indices[family_i].all(
-                            lambda sample_i: self.GENOTYPE_QUERY_MAP[genotype](family_samples[sample_i].GT)
-                        ), family_samples,
-                    ),
-                )
-            })
-
-        return ht
+        return entry_indices_by_gt
 
     def _get_family_passes_quality_filter(self, quality_filter, ht, **kwargs):
         quality_filter = quality_filter or {}

hail_search/queries/mito.py

@@ -209,34 +209,55 @@ def _filter_entries_ht_both_sample_types(
         for sample_type, sorted_family_sample_data in sample_types:
             ht, ch_ht = self._filter_inheritance(
                 ht, ch_ht, inheritance_filter, sorted_family_sample_data,
-                annotation=sample_type.passes_inheritance_field, entries_ht_field=sample_type.family_entries_field
+                annotate_func=self._annotate_failed_family_samples_inheritance,
+                annotation=sample_type.failed_family_sample_field, entries_ht_field=sample_type.family_entries_field,
             )
             for family_idx, samples in enumerate(sorted_family_sample_data):
                 family_guid = samples[0]['familyGuid']
                 family_guid_idx_map[family_guid][sample_type.value] = family_idx
 
-        family_idx_map = hl.dict(family_guid_idx_map)
-        ht = self._apply_multi_sample_type_entry_filters(ht, family_idx_map)
-        ch_ht = self._apply_multi_sample_type_entry_filters(ch_ht, family_idx_map)
+        family_guid_idx_map = hl.dict(family_guid_idx_map)
+        ht = self._apply_multi_sample_type_entry_filters(ht, family_guid_idx_map)
+        ch_ht = self._apply_multi_sample_type_entry_filters(ch_ht, family_guid_idx_map)
         return ht, ch_ht
 
+    def _annotate_failed_family_samples_inheritance(
+        self, ht, inheritance_mode, inheritance_filter, sorted_family_sample_data, annotation, entries_ht_field
+    ):
+        entry_indices_by_gt = self._get_entry_indices_by_gt_map(
+           inheritance_filter, inheritance_mode, sorted_family_sample_data
+        )
+
+        if ht is None:
+            return ht
+
+        # Initialize empty array
+        ht = ht.annotate(**{annotation: ht[entries_ht_field].map(lambda x: hl.empty_array(hl.tstr))})
+
+        # Add failed genotype samples
+        for genotype, entry_indices in entry_indices_by_gt.items():
+            if not entry_indices:
+                continue
+
+            entry_indices = hl.dict(entry_indices)
+            ht = ht.annotate(**{annotation: hl.enumerate(ht[entries_ht_field]).starmap(
+                lambda family_idx, entries: hl.bind(
+                    lambda failed_samples: ht[annotation][family_idx].extend(failed_samples),
+                    entry_indices.get(family_idx).filter(
+                        lambda sample_idx: ~self.GENOTYPE_QUERY_MAP[genotype](entries[sample_idx].GT)
+                    ).map(lambda sample_idx: entries[sample_idx]['sampleId'])
+                )
+            )})
+
+        return ht
+
     def _apply_multi_sample_type_entry_filters(self, ht, family_idx_map):
         if ht is None:
             return ht
 
-        # Keep family from both sample types if either passes quality AND inheritance
         for sample_type in SampleType:
-            ht = ht.annotate(**{
-                sample_type.family_entries_field: hl.enumerate(ht[sample_type.family_entries_field]).starmap(
-                    lambda i, family_samples: hl.or_missing(
-                        hl.bind(
-                            lambda other_sample_type_idx: (
-                                self._family_has_valid_sample_type_entries(ht, sample_type, i) |
-                                self._family_has_valid_sample_type_entries(ht, sample_type.other_sample_type, other_sample_type_idx)
-                            ),
-                            family_idx_map.get(hl.coalesce(family_samples)[0]['familyGuid']).get(sample_type.other_sample_type.value),
-                       ), family_samples)
-                )})
+            ht = self._apply_quality_entry_filters(ht, sample_type, family_idx_map)
+            ht = self._apply_inheritance_entry_filters(ht, sample_type, family_idx_map)
 
         # Merge family entries and filters from both sample types
         ht = ht.transmute(
@@ -252,15 +273,49 @@ def _apply_multi_sample_type_entry_filters(self, ht, family_idx_map):
         # Filter out families with no valid entries in either sample type
         return ht.filter(ht.family_entries.any(hl.is_defined))
 
+    def _apply_quality_entry_filters(self, ht, sample_type, family_idx_map):
+        return ht.annotate(**{
+            sample_type.family_entries_field: hl.enumerate(ht[sample_type.family_entries_field]).starmap(
+                lambda i, family_samples: hl.or_missing(
+                    hl.bind(lambda other_sample_type_idx: (
+                        self._family_has_valid_quality(ht, sample_type, i) |
+                        self._family_has_valid_quality(ht, sample_type.other_sample_type, other_sample_type_idx)
+                    ), family_idx_map.get(hl.coalesce(family_samples)[0]['familyGuid']).get(sample_type.other_sample_type.value),
+                ), family_samples)
+            )})
+
     @staticmethod
-    def _family_has_valid_sample_type_entries(ht, sample_type, sample_type_family_idx):
-        # Note: This logic does not sufficiently handle case 2 here https://docs.google.com/presentation/d/1hqDV8ulhviUcR5C4PtNUqkCLXKDsc6pccgFVlFmWUAU/edit?usp=sharing
-        # and will need to be changed to support it - https://github.com/broadinstitute/seqr/issues/4403
+    def _family_has_valid_quality(ht, sample_type, sample_type_family_idx):
         return (
             hl.is_defined(sample_type_family_idx) &
-            hl.is_defined(ht[sample_type.passes_quality_field][sample_type_family_idx]) &
-            hl.is_defined(ht[sample_type.passes_inheritance_field][sample_type_family_idx])
+            hl.is_defined(ht[sample_type.passes_quality_field][sample_type_family_idx])
+        )
+
+    @staticmethod
+    def _apply_inheritance_entry_filters(ht, sample_type, family_idx_map):
+        ht = ht.annotate(
+            **{sample_type.family_entries_field: hl.enumerate(ht[sample_type.family_entries_field]).starmap(
+                lambda family_idx, family_samples: hl.or_missing(
+                    hl.bind(lambda other_sample_type_family_idx: (
+                        hl.bind(
+                            lambda other_sample_type_pass_samples, sample_type_pass_samples: (
+                                ht[sample_type.failed_family_sample_field][family_idx].all(
+                                    other_sample_type_pass_samples.contains
+                                ) & ht[sample_type.other_sample_type.failed_family_sample_field][other_sample_type_family_idx].all(
+                                    sample_type_pass_samples.contains
+                                )),
+                            ht[sample_type.other_sample_type.family_entries_field][other_sample_type_family_idx].filter(
+                                lambda s: ~ht[sample_type.other_sample_type.failed_family_sample_field][other_sample_type_family_idx].contains(s['sampleId'])
+                            ).map(lambda s: s['sampleId']),
+                            ht[sample_type.other_sample_type.family_entries_field][family_idx].filter(
+                                lambda s: ~ht[sample_type.failed_family_sample_field][family_idx].contains(s['sampleId'])
+                            ).map(lambda s: s['sampleId']),
+                        )
+                    ), family_idx_map.get(hl.coalesce(family_samples)[0]['familyGuid']).get(sample_type.other_sample_type.value)
+                ), family_samples)
+            )}
         )
+        return ht
 
     def _get_sample_genotype(self, samples, r=None, include_genotype_overrides=False, select_fields=None, **kwargs):
         if not self._has_both_sample_types:

hail_search/test_search.py

@@ -388,10 +388,11 @@ async def test_both_sample_types_search(self):
             [VARIANT2_BOTH_SAMPLE_TYPES], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS,
             inheritance_mode=inheritance_mode, **COMP_HET_ALL_PASS_FILTERS, intervals=[variant2_interval]
         )
-        # Genome passes quality and inheritance exome fails inheritance (parental data shows variant is inherited).
+        # Genome passes quality and inheritance but exome fails inheritance (parental data shows variant is inherited).
+        # Variant is excluded from search results.
         inheritance_mode = 'de_novo'
         await self._assert_expected_search(
-            [VARIANT2_BOTH_SAMPLE_TYPES], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS,
+            [], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS,
             inheritance_mode=inheritance_mode, intervals=[variant2_interval]
         )
 

hail_search/test_utils.py

@@ -165,7 +165,7 @@
 VARIANT1_BOTH_SAMPLE_TYPES['genotypes'] = {
     'I000004_hg00731': [
         genotypes['I000004_hg00731'],
-        {**genotypes['I000004_hg00731'], 'sampleType': 'WGS'}
+        {**genotypes['I000004_hg00731'], 'numAlt': 2, 'sampleType': 'WGS'}
     ],
     'I000005_hg00732': [
         genotypes['I000005_hg00732'],