fix return type hints and docstrings to match actual return values

Author: ShaokunAn

SpatialQuery/spatial_gene_covarying.py

@@ -1,5 +1,5 @@
 
-from typing import List, Union, Optional, Literal, Dict
+from typing import List, Union, Optional, Literal, Dict, Tuple
 
 import numpy as np
 import pandas as pd
@@ -586,9 +586,9 @@ def compute_gene_gene_correlation_adata(sq_obj,
                                         min_size: int = 0,
                                         min_nonzero: int = 10,
                                         alpha: Optional[float] = None
-                                        ) -> pd.DataFrame:
+                                        ) -> Tuple[pd.DataFrame, dict]:
     """
-    Compute gene-gene cross correlation between anchor and neighboring motif cells. Only considers inter-cell-type interactions. 
+    Compute gene-gene cross correlation between anchor and neighboring motif cells. Only considers inter-cell-type interactions.
     After finding neighbors using the full motif, removes all cells of the center cell type from both neighbor and
     non-neighbor groups. For Pearson correlation, uses shifted correlation (subtract cell type mean) to enable
     comparison across different niches/motifs.
@@ -1017,7 +1017,7 @@ def compute_gene_gene_correlation_binary(sq_obj,
                                          min_size: int = 0,
                                          min_nonzero: int = 10,
                                          alpha: Optional[float] = None
-                                         ) -> pd.DataFrame:
+                                         ) -> Tuple[pd.DataFrame, dict]:
     """
     Compute gene-gene correlation using deviation from global binary mean for binary expression data.
 

SpatialQuery/spatial_query.py

@@ -1348,8 +1348,8 @@ def plot_gene_pair_heatmap(self,
                                title: Optional[str] = None,
                                cmap: str = 'GnBu'):
         """
-        Plot a heatmap showing the cross-varying gene pairs.
-        
+        Plot a heatmap showing the cross-varying gene pairs and return cluster assignments.
+
         Parameter
         ---------
         gene_pair_df:
@@ -1362,10 +1362,14 @@ def plot_gene_pair_heatmap(self,
             Figure title. If None, will use a default title based on center cell type.
         cmap:
             Colormap for the heatmap, default is 'GnBu'
-        
+
         Return
         ------
-        A figure showing the heatmap of motif cell type distribution.
+        pd.DataFrame
+            Displays a biclustered heatmap of gene pairs and returns a DataFrame with
+            cluster assignments. Columns: gene_center, gene_motif, combined_score,
+            cluster_type, cluster_row, cluster_col. If input has 'cell_type' column,
+            it will also be included.
         """
         return plotting.plot_gene_pair_heatmap(gene_pair_df=gene_pair_df, figsize=figsize,
                                                        save_path=save_path)

SpatialQuery/spatial_query_multiple_fov.py

@@ -1629,9 +1629,6 @@ def compute_gene_gene_correlation(self,
             DataFrame with correlation results. Columns: gene_center, gene_motif, corr_neighbor,
             corr_non_neighbor, corr_center_no_motif, p_value_test1, p_value_test2,
             delta_corr_test1, delta_corr_test2, combined_score, adj-pval-test1, adj-pval-test2.
-        fov_info : dict
-            FOV-level information with keys: fov_statistics, total_pairs_neighbor,
-            total_pairs_non_neighbor, total_pairs_no_motif, n_fovs_analyzed.
         """
         if not self.build_gene_index:
             print('Computing covarying genes using expression data ...')

SpatialQuery/utils.py

@@ -57,7 +57,7 @@ def retrieve_niche_pattern_freq(fp, sp, ct, max_dist):
         radius of neighborhood
 
     Return:
-        AnnData with central cell type and neighboring cells in fp
+        pd.DataFrame with frequency of each cell type in fp around the central cell type
     """
     if ct not in sp.labels.unique():
         print(f"{ct} does not exist in FOV.")