chore: fix spatialdata access in squidpy via a helper function

docs/conf.py

@@ -49,7 +49,7 @@
     numpy=("https://numpy.org/doc/stable", None),
     statsmodels=("https://www.statsmodels.org/stable", None),
     scipy=("https://docs.scipy.org/doc/scipy", None),
-    pandas=("https://pandas.pydata.org/pandas-docs/stable", None),
+    pandas=("https://pandas.pydata.org/docs", None),
     anndata=("https://anndata.readthedocs.io/en/stable", None),
     scanpy=("https://scanpy.readthedocs.io/en/stable", None),
     matplotlib=("https://matplotlib.org/stable", None),

src/squidpy/_constants/_pkg_constants.py

@@ -196,10 +196,18 @@ def _sort_haystack(
             return haystack
 
     class obsp:
+        @staticmethod
+        def _spatial_key(value: str | None, suffix: str) -> str:
+            if value is None:
+                return f"{Key.obsm.spatial}_{suffix}"
+            if value.endswith(f"_{suffix}"):
+                return value
+            return f"{value}_{suffix}"
+
         @classmethod
         def spatial_dist(cls, value: str | None = None) -> str:
-            return f"{Key.obsm.spatial}_distances" if value is None else f"{value}_distances"
+            return cls._spatial_key(value, "distances")
 
         @classmethod
         def spatial_conn(cls, value: str | None = None) -> str:
-            return f"{Key.obsm.spatial}_connectivities" if value is None else f"{value}_connectivities"
+            return cls._spatial_key(value, "connectivities")

src/squidpy/gr/_build.py

@@ -47,6 +47,7 @@
     _assert_categorical_obs,
     _assert_spatial_basis,
     _save_data,
+    extract_adata,
 )
 
 __all__ = ["spatial_neighbors"]
