Implement modality structure representation

Resolves #44

Author: gtca

mudata/_core/io.py

@@ -19,7 +19,7 @@
 from pathlib import Path
 from scipy import sparse
 
-from mudata import MuData
+from .mudata import ModDict, MuData
 from .file_backing import MuDataFileManager, AnnDataFileManager
 
 #
@@ -374,7 +374,7 @@ def read_h5mu(filename: PathLike, backed: Union[str, bool, None] = None):
             if k in ["obs", "var"]:
                 d[k] = read_dataframe(f[k])
             if k == "mod":
-                mods = {}
+                mods = ModDict()
                 gmods = f[k]
                 for m in gmods.keys():
                     ad = _read_h5mu_mod(gmods[m], manager, backed not in (None, False))

mudata/_core/mudata.py

@@ -50,6 +50,74 @@ class MuAxisArrays(AxisArrays):
     _view_class = MuAxisArraysView
 
 
+class ModDict(dict):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    def _repr_hierarchy(
+        self, nest_level: int = 0, is_last: bool = False, active_levels: Optional[List[int]] = None
+    ) -> str:
+        descr = ""
+        active_levels = active_levels or []
+        for i, kv in enumerate(self.items()):
+            k, v = kv
+            indent = ("   " * nest_level) + ("└─ " if i == len(self) - 1 else "├─ ")
+
+            if len(active_levels) > 0:
+                indent_list = list(indent)
+                for level in active_levels:
+                    indent_list[level * 3] = "│"
+                indent = "".join(indent_list)
+
+            is_view = " view" if v.is_view else ""
+            backed_at = f" backed at {str(v.filename)!r}" if v.isbacked else ""
+
+            if isinstance(v, MuData):
+                maybe_axis = (
+                    (
+                        f" [shared obs] "
+                        if v.axis == 0
+                        else f" [shared var] "
+                        if v.axis == 1
+                        else f" [shared obs and var] "
+                    )
+                    if hasattr(v, "axis")
+                    else ""
+                )
+                descr += (
+                    f"\n{indent}{k} MuData{maybe_axis}({v.n_obs} × {v.n_vars}){backed_at}{is_view}"
+                )
+
+                if i != len(self) - 1:
+                    levels = [nest_level] + [level for level in active_levels]
+                else:
+                    levels = [level for level in active_levels if level != nest_level]
+                descr += v.mod._repr_hierarchy(nest_level=nest_level + 1, active_levels=levels)
+            elif isinstance(v, AnnData):
+                descr += f"\n{indent}{k} AnnData ({v.n_obs} x {v.n_vars}){backed_at}{is_view}"
+            else:
+                continue
+
+        return descr
+
+    def __repr__(self) -> str:
+        """
+        Represent the hierarchy of the modalities in the object.
+
+        A MuData object with two modalities, protein and RNA,
+        with the latter being a MuData containing raw, QC'ed and hvg-filtered AnnData objects,
+        will be represented as:
+
+        root MuData (axis=0) (5000 x 20100)
+        ├── protein AnnData (5000 x 100)
+        └── rna MuData (axis=-1) (5000 x 20000)
+            ├── raw AnnData (5000 x 20000)
+            ├── quality-filtered AnnData (3000 x 20000)
+            └── hvg-filtered AnnData (3000 x 4000)
+        """
+        return "MuData" + self._repr_hierarchy()
+
+
 class MuData:
     """
     Multimodal data object
@@ -81,7 +149,7 @@ def __init__(
             return
 
         # Add all modalities to a MuData object
-        self.mod = dict()
+        self.mod = ModDict()
         if isinstance(data, abc.Mapping):
             for k, v in data.items():
                 self.mod[k] = v
@@ -185,7 +253,7 @@ def _init_as_view(self, mudata_ref: "MuData", index):
         if isinstance(varidx, Integral):
             varidx = slice(varidx, varidx + 1)
 
-        self.mod = dict()
+        self.mod = ModDict()
         for m, a in mudata_ref.mod.items():
             cobsidx, cvaridx = mudata_ref.obsmap[m][obsidx], mudata_ref.varmap[m][varidx]
             cobsidx, cvaridx = cobsidx[cobsidx > 0] - 1, cvaridx[cvaridx > 0] - 1
@@ -1239,7 +1307,17 @@ def _gen_repr(self, n_obs, n_vars, extensive: bool = False, nest_level: int = 0)
         indent = "    " * nest_level
         backed_at = f" backed at {str(self.filename)!r}" if self.isbacked else ""
         view_of = "View of " if self.is_view else ""
-        maybe_axis = f" (axis={self.axis}) " if hasattr(self, "axis") and self.axis != 0 else ""
+        maybe_axis = (
+            (
+                f" (shared obs) "
+                if self.axis == 0
+                else f" (shared var) "
+                if self.axis == 1
+                else f" (shared obs and var) "
+            )
+            if hasattr(self, "axis")
+            else ""
+        )
         descr = f"{view_of}MuData object with n_obs × n_vars = {n_obs} × {n_vars}{maybe_axis}{backed_at}"
         for attr in ["obs", "var", "uns", "obsm", "varm", "obsp", "varp"]:
             if hasattr(self, attr) and getattr(self, attr) is not None:

tests/test_repr.py

@@ -0,0 +1,46 @@
+import unittest
+import pytest
+
+import numpy as np
+from anndata import AnnData
+from mudata import MuData
+
+
+# Dimensions
+N = 100
+D1, D2 = 10, 20
+D = D1 + D2
+
+
+@pytest.fixture()
+def mdata():
+    mod1 = AnnData(np.arange(0, 100, 0.1).reshape(-1, D1))
+    mod1.obs_names = [f"obs{i}" for i in range(mod1.n_obs)]
+    mod1.var_names = [f"var{i}" for i in range(D1)]
+
+    mod21 = AnnData(np.arange(3101, 5101, 1).reshape(-1, D2))
+    mod22 = AnnData(np.arange(3101, 5101, 1).reshape(-1, D2))
+    # Same obs_names and var_names
+    mod21.obs_names = mod1.obs_names.copy()
+    mod22.obs_names = mod1.obs_names.copy()
+    mod21.var_names = [f"var{i}" for i in range(D1, D)]
+    mod22.var_names = [f"var{i}" for i in range(D1, D)]
+    mod2 = MuData({"mod21": mod21, "mod22": mod22}, axis=-1)
+
+    mdata = MuData({"mod1": mod1, "mod2": mod2})
+    yield mdata
+
+
+class TestMuData:
+    def test_nested_mudata(self, mdata):
+        assert mdata.shape == (N, D)
+        assert mdata["mod1"].shape == (N, D1)
+        assert mdata["mod2"].shape == (N, D2)
+        assert mdata.axis == 0
+        assert mdata["mod2"].axis == -1
+
+    def test_mod_repr(self, mdata):
+        assert (
+            mdata.mod.__repr__()
+            == f"MuData\n├─ mod1 AnnData ({N} x {D1})\n└─ mod2 MuData [shared obs and var] ({N} × 20)\n   ├─ mod21 AnnData ({N} x {D2})\n   └─ mod22 AnnData ({N} x {D2})"
+        )