feat: add partition_counts for conjugate node protein sums

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.partition_counts` to sum node protein counts by partition group (cell1 / cell2 / interface / other) on a `PNAGraph`.
 
 ### Changed
 - `density_scatter_plot` now lives in `pixelator.plot` (previously `pixelator.mpx.plot`).

docs/api/overview.rst

@@ -38,3 +38,4 @@ for usage examples.
 
 * :func:`pixelator.pna.analysis.calculate_differential_proximity`
 * :func:`pixelator.pna.analysis.summarize_proximity_scores`
+* :func:`pixelator.pna.analysis.partition_counts`

src/pixelator/pna/analysis/__init__.py

@@ -7,8 +7,13 @@
     calculate_differential_proximity,
     summarize_proximity_scores,
 )
+from pixelator.pna.analysis.segmentation import partition_counts
 
-__all__ = ["calculate_differential_proximity", "summarize_proximity_scores"]
+__all__ = [
+    "calculate_differential_proximity",
+    "partition_counts",
+    "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.partition import partition_counts
+
+__all__ = ["partition_counts"]

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

tests/pna/analysis/test_partition_counts.py

@@ -0,0 +1,156 @@
+"""Tests for `pixelator.pna.analysis.segmentation.partition_counts`.
+
+Copyright © 2026 Pixelgen Technologies AB.
+"""
+
+from __future__ import annotations
+
+import networkx as nx
+import pandas as pd
+import pytest
+from pandas.testing import assert_frame_equal
+
+from pixelator.pna.analysis.segmentation import partition_counts
+from pixelator.pna.graph import PNAGraph
+
+CELL1 = "cell1"
+CELL2 = "cell2"
+INTERFACE = "interface"
+
+
+def _tiny_graph() -> PNAGraph:
+    """Two small communities plus two interface nodes, exclusive markers."""
+    edges = pd.DataFrame(
+        {
+            "umi1": ["a1", "a2", "a3"],
+            "umi2": ["b1", "b2", "b3"],
+            "read_count": [1, 1, 1],
+            "marker_1": ["CD3e", "CD20", "HLA-ABC"],
+            "marker_2": ["CD4", "CD19", "HLA-ABC"],
+        }
+    )
+    return PNAGraph.from_edgelist(edges)
+
+
+@pytest.fixture
+def graph() -> PNAGraph:
+    return _tiny_graph()
+
+
+@pytest.fixture
+def node_order(graph: PNAGraph) -> list:
+    return list(graph.raw.nodes())
+
+
+@pytest.fixture
+def labels(node_order: list) -> list[str]:
+    by_node = {
+        "a1": CELL1,
+        "b1": CELL1,
+        "a2": CELL2,
+        "b2": CELL2,
+        "a3": INTERFACE,
+        "b3": INTERFACE,
+    }
+    return [by_node[node] for node in node_order]
+
+
+@pytest.fixture
+def expected_counts() -> pd.DataFrame:
+    return pd.DataFrame(
+        {
+            "CD19": [0, 1, 0],
+            "CD20": [0, 1, 0],
+            "CD3e": [1, 0, 0],
+            "CD4": [1, 0, 0],
+            "HLA-ABC": [0, 0, 2],
+        },
+        index=pd.Index([CELL1, CELL2, INTERFACE], name="partition"),
+    )
+
+
+def _align_expected(result: pd.DataFrame, expected: pd.DataFrame) -> pd.DataFrame:
+    return expected.reindex(index=result.index, columns=result.columns)
+
+
+def test_partition_counts_requires_exactly_one_argument(graph):
+    with pytest.raises(
+        ValueError, match="Either `partition` or `partition_column` must be provided"
+    ):
+        partition_counts(graph)
+
+    with pytest.raises(
+        ValueError,
+        match="One of `partition` or `partition_column` must be provided, not both",
+    ):
+        partition_counts(graph, partition=["a"], partition_column="compartment")
+
+
+def test_partition_counts_aggregates_by_partition_vector(
+    graph, labels, expected_counts
+):
+    result = partition_counts(graph, partition=labels)
+    expected = _align_expected(result, expected_counts)
+    assert_frame_equal(result, expected, check_dtype=False)
+    assert list(result.index) == [CELL1, CELL2, INTERFACE]
+
+
+def test_partition_counts_aggregates_by_partition_column(
+    graph, labels, expected_counts
+):
+    nx.set_node_attributes(
+        graph.raw, dict(zip(graph.raw.nodes(), labels)), "compartment"
+    )
+    result = partition_counts(graph, partition_column="compartment")
+    expected = _align_expected(result, expected_counts)
+    assert_frame_equal(result, expected, check_dtype=False)
+
+
+def test_partition_counts_length_mismatch_raises(graph, labels):
+    with pytest.raises(
+        ValueError, match="Length of `partition` must match the number of nodes"
+    ):
+        partition_counts(graph, partition=labels[:2])
+
+
+def test_partition_counts_missing_column_raises(graph):
+    with pytest.raises(
+        ValueError, match="Column 'compartment' not found in cell graph node attributes"
+    ):
+        partition_counts(graph, partition_column="compartment")
+
+
+def test_partition_counts_wrong_graph_type_raises():
+    with pytest.raises(TypeError, match="graph must be a PNAGraph"):
+        partition_counts("not a graph", partition=["a"])
+
+
+def test_partition_counts_exported_from_analysis():
+    from pixelator.pna.analysis import partition_counts as exported
+
+    assert exported is partition_counts
+
+
+def test_partition_counts_keeps_categorical_na_group(graph, node_order):
+    labels = pd.Series(
+        pd.Categorical(
+            [
+                CELL1
+                if node in {"a1", "b1"}
+                else CELL2
+                if node in {"a2", "b2"}
+                else pd.NA
+                for node in node_order
+            ],
+            categories=[CELL1, CELL2, INTERFACE],
+        ),
+        index=node_order,
+    )
+    result = partition_counts(graph, partition=labels)
+
+    assert list(result.index[:3]) == [CELL1, CELL2, INTERFACE]
+    assert result.index.isna().any()
+    assert result.loc[CELL1, "CD3e"] == 1
+    assert result.loc[CELL2, "CD20"] == 1
+    assert (result.loc[INTERFACE] == 0).all()
+    assert result.loc[result.index.isna(), "HLA-ABC"].item() == 2