add Layers to store expression matrices for different stages

stereo/algorithm/total_vi.py

@@ -250,7 +250,7 @@ def save_result(
 
         if h5mu_file_name is None:
             h5mu_file_name = f'{self._rna_data.sn}_{self._rna_data.bin_size}.h5mu'
-        mudata.write(f"{out_dir}/{h5mu_file_name}", mdata)
+        mudata.write(f"{out_dir}/{h5mu_file_name}", mdata, compression='gzip')
 
         self._differential_expression = de_df
         self._use_cluster_res_key = use_cluster_res_key

stereo/core/cell.py

@@ -57,6 +57,24 @@ def __getitem__(self, key):
     def __len__(self):
         return self.size
 
+    def get_index(self, cell_list=None):
+        if cell_list is None:
+            return np.arange(self.size)
+
+        if isinstance(cell_list, (int, np.integer, str)):
+            cell_list = [cell_list]
+
+        if isinstance(cell_list, (list, np.ndarray, pd.Index)):
+            cell_list = np.array(cell_list)
+            if isinstance(cell_list[0], (int, np.integer)):
+                cell_index = cell_list[(cell_list >= 0) & (cell_list < self.size)]
+            else:
+                cell_index = self.obs.index.get_indexer(cell_list)
+                cell_index = cell_index[cell_index >= 0]
+            return cell_index
+        else:
+            raise TypeError('cell_list must be a int or str or list or np.ndarray or pd.Index object.')
+
     @property
     def matrix(self):
         return self._matrix

stereo/core/data_component.py

@@ -1,7 +1,5 @@
-from typing import Union
 from copy import deepcopy
 
-from anndata import AnnData
 import numpy as np
 from scipy.sparse import spmatrix
 import pandas as pd
@@ -50,7 +48,7 @@ def __setitem__(self, key, value):
         super().__setitem__(key, value)
 
     def __str__(self) -> str:
-        info = f"layers with keys {list(self.keys())}."
+        info = f"Layers with keys: {', '.join(list(self.keys()))}"
         return info
 
     def __repr__(self) -> str:

stereo/core/gene.py

@@ -46,6 +46,35 @@ def __getitem__(self, key):
     def __len__(self):
         return self.size
 
+    def get_index(self, gene_list=None, only_highly_genes=False):
+        only_highly_genes = only_highly_genes and 'highly_variable' in self.var.columns
+        if only_highly_genes:
+            highly_variable = self.var['highly_variable'].fillna(False).to_numpy()
+            used_gene_index = np.arange(self.size)[highly_variable]
+        else:
+            used_gene_index = np.arange(self.size)
+
+        if gene_list is None:
+            return used_gene_index
+
+        if isinstance(gene_list, (int, np.integer, str)):
+            gene_list = [gene_list]
+
+        if isinstance(gene_list, (list, np.ndarray, pd.Index)):
+            gene_list = np.array(gene_list)
+            if isinstance(gene_list[0], (int, np.integer)):
+                # gene_index = np.intersect1d(used_gene_index, gene_list)
+                isin = np.isin(gene_list, used_gene_index)
+                gene_index = gene_list[isin]
+            else:
+                gene_index = self.var.index.get_indexer(gene_list)
+                gene_index = gene_index[gene_index != -1]
+                isin = np.isin(gene_index, used_gene_index)
+                gene_index = gene_index[isin]
+            return gene_index
+        else:
+            raise TypeError('gene_list must be a int or str or list or np.ndarray or pd.Index object.')
+
     @property
     def matrix(self):
         return self._matrix

stereo/core/result.py

@@ -664,7 +664,8 @@ def _get_marker_genes_res(self, name):
             'method': method,
             'control_groups': marker_genes_result['params']['reference'],
             'corr_method': marker_genes_result['params']['corr_method'],
-            'use_raw': marker_genes_result['params']['use_raw']
+            'use_raw': marker_genes_result['params']['use_raw'],
+            'layer': marker_genes_result['params']['layer']
         }
         if 'marker_genes_res_key' in marker_genes_result['params']:
             marker_genes_result_reconstructed['parameters']['marker_genes_res_key'] = \