Repository: scverse/scanpy

Issue:
aggregate throws error when aggregating `obsm` or `varm`
### Please make sure these conditions are met

- [X] I have checked that this issue has not already been reported.
- [X] I have confirmed this bug exists on the latest version of scanpy.
- [X] (optional) I have confirmed this bug exists on the main branch of scanpy.

### What happened?

aggregate throws error when aggregating `obsm` or `varm`

### Minimal code sample

```python
import scanpy as sc

adata = sc.datasets.pbmc68k_reduced()
sc.get.aggregate(adata, by="louvain", func="mean", obsm="X_umap")
```


### Error output

```pytb
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[3], line 1
----> 1 sc.get.aggregate(pbmc, by="louvain", func="mean", obsm="X_umap")

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/functools.py:909, in singledispatch.<locals>.wrapper(*args, **kw)
    905 if not args:
    906     raise TypeError(f'{funcname} requires at least '
    907                     '1 positional argument')
--> 909 return dispatch(args[0].__class__)(*args, **kw)

File /mnt/workspace/repos/scanpy/scanpy/get/_aggregated.py:272, in aggregate(adata, by, func, axis, mask, dof, layer, obsm, varm)
    264 # Actual computation
    265 layers = aggregate(
    266     data,
    267     by=categorical,
   (...)
    270     dof=dof,
    271 )
--> 272 result = AnnData(
    273     layers=layers,
    274     obs=new_label_df,
    275     var=getattr(adata, "var" if axis == 0 else "obs"),
    276 )
    278 if axis == 1:
    279     return result.T

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/site-packages/anndata/_core/anndata.py:271, in AnnData.__init__(self, X, obs, var, uns, obsm, varm, layers, raw, dtype, shape, filename, filemode, asview, obsp, varp, oidx, vidx)
    269     self._init_as_view(X, oidx, vidx)
    270 else:
--> 271     self._init_as_actual(
    272         X=X,
    273         obs=obs,
    274         var=var,
    275         uns=uns,
    276         obsm=obsm,
    277         varm=varm,
    278         raw=raw,
    279         layers=layers,
    280         dtype=dtype,
    281         shape=shape,
    282         obsp=obsp,
    283         varp=varp,
    284         filename=filename,
    285         filemode=filemode,
    286     )

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/site-packages/anndata/_core/anndata.py:501, in AnnData._init_as_actual(self, X, obs, var, uns, obsm, varm, varp, obsp, raw, layers, dtype, shape, filename, filemode)
    498 self._clean_up_old_format(uns)
    500 # layers
--> 501 self._layers = Layers(self, layers)

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:331, in Layers.__init__(self, parent, vals)
    329 self._data = dict()
    330 if vals is not None:
--> 331     self.update(vals)

File <frozen _collections_abc>:949, in update(self, other, **kwds)

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:199, in AlignedActualMixin.__setitem__(self, key, value)
    198 def __setitem__(self, key: str, value: V):
--> 199     value = self._validate_value(value, key)
    200     self._data[key] = value

File /mnt/workspace/mambaforge/envs/scanpy-dev/lib/python3.11/site-packages/anndata/_core/aligned_mapping.py:89, in AlignedMapping._validate_value(self, val, key)
     83         dims = tuple(("obs", "var")[ax] for ax in self.axes)
     84         msg = (
     85             f"Value passed for key {key!r} is of incorrect shape. "
     86             f"Values of {self.attrname} must match dimensions {dims} of parent. "
     87             f"Value had shape {actual_shape} while it should have had {right_shape}."
     88         )
---> 89     raise ValueError(msg)
     91 if not self._allow_df and isinstance(val, pd.DataFrame):
     92     name = self.attrname.title().rstrip("s")

ValueError: Value passed for key 'mean' is of incorrect shape. Values of layers must match dimensions ('obs', 'var') of parent. Value had shape (11, 2) while it should have had (11, 765).
```


### Versions

<details>

