From e2293f36edc2dd8117ad7647f76d8195c975d3ee Mon Sep 17 00:00:00 2001
From: Julia Klugherz <juliaklugherz@gmail.com>
Date: Tue, 12 Nov 2024 10:30:36 -0500
Subject: [PATCH] handle wes/wgs inheritance edge case  (#4440)

* almost working

* updates

* or_missing

* fix sample tracking logic

* f string

* another pass

* minor things

* less code

* reuse as much code from _annotate_families_inheritance

* rstructure passes inheritance field

* fix _family_has_valid_inheritance

* fix up test cases

* oop

* typo

* typo 2

* PR comments

* comments

* extra lines
---
 .../families/WGS/F000002_2.ht/.README.txt.crc | Bin 12 -> 12 bytes
 .../families/WGS/F000002_2.ht/README.txt      |   2 +-
 .../.index.crc                                | Bin
 .../.metadata.json.gz.crc                     | Bin
 .../index                                     | Bin
 .../metadata.json.gz                          | Bin
 .../F000002_2.ht/rows/.metadata.json.gz.crc   | Bin 16 -> 16 bytes
 .../WGS/F000002_2.ht/rows/metadata.json.gz    | Bin 627 -> 628 bytes
 ...0-5efaaf66-a01a-4640-a3e4-aef656269368.crc | Bin 12 -> 0 bytes
 ...0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.crc | Bin 0 -> 12 bytes
 ...art-0-5efaaf66-a01a-4640-a3e4-aef656269368 | Bin 206 -> 0 bytes
 ...art-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac | Bin 0 -> 207 bytes
 hail_search/queries/base.py                   |  58 ++++++-----
 hail_search/queries/mito.py                   |  77 +++++++++++----
 hail_search/test_search.py                    |  40 ++++----
 hail_search/test_utils.py                     |  93 +++++++++++++++---
 16 files changed, 193 insertions(+), 77 deletions(-)
 rename hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/{part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx => part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx}/.index.crc (100%)
 rename hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/{part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx => part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx}/.metadata.json.gz.crc (100%)
 rename hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/{part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx => part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx}/index (100%)
 rename hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/{part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx => part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx}/metadata.json.gz (100%)
 delete mode 100644 hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-5efaaf66-a01a-4640-a3e4-aef656269368.crc
 create mode 100644 hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.crc
 delete mode 100644 hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-5efaaf66-a01a-4640-a3e4-aef656269368
 create mode 100644 hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/.README.txt.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/.README.txt.crc
index f7799b94778f02019a04657be1cc91ebeb101949..3b6f654d6c82ab84756a952e252b70193d3aaff9 100644
GIT binary patch
literal 12
TcmYc;N@ieSU}AVIzG@->6E*{I

literal 12
TcmYc;N@ieSU}6yY%H|FL5ds2{

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/README.txt b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/README.txt
index 3fea5d0da1..6eb9cb83f0 100644
--- a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/README.txt
+++ b/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
\ No newline at end of file
+  Created at 2024/11/04 13:45:23
\ No newline at end of file
diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/.index.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/.index.crc
similarity index 100%
rename from hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/.index.crc
rename to hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/.index.crc
diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/.metadata.json.gz.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/.metadata.json.gz.crc
similarity index 100%
rename from hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/.metadata.json.gz.crc
rename to hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/.metadata.json.gz.crc
diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/index b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/index
similarity index 100%
rename from hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/index
rename to hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/index
diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/metadata.json.gz b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/metadata.json.gz
similarity index 100%
rename from hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-5efaaf66-a01a-4640-a3e4-aef656269368.idx/metadata.json.gz
rename to hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/index/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.idx/metadata.json.gz
diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/.metadata.json.gz.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/.metadata.json.gz.crc
index caf4248bc39fd6678d852baffee083d06b11a820..6a1017e8ebf299656da8b2b4dee4cd373e563000 100644
GIT binary patch
literal 16
XcmYc;N@ieSU}8A2OZP$7&J9}uDVPRw

