Repository: scverse/muon

Issue:
filter_obs is broken with latest AnnData
**Describe the bug**
Running `mu.pp.filter_obs` results in a `ValueError` with the latest AnnData

**To Reproduce**
```python
import scanpy as sc
import muon as mu

adata = sc.datasets.pbmc3k_processed()
mdata = mu.MuData({
    "A": adata[:500, ].copy(),
    "B": adata[500:, ].copy()
})
mu.pp.filter_obs(mdata, "A:louvain", lambda x: x == "B cells")
```

```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[22], [line 1](vscode-notebook-cell:?execution_count=22&line=1)
----> [1](vscode-notebook-cell:?execution_count=22&line=1) mu.pp.filter_obs(mdata, "A:louvain", lambda x: x == "B cells")

File ~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:763, in filter_obs(data, var, func)
    [761](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:761) obsmap = data.obsmap[m][obs_subset]
    [762](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:762) obsidx = obsmap > 0
--> [763](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:763) filter_obs(mod, mod.obs_names[obsmap[obsidx] - 1])
    [764](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:764) maporder = np.argsort(obsmap[obsidx])
    [765](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:765) nobsmap = np.empty(maporder.size)

File ~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:728, in filter_obs(data, var, func)
    [725](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:725) data._n_obs = data.obs.shape[0]
    [727](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:727) # Subset .obsm
--> [728](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:728) for k, v in data.obsm.items():
    [729](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:729)     data.obsm[k] = v[obs_subset]
    [731](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/muon/_core/preproc.py:731) # Subset .obsp

File ~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:422, in AlignedMappingProperty.__get__(self, obj, objtype)
    [420](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:420)     return self  # type: ignore
    [421](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:421) if not obj.is_view:
--> [422](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:422)     return self.construct(obj, store=getattr(obj, f"_{self.name}"))
    [423](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:423) parent_anndata = obj._adata_ref
    [424](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:424) idxs = (obj._oidx, obj._vidx)
...
---> [97](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:97)     raise ValueError(msg)
     [99](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:99) name = f"{self.attrname.title().rstrip('s')} {key!r}"
    [100](https://file+.vscode-resource.vscode-cdn.net/home/sturm/projects/2020/scirpy/docs/tutorials/~/apps/micromamba/envs/test_scirpy/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:100) return coerce_array(val, name=name, allow_df=self._allow_df)

ValueError: Value passed for key 'X_pca' is of incorrect shape. Values of obsm must match dimensions ('obs',) of parent. Value had shape (500,) while it should have had (69,).
```

**Expected behaviour**
Filter the mudata object as this has worked previously

**System**
```
anndata==0.10.9
mudata==0.3.1
muon==0.1.6
scanpy==1.10.2
```




Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 81d184d..32b5732 100644
--- a/.gitignore
+++ b/.gitignore
@@ -6,6 +6,9 @@ __pycache__/
 # C extensions
 *.so
 
+# cached data
+data/
+
 # Distribution / packaging
 .Python
 build/
diff --git a/muon/_core/preproc.py b/muon/_core/preproc.py
index 974cbae..04c0735 100644
--- a/muon/_core/preproc.py
+++ b/muon/_core/preproc.py
@@ -657,23 +657,24 @@ def intersect_obs(mdata: MuData):
     return
 
 
