From 53a717f615f9f4ef0d2c60c1acfd8bc27a0092c0 Mon Sep 17 00:00:00 2001
From: ArneD <arnedefauw@gmail.com>
Date: Mon, 6 Jan 2025 13:04:31 +0100
Subject: [PATCH] SC_42 visium hd io (#80)

---
 src/harpy/io/_visium_hd.py | 30 ++++++++++++++++++++----------
 1 file changed, 20 insertions(+), 10 deletions(-)

diff --git a/src/harpy/io/_visium_hd.py b/src/harpy/io/_visium_hd.py
index 3b1f92d..845e0ad 100644
--- a/src/harpy/io/_visium_hd.py
+++ b/src/harpy/io/_visium_hd.py
@@ -52,14 +52,13 @@ 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,
@@ -67,13 +66,24 @@ def visium_hd(
             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)