methods for calculating radial expansion

src/celldega/nbhd/__init__.py

@@ -10,6 +10,10 @@
     _get_df_cell,
     _get_gdf_cell,
     _get_gdf_trx,
+    make_column_names_unique_fast,
+    safe_polygon,
+    simple_format,
+    transform_polygon,
 )
 
 
@@ -25,4 +29,8 @@
     "filter_alpha_shapes",
     "generate_hextile",
     "hextile_niche",
+    "make_column_names_unique_fast",
+    "safe_polygon",
+    "simple_format",
+    "transform_polygon",
 ]

src/celldega/nbhd/collection.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+from collections.abc import Sequence
 from pathlib import Path
 from typing import Any
 
@@ -313,6 +314,89 @@ def calc_gradient(
         )
         return type(self)(gdf=gdf_rings, nbhd_type=nbhd_type, **kwargs)
 
+    def calc_expansion(
+        self,
+        gdf_bounds: gpd.GeoDataFrame,
+        radii_um: Sequence[float] = (0.5, 1, 1.5, 2, 2.5),
+        nbhd_type: str = "expansion",
+        *,
+        technology: str | None = None,
+        scale_um_per_pixel: float | None = None,
+        join_style: int = 2,
+        mitre_limit: float = 5.0,
+        add_colors: bool = True,
+        **kwargs: Any,
+    ) -> dict[float, NeighborhoodCollection]:
+        """Buffer every entity in this collection outward, clipped to its own bound.
+
+        Unlike :meth:`calc_gradient` (concentric rings from ONE dissolved ROI),
+        this grows **every** neighborhood independently — e.g. a segmented
+        nucleus growing into its cell — clipping each to a matching row in
+        ``gdf_bounds`` (joined by ``self.nbhd_col``). Returns one new
+        ``NeighborhoodCollection`` per radius, sharing the same observation axis
+        so downstream results stay comparable across radii.
+
+        Args:
+            gdf_bounds: Per-entity clipping boundary, with a column named
+                ``self.nbhd_col`` and a ``geometry`` column.
+            radii_um: Buffer distances in microns. ``0`` returns the original
+                (validity-repaired) entity geometry, clipped to its bound.
+            nbhd_type: Label recorded on each returned collection.
+            technology: Imaging platform used to look up ``scale_um_per_pixel``
+                (e.g. ``"Xenium"``). Ignored if ``scale_um_per_pixel`` is given.
+            scale_um_per_pixel: Microns per pixel — a micron distance is
+                *divided* by this to get the geometry's native units. Defaults
+                to ``1.0`` (geometry already in microns, i.e. no conversion).
+                If this collection's geometry is in pixel space and you only
+                have a pixels-per-micron factor, pass its reciprocal
+                (``1 / pixels_per_micron``).
+            join_style: Shapely buffer join style (``1``=round, ``2``=mitre
+                (default), ``3``=bevel).
+            mitre_limit: Shapely mitre limit, used when ``join_style=2``.
+            add_colors: If ``True`` (default), add a ``color`` column — one
+                shade per radius — for visualization.
+            **kwargs: Forwarded to each new ``NeighborhoodCollection``.
+
+        Returns:
+            A dict mapping each radius to a new ``NeighborhoodCollection`` of
+            that radius's buffered, clipped geometries.
+
+        Raises:
+            ValueError: If this collection has no geometry, or if ids fail to
+                match ``gdf_bounds``.
+
+        Examples:
+            >>> nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_col="cell_id")
+            >>> series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[1, 2, 3])
+            >>> for radius, nbhd in series.items():
+            ...     nbhd.calc_signature(by="cell-free", data_dir=data_dir, drop_missing=False)
+        """
+        from celldega.nbhd.expansion import _calc_expansion
+
+        if self.gdf is None:
+            raise ValueError("gdf or geometry is required to calculate an expansion series")
+
+        if self.transformation_matrix is not None and "transformation_matrix" not in kwargs:
+            kwargs["transformation_matrix"] = self.transformation_matrix
+
+        per_radius_gdf = _calc_expansion(
+            self.gdf,
+            gdf_bounds,
+            radii_um=radii_um,
+            id_col=self.nbhd_col,
+            technology=technology,
+            scale_um_per_pixel=scale_um_per_pixel,
+            join_style=join_style,
+            mitre_limit=mitre_limit,
+            add_colors=add_colors,
+        )
+        return {
+            radius: type(self)(
+                gdf=gdf_radius, nbhd_type=nbhd_type, nbhd_col=self.nbhd_col, **kwargs
+            )
+            for radius, gdf_radius in per_radius_gdf.items()
+        }
+
     @property
     def geometry(self) -> gpd.GeoDataFrame | None:
         """Neighborhood geometry. Alias of :attr:`gdf` (single source of truth)."""
@@ -450,6 +534,9 @@ def calc_signature(
         modality_name: str | None = None,
         min_cells: int = 1,
         data_dir: str | None = None,
+        feature_col: str = "feature_name",
+        x_col: str = "x_location",
+        y_col: str = "y_location",
         drop_missing: bool = True,
     ) -> None:
         """Calculate a neighborhood-by-gene modality and attach it to ``self.mod``.
@@ -465,8 +552,18 @@ def calc_signature(
             modality_name: Key for the modality; defaults to ``"gene"``
                 (cell-derived) or ``"gene_cell_free"`` (transcript-derived).
             min_cells: Minimum cells/transcripts for a neighborhood to be kept.
-            data_dir: Transcript directory for ``by="cell-free"``; defaults to
-                ``self.data_dir``.
+            data_dir: Directory containing a transcripts parquet file — any
+                file whose name ends with ``transcripts.parquet`` (e.g.
+                ``transcripts.parquet``, ``data1_transcripts.parquet``), with
+                columns named ``feature_col``/``x_col``/``y_col`` (Xenium
+                convention by default), streamed in batches; defaults to
+                ``self.data_dir``. Required for ``by="cell-free"``.
+            feature_col: Gene/feature column in ``data_dir``'s
+                ``transcripts.parquet`` (default ``"feature_name"``).
+            x_col: Transcript x-coordinate column in ``data_dir``'s
+                ``transcripts.parquet`` (default ``"x_location"``).
+            y_col: Transcript y-coordinate column in ``data_dir``'s
+                ``transcripts.parquet`` (default ``"y_location"``).
             drop_missing: When ``True`` (default), neighborhoods with fewer than
                 ``min_cells`` cells (or transcripts) are removed from the
                 collection entirely. When ``False``, the collection keeps all
@@ -477,8 +574,8 @@ def calc_signature(
             ``None`` — the modality is attached to ``self.mod``.
 
         Raises:
-            ValueError: If ``adata`` is missing for ``by="cell"``, or ``data_dir``
-                is missing for ``by="cell-free"``.
+            ValueError: If ``adata`` is missing for ``by="cell"``, or
+                ``data_dir`` is missing for ``by="cell-free"``.
         """
         from celldega.nbhd.neighborhoods import (
             _calc_nbhd_by_gene,
@@ -499,6 +596,9 @@ def calc_signature(
             by=by,
             adata=adata,
             data_dir=resolved_data_dir,
+            feature_col=feature_col,
+            x_col=x_col,
+            y_col=y_col,
             nbhd_col=self.nbhd_col,
             min_cells=min_cells,
         )
@@ -592,8 +692,10 @@ def calc_transcript_assignment(
     ) -> None:
         """Add per-neighborhood transcript-assignment columns to ``obs``.
 
-        From ``transcripts.parquet`` in ``data_dir``, adds three ``obs`` columns
-        (on the underlying MuData) for each neighborhood:
+        From the transcripts parquet file in ``data_dir`` (any file whose name
+        ends with ``transcripts.parquet``, e.g. ``transcripts.parquet`` or
+        ``data1_transcripts.parquet``), adds three ``obs`` columns (on the
+        underlying MuData) for each neighborhood:
 
         - ``total_transcripts`` — transcripts falling inside the neighborhood.
         - ``unassigned_transcripts`` — those with ``cell_id == "UNASSIGNED"``.
@@ -606,7 +708,8 @@ def calc_transcript_assignment(
         Only transcripts are needed — no ``adata`` or cell polygons.
 
         Args:
-            data_dir: Directory containing ``transcripts.parquet``; defaults to
+            data_dir: Directory containing a transcripts parquet file (any
+                name ending with ``transcripts.parquet``); defaults to
                 ``self.data_dir``.
 
         Returns:

src/celldega/nbhd/expansion.py

@@ -0,0 +1,155 @@
+"""Expansion: per-entity buffering clipped to a matching bounding geometry.
+
+Unlike :mod:`celldega.nbhd.gradient` (concentric rings from ONE dissolved ROI),
+this grows **every** entity in a collection independently, clipping each to a
+matching row of a per-entity bounding GeoDataFrame so it never grows past its
+own outer limit — e.g. a segmented nucleus growing outward until it reaches its
+corresponding cell boundary.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+
+import geopandas as gpd
+from shapely.validation import make_valid
+
+from .gradient import _get_micron_per_pixel, _ring_colors
+
+
+_DEFAULT_RADII_UM: tuple[float, ...] = (0.5, 1, 1.5, 2, 2.5)
+
+
+def _calc_expansion(
+    gdf_source: gpd.GeoDataFrame,
+    gdf_bounds: gpd.GeoDataFrame,
+    radii_um: Sequence[float] = _DEFAULT_RADII_UM,
+    *,
+    id_col: str = "id",
+    technology: str | None = None,
+    scale_um_per_pixel: float | None = None,
+    join_style: int = 2,
+    mitre_limit: float = 5.0,
+    add_colors: bool = True,
+) -> dict[float, gpd.GeoDataFrame]:
+    """Engine behind :meth:`NeighborhoodCollection.calc_expansion`.
+
+    For each radius in ``radii_um``, buffers every entity in ``gdf_source``
+    outward and intersects the result with the matching row (by ``id_col``) in
+    ``gdf_bounds``, so growth stops at that entity's own bound. Invalid input
+    geometries are repaired with ``shapely.make_valid`` first.
+
+    Args:
+        gdf_source: One row per entity to expand, with an ``id_col`` column and
+            a ``geometry`` column.
+        gdf_bounds: One row per entity's clipping boundary, with a matching
+            ``id_col`` column and a ``geometry`` column.
+        radii_um: Buffer distances in microns. ``0`` returns the original
+            (validity-repaired) source geometry, clipped to its bound.
+        id_col: Column identifying each entity, shared by both frames.
+        technology: Imaging platform (e.g. ``"Xenium"``) used to look up
+            ``scale_um_per_pixel``. Ignored if ``scale_um_per_pixel`` is given.
+        scale_um_per_pixel: Microns per pixel — a micron distance is *divided*
+            by this to get the geometry's native units. Defaults to ``1.0``
+            (geometry already in microns, i.e. no conversion). If your
+            geometry is in pixel space and you only have a pixels-per-micron
+            factor, pass its reciprocal (``1 / pixels_per_micron``).
+        join_style: Shapely buffer join style (``1``=round, ``2``=mitre
+            (default), ``3``=bevel).
+        mitre_limit: Shapely mitre limit, used when ``join_style=2``.
+        add_colors: If ``True`` (default), add a ``color`` column — one shade
+            per radius — for visualization.
+
+    Returns:
+        A dict mapping each radius to a ``GeoDataFrame`` of that radius's
+        buffered, clipped entities (``id_col``, ``geometry``, ``radius_um``,
+        ``center_x``/``center_y``, ``area``/``area_um2``/``area_px2``, and
+        ``color`` if requested). Entities that vanish at a given radius are
+        dropped from that radius's frame.
+
+    Raises:
+        KeyError: If ``id_col`` is missing from either frame.
+        ValueError: If ids are duplicated in ``gdf_bounds`` or fail to match
+            between frames.
+
+    Examples:
+        >>> series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[1, 2, 3])
+
+        If geometry is in pixel space (e.g. an OME-XML ``PhysicalSizeX``, or
+        the reciprocal of a notebook's own ``high_res_scale``)::
+
+            >>> series = nbhd_nuclei.calc_expansion(
+            ...     gdf_cells, radii_um=[1, 2, 3], scale_um_per_pixel=1 / high_res_scale,
+            ... )
+    """
+    if id_col not in gdf_source.columns:
+        raise KeyError(f"gdf_source missing '{id_col}'")
+    if id_col not in gdf_bounds.columns:
+        raise KeyError(f"gdf_bounds missing '{id_col}'")
+
+    if scale_um_per_pixel is None:
+        scale_um_per_pixel = _get_micron_per_pixel(technology) if technology is not None else 1.0
+
+    source = gdf_source[[id_col, "geometry"]].copy()
+    source[id_col] = source[id_col].astype(str)
+    if source[id_col].duplicated().any():
+        dupes = source.loc[source[id_col].duplicated(), id_col].unique()[:5]
+        raise ValueError(f"gdf_source has duplicate '{id_col}' values, e.g. {list(dupes)}")
+    source["geometry"] = source["geometry"].apply(make_valid)
+
+    bounds = gdf_bounds[[id_col, "geometry"]].copy()
+    bounds[id_col] = bounds[id_col].astype(str)
+    if bounds[id_col].duplicated().any():
+        dupes = bounds.loc[bounds[id_col].duplicated(), id_col].unique()[:5]
+        raise ValueError(f"gdf_bounds has duplicate '{id_col}' values, e.g. {list(dupes)}")
+    bounds_lookup = bounds.set_index(id_col)["geometry"].apply(make_valid)
+
+    missing = set(source[id_col]) - set(bounds_lookup.index)
+    if missing:
+        example = sorted(missing)[:5]
+        raise ValueError(
+            f"{len(missing)} entities have no matching row in gdf_bounds (by '{id_col}'), "
+            f"e.g. {example}"
+        )
+
+    radii_sorted = sorted({float(r) for r in radii_um})
+    colors = (
+        _ring_colors("viridis", len(radii_sorted)) if add_colors else [None] * len(radii_sorted)
+    )
+    color_by_radius = dict(zip(radii_sorted, colors, strict=True))
+
+    results: dict[float, gpd.GeoDataFrame] = {}
+    for radius_um in radii_sorted:
+        radius_native = radius_um / scale_um_per_pixel
+
+        buffered = source["geometry"].buffer(
+            radius_native, join_style=join_style, mitre_limit=mitre_limit
+        )
+        clipped = [
+            geom.intersection(bounds_lookup.loc[eid])
+            for eid, geom in zip(source[id_col], buffered, strict=True)
+        ]
+
+        gdf_radius = gpd.GeoDataFrame(
+            {id_col: source[id_col].to_numpy()},
+            geometry=clipped,
+            crs=gdf_source.crs,
+        )
+        gdf_radius = gdf_radius[~gdf_radius.geometry.is_empty].reset_index(drop=True)
+        gdf_radius["radius_um"] = radius_um
+        gdf_radius["center_x"] = gdf_radius.centroid.x
+        gdf_radius["center_y"] = gdf_radius.centroid.y
+
+        # area_native is in the geometry's own (native/pixel) units; scale_um_per_pixel
+        # converts to microns (identity when geometry is already in microns).
+        area_native = gdf_radius.geometry.area
+        gdf_radius["area_px2"] = area_native
+        gdf_radius["area_um2"] = area_native * (scale_um_per_pixel**2)
+        gdf_radius["area"] = gdf_radius["area_um2"]
+
+        if add_colors:
+            gdf_radius["color"] = color_by_radius[radius_um]
+
+        results[radius_um] = gdf_radius
+
+    return results