@@ -142,54 +143,14 @@ def spatial_neighbors(
         - :attr:`anndata.AnnData.obsp` ``['{{key_added}}_distances']`` - the spatial distances.
         - :attr:`anndata.AnnData.uns`  ``['{{key_added}}']`` - :class:`dict` containing parameters.
     """
-    if isinstance(adata, SpatialData):
-        assert elements_to_coordinate_systems is not None, (
-            "Since `adata` is a :class:`spatialdata.SpatialData`, `elements_to_coordinate_systems` must not be `None`."
-        )
-        assert table_key is not None, (
-            "Since `adata` is a :class:`spatialdata.SpatialData`, `table_key` must not be `None`."
-        )
-        elements, table = match_element_to_table(adata, list(elements_to_coordinate_systems), table_key)
-        assert table.obs_names.equals(adata.tables[table_key].obs_names), (
-            "The spatialdata table must annotate all elements keys. Some elements are missing, please check the `elements_to_coordinate_systems` dictionary."
-        )
-        regions, region_key, instance_key = get_table_keys(adata.tables[table_key])
-        regions = [regions] if isinstance(regions, str) else regions
-        ordered_regions_in_table = adata.tables[table_key].obs[region_key].unique()
-
-        # TODO: remove this after https://github.com/scverse/spatialdata/issues/614
-        remove_centroids = {}
-        elem_instances = []
-        for e in regions:
-            schema = get_model(elements[e])
-            element_instances = get_element_instances(elements[e]).to_series()
-            if np.isin(0, element_instances.values) and (schema in (Labels2DModel, Labels3DModel)):
-                element_instances = element_instances.drop(index=0)
-                remove_centroids[e] = True
-            else:
-                remove_centroids[e] = False
-            elem_instances.append(element_instances)
-
-        element_instances = pd.concat(elem_instances)
-        if (not np.all(element_instances.values == adata.tables[table_key].obs[instance_key].values)) or (
-            not np.all(ordered_regions_in_table == regions)
-        ):
-            raise ValueError(
-                "The spatialdata table must annotate all elements keys. Some elements are missing or not ordered correctly, please check the `elements_to_coordinate_systems` dictionary."
-            )
-        centroids = []
-        for region_ in ordered_regions_in_table:
-            cs = elements_to_coordinate_systems[region_]
-            centroid = get_centroids(adata[region_], coordinate_system=cs)[["x", "y"]].compute()
 
-            # TODO: remove this after https://github.com/scverse/spatialdata/issues/614
-            if remove_centroids[region_]:
-                centroid = centroid[1:].copy()
-            centroids.append(centroid)
-
-        adata.tables[table_key].obsm[spatial_key] = np.concatenate(centroids)
-        adata = adata.tables[table_key]
-        library_key = region_key
+    adata, library_key = _resolve_sdata(
+        adata=adata,
+        elements_to_coordinate_systems=elements_to_coordinate_systems,
+        table_key=table_key,
+        spatial_key=spatial_key,
+        library_key=library_key,
+    )
 
     assert_positive(n_rings, name="n_rings")
     assert_positive(n_neighs, name="n_neighs")
@@ -436,6 +397,62 @@ def _build_connectivity(
     return Adj
 
 
+def _resolve_sdata(
+    adata: AnnData | SpatialData,
+    elements_to_coordinate_systems: dict[str, str] | None,
+    table_key: str = "table",
+    spatial_key: str = Key.obsm.spatial,
+    library_key: str | None = None,
+) -> tuple[AnnData, str | None]:
+    if not isinstance(adata, SpatialData):
+        return adata, library_key
+
+    assert elements_to_coordinate_systems is not None, (
+        "Since `adata` is a :class:`spatialdata.SpatialData`, `elements_to_coordinate_systems` must not be `None`."
+    )
+    table = extract_adata(adata, table_key=table_key)
+    elements, matched_table = match_element_to_table(adata, list(elements_to_coordinate_systems), table_key)
+    assert matched_table.obs_names.equals(table.obs_names), (
+        "The spatialdata table must annotate all elements keys. Some elements are missing, please check the `elements_to_coordinate_systems` dictionary."
+    )
+    regions, region_key, instance_key = get_table_keys(table)
+    regions = [regions] if isinstance(regions, str) else regions
+    ordered_regions_in_table = table.obs[region_key].unique()
+
+    # TODO: remove this after https://github.com/scverse/spatialdata/issues/614
+    remove_centroids = {}
+    elem_instances = []
+    for e in regions:
+        schema = get_model(elements[e])
+        element_instances = get_element_instances(elements[e]).to_series()
+        if np.isin(0, element_instances.values) and (schema in (Labels2DModel, Labels3DModel)):
+            element_instances = element_instances.drop(index=0)
+            remove_centroids[e] = True
+        else:
+            remove_centroids[e] = False
+        elem_instances.append(element_instances)
+
+    element_instances = pd.concat(elem_instances)
+    if (not np.all(element_instances.values == table.obs[instance_key].values)) or (
+        not np.all(ordered_regions_in_table == regions)
+    ):
+        raise ValueError(
+            "The spatialdata table must annotate all elements keys. Some elements are missing or not ordered correctly, please check the `elements_to_coordinate_systems` dictionary."
+        )
+    centroids = []
+    for region_ in ordered_regions_in_table:
+        cs = elements_to_coordinate_systems[region_]
+        centroid = get_centroids(adata[region_], coordinate_system=cs)[["x", "y"]].compute()
+
+        # TODO: remove this after https://github.com/scverse/spatialdata/issues/614
+        if remove_centroids[region_]:
+            centroid = centroid[1:].copy()
+        centroids.append(centroid)
+
+    table.obsm[spatial_key] = np.concatenate(centroids)
+    return table, region_key
+
+
 @njit
 def _csr_bilateral_diag_scale_helper(
     mat: csr_array | csr_matrix,
@@ -559,8 +576,7 @@ def mask_graph(
     if not isinstance(polygon_mask, Polygon | MultiPolygon):
         raise ValueError(f"`polygon_mask` should be of type `Polygon` or `MultiPolygon`, got {type(polygon_mask)}")
 
-    # get elements
-    table = sdata.tables[table_key]
+    table = extract_adata(sdata, table_key=table_key)
     coords = table.obsm[spatial_key]
     Adj = table.obsp[conns_key]
     Dst = table.obsp[dists_key]

src/squidpy/gr/_ligrec.py

@@ -26,6 +26,7 @@
     _assert_categorical_obs,
     _genesymbols,
     _save_data,
+    extract_adata,
 )
 
 __all__ = ["ligrec", "PermutationTest"]
@@ -641,6 +642,8 @@ def ligrec(
     copy: bool = False,
     key_added: str | None = None,
     gene_symbols: str | None = None,
+    *,
+    table_key: str = "table",
     **kwargs: Any,
 ) -> Mapping[str, pd.DataFrame] | None:
     """
@@ -658,8 +661,7 @@ def ligrec(
     -------
     %(ligrec_test_returns)s
     """  # noqa: D400
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     with _genesymbols(adata, key=gene_symbols, use_raw=use_raw, make_unique=False):
         return (  # type: ignore[no-any-return]
             PermutationTest(adata, use_raw=use_raw)

src/squidpy/gr/_nhood.py

@@ -27,6 +27,7 @@
     _assert_connectivity_key,
     _save_data,
     _shuffle_group,
+    extract_adata,
 )
 
 __all__ = ["nhood_enrichment", "centrality_scores", "interaction_matrix"]
@@ -146,6 +147,8 @@ def nhood_enrichment(
     n_jobs: int | None = None,
     backend: str = "loky",
     show_progress_bar: bool = True,
+    *,
+    table_key: str = "table",
 ) -> NhoodEnrichmentResult | None:
     """
     Compute neighborhood enrichment by permutation test.
@@ -171,8 +174,7 @@ def nhood_enrichment(
         - :attr:`anndata.AnnData.uns` ``['{cluster_key}_nhood_enrichment']['zscore']`` - the enrichment z-score.
         - :attr:`anndata.AnnData.uns` ``['{cluster_key}_nhood_enrichment']['count']`` - the enrichment count.
     """
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     connectivity_key = Key.obsp.spatial_conn(connectivity_key)
     _assert_categorical_obs(adata, cluster_key)
     _assert_connectivity_key(adata, connectivity_key)
@@ -239,6 +241,8 @@ def centrality_scores(
     n_jobs: int | None = None,
     backend: str = "loky",
     show_progress_bar: bool = False,
+    *,
+    table_key: str = "table",
 ) -> pd.DataFrame | None:
     """
     Compute centrality scores per cluster or cell type.
@@ -268,8 +272,7 @@ def centrality_scores(
         - :attr:`anndata.AnnData.uns` ``['{{cluster_key}}_centrality_scores']`` - the centrality scores,
           as mentioned above.
     """
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     connectivity_key = Key.obsp.spatial_conn(connectivity_key)
     _assert_categorical_obs(adata, cluster_key)
     _assert_connectivity_key(adata, connectivity_key)
@@ -333,6 +336,8 @@ def interaction_matrix(
     normalized: bool = False,
     copy: bool = False,
     weights: bool = False,
+    *,
+    table_key: str = "table",
 ) -> NDArrayA | None:
     """
     Compute interaction matrix for clusters.
@@ -356,8 +361,7 @@ def interaction_matrix(
 
         - :attr:`anndata.AnnData.uns` ``['{cluster_key}_interactions']`` - the interaction matrix.
     """
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     connectivity_key = Key.obsp.spatial_conn(connectivity_key)
     _assert_categorical_obs(adata, cluster_key)
     _assert_connectivity_key(adata, connectivity_key)

src/squidpy/gr/_niche.py

@@ -21,6 +21,7 @@
 from squidpy._constants._constants import NicheDefinitions
 from squidpy._docs import d, inject_docs
 from squidpy._validators import assert_isinstance, assert_key_in_adata, assert_one_of
+from squidpy.gr._utils import extract_adata
 
 __all__ = ["calculate_niche"]
 
@@ -169,12 +170,8 @@ def calculate_niche(
     if resolutions is None:
         resolutions = [0.5]
 
-    if isinstance(data, SpatialData):
-        orig_adata = data.tables[table_key]
-        adata = orig_adata.copy()
-    else:
-        orig_adata = data
-        adata = data.copy()
+    orig_adata = extract_adata(data, table_key=table_key or "table")
+    adata = orig_adata.copy()
 
     assert_key_in_adata(
         adata,
@@ -843,18 +840,14 @@ def _validate_niche_args(
 
     assert_one_of(flavor, ["neighborhood", "utag", "cellcharter", "spatialleiden"], name="flavor")
 
+    if isinstance(data, SpatialData) and table_key is None:
+        raise ValueError("'table_key' is required when 'data' is a SpatialData object")
+
     if library_key is not None:
         assert_isinstance(library_key, str, name="library_key")
-        if isinstance(data, AnnData):
-            if library_key not in data.obs.columns:
-                raise ValueError(f"'library_key' must be a column in 'adata.obs', got {library_key}")
-        elif isinstance(data, SpatialData):
-            if table_key is None:
-                raise ValueError("'table_key' is required when 'data' is a SpatialData object")
-            if table_key not in data.tables:
-                raise ValueError(f"'table_key' must be a valid table key in 'data', got {table_key}")
-            if library_key not in data.tables[table_key].obs.columns:
-                raise ValueError(f"'library_key' must be a column in 'adata.obs', got {library_key}")
+        adata = extract_adata(data, table_key=table_key or "table")
+        if library_key not in adata.obs.columns:
+            raise ValueError(f"'library_key' must be a column in 'adata.obs', got {library_key}")
 
     if n_neighbors is not None:
         assert_isinstance(n_neighbors, int, name="n_neighbors")

src/squidpy/gr/_ppatterns.py

@@ -30,6 +30,7 @@
     _assert_connectivity_key,
     _assert_spatial_basis,
     _save_data,
+    extract_adata,
 )
 
 __all__ = ["spatial_autocorr", "co_occurrence"]
@@ -62,6 +63,8 @@ def spatial_autocorr(
     n_jobs: int | None = None,
     backend: str = "loky",
     show_progress_bar: bool = True,
+    *,
+    table_key: str = "table",
 ) -> pd.DataFrame | None:
     """
     Calculate Global Autocorrelation Statistic (Moran’s I  or Geary's C).
@@ -128,8 +131,7 @@ def spatial_autocorr(
         - :attr:`anndata.AnnData.uns` ``['moranI']`` - the above mentioned dataframe, if ``mode = {sp.MORAN.s!r}``.
         - :attr:`anndata.AnnData.uns` ``['gearyC']`` - the above mentioned dataframe, if ``mode = {sp.GEARY.s!r}``.
     """
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     _assert_connectivity_key(adata, connectivity_key)
 
     def extract_X(adata: AnnData, genes: str | Sequence[str] | None) -> tuple[NDArrayA | spmatrix, Sequence[Any]]:
@@ -348,6 +350,8 @@ def co_occurrence(
     spatial_key: str = Key.obsm.spatial,
     interval: int | NDArrayA = 50,
     copy: bool = False,
+    *,
+    table_key: str = "table",
 ) -> tuple[NDArrayA, NDArrayA] | None:
     """
     Compute co-occurrence probability of clusters.
@@ -373,9 +377,7 @@ def co_occurrence(
         - :attr:`anndata.AnnData.uns` ``['{cluster_key}_co_occurrence']['interval']`` - the distance thresholds
           computed at ``interval``.
     """
-
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     _assert_categorical_obs(adata, key=cluster_key)
     _assert_spatial_basis(adata, key=spatial_key)
 

src/squidpy/gr/_ripley.py

@@ -19,7 +19,7 @@
 from squidpy._constants._pkg_constants import Key
 from squidpy._docs import d, inject_docs
 from squidpy._utils import NDArrayA
-from squidpy.gr._utils import _assert_categorical_obs, _assert_spatial_basis, _save_data
+from squidpy.gr._utils import _assert_categorical_obs, _assert_spatial_basis, _save_data, extract_adata
 
 __all__ = ["ripley"]
 
@@ -39,6 +39,8 @@ def ripley(
     n_steps: int = 50,
     seed: int | None = None,
     copy: bool = False,
+    *,
+    table_key: str = "table",
 ) -> dict[str, pd.DataFrame | NDArrayA]:
     r"""
     Calculate various Ripley's statistics for point processes.
@@ -104,8 +106,7 @@ def ripley(
     `Wikipedia <https://en.wikipedia.org/wiki/Spatial_descriptive_statistics#Ripley's_K_and_L_functions>`_
     or :cite:`Baddeley2015-lm`.
     """
-    if isinstance(adata, SpatialData):
-        adata = adata.table
+    adata = extract_adata(adata, table_key=table_key)
     _assert_categorical_obs(adata, key=cluster_key)
     _assert_spatial_basis(adata, key=spatial_key)
     coordinates = adata.obsm[spatial_key]