feat: add distance_from_node_set

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`).

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`

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

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"]

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

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
+    )