literal 16
XcmYc;N@ieSU}E_3?py2Xq+2ZjFBJy|

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/metadata.json.gz b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/metadata.json.gz
index c18fceb46be2b260ee2b9874e13c6452c7f4e278..eeeb78776570901cc864e22a18bd0d1ed5646f88 100644
GIT binary patch
delta 616
zcmV-u0+;>s1oQ-u7k@4Z5J7H0DMG4N-GcV86+*5(0dG?WIc8NB<-d2F?*vvVm3k=$
zl)M?wyf=Q4eI$}T1M$gRa(DrL|Lc0O1LhKote@aO{09<7Y#-x04MPYaCMyiFhyyc;
z#6K1lmWLhiNg2Q_Fc3z*O~*8bg;sl}k7+r-0kKGGu`!sCRDT_;T@(u~B8kXc>aLqu
zd2a+pO3}kZk-D-15o%Dee-Z$xm4Gr(?>GDDXN9fN*q`2n5!GJLy1pyW*)o4hR*BK^
zH`mcnaBKd*%<d=K?%-=lFl&a1LT)N@3(zH2@G3KNMkS#dpOVf5>*r)1_dx)vHT&D?
z*l`ErDr6D16@MYGrHqmG6EbKdJsNN!{$YaG?{~e^Wxr$1)1?=#Lt#Ll+ba|ELlMHb
zt)*dJV5#{Ogm+YL5x=_no~9c#1RA6kdz0_#dmacC>$Tu7nW=ETXzzxhs$7itAL!Y{
zih>{pT{59kMw;fUYdw<l{;Q<agf#4A1YBy?lGO>P7Jr*hVnF;bB@KU|Z{BC9w4VZb
z-a@Hk*ZBRzy0zJ-M$C{8H(l070gDtrPwGM+mJ?`r?^JnPcU$%4PCX~(KSDmEd2wj(
zJdx5FRMTlzNr)-Fqe|ZhbI{}2Z(i(<&AkEVF74VK>J03j%ZK)md#=rU0rPsE=X5Ab
z?*)JEsbSo}IND2Y)*nf%v4Rnn-6HTI^DuU0Tjd~~!o$ROI=${NsV*AWe>B025y34d
zAo&~X44uB?cmuClYKb|sw33i4+;cQ{N)5F`iA_xtgCs@`)moWeN%9JAp*YYd2><}3
CHzjxg

delta 615
zcmV-t0+{{u1oH%t7k`!zupl==DMG4NT|j%-3L!U{fVZiG9J4Bm^4~kocLJ44rC!Pb
zC2z(v?~Pw%ABkYkKpZj?KD>ZG|9!LAf!8GlSwF#nI1eO@ynT%CEDRxpn5;15MLu{s
zkvPYq!t$^Ke_RGI3p9k09}}0wuuyXEsbgBsX+T~i6yIo!3x6gL#x9Dv;*mgPCY0Sx
zjJ!7jBPrS8p-5dB0S^_(*f|M+P?AF#sP~)xa<an4$aN;SVZ;>e88`O@I$h>Z$tux0
z{^dFv3hwm(x9NPmwFlozf>|?6<YH5iTYxIDf>)WDGb#zy_>^?USU)H8xDNu5t=ZpA
zT&p{9tB^swEq@RFT8fLbpO8T#=~06V{vQ))zi-piWq)Gjr%Nwfhr)n8w^t_mha$Lk
zTT8>Nz*6%`2=AHNB7Svsou+Fv1Qw(gdzbI(dmeBZt2OsuGE?Dv(cTS1Rk;{(KGCy@
z5d}dEs${~1h!o3L*Lo!9{Z~n;32E3#3%Jy*C94xoEq^v&#DMr=3KsrE-+ataX+H_#
zyoFN7uJQYab!)RvjhG>yZn~_E1TT{Qc~Tepu$(}{d#B3Vy4$KRcj`GQZwUE{=Eb3X
z@I*>yP)(;<B_X=_j!AXP^+AtkzkV@0W*-8^0!mHRv6$JTy^hIj=$Q-x+NWJQwCNzd
z7yO+wv0wo=+6!N=KayBs1uZPQMc_l`VeIN{rG<0~4-?1g^z30$T{N)&Xp9*nf?G~N
z`0uPUwEC7s2eesgi8<4>5*J>$_tD%bHPj9zHqlKql4vzdDXDuU$tx--88asd008!1
BAaVcz

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-5efaaf66-a01a-4640-a3e4-aef656269368.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-5efaaf66-a01a-4640-a3e4-aef656269368.crc
deleted file mode 100644
index d4dcde073f96bc3f5e9908e425b67b6ab38d26da..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 12
TcmYc;N@ieSU}E^a_mM3C6{G|b

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.crc b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/.part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac.crc
new file mode 100644
index 0000000000000000000000000000000000000000..4155311e5a4dac9c8a79c47c1560d5917061165e
GIT binary patch
literal 12
TcmYc;N@ieSU}88fdg&Yh6HNoR

literal 0
HcmV?d00001

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-5efaaf66-a01a-4640-a3e4-aef656269368 b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-5efaaf66-a01a-4640-a3e4-aef656269368
deleted file mode 100644
index e113fa74e796762da26b63b59a9957b9f99af9dd..0000000000000000000000000000000000000000
GIT binary patch
literal 0
HcmV?d00001

literal 206
zcmdnVz`*bWh&8tA|5f;r%F0l}!N|apoKa+WJdlaW(V5Yafti6pgu#Y|A(=sf=Yc&J
z&*5Xcg_sya7~M4#pV%vWuvZaj{CGi|iP1e-f8pMiZ$h{1rFj^VYp?MxJCT0ZUX|s8
zy++&ii`tCRVS&LxAw`LK$r-7^egT%2dX}ao`Gy7tdIko}x-1M5EDTmPQi(g43m&tV
z<zWzDSiqCl@IWSmWshK7H?x_7Z~)sB<|A9Uj8@Ny=U|IG&mhbTbT|lr9nZkXzyJUk
CN;)zC

diff --git a/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac b/hail_search/fixtures/GRCh38/SNV_INDEL/families/WGS/F000002_2.ht/rows/parts/part-0-d68dd2a3-0a18-42d9-9d62-d4fc646610ac
new file mode 100644
index 0000000000000000000000000000000000000000..4317a07a18b36279fdb7e877ea6f0a02e24cf647
GIT binary patch
literal 207
zcmdnRz`*bWh&8tA|5f;r%gWHg!N|apoKa+WJdlaW(V5Yafti6pgu#Y|A(=sf=Yc(!
z$l+tVg_sya7~M4#pV%v~e6Uw({CGi|iP1e-f8pMiZ$h{1rFj^VYp?MxJCT0ZUKOZ7
zgQso#MQujuu)yG;kfOxA<c!o{zW_^1JxkM)d_w~RJp%(~T^0rj76z*tsl=Vj1&`Uw
z@-PT6Ea1s&cp#I(Vj&dQ&1}}d<G?Y6`N$S7qt$cbIoKl4GYInnT@C_Z*E29OFaQ8P
CygC;E

literal 0
HcmV?d00001

diff --git a/hail_search/queries/base.py b/hail_search/queries/base.py
index d30a9b8fdc..0da165bcb4 100644
--- a/hail_search/queries/base.py
+++ b/hail_search/queries/base.py
@@ -472,9 +472,7 @@ def _apply_entry_filters(ht):
     def _filter_single_entries_table(self, ht, project_families, inheritance_filter=None, quality_filter=None, is_merged_ht=False, **kwargs):
         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, ch_ht = self._filter_inheritance(ht, None, inheritance_filter, sorted_family_sample_data)
         ht = self._apply_entry_filters(ht)
         ch_ht = self._apply_entry_filters(ch_ht)
 
@@ -574,7 +572,7 @@ def _get_sample_type(cls, family_index, ht_globals):
 
     def _filter_inheritance(
         self, ht, comp_het_ht, inheritance_filter, sorted_family_sample_data,
-        annotation='family_entries', entries_ht_field='family_entries'
+        annotation='family_entries', entries_ht_field='family_entries', **kwargs
     ):
         any_valid_entry = lambda x: self.GENOTYPE_QUERY_MAP[HAS_ALT](x.GT)
 
@@ -584,14 +582,14 @@ 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 if comp_het_ht is not None else ht, COMPOUND_HET, inheritance_filter,
-                sorted_family_sample_data, annotation, entries_ht_field
+                sorted_family_sample_data, annotation, entries_ht_field, **kwargs
             )
 
         if is_any_affected or not (inheritance_filter or self._inheritance_mode):
@@ -600,15 +598,43 @@ def _filter_inheritance(
 
         ht = None if self._inheritance_mode == COMPOUND_HET else self._annotate_families_inheritance(
             ht, self._inheritance_mode, inheritance_filter, sorted_family_sample_data,
-            annotation, entries_ht_field
+            annotation, entries_ht_field, **kwargs
         )
 
         return ht, comp_het_ht
 
     def _annotate_families_inheritance(
         self, ht, inheritance_mode, inheritance_filter, sorted_family_sample_data,
-        annotation, entries_ht_field,
+        annotation, entries_ht_field, family_passes_inheritance_filter = None
     ):
+        if not family_passes_inheritance_filter:
+            family_passes_inheritance_filter = self._get_family_passes_inheritance_filter
+
+        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_idx, family_samples: family_passes_inheritance_filter(
+                        entry_indices, family_idx, genotype, family_samples, ht, annotation
+                    )
+                )
+            })
+
+        return ht
+
+    def _get_family_passes_inheritance_filter(self, entry_indices, family_idx, genotype, family_samples, *args):
+        return hl.or_missing(
+            ~entry_indices.contains(family_idx) | entry_indices[family_idx].all(
+                lambda sample_i: self.GENOTYPE_QUERY_MAP[genotype](family_samples[sample_i].GT)
+        ), family_samples)
+
+    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 +656,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 {}
diff --git a/hail_search/queries/mito.py b/hail_search/queries/mito.py
index 6b46ada427..73263e8098 100644
--- a/hail_search/queries/mito.py
+++ b/hail_search/queries/mito.py
@@ -205,37 +205,62 @@ def _filter_entries_ht_both_sample_types(
             )
 
         ch_ht = None
