compute graph base on spatialdata regions and add mask graph based on polygon (#842)

* add tests

* pin up spatialdata

* fix test

* add test for spatialdata graph

* address comments, pass shapely polygon and not geodataframe to mask_graph

* add note

* add release note and add mask_graph to API

* add shapely to intersphinx

Author: giovp

docs/api.rst

@@ -14,6 +14,7 @@ Graph
     :toctree: api
 
     gr.spatial_neighbors
+    gr.mask_graph
     gr.nhood_enrichment
     gr.co_occurrence
     gr.centrality_scores

docs/conf.py

@@ -94,6 +94,7 @@
     omnipath=("https://omnipath.readthedocs.io/en/latest", None),
     napari=("https://napari.org/", None),
     spatialdata=("https://spatialdata.scverse.org/en/latest", None),
+    shapely=("https://shapely.readthedocs.io/en/stable", None),
 )
 
 # Add any paths that contain templates here, relative to this directory.

docs/release/notes-dev.rst

@@ -1,2 +1,9 @@
 Squidpy dev (the-future)
 ========================
+
+Features
+--------
+- Now :func:`squidpy.gr.spatial_graph` can also be used on :class:`spatialdata.SpatialData` objects.
+- Add :func:`squidpy.gr.mask_graph` to mask a spatial graph based on :class:`shapely.Polygon` or :class:`shapely.MultiPolygon`
+  `@giovp <https://github.com/giovp>`__
+  `#842 <https://github.com/scverse/squidpy/pull/842>`__

pyproject.toml

@@ -72,7 +72,7 @@ dependencies = [
    "validators>=0.18.2",
    "xarray>=0.16.1",
    "zarr>=2.6.1",
-   "spatialdata",
+   "spatialdata>=0.2.0",
 ]
 
 [project.optional-dependencies]

src/squidpy/gr/__init__.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from squidpy.gr._build import spatial_neighbors
+from squidpy.gr._build import mask_graph, spatial_neighbors
 from squidpy.gr._ligrec import ligrec
 from squidpy.gr._nhood import centrality_scores, interaction_matrix, nhood_enrichment
 from squidpy.gr._ppatterns import co_occurrence, spatial_autocorr

src/squidpy/gr/_build.py

@@ -6,10 +6,14 @@
 from collections.abc import Iterable  # noqa: F401
 from functools import partial
 from itertools import chain
+from typing import Any
 
+import geopandas as gpd
 import numpy as np
+import pandas as pd
 from anndata import AnnData
 from anndata.utils import make_index_unique
+from geopandas import GeoDataFrame
 from numba import njit
 from scanpy import logging as logg
 from scipy.sparse import (
@@ -20,9 +24,28 @@
     spmatrix,
 )
 from scipy.spatial import Delaunay
+from shapely import LineString, MultiPolygon, Point, Polygon, distance
 from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
 from sklearn.neighbors import NearestNeighbors
 from spatialdata import SpatialData
+from spatialdata._core.centroids import get_centroids
+from spatialdata._core.query.relational_query import (
+    get_element_instances,
+    match_element_to_table,
+)
+from spatialdata.models import SpatialElement, get_table_keys
+from spatialdata.models.models import (
+    Image2DModel,
+    Image3DModel,
+    Labels2DModel,
+    Labels3DModel,
+    PointsModel,
+    RasterSchema,
+    ShapesModel,
+    TableModel,
+    get_axes_names,
+    get_model,
+)
 
 from squidpy._constants._constants import CoordType, Transform
 from squidpy._constants._pkg_constants import Key
