do some updates.

STOmics · Sep 5, 2024 · 8e61353 · 8e61353
1 parent 2063880
commit 8e61353
Show file tree

Hide file tree

Showing 3 changed files with 23 additions and 9 deletions.
diff --git a/stereo/core/result.py b/stereo/core/result.py
@@ -573,14 +573,16 @@ def _set_connectivities_res(self, key, value):
         assert type(value) is dict and not {'connectivities', 'nn_dist'} - set(value.keys()), \
             'not enough key to set connectivities'
         self.__based_ann_data.uns[key] = {
-            'params': {
-                'method': value['method'],
-                'n_neighbors': value['n_neighbors'],
-                'metric': value['metric'],
-            },
+            'params': {},
             'source': 'stereopy',
             'method': 'neighbors'
         }
+        if 'method' in value:
+            self.__based_ann_data.uns[key]['params']['method'] = value['method']
+        if 'n_neighbors' in value:
+            self.__based_ann_data.uns[key]['params']['n_neighbors'] = value['n_neighbors']
+        if 'metric' in value:
+            self.__based_ann_data.uns[key]['params']['metric'] = value['metric']
         if key == 'neighbors':
             self.__based_ann_data.uns[key]['connectivities_key'] = 'connectivities'
             self.__based_ann_data.uns[key]['distances_key'] = 'distances'

diff --git a/stereo/core/st_pipeline.py b/stereo/core/st_pipeline.py
@@ -85,6 +85,10 @@ def __getattr__(self, item):
     def key_record(self):
         return self._key_record
 
+    @key_record.setter
+    def key_record(self, key_record):
+        self._key_record = key_record
+
     @property
     def raw(self) -> Union[StereoExpData, AnnBasedStereoExpData]:
         """
@@ -921,12 +925,15 @@ def spatial_neighbors(self,
         from ..io.reader import stereo_to_anndata
         import squidpy as sq
         neighbor, connectivities, dists = copy.deepcopy(self.get_neighbors_res(neighbors_res_key))
-        adata = stereo_to_anndata(self.data, split_batches=False)
+        if isinstance(self.data, AnnBasedStereoExpData):
+            adata = self.data.adata
+        else:
+            adata = stereo_to_anndata(self.data, split_batches=False)
         sq.gr.spatial_neighbors(adata, n_neighs=n_neighbors)
         connectivities.data[connectivities.data > 0] = 1
         adj = connectivities + adata.obsp['spatial_connectivities']
         adj.data[adj.data > 0] = 1
-        res = {'neighbor': neighbor, 'connectivities': adj, 'nn_dist': dists}
+        res = {'neighbor': neighbor, 'connectivities': adj, 'nn_dist': dists, 'n_neighbors': n_neighbors}
         self.result[res_key] = res
         key = 'neighbors'
         self.reset_key_record(key, res_key)
@@ -1669,3 +1676,8 @@ def key_record(self):
         if 'key_record' not in self.data.adata.uns:
             self.data.adata.uns['key_record'] = self._key_record
         return self.data.adata.uns['key_record']
+
+    @key_record.setter
+    def key_record(self, key_record):
+        self._key_record = key_record
+        self.data.adata.uns['key_record'] = key_record
diff --git a/stereo/utils/data_helper.py b/stereo/utils/data_helper.py
@@ -601,8 +601,8 @@ def split_for_ann_based_stereo_exp_data(data: AnnBasedStereoExpData = None):
     # data.array2sparse()
     batch = np.unique(data.cells.batch)
     for bno in batch:
-        adata = data.adata[data.adata.obs['batch'] == bno]
-        adata.uns = adata.uns.copy()
+        adata = data.adata[data.adata.obs['batch'] == bno].copy()
+        # adata.uns = adata.uns.copy()
         new_data = AnnBasedStereoExpData(based_ann_data=adata, spatial_key=data.spatial_key)
         new_data.tl.key_record = deepcopy(data.tl.key_record)
         new_data.sn = data.sn[bno]