-        family_guid_idx_map = defaultdict(dict)
+        family_idx_map = defaultdict(dict)
         for sample_type, sorted_family_sample_data in sample_types:
+            ht = self._annotate_initial_passes_inheritance(ht, sample_type)
+            ch_ht = self._annotate_initial_passes_inheritance(ch_ht, sample_type)
             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
+                annotation=sample_type.passes_inheritance_field, entries_ht_field=sample_type.family_entries_field,
+                family_passes_inheritance_filter=self._get_family_passes_inheritance_filter_both_sample_types
             )
             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[family_guid][sample_type.value] = family_idx
 
-        family_idx_map = hl.dict(family_guid_idx_map)
+        family_idx_map = hl.dict(family_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)
         return ht, ch_ht
 
+    @staticmethod
+    def _annotate_initial_passes_inheritance(ht, sample_type):
+        if ht is None:
+            return ht
+
+        return ht.annotate(**{
+            sample_type.passes_inheritance_field: ht[sample_type.family_entries_field].map(
+                lambda family_entries: hl.array(
+                    hl.range(0, hl.len(family_entries)).map(lambda _: True)
+                )
+            )})
+
+    def _get_family_passes_inheritance_filter_both_sample_types(
+        self, entry_indices, family_idx, genotype, family_samples, ht, annotation
+    ):
+        return hl.enumerate(ht[annotation][family_idx]).starmap(
+            lambda sample_idx, passes: (hl.case()
+                .when(~entry_indices.get(family_idx).contains(sample_idx), passes)
+                .when(~self.GENOTYPE_QUERY_MAP[genotype](family_samples[sample_idx].GT), False)
+                .default(passes))
+        )
+
     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)
+                    lambda family_idx, family_samples: hl.or_missing(
+                        hl.bind(lambda other_sample_type_family_idx: ((
+                            self._family_has_valid_quality(ht, sample_type, family_idx) |
+                            self._family_has_valid_quality(ht, sample_type.other_sample_type, other_sample_type_family_idx)
+                             ) &
+                            self._family_has_valid_inheritance(ht, sample_type, family_idx, other_sample_type_family_idx) &
+                            self._family_has_valid_inheritance(ht, sample_type.other_sample_type, other_sample_type_family_idx, family_idx)
+                        ), family_idx_map.get(hl.coalesce(family_samples)[0]['familyGuid']).get(sample_type.other_sample_type.value),
+                    ), family_samples)
                 )})
 
         # Merge family entries and filters from both sample types
