From 7fdc045f41a3c79df299636bef8d25379d3c00e2 Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 18:40:19 -0400 Subject: [PATCH 1/7] Add pycea.tl.parsimony and pycea.tl.fitch_count Migrate small-parsimony scoring and the FitchCount transition-count algorithm from Cassiopeia into first-class pycea.tl functions operating on TreeData. Reuse pycea's Fitch-Hartigan machinery by extracting a shared _fitch_hartigan_downpass helper in ancestral_states.py, used by both the existing reconstruction and fitch_count's state-set computation. - parsimony: reconstructs states (via ancestral_states) or uses existing ones; returns an int for a single tree, a Series across trees. - fitch_count: returns an MxM count matrix, summed across trees. Wire both into the API, docs, changelog, and add the Quinn 2021 citation. Includes tests in tests/test_parsimony.py. Co-Authored-By: Claude Opus 4.8 (1M context) --- CHANGELOG.md | 1 + docs/api.md | 2 + docs/references.bib | 15 ++ src/pycea/tl/__init__.py | 1 + src/pycea/tl/ancestral_states.py | 42 +++- src/pycea/tl/parsimony.py | 382 +++++++++++++++++++++++++++++++ tests/test_parsimony.py | 101 ++++++++ 7 files changed, 533 insertions(+), 11 deletions(-) create mode 100755 src/pycea/tl/parsimony.py create mode 100644 tests/test_parsimony.py diff --git a/CHANGELOG.md b/CHANGELOG.md index c6e1bd0..43320d6 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,6 +11,7 @@ and this project adheres to [Semantic Versioning][]. ## Unreleased ### Added +- `pycea.tl.parsimony` and `pycea.tl.fitch_count` for small-parsimony scoring and the FitchCount transition-count algorithm (migrated from Cassiopeia) - `pycea.pl.ancestral_linkage` for plotting the pairwise linkage matrix as a clustered heatmap (#58) - Added fast `"sum"` method to `pycea.tl.ancestral_states` (#54, #56) - Added additional node ploting customization - `outline_width` and option to directly specify color with hex code (#53) diff --git a/docs/api.md b/docs/api.md index edc92f9..17a9119 100644 --- a/docs/api.md +++ b/docs/api.md @@ -34,6 +34,8 @@ tl.fitness tl.partition_test tl.expansion_test + tl.parsimony + tl.fitch_count ``` ## Plotting diff --git a/docs/references.bib b/docs/references.bib index 2183930..8e854f2 100644 --- a/docs/references.bib +++ b/docs/references.bib @@ -102,3 +102,18 @@ @article{Griffiths_1998 keywords = {Age of a mutation, Coalescent process, Population Genetics, Samples of DNA}, pages = {273--295}, } + +@article{Quinn_2021, + title = {Single-cell lineages reveal the rates, routes, and drivers of metastasis in cancer xenografts}, + doi = {10.1126/science.abc1944}, + abstract = {Detailed phylogenies of tumor populations can recount the history and chronology of critical events during cancer progression, such as metastatic dissemination. We applied a Cas9-based, single-cell lineage tracer to study the rates, routes, and drivers of metastasis in a lung cancer xenograft mouse model. We report deeply resolved phylogenies for tens of thousands of cancer cells traced over months of growth and dissemination. This revealed stark heterogeneity in metastatic capacity, arising from preexisting and heritable differences in gene expression. We demonstrate that these identified genes can drive invasiveness and uncovered an unanticipated suppressive role for KRT17. We also show that metastases disseminated via multidirectional tissue routes and complex seeding topologies. Overall, we demonstrate the power of tracing cancer progression at subclonal resolution and vast scale.}, + language = {en}, + journal = {Science}, + volume = {371}, + number = {6532}, + pages = {eabc1944}, + year = {2021}, + month = feb, + author = {Quinn, Jeffrey J. and Jones, Matthew G. and Okimoto, Ross A. and Nanjo, Shigeki and Chan, Michelle M. and Yosef, Nir and Bivona, Trever G. and Weissman, Jonathan S.}, + +} diff --git a/src/pycea/tl/__init__.py b/src/pycea/tl/__init__.py index 3a63b42..0130fa6 100644 --- a/src/pycea/tl/__init__.py +++ b/src/pycea/tl/__init__.py @@ -6,6 +6,7 @@ from .fitness import fitness from .n_extant import n_extant from .neighbor_distance import neighbor_distance +from .parsimony import fitch_count, parsimony from .partition_test import partition_test from .sort import sort from .topology import expansion_test diff --git a/src/pycea/tl/ancestral_states.py b/src/pycea/tl/ancestral_states.py index 3826f61..d0d6278 100755 --- a/src/pycea/tl/ancestral_states.py +++ b/src/pycea/tl/ancestral_states.py @@ -44,39 +44,59 @@ def _remove_node_attributes(tree: nx.DiGraph, key: str) -> None: del tree.nodes[node][key] -def _reconstruct_fitch_hartigan( - tree: nx.DiGraph, key: str, missing: str | None = None, index: int | None = None, fixed_nodes: set | None = None +def _fitch_hartigan_downpass( + tree: nx.DiGraph, + key: str, + missing: str | None = None, + index: int | None = None, + fixed_nodes: set | None = None, + set_key: str = "value_set", ) -> None: - """Reconstructs ancestral states using the Fitch-Hartigan algorithm.""" + """Populate ``tree.nodes[node][set_key]`` with Fitch-Hartigan state sets (down-pass). + + Each node is annotated with the set of states it may take on in an + equally-parsimonious solution. Nodes for which all descendants are missing are + annotated with the ``missing`` sentinel instead of a set. + """ def _is_fixed(node, value): return fixed_nodes is not None and node in fixed_nodes and value is not None and value != missing - # Recursive function to calculate the downpass def downpass(node): value = _get_node_value(tree, node, key, index) # Base case: leaf or fixed observed internal node if tree.out_degree(node) == 0 or _is_fixed(node, value): if value == missing: - tree.nodes[node]["value_set"] = missing + tree.nodes[node][set_key] = missing else: - tree.nodes[node]["value_set"] = {value} + tree.nodes[node][set_key] = {value} # Recursive case: internal node else: value_sets = [] for child in tree.successors(node): downpass(child) - value_set = tree.nodes[child]["value_set"] + value_set = tree.nodes[child][set_key] if value_set != missing: value_sets.append(value_set) if len(value_sets) > 0: intersection = set.intersection(*value_sets) if intersection: - tree.nodes[node]["value_set"] = intersection + tree.nodes[node][set_key] = intersection else: - tree.nodes[node]["value_set"] = set.union(*value_sets) + tree.nodes[node][set_key] = set.union(*value_sets) else: - tree.nodes[node]["value_set"] = missing + tree.nodes[node][set_key] = missing + + downpass(get_root(tree)) + + +def _reconstruct_fitch_hartigan( + tree: nx.DiGraph, key: str, missing: str | None = None, index: int | None = None, fixed_nodes: set | None = None +) -> None: + """Reconstructs ancestral states using the Fitch-Hartigan algorithm.""" + + def _is_fixed(node, value): + return fixed_nodes is not None and node in fixed_nodes and value is not None and value != missing # Recursive function to calculate the uppass def uppass(node, parent_state=None): @@ -100,7 +120,7 @@ def uppass(node, parent_state=None): # Run the algorithm root = get_root(tree) - downpass(root) + _fitch_hartigan_downpass(tree, key, missing, index, fixed_nodes) uppass(root) # Clean up for node in tree.nodes: diff --git a/src/pycea/tl/parsimony.py b/src/pycea/tl/parsimony.py new file mode 100755 index 0000000..b82f255 --- /dev/null +++ b/src/pycea/tl/parsimony.py @@ -0,0 +1,382 @@ +from __future__ import annotations + +import itertools +from collections.abc import Callable, Sequence +from typing import Any, Literal, overload + +import networkx as nx +import numpy as np +import pandas as pd +import treedata as td + +from pycea.utils import _check_tree_overlap, check_tree_has_key, get_keyed_obs_data, get_root, get_trees + +from .ancestral_states import _fitch_hartigan_downpass, ancestral_states + + +@overload +def parsimony( + tdata: td.TreeData, + key: str, + method: str | Callable = "fitch_hartigan", + missing_state: str | None = None, + costs: pd.DataFrame | None = None, + reconstruct: bool = True, + tree: str | Sequence[str] | None = None, + key_added: str = "parsimony", + copy: Literal[True] = True, +) -> int | pd.Series: ... +@overload +def parsimony( + tdata: td.TreeData, + key: str, + method: str | Callable = "fitch_hartigan", + missing_state: str | None = None, + costs: pd.DataFrame | None = None, + reconstruct: bool = True, + tree: str | Sequence[str] | None = None, + key_added: str = "parsimony", + copy: Literal[False] = False, +) -> None: ... +def parsimony( + tdata: td.TreeData, + key: str, + method: str | Callable = "fitch_hartigan", + missing_state: str | None = None, + costs: pd.DataFrame | None = None, + reconstruct: bool = True, + tree: str | Sequence[str] | None = None, + key_added: str = "parsimony", + copy: Literal[True, False] = True, +) -> int | pd.Series | None: + """Computes the small-parsimony score of a tree. + + The parsimony score is the number of edges along which the state of a + categorical attribute changes. Ancestral states can be reconstructed on the + fly (default) or read from states already present on the tree. + + Parameters + ---------- + tdata + TreeData object. + key + An `obs.keys()` corresponding to a categorical variable. + method + Method used to reconstruct ancestral states when ``reconstruct=True``. + Passed to :func:`pycea.tl.ancestral_states`. Typically ``'fitch_hartigan'`` + or ``'sankoff'``. + missing_state + The state to consider as missing data. Edges touching a missing state are + not counted. + costs + A :class:`DataFrame ` with the costs of changing states + (from rows to columns). Only used if ``method='sankoff'``. + reconstruct + If True, reconstruct ancestral states with :func:`pycea.tl.ancestral_states` + before scoring. If False, use states already stored under ``key`` on the tree. + tree + The `obst` key or keys of the trees to use. If `None`, all trees are used. + key_added + Key in ``tdata.uns`` where the parsimony score(s) are stored. + copy + If True, returns the parsimony score(s). + + Returns + ------- + Returns `None` if `copy=False`. Otherwise returns an ``int`` for a single tree, + or a :class:`Series ` of scores indexed by tree for multiple trees. + + Sets the following fields: + + * `tdata.uns[key_added]` : `int` | :class:`Series ` + - Parsimony score for a single tree, or a Series of scores for multiple trees. + + Examples + -------- + Compute the parsimony score of a categorical character: + + >>> py.tl.parsimony(tdata, key="clone") + """ + tree_keys = tree + _check_tree_overlap(tdata, tree_keys) + trees = get_trees(tdata, tree_keys) + if reconstruct: + ancestral_states(tdata, keys=key, method=method, missing_state=missing_state, costs=costs, tree=tree_keys) + scores = {} + for name, t in trees.items(): + if not reconstruct: + check_tree_has_key(t, key) + score = 0 + for parent, child in nx.dfs_edges(t, source=get_root(t)): + parent_state = t.nodes[parent].get(key) + child_state = t.nodes[child].get(key) + if parent_state is None or child_state is None: + continue + if missing_state is not None and (parent_state == missing_state or child_state == missing_state): + continue + if parent_state != child_state: + score += 1 + scores[name] = score + result: int | pd.Series + result = next(iter(scores.values())) if len(scores) == 1 else pd.Series(scores, name=key_added) + tdata.uns[key_added] = result + if copy: + return result + + +@overload +def fitch_count( + tdata: td.TreeData, + key: str, + states: Sequence[str] | None = None, + missing_state: str | None = None, + root: str | None = None, + tree: str | Sequence[str] | None = None, + key_added: str = "fitch_count", + copy: Literal[True] = True, +) -> pd.DataFrame: ... +@overload +def fitch_count( + tdata: td.TreeData, + key: str, + states: Sequence[str] | None = None, + missing_state: str | None = None, + root: str | None = None, + tree: str | Sequence[str] | None = None, + key_added: str = "fitch_count", + copy: Literal[False] = False, +) -> None: ... +def fitch_count( + tdata: td.TreeData, + key: str, + states: Sequence[str] | None = None, + missing_state: str | None = None, + root: str | None = None, + tree: str | Sequence[str] | None = None, + key_added: str = "fitch_count", + copy: Literal[True, False] = True, +) -> pd.DataFrame | None: + """Runs the FitchCount algorithm. + + Performs the FitchCount algorithm for inferring the number of times that two + states transition to one another across all equally-parsimonious solutions + returned by the Fitch-Hartigan algorithm. The original algorithm was described + in :cite:`Quinn_2021`. The output is an MxM count matrix, where the values + indicate the number of times that ``m1`` transitioned to ``m2`` along an edge in + a Fitch-Hartigan solution. To obtain probabilities ``P(m1 -> m2)``, divide each + row by its row-sum. + + This procedure will only work on categorical data. + + Parameters + ---------- + tdata + TreeData object. + key + An `obs.keys()` corresponding to a categorical variable. + states + State space that can be optionally provided by the user. If not provided, + the unique non-missing values in ``tdata.obs[key]`` are used. + missing_state + The state to consider as missing data. Missing leaves are treated as + wildcards (may take on any state in ``states``). + root + Node to treat as the root. Only the subtree below this node is considered. + Only valid when a single tree is used. + tree + The `obst` key or keys of the trees to use. If `None`, all trees are used. + key_added + Key in ``tdata.uns`` where the count matrix is stored. + copy + If True, returns the count matrix. + + Returns + ------- + Returns `None` if `copy=False`. Otherwise returns an MxM + :class:`DataFrame `. When multiple trees are used, the + per-tree count matrices are summed into a single matrix. + + Sets the following fields: + + * `tdata.uns[key_added]` : :class:`DataFrame ` + - Count matrix, summed across all trees used. + + Examples + -------- + Count state transitions across equally-parsimonious solutions: + + >>> py.tl.fitch_count(tdata, key="tissue", copy=True) + """ + tree_keys = tree + _check_tree_overlap(tdata, tree_keys) + trees = get_trees(tdata, tree_keys) + if root is not None and len(trees) > 1: + raise ValueError("`root` can only be specified when a single tree is used.") + data, _, _ = get_keyed_obs_data(tdata, [key]) + values = data[key] + observed = pd.unique(values[values != missing_state].dropna()) + if states is None: + states = list(observed) + else: + states = list(states) + if len(np.setdiff1d(np.asarray(observed), np.asarray(states))) > 0: + raise ValueError("Specified `states` do not span the set of states that appear in the data.") + + matrices = [] + for _name, t in trees.items(): + g = t + if root is not None: + g = nx.DiGraph(g.subgraph(nx.descendants(g, root) | {root}).copy()) + actual_root = root if root is not None else get_root(g) + # Set leaf states then compute Fitch-Hartigan state sets + leaf_states = {node: values[node] for node in g.nodes if g.out_degree(node) == 0 and node in values.index} + nx.set_node_attributes(g, leaf_states, key) + _fitch_hartigan_downpass(g, key, missing_state, set_key="_fitch_set") + # Treat missing (wildcard) sets as the full state space + for node in g.nodes: + if g.nodes[node]["_fitch_set"] == missing_state: + g.nodes[node]["_fitch_set"] = set(states) + + bfs_nodes = [actual_root] + [v for _, v in nx.bfs_edges(g, actual_root)] + node_to_i = dict(zip(bfs_nodes, range(len(bfs_nodes)), strict=False)) + label_to_j = dict(zip(states, range(len(states)), strict=False)) + + N = _N_fitch_count(g, states, node_to_i, label_to_j, "_fitch_set") + C = _C_fitch_count(g, N, states, node_to_i, label_to_j, "_fitch_set") + + M = pd.DataFrame(np.zeros((len(states), len(states))), index=states, columns=states) + for s1 in states: + for s2 in states: + M.loc[s1, s2] = np.sum(C[node_to_i[actual_root], :, label_to_j[s1], label_to_j[s2]]) + matrices.append(M) + + result = sum(matrices) + tdata.uns[key_added] = result + if copy: + return result # type: ignore + + +def _N_fitch_count( + g: nx.DiGraph, + unique_states: Sequence[str], + node_to_i: dict[Any, int], + label_to_j: dict[str, int], + state_key: str = "S1", +) -> np.ndarray: + """Fill in the dynamic programming table N for FitchCount. + + Computes N[v, s], corresponding to the number of solutions below + a node v in the tree given v takes on the state s. + + Parameters + ---------- + g + Directed graph with state_key attribute set on all nodes. + unique_states + The state space that a node can take on. + node_to_i + Mapping of each node to a unique integer. + label_to_j + Mapping of each unique state to a unique integer. + state_key + Node attribute storing the possible states for each node. + + Returns + ------- + A 2-dimensional array storing N[v, s]. + """ + + def _fill(v: str, s: str) -> float: + if g.out_degree(v) == 0: + return 1 + children = list(g.successors(v)) + A = np.zeros(len(children)) + for i, u in enumerate(children): + if s not in g.nodes[u][state_key]: + legal_states = g.nodes[u][state_key] + else: + legal_states = [s] + A[i] = np.sum([N[node_to_i[u], label_to_j[sp]] for sp in legal_states]) + return float(np.prod(A)) + + N = np.full((len(g.nodes), len(unique_states)), 0.0) + root = next(n for n in g.nodes if g.in_degree(n) == 0) + for n in nx.dfs_postorder_nodes(g, source=root): + for s in g.nodes[n][state_key]: + N[node_to_i[n], label_to_j[s]] = _fill(n, s) + + return N + + +def _C_fitch_count( + g: nx.DiGraph, + N: np.ndarray, + unique_states: Sequence[str], + node_to_i: dict[Any, int], + label_to_j: dict[str, int], + state_key: str = "S1", +) -> np.ndarray: + """Fill in the dynamic programming table C for FitchCount. + + Computes C[v, s, s1, s2], the number of transitions from state s1 to + state s2 in the subtree rooted at v, given that state v takes on the + state s. + + Parameters + ---------- + g + Directed graph with state_key attribute set on all nodes. + N + N array computed during FitchCount. + unique_states + The state space that a node can take on. + node_to_i + Mapping of each node to a unique integer. + label_to_j + Mapping of each unique state to a unique integer. + state_key + Node attribute storing the possible states for each node. + + Returns + ------- + A 4-dimensional array storing C[v, s, s1, s2]. + """ + + def _fill(v: str, s: str, s1: str, s2: str) -> float: + if g.out_degree(v) == 0: + return 0 + + children = list(g.successors(v)) + A = np.zeros(len(children)) + LS = [[]] * len(children) + + for i, u in enumerate(children): + if s in g.nodes[u][state_key]: + LS[i] = [s] + else: + LS[i] = g.nodes[u][state_key] + + A[i] = np.sum([C[node_to_i[u], label_to_j[sp], label_to_j[s1], label_to_j[s2]] for sp in LS[i]]) + + if s1 == s and s2 in LS[i]: + A[i] += N[node_to_i[u], label_to_j[s2]] + + parts = [] + for i, u in enumerate(children): + prod = 1 + for k, up in enumerate(children): + if up == u: + continue + prod *= sum(N[node_to_i[up], label_to_j[sp]] for sp in LS[k]) + parts.append(A[i] * prod) + + return np.sum(parts) + + C = np.zeros((len(g.nodes), N.shape[1], N.shape[1], N.shape[1])) + root = next(n for n in g.nodes if g.in_degree(n) == 0) + for n in nx.dfs_postorder_nodes(g, source=root): + for s in g.nodes[n][state_key]: + for s1, s2 in itertools.product(unique_states, repeat=2): + C[node_to_i[n], label_to_j[s], label_to_j[s1], label_to_j[s2]] = _fill(n, s, s1, s2) + + return C diff --git a/tests/test_parsimony.py b/tests/test_parsimony.py new file mode 100644 index 0000000..5866f3e --- /dev/null +++ b/tests/test_parsimony.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import networkx as nx +import pandas as pd +import pytest +import treedata as td + +from pycea.tl.ancestral_states import ancestral_states +from pycea.tl.parsimony import fitch_count, parsimony + + +@pytest.fixture +def tdata(): + # tree1: root -> A, B; B -> C, D (leaves A, C, D) + # tree2: root -> E, F (leaves E, F) + tree1 = nx.DiGraph([("root", "A"), ("root", "B"), ("B", "C"), ("B", "D")]) + tree2 = nx.DiGraph([("root", "E"), ("root", "F")]) + tdata = td.TreeData( + obs=pd.DataFrame( + {"clone": ["x", "x", "y", "x", "y"]}, + index=["A", "C", "D", "E", "F"], + ), + obst={"tree1": tree1, "tree2": tree2}, + ) + yield tdata + + +def test_parsimony_reconstruct(tdata): + score = parsimony(tdata, "clone", tree="tree1", copy=True) + assert score == 1 + assert tdata.uns["parsimony"] == 1 + # states were reconstructed on the tree + assert tdata.obst["tree1"].nodes["root"]["clone"] == "x" + + +def test_parsimony_multiple_trees(tdata): + scores = parsimony(tdata, "clone", copy=True) + assert scores.loc["tree1"] == 1 + assert scores.loc["tree2"] == 1 + + +def test_parsimony_no_reconstruct(tdata): + ancestral_states(tdata, "clone", method="fitch_hartigan", tree="tree1") + score = parsimony(tdata, "clone", reconstruct=False, tree="tree1", copy=True) + assert score == 1 + + +def test_parsimony_no_reconstruct_missing_key(tdata): + with pytest.raises(ValueError): + parsimony(tdata, "clone", reconstruct=False, tree="tree1") + + +def test_fitch_count_single_tree(tdata): + M = fitch_count(tdata, "clone", tree="tree1", copy=True) + assert isinstance(M, pd.DataFrame) + assert M.loc["x", "x"] == 3 + assert M.loc["x", "y"] == 1 + assert M.loc["y", "x"] == 0 + assert M.loc["y", "y"] == 0 + # stored in uns + assert isinstance(tdata.uns["fitch_count"], pd.DataFrame) + + +def test_fitch_count_two_mp_solutions(tdata): + M = fitch_count(tdata, "clone", tree="tree2", copy=True) + # root has both x and y in its Fitch set -> two equally parsimonious solutions + assert M.loc["x", "x"] == 1 + assert M.loc["x", "y"] == 1 + assert M.loc["y", "x"] == 1 + assert M.loc["y", "y"] == 1 + + +def test_fitch_count_multiple_trees(tdata): + # tree1: [[3,1],[0,0]] + tree2: [[1,1],[1,1]] summed across trees + result = fitch_count(tdata, "clone", copy=True) + assert isinstance(result, pd.DataFrame) + assert result.loc["x", "x"] == 4 + assert result.loc["x", "y"] == 2 + assert result.loc["y", "x"] == 1 + assert result.loc["y", "y"] == 1 + + +def test_fitch_count_explicit_states(tdata): + M = fitch_count(tdata, "clone", states=["x", "y", "z"], tree="tree1", copy=True) + assert list(M.index) == ["x", "y", "z"] + assert M.loc["x", "x"] == 3 + assert M.loc["z"].sum() == 0 + + +def test_fitch_count_invalid_states(tdata): + with pytest.raises(ValueError): + fitch_count(tdata, "clone", states=["x"], tree="tree1") + + +def test_fitch_count_root_requires_single_tree(tdata): + with pytest.raises(ValueError): + fitch_count(tdata, "clone", root="B") + + +if __name__ == "__main__": + pytest.main(["-v", __file__]) From 968af2a0fc2e60b45cf32f5bb566e89d304c072c Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 18:40:30 -0400 Subject: [PATCH 2/7] Suppress expected warnings in plot tests Add module-level filterwarnings markers for warnings the tests intentionally trigger: the "no artists with labels" legend warning in the hex-color tree plots, and the small-category / multi-threaded-fork warnings in the ancestral_linkage plot tests. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_plot_ancestral_linkage.py | 6 ++++++ tests/test_plot_tree.py | 3 +++ 2 files changed, 9 insertions(+) diff --git a/tests/test_plot_ancestral_linkage.py b/tests/test_plot_ancestral_linkage.py index c52d659..f5df48e 100644 --- a/tests/test_plot_ancestral_linkage.py +++ b/tests/test_plot_ancestral_linkage.py @@ -18,6 +18,12 @@ plot_path = Path(__file__).parent / "plots" +# Expected warnings from intentionally small test fixtures / parallel workers. +pytestmark = [ + pytest.mark.filterwarnings("ignore:Categories with fewer than 10 cells"), + pytest.mark.filterwarnings("ignore:This process .* is multi-threaded"), +] + @pytest.fixture def tdata(): diff --git a/tests/test_plot_tree.py b/tests/test_plot_tree.py index 7a1e6e6..efbf89d 100755 --- a/tests/test_plot_tree.py +++ b/tests/test_plot_tree.py @@ -9,6 +9,9 @@ plot_path = Path(__file__).parent / "plots" +# Hex-color tests plot without labels, which matplotlib warns about; expected here. +pytestmark = pytest.mark.filterwarnings("ignore:No artists with labels found to put in legend") + @pytest.fixture def tdata() -> td.TreeData: From 92aaf0b332262fb22689fbc66f790c29e827bfd3 Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 18:40:31 -0400 Subject: [PATCH 3/7] Consolidate ancestral_linkage tests Reduce test_ancestral_linkage from 71 test functions to 22, grouping redundant per-assertion tests into coherent per-feature tests (pairwise values, output structure, min_size, symmetrize, single-target, by_tree, permutation, alternative, permutation_mode, normalize) and parametrizing the brute-force regression checks over the ultrametric and non-ultrametric fixtures. Also add filterwarnings markers for the expected small-category and multi-threaded-fork warnings. Coverage of distinct behaviors is preserved. Co-Authored-By: Claude Opus 4.8 (1M context) --- tests/test_ancestral_linkage.py | 1173 +++++++++---------------------- 1 file changed, 330 insertions(+), 843 deletions(-) diff --git a/tests/test_ancestral_linkage.py b/tests/test_ancestral_linkage.py index 2b9a042..6f17674 100644 --- a/tests/test_ancestral_linkage.py +++ b/tests/test_ancestral_linkage.py @@ -1,5 +1,6 @@ """Tests for tl.ancestral_linkage.""" +import warnings from collections import defaultdict import networkx as nx @@ -10,6 +11,12 @@ import pycea.tl as tl +# Expected warnings from intentionally small test fixtures / parallel workers. +pytestmark = [ + pytest.mark.filterwarnings("ignore:Categories with fewer than 10 cells"), + pytest.mark.filterwarnings("ignore:This process .* is multi-threaded"), +] + # ── fixtures ────────────────────────────────────────────────────────────────── @@ -27,58 +34,23 @@ def balanced_tdata(): ├── b1 (depth=1.0) cat=B └── b2 (depth=1.0) cat=B - LCA depths: - LCA(a1, a2) = 0.5 (within A) - LCA(b1, b2) = 0.5 (within B) - LCA(a*, b*) = 0.0 (between A and B) - - Path distances: - path(a1, a2) = 2*(1.0 - 0.5) = 1.0 - path(b1, b2) = 1.0 - path(a*, b*) = 2*(1.0 - 0.0) = 2.0 + LCA depths: within = 0.5, between = 0.0 + Path dists: within = 1.0, between = 2.0 """ t = nx.DiGraph() - nodes = { - "root": 0.0, - "n1": 0.5, - "n2": 0.5, - "a1": 1.0, - "a2": 1.0, - "b1": 1.0, - "b2": 1.0, - } + nodes = {"root": 0.0, "n1": 0.5, "n2": 0.5, "a1": 1.0, "a2": 1.0, "b1": 1.0, "b2": 1.0} for node, depth in nodes.items(): t.add_node(node, depth=depth) - edges = [ - ("root", "n1"), ("root", "n2"), - ("n1", "a1"), ("n1", "a2"), - ("n2", "b1"), ("n2", "b2"), - ] - t.add_edges_from(edges) - - obs = pd.DataFrame( - {"celltype": ["A", "A", "B", "B"]}, - index=["a1", "a2", "b1", "b2"], + t.add_edges_from( + [("root", "n1"), ("root", "n2"), ("n1", "a1"), ("n1", "a2"), ("n2", "b1"), ("n2", "b2")] ) + obs = pd.DataFrame({"celltype": ["A", "A", "B", "B"]}, index=["a1", "a2", "b1", "b2"]) return td.TreeData(obs=obs, obst={"tree": t}) @pytest.fixture def three_cat_tdata(): - """Six-leaf tree with categories A, B, C for fuller pairwise tests. - - Structure: - root (0) - ├── n1 (0.4) - │ ├── a1 (1.0) A - │ └── a2 (1.0) A - ├── n2 (0.4) - │ ├── b1 (1.0) B - │ └── b2 (1.0) B - └── n3 (0.4) - ├── c1 (1.0) C - └── c2 (1.0) C - """ + """Six-leaf tree with categories A, B, C (n=0.4 for internal nodes).""" t = nx.DiGraph() for node, depth in [ ("root", 0.0), ("n1", 0.4), ("n2", 0.4), ("n3", 0.4), @@ -87,9 +59,7 @@ def three_cat_tdata(): t.add_node(node, depth=depth) for u, v in [ ("root", "n1"), ("root", "n2"), ("root", "n3"), - ("n1", "a1"), ("n1", "a2"), - ("n2", "b1"), ("n2", "b2"), - ("n3", "c1"), ("n3", "c2"), + ("n1", "a1"), ("n1", "a2"), ("n2", "b1"), ("n2", "b2"), ("n3", "c1"), ("n3", "c2"), ]: t.add_edge(u, v) obs = pd.DataFrame( @@ -101,960 +71,477 @@ def three_cat_tdata(): @pytest.fixture def two_tree_tdata(): - """Two separate trees, each with leaves from categories A and B. - - tree1: a1 (A), b1 (B) — root depth 0, leaves depth 1, n1 depth 0.5 - tree2: a2 (A), b2 (B) — root depth 0, leaves depth 1, n2 depth 0.5 - """ + """Two separate trees, each with one A leaf and one B leaf.""" t1 = nx.DiGraph() for node, depth in [("r1", 0.0), ("n1", 0.5), ("a1", 1.0), ("b1", 1.0)]: t1.add_node(node, depth=depth) t1.add_edges_from([("r1", "n1"), ("n1", "a1"), ("n1", "b1")]) - t2 = nx.DiGraph() for node, depth in [("r2", 0.0), ("n2", 0.5), ("a2", 1.0), ("b2", 1.0)]: t2.add_node(node, depth=depth) t2.add_edges_from([("r2", "n2"), ("n2", "a2"), ("n2", "b2")]) - - obs = pd.DataFrame( - {"celltype": ["A", "B", "A", "B"]}, - index=["a1", "b1", "a2", "b2"], - ) + obs = pd.DataFrame({"celltype": ["A", "B", "A", "B"]}, index=["a1", "b1", "a2", "b2"]) return td.TreeData(obs=obs, obst={"tree1": t1, "tree2": t2}) -# ── pairwise mode tests ─────────────────────────────────────────────────────── +@pytest.fixture +def small_category_tdata(): + """Star tree with a large category A (12 cells) and a small category B (3 cells).""" + t = nx.DiGraph() + t.add_node("root", depth=0.0) + labels = {} + for i in range(12): + t.add_node(f"a{i}", depth=1.0) + t.add_edge("root", f"a{i}") + labels[f"a{i}"] = "A" + for i in range(3): + t.add_node(f"b{i}", depth=1.0) + t.add_edge("root", f"b{i}") + labels[f"b{i}"] = "B" + obs = pd.DataFrame({"celltype": list(labels.values())}, index=list(labels)) + return td.TreeData(obs=obs, obst={"tree": t}) -def test_pairwise_path_min_within_greater_than_between(balanced_tdata): - """Within-category path distance (min) should be smaller than between-category.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", normalize=False) - mat = tdata.uns["celltype_linkage"] - assert mat.loc["A", "A"] < mat.loc["A", "B"] - assert mat.loc["B", "B"] < mat.loc["B", "A"] +@pytest.fixture +def non_ultrametric_tdata(): + """Tree with leaves at unequal depths (non-ultrametric), two categories A and B. + root(0.0) + ├── n1(0.3): a1(0.5) A, a2(2.0) A + └── n2(0.6): b1(1.0) B, b2(1.5) B + """ + t = nx.DiGraph() + for node, depth in [ + ("root", 0.0), ("n1", 0.3), ("n2", 0.6), ("a1", 0.5), ("a2", 2.0), ("b1", 1.0), ("b2", 1.5), + ]: + t.add_node(node, depth=depth) + t.add_edges_from( + [("root", "n1"), ("root", "n2"), ("n1", "a1"), ("n1", "a2"), ("n2", "b1"), ("n2", "b2")] + ) + obs = pd.DataFrame({"celltype": ["A", "A", "B", "B"]}, index=["a1", "a2", "b1", "b2"]) + return td.TreeData(obs=obs, obst={"tree": t}) -def test_pairwise_lca_max_within_greater_than_between(balanced_tdata): - """Within-category LCA depth (max) should be larger than between-category.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", normalize=False) - mat = tdata.uns["celltype_linkage"] - assert mat.loc["A", "A"] > mat.loc["A", "B"] - assert mat.loc["B", "B"] > mat.loc["B", "A"] +def _bruteforce_lca_max(tdata, groupby="celltype"): + """Reference pairwise max-LCA-depth linkage via explicit LCA on the tree.""" + t = list(tdata.obst.values())[0] + depth = dict(t.nodes(data="depth")) + cats = defaultdict(list) + for leaf, c in tdata.obs[groupby].items(): + cats[c].append(leaf) -def test_pairwise_lca_max_known_values(balanced_tdata): - """Verify exact values for lca+max aggregate.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", normalize=False) - mat = tdata.uns["celltype_linkage"] - # within: self-pair gives lca = (1.0+1.0-0)/2 = 1.0 (self is always max) - assert np.isclose(mat.loc["A", "A"], 1.0) - assert np.isclose(mat.loc["B", "B"], 1.0) - # between: best LCA from a1/a2 to any b = root.depth = 0.0 - assert np.isclose(mat.loc["A", "B"], 0.0) + def lca_depth(i, j): + anc_i = nx.ancestors(t, i) | {i} + anc_j = nx.ancestors(t, j) | {j} + return max(depth[a] for a in anc_i & anc_j) + out = {} + for s in cats: + row = {} + for tcat in cats: + per_cell = [max(depth[a] if a == b else lca_depth(a, b) for b in cats[tcat]) for a in cats[s]] + row[tcat] = float(np.mean(per_cell)) + out[s] = row + return pd.DataFrame(out).T -def test_pairwise_path_min_known_values(balanced_tdata): - """Verify exact values for path+min aggregate.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", normalize=False) - mat = tdata.uns["celltype_linkage"] - # Within A: min dist is self (0.0) - assert np.isclose(mat.loc["A", "A"], 0.0) - # path(a*, b*) = |1.0 + 1.0 - 2*0.0| = 2.0 - assert np.isclose(mat.loc["A", "B"], 2.0) +def _bruteforce_path_min(tdata, groupby="celltype"): + """Reference pairwise min-path-distance linkage via explicit LCA on the tree.""" + t = list(tdata.obst.values())[0] + depth = dict(t.nodes(data="depth")) + cats = defaultdict(list) + for leaf, c in tdata.obs[groupby].items(): + cats[c].append(leaf) -def test_pairwise_mean_aggregate(balanced_tdata): - """mean aggregate should equal mean of all cross-category path distances.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="mean", metric="path", normalize=False) - mat = tdata.uns["celltype_linkage"] - # All a*-b* pairs have path distance 2.0 → mean = 2.0 - assert np.isclose(mat.loc["A", "B"], 2.0) - # Within A: mean([path(a1,a1), path(a1,a2), path(a2,a1), path(a2,a2)]) = mean([0, 1, 1, 0]) = 0.5 - assert np.isclose(mat.loc["A", "A"], 0.5) + def path(i, j): + if i == j: + return 0.0 + anc_i = nx.ancestors(t, i) | {i} + anc_j = nx.ancestors(t, j) | {j} + lca = max(depth[a] for a in anc_i & anc_j) + return depth[i] + depth[j] - 2 * lca + out = {} + for s in cats: + row = {} + for tcat in cats: + per_cell = [min(path(a, b) for b in cats[tcat]) for a in cats[s]] + row[tcat] = float(np.mean(per_cell)) + out[s] = row + return pd.DataFrame(out).T -def test_pairwise_stored_in_uns(balanced_tdata): - """Results stored in tdata.uns with expected keys.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype") - assert "celltype_linkage" in tdata.uns - assert "celltype_linkage_params" in tdata.uns - assert "celltype_linkage_stats" in tdata.uns - df = tdata.uns["celltype_linkage"] - assert isinstance(df, pd.DataFrame) - assert set(df.index) == {"A", "B"} - assert set(df.columns) == {"A", "B"} - - -def test_pairwise_stats_has_n_columns(balanced_tdata): - """Stats DataFrame has source_n and target_n columns (always).""" + +# ── pairwise mode ─────────────────────────────────────────────────────────── + + +def test_pairwise_known_values(balanced_tdata): + """Exact linkage values and within/between relationships for each metric/aggregate.""" tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype") - stats = tdata.uns["celltype_linkage_stats"] - assert isinstance(stats, pd.DataFrame) - assert "source_n" in stats.columns - assert "target_n" in stats.columns - assert "source" in stats.columns - assert "target" in stats.columns - assert "value" in stats.columns - # 2 categories → 4 rows - assert len(stats) == 4 - # source_n for category A = 2 leaves - assert stats.loc[stats["source"] == "A"].iloc[0]["source_n"] == 2 + # lca + max: within (self) = 1.0, between = root depth 0.0 + lca = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", + normalize=False, copy=True) + assert np.isclose(lca.loc["A", "A"], 1.0) and np.isclose(lca.loc["B", "B"], 1.0) + assert np.isclose(lca.loc["A", "B"], 0.0) + assert lca.loc["A", "A"] > lca.loc["A", "B"] + # path + min: within (self) = 0.0, between = 2.0 + path = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", + normalize=False, copy=True) + assert np.isclose(path.loc["A", "A"], 0.0) and np.isclose(path.loc["A", "B"], 2.0) + assert path.loc["A", "A"] < path.loc["A", "B"] + # mean path: between = 2.0, within = mean([0, 1, 1, 0]) = 0.5 + mean = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="mean", metric="path", + normalize=False, copy=True) + assert np.isclose(mean.loc["A", "B"], 2.0) and np.isclose(mean.loc["A", "A"], 0.5) -def test_pairwise_copy_returns_dataframe(balanced_tdata): - """copy=True returns DataFrame and also stores in uns.""" +def test_pairwise_output_structure(balanced_tdata): + """uns keys, matrix shape, stats columns, key_added, and copy return type.""" tdata = balanced_tdata - result = tl.ancestral_linkage(tdata, groupby="celltype", copy=True) + result = tl.ancestral_linkage(tdata, groupby="celltype", key_added="mykey", copy=True) + # copy returns the linkage matrix (no test) assert isinstance(result, pd.DataFrame) - assert "celltype_linkage" in tdata.uns + assert set(result.index) == {"A", "B"} and set(result.columns) == {"A", "B"} + # uns keys follow the key_added prefix + for suffix in ("_linkage", "_linkage_params", "_linkage_stats"): + assert f"mykey{suffix}" in tdata.uns + # stats table structure + stats = tdata.uns["mykey_linkage_stats"] + assert {"source", "target", "source_n", "target_n", "value"} <= set(stats.columns) + assert "tree" not in stats.columns # by_tree=False + assert len(stats) == 4 # 2 categories -> 4 ordered pairs + assert stats.loc[stats["source"] == "A"].iloc[0]["source_n"] == 2 -def test_pairwise_key_added(balanced_tdata): - """key_added controls storage key.""" +def test_custom_callable_aggregate(balanced_tdata): + """A custom callable aggregate matches the equivalent named aggregate.""" tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", key_added="mykey") - assert "mykey_linkage" in tdata.uns - assert "mykey_linkage_params" in tdata.uns - assert "mykey_linkage_stats" in tdata.uns + mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate=np.mean, metric="path", copy=True) + ref = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="mean", metric="path", + key_added="ref", copy=True) + pd.testing.assert_frame_equal(mat, ref) -# ── min_size tests ──────────────────────────────────────────────────────────── +# ── min_size ────────────────────────────────────────────────────────────── -def test_min_size_excludes_small_categories(three_cat_tdata): - """Categories with fewer than min_size cells are dropped from the pairwise matrix.""" +def test_min_size(three_cat_tdata): + """min_size drops small categories (matrix + stats), is recorded, and works with a test.""" tdata = three_cat_tdata - # Give category C a single cell by relabeling c2 → B; C now has 1 cell. + # default keeps everything + mat = tl.ancestral_linkage(tdata, groupby="celltype", min_size=1, copy=True) + assert set(mat.index) == {"A", "B", "C"} + # relabel so C has a single cell, then exclude it tdata.obs["celltype"] = ["A", "A", "B", "B", "C", "B"] - tl.ancestral_linkage(tdata, groupby="celltype", min_size=2) + tl.ancestral_linkage(tdata, groupby="celltype", min_size=2, test="permutation", + n_permutations=20, random_state=0) mat = tdata.uns["celltype_linkage"] - assert "C" not in mat.index - assert "C" not in mat.columns - assert set(mat.index) == {"A", "B"} - # Stats table also excludes C stats = tdata.uns["celltype_linkage_stats"] - assert "C" not in stats["source"].values - assert "C" not in stats["target"].values - - -def test_min_size_default_keeps_all(three_cat_tdata): - """min_size=1 (default) keeps every category.""" - tdata = three_cat_tdata - tl.ancestral_linkage(tdata, groupby="celltype", min_size=1) - mat = tdata.uns["celltype_linkage"] - assert set(mat.index) == {"A", "B", "C"} - - -def test_min_size_stored_in_params(balanced_tdata): - """min_size is recorded in the linkage params dict.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", min_size=2) + assert set(mat.index) == {"A", "B"} and "C" not in mat.columns + assert "C" not in stats["source"].values and "C" not in stats["target"].values + assert stats["p_value"].between(0, 1).all() assert tdata.uns["celltype_linkage_params"]["min_size"] == 2 -def test_min_size_invalid_raises(balanced_tdata): - """Non-positive or non-integer min_size raises ValueError.""" +def test_min_size_invalid_raises(balanced_tdata, three_cat_tdata): + """Non-positive min_size and a min_size that excludes everything both raise.""" with pytest.raises(ValueError, match="min_size"): tl.ancestral_linkage(balanced_tdata, groupby="celltype", min_size=0) - - -def test_min_size_all_excluded_raises(three_cat_tdata): - """If no category meets min_size, a ValueError is raised.""" - tdata = three_cat_tdata with pytest.raises(ValueError, match="min_size"): - tl.ancestral_linkage(tdata, groupby="celltype", min_size=100) + tl.ancestral_linkage(three_cat_tdata, groupby="celltype", min_size=100) -def test_min_size_with_permutation(three_cat_tdata): - """min_size works together with a permutation test.""" - tdata = three_cat_tdata - tdata.obs["celltype"] = ["A", "A", "B", "B", "C", "B"] - tl.ancestral_linkage( - tdata, groupby="celltype", min_size=2, test="permutation", - n_permutations=20, random_state=0, - ) - stats = tdata.uns["celltype_linkage_stats"] - assert set(stats["source"]) == {"A", "B"} - assert (stats["p_value"].between(0, 1)).all() +# ── warnings ──────────────────────────────────────────────────────────────── -@pytest.fixture -def small_category_tdata(): - """Star tree with a large category A (12 cells) and a small category B (3 cells).""" - t = nx.DiGraph() - t.add_node("root", depth=0.0) - labels = {} - for i in range(12): - t.add_node(f"a{i}", depth=1.0) - t.add_edge("root", f"a{i}") - labels[f"a{i}"] = "A" - for i in range(3): - t.add_node(f"b{i}", depth=1.0) - t.add_edge("root", f"b{i}") - labels[f"b{i}"] = "B" - obs = pd.DataFrame({"celltype": list(labels.values())}, index=list(labels)) - return td.TreeData(obs=obs, obst={"tree": t}) - - -def test_small_category_warning(small_category_tdata): - """Categories with fewer than 10 cells trigger a warning that lists them.""" +def test_warnings(balanced_tdata, small_category_tdata): + """Small categories warn (and can be silenced by min_size); lca+min warns.""" + # Small category B (3 cells) is flagged, large category A (12) is not. with pytest.warns(UserWarning, match="noisy linkage") as record: tl.ancestral_linkage(small_category_tdata, groupby="celltype") msg = next(str(w.message) for w in record if "noisy linkage" in str(w.message)) - assert "'B'" in msg # the small category is listed - assert "'A'" not in msg # A (12 cells) is not flagged - - -def test_small_category_warning_silenced_by_min_size(small_category_tdata): - """Excluding the small category via min_size removes the noisy-linkage warning.""" - import warnings as _warnings - - with _warnings.catch_warnings(record=True) as rec: - _warnings.simplefilter("always") + assert "'B'" in msg and "'A'" not in msg + # Excluding the small category via min_size removes the warning. + with warnings.catch_warnings(record=True) as rec: + warnings.simplefilter("always") tl.ancestral_linkage(small_category_tdata, groupby="celltype", min_size=10) assert not any("noisy linkage" in str(w.message) for w in rec) + # aggregate='min' with metric='lca' warns about shallowest ancestor. + with pytest.warns(UserWarning, match="shallowest"): + tl.ancestral_linkage(balanced_tdata, groupby="celltype", aggregate="min", metric="lca") -# ── symmetrize tests ────────────────────────────────────────────────────────── +# ── symmetrize ────────────────────────────────────────────────────────────── -def test_symmetrize_mean(three_cat_tdata): - """After symmetrize='mean', M[i,j] == M[j,i].""" +def test_symmetrize(three_cat_tdata): + """symmetrize='mean'/'max' make the matrix symmetric; 'max' >= raw; invalid raises.""" tdata = three_cat_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", - symmetrize="mean") - mat = tdata.uns["celltype_linkage"] - for i in mat.index: - for j in mat.columns: - assert np.isclose(mat.loc[i, j], mat.loc[j, i]) - - -def test_symmetrize_max(three_cat_tdata): - """After symmetrize='max', M[i,j] == M[j,i] and M[i,j] >= original M[i,j].""" - tdata = three_cat_tdata - result_nosym = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", - metric="path", copy=True) - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", - symmetrize="max") - mat = tdata.uns["celltype_linkage"] - for i in mat.index: - for j in mat.columns: - assert np.isclose(mat.loc[i, j], mat.loc[j, i]) - assert mat.loc[i, j] >= result_nosym.loc[i, j] - 1e-9 - - -# ── single-target mode tests ────────────────────────────────────────────────── + raw = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", copy=True) + for mode in ("mean", "max"): + mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", + symmetrize=mode, copy=True) + for i in mat.index: + for j in mat.columns: + assert np.isclose(mat.loc[i, j], mat.loc[j, i]) + if mode == "max": + assert mat.loc[i, j] >= raw.loc[i, j] - 1e-9 + with pytest.raises(ValueError, match="symmetrize"): + tl.ancestral_linkage(tdata, groupby="celltype", symmetrize="meen") -def test_single_target_stores_in_obs(balanced_tdata): - """Single-target result stored in tdata.obs.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", target="B") - assert "B_linkage" in tdata.obs.columns - assert tdata.obs["B_linkage"].notna().all() +# ── single-target mode ──────────────────────────────────────────────────── -def test_single_target_path_known_values(balanced_tdata): - """a1, a2 should have path distance 2.0 to nearest B leaf; b1, b2 distance to self = 0.""" +def test_single_target_known_values(balanced_tdata): + """Single-target stores per-cell scores in obs with exact values, and copy returns a summary.""" tdata = balanced_tdata + # path: a* are 2.0 from nearest B, b* (in target) are 0.0 tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="path", normalize=False) + assert "B_linkage" in tdata.obs.columns and tdata.obs["B_linkage"].notna().all() assert np.isclose(tdata.obs.loc["a1", "B_linkage"], 2.0) - assert np.isclose(tdata.obs.loc["a2", "B_linkage"], 2.0) - # b1, b2 are in target B → distance to self = 0 assert np.isclose(tdata.obs.loc["b1", "B_linkage"], 0.0) - assert np.isclose(tdata.obs.loc["b2", "B_linkage"], 0.0) - - -def test_single_target_lca_known_values(balanced_tdata): - """a1, a2 should have LCA depth 0.0 to best B leaf; b1/b2 in B → LCA with self = depth = 1.0.""" - tdata = balanced_tdata + # lca: a* have LCA depth 0.0 to best B; b1 (in B) has LCA with self = depth 1.0 tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="lca", normalize=False) assert np.isclose(tdata.obs.loc["a1", "B_linkage"], 0.0) - # b1 is in B; nearest is itself (path=0), lca = (1.0 + 1.0 - 0) / 2 = 1.0 assert np.isclose(tdata.obs.loc["b1", "B_linkage"], 1.0) - - -def test_single_target_copy_returns_series_df(balanced_tdata): - """copy=True returns a DataFrame with per-category means.""" - tdata = balanced_tdata - result = tl.ancestral_linkage(tdata, groupby="celltype", target="B", - metric="path", copy=True) + # copy returns a per-category summary DataFrame + result = tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="path", copy=True) assert isinstance(result, pd.DataFrame) - assert "B_linkage" in result.columns - assert "A" in result.index and "B" in result.index - + assert "B_linkage" in result.columns and {"A", "B"} <= set(result.index) -def test_single_target_invalid_target(balanced_tdata): - """Specifying a non-existent target raises ValueError.""" - with pytest.raises(ValueError, match="not found"): - tl.ancestral_linkage(balanced_tdata, groupby="celltype", target="Z") +# ── by_tree ────────────────────────────────────────────────────────────── -# ── by_tree tests ───────────────────────────────────────────────────────────── - -def test_by_tree_stats_has_tree_column(two_tree_tdata): - """When by_tree=True, stats DataFrame has a 'tree' column.""" +def test_by_tree(two_tree_tdata): + """by_tree adds a per-tree 'tree' column with per-tree counts; matrix matches global.""" tdata = two_tree_tdata + global_mat = tl.ancestral_linkage(tdata, groupby="celltype", copy=True) tl.ancestral_linkage(tdata, groupby="celltype", by_tree=True) stats = tdata.uns["celltype_linkage_stats"] assert "tree" in stats.columns - # 2 trees × 2 cats × 2 cats = 8 rows - assert len(stats) == 8 assert set(stats["tree"]) == {"tree1", "tree2"} + assert len(stats) == 8 # 2 trees × 2 × 2 + # each tree has one A leaf and one B leaf + a_row = stats[(stats["tree"] == "tree1") & (stats["source"] == "A")].iloc[0] + assert a_row["source_n"] == 1 and a_row["target_n"] == 1 + # the aggregated linkage matrix is identical to the global (by_tree=False) result + pd.testing.assert_frame_equal(global_mat, tdata.uns["celltype_linkage"]) -def test_by_tree_source_target_n(two_tree_tdata): - """Per-tree source_n and target_n reflect leaves in that tree only.""" - tdata = two_tree_tdata - tl.ancestral_linkage(tdata, groupby="celltype", by_tree=True) - stats = tdata.uns["celltype_linkage_stats"] - # Each tree has 1 A leaf and 1 B leaf - tree1_rows = stats[stats["tree"] == "tree1"] - a_row = tree1_rows[tree1_rows["source"] == "A"].iloc[0] - assert a_row["source_n"] == 1 - assert a_row["target_n"] == 1 - - -def test_by_tree_linkage_matches_global(two_tree_tdata): - """linkage df with by_tree=True is the same as by_tree=False (weighted mean).""" - tdata = two_tree_tdata - result_global = tl.ancestral_linkage(tdata, groupby="celltype", copy=True) - tl.ancestral_linkage(tdata, groupby="celltype", by_tree=True) - result_by_tree = tdata.uns["celltype_linkage"] - pd.testing.assert_frame_equal(result_global, result_by_tree) - - -def test_by_tree_false_no_tree_column(balanced_tdata): - """When by_tree=False (default), stats has no 'tree' column.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype") - stats = tdata.uns["celltype_linkage_stats"] - assert "tree" not in stats.columns - - -# ── permutation test ────────────────────────────────────────────────────────── - - -def test_permutation_test_pairwise_linkage_stores_enrichment(balanced_tdata): - """Pairwise permutation test stores observed - permuted_mean in uns[linkage].""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=20, - random_state=42, - ) - mat = tdata.uns["celltype_linkage"] - assert isinstance(mat, pd.DataFrame) - assert set(mat.index) == {"A", "B"} - assert set(mat.columns) == {"A", "B"} - # Values are finite (observed - permuted_mean, not z-scores) - assert np.isfinite(mat.values).all() - - -def test_permutation_test_pairwise_copy_returns_stats(balanced_tdata): - """copy=True with permutation test returns linkage_stats DataFrame.""" - tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=20, - random_state=42, copy=True - ) - assert isinstance(result, pd.DataFrame) - assert set(result.columns) >= {"source", "target", "value", "permuted_value", "z_score", "p_value"} - assert set(result["source"]) == {"A", "B"} +# ── permutation test ──────────────────────────────────────────────────────── -def test_copy_no_test_returns_linkage_matrix(balanced_tdata): - """copy=True without a test returns the linkage matrix DataFrame.""" +def test_permutation_pairwise(balanced_tdata): + """Pairwise permutation test: enrichment matrix + stats columns; no separate uns keys.""" tdata = balanced_tdata - result = tl.ancestral_linkage(tdata, groupby="celltype", copy=True) + result = tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", + n_permutations=20, random_state=42, copy=True) + # copy returns the stats DataFrame when a test is run assert isinstance(result, pd.DataFrame) - assert set(result.index) == {"A", "B"} - assert set(result.columns) == {"A", "B"} + assert {"source", "target", "value", "permuted_value", "z_score", "p_value"} <= set(result.columns) + assert result["p_value"].between(0, 1).all() + assert result["permuted_value"].notna().all() + # linkage matrix stored, finite; no stray z/p uns keys + mat = tdata.uns["celltype_linkage"] + assert set(mat.index) == {"A", "B"} and np.isfinite(mat.values).all() + assert "celltype_z_score" not in tdata.uns and "celltype_p_value" not in tdata.uns -def test_permutation_test_single_target(balanced_tdata): - """Single-target permutation test returns long DataFrame.""" +def test_permutation_single_target(balanced_tdata): + """Single-target permutation test returns a long DataFrame scoped to the target.""" tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="B", test="permutation", - n_permutations=20, random_state=0, copy=True - ) + result = tl.ancestral_linkage(tdata, groupby="celltype", target="B", test="permutation", + n_permutations=20, random_state=0, copy=True) assert isinstance(result, pd.DataFrame) - assert set(result.columns) >= {"source", "target", "value", "z_score", "p_value"} + assert {"source", "target", "value", "z_score", "p_value", "permuted_value"} <= set(result.columns) assert (result["target"] == "B").all() - - -def test_permutation_test_stored_in_stats(balanced_tdata): - """Permutation test stores z_score, p_value, and permuted_value in stats.""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=10, - random_state=1 - ) - assert "celltype_linkage" in tdata.uns - assert "celltype_linkage_stats" in tdata.uns - # No separate z_score / p_value uns keys - assert "celltype_z_score" not in tdata.uns - assert "celltype_p_value" not in tdata.uns - stats = tdata.uns["celltype_linkage_stats"] - assert "z_score" in stats.columns - assert "p_value" in stats.columns - assert "permuted_value" in stats.columns - assert (stats["p_value"].between(0, 1)).all() - - -def test_permuted_value_in_pairwise_stats(balanced_tdata): - """permuted_value column holds mean null distribution value for each (src, tgt) pair.""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=20, - random_state=42 - ) - stats = tdata.uns["celltype_linkage_stats"] - assert "permuted_value" in stats.columns - assert stats["permuted_value"].notna().all() - - -def test_permuted_value_in_single_target_test(balanced_tdata): - """permuted_value in single-target test DataFrame holds mean null score.""" - tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="B", test="permutation", - n_permutations=20, random_state=0, copy=True - ) - assert "permuted_value" in result.columns assert result["permuted_value"].notna().all() -def test_n_threads_parallel_matches_serial(balanced_tdata): - """n_threads > 1 produces same z-score DataFrame as single-threaded (same seed).""" - def run(n_threads): - tdata = balanced_tdata - return tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=20, - random_state=42, n_threads=n_threads, copy=True, - ) - - serial = run(None) - parallel = run(2) - pd.testing.assert_frame_equal(serial, parallel) - - -def test_permutation_test_stats_not_symmetrized(balanced_tdata): - """Stats z_scores are not symmetrized even when symmetrize is set.""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=20, - random_state=5, symmetrize="mean" - ) - stats = tdata.uns["celltype_linkage_stats"] - # linkage (z_score) should be symmetrized - mat = tdata.uns["celltype_linkage"] - assert np.isclose(mat.loc["A", "B"], mat.loc["B", "A"]) - # stats z_scores might not be symmetric (AB vs BA could differ) - ab = stats[(stats["source"] == "A") & (stats["target"] == "B")]["z_score"].values[0] - ba = stats[(stats["source"] == "B") & (stats["target"] == "A")]["z_score"].values[0] - # We just check both are present; symmetry is not required - assert np.isfinite(ab) or np.isnan(ab) - assert np.isfinite(ba) or np.isnan(ba) - - -def test_permutation_random_state(balanced_tdata): - """Same random_state produces identical output DataFrames.""" - def run(seed): - tdata = balanced_tdata - return tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", n_permutations=30, - random_state=seed, copy=True - ) - - r1 = run(7) - r2 = run(7) - pd.testing.assert_frame_equal(r1, r2) - - -def test_by_tree_permutation_stats_has_z_score(two_tree_tdata): - """by_tree=True with permutation test adds z_score/p_value/permuted_value to stats rows.""" +def test_permutation_by_tree(two_tree_tdata): + """by_tree + permutation stores per-tree z_score/p_value/permuted_value.""" tdata = two_tree_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", by_tree=True, test="permutation", - n_permutations=10, random_state=0, - ) + tl.ancestral_linkage(tdata, groupby="celltype", by_tree=True, test="permutation", + n_permutations=10, random_state=0) stats = tdata.uns["celltype_linkage_stats"] - assert "z_score" in stats.columns - assert "p_value" in stats.columns - assert "permuted_value" in stats.columns - assert (stats["p_value"].between(0, 1)).all() - - -# ── alternative parameter tests ─────────────────────────────────────────────── + assert {"z_score", "p_value", "permuted_value"} <= set(stats.columns) + assert stats["p_value"].between(0, 1).all() -def test_alternative_two_sided_pairwise_p_values_in_range(balanced_tdata): - """Two-sided p-values are in [0, 1].""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", alternative="two-sided", - n_permutations=20, random_state=0, - ) - stats = tdata.uns["celltype_linkage_stats"] - assert (stats["p_value"].between(0, 1)).all() +def test_permutation_reproducible(balanced_tdata): + """Same random_state is deterministic; parallel matches serial.""" + def run(n_threads): + return tl.ancestral_linkage(balanced_tdata, groupby="celltype", test="permutation", + n_permutations=20, random_state=42, n_threads=n_threads, copy=True) + pd.testing.assert_frame_equal(run(None), run(None)) # deterministic + pd.testing.assert_frame_equal(run(None), run(2)) # parallel == serial -def test_alternative_two_sided_symmetric(balanced_tdata): - """Two-sided p-values are symmetric: p(A→B) == p(B→A) when the tree is symmetric.""" +def test_alternative(balanced_tdata): + """`alternative` p-values stay in range, are symmetric on a symmetric tree, and None==default.""" tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", alternative="two-sided", - n_permutations=50, random_state=1, - ) + tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", alternative="two-sided", + n_permutations=50, random_state=1) stats = tdata.uns["celltype_linkage_stats"] + assert stats["p_value"].between(0, 1).all() ab = stats[(stats["source"] == "A") & (stats["target"] == "B")]["p_value"].values[0] ba = stats[(stats["source"] == "B") & (stats["target"] == "A")]["p_value"].values[0] assert np.isclose(ab, ba) - - -def test_alternative_two_sided_single_target(balanced_tdata): - """Two-sided p-values work in single-target mode.""" - tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="B", test="permutation", - alternative="two-sided", n_permutations=20, random_state=0, copy=True, - ) - assert isinstance(result, pd.DataFrame) - assert (result["p_value"].between(0, 1)).all() - - -def test_alternative_none_matches_default(balanced_tdata): - """alternative=None produces identical results to omitting the parameter.""" - def run(alt): - tdata = balanced_tdata - return tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", alternative=alt, - n_permutations=20, random_state=42, copy=True, - ) - - pd.testing.assert_frame_equal(run(None), run(None)) - - -# ── permutation_mode tests ──────────────────────────────────────────────────── - - -def test_permutation_mode_non_target_pairwise_p_in_range(balanced_tdata): - """non_target mode pairwise: p-values in [0, 1].""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", permutation_mode="non_target", - n_permutations=20, random_state=0, - ) - stats = tdata.uns["celltype_linkage_stats"] - assert (stats["p_value"].between(0, 1)).all() - assert "z_score" in stats.columns - assert "permuted_value" in stats.columns - - -def test_permutation_mode_non_target_single_target_p_in_range(balanced_tdata): - """non_target mode single-target: p-values in [0, 1].""" - tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="B", test="permutation", - permutation_mode="non_target", n_permutations=20, random_state=0, copy=True, - ) - assert isinstance(result, pd.DataFrame) - assert (result["p_value"].between(0, 1)).all() - - -def test_permutation_mode_non_target_differs_from_all(balanced_tdata): - """non_target and all modes produce different null means (different null construction).""" - def run(mode): - tdata = balanced_tdata - return tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", permutation_mode=mode, - n_permutations=50, random_state=42, copy=True, - ) - - stats_all = run("all") - stats_non_target = run("non_target") - # permuted_value may differ between modes (different null distributions) - # At minimum both should produce valid DataFrames with the same structure - assert set(stats_all.columns) == set(stats_non_target.columns) - - -def test_permutation_mode_non_target_by_tree(two_tree_tdata): - """non_target mode works with by_tree=True.""" - tdata = two_tree_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", by_tree=True, test="permutation", - permutation_mode="non_target", n_permutations=10, random_state=0, - ) - stats = tdata.uns["celltype_linkage_stats"] - assert "z_score" in stats.columns - assert (stats["p_value"].between(0, 1)).all() - - -# ── warning tests ───────────────────────────────────────────────────────────── - - -def test_lca_min_warning(balanced_tdata): - """aggregate='min' + metric='lca' should emit a UserWarning.""" - with pytest.warns(UserWarning, match="shallowest"): - tl.ancestral_linkage(balanced_tdata, groupby="celltype", - aggregate="min", metric="lca") - - -# ── edge / validation tests ─────────────────────────────────────────────────── - - -def test_invalid_groupby(balanced_tdata): + # two-sided also works in single-target mode + result = tl.ancestral_linkage(tdata, groupby="celltype", target="B", test="permutation", + alternative="two-sided", n_permutations=20, random_state=0, copy=True) + assert result["p_value"].between(0, 1).all() + # alternative=None matches omitting the parameter (the default) + explicit = tl.ancestral_linkage(balanced_tdata, groupby="celltype", test="permutation", + alternative=None, n_permutations=20, random_state=42, copy=True) + default = tl.ancestral_linkage(balanced_tdata, groupby="celltype", test="permutation", + n_permutations=20, random_state=42, copy=True) + pd.testing.assert_frame_equal(explicit, default) + + +def test_permutation_mode_non_target(balanced_tdata, two_tree_tdata): + """permutation_mode='non_target' yields valid p-values in every mode combination.""" + tdata = balanced_tdata + # pairwise: 'all' and 'non_target' produce the same stats schema and valid p-values + all_df = tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", + permutation_mode="all", n_permutations=20, random_state=42, copy=True) + nt_df = tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", + permutation_mode="non_target", n_permutations=20, random_state=42, copy=True) + assert set(all_df.columns) == set(nt_df.columns) + assert {"z_score", "permuted_value"} <= set(nt_df.columns) + assert nt_df["p_value"].between(0, 1).all() + # single-target + result = tl.ancestral_linkage(tdata, groupby="celltype", target="B", test="permutation", + permutation_mode="non_target", n_permutations=20, random_state=0, copy=True) + assert result["p_value"].between(0, 1).all() + # by_tree + tl.ancestral_linkage(two_tree_tdata, groupby="celltype", by_tree=True, test="permutation", + permutation_mode="non_target", n_permutations=10, random_state=0) + bt_stats = two_tree_tdata.uns["celltype_linkage_stats"] + assert bt_stats["p_value"].between(0, 1).all() + + +# ── validation ────────────────────────────────────────────────────────────── + + +def test_invalid_params(balanced_tdata): + """Invalid groupby, aggregate, and target all raise ValueError.""" with pytest.raises(ValueError, match="not found"): tl.ancestral_linkage(balanced_tdata, groupby="nonexistent") - - -def test_invalid_aggregate(balanced_tdata): with pytest.raises(ValueError, match="aggregate"): tl.ancestral_linkage(balanced_tdata, groupby="celltype", aggregate="median") + with pytest.raises(ValueError, match="not found"): + tl.ancestral_linkage(balanced_tdata, groupby="celltype", target="Z") -def test_custom_callable_aggregate(balanced_tdata): - """Custom callable is accepted for aggregate.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", aggregate=np.mean, metric="path") - mat = tdata.uns["celltype_linkage"] - # Should match mean aggregate - tl.ancestral_linkage(tdata, groupby="celltype", aggregate="mean", metric="path", - key_added="ref") - ref = tdata.uns["ref_linkage"] - pd.testing.assert_frame_equal(mat, ref) - - -# ── bug-fix regression tests ────────────────────────────────────────────────── - - -@pytest.fixture -def non_ultrametric_tdata(): - """Tree with leaves at unequal depths (non-ultrametric), two categories A and B. - - root(0.0) - ├── n1(0.3) - │ ├── a1(0.5) A - │ └── a2(2.0) A - └── n2(0.6) - ├── b1(1.0) B - └── b2(1.5) B - - LCA depths (deepest common ancestor): - LCA(a1, a2) = n1 = 0.3 LCA(b1, b2) = n2 = 0.6 - LCA(a*, b*) = root = 0.0 - """ - t = nx.DiGraph() - for node, depth in [ - ("root", 0.0), ("n1", 0.3), ("n2", 0.6), - ("a1", 0.5), ("a2", 2.0), ("b1", 1.0), ("b2", 1.5), - ]: - t.add_node(node, depth=depth) - t.add_edges_from([ - ("root", "n1"), ("root", "n2"), - ("n1", "a1"), ("n1", "a2"), - ("n2", "b1"), ("n2", "b2"), - ]) - obs = pd.DataFrame({"celltype": ["A", "A", "B", "B"]}, index=["a1", "a2", "b1", "b2"]) - return td.TreeData(obs=obs, obst={"tree": t}) - - -def _bruteforce_lca_max(tdata, groupby="celltype"): - """Reference pairwise max-LCA-depth linkage via explicit LCA on the tree.""" - t = list(tdata.obst.values())[0] - depth = dict(t.nodes(data="depth")) - cats = defaultdict(list) - for leaf, c in tdata.obs[groupby].items(): - cats[c].append(leaf) - - def lca_depth(i, j): - anc_i = nx.ancestors(t, i) | {i} - anc_j = nx.ancestors(t, j) | {j} - return max(depth[a] for a in anc_i & anc_j) - - out = {} - for s in cats: - row = {} - for tcat in cats: - per_cell = [max(depth[a] if a == b else lca_depth(a, b) for b in cats[tcat]) for a in cats[s]] - row[tcat] = float(np.mean(per_cell)) - out[s] = row - return pd.DataFrame(out).T - - -def _bruteforce_path_min(tdata, groupby="celltype"): - """Reference pairwise min-path-distance linkage via explicit LCA on the tree.""" - t = list(tdata.obst.values())[0] - depth = dict(t.nodes(data="depth")) - cats = defaultdict(list) - for leaf, c in tdata.obs[groupby].items(): - cats[c].append(leaf) - - def path(i, j): - if i == j: - return 0.0 - anc_i = nx.ancestors(t, i) | {i} - anc_j = nx.ancestors(t, j) | {j} - lca = max(depth[a] for a in anc_i & anc_j) - return depth[i] + depth[j] - 2 * lca - - out = {} - for s in cats: - row = {} - for tcat in cats: - per_cell = [min(path(a, b) for b in cats[tcat]) for a in cats[s]] - row[tcat] = float(np.mean(per_cell)) - out[s] = row - return pd.DataFrame(out).T - - -def test_path_min_ultrametric_matches_bruteforce(three_cat_tdata): - """Ultrametric path+min (walk-up transform) matches explicit min-path brute force.""" - tdata = three_cat_tdata - mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", normalize=False, symmetrize=False, copy=True) - ref = _bruteforce_path_min(tdata) - pd.testing.assert_frame_equal(mat, ref.loc[mat.index, mat.columns]) +# ── correctness vs brute force (ultrametric & non-ultrametric) ──────────────── -def test_path_min_non_ultrametric_matches_bruteforce(non_ultrametric_tdata): - """Non-ultrametric path+min (Dijkstra) matches explicit min-path brute force.""" - tdata = non_ultrametric_tdata - mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", normalize=False, symmetrize=False, copy=True) +@pytest.mark.parametrize("fixture_name", ["three_cat_tdata", "non_ultrametric_tdata"]) +def test_path_min_matches_bruteforce(fixture_name, request): + """path+min matches an explicit min-path brute force (walk-up and Dijkstra paths).""" + tdata = request.getfixturevalue(fixture_name) + mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="min", metric="path", + normalize=False, symmetrize=False, copy=True) ref = _bruteforce_path_min(tdata) pd.testing.assert_frame_equal(mat, ref.loc[mat.index, mat.columns]) -def test_lca_max_non_ultrametric_known_values(non_ultrametric_tdata): - """aggregate='max' + metric='lca' on a non-ultrametric tree computes exact LCA depths.""" - tdata = non_ultrametric_tdata - mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", normalize=False, symmetrize=False, copy=True) - # Within-A: max LCA per cell is its own depth (self) -> mean(0.5, 2.0) = 1.25 - assert np.isclose(mat.loc["A", "A"], 1.25) - # Within-B: mean(1.0, 1.5) = 1.25 - assert np.isclose(mat.loc["B", "B"], 1.25) - # Between: best A<->B LCA is the root, depth 0.0 - assert np.isclose(mat.loc["A", "B"], 0.0) - assert np.isclose(mat.loc["B", "A"], 0.0) - - -def test_lca_max_non_ultrametric_matches_bruteforce(non_ultrametric_tdata): - """Non-ultrametric lca+max matches an explicit LCA brute-force computation.""" - tdata = non_ultrametric_tdata - mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", normalize=False, symmetrize=False, copy=True) +@pytest.mark.parametrize("fixture_name", ["three_cat_tdata", "non_ultrametric_tdata"]) +def test_lca_max_matches_bruteforce(fixture_name, request): + """lca+max matches an explicit LCA brute force on ultrametric and non-ultrametric trees.""" + tdata = request.getfixturevalue(fixture_name) + mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", + normalize=False, symmetrize=False, copy=True) ref = _bruteforce_lca_max(tdata) pd.testing.assert_frame_equal(mat, ref.loc[mat.index, mat.columns]) + if fixture_name == "non_ultrametric_tdata": + # within-A mean(0.5, 2.0) = 1.25, within-B mean(1.0, 1.5) = 1.25, between = root 0.0 + assert np.isclose(mat.loc["A", "A"], 1.25) and np.isclose(mat.loc["B", "B"], 1.25) + assert np.isclose(mat.loc["A", "B"], 0.0) -def test_lca_max_ultrametric_matches_bruteforce(three_cat_tdata): - """On an ultrametric tree the Dijkstra shortcut matches brute force (and walk-up).""" - tdata = three_cat_tdata - mat = tl.ancestral_linkage(tdata, groupby="celltype", aggregate="max", metric="lca", normalize=False, symmetrize=False, copy=True) - ref = _bruteforce_lca_max(tdata) - pd.testing.assert_frame_equal(mat, ref.loc[mat.index, mat.columns]) - +# ── normalize ──────────────────────────────────────────────────────────────── -def test_lca_max_non_ultrametric_with_permutation(non_ultrametric_tdata): - """Non-ultrametric lca+max works end-to-end with a permutation test.""" - tdata = non_ultrametric_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", aggregate="max", metric="lca", - test="permutation", n_permutations=20, random_state=0, - ) - stats = tdata.uns["celltype_linkage_stats"] - assert (stats["p_value"].between(0, 1)).all() - - -def test_symmetrize_invalid_raises(): - """Unknown symmetrize value must raise ValueError, not silently use 'min'.""" - t = nx.DiGraph() - for node, depth in [("root", 0.0), ("a1", 1.0), ("b1", 1.0)]: - t.add_node(node, depth=depth) - t.add_edges_from([("root", "a1"), ("root", "b1")]) - obs = pd.DataFrame({"ct": ["A", "B"]}, index=["a1", "b1"]) - tdata = td.TreeData(obs=obs, obst={"tree": t}) - with pytest.raises(ValueError, match="symmetrize"): - tl.ancestral_linkage(tdata, groupby="ct", symmetrize="meen") - - -# ── permuted-value / normalize-without-test tests ────────────────────────────── - - -def test_permuted_value_without_test(balanced_tdata): - """permuted_value is populated even without a test (single permutation); z/p are not.""" +def test_permuted_value_and_normalize_without_test(balanced_tdata): + """permuted_value is populated without a test; normalize uses it; z/p are absent.""" tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype") + tl.ancestral_linkage(tdata, groupby="celltype", normalize=False) stats = tdata.uns["celltype_linkage_stats"] - assert "permuted_value" in stats.columns - assert stats["permuted_value"].notna().all() - assert "z_score" not in stats.columns - assert "p_value" not in stats.columns - - -def test_normalize_without_test_pairwise(balanced_tdata): - """normalize=True works without test='permutation' (uses the single-permutation mean).""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", normalize=True) linkage = tdata.uns["celltype_linkage"] - stats = tdata.uns["celltype_linkage_stats"] + assert stats["permuted_value"].notna().all() + assert "z_score" not in stats.columns and "p_value" not in stats.columns + # normalize=False stores the raw linkage for _, row in stats.iterrows(): - expected = row["value"] - row["permuted_value"] - assert linkage.loc[row["source"], row["target"]] == pytest.approx(expected, abs=1e-9) - - -def test_no_test_no_normalize_returns_raw(balanced_tdata): - """With normalize=False, the stored matrix is the raw linkage.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", normalize=False) + assert linkage.loc[row["source"], row["target"]] == pytest.approx(row["value"], abs=1e-9) + # normalize=True (no test) stores value - permuted_value + tl.ancestral_linkage(tdata, groupby="celltype", normalize=True) linkage = tdata.uns["celltype_linkage"] stats = tdata.uns["celltype_linkage_stats"] for _, row in stats.iterrows(): - assert linkage.loc[row["source"], row["target"]] == pytest.approx(row["value"], abs=1e-9) - - -def test_normalize_without_test_single_target(balanced_tdata): - """Single-target normalize works without a full permutation test and creates no _test key.""" - tdata = balanced_tdata + assert linkage.loc[row["source"], row["target"]] == pytest.approx(row["value"] - row["permuted_value"], abs=1e-9) + # single-target normalize without a test creates no _test key tl.ancestral_linkage(tdata, groupby="celltype", target="B", normalize=True) - assert "B_linkage" in tdata.obs.columns - assert "celltype_test" not in tdata.uns # _test only stored with test='permutation' - - -# ── normalize parameter tests ────────────────────────────────────────────────── + assert "B_linkage" in tdata.obs.columns and "celltype_test" not in tdata.uns -def test_pairwise_normalize_stores_value_minus_permuted(balanced_tdata): - """normalize=True stores observed - permuted_mean in the pairwise linkage matrix.""" +def test_pairwise_normalize_with_test(balanced_tdata): + """With test='permutation', normalize toggles between enrichment and raw linkage.""" tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", normalize=True, - n_permutations=20, random_state=0, - ) - linkage = tdata.uns["celltype_linkage"] - stats = tdata.uns["celltype_linkage_stats"] + tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", normalize=True, + n_permutations=20, random_state=0) + linkage, stats = tdata.uns["celltype_linkage"], tdata.uns["celltype_linkage_stats"] for _, row in stats.iterrows(): - expected = row["value"] - row["permuted_value"] - assert linkage.loc[row["source"], row["target"]] == pytest.approx(expected, abs=1e-9) - - -def test_pairwise_no_normalize_stores_raw_linkage(balanced_tdata): - """normalize=False (default) stores raw linkage even when test='permutation'.""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", test="permutation", normalize=False, - n_permutations=20, random_state=0, - ) - linkage = tdata.uns["celltype_linkage"] - stats = tdata.uns["celltype_linkage_stats"] + assert linkage.loc[row["source"], row["target"]] == pytest.approx(row["value"] - row["permuted_value"], abs=1e-9) + tl.ancestral_linkage(tdata, groupby="celltype", test="permutation", normalize=False, + n_permutations=20, random_state=0) + linkage, stats = tdata.uns["celltype_linkage"], tdata.uns["celltype_linkage_stats"] for _, row in stats.iterrows(): assert linkage.loc[row["source"], row["target"]] == pytest.approx(row["value"], abs=1e-9) -def test_single_target_normalize_overwrites_linkage(balanced_tdata): - """normalize=True replaces _linkage in obs with score - category_permuted_mean.""" +def test_single_target_normalize(balanced_tdata): + """Single-target normalize overwrites obs with score - category_permuted_mean.""" tdata = balanced_tdata - # Raw linkage without normalize - tl.ancestral_linkage(tdata, groupby="celltype", target="B", test="permutation", + tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="path", test="permutation", normalize=False, n_permutations=30, random_state=1) raw = tdata.obs["B_linkage"].copy() - - # Normalized linkage - tl.ancestral_linkage(tdata, groupby="celltype", target="B", test="permutation", + assert raw.loc["a1"] == pytest.approx(2.0) # raw min-path score + tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="path", test="permutation", normalize=True, n_permutations=30, random_state=1) - norm = tdata.obs["B_linkage"].copy() + norm = tdata.obs["B_linkage"] test_df = tdata.uns["celltype_test"] - - # Normalized values should differ from raw (permuted_mean is non-trivial) - # And for each cell: norm = raw - cat_permuted_mean for cell in ["a1", "a2", "b1", "b2"]: cat = tdata.obs.loc[cell, "celltype"] perm_val = test_df.loc[test_df["source"] == cat, "permuted_value"].iloc[0] assert norm[cell] == pytest.approx(raw[cell] - perm_val, abs=1e-9) - assert "B_norm_linkage" not in tdata.obs.columns -def test_single_target_normalize_copy_matches_obs(balanced_tdata): - """copy=True per-category means match the normalized obs column (test=None default path).""" - tdata = balanced_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="B", metric="path", - normalize=True, random_state=1, copy=True, # test=None -> single-permutation normalization - ) - obs = tdata.obs - for cat in ["A", "B"]: - cells = obs.index[obs["celltype"] == cat] - expected = float(np.nanmean(obs.loc[cells, "B_linkage"])) - assert result.loc[cat, "B_linkage"] == pytest.approx(expected, abs=1e-9) - - -def test_single_target_by_tree_normalize_copy_matches_obs(three_cat_tdata): - """by_tree copy=True per-category means match the normalized obs column (test=None).""" +def test_single_target_by_tree(three_cat_tdata, balanced_tdata): + """by_tree single-target: per-tree test rows, normalized copy matches obs means, non_target.""" tdata = three_cat_tdata - result = tl.ancestral_linkage( - tdata, groupby="celltype", target="A", metric="path", - normalize=True, by_tree=True, random_state=1, copy=True, - ) - obs = tdata.obs + # copy per-category means match the (normalized) obs column + result = tl.ancestral_linkage(tdata, groupby="celltype", target="A", metric="path", + normalize=True, by_tree=True, random_state=1, copy=True) for cat in ["A", "B", "C"]: - cells = obs.index[obs["celltype"] == cat] - expected = float(np.nanmean(obs.loc[cells, "A_linkage"])) + cells = tdata.obs.index[tdata.obs["celltype"] == cat] + expected = float(np.nanmean(tdata.obs.loc[cells, "A_linkage"])) assert result.loc[cat, "A_linkage"] == pytest.approx(expected, abs=1e-9) - - -def test_single_target_no_normalize_keeps_raw(balanced_tdata): - """normalize=False does not overwrite _linkage in obs.""" - tdata = balanced_tdata - tl.ancestral_linkage(tdata, groupby="celltype", target="B", metric="path", test="permutation", - normalize=False, n_permutations=20, random_state=2) - # With normalize=False, a1 (category A, not in B) has raw min-path score 2.0 - assert tdata.obs.loc["a1", "B_linkage"] == pytest.approx(2.0) - - -def test_single_target_by_tree_permutation(three_cat_tdata): - """by_tree=True with single target + permutation runs per-tree (stats has tree column).""" - tdata = three_cat_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", target="A", test="permutation", by_tree=True, - n_permutations=20, random_state=7, - ) - assert "A_linkage" in tdata.obs.columns - assert "A_norm_linkage" not in tdata.obs.columns + # by_tree + permutation writes a per-tree test table + tl.ancestral_linkage(tdata, groupby="celltype", target="A", test="permutation", by_tree=True, + n_permutations=20, random_state=7) test_df = tdata.uns["celltype_test"] assert "tree" in test_df.columns - for tree_key in tdata.obst: - assert tree_key in test_df["tree"].values - - -def test_single_target_by_tree_normalize(three_cat_tdata): - """by_tree=True + normalize=True normalizes each cell using its own tree's null distribution.""" - tdata = three_cat_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", target="A", test="permutation", - by_tree=True, normalize=True, n_permutations=20, random_state=7, - ) - assert "A_linkage" in tdata.obs.columns + assert set(tdata.obst) <= set(test_df["tree"].values) assert "A_norm_linkage" not in tdata.obs.columns - # _linkage values should be finite for leaves with valid scores - assert tdata.obs["A_linkage"].notna().any() + # by_tree + non_target permutation mode also produces a tree column + tl.ancestral_linkage(balanced_tdata, groupby="celltype", target="B", test="permutation", + by_tree=True, permutation_mode="non_target", n_permutations=20, random_state=3) + assert "tree" in balanced_tdata.uns["celltype_test"].columns -def test_single_target_by_tree_perm_non_target(balanced_tdata): - """by_tree + permutation_mode='non_target' in single-target mode produces tree column in stats.""" - tdata = balanced_tdata - tl.ancestral_linkage( - tdata, groupby="celltype", target="B", test="permutation", - by_tree=True, permutation_mode="non_target", n_permutations=20, random_state=3, - ) - assert "B_norm_linkage" not in tdata.obs.columns - test_df = tdata.uns["celltype_test"] - assert "tree" in test_df.columns +if __name__ == "__main__": + pytest.main(["-v", __file__]) From e141785c34ef56c1ba8730c7b7674a6ba3906430 Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 18:40:51 -0400 Subject: [PATCH 4/7] Type-checking and robustness fixes across tl/pl - get_keyed_obs_data: pass through already-categorical array data unchanged instead of recoercing. - ancestral_linkage: accept Mapping trees, map obs via plain dicts, and annotate tree_distance/compute_scores calls for the type checker. - partition_test: support callable metrics, use np.atleast_2d for stat reshaping, and quiet ttest_ind typing. - pl.ancestral_linkage: clearer error when the stats key or figure is missing. - docs: register tl.ancestral_linkage in the API reference. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/api.md | 1 + src/pycea/pl/plot_ancestral_linkage.py | 8 ++++++-- src/pycea/tl/ancestral_linkage.py | 22 +++++++++++----------- src/pycea/tl/partition_test.py | 18 ++++++++++++++---- src/pycea/utils.py | 14 +++++++++----- 5 files changed, 41 insertions(+), 22 deletions(-) diff --git a/docs/api.md b/docs/api.md index 17a9119..92bfb24 100644 --- a/docs/api.md +++ b/docs/api.md @@ -22,6 +22,7 @@ :toctree: generated tl.ancestral_states + tl.ancestral_linkage tl.autocorr tl.clades tl.compare_distance diff --git a/src/pycea/pl/plot_ancestral_linkage.py b/src/pycea/pl/plot_ancestral_linkage.py index 286c0ce..a45b677 100644 --- a/src/pycea/pl/plot_ancestral_linkage.py +++ b/src/pycea/pl/plot_ancestral_linkage.py @@ -142,8 +142,9 @@ def ancestral_linkage( groupby = candidates[0] stats = tdata.uns.get(f"{groupby}{_STATS_SUFFIX}") if stats is None: + key_added = groupby if groupby is not None else "" raise KeyError( - f"{groupby + _STATS_SUFFIX!r} not found in tdata.uns. " + f"{key_added + _STATS_SUFFIX!r} not found in tdata.uns. " f"Run pycea.tl.ancestral_linkage(..., key_added={groupby!r}) first." ) stats = stats.copy() @@ -256,6 +257,9 @@ def ancestral_linkage( ax.set_ylabel("") if cbar: - ax.get_figure().colorbar(mesh, ax=ax) + fig = ax.get_figure() + if fig is None: + raise RuntimeError("no figure associated with ax; cannot add colorbar") + fig.colorbar(mesh, ax=ax) return ax diff --git a/src/pycea/tl/ancestral_linkage.py b/src/pycea/tl/ancestral_linkage.py index a2680a2..277ec8b 100644 --- a/src/pycea/tl/ancestral_linkage.py +++ b/src/pycea/tl/ancestral_linkage.py @@ -4,7 +4,7 @@ import sys import warnings from collections import defaultdict -from collections.abc import Callable, Sequence +from collections.abc import Callable, Mapping, Sequence from typing import Literal, overload import networkx as nx @@ -123,7 +123,7 @@ def _max_lca_depth_scores( def _all_pairs_scores( tdata: td.TreeData, - trees: dict, + trees: dict | Mapping, source_leaves_by_tree: dict, target_cats: list, cat_to_leaves_by_tree: dict, @@ -135,7 +135,7 @@ def _all_pairs_scores( from pycea.tl.tree_distance import tree_distance as _td tree_keys = list(trees.keys()) - result = _td(tdata, depth_key=depth_key, metric=metric, tree=tree_keys, copy=True) + result = _td(tdata, depth_key=depth_key, metric=metric, tree=tree_keys, copy=True) # type: ignore precomputed = result.toarray() if sp.sparse.issparse(result) else result scores: dict = {} @@ -160,7 +160,7 @@ def _all_pairs_scores( def _compute_scores( tdata: td.TreeData, - trees: dict, + trees: dict | Mapping, leaf_to_cat: dict, target_cats: list, aggregate: str | Callable, @@ -416,7 +416,7 @@ def _compute_p_values( def _run_permutation_test( tdata: td.TreeData, - trees: dict, + trees: dict | Mapping, leaf_to_cat: dict, all_cats: list, target_cats: list, @@ -470,7 +470,7 @@ def _run_permutation_test( def _run_permutation_test_non_target( tdata: td.TreeData, - trees: dict, + trees: dict | Mapping, leaf_to_cat: dict, all_cats: list, observed_df: pd.DataFrame, @@ -954,9 +954,9 @@ def _run_single_perm(single_tree, tree_lc, tree_sm, tree_cl, extra_row_fields=No perm_val = cat_null_mean.get(cat, np.nan) if cat is not None else np.nan merged_norm_map[leaf] = (score - perm_val) if not np.isnan(score) else np.nan - tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(merged_score_map, dtype=float)) + tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(merged_score_map, dtype=float).to_dict()) if normalize: - tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(merged_norm_map, dtype=float)) + tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(merged_norm_map, dtype=float).to_dict()) # Return per-category means from whichever map was written to obs. linkage_map = merged_norm_map if normalize else merged_score_map if test == "permutation": @@ -977,9 +977,9 @@ def _run_single_perm(single_tree, tree_lc, tree_sm, tree_cl, extra_row_fields=No else: # Global (non-by_tree) path - all_scores = _compute_scores(tdata, trees, leaf_to_cat, [target], single_agg, metric, depth_key) + all_scores = _compute_scores(tdata, trees, leaf_to_cat, [target], single_agg, metric, depth_key) # type: ignore score_map = {leaf: scores.get(target, np.nan) for leaf, scores in all_scores.items()} - tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(score_map, dtype=float)) + tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(score_map, dtype=float).to_dict()) if run_perm: rows, cat_null_mean = _run_single_perm(trees, leaf_to_cat, score_map, cat_to_leaves) @@ -992,7 +992,7 @@ def _run_single_perm(single_tree, tree_lc, tree_sm, tree_cl, extra_row_fields=No else np.nan for leaf, score in score_map.items() } - tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(score_map, dtype=float)) + tdata.obs[f"{target}_linkage"] = tdata.obs.index.map(pd.Series(score_map, dtype=float).to_dict()) if test == "permutation": test_df = pd.DataFrame(rows) tdata.uns[f"{key_added}_test"] = test_df diff --git a/src/pycea/tl/partition_test.py b/src/pycea/tl/partition_test.py index 64972ef..c437866 100644 --- a/src/pycea/tl/partition_test.py +++ b/src/pycea/tl/partition_test.py @@ -22,7 +22,7 @@ def _run_permutations( data: pd.DataFrame, n_permutations: int, aggregate_fn: _AggregatorFn, - metric_fn: _MetricFn, + metric_fn: MeanDiffMetric | _MetricFn | DistanceMetric, n_right: int, n_left: int, ) -> np.ndarray: @@ -66,7 +66,12 @@ def _run_permutations( left_stat = aggregate_fn(left_df.to_numpy()) right_stat = aggregate_fn(right_df.to_numpy()) - permutation_vals[i] = np.squeeze(metric_fn.pairwise(left_stat.reshape(1, -1), right_stat.reshape(1, -1))) + permutation_vals[i] = np.squeeze( + metric_fn.pairwise( + np.atleast_2d(left_stat), + np.atleast_2d(right_stat), + ) + ) return permutation_vals @@ -249,6 +254,8 @@ def partition_test( if metric == "mean_difference": metric_fn = MeanDiffMetric() + elif callable(metric): + metric_fn = metric else: metric_fn = DistanceMetric.get_metric(metric, **(metric_kwds or {})) @@ -321,7 +328,7 @@ def partition_test( left_stat = aggregate_fn(left_data.to_numpy()) right_stat = aggregate_fn(right_data.to_numpy()) - split_stat = metric_fn.pairwise(left_stat.reshape(1, -1), right_stat.reshape(1, -1)) + split_stat = metric_fn.pairwise(np.atleast_2d(left_stat), np.atleast_2d(right_stat)) nx.set_node_attributes(t, {child: {f"{key_added}_value": left_stat}}) @@ -355,7 +362,10 @@ def partition_test( elif test == "t-test": _, two_sided_pval = ttest_ind( - left_data.to_numpy(), right_data.to_numpy(), axis=None, equal_var=equal_var + left_data.to_numpy(), + right_data.to_numpy(), + axis=None, # type: ignore + equal_var=equal_var, ) nx.set_edge_attributes(t, {(parent, child): {f"{key_added}_pval": two_sided_pval}}) diff --git a/src/pycea/utils.py b/src/pycea/utils.py index f87ea05..7bb2207 100755 --- a/src/pycea/utils.py +++ b/src/pycea/utils.py @@ -230,11 +230,15 @@ def get_keyed_obs_data( data = cast(pd.DataFrame, data) if array_keys and len(set(data.dtypes)) > 1: raise ValueError("Cannot use arrays with mixed dtypes.") - if array_keys and data.iloc[:, 0].dtype.kind in ["b", "O", "S"]: - long_data = data.values.ravel() - categories = _get_categories(pd.Series(long_data), sort) - categorical_type = pd.CategoricalDtype(categories=categories, ordered=True) - data = data.apply(lambda col: col.astype(categorical_type)) + if array_keys: + dtype = data.iloc[:, 0].dtype + if isinstance(dtype, pd.CategoricalDtype): + pass + elif dtype.kind in ["b", "O", "S"]: + long_data = data.values.ravel() + categories = _get_categories(pd.Series(long_data), sort) + categorical_type = pd.CategoricalDtype(categories=categories, ordered=True) + data = data.apply(lambda col: col.astype(categorical_type)) return data, array_keys, is_square From de6ef978375ea7a8ed065bc9bc4dae02fed829c4 Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 18:43:43 -0400 Subject: [PATCH 5/7] v0.3.0 --- CHANGELOG.md | 10 +++++++++- pyproject.toml | 2 +- 2 files changed, 10 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 43320d6..bb69794 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -11,7 +11,15 @@ and this project adheres to [Semantic Versioning][]. ## Unreleased ### Added -- `pycea.tl.parsimony` and `pycea.tl.fitch_count` for small-parsimony scoring and the FitchCount transition-count algorithm (migrated from Cassiopeia) + +### Changed + +### Fixed + +## [0.3.0] - 2026-07-08 + +### Added +- `pycea.tl.parsimony` and `pycea.tl.fitch_count` for small-parsimony scoring and the FitchCount transition-count algorithm (migrated from Cassiopeia) (#59) - `pycea.pl.ancestral_linkage` for plotting the pairwise linkage matrix as a clustered heatmap (#58) - Added fast `"sum"` method to `pycea.tl.ancestral_states` (#54, #56) - Added additional node ploting customization - `outline_width` and option to directly specify color with hex code (#53) diff --git a/pyproject.toml b/pyproject.toml index d8bb570..2aa26d7 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ requires = [ "hatchling" ] [project] name = "pycea-lineage" -version = "0.2.0" +version = "0.3.0" description = "Scverse lineage tracing toolkit" readme = "README.md" license = { file = "LICENSE" } From bd0a21e9c8baaec39bd26923e71151f865c2f65e Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 19:01:41 -0400 Subject: [PATCH 6/7] docs: render overloaded functions as a single entry Set typehints_document_overloads = False so sphinx-autodoc-typehints no longer emits every @typing.overload signature. This restores the short description in the API summary table for overloaded functions (parsimony, fitch_count, autocorr, clades, distance, ...) and shows a single implementation signature on each function's page instead of the copy-only overload stubs. Co-Authored-By: Claude Opus 4.8 (1M context) --- docs/conf.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/docs/conf.py b/docs/conf.py index 4a96c3d..fd25dc5 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -87,6 +87,11 @@ nb_execution_mode = "off" nb_merge_streams = True typehints_defaults = "braces" +# Render @typing.overload'd functions as a single entry (the implementation). +# Otherwise sphinx-autodoc-typehints emits every overload signature, which both +# drops the short description from the API summary table and clutters each +# function's page with the (copy-only) overload stubs. +typehints_document_overloads = False source_suffix = { ".rst": "restructuredtext", From f6e90c0c73dda1a48ca05da143c53928778ca33b Mon Sep 17 00:00:00 2001 From: colganwi Date: Wed, 8 Jul 2026 19:09:34 -0400 Subject: [PATCH 7/7] Address PR review: NaN wildcards in fitch_count, callable metrics in partition_test - fitch_count: coerce pandas NaN leaf values to the missing sentinel so missing observations are treated as Fitch wildcards instead of raising KeyError when missing_state is left at the default. - partition_test: wrap a user-supplied callable metric in a _CallableMetric adapter that provides the .pairwise API, so callable metrics no longer fail with AttributeError. - Add tests covering both paths. Co-Authored-By: Claude Opus 4.8 (1M context) --- src/pycea/tl/_metrics.py | 13 +++++++++++++ src/pycea/tl/parsimony.py | 9 +++++++-- src/pycea/tl/partition_test.py | 6 +++--- tests/test_parsimony.py | 10 ++++++++++ tests/test_partition_test.py | 19 +++++++++++++++++++ 5 files changed, 52 insertions(+), 5 deletions(-) diff --git a/src/pycea/tl/_metrics.py b/src/pycea/tl/_metrics.py index d26883a..d1d9c3f 100755 --- a/src/pycea/tl/_metrics.py +++ b/src/pycea/tl/_metrics.py @@ -40,6 +40,19 @@ def __call__(self, a, b): def pairwise(self, X, Y): return pairwise_distances(X, Y, metric=self.__call__) + +class _CallableMetric: + """Adapts a plain ``metric(a, b) -> float`` callable to the ``.pairwise`` API.""" + + def __init__(self, func: _MetricFn): + self.func = func + + def __call__(self, a, b): + return self.func(a, b) + + def pairwise(self, X, Y): + return pairwise_distances(X, Y, metric=self.func) + def _lca_distance(tree, depth_key, node1, node2, lca): """Compute the lca distance between two nodes in a tree.""" if node1 == node2: diff --git a/src/pycea/tl/parsimony.py b/src/pycea/tl/parsimony.py index b82f255..adddbbf 100755 --- a/src/pycea/tl/parsimony.py +++ b/src/pycea/tl/parsimony.py @@ -228,8 +228,13 @@ def fitch_count( if root is not None: g = nx.DiGraph(g.subgraph(nx.descendants(g, root) | {root}).copy()) actual_root = root if root is not None else get_root(g) - # Set leaf states then compute Fitch-Hartigan state sets - leaf_states = {node: values[node] for node in g.nodes if g.out_degree(node) == 0 and node in values.index} + # Set leaf states then compute Fitch-Hartigan state sets. Pandas NaNs are + # coerced to the missing sentinel so missing leaves are treated as wildcards. + leaf_states = {} + for node in g.nodes: + if g.out_degree(node) == 0 and node in values.index: + val = values[node] + leaf_states[node] = missing_state if pd.isna(val) else val nx.set_node_attributes(g, leaf_states, key) _fitch_hartigan_downpass(g, key, missing_state, set_key="_fitch_set") # Treat missing (wildcard) sets as the full state space diff --git a/src/pycea/tl/partition_test.py b/src/pycea/tl/partition_test.py index c437866..a93c9a9 100644 --- a/src/pycea/tl/partition_test.py +++ b/src/pycea/tl/partition_test.py @@ -14,7 +14,7 @@ from pycea.utils import _check_tree_overlap, _get_descendant_leaves, get_keyed_obs_data, get_trees from ._aggregators import _Aggregator, _AggregatorFn, _get_aggregator -from ._metrics import MeanDiffMetric, _Metric, _MetricFn +from ._metrics import MeanDiffMetric, _CallableMetric, _Metric, _MetricFn from ._utils import _set_random_state @@ -22,7 +22,7 @@ def _run_permutations( data: pd.DataFrame, n_permutations: int, aggregate_fn: _AggregatorFn, - metric_fn: MeanDiffMetric | _MetricFn | DistanceMetric, + metric_fn: MeanDiffMetric | _CallableMetric | DistanceMetric, n_right: int, n_left: int, ) -> np.ndarray: @@ -255,7 +255,7 @@ def partition_test( if metric == "mean_difference": metric_fn = MeanDiffMetric() elif callable(metric): - metric_fn = metric + metric_fn = _CallableMetric(metric) else: metric_fn = DistanceMetric.get_metric(metric, **(metric_kwds or {})) diff --git a/tests/test_parsimony.py b/tests/test_parsimony.py index 5866f3e..daa639b 100644 --- a/tests/test_parsimony.py +++ b/tests/test_parsimony.py @@ -97,5 +97,15 @@ def test_fitch_count_root_requires_single_tree(tdata): fitch_count(tdata, "clone", root="B") +def test_fitch_count_nan_treated_as_wildcard(tdata): + """A pandas NaN leaf (default missing_state) is treated as a wildcard, not a state.""" + import numpy as np + + tdata.obs["clone"] = ["x", "x", np.nan, "x", "y"] # leaf D is missing + M = fitch_count(tdata, "clone", tree="tree1", copy=True) + assert set(M.index) == {"x", "y"} # NaN is not a state + assert np.isfinite(M.values).all() + + if __name__ == "__main__": pytest.main(["-v", __file__]) diff --git a/tests/test_partition_test.py b/tests/test_partition_test.py index 60690a9..3718433 100644 --- a/tests/test_partition_test.py +++ b/tests/test_partition_test.py @@ -264,5 +264,24 @@ def test_nonbinary_negative_control_middle_vs_rest_p_near_one(): assert pytest.approx(pval, rel=0, abs=1e-12) == 1.0 +def test_callable_metric_matches_builtin(): + """A custom callable metric runs (has a working .pairwise) and matches mean_difference.""" + G = _three_way_root_tree(k_per_child=100) + obs = _obs_for_three_way(G, left_val=1.0, middle_val=1.0, right_val=0.0) + tdata = td.TreeData(obs=obs, obst={"tri": G}) + + def mean_diff(a, b): + return np.mean(a - b) + + builtin = partition_test(tdata, keys="x", aggregate="mean", metric="mean_difference", + test="permutation", n_permutations=50, tree="tri", + random_state=0, copy=True) + custom = partition_test(tdata, keys="x", aggregate="mean", metric=mean_diff, + test="permutation", n_permutations=50, tree="tri", + random_state=0, copy=True) + assert custom is not None + pd.testing.assert_frame_equal(builtin, custom) + + if __name__ == "__main__": pytest.main(["-v", __file__])