SC_42 visium hd io (#80)

saeyslab · Jan 6, 2025 · 53a717f · 53a717f
1 parent 052a3b3
commit 53a717f
Showing 1 changed file with 20 additions and 10 deletions.
diff --git a/src/harpy/io/_visium_hd.py b/src/harpy/io/_visium_hd.py
@@ -52,28 +52,38 @@ def visium_hd(
         bins_as_squares=bins_as_squares,
     )
 
-    for tables_layer in [*sdata.tables]:
-        adata = sdata[tables_layer]
+    for table_layer in [*sdata.tables]:
+        adata = sdata[table_layer]
         adata.var_names_make_unique()
         adata.X = adata.X.tocsc()
 
-        adata.obs.rename(
-            columns={VisiumHDKeys.REGION_KEY: _REGION_KEY, VisiumHDKeys.INSTANCE_KEY: _INSTANCE_KEY}, inplace=True
-        )
+        _old_instance_key = sdata[table_layer].uns[TableModel.ATTRS_KEY][TableModel.INSTANCE_KEY]
+        adata.obs.rename(columns={VisiumHDKeys.REGION_KEY: _REGION_KEY, _old_instance_key: _INSTANCE_KEY}, inplace=True)
         adata.uns.pop(TableModel.ATTRS_KEY)
         adata = TableModel.parse(
             adata,
             region_key=_REGION_KEY,
             region=adata.obs[_REGION_KEY].cat.categories.to_list(),
             instance_key=_INSTANCE_KEY,
         )
+        # get the shapes layer for this table layer
+        for _shapes_layer in [*sdata.shapes]:
+            if table_layer in _shapes_layer:
+                shapes_layer = _shapes_layer
+                break
+        assert (
+            len(sdata[shapes_layer]) == len(adata)
+        ), f"Shapes layer containing bins '{shapes_layer}' and corresponding table '{table_layer}' should have same length."
+        sdata[shapes_layer].index = (
+            adata.obs.set_index(VisiumHDKeys.INSTANCE_KEY).loc[sdata[shapes_layer].index, _INSTANCE_KEY].values
+        )
+        if VisiumHDKeys.INSTANCE_KEY in adata.obs.columns:
+            adata.obs.drop(columns=VisiumHDKeys.INSTANCE_KEY, inplace=True)
+        sdata[shapes_layer].index.name = _INSTANCE_KEY
 
-        del sdata[tables_layer]
-
-        sdata[tables_layer] = adata
+        del sdata[table_layer]
 
-    for shapes_layer in [*sdata.shapes]:
-        sdata[shapes_layer].index.name = _INSTANCE_KEY
+        sdata[table_layer] = adata
 
     if output is not None:
         sdata.write(output)