From 3a829e08c1b99ff5725bded9202f70e4928d9850 Mon Sep 17 00:00:00 2001
From: Jan Patrick Pett <jpatrickpett@googlemail.com>
Date: Sun, 10 Mar 2024 15:55:02 +0000
Subject: [PATCH] handle de_groups in score removal

---
 snp2cell/snp2cell_class.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/snp2cell/snp2cell_class.py b/snp2cell/snp2cell_class.py
index 884f288..c7e78d0 100644
--- a/snp2cell/snp2cell_class.py
+++ b/snp2cell/snp2cell_class.py
@@ -55,6 +55,7 @@ def _init_scores(self) -> None:
             self.scores = pd.DataFrame(index=list(self.grn.nodes))
             self.scores_prop = pd.DataFrame(index=list(self.grn.nodes))
             self.scores_rand = {}
+            self.de_groups = {}
         else:
             raise ValueError("No GRN set, add GRN first.")
 
@@ -859,8 +860,11 @@ def remove_scores(self, which: str = "propagated", **kwargs: Any) -> None:
             if kwargs:
                 cols = self.scores.filter(**kwargs).columns  # type: ignore
                 self.scores = self.scores.drop(columns=cols)  # type: ignore
+                for k in self.de_groups:
+                    self.de_groups[k] = [i for i in self.de_groups[k] if i not in cols]
             else:
                 self.scores = pd.DataFrame(index=list(self.grn.nodes))  # type: ignore
+                self.de_groups = {}
 
     def get_components(self, sel_nodes: list[str]) -> tuple[nx.Graph, list[set]]:
         """