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

Author: ilia-kats

CHANGELOG.md

@@ -8,6 +8,12 @@ and this project adheres to [Semantic Versioning][].
 [keep a changelog]: https://keepachangelog.com/en/1.1.0/
 [semantic versioning]: https://semver.org/spec/v2.0.0.html
 
+## [0.3.7]
+
+### Fixed
+
+- MuData objects with modalities containing `.raw` attributes can once again be read from disk.
+
 ## [0.3.6]
 
 ### Fixed
@@ -156,6 +162,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.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

@@ -82,7 +82,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)
@@ -201,7 +206,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)
@@ -527,16 +537,16 @@ def _read_zarr_mod(g: zarr.Group, manager: MuDataFileManager = None, backed: boo
                 d["X"] = read_elem(X)
         elif k != "raw":
             d[k] = read_elem(g[k])
-    ad = AnnData(**d)
+    adata = AnnData(**d)
     if manager is not None:
-        ad.file = AnnDataFileManager(ad, Path(g.name).name, manager)
+        adata.file = AnnDataFileManager(adata, Path(g.name).name, manager)
 
     raw = _read_legacy_raw(
         g, d.get("raw"), read_zarr_dataframe, read_elem, attrs=("var", "varm") if backed else ("var", "varm", "X")
     )
     if raw:
-        ad._raw = ad.Raw(ad, **raw)
-    return ad
+        adata._raw = ad.Raw(adata, **raw)
+    return adata
 
 
 def _read_h5mu_mod(g: h5py.Group, manager: MuDataFileManager = None, backed: bool = False) -> dict:
@@ -551,14 +561,14 @@ def _read_h5mu_mod(g: h5py.Group, manager: MuDataFileManager = None, backed: boo
                 d["X"] = read_elem(X)
         elif k != "raw":
             d[k] = read_elem(g[k])
-    ad = AnnData(**d)
+    adata = AnnData(**d)
     if manager is not None:
-        ad.file = AnnDataFileManager(ad, Path(g.name).name, manager)
+        adata.file = AnnDataFileManager(adata, Path(g.name).name, manager)
 
     raw = _read_raw(g, attrs=("var", "varm") if backed else ("var", "varm", "X"))
     if raw:
-        ad._raw = ad.Raw(ad, **raw)
-    return ad
+        adata._raw = ad.Raw(adata, **raw)
+    return adata
 
 
 def read_h5ad(

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()