diff --git a/docs/examples/brief_notebooks/Nuclear_Expansion_Radial_Buffering.ipynb b/docs/examples/brief_notebooks/Nuclear_Expansion_Radial_Buffering.ipynb new file mode 100644 index 00000000..757d24c9 --- /dev/null +++ b/docs/examples/brief_notebooks/Nuclear_Expansion_Radial_Buffering.ipynb @@ -0,0 +1,545 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "id": "3341e980", + "metadata": {}, + "source": [ + "# Nuclear-to-Cell Expansion with `celldega.nbhd`\n", + "\n", + "This notebook is a runnable companion to a nucleus/cell segmentation-sensitivity\n", + "analysis: starting from a nucleus polygon, grow it outward in fixed steps until it\n", + "reaches the boundary of its corresponding (larger) cell segmentation, and compute a\n", + "cell-by-gene matrix at every step.\n", + "\n", + "That workflow is now a first-class part of Celldega's neighborhood API:\n", + "\n", + "- **`NeighborhoodCollection.calc_expansion`** Give it a\n", + " `NeighborhoodCollection` of *any* entity and a matching per-entity bounding\n", + " `GeoDataFrame`; it buffers every entity outward at each requested radius (in\n", + " microns), clips each one to its own bound so growth never overshoots it, and\n", + " returns one new `NeighborhoodCollection` per radius, all sharing the same\n", + " observation axis so results stay directly comparable across radii. Nucleus ->\n", + " cell is just the running example below -- the same method works for any other\n", + " pair of nested per-entity geometries.\n", + "- **`NeighborhoodCollection.calc_signature(by=\"cell-free\", data_dir=...)`**\n", + " It always streams a `transcripts.parquet` directory in batches\n", + " (narrowing candidate entities per batch with a spatial index before testing\n", + " exact polygons), so a whole file doesn't need to be loaded into memory\n", + " once per radius; `feature_col`/`x_col`/`y_col` name its gene/x/y columns\n", + " (Xenium convention by default, but overridable for any column layout).\n", + "\n", + "Because the real instrument files aren't available here, this notebook builds a small\n", + "**synthetic** nucleus/cell/transcript dataset with the same shape as a real\n", + "segmentation export, so every cell below runs standalone." + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "id": "4d0f46b0", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:45.268799Z", + "iopub.status.busy": "2026-07-17T20:08:45.268702Z", + "iopub.status.idle": "2026-07-17T20:08:48.148942Z", + "shell.execute_reply": "2026-07-17T20:08:48.148095Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "'0.18.0'" + ] + }, + "execution_count": 1, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "import os\n", + "import tempfile\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "import geopandas as gpd\n", + "import matplotlib.pyplot as plt\n", + "from shapely.geometry import Point\n", + "\n", + "import celldega as dega\n", + "\n", + "dega.__version__" + ] + }, + { + "cell_type": "markdown", + "id": "0077007a", + "metadata": {}, + "source": [ + "## 1. Nucleus + cell-boundary polygons\n", + "\n", + "A real pipeline builds these from segmentation contour CSVs (one polygon per cell,\n", + "in each of a nucleus file and an expanded-cell-boundary file). Here we synthesize\n", + "the same shape: two `GeoDataFrame`s sharing a `cell_id` column, one with a small\n", + "nucleus polygon per cell and one with its larger enclosing cell polygon." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "id": "5f2394c8", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.151771Z", + "iopub.status.busy": "2026-07-17T20:08:48.151381Z", + "iopub.status.idle": "2026-07-17T20:08:48.166227Z", + "shell.execute_reply": "2026-07-17T20:08:48.165419Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((120, 2), (120, 2))" + ] + }, + "execution_count": 2, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "rng = np.random.default_rng(0)\n", + "\n", + "N_ROWS, N_COLS = 10, 12\n", + "SPACING_UM = 20.0\n", + "\n", + "records_nuclei, records_cells, cell_meta = [], [], []\n", + "\n", + "cell_id = 0\n", + "for row in range(N_ROWS):\n", + " for col in range(N_COLS):\n", + " cx = col * SPACING_UM + rng.normal(0, 1.5)\n", + " cy = row * SPACING_UM + rng.normal(0, 1.5)\n", + "\n", + " cell_radius = rng.uniform(7.0, 9.0)\n", + " nucleus_radius = rng.uniform(2.5, 3.5)\n", + " jitter = rng.uniform(0, 2.0, size=2)\n", + " nx, ny = cx + jitter[0], cy + jitter[1]\n", + "\n", + " # two synthetic \"cell types\" so downstream clustering has real structure\n", + " cell_type = \"TypeA\" if (row + col) % 2 == 0 else \"TypeB\"\n", + "\n", + " records_nuclei.append(\n", + " {\"cell_id\": cell_id, \"geometry\": Point(nx, ny).buffer(nucleus_radius, resolution=12)}\n", + " )\n", + " records_cells.append(\n", + " {\"cell_id\": cell_id, \"geometry\": Point(cx, cy).buffer(cell_radius, resolution=12)}\n", + " )\n", + " cell_meta.append(\n", + " {\"cell_id\": cell_id, \"cell_type\": cell_type, \"cx\": cx, \"cy\": cy,\n", + " \"nx\": nx, \"ny\": ny, \"nucleus_radius\": nucleus_radius, \"cell_radius\": cell_radius}\n", + " )\n", + " cell_id += 1\n", + "\n", + "gdf_nuclei = gpd.GeoDataFrame(records_nuclei)\n", + "gdf_cells = gpd.GeoDataFrame(records_cells)\n", + "df_cell_meta = pd.DataFrame(cell_meta).set_index(\"cell_id\")\n", + "\n", + "gdf_nuclei.shape, gdf_cells.shape" + ] + }, + { + "cell_type": "markdown", + "id": "69869e12", + "metadata": {}, + "source": [ + "## 2. Wrap the nuclei as a `NeighborhoodCollection`\n", + "\n", + "Each nucleus becomes one observation (\"neighborhood\"), keyed by `cell_id`." + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "id": "85a6f217", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.167907Z", + "iopub.status.busy": "2026-07-17T20:08:48.167766Z", + "iopub.status.idle": "2026-07-17T20:08:48.182136Z", + "shell.execute_reply": "2026-07-17T20:08:48.181644Z" + } + }, + "outputs": [ + { + "data": { + "text/html": [ + "
\n", + "\n", + "\n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + " \n", + "
area_um2
neighborhood_id
019.838659
136.964149
222.426905
321.573983
437.955601
\n", + "
" + ], + "text/plain": [ + " area_um2\n", + "neighborhood_id \n", + "0 19.838659\n", + "1 36.964149\n", + "2 22.426905\n", + "3 21.573983\n", + "4 37.955601" + ] + }, + "execution_count": 3, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "nbhd_nuclei = dega.nbhd.NeighborhoodCollection(\n", + " gdf=gdf_nuclei, nbhd_type=\"nucleus\", nbhd_col=\"cell_id\"\n", + ")\n", + "nbhd_nuclei.obs[[\"area_um2\"]].head()" + ] + }, + { + "cell_type": "markdown", + "id": "2bda07c4", + "metadata": {}, + "source": [ + "## 3. Expansion series\n", + "\n", + "`calc_expansion` buffers every nucleus outward at each radius in\n", + "`radii_um` and intersects it with the matching row of `gdf_cells`, so a nucleus\n", + "never grows past its own cell's membrane. It returns a dict keyed by radius, each\n", + "value a new `NeighborhoodCollection` sharing the same `cell_id` observation axis." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "id": "658414fd", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.184092Z", + "iopub.status.busy": "2026-07-17T20:08:48.183960Z", + "iopub.status.idle": "2026-07-17T20:08:48.268152Z", + "shell.execute_reply": "2026-07-17T20:08:48.267650Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "radius= 0.5 um -> n= 120 mean_area= 38.52 um^2\n", + "radius= 1.0 um -> n= 120 mean_area= 50.28 um^2\n", + "radius= 1.5 um -> n= 120 mean_area= 63.58 um^2\n", + "radius= 2.0 um -> n= 120 mean_area= 78.40 um^2\n", + "radius= 2.5 um -> n= 120 mean_area= 94.52 um^2\n" + ] + } + ], + "source": [ + "radii_um = [0.5, 1, 1.5, 2, 2.5]\n", + "nbhd_series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=radii_um)\n", + "\n", + "for radius, nbhd in nbhd_series.items():\n", + " print(f\"radius={radius:>4} um -> n={len(nbhd.gdf):>4} mean_area={nbhd.gdf['area_um2'].mean():6.2f} um^2\")" + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "id": "140c06c5", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.269740Z", + "iopub.status.busy": "2026-07-17T20:08:48.269622Z", + "iopub.status.idle": "2026-07-17T20:08:48.724554Z", + "shell.execute_reply": "2026-07-17T20:08:48.723975Z" + } + }, + "outputs": [], + "source": [ + "# visual sanity check for one example cell, mirroring the original notebook's plot\n", + "example_id = str(int(df_cell_meta.index[7]))\n", + "\n", + "fig, axes = plt.subplots(1, len(radii_um), figsize=(3 * len(radii_um), 3))\n", + "for ax, radius in zip(axes, radii_um):\n", + " nbhd = nbhd_series[radius]\n", + " gdf_cells[gdf_cells[\"cell_id\"].astype(str) == example_id].boundary.plot(ax=ax, color=\"black\")\n", + " nbhd.gdf.loc[[example_id]].plot(ax=ax, color=\"lightblue\", edgecolor=\"blue\", alpha=0.7)\n", + " ax.set_title(f\"+{radius} um\")\n", + " ax.set_aspect(\"equal\")\n", + " ax.axis(\"off\")\n", + "fig.suptitle(f\"cell_id {example_id}\")\n", + "fig.tight_layout()" + ] + }, + { + "cell_type": "markdown", + "id": "05b9882a", + "metadata": {}, + "source": [ + "## 4. Synthetic transcripts\n", + "\n", + "Stand-in for a `transcripts.parquet` with non-Xenium columns -- `x`, `y`, `name` --\n", + "matching the columns used in the original notebook's\n", + "`assign_trx_to_entity_streaming_parquet_optimized(..., x_col=\"x\", y_col=\"y\",\n", + "gene_col=\"name\")` call. Two gene pairs simulate real biology: `NucGene*`\n", + "transcripts cluster tightly at the nucleus center (captured at every radius), while\n", + "`CytoGene*` and a cell-type marker gene (`MarkerA`/`MarkerB`) scatter through the\n", + "cytoplasm and are only picked up as the buffer radius grows." + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "id": "09ef320b", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.726457Z", + "iopub.status.busy": "2026-07-17T20:08:48.726342Z", + "iopub.status.idle": "2026-07-17T20:08:48.822022Z", + "shell.execute_reply": "2026-07-17T20:08:48.821542Z" + } + }, + "outputs": [ + { + "data": { + "text/plain": [ + "((5490, 3),\n", + " {np.str_('CytoGene1'): 1243,\n", + " np.str_('CytoGene2'): 1157,\n", + " np.str_('NucGene1'): 984,\n", + " np.str_('NucGene2'): 837,\n", + " 'MarkerA': 644,\n", + " 'MarkerB': 625})" + ] + }, + "execution_count": 6, + "metadata": {}, + "output_type": "execute_result" + } + ], + "source": [ + "trx_rows = []\n", + "for cid, meta in df_cell_meta.iterrows():\n", + " n_nuc = rng.poisson(15)\n", + " nuc_xy = rng.normal([meta[\"nx\"], meta[\"ny\"]], meta[\"nucleus_radius\"] / 3, size=(n_nuc, 2))\n", + " nuc_genes = rng.choice([\"NucGene1\", \"NucGene2\"], size=n_nuc)\n", + "\n", + " # rejection-sample points in the cytoplasm annulus (inside cell, outside nucleus)\n", + " cyto_xy = []\n", + " while len(cyto_xy) < 20:\n", + " theta = rng.uniform(0, 2 * np.pi)\n", + " r = meta[\"cell_radius\"] * np.sqrt(rng.uniform(0, 1))\n", + " x, y = meta[\"cx\"] + r * np.cos(theta), meta[\"cy\"] + r * np.sin(theta)\n", + " if (x - meta[\"nx\"]) ** 2 + (y - meta[\"ny\"]) ** 2 > meta[\"nucleus_radius\"] ** 2:\n", + " cyto_xy.append((x, y))\n", + " cyto_xy = np.array(cyto_xy)\n", + " cyto_genes = rng.choice([\"CytoGene1\", \"CytoGene2\"], size=len(cyto_xy))\n", + "\n", + " marker_gene = \"MarkerA\" if meta[\"cell_type\"] == \"TypeA\" else \"MarkerB\"\n", + " marker_xy = cyto_xy[rng.integers(0, len(cyto_xy), size=rng.poisson(10))]\n", + "\n", + " for xy, gene in zip(nuc_xy, nuc_genes):\n", + " trx_rows.append({\"x\": xy[0], \"y\": xy[1], \"name\": gene})\n", + " for xy, gene in zip(cyto_xy, cyto_genes):\n", + " trx_rows.append({\"x\": xy[0], \"y\": xy[1], \"name\": gene})\n", + " for xy in marker_xy:\n", + " trx_rows.append({\"x\": xy[0], \"y\": xy[1], \"name\": marker_gene})\n", + "\n", + "df_trx = pd.DataFrame(trx_rows)\n", + "\n", + "# calc_signature's cell-free mode always streams from a transcripts.parquet on\n", + "# disk, so persist these to a directory rather than keeping them in memory\n", + "trx_dir = tempfile.mkdtemp()\n", + "df_trx.to_parquet(f\"{trx_dir}/transcripts.parquet\")\n", + "df_trx.shape, df_trx[\"name\"].value_counts().to_dict()" + ] + }, + { + "cell_type": "markdown", + "id": "103af3bf", + "metadata": {}, + "source": [ + "## 4. Cell-by-gene matrix at every radius\n", + "\n", + "`calc_signature(by=\"cell-free\", data_dir=...)` spatially joins transcripts to\n", + "each radius's polygons and returns transcript counts as an `AnnData` in\n", + "`nbhd.mod[\"gene_cell_free\"]`. `data_dir` points to a directory containing a `transcripts.parquet`;\n", + "`feature_col`/`x_col`/`y_col` name its gene/x/y columns (Xenium convention by\n", + "default, overridden below for this notebook's custom `name`/`x`/`y` columns).\n", + "The file is always streamed in batches internally -- narrowing candidate\n", + "entities per batch with a spatial index before testing exact polygons -- so a\n", + "whole file doesn't need to be loaded into memory once per radius." + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "id": "10f40fda", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.823730Z", + "iopub.status.busy": "2026-07-17T20:08:48.823611Z", + "iopub.status.idle": "2026-07-17T20:08:48.990788Z", + "shell.execute_reply": "2026-07-17T20:08:48.990170Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Calculating neighborhood-by-gene (cell-free, streaming)\n", + "Calculating neighborhood-by-gene (cell-free, streaming)\n", + "Calculating neighborhood-by-gene (cell-free, streaming)\n", + "Calculating neighborhood-by-gene (cell-free, streaming)\n", + "Calculating neighborhood-by-gene (cell-free, streaming)\n" + ] + } + ], + "source": [ + "for radius, nbhd in nbhd_series.items():\n", + " nbhd.calc_signature(\n", + " by=\"cell-free\", data_dir=trx_dir, feature_col=\"name\", x_col=\"x\", y_col=\"y\",\n", + " drop_missing=False,\n", + " )" + ] + }, + { + "cell_type": "markdown", + "id": "dfa3ca39-f40a-4873-ac59-9bea47541d39", + "metadata": {}, + "source": [ + "## 6. Save NeighborhoodCollections to disk" + ] + }, + { + "cell_type": "code", + "execution_count": 8, + "id": "14cd97bd-fa73-40c3-ab90-49a3328d5efc", + "metadata": { + "execution": { + "iopub.execute_input": "2026-07-17T20:08:48.992366Z", + "iopub.status.busy": "2026-07-17T20:08:48.992255Z", + "iopub.status.idle": "2026-07-17T20:08:49.106001Z", + "shell.execute_reply": "2026-07-17T20:08:49.105495Z" + } + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved resolution 0.5 -> /var/folders/_6/bhs42vt57t1dkb59k4sy0p440000gp/T/tmpyvdqmiu4/nbhd_expansion_0.5.h5mu\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Saved resolution 1.0 -> /var/folders/_6/bhs42vt57t1dkb59k4sy0p440000gp/T/tmpyvdqmiu4/nbhd_expansion_1.0.h5mu\n", + "Saved resolution 1.5 -> /var/folders/_6/bhs42vt57t1dkb59k4sy0p440000gp/T/tmpyvdqmiu4/nbhd_expansion_1.5.h5mu\n", + "Saved resolution 2.0 -> /var/folders/_6/bhs42vt57t1dkb59k4sy0p440000gp/T/tmpyvdqmiu4/nbhd_expansion_2.0.h5mu\n", + "Saved resolution 2.5 -> /var/folders/_6/bhs42vt57t1dkb59k4sy0p440000gp/T/tmpyvdqmiu4/nbhd_expansion_2.5.h5mu\n" + ] + } + ], + "source": [ + "nbhd_dir = tempfile.mkdtemp()\n", + "\n", + "for res, nbhd in nbhd_series.items():\n", + " out_path = f\"{nbhd_dir}/nbhd_expansion_{res}.h5mu\"\n", + " nbhd.write(out_path) # or nbhd.write_h5mu(out_path) if .write isn't available\n", + " print(f\"Saved resolution {res} -> {out_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "id": "4841e593-cad8-4e6f-9178-9caced708205", + "metadata": {}, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 3 (ipykernel)", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.12.11" + }, + "widgets": { + "application/vnd.jupyter.widget-state+json": { + "state": {}, + "version_major": 2, + "version_minor": 0 + } + } + }, + "nbformat": 4, + "nbformat_minor": 5 +} diff --git a/src/celldega/nbhd/__init__.py b/src/celldega/nbhd/__init__.py index a68f3678..339a679f 100644 --- a/src/celldega/nbhd/__init__.py +++ b/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", ] diff --git a/src/celldega/nbhd/collection.py b/src/celldega/nbhd/collection.py index 6aeb269f..8c406361 100644 --- a/src/celldega/nbhd/collection.py +++ b/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: diff --git a/src/celldega/nbhd/expansion.py b/src/celldega/nbhd/expansion.py new file mode 100644 index 00000000..5ed9972f --- /dev/null +++ b/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 diff --git a/src/celldega/nbhd/neighborhoods.py b/src/celldega/nbhd/neighborhoods.py index b82469dc..74ec7c3d 100644 --- a/src/celldega/nbhd/neighborhoods.py +++ b/src/celldega/nbhd/neighborhoods.py @@ -32,6 +32,9 @@ def _calc_nbhd_by_gene( by: str = "cell", adata: AnnData | None = None, data_dir: str | None = None, + feature_col: str = "feature_name", + x_col: str = "x_location", + y_col: str = "y_location", nbhd_col: str = "name", min_cells: int = 1, ) -> AnnData: @@ -41,45 +44,38 @@ def _calc_nbhd_by_gene( Internal spatial-computation kernel. The public entry point is :meth:`NeighborhoodCollection.calc_signature`. - Computes gene expression values for each neighborhood, either from cell-level - expression data (mean expression of cells within each neighborhood) or from - raw transcript counts (cell-free mode). + `by="cell"` averages cell-level expression per neighborhood; `by="cell-free"` + counts transcripts per neighborhood, streamed in batches from `data_dir`'s + `transcripts.parquet`. Parameters ---------- gdf_nbhd : gpd.GeoDataFrame - GeoDataFrame containing neighborhood geometries. Must have a geometry column - and a column specified by `nbhd_col` for neighborhood identifiers. + Neighborhood geometries, with a `geometry` column and a `nbhd_col` id column. by : str, default "cell" - Method for calculating gene expression: - - "cell": Mean expression of cells within each neighborhood (requires `adata`) - - "cell-free": Transcript counts within each neighborhood (requires `data_dir`) + "cell" (requires `adata`) or "cell-free" (requires `data_dir`). adata : AnnData, optional - AnnData object with cell data. Required when `by="cell"`. Must have spatial - coordinates in `obsm["spatial"]`. + Cell-level data with spatial coordinates in `obsm["spatial"]`; required + for `by="cell"`. data_dir : str, optional - Path to directory containing `transcripts.parquet`. Required when - `by="cell-free"`. + 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). Required for `by="cell-free"`. + feature_col : str, default "feature_name" + Gene/feature column in `data_dir`'s transcripts parquet file. + x_col, y_col : str, default "x_location", "y_location" + Transcript coordinate columns in `data_dir`'s `transcripts.parquet`. nbhd_col : str, default "name" - Column in `gdf_nbhd` containing neighborhood identifiers. + Neighborhood id column in `gdf_nbhd`. min_cells : int, default 1 - Minimum number of cells/transcripts required within a neighborhood to - include it in the output. Only applies when `by="cell"`. + Minimum cells/transcripts for a neighborhood to be kept. Returns ------- AnnData - AnnData object with shape (n_neighborhoods, n_genes) where: - - `X`: Matrix of gene expression values (mean for cell-derived, counts for cell-free) - - `obs`: DataFrame indexed by neighborhood names - - `var`: DataFrame indexed by gene names - - `obs["n_cells"]`: Cell count per neighborhood (when `by="cell"`) - - `uns["by"]`: Method used ("cell" or "cell-free") - - Notes - ----- - For cluster-specific gene expression analysis, filter your AnnData object - to include only cells from the desired cluster before calling this function. + Shape (n_neighborhoods, n_genes); `obs["n_cells"]` (`by="cell"`) or + `obs["n_transcripts"]` (`by="cell-free"`). """ if by == "cell": if adata is None: @@ -135,26 +131,19 @@ def _calc_nbhd_by_gene( if data_dir is None: raise ValueError("data_dir is required when by='cell-free'") - print("Calculating neighborhood-by-gene (cell-free)") - - df_trx = pd.read_parquet( - f"{data_dir}/transcripts.parquet", - columns=["feature_name", "x_location", "y_location"], - engine="pyarrow", - ) - geometry = gpd.points_from_xy(df_trx["x_location"], df_trx["y_location"]) - gdf_trx = gpd.GeoDataFrame(df_trx[["feature_name"]], geometry=geometry) - gdf_trx = gdf_trx.sjoin( - _nbhd_geometry_for_join(gdf_nbhd, nbhd_col), - how="left", - predicate="within", - ) + print("Calculating neighborhood-by-gene (cell-free, streaming)") + from celldega.nbhd.trx_streaming import _assign_trx_to_entity_streaming_parquet + from celldega.nbhd.utils import _find_transcripts_parquet df_result = ( - gdf_trx.groupby([nbhd_col, "feature_name"]) - .size() - .unstack(fill_value=0) - .rename_axis(None, axis=1) + _assign_trx_to_entity_streaming_parquet( + _find_transcripts_parquet(data_dir), + gdf_nbhd, + id_col=nbhd_col, + x_col=x_col, + y_col=y_col, + gene_col=feature_col, + ) .reindex(gdf_nbhd[nbhd_col]) .fillna(0) .astype(int) diff --git a/src/celldega/nbhd/trx_streaming.py b/src/celldega/nbhd/trx_streaming.py new file mode 100644 index 00000000..36794529 --- /dev/null +++ b/src/celldega/nbhd/trx_streaming.py @@ -0,0 +1,139 @@ +"""Streaming, spatial-index-accelerated transcript-to-entity assignment. + +Backs :meth:`NeighborhoodCollection.calc_signature`'s ``data_dir=`` cell-free +path: reads a transcripts parquet file in batches via ``pyarrow``, and per batch +only tests entities whose bounding box the batch could plausibly intersect (via +the entity ``GeoDataFrame``'s spatial index), so memory stays bounded by the +batch size regardless of file size — useful for a whole-tile +``transcripts.parquet`` re-joined once per radius in a +:meth:`~celldega.nbhd.collection.NeighborhoodCollection.calc_expansion` series. +""" + +from __future__ import annotations + +from collections import defaultdict + +import geopandas as gpd +import numpy as np +import pandas as pd +import pyarrow.dataset as ds +import shapely + + +def _assign_trx_to_entity_streaming_parquet( + trx_parquet_path: str, + gdf_entity: gpd.GeoDataFrame, + id_col: str, + *, + x_col: str = "x", + y_col: str = "y", + gene_col: str = "gene", + batch_size: int = 1_000_000, + assume_non_overlapping: bool = True, +) -> pd.DataFrame: + """Stream transcripts from ``trx_parquet_path`` and count them per entity/gene. + + For each streamed batch, candidate entities are first narrowed down with + ``gdf_entity``'s spatial index (by the batch's bounding box), then each + candidate's exact polygon is tested with a vectorized point-in-polygon check + (``shapely.contains_xy``). Counts are accumulated across batches. + + Args: + trx_parquet_path: Path to a transcripts parquet file (or dataset) + containing at least ``x_col``, ``y_col``, and ``gene_col``. + gdf_entity: One row per entity to assign transcripts to, with an + ``id_col`` column and a ``geometry`` column. + id_col: Column in ``gdf_entity`` identifying each entity. + x_col: Transcript x-coordinate column in the parquet file. + y_col: Transcript y-coordinate column in the parquet file. + gene_col: Transcript gene/feature column in the parquet file. + batch_size: Rows read per streamed batch. Bounds peak memory use; does + not affect the result. + assume_non_overlapping: If ``True`` (default), a transcript is excluded + from consideration once assigned — valid whenever entities don't + overlap — and lets a batch stop early once every point has a match. + + Returns: + A ``DataFrame`` indexed by entity id (as ``str``) with one integer count + column per gene seen in an assigned transcript. Entities/genes never + seen are simply absent — callers typically reindex/``fillna(0)``. + + Raises: + KeyError: If ``id_col`` is missing from ``gdf_entity``. + ValueError: If ``gdf_entity`` has no valid (non-null) geometries. + """ + if id_col not in gdf_entity.columns: + raise KeyError(f"gdf_entity missing '{id_col}'") + + entity = gdf_entity[[id_col, "geometry"]].copy() + entity = entity[entity.geometry.notna()].reset_index(drop=True) + if entity.empty: + raise ValueError("gdf_entity has no valid geometries") + entity["geometry"] = entity.geometry.buffer(0) + entity[id_col] = entity[id_col].astype(str) + + ids = entity[id_col].to_numpy() + geoms = entity.geometry.to_numpy() + bboxes = np.array([g.bounds for g in geoms], dtype=np.float64) + sindex = entity.sindex + + dataset = ds.dataset(trx_parquet_path, format="parquet") + scanner = dataset.scanner(columns=[x_col, y_col, gene_col], batch_size=batch_size) + + counts: dict[tuple[str, str], int] = defaultdict(int) + + for batch in scanner.to_batches(): + if batch.num_rows == 0: + continue + + x = batch.column(batch.schema.get_field_index(x_col)).to_numpy(zero_copy_only=False) + y = batch.column(batch.schema.get_field_index(y_col)).to_numpy(zero_copy_only=False) + gene = batch.column(batch.schema.get_field_index(gene_col)).to_numpy(zero_copy_only=False) + + valid = np.isfinite(x) & np.isfinite(y) + if not valid.all(): + x, y, gene = x[valid], y[valid], gene[valid] + if len(x) == 0: + continue + + assigned = np.full(len(x), -1, dtype=np.int64) + + candidates = list(sindex.intersection((x.min(), y.min(), x.max(), y.max()))) + for j in candidates: + minx, miny, maxx, maxy = bboxes[j] + cand = (x >= minx) & (x <= maxx) & (y >= miny) & (y <= maxy) + if assume_non_overlapping: + cand &= assigned == -1 + if not cand.any(): + continue + + idx = np.flatnonzero(cand) + inside = shapely.contains_xy(geoms[j], x[idx], y[idx]) + if inside.any(): + assigned[idx[inside]] = j + + if assume_non_overlapping and (assigned != -1).all(): + break + + keep = assigned != -1 + if not keep.any(): + continue + + chunk = pd.DataFrame({id_col: ids[assigned[keep]], gene_col: gene[keep]}) + for (eid, g), c in chunk.value_counts().items(): + counts[(eid, g)] += int(c) + + if not counts: + return pd.DataFrame(index=pd.Index([], name=id_col)) + + df_long = pd.DataFrame( + [(eid, g, c) for (eid, g), c in counts.items()], + columns=[id_col, gene_col, "count"], + ) + return ( + df_long.pivot_table( + index=id_col, columns=gene_col, values="count", fill_value=0, aggfunc="sum" + ) + .rename_axis(None, axis=1) + .astype(int) + ) diff --git a/src/celldega/nbhd/utils.py b/src/celldega/nbhd/utils.py index 66a2a3cd..5a1788c0 100644 --- a/src/celldega/nbhd/utils.py +++ b/src/celldega/nbhd/utils.py @@ -1,14 +1,16 @@ """Helper and utility functions.""" # Standard library imports +from collections import defaultdict from collections.abc import Sequence +from pathlib import Path from typing import Any # Third-party imports import geopandas as gpd import numpy as np import pandas as pd -from shapely.geometry import Point, base +from shapely.geometry import Point, Polygon, base from shapely.ops import transform @@ -105,6 +107,29 @@ def _get_gdf_cell(adata: Any) -> gpd.GeoDataFrame: ) +def _find_transcripts_parquet(data_dir: str) -> str: + """ + Find the transcripts parquet file in `data_dir`. + + Matches any file whose name ends with `transcripts.parquet` (e.g. + `transcripts.parquet`, `data1_transcripts.parquet`, + `aziz_1_20260217_5_transcripts.parquet`), not just the literal Xenium + convention `transcripts.parquet`. + """ + candidates = sorted( + p for p in Path(data_dir).iterdir() if p.name.endswith("transcripts.parquet") + ) + if not candidates: + raise FileNotFoundError(f"No file ending with 'transcripts.parquet' found in '{data_dir}'") + if len(candidates) > 1: + raise ValueError( + f"Multiple files ending with 'transcripts.parquet' found in '{data_dir}': " + f"{[p.name for p in candidates]}. Keep only one, or point data_dir at a " + "directory containing a single transcripts file." + ) + return str(candidates[0]) + + def _get_gdf_trx(data_dir: str) -> gpd.GeoDataFrame: """ Load transcript data as a GeoDataFrame with spatial coordinates. @@ -112,7 +137,7 @@ def _get_gdf_trx(data_dir: str) -> gpd.GeoDataFrame: No CRS is set since coordinates are in micron imaging space, not geospatial. """ df_trx = pd.read_parquet( - f"{data_dir}/transcripts.parquet", + _find_transcripts_parquet(data_dir), columns=["feature_name", "x_location", "y_location", "cell_id"], engine="pyarrow", ) @@ -147,3 +172,50 @@ def round_coords( return (round(x, precision), round(y, precision)) return transform(round_coords, geometry) + + +def safe_polygon(row: pd.Series) -> Polygon: + """Build a `Polygon` from a row's `vertex_x`/`vertex_y` coordinate lists; empty on failure.""" + try: + return Polygon(zip(row["vertex_x"], row["vertex_y"], strict=True)) + except Exception: + return Polygon() + + +def simple_format(geometry: Sequence[Sequence[Sequence[float]]], image_scale: float) -> list: + """Rescale a nested polygon-ring coordinate list by dividing by `image_scale`.""" + return [ + [[coord[0] / image_scale, coord[1] / image_scale] for coord in polygon] + for polygon in geometry + ] + + +def transform_polygon(polygon: Polygon) -> np.ndarray: + """Convert a `Polygon`'s exterior ring into a `[1, n_points, 2]` object array.""" + exterior_coords = polygon.exterior.coords + original_format_coords = np.array([np.array(coord) for coord in exterior_coords]) + return np.array([original_format_coords], dtype=object) + + +def make_column_names_unique_fast(df: pd.DataFrame) -> pd.DataFrame: + """Rename duplicate columns in place (`col`, `col_1`, `col_2`, ...) and return `df`.""" + counts: dict[str, int] = defaultdict(int) + used: set[str] = set() + new_cols = [] + + for col in df.columns: + if col not in used: + new_cols.append(col) + used.add(col) + counts[col] += 1 + else: + while True: + new_name = f"{col}_{counts[col]}" + counts[col] += 1 + if new_name not in used: + new_cols.append(new_name) + used.add(new_name) + break + + df.columns = new_cols + return df diff --git a/tests/unit/test_nbhd/test_expansion.py b/tests/unit/test_nbhd/test_expansion.py new file mode 100644 index 00000000..2447a936 --- /dev/null +++ b/tests/unit/test_nbhd/test_expansion.py @@ -0,0 +1,335 @@ +import geopandas as gpd +import numpy as np +import pandas as pd +import pytest +from shapely.geometry import Polygon + +from celldega.nbhd import NeighborhoodCollection +from celldega.nbhd.expansion import _calc_expansion + + +def _synthetic_nucleus_cell_inputs(): + # cell 1: 10x10 square at origin; nucleus 1: centered 2x2 square (area 4) + # cell 2: 10x10 square offset far away; nucleus 2: centered 2x2 square (area 4) + gdf_nuclei = gpd.GeoDataFrame( + { + "cell_id": ["c1", "c2"], + "geometry": [ + Polygon([(4, 4), (6, 4), (6, 6), (4, 6)]), + Polygon([(24, 24), (26, 24), (26, 26), (24, 26)]), + ], + } + ) + gdf_cells = gpd.GeoDataFrame( + { + "cell_id": ["c1", "c2"], + "geometry": [ + Polygon([(0, 0), (10, 0), (10, 10), (0, 10)]), + Polygon([(20, 20), (30, 20), (30, 30), (20, 30)]), + ], + } + ) + return gdf_nuclei, gdf_cells + + +def test_calc_expansion_grows_and_clips_to_bound(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + + series = _calc_expansion(gdf_nuclei, gdf_cells, radii_um=[0, 1, 5], id_col="cell_id") + + assert list(series.keys()) == [0.0, 1.0, 5.0] + + # radius 0: original 2x2 source polygon, area 4 + gdf_0 = series[0.0] + assert set(gdf_0["cell_id"]) == {"c1", "c2"} + np.testing.assert_allclose(sorted(gdf_0["area_um2"]), [4.0, 4.0]) + + # radius 1: buffered 1 unit on each side -> 4x4 square, area 16, still inside the bound + gdf_1 = series[1.0] + np.testing.assert_allclose(sorted(gdf_1["area_um2"]), [16.0, 16.0]) + + # radius 5: buffer would overshoot the bound -> clipped to the full 10x10 bound + gdf_5 = series[5.0] + np.testing.assert_allclose(sorted(gdf_5["area_um2"]), [100.0, 100.0]) + + +def test_calc_expansion_scale_um_per_pixel_converts_pixel_space_geometry(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + high_res_scale = 2.0 # pixels per micron, e.g. a notebook's own scale variable + scale_um_per_pixel = 1.0 / high_res_scale # microns per pixel -- what calc_expansion wants + + result = _calc_expansion( + gdf_nuclei, + gdf_cells, + radii_um=[1], + id_col="cell_id", + scale_um_per_pixel=scale_um_per_pixel, + ) + + # matches `buffer_dist = expand_um * high_res_scale`: a 2x2 nucleus buffered by + # 1um * 2px/um = 2px on each side -> 6x6 = 36 px^2, well inside the 10x10 bound + gdf_1 = result[1.0] + np.testing.assert_allclose(sorted(gdf_1["area_px2"]), [36.0, 36.0]) + # area_um2 = area_px2 * scale_um_per_pixel**2 = 36 * 0.25 = 9 + np.testing.assert_allclose(sorted(gdf_1["area_um2"]), [9.0, 9.0]) + + +def test_calc_expansion_default_scale_treats_geometry_as_microns(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + + result = _calc_expansion(gdf_nuclei, gdf_cells, radii_um=[1], id_col="cell_id") + + # no conversion: a 2x2 nucleus buffered by 1um -> 4x4 = 16 um^2 + gdf_1 = result[1.0] + np.testing.assert_allclose(sorted(gdf_1["area_um2"]), [16.0, 16.0]) + + +def test_calc_expansion_resolves_scale_from_technology(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + + result = _calc_expansion( + gdf_nuclei, gdf_cells, radii_um=[1], id_col="cell_id", technology="Xenium" + ) + expected = _calc_expansion( + gdf_nuclei, gdf_cells, radii_um=[1], id_col="cell_id", scale_um_per_pixel=0.2125 + ) + pd.testing.assert_frame_equal( + result[1.0].drop(columns="color"), expected[1.0].drop(columns="color") + ) + + +def test_neighborhood_collection_calc_expansion_accepts_scale_um_per_pixel(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[1], scale_um_per_pixel=0.5) + + np.testing.assert_allclose(sorted(series[1.0].gdf["area_px2"]), [36.0, 36.0]) + + +def test_calc_expansion_works_for_non_nucleus_entities(): + # Demonstrates this isn't nucleus/cell-specific: any pair of per-entity + # source/bound geometries with a shared id column works, e.g. a small "core" + # region expanding into a larger parent "zone". + gdf_core = gpd.GeoDataFrame( + { + "region_id": ["r1", "r2"], + "geometry": [ + Polygon([(1, 1), (2, 1), (2, 2), (1, 2)]), + Polygon([(11, 1), (12, 1), (12, 2), (11, 2)]), + ], + } + ) + gdf_zone = gpd.GeoDataFrame( + { + "region_id": ["r1", "r2"], + "geometry": [ + Polygon([(0, 0), (5, 0), (5, 5), (0, 5)]), + Polygon([(10, 0), (15, 0), (15, 5), (10, 5)]), + ], + } + ) + + series = _calc_expansion(gdf_core, gdf_zone, radii_um=[0, 10], id_col="region_id") + + assert list(series.keys()) == [0.0, 10.0] + assert set(series[0.0]["region_id"]) == {"r1", "r2"} + # radius 10 overshoots every zone -> clipped to each 5x5 zone, area 25 + np.testing.assert_allclose(sorted(series[10.0]["area_um2"]), [25.0, 25.0]) + + +def test_calc_expansion_add_colors(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + + with_colors = _calc_expansion(gdf_nuclei, gdf_cells, radii_um=[0, 1, 2], id_col="cell_id") + assert all("color" in gdf.columns for gdf in with_colors.values()) + # one shade per radius, shared across entities within that radius + assert with_colors[0.0]["color"].nunique() == 1 + assert len({gdf["color"].iloc[0] for gdf in with_colors.values()}) == 3 + + without_colors = _calc_expansion( + gdf_nuclei, gdf_cells, radii_um=[0, 1], id_col="cell_id", add_colors=False + ) + assert all("color" not in gdf.columns for gdf in without_colors.values()) + + +def test_calc_expansion_raises_on_duplicate_ids(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + gdf_cells_dup = pd.concat([gdf_cells, gdf_cells.iloc[[0]]], ignore_index=True) + gdf_cells_dup = gpd.GeoDataFrame(gdf_cells_dup, geometry="geometry") + + with pytest.raises(ValueError, match="duplicate"): + _calc_expansion(gdf_nuclei, gdf_cells_dup, radii_um=[0, 1], id_col="cell_id") + + +def test_calc_expansion_raises_on_missing_bound_match(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + gdf_cells_missing = gdf_cells.iloc[[0]].reset_index(drop=True) + + with pytest.raises(ValueError, match="no matching row"): + _calc_expansion(gdf_nuclei, gdf_cells_missing, radii_um=[0, 1], id_col="cell_id") + + +def test_neighborhood_collection_calc_expansion_returns_series(): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[0, 1, 5]) + + assert list(series.keys()) == [0.0, 1.0, 5.0] + for nbhd in series.values(): + assert isinstance(nbhd, NeighborhoodCollection) + assert nbhd.nbhd_col == "cell_id" + assert set(nbhd.obs.index) == {"c1", "c2"} + assert nbhd.nbhd_type == "expansion" + + +def test_calc_signature_cell_free_accepts_custom_columns_across_radii(tmp_path): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[0, 5]) + + # Custom transcript format: "name" for gene, arbitrary x/y columns -- + # exercises the feature_col/x_col/y_col overrides end to end. + pd.DataFrame( + { + "name": ["GeneA", "GeneB", "GeneA"], + "x": [5, 1, 25], + "y": [5, 1, 25], + } + ).to_parquet(tmp_path / "transcripts.parquet") + + nbhd_r0 = series[0.0] + nbhd_r0.calc_signature( + by="cell-free", + data_dir=str(tmp_path), + feature_col="name", + x_col="x", + y_col="y", + drop_missing=False, + ) + modality_r0 = nbhd_r0.mod["gene_cell_free"] + df_r0 = pd.DataFrame(modality_r0.X, index=modality_r0.obs_names, columns=modality_r0.var_names) + # at radius 0 the source polygon doesn't reach (1, 1); only the point inside it counts + # ("GeneB" never falls inside any neighborhood at this radius, so it has no column) + assert df_r0.loc["c1", "GeneA"] == 1 + assert "GeneB" not in df_r0.columns + assert df_r0.loc["c2", "GeneA"] == 1 + + nbhd_r5 = series[5.0] + nbhd_r5.calc_signature( + by="cell-free", + data_dir=str(tmp_path), + feature_col="name", + x_col="x", + y_col="y", + drop_missing=False, + ) + modality_r5 = nbhd_r5.mod["gene_cell_free"] + df_r5 = pd.DataFrame(modality_r5.X, index=modality_r5.obs_names, columns=modality_r5.var_names) + # at radius 5 the source polygon has expanded to the full bound, now capturing (1, 1) too + assert df_r5.loc["c1", "GeneA"] == 1 + assert df_r5.loc["c1", "GeneB"] == 1 + + +def test_calc_signature_cell_free_requires_data_dir(): + gdf_nuclei, _gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + + with pytest.raises(ValueError, match="data_dir is required"): + nbhd.calc_signature(by="cell-free") + + +def test_calc_signature_cell_free_streams_from_data_dir_across_radii(tmp_path): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[0, 5]) + + # data_dir's cell-free path is now backed by the streaming engine internally, + # using the Xenium transcripts.parquet convention. + pd.DataFrame( + { + "feature_name": ["GeneA", "GeneB", "GeneA"], + "x_location": [5, 1, 25], + "y_location": [5, 1, 25], + } + ).to_parquet(tmp_path / "transcripts.parquet") + + nbhd_r0 = series[0.0] + nbhd_r0.calc_signature(by="cell-free", data_dir=str(tmp_path), drop_missing=False) + df_r0 = pd.DataFrame( + nbhd_r0.mod["gene_cell_free"].X, + index=nbhd_r0.mod["gene_cell_free"].obs_names, + columns=nbhd_r0.mod["gene_cell_free"].var_names, + ) + assert df_r0.loc["c1", "GeneA"] == 1 + assert "GeneB" not in df_r0.columns + + nbhd_r5 = series[5.0] + nbhd_r5.calc_signature(by="cell-free", data_dir=str(tmp_path), drop_missing=False) + df_r5 = pd.DataFrame( + nbhd_r5.mod["gene_cell_free"].X, + index=nbhd_r5.mod["gene_cell_free"].obs_names, + columns=nbhd_r5.mod["gene_cell_free"].var_names, + ) + assert df_r5.loc["c1", "GeneA"] == 1 + assert df_r5.loc["c1", "GeneB"] == 1 + + +def test_calc_signature_cell_free_finds_prefixed_transcripts_filename(tmp_path): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[5]) + + # not every transcripts file is literally named "transcripts.parquet" + pd.DataFrame( + { + "feature_name": ["GeneA", "GeneA"], + "x_location": [5, 25], + "y_location": [5, 25], + } + ).to_parquet(tmp_path / "aziz_1_20260217_5_transcripts.parquet") + + nbhd_r5 = series[5.0] + nbhd_r5.calc_signature(by="cell-free", data_dir=str(tmp_path), drop_missing=False) + df_r5 = pd.DataFrame( + nbhd_r5.mod["gene_cell_free"].X, + index=nbhd_r5.mod["gene_cell_free"].obs_names, + columns=nbhd_r5.mod["gene_cell_free"].var_names, + ) + assert df_r5.loc["c1", "GeneA"] == 1 + assert df_r5.loc["c2", "GeneA"] == 1 + + +def test_calc_signature_cell_free_data_dir_accepts_custom_columns(tmp_path): + gdf_nuclei, gdf_cells = _synthetic_nucleus_cell_inputs() + nbhd_nuclei = NeighborhoodCollection(gdf=gdf_nuclei, nbhd_type="nucleus", nbhd_col="cell_id") + series = nbhd_nuclei.calc_expansion(gdf_cells, radii_um=[5]) + + # non-Xenium transcripts.parquet: "name"/"x"/"y" instead of + # "feature_name"/"x_location"/"y_location" + pd.DataFrame( + { + "name": ["GeneA", "GeneB", "GeneA"], + "x": [5, 1, 25], + "y": [5, 1, 25], + } + ).to_parquet(tmp_path / "transcripts.parquet") + + nbhd_r5 = series[5.0] + nbhd_r5.calc_signature( + by="cell-free", + data_dir=str(tmp_path), + feature_col="name", + x_col="x", + y_col="y", + drop_missing=False, + ) + df_r5_custom = pd.DataFrame( + nbhd_r5.mod["gene_cell_free"].X, + index=nbhd_r5.mod["gene_cell_free"].obs_names, + columns=nbhd_r5.mod["gene_cell_free"].var_names, + ) + assert df_r5_custom.loc["c1", "GeneA"] == 1 + assert df_r5_custom.loc["c1", "GeneB"] == 1 + assert df_r5_custom.loc["c2", "GeneA"] == 1 diff --git a/tests/unit/test_nbhd/test_nbhd_collection.py b/tests/unit/test_nbhd/test_nbhd_collection.py index 2ff2924e..7ccef25d 100644 --- a/tests/unit/test_nbhd/test_nbhd_collection.py +++ b/tests/unit/test_nbhd/test_nbhd_collection.py @@ -194,6 +194,25 @@ def test_neighborhood_collection_transcript_assignment(tmp_path): assert obs["transcript_assignment_proportion"].loc["B"] == 0.0 +def test_transcript_assignment_finds_prefixed_transcripts_filename(tmp_path): + gdf, _adata = _synthetic_nbhd_inputs() + trx = pd.DataFrame( + { + "feature_name": ["g"] * 3, + "x_location": [1, 2, 11], + "y_location": [1, 2, 1], + "cell_id": ["c1", "c2", "UNASSIGNED"], + } + ) + # not every transcripts file is literally named "transcripts.parquet" + trx.to_parquet(tmp_path / "aziz_1_20260217_5_transcripts.parquet") + + collection = NeighborhoodCollection(gdf=gdf, nbhd_type="manual") + collection.calc_transcript_assignment(data_dir=str(tmp_path)) + + assert list(collection.obs["total_transcripts"]) == [2, 1] + + def test_transcript_assignment_warns_when_no_unassigned_sentinel(tmp_path): gdf, _adata = _synthetic_nbhd_inputs() # cell_id present but no "UNASSIGNED" sentinel -> warns (may be fully assigned) diff --git a/tests/unit/test_nbhd/test_trx_streaming.py b/tests/unit/test_nbhd/test_trx_streaming.py new file mode 100644 index 00000000..202dbf73 --- /dev/null +++ b/tests/unit/test_nbhd/test_trx_streaming.py @@ -0,0 +1,86 @@ +import geopandas as gpd +import pandas as pd +import pytest +from shapely.geometry import Polygon + +from celldega.nbhd.trx_streaming import _assign_trx_to_entity_streaming_parquet + + +def _synthetic_entities(): + return gpd.GeoDataFrame( + { + "cell_id": ["c1", "c2"], + "geometry": [ + Polygon([(0, 0), (10, 0), (10, 10), (0, 10)]), + Polygon([(20, 20), (30, 20), (30, 30), (20, 30)]), + ], + } + ) + + +def _write_trx_parquet(tmp_path, rows): + path = tmp_path / "transcripts.parquet" + pd.DataFrame(rows, columns=["x", "y", "gene"]).to_parquet(path) + return str(path) + + +def test_streaming_assignment_counts_points_per_entity_and_gene(tmp_path): + gdf_entity = _synthetic_entities() + trx_path = _write_trx_parquet( + tmp_path, + [ + (1, 1, "GeneA"), + (2, 2, "GeneA"), + (3, 3, "GeneB"), + (25, 25, "GeneA"), + (100, 100, "GeneA"), # outside every entity -> dropped + ], + ) + + counts = _assign_trx_to_entity_streaming_parquet(trx_path, gdf_entity, id_col="cell_id") + + assert counts.loc["c1", "GeneA"] == 2 + assert counts.loc["c1", "GeneB"] == 1 + assert counts.loc["c2", "GeneA"] == 1 + assert "c2" not in counts.index or counts.loc["c2"].get("GeneB", 0) == 0 + + +def test_streaming_assignment_batches_across_multiple_reads(tmp_path): + gdf_entity = _synthetic_entities() + rows = [(1, 1, "GeneA") for _ in range(5)] + [(25, 25, "GeneB") for _ in range(3)] + trx_path = _write_trx_parquet(tmp_path, rows) + + counts = _assign_trx_to_entity_streaming_parquet( + trx_path, gdf_entity, id_col="cell_id", batch_size=2 + ) + + assert counts.loc["c1", "GeneA"] == 5 + assert counts.loc["c2", "GeneB"] == 3 + + +def test_streaming_assignment_custom_column_names(tmp_path): + gdf_entity = _synthetic_entities() + path = tmp_path / "custom_trx.parquet" + pd.DataFrame({"xx": [1, 2], "yy": [1, 2], "name": ["GeneA", "GeneA"]}).to_parquet(path) + + counts = _assign_trx_to_entity_streaming_parquet( + str(path), gdf_entity, id_col="cell_id", x_col="xx", y_col="yy", gene_col="name" + ) + + assert counts.loc["c1", "GeneA"] == 2 + + +def test_streaming_assignment_raises_on_missing_id_col(tmp_path): + gdf_entity = _synthetic_entities() + trx_path = _write_trx_parquet(tmp_path, [(1, 1, "GeneA")]) + + with pytest.raises(KeyError): + _assign_trx_to_entity_streaming_parquet(trx_path, gdf_entity, id_col="not_a_column") + + +def test_streaming_assignment_no_matches_returns_empty_frame(tmp_path): + gdf_entity = _synthetic_entities() + trx_path = _write_trx_parquet(tmp_path, [(1000, 1000, "GeneA")]) + + counts = _assign_trx_to_entity_streaming_parquet(trx_path, gdf_entity, id_col="cell_id") + assert counts.empty diff --git a/tests/unit/test_nbhd/test_utils.py b/tests/unit/test_nbhd/test_utils.py new file mode 100644 index 00000000..7d249383 --- /dev/null +++ b/tests/unit/test_nbhd/test_utils.py @@ -0,0 +1,71 @@ +import pandas as pd +import pytest +from shapely.geometry import Polygon + +from celldega.nbhd import ( + make_column_names_unique_fast, + safe_polygon, + simple_format, + transform_polygon, +) +from celldega.nbhd.utils import _find_transcripts_parquet + + +def test_safe_polygon_builds_from_vertex_columns(): + row = pd.Series({"vertex_x": [0, 10, 10, 0], "vertex_y": [0, 0, 10, 10]}) + assert safe_polygon(row).area == 100.0 + + +def test_safe_polygon_returns_empty_on_malformed_row(): + row = pd.Series({"vertex_x": [0, 1], "vertex_y": [0]}) + assert safe_polygon(row).is_empty + + +def test_simple_format_rescales_coordinates(): + geometry = [[[10, 20], [30, 40]]] + assert simple_format(geometry, image_scale=2) == [[[5.0, 10.0], [15.0, 20.0]]] + + +def test_transform_polygon_returns_exterior_as_object_array(): + poly = Polygon([(0, 0), (1, 0), (1, 1)]) + result = transform_polygon(poly) + assert result.shape == (1, 4, 2) + assert list(result[0][0]) == [0, 0] + + +def test_make_column_names_unique_fast_dedupes_columns(): + df = pd.DataFrame([[1, 2, 3]], columns=["gene", "gene", "gene"]) + result = make_column_names_unique_fast(df) + assert list(result.columns) == ["gene", "gene_1", "gene_2"] + + +def test_find_transcripts_parquet_matches_literal_name(tmp_path): + (tmp_path / "transcripts.parquet").write_bytes(b"") + assert _find_transcripts_parquet(str(tmp_path)) == str(tmp_path / "transcripts.parquet") + + +def test_find_transcripts_parquet_matches_prefixed_name(tmp_path): + (tmp_path / "aziz_1_20260217_5_transcripts.parquet").write_bytes(b"") + assert _find_transcripts_parquet(str(tmp_path)) == str( + tmp_path / "aziz_1_20260217_5_transcripts.parquet" + ) + + +def test_find_transcripts_parquet_ignores_unrelated_files(tmp_path): + (tmp_path / "data1_transcripts.parquet").write_bytes(b"") + (tmp_path / "cells.parquet").write_bytes(b"") + (tmp_path / "notes.txt").write_bytes(b"") + assert _find_transcripts_parquet(str(tmp_path)) == str(tmp_path / "data1_transcripts.parquet") + + +def test_find_transcripts_parquet_raises_when_none_found(tmp_path): + (tmp_path / "cells.parquet").write_bytes(b"") + with pytest.raises(FileNotFoundError, match=r"transcripts\.parquet"): + _find_transcripts_parquet(str(tmp_path)) + + +def test_find_transcripts_parquet_raises_when_ambiguous(tmp_path): + (tmp_path / "data1_transcripts.parquet").write_bytes(b"") + (tmp_path / "data2_transcripts.parquet").write_bytes(b"") + with pytest.raises(ValueError, match="Multiple files"): + _find_transcripts_parquet(str(tmp_path))