-# Utility functions: filtering observations
+# Utility functions: filtering observations or variables
 
 
-def filter_obs(
-    data: Union[AnnData, MuData], var: Union[str, Sequence[str]], func: Optional[Callable] = None
+def _filter_attr(
+    data: Union[AnnData, MuData],
+    attr: Literal["obs", "var"],
+    key: Union[str, Sequence[str]],
+    func: Optional[Callable] = None,
 ) -> None:
     """
-    Filter observations (samples or cells) in-place
-    using any column in .obs or in .X.
+    Filter observations or variables in-place.
 
     Parameters
     ----------
     data: AnnData or MuData
             AnnData or MuData object
-    var: str or Sequence[str]
-            Column name in .obs or in .X to be used for filtering.
-            Alternatively, obs_names can be provided directly.
+    key: str or Sequence[str]
+            Names or key to filter
     func
             Function to apply to the variable used for filtering.
             If the variable is of type boolean and func is an identity function,
@@ -694,51 +695,76 @@ def filter_obs(
                 "MuData object is backed. The requested subset of the .X matrices of its modalities will be read into memory, and the object will not be backed anymore."
             )
 
-    if isinstance(var, str):
-        if var in data.obs.columns:
+    assert attr in ("obs", "var"), "Attribute has to be either 'obs' or 'var'."
+
+    df = getattr(data, attr)
+    names = getattr(data, f"{attr}_names")
+    other = "obs" if attr == "var" else "var"
+    other_names = getattr(data, f"{other}_names")
+    attrm = getattr(data, f"{attr}m")
+    attrp = getattr(data, f"{attr}p")
+
+    if isinstance(key, str):
+        if key in df.columns:
             if func is None:
-                if data.obs[var].dtypes.name == "bool":
+                if df[key].dtypes.name == "bool":
 
                     def func(x):
                         return x
 
                 else:
-                    raise ValueError(f"Function has to be provided since {var} is not boolean")
-            obs_subset = func(data.obs[var].values)
-        elif var in data.var_names:
-            obs_subset = func(data.X[:, np.where(data.var_names == var)[0]].reshape(-1))
+                    raise ValueError(f"Function has to be provided since {key} is not boolean")
+            subset = func(df[key].values)
+        elif key in other_names:
+            if attr == "obs":
+                subset = func(data.X[:, np.where(other_names == key)[0]].reshape(-1))
+            else:
+                subset = func(data.X[np.where(other_names == key)[0], :].reshape(-1))
         else:
             raise ValueError(
-                f"Column name from .obs or one of the var_names was expected but got {var}."
+                f"Column name from .{attr} or one of the {other}_names was expected but got {key}."
             )
     else:
         if func is None:
-            if np.array(var).dtype == bool:
-                obs_subset = np.array(var)
+            if np.array(key).dtype == bool:
+                subset = np.array(key)
             else:
-                obs_subset = data.obs_names.isin(var)
+                subset = names.isin(key)
         else:
-            raise ValueError("When providing obs_names directly, func has to be None.")
+            raise ValueError(f"When providing {attr}_names directly, func has to be None.")
 
-    # Subset .obs
-    data._obs = data.obs[obs_subset]
-    data._n_obs = data.obs.shape[0]
+    if isinstance(data, AnnData):
+        # Collect elements to subset
+        # NOTE: accessing them after subsetting .obs/.var
+        # will fail due to _validate_value()
+        attrm = dict(attrm)
+        attrp = dict(attrp)
 
-    # Subset .obsm
-    for k, v in data.obsm.items():
-        data.obsm[k] = v[obs_subset]
+        # Subset .obs/.var
+        setattr(data, f"_{attr}", df[subset])
 
-    # Subset .obsp
-    for k, v in data.obsp.items():
-        data.obsp[k] = v[obs_subset][:, obs_subset]
+        # Subset .obsm/.varm
+        for k, v in attrm.items():
+            attrm[k] = v[subset]
+        setattr(data, f"{attr}m", attrm)
+
+        # Subset .obsp/.obsp
+        for k, v in attrp.items():
+            attrp[k] = v[subset][:, subset]
+        setattr(data, f"{attr}p", attrp)
 
-    if isinstance(data, AnnData):
         # Subset .X
         if data._X is not None:
             try:
-                data._X = data.X[obs_subset, :]
+                if attr == "obs":
+                    data._X = data.X[subset, :]
+                else:
+                    data._X = data.X[:, subset]
             except TypeError:
-                data._X = data.X[np.where(obs_subset)[0], :]
+                if attr == "obs":
+                    data._X = data.X[np.where(subset)[0], :]
+                else:
+                    data._X = data.X[:, np.where(subset)[0]]
                 # For some h5py versions, indexing arrays must have integer dtypes
                 # https://github.com/h5py/h5py/issues/1847
 
@@ -748,29 +774,71 @@ def filter_obs(
 
         # Subset layers
         for layer in data.layers:
-            data.layers[layer] = data.layers[layer][obs_subset, :]
+            if attr == "obs":
+                data.layers[layer] = data.layers[layer][subset, :]
+            else:
+                data.layers[layer] = data.layers[layer][:, subset]
 
-        # Subset raw
-        if data.raw is not None:
-            data.raw._X = data.raw.X[obs_subset, :]
-            data.raw._n_obs = data.raw.X.shape[0]
+        # Subset raw - only when subsetting obs
+        if attr == "obs" and data.raw is not None:
+            data.raw._X = data.raw.X[subset, :]
 
     else:
-        # filter_obs() for each modality
+        attrmap = getattr(data, f"{attr}map")
+
+        # Subset .obs/.var
+        setattr(data, f"_{attr}", df[subset])
+
+        # Subset .obsm/.varm
+        for k, v in attrm.items():
+            attrm[k] = v[subset]
+        setattr(data, f"{attr}m", attrm)
+
+        # Subset .obsp/.varp
+        for k, v in attrp.items():
+            attrp[k] = v[subset][:, subset]
+        setattr(data, f"{attr}p", attrp)
+
+        # _filter_attr() for each modality
         for m, mod in data.mod.items():
-            obsmap = data.obsmap[m][obs_subset]
-            obsidx = obsmap > 0
-            filter_obs(mod, mod.obs_names[obsmap[obsidx] - 1])
-            maporder = np.argsort(obsmap[obsidx])
+            map_subset = attrmap[m][subset]
+            attridx = map_subset > 0
+            orig_attr = getattr(mod, attr).copy()
+            mod_names = getattr(mod, f"{attr}_names")
+            _filter_attr(mod, attr, mod_names[map_subset[attridx] - 1])
+            data.mod[m]._remove_unused_categories(orig_attr, getattr(mod, attr), mod.uns)
+            maporder = np.argsort(map_subset[attridx])
             nobsmap = np.empty(maporder.size)
             nobsmap[maporder] = np.arange(1, maporder.size + 1)
-            obsmap[obsidx] = nobsmap
-            data.obsmap[m] = obsmap
+            map_subset[attridx] = nobsmap
+            getattr(data, f"{attr}map")[m] = map_subset
 
     return
 
 
-# Utility functions: filtering variables
+def filter_obs(
+    data: Union[AnnData, MuData], var: Union[str, Sequence[str]], func: Optional[Callable] = None
+) -> None:
+    """
+    Filter observations (samples or cells) in-place
+    using any column in .obs or in .X.
+
+    Parameters
+    ----------
+    data: AnnData or MuData
+            AnnData or MuData object
+    var: str or Sequence[str]
+            Column name in .obs or in .X to be used for filtering.
+            Alternatively, obs_names can be provided directly.
+    func
+            Function to apply to the variable used for filtering.
+            If the variable is of type boolean and func is an identity function,
+            the func argument can be omitted.
+    """
+
+    _filter_attr(data, "obs", var, func)
+
+    return
 
 
 def filter_var(
@@ -793,84 +861,7 @@ def filter_var(
             the func argument can be omitted.
     """
 
-    if data.is_view:
-        raise ValueError(
-            "The provided adata is a view. In-place filtering does not operate on views."
-        )
-    if data.isbacked:
-        if isinstance(data, AnnData):
-            warnings.warn(
-                "AnnData object is backed. The requested subset of the matrix .X will be read into memory, and the object will not be backed anymore."
-            )
-        else:
-            warnings.warn(
-                "MuData object is backed. The requested subset of the .X matrices of its modalities will be read into memory, and the object will not be backed anymore."
-            )
-
-    if isinstance(var, str):
-        if var in data.var.columns:
-            if func is None:
-                if data.var[var].dtypes.name == "bool":
-
-                    def func(x):
-                        return x
-
-                else:
-                    raise ValueError(f"Function has to be provided since {var} is not boolean")
-            var_subset = func(data.var[var].values)
-        elif var in data.obs_names:
-            var_subset = func(data.X[:, np.where(data.obs_names == var)[0]].reshape(-1))
-        else:
-            raise ValueError(
-                f"Column name from .var or one of the obs_names was expected but got {var}."
-            )
-    else:
-        if func is None:
-            var_subset = var if np.array(var).dtype == bool else data.var_names.isin(var)
-        else:
-            raise ValueError("When providing var_names directly, func has to be None.")
-
-    # Subset .var
-    data._var = data.var[var_subset]
-    data._n_vars = data.var.shape[0]
-
-    # Subset .varm
-    for k, v in data.varm.items():
-        data.varm[k] = v[var_subset]
-
-    # Subset .varp
-    for k, v in data.varp.items():
-        data.varp[k] = v[var_subset][:, var_subset]
-
-    if isinstance(data, AnnData):
-        # Subset .X
-        try:
-            data._X = data.X[:, var_subset]
-        except TypeError:
-            data._X = data.X[:, np.where(var_subset)[0]]
-            # For some h5py versions, indexing arrays must have integer dtypes
-            # https://github.com/h5py/h5py/issues/1847
-        if data.isbacked:
-            data.file.close()
-            data.filename = None
-
-        # Subset layers
-        for layer in data.layers:
-            data.layers[layer] = data.layers[layer][:, var_subset]
-
-        # NOTE: .raw is not subsetted
-
-    else:
-        # filter_var() for each modality
-        for m, mod in data.mod.items():
-            varmap = data.varmap[m][var_subset]
-            varidx = varmap > 0
-            filter_var(mod, mod.var_names[varmap[varidx] - 1])
-            maporder = np.argsort(varmap[varidx])
-            nvarmap = np.empty(maporder.size)
-            nvarmap[maporder] = np.arange(1, maporder.size + 1)
-            varmap[varidx] = nvarmap
-            data.varmap[m] = varmap
+    _filter_attr(data, "var", var, func)
 
     return
 
diff --git a/pyproject.toml b/pyproject.toml
index d322dd7..75f519a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -49,6 +49,7 @@ atac = [
 test = [
     "pytest",
     "flake8",
+    "pytest",
 ]
 
 [tool.flit.metadata.urls]


Test patch:
diff --git a/tests/conftest.py b/tests/conftest.py
index d9ab19a..4369a0d 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -1,4 +1,5 @@
 import pytest
+import scanpy as sc
 
 
 @pytest.fixture(scope="module")
@@ -9,3 +10,8 @@ def filepath_h5mu(tmpdir_factory):
 @pytest.fixture(scope="module")
 def filepath_hdf5(tmpdir_factory):
     yield str(tmpdir_factory.mktemp("tmp_mofa_dir").join("mofa_pytest.hdf5"))
+
+
+@pytest.fixture(scope="module")
+def pbmc3k_processed():
+    yield sc.datasets.pbmc3k_processed()
diff --git a/tests/test_muon_preproc.py b/tests/test_muon_preproc.py
index b4e388c..8a4ed99 100644
--- a/tests/test_muon_preproc.py
+++ b/tests/test_muon_preproc.py
@@ -1,12 +1,11 @@
 import unittest
 import pytest
 
-import os
 from functools import reduce
 
 import numpy as np
-from scipy.sparse import csr_matrix
 from anndata import AnnData
+from anndata.tests.helpers import assert_equal
 from mudata import MuData
 import muon as mu
 
@@ -83,6 +82,21 @@ class TestInPlaceFiltering:
             sub = np.random.binomial(1, 0.5, view.n_obs).astype(bool)
             mu.pp.filter_obs(view, sub)
 
+    def test_filter_obs_with_obsm_obsp(self, pbmc3k_processed):
+        A = pbmc3k_processed[:500,].copy()
+        B = pbmc3k_processed[500:,].copy()
+        A_subset = A[A.obs["louvain"] == "B cells"].copy()
+        B_subset = B[B.obs["louvain"] == "B cells"].copy()
+        mdata = mu.MuData({"A": A, "B": B}, axis=1)
+        mdata.pull_obs("louvain")
+        mu.pp.filter_obs(mdata, "louvain", lambda x: x == "B cells")
+        assert mdata["B"].n_obs == B_subset.n_obs
+        assert mdata["A"].obs["louvain"].unique() == "B cells"
+        assert B.n_obs == B_subset.n_obs
+        assert A.obs["louvain"].unique() == "B cells"
+        assert_equal(mdata["A"], A_subset)
+        assert_equal(mdata["B"], B_subset)
+
     # Variables
 
     def test_filter_var_adata(self, mdata, filepath_h5mu):
@@ -132,6 +146,23 @@ class TestInPlaceFiltering:
             sub = np.random.binomial(1, 0.5, view.n_vars).astype(bool)
             mu.pp.filter_var(view, sub)
 
+    def test_filter_var_with_varm_varp(self, pbmc3k_processed):
+        A = pbmc3k_processed[:, :500].copy()
+        B = pbmc3k_processed[:, 500:].copy()
+        np.random.seed(42)
+        A_var_sel = np.random.choice(np.array([0, 1]), size=A.n_vars, replace=True)
+        B_var_sel = np.random.choice(np.array([0, 1]), size=B.n_vars, replace=True)
+        A.var["sel"] = A_var_sel
+        B.var["sel"] = B_var_sel
+        A_subset = A[:, A_var_sel == 1].copy()
+        B_subset = B[:, B_var_sel == 1].copy()
+        mdata = mu.MuData({"A": A, "B": B})
+        mdata.pull_var("sel")
+        mu.pp.filter_var(mdata, "sel", lambda y: y == 1)
+        assert mdata.shape[1] == int(np.sum(A_var_sel) + np.sum(B_var_sel))
+        assert_equal(mdata["A"], A_subset)
+        assert_equal(mdata["B"], B_subset)
+
 
 @pytest.mark.usefixtures("filepath_h5mu")
 class TestIntersectObs: