diff --git a/CHANGELOG.md b/CHANGELOG.md index 7a41510dd..64150b0d5 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.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`). diff --git a/docs/api/overview.rst b/docs/api/overview.rst index 25cad9d60..6ff12dca3 100644 --- a/docs/api/overview.rst +++ b/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` diff --git a/src/pixelator/pna/analysis/__init__.py b/src/pixelator/pna/analysis/__init__.py index 3689e6533..8db80210d 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 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 diff --git a/src/pixelator/pna/analysis/segmentation/__init__.py b/src/pixelator/pna/analysis/segmentation/__init__.py new file mode 100644 index 000000000..1700315ec --- /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.partition import partition_counts + +__all__ = ["partition_counts"] diff --git a/src/pixelator/pna/analysis/segmentation/partition.py b/src/pixelator/pna/analysis/segmentation/partition.py new file mode 100644 index 000000000..08fd3cb6d --- /dev/null +++ b/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 diff --git a/tests/pna/analysis/test_partition_counts.py b/tests/pna/analysis/test_partition_counts.py new file mode 100644 index 000000000..c062168c2 --- /dev/null +++ b/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