support multiple motifs input in motif_enrichment_knn/dist

Author: ShaokunAn

SpatialQuery/spatial_query.py

@@ -235,7 +235,7 @@ def find_fp_dist(self,
 
     def motif_enrichment_knn(self,
                              ct: str,
-                             motifs: Union[str, List[str]] = None,
+                             motifs: Union[str, List[str], List[List[str]]] = None,
                              k: int = 30,
                              min_support: float = 0.5,
                              max_dist: float = 20,
@@ -295,14 +295,21 @@ def motif_enrichment_knn(self,
             motifs = fp['itemsets']
         else:
             if isinstance(motifs, str):
+                motifs = [[motifs]]
+            elif isinstance(motifs, list) and all(isinstance(m, str) for m in motifs):
                 motifs = [motifs]
+            # else: List[List[str]], keep as is
 
             labels_unique = self.labels.unique()
-            motifs_exc = [m for m in motifs if m not in labels_unique]
-            if len(motifs_exc) != 0:
-                print(f"Found no {motifs_exc} in {self.label_key}. Ignoring them.")
-            motifs = [m for m in motifs if m not in motifs_exc]
-            motifs = [motifs]
+            filtered_motifs = []
+            for motif in motifs:
+                motif_exc = [m for m in motif if m not in labels_unique]
+                if len(motif_exc) > 0:
+                    print(f"Found no {motif_exc} in {self.label_key}. Ignoring them.")
+                valid_motif = [m for m in motif if m in labels_unique]
+                if len(valid_motif) > 0:
+                    filtered_motifs.append(valid_motif)
+            motifs = filtered_motifs
 
         if len(motifs) == 0:
             # Return empty DataFrame with same structure
@@ -401,7 +408,7 @@ def motif_enrichment_knn(self,
 
     def motif_enrichment_dist(self,
                               ct: str,
-                              motifs: Union[str, List[str]] = None,
+                              motifs: Union[str, List[str], List[List[str]]] = None,
                               max_dist: float = 20,
                               min_size: int = 0,
                               min_support: float = 0.5,
@@ -454,14 +461,21 @@ def motif_enrichment_dist(self,
             motifs = fp['itemsets']
         else:
             if isinstance(motifs, str):
+                motifs = [[motifs]]
+            elif isinstance(motifs, list) and all(isinstance(m, str) for m in motifs):
                 motifs = [motifs]
+            # else: List[List[str]], keep as is
 
             labels_unique = self.labels.unique()
-            motifs_exc = [m for m in motifs if m not in labels_unique]
-            if len(motifs_exc) != 0:
-                print(f"Found no {motifs_exc} in {self.label_key}. Ignoring them.")
-            motifs = [m for m in motifs if m not in motifs_exc]
-            motifs = [motifs]
+            filtered_motifs = []
+            for motif in motifs:
+                motif_exc = [m for m in motif if m not in labels_unique]
+                if len(motif_exc) > 0:
+                    print(f"Found no {motif_exc} in {self.label_key}. Ignoring them.")
+                valid_motif = [m for m in motif if m in labels_unique]
+                if len(valid_motif) > 0:
+                    filtered_motifs.append(valid_motif)
+            motifs = filtered_motifs
 
         if len(motifs) == 0:
             # Return empty DataFrame with same structure

SpatialQuery/spatial_query_multiple_fov.py

@@ -300,7 +300,7 @@ def find_fp_dist(self,
 
     def motif_enrichment_knn(self,
                              ct: str,
-                             motifs: Union[str, List[str]] = None,
+                             motifs: Union[str, List[str], List[List[str]]] = None,
                              dataset: Union[str, List[str]] = None,
                              k: int = 30,
                              min_support: float = 0.5,
@@ -374,15 +374,23 @@ def motif_enrichment_knn(self,
             motifs = fp['itemsets'].tolist()
         else:
             if isinstance(motifs, str):
+                motifs = [[motifs]]
+            elif isinstance(motifs, list) and all(isinstance(m, str) for m in motifs):
                 motifs = [motifs]
+            # else: List[List[str]], keep as is
+
+            filtered_motifs = []
+            for motif in motifs:
+                motif_exc = [m for m in motif if m not in labels_unique_all]
+                if len(motif_exc) > 0:
+                    print(f"Found no {motif_exc} in {dataset}. Ignoring them.")
+                valid_motif = [m for m in motif if m in labels_unique_all]
+                if len(valid_motif) > 0:
+                    filtered_motifs.append(valid_motif)
 
-            motifs_exc = [m for m in motifs if m not in labels_unique_all]
-            if len(motifs_exc) != 0:
-                print(f"Found no {motifs_exc} in {dataset}. Ignoring them.")
-            motifs = [m for m in motifs if m not in motifs_exc]
-            if len(motifs) == 0:
+            if len(filtered_motifs) == 0:
                 raise ValueError(f"All cell types in motifs are missed in {self.label_key}.")
-            motifs = [motifs]
+            motifs = filtered_motifs
 
         # Initialize dictionaries to store cell IDs if requested
         motif_cell_ids = {}