@@ -43,6 +66,8 @@
 def spatial_neighbors(
     adata: AnnData | SpatialData,
     spatial_key: str = Key.obsm.spatial,
+    elements_to_coordinate_systems: dict[str, str] | None = None,
+    table_key: str | None = None,
     library_key: str | None = None,
     coord_type: str | CoordType | None = None,
     n_neighs: int = 6,
@@ -62,6 +87,17 @@ def spatial_neighbors(
     ----------
     %(adata)s
     %(spatial_key)s
+        If `adata` is a :class:`spatialdata.SpatialData`, the coordinates of the centroids will be stored in the
+        `adata` with this key.
+    elements_to_coordinate_systems
+        A dictionary mapping element names of the SpatialData object to coordinate systems.
+        The elements can be either Shapes or Labels. For compatibility, the spatialdata table must annotate
+        all regions keys. Must not be `None` if `adata` is a :class:`spatialdata.SpatialData`.
+    table_key
+        Key in :attr:`spatialdata.SpatialData.tables` where the spatialdata table is stored. Must not be `None` if
+        `adata` is a :class:`spatialdata.SpatialData`.
+    mask_polygon
+        The Polygon or MultiPolygon element.
     %(library_key)s
     coord_type
         Type of coordinate system. Valid options are:
@@ -109,7 +145,54 @@ def spatial_neighbors(
         - :attr:`anndata.AnnData.uns`  ``['{{key_added}}']`` - :class:`dict` containing parameters.
     """
     if isinstance(adata, SpatialData):
-        adata = adata.table
+        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
+
     _assert_positive(n_rings, name="n_rings")
     _assert_positive(n_neighs, name="n_neighs")
     _assert_spatial_basis(adata, spatial_key)
@@ -167,7 +250,12 @@ def spatial_neighbors(
     neighbors_dict = {
         "connectivities_key": conns_key,
         "distances_key": dists_key,
-        "params": {"n_neighbors": n_neighs, "coord_type": coord_type.v, "radius": radius, "transform": transform.v},
+        "params": {
+            "n_neighbors": n_neighs,
+            "coord_type": coord_type.v,
+            "radius": radius,
+            "transform": transform.v,
+        },
     }
 
     if copy:
@@ -194,10 +282,21 @@ def _spatial_neighbor(
     with warnings.catch_warnings():
         warnings.simplefilter("ignore", SparseEfficiencyWarning)
         if coord_type == CoordType.GRID:
-            Adj, Dst = _build_grid(coords, n_neighs=n_neighs, n_rings=n_rings, delaunay=delaunay, set_diag=set_diag)
+            Adj, Dst = _build_grid(
+                coords,
+                n_neighs=n_neighs,
+                n_rings=n_rings,
+                delaunay=delaunay,
+                set_diag=set_diag,
+            )
         elif coord_type == CoordType.GENERIC:
             Adj, Dst = _build_connectivity(
-                coords, n_neighs=n_neighs, radius=radius, delaunay=delaunay, return_distance=True, set_diag=set_diag
+                coords,
+                n_neighs=n_neighs,
+                radius=radius,
+                delaunay=delaunay,
+                return_distance=True,
+                set_diag=set_diag,
             )
         else:
             raise NotImplementedError(f"Coordinate type `{coord_type}` is not yet implemented.")
@@ -233,7 +332,11 @@ def _spatial_neighbor(
 
 
 def _build_grid(
-    coords: NDArrayA, n_neighs: int, n_rings: int, delaunay: bool = False, set_diag: bool = False
+    coords: NDArrayA,
+    n_neighs: int,
+    n_rings: int,
+    delaunay: bool = False,
+    set_diag: bool = False,
 ) -> tuple[csr_matrix, csr_matrix]:
     if n_rings > 1:
         Adj: csr_matrix = _build_connectivity(
@@ -258,7 +361,13 @@ def _build_grid(
         Dst = Adj.copy()
         Adj.data[:] = 1.0
     else:
-        Adj = _build_connectivity(coords, n_neighs=n_neighs, neigh_correct=True, delaunay=delaunay, set_diag=set_diag)
+        Adj = _build_connectivity(
+            coords,
+            n_neighs=n_neighs,
+            neigh_correct=True,
+            delaunay=delaunay,
+            set_diag=set_diag,
+        )
         Dst = Adj.copy()
 
     Dst.setdiag(0.0)
@@ -302,14 +411,21 @@ def _build_connectivity(
             if neigh_correct:
                 dist_cutoff = np.median(dists) * 1.3  # there's a small amount of sway
                 mask = dists < dist_cutoff
-                row_indices, col_indices, dists = row_indices[mask], col_indices[mask], dists[mask]
+                row_indices, col_indices, dists = (
+                    row_indices[mask],
+                    col_indices[mask],
+                    dists[mask],
+                )
         else:
             dists, col_indices = tree.radius_neighbors()
             row_indices = np.repeat(np.arange(N), [len(x) for x in col_indices])
             dists = np.concatenate(dists)
             col_indices = np.concatenate(col_indices)
 
-        Adj = csr_matrix((np.ones_like(row_indices, dtype=np.float64), (row_indices, col_indices)), shape=(N, N))
+        Adj = csr_matrix(
+            (np.ones_like(row_indices, dtype=np.float64), (row_indices, col_indices)),
+            shape=(N, N),
+        )
         if return_distance:
             Dst = csr_matrix((dists, (row_indices, col_indices)), shape=(N, N))
 
@@ -349,3 +465,121 @@ def _transform_a_spectral(a: spmatrix) -> spmatrix:
 
 def _transform_a_cosine(a: spmatrix) -> spmatrix:
     return cosine_similarity(a, dense_output=False)
+
+
+@d.dedent
+def mask_graph(
+    sdata: SpatialData,
+    table_key: str,
+    polygon_mask: Polygon | MultiPolygon,
+    negative_mask: bool = False,
+    spatial_key: str = Key.obsm.spatial,
+    key_added: str = "mask",
+    copy: bool = False,
+) -> SpatialData:
+    """
+    Mask the graph based on a polygon mask.
+
+    Given a spatial graph stored in :attr:`anndata.AnnData.obsp` ``['{{key_added}}_{{spatial_key}}_connectivities']`` and spatial coordinates stored in :attr:`anndata.AnnData.obsp` ``['{{spatial_key}}']``, it maskes the graph so that only edges fully contained in the polygons are kept.
+
+    Parameters
+    ----------
+    sdata
+        The spatial data object.
+    table_key:
+        The key of the table containing the spatial data.
+    polygon_mask
+        The :class:`shapely.Polygon` or :class:`shapely.MultiPolygon` to be used as mask.
+    negative_mask
+        Whether to keep the edges within the polygon mask or outside.
+        Note that when ``negative_mask = True``, only the edges fully contained in the polygon are removed.
+        If edges are partially contained in the polygon, they are kept.
+    %(spatial_key)s
+    key_added
+        Key which controls where the results are saved if ``copy = False``.
+    %(copy)s
+
+    Returns
+    -------
+    If ``copy = True``, returns a :class:`tuple` with the masked spatial connectivities and masked distances matrices.
+
+    Otherwise, modifies the ``adata`` with the following keys:
+
+        - :attr:`anndata.AnnData.obsp` ``['{{key_added}}_{{spatial_key}}_connectivities']`` - the spatial connectivities.
+        - :attr:`anndata.AnnData.obsp` ``['{{key_added}}_{{spatial_key}}_distances']`` - the spatial distances.
+        - :attr:`anndata.AnnData.uns`  ``['{{key_added}}_{{spatial_key}}']`` - :class:`dict` containing parameters.
+
+    Notes
+    -----
+    The `polygon_mask` must be in the same `coordinate_systems` of the spatial graph, but no check is performed to assess this.
+    """
+    # we could add this to arg, but I don't see use case for now
+    neighs_key = Key.uns.spatial_neighs(spatial_key)
+    conns_key = Key.obsp.spatial_conn(spatial_key)
+    dists_key = Key.obsp.spatial_dist(spatial_key)
+
+    # check polygon type
+    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]
+    coords = table.obsm[spatial_key]
+    Adj = table.obsp[conns_key]
+    Dst = table.obsp[dists_key]
+
+    # convert edges to lines
+    lines_coords, idx_out = _get_lines_coords(Adj.indices, Adj.indptr, coords)
+    lines_coords, idx_out = np.array(lines_coords), np.array(idx_out)
+    lines_df = gpd.GeoDataFrame(geometry=list(map(LineString, lines_coords)))
+
+    # check that lines overlap with the polygon
+    filt_lines = lines_df.geometry.within(polygon_mask).values
+
+    # ~ within index, and set that to 0
+    if not negative_mask:
+        # keep only the lines that are within the polygon
+        filt_lines = ~filt_lines
+    filt_idx_out = idx_out[filt_lines]
+
+    # filter connectivities
+    Adj[filt_idx_out[:, 0], filt_idx_out[:, 1]] = 0
+    Adj.eliminate_zeros()
+
+    # filter_distances
+    Dst[filt_idx_out[:, 0], filt_idx_out[:, 1]] = 0
+    Dst.eliminate_zeros()
+
+    mask_conns_key = f"{key_added}_{conns_key}"
+    mask_dists_key = f"{key_added}_{dists_key}"
+    mask_neighs_key = f"{key_added}_{neighs_key}"
+
+    neighbors_dict = {
+        "connectivities_key": mask_conns_key,
+        "distances_key": mask_dists_key,
+        "unfiltered_graph_key": conns_key,
+        "params": {
+            "negative_mask": negative_mask,
+            "table_key": table_key,
+        },
+    }
+
+    if copy:
+        return Adj, Dst
+
+    # save back to spatialdata
+    _save_data(table, attr="obsp", key=mask_conns_key, data=Adj)
+    _save_data(table, attr="obsp", key=mask_dists_key, data=Dst, prefix=False)
+    _save_data(table, attr="uns", key=mask_neighs_key, data=neighbors_dict, prefix=False)
+
+
+@njit
+def _get_lines_coords(indices: NDArrayA, indptr: NDArrayA, coords: NDArrayA) -> tuple[list[Any], list[Any]]:
+    lines = []
+    idx_out = []
+    for i in range(len(indptr) - 1):
+        ixs = indices[indptr[i] : indptr[i + 1]]
+        for ix in ixs:
+            lines.append([coords[i], coords[ix]])
+            idx_out.append((i, ix))
+    return lines, idx_out