diff --git a/src/squidpy/tl/_sliding_window.py b/src/squidpy/tl/_sliding_window.py index 5709d8591..710708b1d 100644 --- a/src/squidpy/tl/_sliding_window.py +++ b/src/squidpy/tl/_sliding_window.py @@ -1,6 +1,8 @@ from __future__ import annotations -from itertools import product +import warnings +from itertools import count, product +from typing import Literal import numpy as np import pandas as pd @@ -13,6 +15,10 @@ __all__ = ["sliding_window"] +# Label for cells that fall in no window (only possible for grid ``drop``/``adaptive``); kept as an +# explicit category instead of ``NaN`` so downstream code (and the ordered sort) never sees a float NaN. +UNASSIGNED = "unassigned" + @d.dedent def sliding_window( @@ -23,40 +29,92 @@ def sliding_window( coord_columns: tuple[str, str] = ("globalX", "globalY"), sliding_window_key: str = "sliding_window_assignment", spatial_key: str = "spatial", - drop_partial_windows: bool = False, + drop_partial_windows: bool | None = None, copy: bool = False, *, + method: Literal["grid", "split"] = "grid", + partial_windows: Literal["keep", "drop", "adaptive"] = "keep", + max_nr_cells: int | None = None, table_key: str | None = None, ) -> pd.DataFrame | None: """ - Divide a tissue slice into regulary shaped spatially contiguous regions (windows). + Divide a tissue slice into spatially contiguous regions (windows). + + Two tiling strategies are available via ``method``: + + - ``"grid"`` (default) lays a regular grid of ``window_size`` windows (optionally overlapping). + ``partial_windows`` controls the windows at the tissue edge. + - ``"split"`` recursively splits the cells into windows of roughly equal cell count + (at most ``max_nr_cells`` each), ignoring ``window_size``/``overlap``. Parameters ---------- %(adata)s - %(table_key)s - window_size: int - Size of the sliding window. %(library_key)s - coord_columns: Tuple[str, str] - Tuple of column names in `adata.obs` that specify the coordinates (x, y), e.i. ('globalX', 'globalY') - sliding_window_key: str - Base name for sliding window columns. + window_size: int | None + Size of each grid window (``method="grid"``). Inferred from the extent when ``None``. overlap: int - Overlap size between consecutive windows. (0 = no overlap) + Overlap between consecutive grid windows (0 = no overlap). Only used for ``method="grid"``. + A positive overlap produces one boolean column per window (an ``n_obs`` x ``n_windows`` table), + which can be memory-heavy for grids with many windows; ``overlap=0`` yields a single categorical + column instead. + coord_columns: tuple[str, str] + Column names in ``adata.obs`` holding the ``(x, y)`` coordinates, e.g. ``('globalX', 'globalY')``. + sliding_window_key: str + Base name for the sliding-window column(s) written to ``.obs``. %(spatial_key)s - drop_partial_windows: bool - If True, drop windows that are smaller than the window size at the borders. + drop_partial_windows: bool | None + Deprecated. Use ``partial_windows`` instead. ``True`` maps to ``partial_windows="drop"``. copy: bool - If True, return the result, otherwise save it to the adata object. + If ``True``, return the result; otherwise store it in ``adata.obs``. + method: Literal["grid", "split"] + Tiling strategy. ``"grid"`` for a regular grid, ``"split"`` for equal-cell-count windows. + partial_windows: Literal["keep", "drop", "adaptive"] + Edge-window handling for ``method="grid"`` (ignored for ``"split"``). + ``"keep"`` clips edge windows to the tissue bounds; ``"drop"`` removes windows that would extend + past the bounds (their cells become ``"unassigned"``); ``"adaptive"`` shrinks all windows slightly + so they tile the extent evenly. + max_nr_cells: int | None + Maximum number of cells per window. Required for (and only used by) ``method="split"``. + %(table_key)s Returns ------- - If ``copy = True``, returns the sliding window annotation(s) as pandas dataframe - Otherwise, stores the sliding window annotation(s) in .obs. + If ``copy = True``, returns the sliding-window annotation(s) as a :class:`pandas.DataFrame`. + Otherwise, stores the annotation(s) in ``adata.obs`` and returns ``None``. """ - if overlap < 0: - raise ValueError("Overlap must be non-negative.") + # --- deprecation: drop_partial_windows -> partial_windows --- + if drop_partial_windows is not None: + warnings.warn( + "`drop_partial_windows` is deprecated and will be removed in a future release; " + "use `partial_windows='drop'` (or 'keep') instead.", + FutureWarning, + stacklevel=2, + ) + if partial_windows != "keep": + raise ValueError("Pass either `drop_partial_windows` (deprecated) or `partial_windows`, not both.") + partial_windows = "drop" if drop_partial_windows else "keep" + + # --- validate arguments --- + if method not in ("grid", "split"): + raise ValueError(f"`method` must be 'grid' or 'split', got {method!r}.") + if partial_windows not in ("keep", "drop", "adaptive"): + raise ValueError(f"`partial_windows` must be 'keep', 'drop' or 'adaptive', got {partial_windows!r}.") + + if method == "split": + if max_nr_cells is None: + raise ValueError("`max_nr_cells` must be set when method='split'.") + if max_nr_cells < 1: + raise ValueError("`max_nr_cells` must be >= 1.") + if window_size is not None or overlap != 0 or partial_windows != "keep": + raise ValueError( + "`window_size`, `overlap` and `partial_windows` are not used with method='split'; leave them unset." + ) + else: # grid + if max_nr_cells is not None: + raise ValueError("`max_nr_cells` is only used with method='split'.") + if overlap < 0: + raise ValueError("Overlap must be non-negative.") adata = extract_adata_if_sdata(adata, table_key=table_key) @@ -64,7 +122,7 @@ def sliding_window( if copy: adata = adata.copy() - # extract coordinates of observations + # --- extract coordinates of observations --- x_col, y_col = coord_columns if x_col in adata.obs and y_col in adata.obs: coords = adata.obs[[x_col, y_col]].copy() @@ -79,41 +137,50 @@ def sliding_window( f"Coordinates not found. Provide `{coord_columns}` in `adata.obs` or specify a suitable `spatial_key` in `adata.obsm`." ) - # infer window size if not provided - if window_size is None: - coord_range = max( - coords[x_col].max() - coords[x_col].min(), - coords[y_col].max() - coords[y_col].min(), - ) - # mostly arbitrary choice, except that full integers usually generate windows with 1-2 cells at the borders - window_size = max(int(np.floor(coord_range // 3.95)), 1) - - if window_size <= 0: - raise ValueError("Window size must be larger than 0.") + # --- grid: infer + validate window size --- + if method == "grid": + if window_size is None: + coord_range = max( + coords[x_col].max() - coords[x_col].min(), + coords[y_col].max() - coords[y_col].min(), + ) + # mostly arbitrary choice, except that full integers usually generate windows with 1-2 cells at the borders + window_size = max(int(np.floor(coord_range // 3.95)), 1) + if window_size <= 0: + raise ValueError("Window size must be larger than 0.") + if overlap >= window_size: + raise ValueError("Overlap must be less than the window size.") + if partial_windows == "adaptive" and overlap >= window_size // 2: + raise ValueError("Overlap must be less than `window_size` // 2 when partial_windows='adaptive'.") if library_key is not None and library_key not in adata.obs: raise ValueError(f"Library key '{library_key}' not found in adata.obs") libraries = [None] if library_key is None else adata.obs[library_key].unique() - # Create a DataFrame to store the sliding window assignments - sliding_window_df = pd.DataFrame(index=adata.obs.index) - if sliding_window_key in adata.obs: logg.warning(f"Overwriting existing column '{sliding_window_key}' in adata.obs.") + sliding_window_df = pd.DataFrame(index=adata.obs.index) + # For overlapping grids we emit one boolean column per window. Collect them all and concatenate once + # at the end: adding them one-by-one fragments the frame and is quadratic in the number of windows. + bool_columns: dict[str, pd.Series] = {} + for lib in libraries: - if lib is not None: - lib_mask = adata.obs[library_key] == lib - lib_coords = coords.loc[lib_mask] - else: - lib_mask = np.ones(len(adata), dtype=bool) - lib_coords = coords + lib_coords = coords.loc[adata.obs[library_key] == lib] if lib is not None else coords + lib_key = f"{lib}_" if lib is not None else "" + + if method == "split": + # each cell is assigned to exactly one window (non-overlapping by construction) + labels = _split_cells(lib_coords, coord_columns, max_nr_cells) + for label in np.unique(labels): + obs_indices = lib_coords.index[labels == label] + sliding_window_df.loc[obs_indices, sliding_window_key] = f"{lib_key}window_{label}" + continue min_x, max_x = lib_coords[x_col].min(), lib_coords[x_col].max() min_y, max_y = lib_coords[y_col].min(), lib_coords[y_col].max() - # precalculate windows windows = _calculate_window_corners( min_x=min_x, max_x=max_x, @@ -121,122 +188,198 @@ def sliding_window( max_y=max_y, window_size=window_size, overlap=overlap, - drop_partial_windows=drop_partial_windows, + partial_windows=partial_windows, ) - lib_key = f"{lib}_" if lib is not None else "" - - # assign observations to windows for idx, window in windows.iterrows(): - x_start = window["x_start"] - x_end = window["x_end"] - y_start = window["y_start"] - y_end = window["y_end"] - - mask = ( - (lib_coords[x_col] >= x_start) - & (lib_coords[x_col] <= x_end) - & (lib_coords[y_col] >= y_start) - & (lib_coords[y_col] <= y_end) + mask = _get_window_mask( + coord_columns=coord_columns, + lib_coords=lib_coords, + x_start=window["x_start"], + x_end=window["x_end"], + y_start=window["y_start"], + y_end=window["y_end"], ) obs_indices = lib_coords.index[mask] - if overlap == 0: - mask = ( - (lib_coords[x_col] >= x_start) - & (lib_coords[x_col] <= x_end) - & (lib_coords[y_col] >= y_start) - & (lib_coords[y_col] <= y_end) - ) - obs_indices = lib_coords.index[mask] sliding_window_df.loc[obs_indices, sliding_window_key] = f"{lib_key}window_{idx}" - else: col_name = f"{sliding_window_key}_{lib_key}window_{idx}" - sliding_window_df.loc[obs_indices, col_name] = True - # Avoid chained assignment for pandas CoW compatibility - sliding_window_df[col_name] = sliding_window_df[col_name].fillna(False) - - if overlap == 0: - # create categorical variable for ordered windows - sliding_window_df[sliding_window_key] = pd.Categorical( - sliding_window_df[sliding_window_key], - ordered=True, - categories=sorted( - sliding_window_df[sliding_window_key].unique(), - key=lambda x: int(x.split("_")[-1]), - ), - ) + col = bool_columns.setdefault(col_name, pd.Series(False, index=sliding_window_df.index)) + col.loc[obs_indices] = True - sliding_window_df[x_col] = coords[x_col] - sliding_window_df[y_col] = coords[y_col] + if bool_columns: + sliding_window_df = pd.concat([sliding_window_df, pd.DataFrame(bool_columns)], axis=1) + + if method == "split" or overlap == 0: + # single categorical column: order windows by their trailing index, put unassigned cells last + sliding_window_df[sliding_window_key] = _ordered_window_categorical(sliding_window_df[sliding_window_key]) if copy: return sliding_window_df for col_name, col_data in sliding_window_df.items(): _save_data(adata, attr="obs", key=col_name, data=col_data) + return None + + +def _ordered_window_categorical(values: pd.Series) -> pd.Categorical: + """Ordered categorical of window labels sorted by trailing index; unassigned cells (``NaN``) go last. + + Cells outside every window (grid ``drop``/``adaptive``) arrive as ``NaN``; they become an explicit + ``"unassigned"`` category so the ordered sort never calls ``int(...)`` on a float ``NaN``. + """ + filled = values.fillna(UNASSIGNED) + present = list(pd.unique(filled)) + windows = sorted((c for c in present if c != UNASSIGNED), key=lambda s: int(str(s).split("_")[-1])) + categories = windows + ([UNASSIGNED] if UNASSIGNED in present else []) + return pd.Categorical(filled, ordered=True, categories=categories) + + +def _get_window_mask( + coord_columns: tuple[str, str], + lib_coords: pd.DataFrame, + x_start: float, + x_end: float, + y_start: float, + y_end: float, +) -> pd.Series: + """Boolean mask selecting the rows of ``lib_coords`` inside the (inclusive) window.""" + x_col, y_col = coord_columns + return ( + (lib_coords[x_col] >= x_start) + & (lib_coords[x_col] <= x_end) + & (lib_coords[y_col] >= y_start) + & (lib_coords[y_col] <= y_end) + ) + + +def _split_cells(coords: pd.DataFrame, coord_columns: tuple[str, str], max_cells: int) -> np.ndarray: + """Assign each cell to a window by recursive count-based (median) splitting. + + Each window holds at most ``max_cells`` cells and, unless the whole input is smaller, at least + ``max_cells // 2``. The split is on cell *position* (the median index of the longer axis), so windows + are **non-overlapping by construction** — no cell can land in two windows, and every split strictly + shrinks both halves, so it always terminates (given ``max_cells >= 1``). + + Parameters + ---------- + coords + Coordinates for one library (index-aligned to the cells). + coord_columns + ``(x_col, y_col)`` column names in ``coords``. + max_cells + Maximum number of cells per window. + + Returns + ------- + Integer window label per row of ``coords`` (positional order). + """ + x_col, y_col = coord_columns + x = coords[x_col].to_numpy() + y = coords[y_col].to_numpy() + labels = np.empty(len(coords), dtype=int) + counter = count() + + def recurse(idx: np.ndarray) -> None: + if len(idx) <= max_cells: + labels[idx] = next(counter) + return + xi, yi = x[idx], y[idx] + # split along the axis with the larger spatial extent, at the median cell + if (xi.max() - xi.min()) >= (yi.max() - yi.min()): + order = idx[np.argsort(xi, kind="stable")] + else: + order = idx[np.argsort(yi, kind="stable")] + mid = len(order) // 2 + recurse(order[:mid]) + recurse(order[mid:]) + + recurse(np.arange(len(coords), dtype=int)) + return labels def _calculate_window_corners( - min_x: int, - max_x: int, - min_y: int, - max_y: int, + min_x: float, + max_x: float, + min_y: float, + max_y: float, window_size: int, overlap: int = 0, - drop_partial_windows: bool = False, + partial_windows: Literal["keep", "drop", "adaptive"] = "keep", ) -> pd.DataFrame: """ - Calculate the corner points of all windows covering the area from min_x to max_x and min_y to max_y, - with specified window_size and overlap. + Corner points of a regular grid of windows covering ``[min_x, max_x] x [min_y, max_y]``. Parameters ---------- - min_x: float - minimum X coordinate - max_x: float - maximum X coordinate - min_y: float - minimum Y coordinate - max_y: float - maximum Y coordinate - window_size: float - size of each window - overlap: float - overlap between consecutive windows (must be less than window_size) - drop_partial_windows: bool - if True, drop border windows that are smaller than window_size; - if False, create smaller windows at the borders to cover the remaining space. + min_x, max_x, min_y, max_y + Extent to tile. + window_size + Size of each window. + overlap + Overlap between consecutive windows (must be less than ``window_size``). + partial_windows + Edge handling: ``"keep"`` clips edge windows to the bounds; ``"drop"`` removes windows that would + extend past the bounds; ``"adaptive"`` shrinks all windows slightly to tile the extent evenly. Returns ------- - windows: pandas DataFrame with columns ['x_start', 'x_end', 'y_start', 'y_end'] + DataFrame with columns ``['x_start', 'x_end', 'y_start', 'y_end']``. """ if overlap < 0: raise ValueError("Overlap must be non-negative.") if overlap >= window_size: raise ValueError("Overlap must be less than the window size.") - x_step = window_size - overlap - y_step = window_size - overlap + if partial_windows == "adaptive": + total_width = max_x - min_x + total_height = max_y - min_y + # number of windows per axis; clamp to >= 1 so a library smaller than one window (e.g. span + # <= overlap, common when the global window_size is set from a larger library) yields a single + # window instead of dividing by zero. + number_x_windows = max(int(np.ceil((total_width - overlap) / (window_size - overlap))), 1) + number_y_windows = max(int(np.ceil((total_height - overlap) / (window_size - overlap))), 1) + # window size per axis (integer to avoid float drift) + x_window_size = np.ceil((total_width + (number_x_windows - 1) * overlap) / number_x_windows) + y_window_size = np.ceil((total_height + (number_y_windows - 1) * overlap) / number_y_windows) + else: + x_window_size = window_size + y_window_size = window_size + + x_step = x_window_size - overlap + y_step = y_window_size - overlap - # Generate starting points - x_starts = np.arange(min_x, max_x, x_step) - y_starts = np.arange(min_y, max_y, y_step) + # Generate starting points. A non-positive step means one window already covers the whole span + # (span <= overlap) -> emit a single window at the minimum rather than an empty grid. + x_starts = np.arange(min_x, max_x, x_step) if x_step > 0 else np.array([min_x]) + y_starts = np.arange(min_y, max_y, y_step) if y_step > 0 else np.array([min_y]) # Create all combinations of x and y starting points starts = list(product(x_starts, y_starts)) windows = pd.DataFrame(starts, columns=["x_start", "y_start"]) - windows["x_end"] = windows["x_start"] + window_size - windows["y_end"] = windows["y_start"] + window_size + windows["x_end"] = windows["x_start"] + x_window_size + windows["y_end"] = windows["y_start"] + y_window_size - # Adjust windows that extend beyond the bounds - if not drop_partial_windows: + if partial_windows == "keep": windows["x_end"] = windows["x_end"].clip(upper=max_x) windows["y_end"] = windows["y_end"].clip(upper=max_y) - else: + elif partial_windows == "adaptive": + # the integer window size can exceed max_x/max_y -> clip, then drop degenerate corner slivers. + # Only drop a thin window when its axis has neighbours (>1 window): a sole window covering a + # small library is thin but not redundant, and must be kept. + windows["x_end"] = windows["x_end"].clip(upper=max_x) + windows["y_end"] = windows["y_end"].clip(upper=max_y) + thin_x = (windows["x_end"] - windows["x_start"]) <= overlap + thin_y = (windows["y_end"] - windows["y_start"]) <= overlap + # a thin window is a redundant sliver only if its axis actually has more than one window; + # a lone window covering a small library is thin but must be kept. + redundant_windows = (thin_x & (len(x_starts) > 1)) | (thin_y & (len(y_starts) > 1)) + windows = windows[~redundant_windows] + elif partial_windows == "drop": valid_windows = (windows["x_end"] <= max_x) & (windows["y_end"] <= max_y) windows = windows[valid_windows] + else: + raise ValueError(f"{partial_windows} is not a valid `partial_windows` argument.") windows = windows.reset_index(drop=True) return windows[["x_start", "x_end", "y_start", "y_end"]] diff --git a/tests/tools/test_sliding_window.py b/tests/tools/test_sliding_window.py index 3dd670b00..76ff0f02c 100644 --- a/tests/tools/test_sliding_window.py +++ b/tests/tools/test_sliding_window.py @@ -1,24 +1,40 @@ from __future__ import annotations +import numpy as np +import pandas as pd import pytest from anndata import AnnData from squidpy.tl import _calculate_window_corners, sliding_window +from squidpy.tl._sliding_window import _split_cells + + +def _grid_adata(n_per_side: int = 30, extent: float = 300.0, seed: int = 0, library_key: str | None = None) -> AnnData: + """A uniform point cloud in an ``extent`` x ``extent`` square (optionally split across libraries).""" + rng = np.random.default_rng(seed) + n = n_per_side * n_per_side + xy = rng.uniform(0, extent, size=(n, 2)) + obs = pd.DataFrame({"globalX": xy[:, 0], "globalY": xy[:, 1]}, index=[f"c{i}" for i in range(n)]) + if library_key is not None: + obs[library_key] = rng.choice(["a", "b"], size=n) + return AnnData(X=np.zeros((n, 1), dtype=np.float32), obs=obs) class TestSlidingWindow: @pytest.mark.parametrize( - "windowsize_overlap_drop", + "windowsize_overlap_partial", [ - (300, 0, False), - (300, 50, False), - (300, 50, True), + (300, 0, "keep"), + (300, 50, "keep"), + (300, 50, "drop"), + (300, 0, "adaptive"), + (300, 50, "adaptive"), ], ) def test_sliding_window_several_slices( self, adata_mibitof: AnnData, - windowsize_overlap_drop: tuple[int, int, bool], + windowsize_overlap_partial: tuple[int, int, str], sliding_window_key: str = "sliding_window_key", library_key: str = "library_id", ): @@ -30,7 +46,7 @@ def _count_total_assignments(): total_cells += df[col].sum() return total_cells - window_size, overlap, drop_partial_windows = windowsize_overlap_drop + window_size, overlap, partial_windows = windowsize_overlap_partial df = sliding_window( adata_mibitof, library_key=library_key, @@ -39,20 +55,23 @@ def _count_total_assignments(): coord_columns=("globalX", "globalY"), sliding_window_key=sliding_window_key, copy=True, - drop_partial_windows=drop_partial_windows, + partial_windows=partial_windows, ) if overlap == 0: sliding_window_columns = [col for col in df.columns if sliding_window_key in col] assert len(sliding_window_columns) == 1 # only one sliding window - assert df[sliding_window_key].isnull().sum() == 0 # no unassigned cells + assert df[sliding_window_key].isnull().sum() == 0 # no NaN (unassigned is an explicit category) assert len(df) == adata_mibitof.n_obs # correct amount of rows else: sliding_window_cols = df.columns[df.columns.str.contains("sliding_window")] - if drop_partial_windows: + if partial_windows == "drop": assert len(sliding_window_cols) == 27 assert _count_total_assignments() == 2536 + elif partial_windows == "adaptive": + assert len(sliding_window_cols) == 48 + assert _count_total_assignments() == 4411 else: assert len(sliding_window_cols) == 70 assert _count_total_assignments() == 4569 @@ -92,86 +111,133 @@ def test_sliding_window_invalid_window_size( ): with pytest.raises(ValueError, match="Window size must be larger than 0."): sliding_window( - adata_squaregrid, - window_size=-10, - overlap=0, - coord_columns=("globalX", "globalY"), - sliding_window_key="sliding_window", - copy=True, + adata_squaregrid, window_size=-10, overlap=0, coord_columns=("globalX", "globalY"), copy=True ) with pytest.raises(ValueError, match="Overlap must be non-negative."): sliding_window( - adata_squaregrid, - window_size=10, - overlap=-10, - coord_columns=("globalX", "globalY"), - sliding_window_key="sliding_window", - copy=True, + adata_squaregrid, window_size=10, overlap=-10, coord_columns=("globalX", "globalY"), copy=True ) - def test_calculate_window_corners_overlap(self): - min_x = 0 - max_x = 200 - min_y = 0 - max_y = 200 - window_size = 100 - overlap = 20 + with pytest.raises(ValueError, match="`max_nr_cells` must be set when method='split'."): + sliding_window(adata_squaregrid, method="split", coord_columns=("globalX", "globalY"), copy=True) + + def test_sliding_window_method_validation(self): + adata = _grid_adata() + with pytest.raises(ValueError, match="must be 'grid' or 'split'"): + sliding_window(adata, method="nope", copy=True) # type: ignore[arg-type] + with pytest.raises(ValueError, match=">= 1"): + sliding_window(adata, method="split", max_nr_cells=0, copy=True) + with pytest.raises(ValueError, match="only used with method='split'"): + sliding_window(adata, method="grid", max_nr_cells=100, copy=True) + with pytest.raises(ValueError, match="not used with method='split'"): + sliding_window(adata, method="split", max_nr_cells=100, overlap=5, copy=True) + with pytest.raises(ValueError, match="not used with method='split'"): + sliding_window(adata, method="split", max_nr_cells=100, window_size=50, copy=True) + + def test_sliding_window_drop_overlap0_no_crash(self): + """Regression: drop + overlap==0 used to raise in the categorical sort; unassigned cells are labelled.""" + adata = _grid_adata() + df = sliding_window(adata, window_size=100, overlap=0, partial_windows="drop", copy=True) + col = df["sliding_window_assignment"] + assert col.isnull().sum() == 0 # no NaN + assert "unassigned" in list(col.cat.categories) + assert col.cat.categories[-1] == "unassigned" # ordered last + assert (col == "unassigned").sum() > 0 # drop genuinely strands some edge cells here + + def test_sliding_window_deprecated_drop_partial_windows(self): + adata = _grid_adata() + with pytest.warns(FutureWarning, match="drop_partial_windows"): + deprecated = sliding_window(adata, window_size=100, overlap=0, drop_partial_windows=True, copy=True) + current = sliding_window(adata, window_size=100, overlap=0, partial_windows="drop", copy=True) + assert ( + deprecated["sliding_window_assignment"].astype(str).equals(current["sliding_window_assignment"].astype(str)) + ) + with pytest.raises(ValueError, match="not both"): + sliding_window(adata, drop_partial_windows=True, partial_windows="drop", copy=True) + + def test_sliding_window_adaptive_small_library(self): + """Regression (B5): a library smaller than one window must not divide-by-zero / drop its only window.""" + rng = np.random.default_rng(1) + big = rng.uniform(0, 1000, size=(500, 2)) + tiny = rng.uniform(0, 10, size=(20, 2)) # span (~10) <= overlap (50) + xy = np.vstack([big, tiny]) + obs = pd.DataFrame( + {"globalX": xy[:, 0], "globalY": xy[:, 1], "library_id": ["big"] * 500 + ["tiny"] * 20}, + index=[f"c{i}" for i in range(520)], + ) + adata = AnnData(X=np.zeros((520, 1), dtype=np.float32), obs=obs) + df = sliding_window( + adata, library_key="library_id", window_size=200, overlap=50, partial_windows="adaptive", copy=True + ) + tiny_cols = [c for c in df.columns if "tiny_" in c] + assert tiny_cols # the tiny library still gets a window + assert (df.iloc[500:][tiny_cols].sum(axis=1) > 0).all() # every tiny cell is covered - windows = _calculate_window_corners( - min_x=min_x, - max_x=max_x, - min_y=min_y, - max_y=max_y, - window_size=window_size, - overlap=overlap, - drop_partial_windows=False, + def test_sliding_window_split_nr_cells( + self, + adata_mibitof: AnnData, + sliding_window_key: str = "sliding_window_key", + library_key: str = "library_id", + ): + """Each window holds <= max_nr_cells and >= max_nr_cells // 2 cells (per library).""" + max_nr_cells = 100 + df = sliding_window( + adata_mibitof, + library_key=library_key, + sliding_window_key=sliding_window_key, + method="split", + max_nr_cells=max_nr_cells, + copy=True, ) + counts = df[sliding_window_key].value_counts() + assert df[sliding_window_key].isnull().sum() == 0 # split covers every cell + assert counts.max() <= max_nr_cells + assert (counts >= max_nr_cells // 2).all() + + def test_split_cells_partition_and_bounds(self): + """_split_cells partitions the cells (each in exactly one window) and respects the bounds, even with ties.""" + rng = np.random.default_rng(2) + coords = pd.DataFrame({"globalX": rng.uniform(0, 100, 1000), "globalY": rng.uniform(0, 100, 1000)}) + labels = _split_cells(coords, ("globalX", "globalY"), max_cells=50) + assert len(labels) == len(coords) # one label per cell -> a partition (non-overlapping, full cover) + counts = pd.Series(labels).value_counts() + assert counts.max() <= 50 + assert (counts >= 25).all() + + # degenerate: all cells share a coordinate -> must still terminate and stay bounded + same = pd.DataFrame({"globalX": np.zeros(200), "globalY": np.zeros(200)}) + labels = _split_cells(same, ("globalX", "globalY"), max_cells=50) + assert pd.Series(labels).value_counts().max() <= 50 + def test_calculate_window_corners_overlap(self): + windows = _calculate_window_corners( + min_x=0, max_x=200, min_y=0, max_y=200, window_size=100, overlap=20, partial_windows="keep" + ) assert windows.shape == (9, 4) assert windows.iloc[0].values.tolist() == [0, 100, 0, 100] assert windows.iloc[-1].values.tolist() == [160, 200, 160, 200] def test_calculate_window_corners_no_overlap(self): - min_x = 0 - max_x = 200 - min_y = 0 - max_y = 200 - window_size = 100 - overlap = 0 - windows = _calculate_window_corners( - min_x=min_x, - max_x=max_x, - min_y=min_y, - max_y=max_y, - window_size=window_size, - overlap=overlap, - drop_partial_windows=False, + min_x=0, max_x=200, min_y=0, max_y=200, window_size=100, overlap=0, partial_windows="keep" ) - assert windows.shape == (4, 4) assert windows.iloc[0].values.tolist() == [0, 100, 0, 100] assert windows.iloc[-1].values.tolist() == [100, 200, 100, 200] def test_calculate_window_corners_drop_partial_windows(self): - min_x = 0 - max_x = 200 - min_y = 0 - max_y = 200 - window_size = 100 - overlap = 20 - windows = _calculate_window_corners( - min_x=min_x, - max_x=max_x, - min_y=min_y, - max_y=max_y, - window_size=window_size, - overlap=overlap, - drop_partial_windows=True, + min_x=0, max_x=200, min_y=0, max_y=200, window_size=100, overlap=20, partial_windows="drop" ) - assert windows.shape == (4, 4) assert windows.iloc[0].values.tolist() == [0, 100, 0, 100] assert windows.iloc[-1].values.tolist() == [80, 180, 80, 180] + + def test_calculate_window_corners_adaptive_partial_windows(self): + windows = _calculate_window_corners( + min_x=0, max_x=200, min_y=0, max_y=200, window_size=100, overlap=20, partial_windows="adaptive" + ) + assert windows.shape == (9, 4) + assert windows.iloc[0].values.tolist() == [0, 80, 0, 80] + assert windows.iloc[-1].values.tolist() == [120, 200, 120, 200]