fix reading anndatas with .raw attributes (closes #142)

(cherry picked from commit 9783f74)

Author: ilia-kats

CHANGELOG.md

@@ -17,6 +17,12 @@ and this project adheres to [Semantic Versioning][].
 - The settings API has changed. Use e.g. `mudata.settings.pull_on_update = True` instead of `mudata.set_options(pull_on_update=True)` and use
   `mudata.settings.override` as context manager for local settings overrides.
 
+## [0.3.7]
+
+### Fixed
+
+- MuData objects with modalities containing `.raw` attributes can once again be read from disk.
+
 ## [0.3.6]
 
 ### Fixed
@@ -166,6 +172,7 @@ To copy the annotations explicitly, you will need to use `pull_obs()` and/or `pu
 Initial `mudata` release with `MuData`, previously a part of the `muon` framework.
 
 [0.4.0]: https://github.com/scverse/mudata/releases/tag/v0.4.0
+[0.3.7]: https://github.com/scverse/mudata/releases/tag/v0.3.7
 [0.3.6]: https://github.com/scverse/mudata/releases/tag/v0.3.6
 [0.3.5]: https://github.com/scverse/mudata/releases/tag/v0.3.5
 [0.3.4]: https://github.com/scverse/mudata/releases/tag/v0.3.4

src/mudata/_core/io.py

@@ -67,7 +67,12 @@ def _write_h5mu(file: h5py.File, mdata: MuData, write_data=True, **kwargs):
         if write_data or not adata.isbacked:
             write_elem(group, "X", adata.X, dataset_kwargs=kwargs)
         if adata.raw is not None:
-            write_elem(group, "raw", adata.raw)
+            if not adata.isbacked:
+                write_elem(group, "raw", adata.raw, dataset_kwargs=kwargs)
+            else:
+                rawgrp = group.require_group("raw")
+                write_elem(rawgrp, "var", adata.raw.var, dataset_kwargs=kwargs)
+                write_elem(rawgrp, "varm", dict(adata.raw.varm), dataset_kwargs=kwargs)
 
         write_elem(group, "obs", adata.obs, dataset_kwargs=kwargs)
         write_elem(group, "var", adata.var, dataset_kwargs=kwargs)
@@ -168,7 +173,12 @@ def write_zarr(
                 else:
                     write_elem(group, "X", adata.X, dataset_kwargs=kwargs)
             if adata.raw is not None:
-                write_elem(group, "raw", adata.raw)
+                if write_data or not adata.isbacked:
+                    write_elem(group, "raw", adata.raw, dataset_kwargs=kwargs)
+                else:
+                    rawgrp = group.require_group("raw")
+                    write_elem(rawgrp, "var", adata.raw.var, dataset_kwargs=kwargs)
+                    write_elem(rawgrp, "varm", dict(adata.raw.varm), dataset_kwargs=kwargs)
 
             write_elem(group, "obs", adata.obs, dataset_kwargs=kwargs)
             write_elem(group, "var", adata.var, dataset_kwargs=kwargs)

tests/conftest.py

@@ -34,7 +34,7 @@ def rng() -> np.random.Generator:
 
 
 @pytest.fixture
-def mdata(rng: np.random.Generator, request: pytest.FixtureRequest):
+def mdata(rng: np.random.Generator, request: pytest.FixtureRequest) -> MuData:
     axis = getattr(request, "param", 0)
     mod1 = AnnData(
         np.arange(0, 200, 0.1).reshape(-1, 20), obs=pd.DataFrame(index=rng.choice(150, size=100, replace=False))
@@ -46,6 +46,8 @@ def mdata(rng: np.random.Generator, request: pytest.FixtureRequest):
     mod2.var["assert-bool"] = False
     mod1.var["assert-boolean-1"] = True
     mod2.var["assert-boolean-2"] = False
+
+    mod1.raw = mod1[:, :10].copy()
     mods = {"mod2": mod2, "mod1": mod1}
 
     attr = "obs" if axis == 0 else "var"

tests/test_view_copy.py

@@ -191,12 +191,14 @@ def _test_view_after_setattr(mdata: md.MuData, mdata_ref: md.MuData, skip: Seque
 
     for k, v in mdata.obsm.items():
         if k in mdata_ref.obsm:
-            assert v.shape[1] == mdata_ref.obsm[k].shape[1]
+            if v.ndim > 1:
+                assert v.shape[1] == mdata_ref.obsm[k].shape[1]
             if isinstance(v, pd.DataFrame):
                 assert (v.columns == mdata_ref.obsm[k].columns).all()
     for k, v in mdata.varm.items():
         if k in mdata_ref.varm:
-            assert v.shape[1] == mdata_ref.obsm[k].shape[1]
+            if v.ndim > 1:
+                assert v.shape[1] == mdata_ref.obsm[k].shape[1]
             if isinstance(v, pd.DataFrame):
                 assert (v.columns == mdata_ref.varm[k].columns).all()