PLaying with gridding optimisation.

stlearn/tools/microenv/cci/analysis.py

@@ -69,7 +69,7 @@ def load_lrs(names: Union[str, list, None] = None, species: str = "human") -> np
 
 
 def grid(adata, n_row: int = 10, n_col: int = 10, use_label: str = None,
-         n_cpus: int=1):
+         n_cpus: int=1, verbose: bool=True):
     """Creates a new anndata representing a gridded version of the data; can be
         used upstream of CCI pipeline. NOTE: intended use is for single cell
         spatial data, not Visium or other lower resolution tech.
@@ -92,8 +92,9 @@ def grid(adata, n_row: int = 10, n_col: int = 10, use_label: str = None,
         Equivalent expression data to adata, except values have been summed by
         cells that fall within defined bins.
     """
+    if verbose:
+        print("Gridding...")
 
-    print("Gridding...")
     # Setting threads for paralellisation #
     if type(n_cpus) != type(None):
         numba.set_num_threads( n_cpus )
@@ -110,12 +111,6 @@ def grid(adata, n_row: int = 10, n_col: int = 10, use_label: str = None,
 
     grid_expr = np.zeros((n_squares, adata.shape[1]))
     grid_coords = np.zeros((n_squares, 2))
-    #grid_bcs = []
-    #grid_cell_counts = []
-    #gridded_cells = []
-    #cell_grid = []
-    #grid_bcs = List( types.unicode_type )
-    #grid_cell_counts = List( numba.int64 )
     grid_cell_counts = np.zeros((n_squares), dtype=np.int64)
     # If use_label specified, then will generate deconvolution information
     cell_labels, cell_set, cell_info = None, None, None
@@ -124,46 +119,7 @@ def grid(adata, n_row: int = 10, n_col: int = 10, use_label: str = None,
         cell_set = np.unique(cell_labels).astype(str)
         cell_info = np.zeros((n_squares, len(cell_set)), dtype=np.float64)
 
-    # # generate grids from top to bottom and left to right
-    # n = 0
-    # for i in range(n_col):
-    #     x_left, x_right = xedges[i], xedges[i + 1]
-    #     for j in range(n_row):
-    #         y_down, y_up = yedges[j], yedges[j + 1]
-    #         grid_coords[n, :] = [(x_right + x_left) / 2, (y_up + y_down) / 2]
-    #
-    #         # Now determining the cells within the gridded area #
-    #         if i != n_col - 1 and j == n_row - 1:  # top left corner
-    #             x_true = (xs >= x_left) & (xs < x_right)
-    #             y_true = (ys <= y_up) & (ys > y_down)
-    #         elif i == n_col - 1 and j != n_row - 1:  # bottom righ corner
-    #             x_true = (xs > x_left) & (xs <= x_right)
-    #             y_true = (ys < y_up) & (ys >= y_down)
-    #         else:  # average case
-    #             x_true = (xs >= x_left) & (xs < x_right)
-    #             y_true = (ys < y_up) & (ys >= y_down)
-    #
-    #         grid_cells = cell_bcs[x_true & y_true]
-    #         grid_cells_str = ",".join(grid_cells)
-    #         grid_bcs.append(grid_cells_str)
-    #         grid_cell_counts.append(len(grid_cells))
-    #         gridded_cells.extend(grid_cells)
-    #         cell_grid.extend([f"grid_{n}"] * len(grid_cells))
-    #
-    #         # Summing the expression across these cells to get the grid expression #
-    #         if len(grid_cells) > 0:
-    #             cell_bool = [cell in grid_cells for cell in cell_bcs]
-    #             grid_expr[n, :] = adata.X[cell_bool, :].sum(axis=0)
-    #
-    #         # If we have cell type information, will record #
-    #         if type(use_label) != type(None) and len(grid_cells) > 0:
-    #             grid_cell_types = cell_labels[cell_bool]
-    #             cell_info[n, :] = [
-    #                 len(np.where(grid_cell_types == ct)[0]) / len(grid_cell_types)
-    #                 for ct in cell_set
-    #             ]
-    #
-    #         n += 1
+    # Performing grid operation in parallel #
     grid_parallel(grid_coords, xedges, yedges, n_row, n_col, xs, ys,
               cell_bcs, grid_cell_counts,
               grid_expr, adata.X, type(use_label)!=type(None),
@@ -187,7 +143,6 @@ def grid(adata, n_row: int = 10, n_col: int = 10, use_label: str = None,
                   index=grid_data.obs_names.values.astype(str), columns=cell_set
         )
         max_indices = np.apply_along_axis(np.argmax, 1, cell_info)
-        #cell_set = np.unique(grid_data.uns[use_label].columns.values)
         grid_data.obs[use_label] = [cell_set[index] for index in max_indices]
         grid_data.obs[use_label] = grid_data.obs[use_label].astype('category')
         grid_data.obs[use_label] = grid_data.obs[use_label].cat.set_categories(