adding variants_to_keep functionality back in

PMBio · Jan 23, 2024 · e111ff6 · e111ff6
1 parent 5ef1a36
commit e111ff6
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/deeprvat/data/dense_gt.py b/deeprvat/data/dense_gt.py
@@ -558,6 +558,8 @@ def setup_variants(
 
             af_vals = af_annotation[af_col].to_numpy()
             ids_to_keep = set(af_annotation.loc[(af_vals >= af_threshold) & (af_vals <= 1 - af_threshold), "id"])
+            if self.variants_to_keep is not None: # Caution: This has not been tested yet but should work 
+                ids_to_keep.intersect(self.variants_to_keep)
             mask = variants["id"].isin(ids_to_keep).to_numpy()
 
             logger.debug(f'    {mask.sum()} variants "common" by AF filter')