@@ -253,13 +278,29 @@ def _apply_multi_sample_type_entry_filters(self, ht, family_idx_map):
         return ht.filter(ht.family_entries.any(hl.is_defined))
 
     @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 _family_has_valid_inheritance(ht, sample_type, family_idx, other_sample_type_family_idx):
+        return hl.bind(
+            lambda sample_type_fail_samples, other_sample_type_pass_samples: (
+                sample_type_fail_samples.all(other_sample_type_pass_samples.contains)
+            ), hl.enumerate(ht[sample_type.family_entries_field][family_idx]).starmap(
+                lambda sample_idx, sample: hl.or_missing(
+                    ~ht[sample_type.passes_inheritance_field][family_idx][sample_idx],
+                    sample['sampleId'],
+                )
+            ).filter(hl.is_defined),
+            hl.enumerate(ht[sample_type.other_sample_type.family_entries_field][other_sample_type_family_idx]).starmap(
+                lambda sample_idx, sample: hl.or_missing(
+                    ht[sample_type.other_sample_type.passes_inheritance_field][other_sample_type_family_idx][sample_idx],
+                    sample['sampleId'],
+                )
+            ).filter(hl.is_defined),
         )
 
     def _get_sample_genotype(self, samples, r=None, include_genotype_overrides=False, select_fields=None, **kwargs):
diff --git a/hail_search/test_search.py b/hail_search/test_search.py
index 39743b29e6..8fc9caa68d 100644
--- a/hail_search/test_search.py
+++ b/hail_search/test_search.py
@@ -13,7 +13,10 @@
     FAMILY_2_MITO_SAMPLE_DATA, FAMILY_2_ALL_SAMPLE_DATA, MITO_VARIANT1, MITO_VARIANT2, MITO_VARIANT3, \
     EXPECTED_SAMPLE_DATA_WITH_SEX, SV_WGS_SAMPLE_DATA_WITH_SEX, VARIANT_LOOKUP_VARIANT, \
     MULTI_PROJECT_SAMPLE_TYPES_SAMPLE_DATA, FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA, \
-    VARIANT1_BOTH_SAMPLE_TYPES, VARIANT2_BOTH_SAMPLE_TYPES, FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS
+    VARIANT1_BOTH_SAMPLE_TYPES, VARIANT2_BOTH_SAMPLE_TYPES, FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS, \
+    VARIANT3_BOTH_SAMPLE_TYPES, VARIANT4_BOTH_SAMPLE_TYPES, VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, \
+    VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, \
+    VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY
 from hail_search.web_app import init_web_app, sync_to_async_hail_query
 from hail_search.queries.base import BaseHailTableQuery
 
@@ -365,34 +368,31 @@ async def test_both_sample_types_search(self):
             MULTI_PROJECT_BOTH_SAMPLE_TYPE_VARIANTS, gene_counts=GENE_COUNTS, sample_data=MULTI_PROJECT_SAMPLE_TYPES_SAMPLE_DATA,
         )
 
-        # Variant1 in family_2 is de novo in exome but maternally inherited in genome.
-        # Genome passes quality and inheritance, show genotypes for both sample types.
-        variant1_interval = ['1', 10438, 10440]
+        # Variant 1 is de novo in exome but inherited and homozygous in genome.
+        # Variant 2 is inherited and homozygous in exome and de novo and homozygous in genome.
+        # Variant 3 is inherited in both sample types. Variant 4 is de novo in both sample types.
         inheritance_mode = 'recessive'
         await self._assert_expected_search(
-            [VARIANT1_BOTH_SAMPLE_TYPES], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA, inheritance_mode=inheritance_mode,
-            **COMP_HET_ALL_PASS_FILTERS, intervals=[variant1_interval]
+            [VARIANT1_BOTH_SAMPLE_TYPES, VARIANT2_BOTH_SAMPLE_TYPES, [VARIANT3_BOTH_SAMPLE_TYPES, VARIANT4_BOTH_SAMPLE_TYPES]],
+            sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA, inheritance_mode=inheritance_mode,
+            **COMP_HET_ALL_PASS_FILTERS
         )
-        # Exome passes quality and inheritance, show genotypes for both sample types.
-        inheritance_mode = 'de_novo'
         await self._assert_expected_search(
-            [VARIANT1_BOTH_SAMPLE_TYPES], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA, inheritance_mode=inheritance_mode,
-            intervals=[variant1_interval]
+            [VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY,
+             [VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY]],
+            sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS, inheritance_mode=inheritance_mode,
+            **COMP_HET_ALL_PASS_FILTERS
         )
 
-        # Variant 2 in family_2 is inherited in exome and there is no parental data in genome.
-        # Genome and exome pass quality and inheritance, show genotypes for both sample types.
-        variant2_interval = ['1', 38724418, 38724420]
-        inheritance_mode = 'recessive'
+        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,
-            inheritance_mode=inheritance_mode, **COMP_HET_ALL_PASS_FILTERS, intervals=[variant2_interval]
+            [VARIANT1_BOTH_SAMPLE_TYPES, VARIANT2_BOTH_SAMPLE_TYPES, VARIANT4_BOTH_SAMPLE_TYPES],
+            sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA, inheritance_mode=inheritance_mode,
         )
-        # Genome passes quality and inheritance exome fails inheritance (parental data shows variant is inherited).
-        inheritance_mode = 'de_novo'
+        # Variant 2 fails inheritance when parental data is missing in genome
         await self._assert_expected_search(
-            [VARIANT2_BOTH_SAMPLE_TYPES], sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS,
-            inheritance_mode=inheritance_mode, intervals=[variant2_interval]
+            [VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY, VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY],
+            sample_data=FAMILY_2_BOTH_SAMPLE_TYPE_SAMPLE_DATA_MISSING_PARENTAL_WGS, inheritance_mode=inheritance_mode,
         )
 
     async def test_inheritance_filter(self):
diff --git a/hail_search/test_utils.py b/hail_search/test_utils.py
index 6895b8df25..722e7242a7 100644
--- a/hail_search/test_utils.py
+++ b/hail_search/test_utils.py
@@ -160,23 +160,29 @@
     '_sort': [1000010439],
     'CAID': 'CA16717152',
 }
-VARIANT1_BOTH_SAMPLE_TYPES = deepcopy(VARIANT1)
-genotypes = VARIANT1['genotypes']
-VARIANT1_BOTH_SAMPLE_TYPES['genotypes'] = {
+
+VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY = deepcopy(VARIANT1)
+genotypes = VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes']
+VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes'] = {
     'I000004_hg00731': [
         genotypes['I000004_hg00731'],
-        {**genotypes['I000004_hg00731'], 'sampleType': 'WGS'}
-    ],
-    'I000005_hg00732': [
-        genotypes['I000005_hg00732'],
-        {**genotypes['I000005_hg00732'], 'gq': 99, 'numAlt': 1, 'sampleType': 'WGS'}
-    ],
-    'I000006_hg00733': [
-        genotypes['I000006_hg00733'],
-        {**genotypes['I000006_hg00733'], 'sampleType': 'WGS'}
+        {**genotypes['I000004_hg00731'], 'numAlt': 2, 'sampleType': 'WGS'}
     ],
+    'I000005_hg00732': [genotypes['I000005_hg00732']],
+    'I000006_hg00733': [genotypes['I000006_hg00733']],
 }
 
+VARIANT1_BOTH_SAMPLE_TYPES = deepcopy(VARIANT1_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY)
+genotypes = VARIANT1_BOTH_SAMPLE_TYPES['genotypes']
+VARIANT1_BOTH_SAMPLE_TYPES['genotypes']['I000005_hg00732'] = [
+    *genotypes['I000005_hg00732'],
+    {**genotypes['I000005_hg00732'][0], 'gq': 99, 'numAlt': 1, 'sampleType': 'WGS'}
+]
+VARIANT1_BOTH_SAMPLE_TYPES['genotypes']['I000006_hg00733'] = [
+    *genotypes['I000006_hg00733'],
+    {**genotypes['I000006_hg00733'][0], 'sampleType': 'WGS'}
+]
+
 VARIANT2 = {
     'variantId': '1-38724419-T-G',
     'chrom': '1',
@@ -280,9 +286,10 @@
     '_sort': [1038724419],
     'CAID': None,
 }
