feat: add PNA cell:cell conjugate segmentation

CHANGELOG.md

@@ -10,6 +10,10 @@ 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.cc_protein_weights` to derive rank-2 NMF protein weights for two labeled populations (`mode="cell_abundance"`), for use as `w` in cell:cell conjugate segmentation.
+- `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).
+- `pixelator.pna.analysis.partition_counts` to sum node protein counts by partition group (cell1 / cell2 / interface / other) on a `PNAGraph`.
+- `pixelator.pna.analysis.segment_cell` to classify conjugate-graph nodes into two cell types (plus optional interface / `other`) using NMF weights from `cc_protein_weights`.
 - `pixelator.pna.plot.proximity_heatmap` to plot a clustered heatmap or dot plot of a summary proximity statistic between marker pairs.
 - `pixelator.pna.analysis.filter_proximity_scores` to filter a proximity score table by marker abundance (Python analog of pixelatorR `FilterProximityScores`).
 

docs/api/overview.rst

@@ -38,5 +38,9 @@ for usage examples.
 **Analysis**
 
 * :func:`pixelator.pna.analysis.calculate_differential_proximity`
+* :func:`pixelator.pna.analysis.cc_protein_weights`
+* :func:`pixelator.pna.analysis.distance_from_node_set`
 * :func:`pixelator.pna.analysis.filter_proximity_scores`
+* :func:`pixelator.pna.analysis.partition_counts`
+* :func:`pixelator.pna.analysis.segment_cell`
 * :func:`pixelator.pna.analysis.summarize_proximity_scores`

src/pixelator/pna/analysis/__init__.py

@@ -8,9 +8,19 @@
     filter_proximity_scores,
     summarize_proximity_scores,
 )
+from pixelator.pna.analysis.segmentation import (
+    cc_protein_weights,
+    distance_from_node_set,
+    partition_counts,
+    segment_cell,
+)
 
 __all__ = [
     "calculate_differential_proximity",
+    "cc_protein_weights",
+    "distance_from_node_set",
+    "partition_counts",
+    "segment_cell",
     "filter_proximity_scores",
     "summarize_proximity_scores",
 ]

src/pixelator/pna/analysis/segmentation/__init__.py

@@ -0,0 +1,16 @@
+"""Helpers for segmenting cell:cell conjugates.
+
+Copyright © 2026 Pixelgen Technologies AB.
+"""
+
+from pixelator.pna.analysis.segmentation.distance import distance_from_node_set
+from pixelator.pna.analysis.segmentation.partition import partition_counts
+from pixelator.pna.analysis.segmentation.protein_weights import cc_protein_weights
+from pixelator.pna.analysis.segmentation.segment import segment_cell
+
+__all__ = [
+    "cc_protein_weights",
+    "distance_from_node_set",
+    "partition_counts",
+    "segment_cell",
+]

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

src/pixelator/pna/analysis/segmentation/partition.py

