From 4a4bdc5fbf7d88a8bdf118dab0e7de675a8f26d3 Mon Sep 17 00:00:00 2001 From: Dominik Date: Sat, 10 Oct 2026 01:18:58 -0700 Subject: [PATCH 1/2] find_landmarks: build-only kNN graph by default, exact and order-invariant options build_graph built a pynndescent index and then queried it for every training point. On large data with low intrinsic dimension and strongly varying density that query is orders of magnitude slower than the build (synthetic branched curve, 19 dims, 8 CPUs: 1.5 s vs 217 s at 200k cells, 10 s vs > 44 min at 1.1M). The default graph is now the index's own neighbor graph, with no re-query; it is also slightly more accurate. - knn_method="nndescent" (default): build-only pynndescent graph. - knn_method="exact": scipy cKDTree kNN; deterministic and independent of row order and thread count, but slow on high-intrinsic-dimension data. - knn_method="nndescent_query": the 0.9.0 graph, to reproduce landmarks. - find_landmarks(order_invariant=True): canonical row sort + exact graph, indices mapped back; requires an int random_state. - Edge list built without a Python loop; returned as an (n_edges, 2) array. - First Leiden resolution clamped to the search bounds [0.01, 1000]. Changes which landmarks are selected relative to 0.9.0; see CHANGELOG. --- CHANGELOG.md | 96 +++++++++++ kompot/utils.py | 265 ++++++++++++++++++++++++----- tests/test_find_landmarks_graph.py | 256 ++++++++++++++++++++++++++++ tests/test_landmark_snap.py | 6 +- 4 files changed, 581 insertions(+), 42 deletions(-) create mode 100644 tests/test_find_landmarks_graph.py diff --git a/CHANGELOG.md b/CHANGELOG.md index ba78d57..14b3242 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,102 @@ All notable changes to this project will be documented in this file. +## [Unreleased] + +### Fixed — `find_landmarks` no longer stalls on large data with low intrinsic dimension + +`find_landmarks` joins the cells into a k-nearest-neighbor graph before +clustering. Through 0.9.0, `build_graph` built a pynndescent index and then +queried that index for every training point, although the build had already +produced a neighbor graph. On data with low intrinsic dimension and strongly +varying density, the kind of geometry diffusion components and other +trajectory-like embeddings have, that query takes orders of magnitude longer +than the build, and its cost grows much faster than the number of cells. + +The default graph is now the one the pynndescent build produces +(`knn_method="nndescent"`), with no re-query. Measured on synthetic data, 19 +dimensions, 8 CPUs, k = 15, one run per cell: + +| data | cells | build only (new default) | build + query (0.9.0) | exact kd-tree | +|---|---|---|---|---| +| branched curve, varying density | 200 000 | 1.5 s | 217 s | 5.9 s | +| branched curve, varying density | 500 000 | 4.7 s | 1 133 s | 20 s | +| branched curve, varying density | 1 100 000 | 10 s | > 44 min (stopped at a 45-min cap) | 130 s | +| isotropic Gaussian mixture | 200 000 | 3.2 s | 6.6 s | 32 s | +| isotropic Gaussian mixture | 500 000 | 5.1 s | 22 s | 102 s | +| isotropic Gaussian mixture | 1 100 000 | 12 s | 89 s | 837 s | + +The curve data are five 1-d curves embedded nonlinearly in 19 dimensions, with +the cells concentrated toward one end of each (`t = u**4`) and little noise: +low intrinsic dimension with strongly varying density. The mixture has 30 +isotropic Gaussian components. Duplicate rows (2% or 20%) and 100x tighter +clusters did not slow the query; varying density on a low-dimensional curve +did. The 0.9.0 column times the pynndescent build and query directly, except +at 1 100 000 mixture cells, where it is the 0.9.0 `build_graph` end to end +(including its Python edge loop). The runs used several cluster nodes and one run per cell, so treat the +figures as orders of magnitude (the same exact search took 18 s and 32 s at +200 000 mixture cells on two nodes). The default graph is also slightly more +accurate than the 0.9.0 query: recall against the exact neighbors at 200 000 +cells was 0.956 against 0.940 on the mixture and 0.992 against 0.990 on the +curve. + +**This changes which landmarks are selected relative to 0.9.0.** Leiden is +sensitive to small perturbations of its graph, so a graph that differs in a +minority of its edges gives a largely different landmark set. On the +200 000-cell mixture above with `n_clusters=200` and `random_state=0`, the new +default and the 0.9.0 graph shared 65 landmarks (of 178 and 175); for scale, +changing the seed to 1 under the new default left 61 shared (of 178 and 192), +while rerunning with the same seed reproduced the landmarks exactly. Every +landmark is still the cell nearest its cluster's centroid; only the choice +among comparably good landmark sets moves. Unseeded +calls (`random_state=None`, the default) were not reproducible before and are +not now. Kompot's own `da()`, `de()` and the differential classes choose +landmarks with mellon, not `find_landmarks`, and are not affected. + +**To reproduce landmarks from 0.9.0**, pass +`find_landmarks(..., knn_method="nndescent_query")` with the same +`random_state`. To reproduce landmarks from 0.8.0, also pass +`exact_snap=False`. One exception for both: with more than 1000 cells per +requested landmark (for example more than 200 000 cells at the default +`n_clusters=200`) the resolution search now starts from a different point (see +below), so the old landmarks are not reproduced exactly there. + +`build_graph` now returns its edges as an integer array of shape +`(n_edges, 2)` instead of a list of tuples, built without a Python loop. + +### Added — `knn_method="exact"` and `find_landmarks(..., order_invariant=True)` + +`knn_method="exact"` builds the graph with an exact Euclidean kd-tree search +(`scipy.spatial.cKDTree`). Unlike the pynndescent graph, which changes with the +number of numba threads, it depends only on the data. It is fast on data with +low intrinsic dimension and slow on data without it (see the table above); it +is not the default for that reason. `exact_snap=False` cannot be combined with +it, because it builds no index to snap with. + +The landmarks also depended on the order of the rows of `X`, through the +graph and through Leiden, which visits vertices in index order; an exact graph +alone does not remove the second. With `order_invariant=True` the rows are +sorted into a canonical (lexicographic) order before the graph and the +clustering, the exact graph is used, and the landmark indices are mapped back +to the caller's order. The same cells in any order, with the same seed, then +give the same landmarks on any machine. It requires an int `random_state` +(without one Leiden is unseeded and the landmarks change from run to run +anyway) and raises a `ValueError` without one, or with a `knn_method` other +than `"exact"`. The sort and the copy of `X` it needs cost seconds, against the +exact search's minutes. + +This buys reproducibility, not stability: changing `random_state` moves the +landmarks about as much as permuting the rows does. Exact duplicate rows are +interchangeable under the sort, so a landmark on a duplicated row keeps its +coordinates in every order but may point at a different copy. + +### Fixed — the first Leiden resolution stays inside the search bounds + +The resolution search started at `n_cells / n_clusters`, which exceeds its own +upper bound of 1000 once there are more than 1000 cells per requested cluster. +If that first step gave too few clusters, the bisection then ran with its lower +bound above its upper one. The start is now clamped to `[0.01, 1000]`. + ## [0.9.0] - 2026-09-15 ### Fixed — sample variance no longer costs more memory than it needs to diff --git a/kompot/utils.py b/kompot/utils.py index c255fa4..a3a2b26 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -1,5 +1,7 @@ """Utility functions for Kompot package.""" +import os + import numpy as np import jax import jax.numpy as jnp @@ -144,47 +146,143 @@ def __getitem__(self, key): # The functions validate_field_run_id and get_run_from_history have been moved to anndata.utils +_KNN_METHODS = ("nndescent", "exact", "nndescent_query") + + +def _available_cpus() -> int: + """Number of CPUs this process may run on (its affinity mask, not the machine).""" + try: + return max(1, len(os.sched_getaffinity(0))) + except (AttributeError, OSError): + return max(1, os.cpu_count() or 1) + + +def _knn_edges(indices: np.ndarray) -> np.ndarray: + """ + Directed kNN edges ``(i, j)`` for every ``j`` in row ``i`` of ``indices``, without self-loops. + + Negative entries, which ``pynndescent`` uses to pad a row when it found + fewer neighbors than requested, are dropped. + + The edges come out in row-major order, the order a loop over the rows and + then over each row's neighbors would produce. + + Parameters + ---------- + indices : np.ndarray + Integer array of shape (n_samples, k); row ``i`` lists the neighbors of + sample ``i``. + + Returns + ------- + np.ndarray + Integer array of shape (n_edges, 2). + """ + indices = np.asarray(indices) + n_obs, k = indices.shape + sources = np.repeat(np.arange(n_obs, dtype=np.int64), k) + targets = indices.reshape(-1).astype(np.int64, copy=False) + keep = (sources != targets) & (targets >= 0) + return np.column_stack((sources[keep], targets[keep])) + + +def _exact_knn(X: np.ndarray, n_neighbors: int) -> np.ndarray: + """ + Exact Euclidean ``n_neighbors`` nearest rows of ``X`` for every row of ``X``, itself included. + + A kd-tree search (``scipy.spatial.cKDTree``). The result depends only on + ``X``: it does not depend on a seed or on the number of threads, because + each row's query is answered independently. + + Returns + ------- + np.ndarray + Integer array of shape (n_samples, n_neighbors). + """ + from scipy.spatial import cKDTree + + X = np.asarray(X, dtype=np.float64) + tree = cKDTree(X) + _, indices = tree.query(X, k=n_neighbors, workers=_available_cpus()) + return np.asarray(indices).reshape(X.shape[0], n_neighbors) + + def build_graph( - X: np.ndarray, n_neighbors: int = 15, random_state: Optional[int] = 42 -) -> Tuple[List[Tuple[int, int]], pynndescent.NNDescent]: + X: np.ndarray, + n_neighbors: int = 15, + random_state: Optional[int] = 42, + knn_method: str = "nndescent", +) -> Tuple[np.ndarray, Optional[pynndescent.NNDescent]]: """ - Build a graph from a dataset using approximate nearest neighbors. + Build a k-nearest-neighbor graph from a dataset. Parameters ---------- X : np.ndarray Data matrix of shape (n_samples, n_features). n_neighbors : int, optional - Number of neighbors for graph construction, by default 15. + Number of neighbors per sample, the sample itself included, by default 15. + The exact search uses at most ``n_samples``. random_state : int or None, optional - Seed passed to ``pynndescent.NNDescent`` for the approximate - nearest-neighbor construction, by default 42. Passing ``None`` lets - pynndescent draw its own entropy (non-reproducible). + Seed passed to ``pynndescent.NNDescent``, by default 42. Not used by + ``knn_method="exact"``, which has no randomness. With ``None`` + pynndescent draws its own entropy (non-reproducible). + knn_method : {"nndescent", "exact", "nndescent_query"}, optional + How the neighbors are found, by default ``"nndescent"``. + + - ``"nndescent"``: the approximate graph that ``pynndescent.NNDescent`` + builds (its ``neighbor_graph``). Fast on all data we measured. It + depends on the row order and on the number of numba threads. + - ``"exact"``: an exact Euclidean kd-tree search + (``scipy.spatial.cKDTree``). Deterministic, and independent of row + order and thread count. Fast when the data have low intrinsic + dimension; much slower than ``"nndescent"`` when they do not (on a + synthetic isotropic Gaussian mixture of 1.1M points in 19 + dimensions, 8 CPUs: 837 s against 11.5 s). + - ``"nndescent_query"``: what kompot 0.9.0 and earlier did: build the + ``pynndescent`` index, then query it for every training point. The + query can be orders of magnitude slower than the build, on data + with low intrinsic dimension and strongly varying density. Use it + only to reproduce landmarks from earlier versions. Returns ------- - Tuple[List[Tuple[int, int]], pynndescent.NNDescent] + Tuple[np.ndarray, pynndescent.NNDescent or None] A tuple containing: - - edges: List of (source, target) tuples defining the graph - - index: The nearest neighbor index for future queries + - edges: integer array of shape (n_edges, 2) of directed + (source, target) pairs, self-loops removed + - index: the ``pynndescent`` index, or None with ``knn_method="exact"`` """ - # Build the nearest neighbor index - index = pynndescent.NNDescent( - X, n_neighbors=n_neighbors, random_state=random_state + if knn_method not in _KNN_METHODS: + raise ValueError( + f"knn_method must be one of {_KNN_METHODS}, got {knn_method!r}" + ) + if n_neighbors < 1: + raise ValueError(f"n_neighbors must be at least 1, got {n_neighbors}") + logger.info( + f"Building the {n_neighbors}-nearest-neighbor graph (knn_method={knn_method!r})" ) - # Query for nearest neighbors - indices, _ = index.query(X, k=n_neighbors) - - # Convert to edges - n_obs = X.shape[0] - edges = [] - for i in range(n_obs): - for j in indices[i]: - if i != j: # Avoid self-loops - edges.append((i, j)) + index = None + if knn_method == "exact": + n_obs = X.shape[0] + if n_neighbors > n_obs: + logger.warning( + f"n_neighbors={n_neighbors} exceeds the number of samples " + f"({n_obs}); using {n_obs}." + ) + n_neighbors = n_obs + indices = _exact_knn(X, n_neighbors) + else: + index = pynndescent.NNDescent( + X, n_neighbors=n_neighbors, random_state=random_state + ) + if knn_method == "nndescent": + indices = index.neighbor_graph[0] + else: + indices, _ = index.query(X, k=n_neighbors) - return edges, index + return _knn_edges(indices), index def compute_mahalanobis_distances( @@ -627,7 +725,7 @@ def compute_mahalanobis_distance( def find_optimal_resolution( - edges: List[Tuple[int, int]], + edges: Union[np.ndarray, List[Tuple[int, int]]], n_obs: int, n_clusters: int, tol: float = 0.1, @@ -639,8 +737,9 @@ def find_optimal_resolution( Parameters ---------- - edges : List[Tuple[int, int]] - List of edges defining the graph. + edges : np.ndarray or List[Tuple[int, int]] + Edges defining the graph: an integer array of shape (n_edges, 2), as + ``build_graph`` returns, or a list of (source, target) pairs. n_obs : int Number of observations (nodes) in the graph. n_clusters : int @@ -669,10 +768,13 @@ def find_optimal_resolution( G_igraph = ig.Graph(edges=edges, directed=False) G_igraph.vs["name"] = [str(i) for i in range(n_obs)] - # Initial heuristic - initial_resolution = n_obs / n_clusters - resolution = initial_resolution + # Initial heuristic, kept inside the search bounds: n_obs / n_clusters + # exceeds the upper bound once there are more than 1000 samples per + # requested cluster, and a start outside [lower, upper] breaks the + # bisection below (a too-coarse first step would set lower above upper). lower, upper = 0.01, 1000.0 + initial_resolution = n_obs / n_clusters + resolution = min(max(initial_resolution, lower), upper) best_partition = None @@ -842,10 +944,16 @@ def find_landmarks( max_iter: int = 10, random_state: Optional[int] = None, exact_snap: bool = True, + knn_method: Optional[str] = None, + order_invariant: bool = False, ) -> Tuple[np.ndarray, np.ndarray]: """ Identify landmark points representing clusters in the dataset. + The cells are joined into a k-nearest-neighbor graph, the graph is + clustered with Leiden at a resolution searched to give about + ``n_clusters`` clusters, and each cluster centroid is snapped to a cell. + Parameters ---------- X : np.ndarray @@ -863,12 +971,13 @@ def find_landmarks( community detection underlying landmark discovery draws from igraph's global random number generator, which is otherwise left unseeded, so the returned landmarks vary run-to-run even for identical input. - Passing an int seeds both the nearest-neighbor construction and the - Leiden step, making ``find_landmarks`` fully reproducible through the - public API (same ``X`` + same ``random_state`` yields identical - landmark indices and coordinates). The default (``None``) preserves the - historical non-deterministic behavior, so existing callers are - unaffected. + Passing an int seeds the Leiden step and the ``pynndescent`` + nearest-neighbor construction, making ``find_landmarks`` reproducible + through the public API: the same ``X`` and the same ``random_state`` + yield identical landmark indices and coordinates, on the same machine + with the same number of threads (with ``knn_method="exact"`` on any + machine). The default (``None``) preserves the historical + non-deterministic behavior. Required when ``order_invariant=True``. exact_snap : bool, optional How each cluster centroid is snapped to a cell, by default True. With True the landmark is the cell exactly nearest the centroid (Euclidean, @@ -877,7 +986,37 @@ def find_landmarks( nearest-neighbor index built for the graph, as kompot 0.8.0 and earlier did; on large, high-dimensional data that index can return a nearby cell instead of the nearest one (24% of centroids on a synthetic - 1M-cell, 100-dimensional mixture). Set False only to reproduce landmarks from earlier versions. + 1M-cell, 100-dimensional mixture). Set False only to reproduce + landmarks from earlier versions (with ``knn_method="nndescent_query"``); + it cannot be combined with ``knn_method="exact"``, which builds no + index. + knn_method : {"nndescent", "exact", "nndescent_query"} or None, optional + How the nearest-neighbor graph is built. By default (None), + ``"exact"`` when ``order_invariant=True`` and ``"nndescent"`` + otherwise. ``"nndescent"`` is the approximate graph ``pynndescent`` builds, which is fast on all data + we measured. ``"exact"`` is an exact kd-tree search: deterministic and + independent of row order and thread count, but much slower on data + with high intrinsic dimension. ``"nndescent_query"`` is the graph of + kompot 0.9.0 and earlier, which can be orders of magnitude slower on + large data with low intrinsic dimension; pass it only to reproduce + landmarks from earlier versions. See :func:`build_graph` for timings. + order_invariant : bool, optional + Make the landmarks independent of the order of the rows of ``X``, by + default False. The rows are put into a canonical order (sorted + lexicographically by their coordinates) before the graph and the + clustering are computed, and the landmark indices are mapped back to + the original row order. The same rows in any order, with the same + ``random_state``, then give the same landmarks, on any machine and + with any number of threads. Requires an int ``random_state`` and uses + the exact graph; a ``ValueError`` is raised without a seed or with + another ``knn_method``. Costs one sort and one copy of ``X``. + + This buys reproducibility, not stability. The landmark set is one draw + from a procedure that is sensitive to every arbitrary choice: changing + ``random_state`` moves the landmarks about as much as permuting the rows + does. Exact duplicate rows are interchangeable under the sort, so a + landmark that falls on a duplicated row has the same coordinates in + every order but may point at a different copy. Returns ------- @@ -886,15 +1025,54 @@ def find_landmarks( - landmarks: Matrix of shape (n_clusters, n_features) containing landmark coordinates - landmark_indices: Indices of landmarks in the original dataset """ + if knn_method is None: + knn_method = "exact" if order_invariant else "nndescent" + if knn_method not in _KNN_METHODS: + raise ValueError( + f"knn_method must be one of {_KNN_METHODS}, got {knn_method!r}" + ) + if not exact_snap and knn_method == "exact": + raise ValueError( + "exact_snap=False snaps with the pynndescent index, which " + "knn_method='exact' does not build. To reproduce landmarks from " + "kompot 0.8.0, pass exact_snap=False together with " + "knn_method='nndescent_query'." + ) + if order_invariant: + if random_state is None: + raise ValueError( + "order_invariant=True requires an int random_state: without a " + "seed the Leiden step is unseeded and the landmarks change from " + "run to run whatever the row order." + ) + if knn_method != "exact": + raise ValueError( + "order_invariant=True requires knn_method='exact': the " + "pynndescent graph depends on the number of threads, so the " + "landmarks would not be reproducible across machines." + ) + + order = None + if order_invariant: + X_input = X + # Canonical row order: sort by the first column, ties by the second, + # and so on. np.lexsort treats its LAST key as the primary one. + X = np.asarray(X) + order = np.lexsort(X.T[::-1]) + X = X[order] + # When no seed is requested, keep the historical nearest-neighbor seed (42) - # so unseeded callers observe unchanged NN behavior; the run-to-run variation - # comes solely from the unseeded Leiden step below. A supplied seed threads - # into both stages for end-to-end reproducibility. + # so unseeded callers observe unchanged pynndescent behavior; the + # run-to-run variation comes solely from the unseeded Leiden step below. + # The exact search ignores the seed. nn_random_state = 42 if random_state is None else random_state # Build graph edges, index = build_graph( - X, n_neighbors=n_neighbors, random_state=nn_random_state + X, + n_neighbors=n_neighbors, + random_state=nn_random_state, + knn_method=knn_method, ) n_obs = X.shape[0] @@ -902,6 +1080,7 @@ def find_landmarks( optimal_resolution, partition = find_optimal_resolution( edges, n_obs, n_clusters, tol=tol, max_iter=max_iter, random_state=random_state ) + del edges clusters = np.array(partition.membership) cluster_ids = np.unique(clusters) @@ -914,6 +1093,10 @@ def find_landmarks( else: landmark_indices, _ = index.query(centroids, k=1) landmark_indices = landmark_indices.flatten() + + if order is not None: + landmark_indices = order[landmark_indices] + X = X_input landmarks = X[landmark_indices] logger.info( diff --git a/tests/test_find_landmarks_graph.py b/tests/test_find_landmarks_graph.py new file mode 100644 index 0000000..426df10 --- /dev/null +++ b/tests/test_find_landmarks_graph.py @@ -0,0 +1,256 @@ +"""The kNN graph, resolution search and row-order handling behind ``find_landmarks``. + +Through 0.9.0 ``build_graph`` built a pynndescent index and then queried it +for every training point. On large data with low intrinsic dimension and +varying density that query takes orders of magnitude longer than the build. +The default graph is now the one the build produces, ``knn_method="exact"`` +gives a deterministic kd-tree kNN, ``order_invariant=True`` makes the +landmarks independent of row order, the edge list is built without a Python +loop, and the first Leiden resolution is kept inside the search bounds. +""" + +import numpy as np +import pynndescent +import pytest + +import kompot.utils as kutils +from kompot.utils import build_graph, find_landmarks, find_optimal_resolution + + +def _mixture(n, d, k, seed): + rng = np.random.default_rng(seed) + centres = rng.normal(size=(k, d)) * 4.0 + labels = rng.integers(0, k, n) + return centres[labels] + rng.normal(size=(n, d)) + + +def _brute_force_knn(X, k): + """Reference: full float64 distance matrix, k smallest per row in distance order.""" + sq = ((X[:, None, :] - X[None, :, :]) ** 2).sum(axis=2) + return np.argsort(sq, axis=1, kind="stable")[:, :k] + + +def _loop_edges(indices): + """The pre-0.10 edge construction, kept verbatim as the reference.""" + edges = [] + for i in range(indices.shape[0]): + for j in indices[i]: + if i != j: + edges.append((i, j)) + return edges + + +class TestExactGraph: + def test_exact_graph_matches_brute_force(self): + """``knn_method="exact"`` equals a brute-force kNN, edge for edge. + + 3 000 points in 30 dimensions with k = 15: here the pynndescent graph + misses some true neighbors, so this fails on an approximate graph. + """ + rng = np.random.default_rng(0) + X = rng.normal(size=(3000, 30)) + k = 15 + + edges, index = build_graph(X, n_neighbors=k, knn_method="exact") + + expected = np.asarray(_loop_edges(_brute_force_knn(X, k))) + np.testing.assert_array_equal(np.asarray(edges), expected) + assert index is None + + def test_exact_graph_does_not_depend_on_thread_count(self, monkeypatch): + X = _mixture(4000, 12, 10, seed=1) + monkeypatch.setattr(kutils, "_available_cpus", lambda: 1) + one, _ = build_graph(X, n_neighbors=15, knn_method="exact") + monkeypatch.setattr(kutils, "_available_cpus", lambda: 4) + four, _ = build_graph(X, n_neighbors=15, knn_method="exact") + np.testing.assert_array_equal(one, four) + + def test_default_graph_is_the_nndescent_construction_graph(self): + """By default the graph is the index's own neighbor graph, with no re-query. + + On 3 000 points in 30 dimensions the construction graph and a query of + the index differ, so this fails on the 0.9.0 re-query. + """ + X = np.random.default_rng(0).normal(size=(3000, 30)) + + edges, index = build_graph(X, n_neighbors=15, random_state=42) + + ref = pynndescent.NNDescent(X, n_neighbors=15, random_state=42) + assert isinstance(index, pynndescent.NNDescent) + np.testing.assert_array_equal( + np.asarray(edges), np.asarray(_loop_edges(ref.neighbor_graph[0])) + ) + + def test_nndescent_query_reproduces_the_old_graph(self): + """``knn_method="nndescent_query"`` returns exactly the graph 0.9.0 built.""" + X = _mixture(2000, 10, 8, seed=2) + + edges, index = build_graph( + X, n_neighbors=15, random_state=42, knn_method="nndescent_query" + ) + + old_index = pynndescent.NNDescent(X, n_neighbors=15, random_state=42) + old_indices, _ = old_index.query(X, k=15) + assert isinstance(index, pynndescent.NNDescent) + np.testing.assert_array_equal( + np.asarray(edges), np.asarray(_loop_edges(old_indices)) + ) + + def test_rejects_unknown_method_and_caps_neighbors(self): + X = np.random.default_rng(3).normal(size=(10, 3)) + with pytest.raises(ValueError, match="knn_method"): + build_graph(X, knn_method="annoy") + with pytest.raises(ValueError, match="n_neighbors"): + build_graph(X, n_neighbors=0) + # More neighbors than samples: every other sample is a neighbor. + edges, _ = build_graph(X, n_neighbors=15, knn_method="exact") + assert len(edges) == 10 * 9 + + +class TestVectorisedEdges: + def test_edges_equal_the_loop(self): + """Same edges, same order, as the loop, including rows without a self entry.""" + rng = np.random.default_rng(4) + n, k = 500, 7 + indices = rng.integers(0, n, size=(n, k)) + # Self at the front of most rows, somewhere else in some, absent in a few. + indices[:400, 0] = np.arange(400) + indices[400:450, 3] = np.arange(400, 450) + + got = kutils._knn_edges(indices) + + expected = np.asarray(_loop_edges(indices)) + assert got.shape == expected.shape + np.testing.assert_array_equal(got, expected) + + def test_padding_entries_are_dropped(self): + """pynndescent pads rows with -1 when it finds fewer neighbors than asked for.""" + indices = np.array([[0, 1, -1], [1, 0, -1], [2, -1, -1]]) + got = kutils._knn_edges(indices) + np.testing.assert_array_equal(got, [[0, 1], [1, 0]]) + + +def _recording_leiden(real, resolutions): + """Wrap ``community_leiden`` so that it records the resolutions it is asked for.""" + + def leiden(graph, *args, **kwargs): + resolutions.append(kwargs["resolution"]) + return real(graph, *args, **kwargs) + + return leiden + + +class TestResolutionClamp: + @staticmethod + def _ring(n): + src = np.arange(n) + return np.column_stack((src, (src + 1) % n)) + + @pytest.mark.parametrize( + "n_obs, n_clusters, expected_first", + [ + (5000, 2, 1000.0), # n_obs / n_clusters = 2500 > upper bound + (10, 2000, 0.01), # n_obs / n_clusters = 0.005 < lower bound + (1000, 100, 10.0), # inside the bounds: unchanged + ], + ) + def test_first_resolution_is_inside_the_bounds( + self, monkeypatch, n_obs, n_clusters, expected_first + ): + import igraph as ig + + resolutions = [] + monkeypatch.setattr( + ig.Graph, + "community_leiden", + _recording_leiden(ig.Graph.community_leiden, resolutions), + ) + + find_optimal_resolution( + self._ring(n_obs), n_obs, n_clusters, max_iter=1, random_state=0 + ) + + assert resolutions[0] == pytest.approx(expected_first) + + +class TestOrderInvariant: + def test_landmarks_do_not_depend_on_row_order(self): + """Three row permutations give the same landmarks, as rows of the original.""" + X = _mixture(3000, 8, 12, seed=5) + ref_landmarks, ref_idx = find_landmarks( + X, n_clusters=40, random_state=0, order_invariant=True + ) + rng = np.random.default_rng(6) + for _ in range(3): + perm = rng.permutation(X.shape[0]) + landmarks, idx = find_landmarks( + X[perm], n_clusters=40, random_state=0, order_invariant=True + ) + # idx indexes X[perm]; perm[idx] are the same rows of X. + assert sorted(perm[idx].tolist()) == sorted(ref_idx.tolist()) + np.testing.assert_array_equal(landmarks, X[perm][idx]) + + def test_indices_point_into_the_input(self): + """The returned indices refer to the caller's row order, not the sorted one. + + On input already in the canonical order the sort is the identity, so + that run's landmarks are the reference; a shuffled copy must give the + same landmark coordinates, found at the returned rows of the shuffled + input. + """ + X = _mixture(1500, 5, 6, seed=7) + X_sorted = X[np.lexsort(X.T[::-1])] + X_shuffled = X_sorted[np.random.default_rng(8).permutation(len(X_sorted))] + + ref_landmarks, _ = find_landmarks( + X_sorted, n_clusters=15, random_state=0, order_invariant=True + ) + landmarks, idx = find_landmarks( + X_shuffled, n_clusters=15, random_state=0, order_invariant=True + ) + + np.testing.assert_array_equal(landmarks, X_shuffled[idx]) + + def rows(a): + return a[np.lexsort(a.T[::-1])] + + np.testing.assert_array_equal(rows(landmarks), rows(ref_landmarks)) + + def test_requires_a_seed(self): + X = _mixture(200, 4, 3, seed=8) + with pytest.raises(ValueError, match="random_state"): + find_landmarks(X, n_clusters=5, order_invariant=True) + + def test_uses_the_exact_graph(self, monkeypatch): + """``order_invariant`` selects the exact graph and refuses any other.""" + X = _mixture(200, 4, 3, seed=8) + for method in ("nndescent", "nndescent_query"): + with pytest.raises(ValueError, match="knn_method='exact'"): + find_landmarks( + X, + n_clusters=5, + random_state=0, + order_invariant=True, + knn_method=method, + ) + + seen = [] + real = kutils.build_graph + + def spy(*args, **kwargs): + seen.append(kwargs.get("knn_method")) + return real(*args, **kwargs) + + monkeypatch.setattr(kutils, "build_graph", spy) + find_landmarks(X, n_clusters=5, random_state=0, order_invariant=True) + find_landmarks(X, n_clusters=5, random_state=0) + assert seen == ["exact", "nndescent"] + + +def test_approximate_snap_needs_a_pynndescent_index(): + """``exact_snap=False`` uses the pynndescent index, which the exact graph does not build.""" + X = _mixture(200, 4, 3, seed=9) + with pytest.raises(ValueError, match="knn_method='nndescent_query'"): + find_landmarks( + X, n_clusters=5, random_state=0, exact_snap=False, knn_method="exact" + ) diff --git a/tests/test_landmark_snap.py b/tests/test_landmark_snap.py index 92ca06b..301999a 100644 --- a/tests/test_landmark_snap.py +++ b/tests/test_landmark_snap.py @@ -91,7 +91,11 @@ def boom(*args, **kwargs): rng = np.random.default_rng(0) X = rng.normal(size=(300, 5)) landmarks, indices = find_landmarks( - X, n_clusters=10, random_state=0, exact_snap=False + X, + n_clusters=10, + random_state=0, + exact_snap=False, + knn_method="nndescent_query", ) assert len(indices) == landmarks.shape[0] np.testing.assert_array_equal(landmarks, X[indices]) From 72288b759ff9d8bcaca6178686dec0c81b386810 Mon Sep 17 00:00:00 2001 From: Dominik Date: Sat, 10 Oct 2026 10:01:49 -0700 Subject: [PATCH 2/2] find_landmarks: drop the resolution clamp; qualify the order-invariance guarantee The clamp kept the first Leiden resolution inside [0.01, 1000], guarding a bisection branch (a first step with too FEW clusters at resolution > 1000) that does not occur on kNN graphs: such a resolution always yields too many clusters. It was the only thing that stopped knn_method="nndescent_query" from reproducing 0.9.0 landmarks above 1000 cells per landmark. Without it, nndescent_query output is bit-identical to 0.9.0 (edges, landmark indices and coordinates) at 1 500 to 6 000 cells per landmark as well. order_invariant's guarantee is now stated as holding for the same numpy, scipy and python-igraph versions. --- CHANGELOG.md | 15 ++------- kompot/utils.py | 16 ++++------ tests/test_find_landmarks_graph.py | 51 +++--------------------------- 3 files changed, 14 insertions(+), 68 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 14b3242..a99c2bb 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -57,10 +57,7 @@ landmarks with mellon, not `find_landmarks`, and are not affected. **To reproduce landmarks from 0.9.0**, pass `find_landmarks(..., knn_method="nndescent_query")` with the same `random_state`. To reproduce landmarks from 0.8.0, also pass -`exact_snap=False`. One exception for both: with more than 1000 cells per -requested landmark (for example more than 200 000 cells at the default -`n_clusters=200`) the resolution search now starts from a different point (see -below), so the old landmarks are not reproduced exactly there. +`exact_snap=False`. `build_graph` now returns its edges as an integer array of shape `(n_edges, 2)` instead of a list of tuples, built without a Python loop. @@ -80,7 +77,8 @@ alone does not remove the second. With `order_invariant=True` the rows are sorted into a canonical (lexicographic) order before the graph and the clustering, the exact graph is used, and the landmark indices are mapped back to the caller's order. The same cells in any order, with the same seed, then -give the same landmarks on any machine. It requires an int `random_state` +give the same landmarks on any machine, given the same versions of numpy, +scipy and python-igraph. It requires an int `random_state` (without one Leiden is unseeded and the landmarks change from run to run anyway) and raises a `ValueError` without one, or with a `knn_method` other than `"exact"`. The sort and the copy of `X` it needs cost seconds, against the @@ -91,13 +89,6 @@ landmarks about as much as permuting the rows does. Exact duplicate rows are interchangeable under the sort, so a landmark on a duplicated row keeps its coordinates in every order but may point at a different copy. -### Fixed — the first Leiden resolution stays inside the search bounds - -The resolution search started at `n_cells / n_clusters`, which exceeds its own -upper bound of 1000 once there are more than 1000 cells per requested cluster. -If that first step gave too few clusters, the bisection then ran with its lower -bound above its upper one. The start is now clamped to `[0.01, 1000]`. - ## [0.9.0] - 2026-09-15 ### Fixed — sample variance no longer costs more memory than it needs to diff --git a/kompot/utils.py b/kompot/utils.py index a3a2b26..f29e58a 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -768,13 +768,10 @@ def find_optimal_resolution( G_igraph = ig.Graph(edges=edges, directed=False) G_igraph.vs["name"] = [str(i) for i in range(n_obs)] - # Initial heuristic, kept inside the search bounds: n_obs / n_clusters - # exceeds the upper bound once there are more than 1000 samples per - # requested cluster, and a start outside [lower, upper] breaks the - # bisection below (a too-coarse first step would set lower above upper). - lower, upper = 0.01, 1000.0 + # Initial heuristic initial_resolution = n_obs / n_clusters - resolution = min(max(initial_resolution, lower), upper) + resolution = initial_resolution + lower, upper = 0.01, 1000.0 best_partition = None @@ -976,7 +973,7 @@ def find_landmarks( through the public API: the same ``X`` and the same ``random_state`` yield identical landmark indices and coordinates, on the same machine with the same number of threads (with ``knn_method="exact"`` on any - machine). The default (``None``) preserves the historical + machine), given the same library versions. The default (``None``) preserves the historical non-deterministic behavior. Required when ``order_invariant=True``. exact_snap : bool, optional How each cluster centroid is snapped to a cell, by default True. With @@ -1006,8 +1003,9 @@ def find_landmarks( lexicographically by their coordinates) before the graph and the clustering are computed, and the landmark indices are mapped back to the original row order. The same rows in any order, with the same - ``random_state``, then give the same landmarks, on any machine and - with any number of threads. Requires an int ``random_state`` and uses + ``random_state``, then give the same landmarks on any machine and + with any number of threads, given the same versions of numpy, scipy + and python-igraph. Requires an int ``random_state`` and uses the exact graph; a ``ValueError`` is raised without a seed or with another ``knn_method``. Costs one sort and one copy of ``X``. diff --git a/tests/test_find_landmarks_graph.py b/tests/test_find_landmarks_graph.py index 426df10..10f87c8 100644 --- a/tests/test_find_landmarks_graph.py +++ b/tests/test_find_landmarks_graph.py @@ -1,12 +1,12 @@ -"""The kNN graph, resolution search and row-order handling behind ``find_landmarks``. +"""The kNN graph and row-order handling behind ``find_landmarks``. Through 0.9.0 ``build_graph`` built a pynndescent index and then queried it for every training point. On large data with low intrinsic dimension and varying density that query takes orders of magnitude longer than the build. The default graph is now the one the build produces, ``knn_method="exact"`` gives a deterministic kd-tree kNN, ``order_invariant=True`` makes the -landmarks independent of row order, the edge list is built without a Python -loop, and the first Leiden resolution is kept inside the search bounds. +landmarks independent of row order, and the edge list is built without a +Python loop. """ import numpy as np @@ -14,7 +14,7 @@ import pytest import kompot.utils as kutils -from kompot.utils import build_graph, find_landmarks, find_optimal_resolution +from kompot.utils import build_graph, find_landmarks def _mixture(n, d, k, seed): @@ -130,49 +130,6 @@ def test_padding_entries_are_dropped(self): np.testing.assert_array_equal(got, [[0, 1], [1, 0]]) -def _recording_leiden(real, resolutions): - """Wrap ``community_leiden`` so that it records the resolutions it is asked for.""" - - def leiden(graph, *args, **kwargs): - resolutions.append(kwargs["resolution"]) - return real(graph, *args, **kwargs) - - return leiden - - -class TestResolutionClamp: - @staticmethod - def _ring(n): - src = np.arange(n) - return np.column_stack((src, (src + 1) % n)) - - @pytest.mark.parametrize( - "n_obs, n_clusters, expected_first", - [ - (5000, 2, 1000.0), # n_obs / n_clusters = 2500 > upper bound - (10, 2000, 0.01), # n_obs / n_clusters = 0.005 < lower bound - (1000, 100, 10.0), # inside the bounds: unchanged - ], - ) - def test_first_resolution_is_inside_the_bounds( - self, monkeypatch, n_obs, n_clusters, expected_first - ): - import igraph as ig - - resolutions = [] - monkeypatch.setattr( - ig.Graph, - "community_leiden", - _recording_leiden(ig.Graph.community_leiden, resolutions), - ) - - find_optimal_resolution( - self._ring(n_obs), n_obs, n_clusters, max_iter=1, random_state=0 - ) - - assert resolutions[0] == pytest.approx(expected_first) - - class TestOrderInvariant: def test_landmarks_do_not_depend_on_row_order(self): """Three row permutations give the same landmarks, as rows of the original."""