-VARIANT2_BOTH_SAMPLE_TYPES = deepcopy(VARIANT2)
-genotypes = VARIANT2_BOTH_SAMPLE_TYPES['genotypes']
-VARIANT2_BOTH_SAMPLE_TYPES['genotypes'] = {
+
+VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY = deepcopy(VARIANT2)
+genotypes = VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes']
+VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes'] = {
     'I000004_hg00731': [
         genotypes['I000004_hg00731'],
         {**genotypes['I000004_hg00731'], 'sampleType': 'WGS'}
@@ -291,6 +298,17 @@
     'I000006_hg00733': [genotypes['I000006_hg00733']],
 }
 
+VARIANT2_BOTH_SAMPLE_TYPES = deepcopy(VARIANT2_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY)
+genotypes = VARIANT2_BOTH_SAMPLE_TYPES['genotypes']
+VARIANT2_BOTH_SAMPLE_TYPES['genotypes']['I000005_hg00732'] = [
+    *genotypes['I000005_hg00732'],
+    {**genotypes['I000005_hg00732'][0], 'numAlt': 0, 'sampleType': 'WGS'}
+]
+VARIANT2_BOTH_SAMPLE_TYPES['genotypes']['I000006_hg00733'] = [
+    *genotypes['I000006_hg00733'],
+    {**genotypes['I000006_hg00733'][0], 'sampleType': 'WGS'}
+]
+
 VARIANT3 = {
     'variantId': '1-91502721-G-A',
     'chrom': '1',
@@ -361,6 +379,29 @@
     '_sort': [1091502721],
     'CAID': 'CA10960369',
 }
+
+VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY = deepcopy(VARIANT3)
+genotypes = VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes']
+VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes'] = {
+    'I000004_hg00731': [
+        genotypes['I000004_hg00731'],
+        {**genotypes['I000004_hg00731'], 'sampleType': 'WGS'}
+    ],
+    'I000005_hg00732': [genotypes['I000005_hg00732']],
+    'I000006_hg00733': [genotypes['I000006_hg00733']],
+}
+
+VARIANT3_BOTH_SAMPLE_TYPES = deepcopy(VARIANT3_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY)
+genotypes = VARIANT3_BOTH_SAMPLE_TYPES['genotypes']
+VARIANT3_BOTH_SAMPLE_TYPES['genotypes']['I000005_hg00732'] = [
+    *genotypes['I000005_hg00732'],
+    {**genotypes['I000005_hg00732'][0], 'sampleType': 'WGS'}
+]
+VARIANT3_BOTH_SAMPLE_TYPES['genotypes']['I000006_hg00733'] = [
+    *genotypes['I000006_hg00733'],
+    {**genotypes['I000006_hg00733'][0], 'sampleType': 'WGS'}
+]
+
 VARIANT4 = {
     'variantId': '1-91511686-T-G',
     'chrom': '1',
@@ -434,6 +475,28 @@
     'CAID': 'CA341062623',
 }
 
+VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY = deepcopy(VARIANT4)
+genotypes = VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes']
+VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY['genotypes'] = {
+    'I000004_hg00731': [
+        genotypes['I000004_hg00731'],
+        {**genotypes['I000004_hg00731'], 'sampleType': 'WGS'}
+    ],
+    'I000005_hg00732': [genotypes['I000005_hg00732']],
+    'I000006_hg00733': [genotypes['I000006_hg00733']],
+}
+
+VARIANT4_BOTH_SAMPLE_TYPES = deepcopy(VARIANT4_BOTH_SAMPLE_TYPES_PROBAND_WGS_ONLY)
+genotypes = VARIANT4_BOTH_SAMPLE_TYPES['genotypes']
+VARIANT4_BOTH_SAMPLE_TYPES['genotypes']['I000005_hg00732'] = [
+    *genotypes['I000005_hg00732'],
+    {**genotypes['I000005_hg00732'][0], 'sampleType': 'WGS'}
+]
+VARIANT4_BOTH_SAMPLE_TYPES['genotypes']['I000006_hg00733'] = [
+    *genotypes['I000006_hg00733'],
+    {**genotypes['I000006_hg00733'][0], 'sampleType': 'WGS'}
+]
+
 VARIANT_LOOKUP_VARIANT = {
     **VARIANT1,
     'familyGenotypes': {