@@ -0,0 +1,159 @@
+"""Protein counts aggregated by node partition.
+
+Copyright © 2026 Pixelgen Technologies AB.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from typing import Any
+
+import networkx as nx
+import pandas as pd
+
+from pixelator.pna.graph import PNAGraph
+
+
+def partition_counts(
+    graph: PNAGraph,
+    partition: Sequence[Any] | pd.Series | None = None,
+    partition_column: str | None = None,
+) -> pd.DataFrame:
+    """Sum node protein counts by partition group.
+
+    Each node belongs to one group (for example ``cell1``, ``cell2``,
+    ``interface``, or ``other`` after conjugate segmentation). This returns
+    the protein count matrix collapsed to those groups: one row per partition
+    and one column per protein.
+
+    Provide exactly one of ``partition`` or ``partition_column``. A positional
+    ``partition`` vector is aligned to graph node order. A :class:`~pandas.Series`
+    whose index matches the node names is aligned by name.     A pandas
+    :class:`~pandas.Categorical` keeps its category order, including unused
+    levels as all-zero rows. Missing labels (``NA``) are kept as their own
+    row so those nodes still contribute to the totals.
+
+    Args:
+        graph: A :class:`~pixelator.pna.graph.PNAGraph` with node marker
+            counts, typically a single component from
+            ``dataset.edgelist().iterator()``.
+        partition: Labels for every node. Either this or
+            ``partition_column`` must be provided.
+        partition_column: Name of a node attribute that holds the partition
+            labels.
+
+    Returns:
+        A DataFrame with one row per partition group and one column per
+        protein. Values are the summed node marker counts in that group.
+
+    Raises:
+        TypeError: If ``graph`` is not a :class:`~pixelator.pna.graph.PNAGraph`.
+        ValueError: If neither or both of ``partition`` and
+            ``partition_column`` are given, if ``partition`` has the wrong
+            length, or if ``partition_column`` is missing from the graph.
+
+    Examples:
+        Sum markers by a vector of node labels, or by a node attribute::
+
+            from pixelator.pna.analysis import partition_counts
+
+            counts = partition_counts(graph, partition=labels)
+            counts = partition_counts(graph, partition_column="compartment")
+
+    See Also:
+        ``partition_counts`` in pixelatorR, the equivalent function for
+        R users.
+
+    """
+    if not isinstance(graph, PNAGraph):
+        raise TypeError("graph must be a PNAGraph.")
+    if partition is None and partition_column is None:
+        raise ValueError("Either `partition` or `partition_column` must be provided.")
+    if partition is not None and partition_column is not None:
+        raise ValueError(
+            "One of `partition` or `partition_column` must be provided, not both."
+        )
+
+    node_order = list(graph.raw.nodes())
+    counts = graph.node_marker_counts.reindex(node_order)
+    if partition_column is not None:
+        labels = _labels_from_column(graph, partition_column, node_order)
+    else:
+        labels = _align_partition(partition, node_order)
+
+    grouped = counts.groupby(labels, sort=False, observed=False, dropna=False).sum()
+    return _reindex_groups(grouped, _group_levels(labels))
+
+
+def _labels_from_column(
+    graph: PNAGraph, partition_column: str, node_order: list[Any]
+) -> pd.Series:
+    if partition_column not in graph.vs.attributes():
+        raise ValueError(
+            f"Column '{partition_column}' not found in cell graph node attributes."
+        )
+    attrs = nx.get_node_attributes(graph.raw, partition_column)
+    missing = [node for node in node_order if node not in attrs]
+    if missing:
+        raise ValueError(f"Column '{partition_column}' is missing on some graph nodes.")
+    return pd.Series([attrs[node] for node in node_order], index=node_order)
+
+
+def _align_partition(
+    partition: Sequence[Any] | pd.Series, node_order: list[Any]
+) -> pd.Series:
+    n_nodes = len(node_order)
+    if isinstance(partition, pd.Series) and set(partition.index) == set(node_order):
+        return partition.reindex(node_order)
+
+    if isinstance(partition, pd.Series):
+        values = partition.tolist()
+        categories = (
+            partition.cat.categories
+            if isinstance(partition.dtype, pd.CategoricalDtype)
+            else None
+        )
+    elif isinstance(partition, pd.Categorical):
+        values = partition.tolist()
+        categories = partition.categories
+    else:
+        values = list(partition)
+        categories = None
+
+    if len(values) != n_nodes:
+        raise ValueError(
+            "Length of `partition` must match the number of nodes in the cell graph."
+        )
+    if categories is not None:
+        return pd.Series(
+            pd.Categorical(values, categories=categories), index=node_order
+        )
+    return pd.Series(values, index=node_order)
+
+
+def _group_levels(labels: pd.Series) -> pd.Index:
+    if isinstance(labels.dtype, pd.CategoricalDtype):
+        levels = pd.Index(labels.cat.categories)
+        if labels.isna().any():
+            return levels.append(pd.Index([pd.NA]))
+        return levels
+    return pd.Index(pd.unique(labels))
+
+
+def _reindex_groups(grouped: pd.DataFrame, levels: pd.Index) -> pd.DataFrame:
+    """Order group sums by ``levels``, keeping an NA row when it is a level."""
+    non_na_levels = levels[~pd.isna(levels)]
+    result = grouped.loc[~grouped.index.isna()].reindex(non_na_levels, fill_value=0)
+    if pd.isna(levels).any():
+        na_vals = grouped.loc[grouped.index.isna()]
+        if na_vals.empty:
+            na_row = pd.DataFrame(0, index=pd.Index([pd.NA]), columns=grouped.columns)
+        else:
+            na_row = pd.DataFrame(
+                na_vals.to_numpy(),
+                index=pd.Index([pd.NA]),
+                columns=grouped.columns,
+            )
+        result = pd.concat([result, na_row])
+    result.index.name = "partition"
+    return result