From 0c84c269392099461b039f0dac5fded9cc26921b Mon Sep 17 00:00:00 2001
From: tanliwei <liweitan@aliyun.com>
Date: Fri, 28 Jun 2024 15:34:45 +0800
Subject: [PATCH] update for IO.

---
 stereo/io/reader.py | 4 +++-
 stereo/io/writer.py | 5 +++--
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/stereo/io/reader.py b/stereo/io/reader.py
index 9948d599..3a2ed8d7 100644
--- a/stereo/io/reader.py
+++ b/stereo/io/reader.py
@@ -472,7 +472,9 @@ def read_h5ms(file_path, use_raw=True, use_result=True):
         # result = {}
         for k in f.keys():
             if k == 'sample':
-                for one_slice_key in f[k].keys():
+                slice_keys = list(f[k].keys())
+                slice_keys.sort(key=lambda k: int(k.split('_')[1]))
+                for one_slice_key in slice_keys:
                     data = _read_stereo_h5ad_from_group(f[k][one_slice_key], StereoExpData(), use_raw, use_result)
                     # encoding_type = f[k][one_slice_key].attrs.get('encoding-type', 'stereo_exp_data')
                     # if encoding_type == 'stereo_exp_data':
diff --git a/stereo/io/writer.py b/stereo/io/writer.py
index 9f2f2691..c2ca6ee6 100644
--- a/stereo/io/writer.py
+++ b/stereo/io/writer.py
@@ -122,6 +122,7 @@ def _write_one_h5ad(f: h5py.File, data: StereoExpData, use_raw=False, use_result
     # h5ad.write(data.bin_size, f, 'bin_size')
     # h5ad.write(data.merged, f, 'merged')
 
+    use_raw = use_raw and data.tl.raw is not None
     if use_raw is True:
         same_genes = np.array_equal(data.tl.raw.gene_names, data.gene_names)
         same_cells = np.array_equal(data.tl.raw.cell_names, data.cell_names)
@@ -300,7 +301,7 @@ def write_h5ms(ms_data, output: str):
         f.create_group('sample')
         for idx, data in enumerate(ms_data._data_list):
             f['sample'].create_group(f'sample_{idx}')
-            _write_one_h5ad(f['sample'][f'sample_{idx}'], data)
+            _write_one_h5ad(f['sample'][f'sample_{idx}'], data, use_raw=True, use_result=True)
             # if isinstance(data, AnnBasedStereoExpData):
             #     _write_one_anndata(f['sample'][f'sample_{idx}'], data)
             # else:
@@ -314,7 +315,7 @@ def write_h5ms(ms_data, output: str):
                 g = f['sample_merged'].create_group(scope_key)
                 if ms_data.merged_data and id(ms_data.merged_data) == id(merged_data):
                     g.attrs['merged_from_all'] = True
-                _write_one_h5ad(g, merged_data)
+                _write_one_h5ad(g, merged_data, use_raw=True, use_result=True)
                 # if isinstance(merged_data, AnnBasedStereoExpData):
                 #     _write_one_anndata(g, merged_data)
                 # else: