diff --git a/src/hest/readers.py b/src/hest/readers.py index 6d9b1cb..0158284 100644 --- a/src/hest/readers.py +++ b/src/hest/readers.py @@ -1241,6 +1241,7 @@ def pool_bins_visiumhd(adata: sc.AnnData, pixel_size: float, dst_bin_size_um=128 sc.AnnData: AnnData object, each row in .obs represents a bin, each row in `.X` represents the sum of visium-hd sub-bins (of size `src_bin_size_um`) within that larger bin (of size `dst_bin_size_um`). Center coordinates of each bin (in pixel on WSI) are in adata.obsm['spatial'] """ import scanpy as sc + from scipy.sparse import issparse if src_bin_size_um >= dst_bin_size_um: raise ValueError("dst_bin_size_um needs to be larger than src_bin_size_um") @@ -1278,7 +1279,11 @@ def pool_bins_visiumhd(adata: sc.AnnData, pixel_size: float, dst_bin_size_um=128 for i in range(nb_chunks): start, end = i * chunk_len, min((i + 1) * chunk_len, len(c)) - spot_grid_np[c[start:end]] += adata.X[start:end] + counts = adata.X[start:end] + if issparse(counts): + counts = counts.toarray() + # Repeated destination indices must accumulate every source bin. + np.add.at(spot_grid_np, c[start:end], counts) expression_df = pd.DataFrame(spot_grid_np, columns=spot_grid.columns) @@ -1460,4 +1465,4 @@ def process_raw_samples( patching_kwargs=patching_kwargs ) - save_meta(dict(row)) \ No newline at end of file + save_meta(dict(row)) diff --git a/tests/hest_tests.py b/tests/hest_tests.py index 73710c8..a001881 100644 --- a/tests/hest_tests.py +++ b/tests/hest_tests.py @@ -19,7 +19,7 @@ raise ImportError(f"Importing 'hest' took too long ({elapsed_time:.2f} seconds). Maximum allowed time is {MAX_HEST_IMPORT_S} seconds. Please, keep large large imports conditional") from hest.autoalign import autoalign_visium -from hest.readers import VisiumReader +from hest.readers import VisiumReader, pool_bins_visiumhd from hest.HESTData import ensembl_id_to_gene from hest.utils import load_image @@ -37,6 +37,59 @@ def download_hest(patterns, local_dir): zip_ref.extractall(seg_dir) +class TestVisiumHDPooling(unittest.TestCase): + def test_count_conservation(self): + from itertools import product + from tempfile import TemporaryDirectory + + import anndata as ad + import numpy as np + import pandas as pd + from scipy.sparse import csr_matrix, csc_matrix, issparse + + # Each destination has repeated indices and a final zero-count source bin. + counts = np.array([ + [1, 10, 0], [2, 20, 0], [3, 30, 0], [0, 0, 0], + [4, 40, 0], [5, 50, 0], [6, 60, 0], [0, 0, 0], + ], dtype=np.float32) + obs = pd.DataFrame({ + 'pxl_col_in_fullres': [8, 24, 8, 24, 136, 152, 136, 152], + 'pxl_row_in_fullres': [8, 8, 24, 24, 8, 8, 24, 24], + }, index=[str(i) for i in range(8)]) + genes = ['gene_b', 'gene_a', 'zero_gene'] + orders = [np.arange(8), np.arange(8)[::-1], np.array([0, 4, 1, 5, 2, 6, 3, 7])] + + with TemporaryDirectory() as tmp: + for storage, backed, chunk_len, order in product( + (np.array, csr_matrix, csc_matrix), (False, True), (1, 2, 3, 50000), orders + ): + with self.subTest(storage=storage.__name__, backed=backed, + chunk_len=chunk_len, order=order.tolist()): + adata = ad.AnnData( + storage(counts[order]), obs=obs.iloc[order].copy(), + var=pd.DataFrame(index=genes), + ) + if backed: + path = os.path.join(tmp, 'counts.h5ad') + adata.write_h5ad(path) + adata = ad.read_h5ad(path, backed='r') + try: + pooled = pool_bins_visiumhd(adata, pixel_size=1.0, chunk_len=chunk_len) + np.testing.assert_array_equal(pooled.X, [[6, 60, 0], [15, 150, 0]]) + self.assertEqual(pooled.X.sum(), counts.sum()) + self.assertEqual(pooled.var_names.tolist(), genes) + np.testing.assert_array_equal(pooled.obsm['spatial'], [[72, 72], [200, 72]]) + np.testing.assert_array_equal(pooled.obs['array_col'], [0, 1]) + np.testing.assert_array_equal(pooled.obs['array_row'], [0, 0]) + remaining = adata.X[:] + if issparse(remaining): + remaining = remaining.toarray() + np.testing.assert_array_equal(remaining, counts[order]) + finally: + if backed: + adata.file.close() + + class TestHESTReader(unittest.TestCase): @classmethod @@ -227,6 +280,7 @@ def test_saving(self): loader = unittest.TestLoader() suite = loader.loadTestsFromTestCase(TestHESTData) + suite.addTests(loader.loadTestsFromTestCase(TestVisiumHDPooling)) # suite = unittest.TestSuite() #suite.addTest(TestHESTData('test_spatialdata')) result = unittest.TextTestRunner(verbosity=2).run(suite)