```
-----
anndata     0.10.6
scanpy      1.10.0rc2.dev19+ga6126980
-----
IPython             8.20.0
PIL                 10.2.0
asciitree           NA
asttokens           NA
cffi                1.16.0
cloudpickle         3.0.0
cycler              0.12.1
cython_runtime      NA
dask                2024.3.0
dateutil            2.8.2
decorator           5.1.1
defusedxml          0.7.1
executing           2.0.1
h5py                3.10.0
igraph              0.11.3
jedi                0.19.1
jinja2              3.1.3
joblib              1.3.2
kiwisolver          1.4.5
legacy_api_wrap     NA
leidenalg           0.10.2
llvmlite            0.41.1
markupsafe          2.1.4
matplotlib          3.8.3
mpl_toolkits        NA
msgpack             1.0.7
natsort             8.4.0
numba               0.58.1
numcodecs           0.12.1
numpy               1.26.3
packaging           23.2
pandas              2.2.1
parso               0.8.3
pexpect             4.9.0
prompt_toolkit      3.0.43
psutil              5.9.8
ptyprocess          0.7.0
pure_eval           0.2.2
pyarrow             15.0.1
pygments            2.17.2
pyparsing           3.1.1
pytz                2023.4
scipy               1.12.0
session_info        1.0.0
six                 1.16.0
sklearn             1.4.0
sparse              0.15.1
stack_data          0.6.3
tblib               3.0.0
texttable           1.7.0
threadpoolctl       3.2.0
tlz                 0.12.1
toolz               0.12.1
traitlets           5.14.1
typing_extensions   NA
wcwidth             0.2.13
yaml                6.0.1
zarr                2.17.1
zipp                NA
-----
Python 3.11.7 | packaged by conda-forge | (main, Dec 23 2023, 14:43:09) [GCC 12.3.0]
Linux-5.15.0-101-generic-x86_64-with-glibc2.35
-----
Session information updated at 2024-03-19 13:05
```

</details>


Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/scverse/scanpy/pull/2934?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) Report
Attention: Patch coverage is `85.71429%` with `1 lines` in your changes are missing coverage. Please review.
> Project coverage is 73.22%. Comparing base [(`6542113`)](https://app.codecov.io/gh/scverse/scanpy/commit/6542113d9e7f6a9e1a287aa940ec5564b60a247d?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) to head [(`7b9333c`)](https://app.codecov.io/gh/scverse/scanpy/pull/2934?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##             main    #2934      +/-   ##
==========================================
- Coverage   75.24%   73.22%   -2.02%     
==========================================
  Files         116      116              
  Lines       12802    12771      -31     
==========================================
- Hits         9633     9352     -281     
- Misses       3169     3419     +250     
```


| [Files](https://app.codecov.io/gh/scverse/scanpy/pull/2934?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse) | Coverage Δ | |
|---|---|---|
| [scanpy/get/\_aggregated.py](https://app.codecov.io/gh/scverse/scanpy/pull/2934?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse#diff-c2NhbnB5L2dldC9fYWdncmVnYXRlZC5weQ==) | `94.48% <85.71%> (-0.56%)` | :arrow_down: |

... and [27 files with indirect coverage changes](https://app.codecov.io/gh/scverse/scanpy/pull/2934/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=scverse)

</details>

Patch:
diff --git a/scanpy/get/_aggregated.py b/scanpy/get/_aggregated.py
index ffc04135..6d603287 100644
--- a/scanpy/get/_aggregated.py
+++ b/scanpy/get/_aggregated.py
@@ -267,11 +267,20 @@ def aggregate(
         mask=mask,
         dof=dof,
     )
-    result = AnnData(
-        layers=layers,
-        obs=new_label_df,
-        var=getattr(adata, "var" if axis == 0 else "obs"),
-    )
+
+    # Define new var dataframe
+    if obsm or varm:
+        if isinstance(data, pd.DataFrame):
+            # Check if there could be labels
+            var = pd.DataFrame(index=data.columns)
+        else:
+            # Create them otherwise
+            var = pd.DataFrame(index=pd.RangeIndex(data.shape[1]).astype(str))
+    else:
+        var = getattr(adata, "var" if axis == 0 else "obs")
+
+    # It's all coming together
+    result = AnnData(layers=layers, obs=new_label_df, var=var)
 
     if axis == 1:
         return result.T
@@ -279,6 +288,11 @@ def aggregate(
         return result
 
 
+@aggregate.register(pd.DataFrame)
+def aggregate_df(data, by, func, *, mask=None, dof=1):
+    return aggregate(data.values, by, func, mask=mask, dof=dof)
+
+
 @aggregate.register(np.ndarray)
 @aggregate.register(sparse.spmatrix)
 def aggregate_array(


Test patch:
diff --git a/scanpy/tests/test_aggregated.py b/scanpy/tests/test_aggregated.py
index c3b5331a..d0e01bb6 100644
--- a/scanpy/tests/test_aggregated.py
+++ b/scanpy/tests/test_aggregated.py
@@ -5,7 +5,7 @@ import numpy as np
 import pandas as pd
 import pytest
 from packaging.version import Version
-from scipy.sparse import csr_matrix
+from scipy import sparse
 
 import scanpy as sc
 from scanpy._utils import _resolve_axis
@@ -61,10 +61,10 @@ def gen_adata(data_key, dim, df_base, df_groupby, X):
     obs_df, var_df = (df_groupby, df_base) if dim == "obs" else (df_base, df_groupby)
     data = X.T if dim == "var" and data_key != "varm" else X
     if data_key != "X":
-        data_dict_sparse = {data_key: {"test": csr_matrix(data)}}
+        data_dict_sparse = {data_key: {"test": sparse.csr_matrix(data)}}
         data_dict_dense = {data_key: {"test": data}}
     else:
-        data_dict_sparse = {data_key: csr_matrix(data)}
+        data_dict_sparse = {data_key: sparse.csr_matrix(data)}
         data_dict_dense = {data_key: data}
 
     adata_sparse = ad.AnnData(obs=obs_df, var=var_df, **data_dict_sparse)
@@ -390,3 +390,67 @@ def test_aggregate_arraytype(array_type, metric):
     adata.X = array_type(adata.X)
     aggregate = sc.get.aggregate(adata, ["louvain"], metric)
     assert isinstance(aggregate.layers[metric], np.ndarray)
+
+
+def test_aggregate_obsm_varm():
+    adata_obsm = sc.datasets.blobs()
+    adata_obsm.obs["blobs"] = adata_obsm.obs["blobs"].astype(str)
+    adata_obsm.obsm["test"] = adata_obsm.X[:, ::2].copy()
+    adata_varm = adata_obsm.T.copy()
+
+    result_obsm = sc.get.aggregate(adata_obsm, "blobs", ["sum", "mean"], obsm="test")
+    result_varm = sc.get.aggregate(adata_varm, "blobs", ["sum", "mean"], varm="test")
+
+    assert_equal(result_obsm, result_varm.T)
+
+    expected_sum = (
+        pd.DataFrame(adata_obsm.obsm["test"], index=adata_obsm.obs_names)
+        .groupby(adata_obsm.obs["blobs"], observed=True)
+        .sum()
+    )
+    expected_mean = (
+        pd.DataFrame(adata_obsm.obsm["test"], index=adata_obsm.obs_names)
+        .groupby(adata_obsm.obs["blobs"], observed=True)
+        .mean()
+    )
+
+    assert_equal(expected_sum.values, result_obsm.layers["sum"])
+    assert_equal(expected_mean.values, result_obsm.layers["mean"])
+
+
+def test_aggregate_obsm_labels():
+    from itertools import chain, repeat
+
+    label_counts = [("a", 5), ("b", 3), ("c", 4)]
+    blocks = [np.ones((n, 1)) for _, n in label_counts]
+    obs_names = pd.Index(
+        [f"cell_{i:02d}" for i in range(sum(b.shape[0] for b in blocks))]
+    )
+    entry = pd.DataFrame(
+        sparse.block_diag(blocks).toarray(),
+        columns=[f"dim_{i}" for i in range(len(label_counts))],
+        index=obs_names,
+    )
+
+    adata = ad.AnnData(
+        obs=pd.DataFrame(
+            {
+                "labels": list(
+                    chain.from_iterable(repeat(l, n) for (l, n) in label_counts)
+                )
+            },
+            index=obs_names,
+        ),
+        var=pd.DataFrame(index=["gene_0"]),
+        obsm={"entry": entry},
+    )
+
+    expected = ad.AnnData(
+        obs=pd.DataFrame({"labels": pd.Categorical(list("abc"))}, index=list("abc")),
+        var=pd.DataFrame(index=[f"dim_{i}" for i in range(3)]),
+        layers={
+            "sum": np.diag([n for _, n in label_counts]),
+        },
+    )
+    result = sc.get.aggregate(adata, by="labels", func="sum", obsm="entry")
+    assert_equal(expected, result)