From f35750a9b454bf46261533aa48532730a90c6f6a Mon Sep 17 00:00:00 2001 From: Erik Date: Fri, 4 Sep 2026 16:06:33 +0200 Subject: [PATCH] add distance_from_node_set for hop distances from seed nodes on a PNAGraph Co-authored-by: Cursor --- CHANGELOG.md | 1 + docs/api/overview.rst | 1 + src/pixelator/pna/analysis/__init__.py | 7 +- .../pna/analysis/segmentation/__init__.py | 8 + .../pna/analysis/segmentation/distance.py | 145 +++++++++++++++++ .../analysis/test_distance_from_node_set.py | 153 ++++++++++++++++++ 6 files changed, 314 insertions(+), 1 deletion(-) create mode 100644 src/pixelator/pna/analysis/segmentation/__init__.py create mode 100644 src/pixelator/pna/analysis/segmentation/distance.py create mode 100644 tests/pna/analysis/test_distance_from_node_set.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a41510dd..b76a0b8f9 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Added - `PNAPixelDataset.layouts()` computes Layouts on the fly, making it easier to work with cell layouts. - `pixelator.pna.analysis.summarize_proximity_scores` to collapse a per-component proximity score table into one row per marker pair. +- `pixelator.pna.analysis.distance_from_node_set` to compute integer hop distances from a set of seed nodes on a `PNAGraph` (unreached nodes stay missing). ### Changed - `density_scatter_plot` now lives in `pixelator.plot` (previously `pixelator.mpx.plot`). diff --git a/docs/api/overview.rst b/docs/api/overview.rst index 25cad9d60..959c9d3dc 100644 --- a/docs/api/overview.rst +++ b/docs/api/overview.rst @@ -37,4 +37,5 @@ for usage examples. **Analysis** * :func:`pixelator.pna.analysis.calculate_differential_proximity` +* :func:`pixelator.pna.analysis.distance_from_node_set` * :func:`pixelator.pna.analysis.summarize_proximity_scores` diff --git a/src/pixelator/pna/analysis/__init__.py b/src/pixelator/pna/analysis/__init__.py index 3689e6533..cb20977a1 100644 --- a/src/pixelator/pna/analysis/__init__.py +++ b/src/pixelator/pna/analysis/__init__.py @@ -7,8 +7,13 @@ calculate_differential_proximity, summarize_proximity_scores, ) +from pixelator.pna.analysis.segmentation import distance_from_node_set -__all__ = ["calculate_differential_proximity", "summarize_proximity_scores"] +__all__ = [ + "calculate_differential_proximity", + "distance_from_node_set", + "summarize_proximity_scores", +] # Note: pixelator.pna.analysis.comparison is intentionally not imported here. # It depends on pixelator.pna.pixeldataset, which itself imports diff --git a/src/pixelator/pna/analysis/segmentation/__init__.py b/src/pixelator/pna/analysis/segmentation/__init__.py new file mode 100644 index 000000000..41d0d02f5 --- /dev/null +++ b/src/pixelator/pna/analysis/segmentation/__init__.py @@ -0,0 +1,8 @@ +"""Helpers for segmenting cell:cell conjugates. + +Copyright © 2026 Pixelgen Technologies AB. +""" + +from pixelator.pna.analysis.segmentation.distance import distance_from_node_set + +__all__ = ["distance_from_node_set"] diff --git a/src/pixelator/pna/analysis/segmentation/distance.py b/src/pixelator/pna/analysis/segmentation/distance.py new file mode 100644 index 000000000..840c1ec90 --- /dev/null +++ b/src/pixelator/pna/analysis/segmentation/distance.py @@ -0,0 +1,145 @@ +"""Hop distance from a set of seed nodes on a cell graph. + +Copyright © 2026 Pixelgen Technologies AB. +""" + +from __future__ import annotations + +from collections.abc import Hashable, Sequence +from typing import Any + +import networkx as nx +import numpy as np + +from pixelator.common.utils import logger +from pixelator.pna.graph import PNAGraph + +_DISTANCE_ATTR = "distance_from_seed" + + +def distance_from_node_set( + graph: PNAGraph, + seed_nodes: Hashable | Sequence[Hashable], + max_iter: int = 40, + verbose: bool = False, +) -> PNAGraph: + """Compute integer hop distance from seed nodes on a cell graph. + + Runs a multi-source breadth-first search on ``graph``. Each seed has + distance 0. Every other node gets the length of the shortest unweighted + path to the nearest seed, up to ``max_iter`` hops. Nodes that are never + reached keep a missing value. + + The result is stored as the node attribute ``distance_from_seed``, + replacing that attribute if it already exists. The same ``PNAGraph`` + instance is updated in place and returned. + + Args: + graph: Component graph to annotate, typically + ``component.graph`` from a ``PNAPixelDataset`` edgelist + iterator. + seed_nodes: One node name or a sequence of node names that must + all be present in ``graph``. + max_iter: Maximum hop distance to compute. Nodes farther than this + (and disconnected nodes) stay missing. Default 40. + verbose: If True, log how many new nodes are reached at each + iteration. Default False. + + Returns: + The same ``PNAGraph``, with integer ``distance_from_seed`` on + reached nodes and ``None`` on unreached nodes. + + Raises: + TypeError: If ``graph`` is not a ``PNAGraph``, or if ``max_iter`` + or ``verbose`` have the wrong type. + ValueError: If ``seed_nodes`` is empty, a seed is missing from + the graph, or ``max_iter`` is negative. + + Examples: + Distances from one node on a component graph:: + + from pixelator.pna.analysis import distance_from_node_set + from pixelator.pna.pixeldataset import read + + component = next(read("sample.pxl").edgelist().iterator()) + seed = next(iter(component.graph.raw.nodes)) + distance_from_node_set(component.graph, seed) + + See Also: + ``distance_from_node_set`` in pixelatorR, the equivalent function + for R users. + + """ + seeds = _validate_distance_from_node_set_params( + graph=graph, + seed_nodes=seed_nodes, + max_iter=max_iter, + verbose=verbose, + ) + max_iter = int(max_iter) + + raw = graph.raw + distances: dict[Any, int | None] = {node: None for node in raw.nodes} + for seed in seeds: + distances[seed] = 0 + + frontier = list(dict.fromkeys(seeds)) + for iteration in range(1, max_iter + 1): + next_frontier: list[Any] = [] + seen_next: set[Any] = set() + for node in frontier: + for neighbor in raw.neighbors(node): + if distances[neighbor] is None and neighbor not in seen_next: + distances[neighbor] = iteration + next_frontier.append(neighbor) + seen_next.add(neighbor) + if not next_frontier: + break + if verbose: + logger.info( + "Iteration %s: %s new nodes reached.", + iteration, + len(next_frontier), + ) + frontier = next_frontier + + nx.set_node_attributes(raw, distances, _DISTANCE_ATTR) + return graph + + +def _validate_distance_from_node_set_params( + *, + graph: PNAGraph, + seed_nodes: Hashable | Sequence[Hashable], + max_iter: int, + verbose: bool, +) -> list[Hashable]: + if not isinstance(graph, PNAGraph): + raise TypeError("graph must be a PNAGraph.") + if not isinstance(max_iter, (int, np.integer)) or isinstance(max_iter, bool): + raise TypeError("max_iter must be an int.") + if int(max_iter) < 0: + raise ValueError("max_iter must be >= 0.") + if not isinstance(verbose, bool): + raise TypeError("verbose must be a bool.") + + seeds = _as_seed_list(seed_nodes) + missing = [seed for seed in seeds if seed not in graph.raw] + if missing: + raise ValueError( + "All seed nodes must be present in the graph. " + f"The following seed nodes are not present in the graph: {missing}" + ) + return seeds + + +def _as_seed_list(seed_nodes: Hashable | Sequence[Hashable]) -> list[Hashable]: + if isinstance(seed_nodes, (str, bytes)): + seeds: list[Hashable] = [seed_nodes] + elif isinstance(seed_nodes, Sequence): + seeds = list(seed_nodes) + else: + seeds = [seed_nodes] + if not seeds: + raise ValueError("seed_nodes must contain at least one node.") + return seeds diff --git a/tests/pna/analysis/test_distance_from_node_set.py b/tests/pna/analysis/test_distance_from_node_set.py new file mode 100644 index 000000000..aa97b83ff --- /dev/null +++ b/tests/pna/analysis/test_distance_from_node_set.py @@ -0,0 +1,153 @@ +"""Tests for `pixelator.pna.analysis.segmentation.distance_from_node_set`. + +Copyright © 2026 Pixelgen Technologies AB. +""" + +from __future__ import annotations + +import inspect + +import networkx as nx +import numpy as np +import polars as pl +import pytest + +from pixelator.pna.analysis.segmentation import distance_from_node_set +from pixelator.pna.graph import PNAGraph + +# Bipartite line a-b-c-d-e plus a disconnected edge f-g: +# a -- b -- c -- d -- e f -- g +_EDGES = { + "umi1": ["a", "c", "c", "e", "f"], + "umi2": ["b", "b", "d", "d", "g"], + "marker_1": ["MA", "MC", "MC", "ME", "MF"], + "marker_2": ["MB", "MB", "MD", "MD", "MG"], + "read_count": [1, 1, 1, 1, 1], +} + + +def _synthetic_graph() -> PNAGraph: + return PNAGraph.from_edgelist(pl.DataFrame(_EDGES).lazy()) + + +def _distances(graph: PNAGraph) -> dict: + return nx.get_node_attributes(graph.raw, "distance_from_seed") + + +@pytest.fixture +def graph() -> PNAGraph: + return _synthetic_graph() + + +def test_distance_from_node_set_single_seed_exact_hops(graph): + result = distance_from_node_set(graph, "a") + + assert result is graph + distances = _distances(graph) + assert distances == { + "a": 0, + "b": 1, + "c": 2, + "d": 3, + "e": 4, + "f": None, + "g": None, + } + assert all(isinstance(distances[n], int) for n in "abcde") + + +def test_distance_from_node_set_multiple_seeds(graph): + distance_from_node_set(graph, ["a", "e"]) + distances = _distances(graph) + + assert distances["a"] == 0 + assert distances["e"] == 0 + assert distances["b"] == 1 + assert distances["d"] == 1 + assert distances["c"] == 2 + assert distances["f"] is None + assert distances["g"] is None + assert sum(d == 0 for d in distances.values() if d is not None) == 2 + + +def test_distance_from_node_set_respects_max_iter(graph): + distance_from_node_set(graph, "a", max_iter=2) + distances = _distances(graph) + + assert distances["a"] == 0 + assert distances["b"] == 1 + assert distances["c"] == 2 + assert distances["d"] is None + assert distances["e"] is None + assert distances["f"] is None + assert distances["g"] is None + reached = [d for d in distances.values() if d is not None] + assert max(reached) == 2 + + +def test_distance_from_node_set_max_iter_zero(graph): + distance_from_node_set(graph, ["a", "c"], max_iter=0) + distances = _distances(graph) + + assert distances["a"] == 0 + assert distances["c"] == 0 + assert all(distances[n] is None for n in "bdefg") + + +def test_distance_from_node_set_replaces_existing_attribute(graph): + nx.set_node_attributes(graph.raw, 999, "distance_from_seed") + distance_from_node_set(graph, "a") + distances = _distances(graph) + + assert 999 not in distances.values() + assert distances["a"] == 0 + assert distances["f"] is None + + +def test_distance_from_node_set_missing_seed_raises(graph): + with pytest.raises(ValueError, match="seed nodes must be present"): + distance_from_node_set(graph, "not_a_real_node") + + +def test_distance_from_node_set_partially_missing_seeds_raises(graph): + with pytest.raises(ValueError, match="seed nodes must be present"): + distance_from_node_set(graph, ["a", "missing"]) + + +def test_distance_from_node_set_empty_seeds_raises(graph): + with pytest.raises(ValueError, match="at least one node"): + distance_from_node_set(graph, []) + + +def test_distance_from_node_set_invalid_graph_raises(graph): + with pytest.raises(TypeError, match="PNAGraph"): + distance_from_node_set("not a graph", "a") + + +def test_distance_from_node_set_invalid_max_iter_raises(graph): + with pytest.raises(TypeError, match="max_iter"): + distance_from_node_set(graph, "a", max_iter=1.5) + with pytest.raises(ValueError, match="max_iter"): + distance_from_node_set(graph, "a", max_iter=-1) + + +def test_distance_from_node_set_numpy_integer_max_iter(graph): + distance_from_node_set(graph, "a", max_iter=np.int64(2)) + assert _distances(graph)["c"] == 2 + assert _distances(graph)["d"] is None + + +def test_distance_from_node_set_verbose_logs(graph, caplog): + with caplog.at_level("INFO"): + distance_from_node_set(graph, "a", max_iter=2, verbose=True) + + assert any("Iteration" in rec.message for rec in caplog.records) + + +def test_distance_from_node_set_exported_from_analysis(): + from pixelator.pna.analysis import distance_from_node_set as exported + + assert exported is distance_from_node_set + assert ( + inspect.signature(distance_from_node_set).parameters["max_iter"].default == 40 + )