diff --git a/CHANGELOG.md b/CHANGELOG.md index e91d7c9..ba78d57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,20 +2,180 @@ All notable changes to this project will be documented in this file. -## [Unreleased] +## [0.9.0] - 2026-09-15 + +### Fixed — sample variance no longer costs more memory than it needs to + +Supplying `sample_col` gives every gene its own `(n_landmarks, n_landmarks)` +covariance matrix per condition. Kompot summed those two tensors into a **third** +dense tensor before use, and then added the shared posterior covariance into it +gene by gene with `__setitem__`. Three consequences, all silent: + + - `StorageSettings(store_arrays_on_disk=True)` did not keep the tensors out of + memory. With `np.memmap` inputs the sum materialised a full + `(n_landmarks, n_landmarks, n_genes)` array in RAM — precisely the array the + flag exists to avoid. + - `compute_mahalanobis_distances` then ran `cov = jnp.array(covariance)` inside + its gene-specific branch and never read `cov`: a second full-size copy of the + tensor, for nothing. + - Under Dask the per-gene `__setitem__` rebuilt the task graph once per gene, + and the delayed per-gene tasks nested with the lazy slices they materialised, + so peak memory scaled with concurrency rather than with one gene. + +The sum is now assembled one gene at a time by `kompot.utils.LazyGeneCovariance`, +which holds references to the terms and materialises a single +`(n_landmarks, n_landmarks)` matrix when the Mahalanobis step asks for a gene. + +Measured on 1 500 cells / 150 genes / 600 landmarks (412 MiB per tensor), same +machine and same synthetic input on both trees, peak **anonymous** memory from +`/proc/self/smaps_rollup`. Every figure in this table comes from one paired +before/after run, so the columns are comparable to each other: + +| run | 0.8.0 | 0.9.0 | written | wall, 0.8.0 → 0.9.0 | +|---|---|---|---|---| +| no sample variance | 949 MiB | 972 MiB | 0 → 0 | 33.0 s → 32.6 s | +| sample variance, in memory | 2 401 MiB | 1 973 MiB | 0 → 0 | 45.7 s → 46.0 s | +| `store_arrays_on_disk`, with `dask` | 2 782 MiB | **1 153 MiB** | 0 → 0 | 231.8 s → 47.8 s | +| `store_arrays_on_disk`, no `dask` | 2 284 MiB | 1 170 MiB | 824 → 824 MiB | 48.2 s → 49.2 s | + +**With `dask` — the recommended path — offloading is now cheaper than holding +the tensors in memory**, where on 0.8.0 it was dearer, and that path is also +~4.8x faster than 0.8.0 at matched thread counts. + +**Anonymous memory alone does not describe the `dask`-less path, and this table +should not be read as if it did.** That path memory-maps the tensors, so their +pages become resident *file-backed* memory rather than private heap: the +improvement above is real on `Anonymous` and largely absent on `Rss`, where the +run costs about as much as holding the tensors in memory. What changes is the +*kind* of page — reclaimable rather than anonymous, which helps you survive a +memory squeeze — not the resident total. Under a cgroup (Slurm `--mem`, a +container) page cache is charged against the same budget anyway, so if you are +under a hard cap, install `dask`. + +The two-instrument treatment, with `Rss` and file-backed residency measured +side by side, is in +[Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html#disk-offload). +It is deliberately not duplicated here: mixing instruments from separate runs +in one table is how the numbers stop being comparable. + +One cost moved the other way, on that same non-recommended path. A gene slice +of a C-contiguous `(n_points, n_points, n_genes)` memory map is strided, so +reading one gene at a time touches the whole file where the old code did one +sequential read into RAM: measured 74 s against 66 s on a clean 900-landmark +pair, about 12% slower. Wall-clock figures here come from a shared machine +under load and should be read as orders of magnitude. + +**Results are unchanged.** Distances differ only by floating-point summation +order: measured maximum relative difference 1.7e-12 across the in-memory, +Dask and memmap paths, which is the same order as the pre-existing difference +*between* those paths on 0.8.0 alone. + +In-memory cost per gene therefore drops from `3 x n_landmarks^2 x 8` to +`2 x n_landmarks^2 x 8` bytes — about 0.37 GiB rather than 0.56 GiB per gene at +the default `n_landmarks=5000`. + +Fixes settylab/kompot#26. + +### Fixed — the dry run now prices the run it is pricing + +`kompot.de(..., dry_run=True)` built its plan *before* `null_genes="auto"` was +resolved, so the estimator received the literal string. That matched neither its +integer nor its list branch and fell through to zero, omitting the 2 000 null +genes the run would add. Measured on a 4 000 x 500 input: the default dry run +reported 0.485 GiB for a run the explicit `null_genes=2000` plan priced at +2.156 GiB — **4.44x optimistic**, and the gap grows with `n_cells`. + +Resolution now happens above the dry-run branch, and +`estimate_differential_expression_resources` raises on a `null_genes` it cannot +interpret rather than counting zero. + +Fixes settylab/kompot#25. + +### Changed — the resource plan describes the code that runs + +Following the fix above, `estimate_differential_expression_resources`: + + - no longer charges a third `(n_landmarks, n_landmarks, n_genes)` tensor that + is never built; + - charges the covariance tensors to **disk** only on the path that writes + them. With `dask` installed nothing reaches `disk_storage_dir`, and the plan + now reports that instead of demanding scratch the run does not use; + - charges the per-gene working set that replaced the dense sum + (`3 x n_landmarks^2 x 8` bytes, one gene at a time); + - reworded the in-memory warning to name the levers that shrink the tensors + (`genes`, linear; `n_landmarks`, quadratic) rather than only pointing at disk + storage. + +`store_arrays_on_disk` semantics are unchanged; only the accounting moved. -### Added +### Documentation + + - **New guide: [Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html)**, + the canonical explanation of what `sample_col` costs and how to run it, with + plans produced by `dry_run=True` at realistic sizes. It prescribes the + two-pass workflow — a cheap first pass over all genes, then sample variance + restricted to the top ~200 genes — ranks the levers, and shows how to price a run + before committing to it. The recommended budget for that second pass is on the order of + **200 genes**; every figure quoted beside it is priced for 200. + - The same warning now appears where the decision is made: `kompot.de`'s + docstring (so `help()` carries it), `GPSettings.n_landmarks`, + `GPSettings.batch_size`, `StorageSettings`, `SampleVarianceEstimator`, + `kompot de --sample-col`, the DE config template, the README and the tutorial + notebooks. + - The CLI "Complete Analysis" example no longer pairs `--sample-col` with + `--n-landmarks 5000` over every gene, which was the most expensive + configuration the package can express; it now shows the restricted second + pass through a config file. + - Corrected `StorageSettings.max_memory_ratio`, documented as "Fraction of RAM + before triggering disk storage". It triggers nothing: `store_arrays_on_disk` + defaults to `None`, which resolves to `disk_storage_dir is not None`, and the + ratio only sets the threshold at which the estimator escalates warnings. The + DE config template's `store_arrays_on_disk: null # (null = auto)` was + misleading for the same reason. + +Fixes settylab/kompot#27. + +### Added — `GPSettings.param_scheme` - **`GPSettings.param_scheme`** chooses which cells the length scale is estimated from: `"condition1"`, `"condition2"`, `"symmetric"`, `"pooled"` or `"separate"`. In `da()` it covers `d` and `mu` as well. The default, `None`, keeps current behaviour: `"condition1"` in `de()`, `"separate"` in `da()`. `"symmetric"` is swap-invariant. Both CLIs accept it as a config key. +### Fixed — the landmark snap is now exact + +`find_landmarks` snaps each cluster centroid to a cell. It used the approximate +pynndescent index built for the graph, and a centroid usually sits off the data, +where that search is least reliable. On synthetic data with 1M cells in 100 +dimensions (a Gaussian mixture with 150 clusters), 36 of 150 centroids (24%) +snapped to a cell that was not the nearest; the snapped cell was a median 1.3% +(max 11%) farther from the centroid than the nearest one. On one real dataset +of similar size the share was 27%. On real datasets of ~75k cells in 19 +dimensions it was at most one centroid in ~130. + +The snap is now an exact Euclidean search, chunked over the cells so that +memory stays bounded (about 64 MiB per working array, plus a shortlist of +near-tied cells that stays small unless the data sit far from the origin +relative to their spread). Exact ties, such as +duplicate cells, go to the lowest cell index, so the result is deterministic. +It costs about 1.2 s on 8 cores for 150 centroids against 1M cells in 100 +dimensions, against ~30 s to build the nearest-neighbor index on the same data. + +**This can change landmarks relative to 0.8.0**, and with them, at the +margin, any result you computed from them: a changed landmark is a nearby cell +that is closer to its cluster's centroid than the old one. Landmarks that were +already the nearest cell do not change, and the clustering is untouched. 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.8.0, +pass `find_landmarks(..., exact_snap=False)`. + +Fixes settylab/kompot#34. + ### Removed - `sync_parameters` on `DifferentialAbundance.fit` and `da()`; use `param_scheme="pooled"`. -### Fixed +### Fixed — other - `da()` now honours `GPSettings.ls`, and warns when the expression-only fields `sigma`, `eps` or `use_empirical_variance` are set. @@ -28,13 +188,45 @@ All notable changes to this project will be documented in this file. where they differed by a median of 82%), and fold changes and densities by at most 6e-4. `da()` without `gp`, `DifferentialAbundance()`, and every run with fewer landmarks than cells are unchanged. - -### Documentation + - **The CLI can write its output under anndata 0.11 and 0.12** (settylab/kompot#23). Those + versions refuse to write a pandas nullable-string array unless + `anndata.settings.allow_write_nullable_strings` is set, and pandas 3 produces one for every + ordinary string column, so `kompot de`, `da`, `smooth` and `dm` failed at the final write. The + CLI now opts in for the duration of its own write only; library calls and your global anndata + settings are untouched. + - **The provenance sha resolves in a linked git worktree checked out on a branch** + (settylab/kompot#28, settylab/kompot#33). Branch refs and `packed-refs` live in the shared + repository directory named by the worktree's `commondir` file, and resolution looked only in + the per-worktree directory, so every run from such a checkout was stamped + `kompot_git_sha: null`. Refs are now looked up in both, and a sha that cannot be resolved + inside a git checkout is logged as a warning rather than recorded silently. + A Kompot source tree vendored into another repository, nested below that repository's root, + no longer stamps that repository's commit as `kompot_git_sha`. A `kompot/` package copied to + the root of another repository still does: by position it cannot be told from Kompot's own. + +### Deprecated + + - **`SampleVarianceEstimator(dask_num_workers=...)`** had no effect and said it did + (settylab/kompot#31). It wrote the Dask config key `pool.num-workers`, which no scheduler + reads, changed Dask's global config as a side effect, and logged "Configured Dask to use N + workers". The Dask covariance tensor is evaluated one gene at a time, so there is no worker + pool for it to bound. Passing it now emits a `FutureWarning` and touches no configuration. + To bound Dask yourself, use `with dask.config.set(num_workers=N): ...`. + +### Documentation — argument order and the docs build - At the default, `de()` depends on argument order: the shared length scale comes from condition 1 alone. See `DifferentialExpression.fit`. - Automatic landmarks and `"pooled"` make `da()` depend on argument order; see `DifferentialAbundance.fit`. + - The docs build now fails when reStructuredText markup survives into the rendered HTML + (settylab/kompot#30). Nested inline markup such as a literal inside bold renders as literal + backticks and `sphinx-build` exits 0, so it was reaching published pages unnoticed. A + `build-finished` check scans every rendered page for unrendered literals, roles and directives. + It found 25 such constructs in the tutorial notebooks, all rewritten: 9 literals nested in bold, + 2 links whose text contained a literal, and 14 links whose whole text was a literal (Markdown + ``[`x`](url)``, which pandoc turns into nested RST that renders as ``x `__``). An earlier revision of the check could not + see the second kind, because the leaked literal renders as a tag and is stripped before matching. ## [0.8.0] - 2026-07-28 diff --git a/README.md b/README.md index 241692d..5412696 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ Key features: - Mahalanobis distance calculation for differential expression significance - JAX-accelerated computations with optional GPU support - Disk-backed covariance storage for sample variance estimation +- **Resource estimation and dry run** for planning large analyses - **Full scverse compatibility** with direct AnnData integration - **Visualization tools** for volcano plots, heatmaps, and embeddings - **Command-line interface** for pipeline integration @@ -66,6 +67,45 @@ kompot.de( ) ``` +### Sample variance: run it as a second pass + +Passing `sample_col` turns on **sample variance**, which replaces the single +shared posterior covariance with **one `(n_landmarks, n_landmarks)` covariance +matrix per gene, per condition**. Two costs follow, and they respond to +different levers: + +- **Memory**, `2 x n_landmarks^2 x n_genes x 8` bytes — about **0.37 GiB per + gene** at the default 5 000 landmarks. `StorageSettings(store_arrays_on_disk=True)` + removes this almost entirely, and `n_landmarks` shrinks it quadratically. +- **Compute**, one Cholesky factorisation **per gene** instead of one in total. + `store_arrays_on_disk` does not help here, but lowering `n_landmarks` does + (~0.016 s/gene at 500 against ~2.1 s at 5 000, single-threaded), as does + analysing fewer genes. + +So run it in two passes, and price the second one first: + +```python +# Pass 1 — all genes, no sample variance +kompot.de(adata, "condition", "Young", "Old") + +mahal = "kompot_de_Young_to_Old_mahalanobis" +top_genes = adata.var.sort_values(mahal, ascending=False).head(200).index + +# Pass 2 — sample variance, restricted to the top genes +plan = kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, + gp=kompot.GPSettings(n_landmarks=2000), # cost is quadratic in this + dry_run=True, # drop once the plan fits +) +``` + +`dry_run=True` returns a full resource plan (per-array memory, disk, output +fields, feasibility) without running anything; `kompot de --dry-run` does the +same from the CLI. Details, measured plans, and the remaining levers: +[Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html). + ### Command-Line Interface ```bash @@ -79,6 +119,7 @@ kompot de input.h5ad -o output.h5ad \ ## Documentation - [Full Documentation](https://kompot.readthedocs.io) +- [Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html) — what sample variance costs and the two-pass workflow - [Tutorial Notebooks](https://github.com/settylab/kompot/tree/main/examples) - [Getting Started](https://github.com/settylab/kompot/blob/main/examples/01_getting_started.ipynb) — differential expression, end to end - [Advanced Differential Expression](https://github.com/settylab/kompot/blob/main/examples/02_differential_expression_detailed.ipynb) — tuning, multiple comparisons, run tracking, resource planning diff --git a/docs/source/anndata.rst b/docs/source/anndata.rst index 04b2c44..bd014d9 100644 --- a/docs/source/anndata.rst +++ b/docs/source/anndata.rst @@ -45,14 +45,16 @@ Smooth Expression Resource Estimation ------------------- -Before running resource-intensive differential expression analyses, you can use the dry run utility to estimate memory and disk requirements, check for field overwrites, and verify parameters. +Before running resource-intensive differential expression analyses, use +``dry_run=True`` to estimate memory and disk requirements, check for field +overwrites, and verify parameters. **Key features:** - **Memory and disk estimation**: Calculates expected resource usage for all intermediate arrays and final results -- **Null genes accounting**: Correctly estimates resource inflation from null distribution genes (default 2000 additional genes) +- **Null genes accounting**: Estimates resource inflation from null distribution genes, resolving ``null_genes`` exactly as the run would, including the ``"auto"`` default - **Field overwrite detection**: Shows which fields will be overwritten, including their run_id and previous run details -- **Sample variance impact**: Estimates additional memory for sample-specific covariance tensors +- **Sample variance impact**: Estimates the per-gene covariance tensors, which are the dominant allocation whenever ``sample_col`` is set - **Disk storage planning**: Estimates disk space needed when using ``store_arrays_on_disk=True`` .. code-block:: python @@ -66,6 +68,7 @@ Before running resource-intensive differential expression analyses, you can use condition1='Young', condition2='Old', sample_col='donor_id', + genes=top_genes, dry_run=True, ) @@ -78,6 +81,12 @@ The dry run output shows: - **Warnings**: Field overwrite information showing previous run timestamp, conditions, and parameters - **Status**: Whether the analysis is feasible given available resources +.. seealso:: + + :doc:`Planning Memory and Disk ` explains what drives + the numbers, why ``sample_col`` must be a second pass over a restricted + gene list, and where the estimate is known to be optimistic. + Utilities --------- diff --git a/docs/source/cli.rst b/docs/source/cli.rst index 6929d90..5e0f7d3 100644 --- a/docs/source/cli.rst +++ b/docs/source/cli.rst @@ -242,13 +242,39 @@ Example: Complete Analysis --condition2 Old \ --obsm-key DM_EigenVectors \ --layer logged_counts \ - --sample-col Sample \ --n-landmarks 5000 \ --batch-size 100 \ --fdr-threshold 0.05 \ --null-genes 2000 \ --store-additional-stats +.. warning:: + + ``--sample-col`` is deliberately absent above. It turns on sample variance, + which gives every gene its own ``(n_landmarks, n_landmarks)`` covariance + matrix and its own Cholesky factorisation: about 0.37 GiB and ~2 s per gene + at ``--n-landmarks 5000``. Run it as a **second pass** over + a restricted gene list, which the CLI takes through the ``genes:`` key of a + config file: + + .. code-block:: yaml + + # sample_variance.yaml + sample_col: "Sample" + genes: ["GATA1", "KLF1", "..."] # top ~200 from the first pass + n_landmarks: 2000 # cost is QUADRATIC in this + null_genes: 0 # not calibrated for sample variance + store_arrays_on_disk: true + disk_storage_dir: "/scratch/kompot" + + .. code-block:: bash + + kompot de bone_marrow.h5ad -o results_sv.h5ad -c sample_variance.yaml \ + --groupby Age --condition1 Young --condition2 Old \ + --result-key kompot_de_sv --dry-run # drop --dry-run once it fits + + See :doc:`Planning Memory and Disk `. + Example: Dry Run (Resource Estimation) ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^ @@ -272,7 +298,13 @@ goes to stderr. Use the same arguments as the real run (``-o`` is ignored). kompot de input.h5ad --dry-run --groupby cond --condition1 A --condition2 B \ 2>/dev/null | jq '.memory.total_human' -The exit code is ``0`` if the plan is feasible, ``1`` if not. +The exit code is ``0`` if the plan is feasible, ``1`` if not, which makes +``--dry-run`` usable as a gate in a pipeline. + +Always dry-run before a ``--sample-col`` run: that is the configuration whose +cost is multiplied by the gene count. See +:doc:`Planning Memory and Disk ` for how to read the plan +and what the levers are. Differential Abundance Command ------------------------------- @@ -689,16 +721,26 @@ For large datasets: .. code-block:: bash - # Reduce batch size + # Reduce batch size (bounds prediction temporaries) kompot de input.h5ad -o output.h5ad ... --batch-size 50 - # Use fewer landmarks + # Use fewer landmarks (cost of sample variance is QUADRATIC in this) kompot de input.h5ad -o output.h5ad ... --n-landmarks 3000 # Enable disk storage (requires config file) # In config.yaml: # store_arrays_on_disk: true - # disk_storage_dir: "/tmp/kompot_cache" + # disk_storage_dir: "/scratch/kompot" # real scratch, not tmpfs + +With ``--sample-col`` the dominant allocation is +``2 x n_landmarks^2 x n_genes x 8`` bytes, so the memory levers are the +**gene list** (linear) and ``--n-landmarks`` (quadratic); setting +``store_arrays_on_disk: true`` in a config file keeps it out of memory +altogether. ``--batch-size`` does not bound that term, and +``store_arrays_on_disk`` does not touch the per-gene factorisation at all -- +but lowering ``--n-landmarks`` reduces it as well as the memory, and a shorter +gene list reduces both linearly. Full treatment: +:doc:`Planning Memory and Disk `. Speed Optimization ^^^^^^^^^^^^^^^^^^ @@ -750,7 +792,10 @@ Common Issues MemoryError or JAX out of memory -*Solution:* Reduce ``--batch-size`` and ``--n-landmarks`` +*Solution:* Reduce ``--batch-size`` and ``--n-landmarks``. If the run uses +``--sample-col``, that is almost certainly the cause: restrict the analysis to +a top-gene list and cut ``--n-landmarks``, then re-check with ``--dry-run`` +(:doc:`Planning Memory and Disk `). Getting Help ^^^^^^^^^^^^ @@ -811,5 +856,6 @@ See Also - :doc:`Python API Documentation ` - :doc:`Getting Started Tutorial ` +- :doc:`Planning Memory and Disk ` - :doc:`Sample Variance Guide ` - `GitHub Repository `_ diff --git a/docs/source/conf.py b/docs/source/conf.py index 2128112..17b1835 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -114,6 +114,12 @@ def copy_nbsphinx_images_and_fix_paths(app, exception): def setup(app): """Setup function to register Sphinx event handlers.""" app.connect("build-finished", copy_nbsphinx_images_and_fix_paths) + # Fail the build if any markup survived into the HTML as text: sphinx + # itself exits 0 on nested inline markup (settylab/kompot#30). + sys.path.insert(0, str(this_directory)) + from markup_leak_check import check_after_build + + app.connect("build-finished", check_after_build) # Add a prolog to control figure rendering diff --git a/docs/source/index.rst b/docs/source/index.rst index 0a5830e..7b33fea 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -6,6 +6,7 @@ :caption: User Guide: Installation + Planning Memory and Disk Command-Line Interface .. toctree:: @@ -78,6 +79,7 @@ Key features: - Empirical variance estimation - **Resource estimation and dry run** for planning large analyses - Disk-backed covariance storage for sample variance estimation + (see :doc:`Planning Memory and Disk `) - Visualization tools (volcano plots, heatmaps, expression plots) - **Full scverse compatibility with direct AnnData integration** - **Command-line interface for pipeline integration** and batch processing @@ -126,13 +128,22 @@ Quick Start # Customize GP and FDR settings kompot.de( adata, "condition", "Young", "Old", - sample_col="donor_id", gp=kompot.GPSettings(sigma=0.5, use_empirical_variance=True), fdr=kompot.FDRSettings(threshold=0.01), ) See :doc:`simplified` for the full API and all available settings. +.. warning:: + + Adding ``sample_col`` turns on sample variance, which gives every gene its + own landmark covariance matrix and its own Cholesky factorisation. At + default settings that is roughly **0.37 GiB and ~2 s per gene**, so it + belongs in a second pass over a restricted gene list rather than a + whole-transcriptome run. Read + :doc:`Planning Memory and Disk ` before your first + sample-variance run. + **New to Kompot?** Start with the :doc:`Getting Started ` tutorial: differential expression, end to end, with rich visualizations. **Ready for more?** Explore advanced topics: @@ -141,6 +152,7 @@ See :doc:`simplified` for the full API and all available settings. - :doc:`DE with Sample Variance ` - Replicate-aware significance in multi-sample studies - :doc:`Differential Abundance ` - Cell-state frequency changes, including its sample-variance treatment - :doc:`Smoothing Expression ` - The expression function underneath DE: reading its two uncertainties, and fitting on some cells to predict on others +- :doc:`Planning Memory and Disk ` - What sample variance costs, the two-pass workflow, and pricing a run with ``dry_run=True`` Command-Line Interface ^^^^^^^^^^^^^^^^^^^^^^^ diff --git a/docs/source/markup_leak_check.py b/docs/source/markup_leak_check.py new file mode 100644 index 0000000..c14e3da --- /dev/null +++ b/docs/source/markup_leak_check.py @@ -0,0 +1,109 @@ +"""Fail the docs build when reStructuredText markup survives into the HTML. + +Nested inline markup (``**the ``dask`` path**``) is invalid RST, and docutils +does not warn about it: it stops treating the inner markers as markup and the +page renders literal backticks, while ``sphinx-build`` exits 0 +(settylab/kompot#30). The same silent pass covers a role or directive written +where it is not interpreted, such as inside a code-block comment. + +Rather than a rule per construct, this asserts the one property all of them +violate: nothing that is supposed to be markup appears as text in the output. +Run automatically from ``conf.py`` at ``build-finished``; also usable +standalone as ``python markup_leak_check.py ``. + +What it cannot see is markup whose leaked form contains no marker, for +example a single ``*`` of italic nested in bold: a bare asterisk is too +common in real text to flag. +""" + +import html as _html +import re +import sys +from pathlib import Path + +# Markup that must never survive rendering, as it would appear in page text. +LEAK_PATTERNS = { + # One hit per literal: a closed ``pair`` on one line, or a stray opener. + "double-backtick literal": re.compile(r"``(?:[^`\n]*``)?"), + # A hyperlink whose text was a literal: pandoc turns Markdown [`x`](url) + # into nested RST, which renders as `x `__. The inner + # literal becomes a tag and is stripped, so only this tail is left as text. + "hyperlink": re.compile(r"<[^<>\s]+>`__?"), + "role": re.compile(r":(?:doc|ref|func|class|meth|mod|attr|obj|data|exc|term):`"), + "directive": re.compile( + r"^\s*\.\. (?:warning|note|admonition|code-block|list-table|toctree|" + r"autofunction|autoclass|image|figure|math|seealso|versionadded|" + r"deprecated)::", + re.MULTILINE, + ), +} + +# Sphinx's own search index and sources copies are not rendered pages. +_SKIP_DIRS = {"_sources", "_static", "_images", "_modules"} + +_TAG = re.compile(r"<[^>]+>") +_SCRIPT_STYLE = re.compile(r"<(script|style)\b.*?", re.DOTALL | re.IGNORECASE) + + +def page_text(markup: str) -> str: + """The visible text of an HTML page: tags removed, entities decoded.""" + markup = _SCRIPT_STYLE.sub("", markup) + return _html.unescape(_TAG.sub("", markup)) + + +def find_leaks(markup: str): + """Return ``[(kind, line_text), ...]`` for every leaked construct in *markup*.""" + text = page_text(markup) + leaks = [] + for kind, pattern in LEAK_PATTERNS.items(): + for match in pattern.finditer(text): + start = text.rfind("\n", 0, match.start()) + 1 + end = text.find("\n", match.end()) + line = text[start : end if end != -1 else len(text)].strip() + leaks.append((kind, line[:160])) + return leaks + + +def scan_build(outdir): + """Scan every rendered page under *outdir*; return ``{relpath: leaks}``.""" + outdir = Path(outdir) + found = {} + for page in sorted(outdir.rglob("*.html")): + rel = page.relative_to(outdir) + if _SKIP_DIRS.intersection(rel.parts): + continue + leaks = find_leaks(page.read_text(encoding="utf-8", errors="replace")) + if leaks: + found[str(rel)] = leaks + return found + + +def format_report(found) -> str: + lines = [ + f"{sum(len(v) for v in found.values())} unrendered markup construct(s) " + f"in {len(found)} page(s) -- nested inline markup, or a role/directive " + "where it is not interpreted (settylab/kompot#30):" + ] + for page, leaks in found.items(): + for kind, line in leaks: + lines.append(f" {page}: {kind}: {line}") + return "\n".join(lines) + + +def check_after_build(app, exception): + """Sphinx ``build-finished`` hook: fail an HTML build that leaked markup.""" + if exception is not None or app.builder.format != "html": + return + found = scan_build(app.outdir) + if found: + from sphinx.errors import ExtensionError + + raise ExtensionError(format_report(found)) + + +if __name__ == "__main__": + result = scan_build(sys.argv[1]) + if result: + print(format_report(result)) + sys.exit(1) + print("no unrendered markup found") diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst new file mode 100644 index 0000000..e79a299 --- /dev/null +++ b/docs/source/resource_planning.rst @@ -0,0 +1,701 @@ +Planning Memory and Disk +======================== + +Most Kompot runs are unremarkable in their resource use. One option is not: +supplying ``sample_col`` to :func:`kompot.de` turns on **sample variance**, +which replaces a single shared covariance matrix with **one covariance matrix +per gene**. Both the memory and the compute of that step are multiplied by the +number of genes analysed. + +This page explains where that cost comes from, prescribes the two-pass workflow +that keeps it bounded, and shows how to price a run with ``dry_run=True`` +before committing to it. + +.. _two-pass-workflow: + +The prescribed workflow: two passes +----------------------------------- + +**Never run sample variance over the whole transcriptome.** Run it twice +instead: + +1. **Pass 1, all genes, no sample variance.** Cheap, and it gives you the + Mahalanobis ranking. +2. **Pass 2, top genes only, with sample variance.** Restrict to the genes that + survived pass 1 (on the order of 200; see :ref:`how-many-genes`), offload + the covariance tensors to disk, and cut ``n_landmarks`` if the run is still + too slow. + +.. code-block:: python + + import kompot + + # ---- Pass 1: all genes, no sample variance ------------------------- + kompot.de(adata, "condition", "Young", "Old") + + # ---- Pick the genes worth refining --------------------------------- + mahal = "kompot_de_Young_to_Old_mahalanobis" + top_genes = adata.var.sort_values(mahal, ascending=False).head(200).index + + # ---- Pass 2: sample variance, restricted and offloaded ------------- + kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, # linear in both memory and time + gp=kompot.GPSettings(n_landmarks=2000), # memory is QUADRATIC in this + storage=kompot.StorageSettings( + store_arrays_on_disk=True, # keeps the tensors out of RAM + disk_storage_dir="/scratch/kompot", # large, fast, node-local + ), + ) + +Price it before you run it: a dry run of exactly that call takes a second and +tells you whether it fits (see :ref:`dry-run`). + +The two costs respond to different levers, which is why the split is worth +making deliberately rather than by habit. ``store_arrays_on_disk`` removes the +memory term almost entirely (:ref:`disk-offload`) and leaves compute alone; +``n_landmarks`` reduces **both**, memory quadratically and the per-gene +factorisation by roughly the same power (see :ref:`other-levers`); and +``genes`` is the only lever that is linear in both at once. Restricting +``genes`` is therefore the one that always moves the total, but it is not the +only one that touches compute. + +Held in memory, the gene list is a hard constraint: a whole transcriptome at +default landmarks asks for 7 552 GiB. Offloaded, it is a budget rather than a +wall — 101 GiB and about 11.6 hours of single-threaded factorisation. See +:ref:`how-many-genes` for which of those you are in. + +The two passes compose rather than collide, which is what makes the split +practical. Measured on a two-pass run where pass 2 analysed 10 of 60 genes: + +* the sample-variance score is written to a **new** column with a + ``_sample_var`` suffix + (``kompot_de_Young_to_Old_mahalanobis_sample_var``), populated for the genes + pass 2 analysed and ``NaN`` elsewhere, while pass 1's + ``..._mahalanobis`` column stays populated for **all** genes. The two sit + side by side in ``adata.var`` and can be compared directly; +* ``..._mean_lfc`` is unchanged. Fold changes are unaffected by sample + variance; only the significance moves; +* the shared layers (``..._smoothed``, ``..._fold_change``) are rewritten for + the genes pass 2 analysed and left untouched for the rest. Pass 2 refits the + GP on its gene subset, so those columns change slightly. If you need pass 1's + layers preserved byte for byte, give pass 2 its own + ``StorageSettings(result_key=...)``; +* two new layers appear, ``...__std``, one per condition. + +FDR is disabled automatically on the second pass: with ``sample_col`` set, +``FDRSettings(null_genes="auto")`` resolves to ``0``, because the null +calibration is not yet validated for sample variance. That is also what you +want for cost, since null genes are charged at the same per-gene rate as real +ones (see :ref:`the null-gene note `). + + +.. _where-the-cost-is: + +Where the cost comes from +------------------------- + +Without sample variance, the Mahalanobis distance for every gene is taken +against **one** posterior covariance matrix of shape ``(n_landmarks, +n_landmarks)``, factorised once. + +With sample variance, Kompot fits a per-sample expression predictor, evaluates +it at the landmarks, and forms the sample-to-sample covariance between landmark +pairs **separately for every gene**. That array has shape + +.. code-block:: text + + (n_landmarks, n_landmarks, n_genes) + +and there is one per condition. Their sum is never materialised: the +Mahalanobis step assembles a single ``(n_landmarks, n_landmarks)`` matrix per +gene as it goes. So the dominant term, when the tensors are held in memory, is + +.. code-block:: text + + bytes = 2 x n_landmarks^2 x n_genes x 8 (float64) + +Read the exponents carefully. The cost is **linear in genes** and **quadratic +in landmarks**. At Kompot's default ``n_landmarks=5000``, a single landmark +covariance matrix is ``5000^2 x 8 B = 190 MiB``, so: + +.. admonition:: The number to remember + :class: warning + + At the default 5 000 landmarks, **every gene added to an in-memory + sample-variance run costs about 0.37 GiB**. Two hundred genes is roughly + 75 GiB; the whole transcriptome is over seven terabytes. Nothing else in a + Kompot run behaves this way. + + :ref:`disk-offload` removes that term from memory entirely, which is what + makes large gene sets possible at all. It does not remove the **compute** + below, which is why the two-pass workflow remains the prescription either + way. + +.. note:: + + Before Kompot 0.9.0 the two tensors were summed into a third dense array + before use, so the figure above was ``3 x n_landmarks^2 x n_genes x 8`` + bytes, or about 0.56 GiB per gene, and ``store_arrays_on_disk`` did not + reduce peak memory at all. If you are reading a plan produced by 0.8.0 or + earlier, those are the numbers it will show + (`settylab/kompot#26 `_). + +Compute scales with genes as well, and it scales worse than the shared path +does. With one shared covariance Kompot performs a single Cholesky +factorisation and solves against it for all genes in vectorised batches; the +total is essentially flat in the gene count. With per-gene covariances it +factorises **once per gene**, in a Python loop over genes +(``kompot/utils.py``), and ``GPSettings.batch_size`` does not apply to that +loop. + +Measured inside a real ``kompot.de`` sample-variance run, 20 genes, with +``OMP_NUM_THREADS=1``: + +.. list-table:: + :header-rows: 1 + :widths: 24 26 24 26 + + * - ``n_landmarks`` + - per gene + - 200 genes + - 20 000 genes + * - 500 + - 0.016 s + - 3 s + - 5 min + * - 1 000 + - 0.062 s + - 12 s + - 21 min + * - 2 000 + - 0.246 s + - 49 s + - 1.4 h + * - 5 000 (default) + - 2.08 s + - 7 min + - 11.6 h + +.. warning:: + + **Pin your BLAS thread count for sample-variance runs on a shared node.** + The per-gene factorisation is a single multi-threaded LAPACK call on a + matrix that is small relative to the core count, so on a busy machine it + spends its time in thread contention rather than in arithmetic. Measured on + the same host, same commit, same run, ``n_landmarks=2000``: + + .. code-block:: text + + OMP_NUM_THREADS=1 0.246 s/gene (whole run 52.7 s) + unrestricted 6.135 s/gene (whole run 280.8 s) + + A bare ``numpy.linalg.cholesky`` at n=2000 shows the same thing from the + other side: 0.125 s at 21.3 GFLOPS pinned, against 5.20 s at **0.5 GFLOPS** + unrestricted. A throughput figure that *collapses* as threads are added is + contention, not work. + + So export ``OMP_NUM_THREADS=1`` (and ``OPENBLAS_NUM_THREADS`` / + ``MKL_NUM_THREADS`` to match) before a sample-variance run you are sharing + a node for. On an idle machine the unrestricted path would win; the point is + that most people run this on a machine that is not idle. + + Pinning also makes the figure *reproducible*, which is the other reason to + do it. Two independent measurements of the same configuration at 5 000 + landmarks, single-threaded, taken at 1-minute loads of 45.9 and 62.8 — a + 37% difference in load — came out at 2.084 and 2.113 s per gene, **1.4% + apart**. Pin the threads and the figure stops depending on what else the + machine is doing. + + The 25x above is a different quantity and worth not confusing with this + one: it is pinned *against* unrestricted at one moment, so it says how much + the level moves, not how much an unpinned measurement varies between runs. + The two unrestricted measurements taken here agreed to within 6-10% of each + other; what makes them unusable is not that they scatter but that they are + all wrong by more than an order of magnitude. + +The table above is therefore a **reproducible floor**, not a forecast: it is +what the step costs when it is not competing for cores. Your own figure depends +on your BLAS build, your core count and your node's load, so time a handful of +genes at your own ``n_landmarks`` before committing to a long run. + +Two smaller terms also grow with the gene count under sample variance, and +neither is affected by ``store_arrays_on_disk``: + +* **Per-sample imputations**, ``2 x n_samples x n_landmarks x n_genes x 8`` B. + At 5 000 landmarks, 200 genes and 6 donors this is 0.09 GiB. +* **Two extra layers**, ``__std``, each + ``n_cells x n_genes x 8`` B, written into ``adata.layers``. + +Measured plans +-------------- + +Kompot's own ``dry_run=True`` output for a synthetic AnnData of 20 000 cells +and 20 000 genes, two conditions, 6 donors, at the default +``n_landmarks=5000`` and ``null_genes=0``. These are totals for the whole plan, +not just the covariance term. (Kompot's report labels these binary magnitudes +``GB``; they are the same numbers.) + +.. list-table:: + :header-rows: 1 + :widths: 12 18 20 22 18 + + * - genes + - no sample variance + - sample variance, in memory + - sample variance, ``store_arrays_on_disk`` + - disk used + * - 200 + - 3.1 GiB + - 78.3 GiB + - 3.8 GiB + - 0 + * - 1 000 + - 6.4 GiB + - 380.2 GiB + - 7.7 GiB + - 0 + * - 2 000 + - 10.6 GiB + - 757.7 GiB + - 12.6 GiB + - 0 + * - 20 000 (all) + - 85.7 GiB + - 7 552 GiB + - 101.2 GiB + - 0 + +Three things to read off this. + +**The in-memory column is the one that explodes.** 200 genes asks for +78 GiB, and the whole transcriptome for seven and a half terabytes, because +that column carries ``2 x n_landmarks^2 x n_genes x 8`` bytes of covariance. + +**Disk offload collapses it.** The same 200 genes plans at 3.8 GiB, barely +above the 3.1 GiB of a run with no sample variance at all, because the tensors +are never held whole. The disk column reads zero because ``dask`` is installed +in the environment that produced the table, and on that path the tensors are +evaluated lazily rather than written; see :ref:`disk-offload` for the +``dask``-less path, where the same figure is written to +``disk_storage_dir`` instead. + +**Which constraint binds depends entirely on the storage mode, and the gap is +three orders of magnitude.** Held in memory, the whole transcriptome asks for +7 552 GiB and is simply out of reach. Offloaded, it asks for 101 GiB and a +large node has that. So :ref:`the two-pass workflow ` is a +hard requirement in the first case and a matter of value for money in the +second — see :ref:`how-many-genes`, which says which is which rather than +implying a wall that is no longer there. + +Reproduce any row on your own data by substituting your ``adata``: + +.. code-block:: python + + for on_disk in (False, True): + plan = kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, + storage=kompot.StorageSettings(store_arrays_on_disk=on_disk), + fdr=kompot.FDRSettings(null_genes=0), + dry_run=True, + ) + print(on_disk, plan.total_memory_required, plan.total_disk_required) + + +.. _dry-run: + +The dry run is the planning tool +-------------------------------- + +``kompot.de(..., dry_run=True)`` estimates the whole plan and returns a +:class:`~kompot.resource_estimation.ResourcePlan` instead of running anything. +It prints a report and exposes the numbers programmatically: + +.. code-block:: python + + plan = kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, + storage=kompot.StorageSettings(store_arrays_on_disk=True), + dry_run=True, + ) + + plan.total_memory_required # bytes + plan.total_disk_required # bytes + plan.is_feasible # False if it cannot fit + plan.warnings # overwrite and headroom warnings + for r in plan.requirements: + print(r.resource_type, r.shape, r.size_human, r.name) + +The report breaks the plan down per array, so the sample-variance term is +visible by name (``Sample covariances (per condition, N samples)``) with its +``(n_landmarks, n_landmarks, n_genes)`` shape. It also lists the AnnData +fields that would be created and flags any that an earlier run already wrote, +with that run's ``run_id``. + +The intended use is to price the options against each other before committing +to any of them. Hold the gene set fixed to read what sample variance itself +costs: + +.. code-block:: python + + # null_genes is pinned on BOTH arms. Left at "auto" it resolves to 2 000 + # without sample_col and to 0 with it, so the comparison would silently be + # between two different gene counts. + base = dict(groupby="condition", condition1="Young", condition2="Old", + genes=top_genes, fdr=kompot.FDRSettings(null_genes=0), + dry_run=True) + + without = kompot.de(adata, **base) + with_sv = kompot.de(adata, sample_col="donor_id", **base) + + print(with_sv.total_memory_required / without.total_memory_required) + +then vary the levers, one at a time, until the plan fits: + +.. code-block:: python + + for n_landmarks in (1000, 2000, 5000): + plan = kompot.de(adata, sample_col="donor_id", + gp=kompot.GPSettings(n_landmarks=n_landmarks), **base) + print(n_landmarks, plan.total_memory_required, plan.is_feasible) + +The same estimate is available from the command line as ``kompot de +--dry-run``, which writes JSON to stdout and the human-readable report to +stderr, and exits non-zero when the plan is infeasible. See :doc:`cli`. + +.. warning:: + + **The plan is a floor, not a total.** It sums the arrays Kompot itself + allocates; it does not model the Python interpreter, NumPy/BLAS scratch, + JAX's device pools or Dask's scheduler. Measured on the 600-landmark / + 150-gene configuration used elsewhere on this page: a run with no sample + variance plans at 69 MiB and peaks at 958 MiB of anonymous memory, and the + *difference* sample variance makes plans at 16.7 MiB against a measured + 222 MiB — about 13x optimistic. + + That gap is roughly constant for a given dataset shape rather than growing + with the covariance term, which is what makes the plan useful anyway: the + figures it gets *exactly* right are the ones that explode, the + ``2 x n_landmarks^2 x n_genes x 8`` tensors and the disk footprint. Use it + to compare configurations and to catch the terabyte-scale plans, and leave + real headroom above whatever it reports. + +.. note:: + + ``kompot.resource_estimation.dry_run_differential_expression()`` is + deprecated. Use ``kompot.de(..., dry_run=True)``. + +.. _null-genes-cost: + +.. note:: + + Null genes are charged at the same per-gene rate as real ones, and the dry + run accounts for them: it resolves ``null_genes`` exactly as the run would, + including the ``"auto"`` default, so the plan prices the run it describes. + If you override the default and request FDR alongside sample variance, each + null gene gets its own landmark covariance matrix too — 200 real genes + plus 2 000 null genes is a 2 200-gene tensor, eleven times the cost of the + genes you asked for. + + Before Kompot 0.9.0 the estimate was built before ``"auto"`` was resolved + and silently counted zero null genes, so a default-settings plan without + ``sample_col`` under-reported by 2 000 genes + (`settylab/kompot#25 `_). + + +.. _disk-offload: + +Disk offload +------------ + +``StorageSettings(store_arrays_on_disk=True)`` keeps the per-gene covariance +tensors out of memory. It is the single most effective lever once you have +decided which genes to refine, and since Kompot 0.9.0 it is also **cheaper +than the in-memory path rather than dearer**. + +.. code-block:: python + + kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, + storage=kompot.StorageSettings( + store_arrays_on_disk=True, + disk_storage_dir="/scratch/kompot", # large, fast, node-local + ), + ) + +What it does mechanically depends on whether ``dask`` is installed: + +* **With** ``dask`` (``pip install 'kompot[dask]'``) each condition's tensor is + a lazy Dask graph, evaluated one gene at a time as the Mahalanobis step + reaches it. **Nothing is written to disk**, so ``disk_storage_dir`` is used + only for its free-space check. This is the faster path and the recommended + one. +* **Without** ``dask`` each condition's tensor is written as a memory-mapped + ``covariance_matrix.npy`` under ``disk_storage_dir``, so the dry run's disk + figure is what actually lands there. The per-gene work is then sequential, + and the dry run warns that ``dask`` is missing. + + **This path is also slower than holding the tensors in memory**, by about + 12% on a measured 900-landmark pair (74 s against 66 s). A gene slice of a + C-contiguous ``(n_points, n_points, n_genes)`` memory map is **strided** — + its elements sit ``n_genes * 8`` bytes apart — so reading one gene at a time + touches the whole file, where holding the tensor reads it once sequentially. + That is the price of not holding it, you pay it only without ``dask``, and + installing ``dask`` removes it. + +Measured on 1 500 cells, 150 genes, 600 landmarks and 4 donors (412 MiB per +tensor), four BLAS threads. **Two memory instruments, because one is not +enough**: ``Anonymous`` counts private heap pages, ``Rss`` counts those *plus* +resident file-backed pages, and the difference is exactly where a memory map +puts its data. + +.. list-table:: + :header-rows: 1 + :widths: 28 13 13 15 12 12 + + * - run + - anon + - Rss + - file-backed + - written + - wall + * - no sample variance + - 958 MiB + - 1 218 MiB + - 260 MiB + - 0 + - 36 s + * - sample variance, in memory + - 1 975 MiB + - 2 238 MiB + - 263 MiB + - 0 + - 49 s + * - ``store_arrays_on_disk``, with ``dask`` + - 1 180 MiB + - 1 442 MiB + - 262 MiB + - **0** + - 50 s + * - ``store_arrays_on_disk``, no ``dask`` + - 1 139 MiB + - **2 224 MiB** + - **1 085 MiB** + - 824 MiB + - 51 s + +As a share of the extra memory the in-memory run costs over a run with no +sample variance at all: + +.. list-table:: + :header-rows: 1 + :widths: 34 22 22 + + * - path + - on ``Anonymous`` + - on ``Rss`` + * - with ``dask`` + - ~1/5 + - **~1/5** + * - no ``dask`` + - ~1/5 + - **~97%** + +**These are deliberately given to one figure.** The share is a ratio of two +differences — an extra of roughly 200 MiB between peaks of 950 and 1 150 MiB — +so run-to-run noise in either peak swings it hard. Measured across twelve runs +on this host, the ``no sample variance`` baseline spans 938-981 MiB and the +``dask`` peak spans 1 134-1 209 MiB, which brackets that share anywhere from +**14% to 27%** without anything changing but the run. Two independent rigs +landing on 17% and 22% is the same noise, not a methodology difference. + +Pinning does not tighten it. Fixing Dask's threaded-scheduler pool at 1, at 4 +and leaving it at the default (36) gave 1 173/1 209, 1 152/1 163 and +1 134/1 138 MiB — all inside the spread, and the single-worker runs were if +anything the highest. + +What the measurement *does* support is the separation, and it is not close: +the ``dask``-less path costs **about a fifth on Anonymous and essentially all +of it on Rss**, a factor of five apart on the same run. That is the finding; +the second significant figure never was. + +The ``dask`` path **wins on every instrument**, and that is the claim to rely +on: about a fifth of the extra memory, nothing written to disk, nothing +file-backed, and 4.8x faster than 0.8.0 at matched thread counts. It is the +recommended path for exactly this reason. + +The ``dask``-less path is a different bargain, and **should not be described as +a smaller footprint**. It is 99% of the in-memory extra on ``Rss``: what it +changes is the *kind* of page, converting private anonymous memory into +resident page cache backed by a file on disk. Those pages are reclaimable — the +kernel can evict them under pressure and read them back — so it is a genuine +win for surviving a memory squeeze, and not a reduction in resident footprint. + +.. warning:: + + **Under a cgroup, page cache is charged to you.** Slurm's ``--mem``, a + container limit and any other cgroup-based cap count resident file-backed + pages against the same budget as anonymous ones. So the ``dask``-less path's + advantage may not exist in exactly the environment most people run this in. + The figures above were taken on an unconstrained host, so they show what is + resident, not what an OOM killer would decide; treat the reclaimability + argument as a mechanism, not as a verdict. If you are under a hard cap, + install ``dask``. + +.. note:: + + Earlier revisions of this page reported only ``Anonymous`` and justified it + as keeping page cache from "flattering the reading". That justification was + backwards: choosing ``Anonymous`` is precisely what *excludes* the 823 MiB + the ``dask``-less path moves into file-backed residency, which is why it + appeared to cost a sixth of the in-memory run rather than nearly all of it. + +Either way, **set** ``disk_storage_dir``, and create it before you plan against +it. A real run creates a missing directory, but the **dry run raises** +``FileNotFoundError`` on one, so the estimate fails on exactly the +configuration you were trying to price. Left unset, Kompot uses the system +temporary directory (honouring ``TMPDIR``), which on a shared cluster is +frequently small, RAM-backed, or both. Kompot creates a unique per-run +subdirectory inside it and removes that subdirectory when the estimator is +collected, so space is reclaimed at the end of the run rather than during it. + +``store_arrays_on_disk`` defaults to ``None``, which means *on if and only if* +``disk_storage_dir`` is set. Nothing turns it on automatically in response to +memory pressure. ``max_memory_ratio`` only sets the threshold at which the +plan's warnings escalate; it does not switch storage modes. + +.. note:: + + Before Kompot 0.9.0 this flag did not reduce peak memory, because the two + tensors were summed into a dense array before use whatever the storage + mode. On the same configuration as the table above, 0.8.0 measured + 2 782 MiB with ``dask`` and 2 284 MiB without, against 2 401 MiB for the + in-memory run — that is, offloading made the run **more** expensive, and + the ``dask`` path was 4.9x slower as well + (`settylab/kompot#26 `_). + +.. _other-levers: + +Other levers +------------ + +``n_landmarks`` **is the strongest memory lever**, because the covariance term +is quadratic in it. Same 200 genes, from the dry run: + +.. list-table:: + :header-rows: 1 + :widths: 24 26 26 + + * - ``n_landmarks`` + - sample variance, in memory + - with ``store_arrays_on_disk`` + * - 1 000 + - 5.4 GiB + - 2.5 GiB + * - 2 000 + - 14.6 GiB + - 2.6 GiB + * - 3 000 + - 29.7 GiB + - 2.9 GiB + * - 5 000 (default) + - 78.3 GiB + - 3.8 GiB + +Halving the landmark count quarters the covariance footprint, which is exact +arithmetic rather than a measurement. It cuts the per-gene factorisation too: +pinned to one thread, 0.016 s per gene at 500 landmarks against 2.08 s at +5 000. + +**That is roughly quadratic, not cubic**, which is worth stating because a +Cholesky's O(n³) flop count invites the opposite guess. Over those four points +the fitted exponent is **2.11**, and it is consistent interval to interval +(1.94, 2.00, 2.33). Cubic would predict 16 s per gene at 5 000 landmarks +against the 2.08 s measured — a 7.7x overshoot. + +Two things eat the third power at these sizes. The step is not only the +factorisation: materialising each gene's matrix and the triangular solve +against it are both O(n²), and at n=500 they are a large share of the work. +And the factorisation itself gets *more efficient* as n grows — measured +single-threaded throughput climbs from 8.2 GFLOPS at n=500 to 37.7 at +n=5 000, a 4.6x improvement that cancels much of the extra work. + +Whether the exponent eventually climbs toward 3 once that efficiency headroom +is exhausted is **not something to assume**: measured out to n=11 000 it has +not started to, and throughput was still rising (37.8 → 42.3 → 53.1 GFLOPS at +n=5 000, 8 000, 11 000). Treat 2.11 as a description of the range measured, not +as a curve to extrapolate along. + +Four timings on one machine are still not a scaling law: they move with your +BLAS build and your node's load. Take them as a floor and time your own. + +Landmarks control the resolution of the cell-state approximation, so reducing +them is a genuine accuracy trade-off rather than a free saving. But 5 000 is +rarely required to resolve a covariance structure that is being summarised gene +by gene. + +``batch_size`` (in :class:`~kompot.GPSettings`) bounds the temporary arrays +during prediction and the gene batches of the *shared*-covariance Mahalanobis +computation. It does **not** bound the per-gene covariance loop, which +processes one gene at a time regardless, so it is a lever for pass 1 and for +peak prediction memory, not for sample variance. + +``genes`` is the lever that matters for pass 2, and it is the subject of the +next section. + +.. _how-many-genes: + +How many genes in pass 2? +------------------------- + +Both costs are linear in the gene count, so this is a budget decision rather +than a statistical one, and Kompot supplies no default: ``genes=None`` means +*all of them*, which is the one setting to avoid here. + +**Roughly 200 genes is a recommendation chosen from the cost curves, not a +tuned or validated parameter — and it is a budget, not a wall.** Be clear about +which constraint you are actually under, because they differ by orders of +magnitude: + +* **In memory**, the whole transcriptome is out of reach: 7 552 GiB at default + landmarks. This is exact arithmetic and it is the hard limit. +* **With** ``store_arrays_on_disk=True``, it is not: 101 GiB, which a large + node has. The memory wall is the one the offload removes. +* **In time**, the whole transcriptome is a long job rather than an impossible + one: ~11.6 h at 5 000 landmarks, ~1.4 h at 2 000, single-threaded. + +So a whole-transcriptome sample-variance run is *feasible* once the tensors are +offloaded. Restricting to ~200 genes is still the right default, for reasons +that are about value rather than capacity: + +* pass 1 has to run over all genes anyway, and it is what ranks them; +* 200 genes costs about a hundredth of the whole transcriptome on both axes, + for the genes you were going to read; +* FDR is not calibrated for sample variance (``null_genes`` resolves to ``0`` + when ``sample_col`` is set), so the extra 19 800 genes buy no additional + testable calls. + +If you do want wider coverage, cut ``n_landmarks`` in the same breath: it +reduces the per-gene time and the memory together. + +Price it rather than guessing. A :ref:`dry run ` gives the memory and +disk figures for your dimensions in a second; for the time, multiply your gene +count by a single-gene timing at your landmark count, pinned. + +See also +-------- + +* :doc:`DE with Sample Variance ` for the + two-pass workflow end to end on real data. +* :doc:`Advanced Differential Expression ` + for resource planning alongside the other tuning knobs. +* :doc:`simplified` for :class:`~kompot.StorageSettings` and + :class:`~kompot.GPSettings` in full. +* :doc:`cli` for ``kompot de --dry-run``. diff --git a/docs/source/simplified.rst b/docs/source/simplified.rst index f0067cd..fb4979f 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -44,6 +44,21 @@ Differential Expression fdr=kompot.FDRSettings(null_genes=0), ) +.. warning:: + + ``sample_col`` is the one expensive option. It gives every gene **its own** + ``(n_landmarks, n_landmarks)`` covariance matrix per condition, so the + dominant allocation becomes ``2 x n_landmarks^2 x n_genes x 8`` bytes — + about 0.37 GiB per gene at the default 5 000 landmarks — and the + Mahalanobis step factorises once per gene instead of once in total. + ``StorageSettings(store_arrays_on_disk=True)`` removes the memory term and + leaves compute alone; lowering ``GPSettings(n_landmarks=...)`` reduces both; + only a shorter gene list is linear in both at once. + Always run it as a second pass over a restricted gene list, after a cheap + first pass with no ``sample_col``. The full treatment, with measured plans + and the other levers, is in + :doc:`Planning Memory and Disk `. + .. autofunction:: kompot.de @@ -166,7 +181,9 @@ Resource Estimation Before running resource-intensive differential expression analyses, pass ``dry_run=True`` to estimate memory and disk requirements without running -the actual computation. +the actual computation. It returns a +:class:`~kompot.resource_estimation.ResourcePlan` and prints a per-array +report. .. code-block:: python @@ -176,9 +193,17 @@ the actual computation. condition1="Young", condition2="Old", sample_col="donor_id", + genes=top_genes, dry_run=True, ) + plan.total_memory_required # bytes + plan.total_disk_required # bytes + plan.is_feasible + +How to read the report and what to compare: +:doc:`Planning Memory and Disk `. + Run Tracking ------------ diff --git a/examples/01_getting_started.ipynb b/examples/01_getting_started.ipynb index 5b1e854..3b0f2f2 100644 --- a/examples/01_getting_started.ipynb +++ b/examples/01_getting_started.ipynb @@ -223,7 +223,7 @@ "## Differential Expression\n", "\n", "One call fits the models, scores every gene, and calibrates an FDR against a null\n", - "of shuffled genes. See [`kompot.de`](https://kompot.readthedocs.io/en/latest/simplified.html#kompot.de)." + "of shuffled genes. See [kompot.de](https://kompot.readthedocs.io/en/latest/simplified.html#kompot.de)." ] }, { @@ -586,7 +586,7 @@ "source": [ "### Volcano plot\n", "\n", - "Effect size vs. significance, one point per gene ([`volcano_de`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.volcano_de)):" + "Effect size vs. significance, one point per gene ([volcano_de](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.volcano_de)):" ] }, { @@ -791,7 +791,7 @@ "source": [ "### Where the signal sits\n", "\n", - "Project a top gene's smoothed expression and fold change back onto the UMAP with [`plot_gene_expression`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.plot_gene_expression):" + "Project a top gene's smoothed expression and fold change back onto the UMAP with [plot_gene_expression](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.plot_gene_expression):" ] }, { @@ -901,7 +901,7 @@ "### Fold-change heatmap\n", "\n", "The top genes summarized per cell type. `fold_change_mode=True` shows the\n", - "log fold change directly ([`heatmap`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.heatmap)):" + "log fold change directly ([heatmap](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.heatmap)):" ] }, { @@ -1013,7 +1013,7 @@ "\n", "The same summary with a second dimension: **color** = mean log fold change,\n", "**dot size** = fraction of cells expressing\n", - "([`dotplot`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.dotplot)):" + "([dotplot](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.dotplot)):" ] }, { @@ -2143,7 +2143,7 @@ "id": "cell-027", "metadata": {}, "source": [ - "Turn the enriched terms into a compact lollipop plot ([`lollipop`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.lollipop)):" + "Turn the enriched terms into a compact lollipop plot ([lollipop](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.lollipop)):" ] }, { @@ -2183,7 +2183,7 @@ "### Run history\n", "\n", "Kompot records every run's parameters and environment. Inspect the latest with\n", - "[`RunInfo`](https://kompot.readthedocs.io/en/latest/anndata.html#kompot.anndata.utils.RunInfo):" + "[RunInfo](https://kompot.readthedocs.io/en/latest/anndata.html#kompot.anndata.utils.RunInfo):" ] }, { @@ -2378,7 +2378,7 @@ "source": [ "### Save\n", "\n", - "Smoothed-expression layers are large; [`cleanup`](https://kompot.readthedocs.io/en/latest/anndata.html#kompot.cleanup)\n", + "Smoothed-expression layers are large; [cleanup](https://kompot.readthedocs.io/en/latest/anndata.html#kompot.cleanup)\n", "drops them while keeping the statistics in `adata.var`:" ] }, diff --git a/examples/02_differential_expression_detailed.ipynb b/examples/02_differential_expression_detailed.ipynb index 38cbe92..a8808ab 100644 --- a/examples/02_differential_expression_detailed.ipynb +++ b/examples/02_differential_expression_detailed.ipynb @@ -151,12 +151,23 @@ "Advanced options are grouped into Settings dataclasses; any field left at its default\n", "is the same as omitting it.\n", "\n", - "- **`GPSettings`** — GP model: `sigma` (noise level), `ls_factor` (smoothness of the\n", + "- `GPSettings` — GP model: `sigma` (noise level), `ls_factor` (smoothness of the\n", " expression function — see *Controlling smoothness* below), `n_landmarks`\n", - " (accuracy vs. speed), `batch_size` (peak memory; `0` = all cells at once).\n", - "- **`FDRSettings`** — `null_genes` (permuted genes for FDR; more = better estimate, slower),\n", + " (accuracy vs. speed, and the *quadratic* driver of sample-variance memory),\n", + " `batch_size` (peak memory; `0` = all cells at once).\n", + "- `FDRSettings` — `null_genes` (permuted genes for FDR; more = better estimate, slower),\n", " `threshold` (the FDR used for `is_de`).\n", - "- **`StorageSettings`** — `result_key` (field-name prefix; change it to keep several runs side by side)." + "- `StorageSettings` — `result_key` (field-name prefix; change it to keep several runs\n", + " side by side), `store_arrays_on_disk` (keeps the sample-variance covariance tensors out\n", + " of memory).\n", + "\n", + "One argument is not a Settings field and is in a cost class of its own: `sample_col`.\n", + "It enables sample variance, which gives every gene its own `(n_landmarks, n_landmarks)`\n", + "covariance matrix per condition — about **0.37 GiB per gene** at the default 5 000\n", + "landmarks — and one Cholesky factorisation per gene. It belongs in a second pass over a\n", + "restricted gene list, never in a whole-transcriptome run. See\n", + "[DE with sample variance](./03_sample_variance.ipynb) and\n", + "[Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html).\n" ] }, { @@ -1328,7 +1339,14 @@ "## Resource Planning\n", "\n", "Before a long or memory-heavy job, `dry_run=True` estimates memory, disk, and the\n", - "fields that would be written — without running anything:" + "fields that would be written — without running anything. The returned `ResourcePlan`\n", + "carries `total_memory_required`, `total_disk_required`, `is_feasible` and a per-array\n", + "`requirements` breakdown, so you can price two configurations against each other before\n", + "committing to either.\n", + "\n", + "The estimate resolves `null_genes` exactly as the run would, including the `\"auto\"`\n", + "default, so the plan prices the run it describes. `null_genes=2000` is passed below to\n", + "make the configuration explicit, not because the default would be mispriced.\n" ] }, { diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index ba48efb..e3eeaec 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -15,7 +15,25 @@ "into the Mahalanobis distance \u2014 yielding conservative, replicate-aware calls. The\n", "derivation is in the [paper](https://www.biorxiv.org/content/10.1101/2025.06.03.657769).\n", "\n", - "This dataset has **3 replicates per age group**." + "This dataset has **3 replicates per age group**.\n", + "\n", + "> **Cost warning \u2014 sample variance is priced per gene, twice over.**\n", + "> Without it, every gene is scored against a *single* shared posterior covariance of\n", + "> shape `(n_landmarks, n_landmarks)`. With it, every gene gets **its own** such matrix\n", + "> per condition, and two costs follow that respond to different levers:\n", + ">\n", + "> - **memory**, `2 \u00d7 n_landmarks\u00b2 \u00d7 n_genes \u00d7 8` bytes \u2014 about **0.37 GiB per gene** at\n", + "> the default `n_landmarks=5000`. `StorageSettings(store_arrays_on_disk=True)` keeps\n", + "> this out of RAM entirely, and `n_landmarks` shrinks it quadratically;\n", + "> - **compute**, one Cholesky factorisation **per gene** instead of one in total.\n", + "> `store_arrays_on_disk` does not touch this, but lowering `n_landmarks` does\n", + "> (~0.016 s/gene at 500 against ~2.1 s at 5 000, single-threaded), as does\n", + "> analysing fewer genes.\n", + ">\n", + "> So the prescription is a **cheap first pass over all genes**, then a **sample-variance\n", + "> pass restricted to the top ones**, which is what this notebook does. The levers and\n", + "> measured plans at realistic sizes:\n", + "> [Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html).\n" ] }, { @@ -223,9 +241,10 @@ "id": "cell-066", "metadata": {}, "source": [ - "## Baseline: DE Without Sample Variance\n", + "## Pass 1 \u2014 Baseline DE Without Sample Variance\n", "\n", - "Run standard DE first as the reference (`run_id=0`):" + "Run standard DE over **all genes** first. This is the cheap pass, and it produces the\n", + "Mahalanobis ranking that pass 2 is restricted to (`run_id=0`):\n" ] }, { @@ -308,21 +327,50 @@ "id": "cell-068", "metadata": {}, "source": [ - "## DE With Sample Variance\n", + "## Pass 2 \u2014 DE With Sample Variance, Restricted to the Top Genes\n", "\n", - "Add `sample_col`. Kompot now computes gene-specific covariance matrices, which costs\n", - "memory, so we use three levers:\n", + "Add `sample_col`. Kompot now builds one covariance matrix per gene per condition, so the\n", + "run is priced per gene and we pull every lever available:\n", "\n", - "1. **subset to top genes** from the standard run,\n", - "2. **store arrays on disk** (`StorageSettings`),\n", - "3. **skip FDR** (`null_genes=0`) \u2014 not yet calibrated for sample variance.\n", + "1. **subset to top genes** from pass 1 \u2014 this is the only lever that touches *both* the\n", + " memory and the per-gene factorisation, which is why it is the one that decides whether\n", + " the run is possible at all. Fifty genes keeps this tutorial fast; on a real study, on\n", + " the order of **200** is the working recommendation;\n", + "2. `store_arrays_on_disk=True` (`StorageSettings`) \u2014 Kompot then assembles one gene's\n", + " covariance at a time instead of holding the tensors, which takes memory almost back to\n", + " the no-sample-variance level. With `dask` installed nothing is even written to disk;\n", + "3. **lower** `n_landmarks` if the run is still too slow \u2014 memory is *quadratic* in it, and\n", + " it cuts the per-gene factorisation too, so it buys on both axes;\n", + "4. **skip FDR** (`null_genes=0`) \u2014 not yet calibrated for sample variance, and null genes\n", + " would be charged at the same per-gene rate as real ones. With `sample_col` set this is\n", + " already the default.\n", "\n", - "Check the plan first with `dry_run=True`:" + "**Always check the plan first** with `dry_run=True`, which reports memory, disk, and the\n", + "fields that would be written, without running anything:\n" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "> **The stored outputs of the next two cells were cleared, deliberately.**\n", + ">\n", + "> They were captured on Kompot 0.8.0 and they assert a resource plan the fixed\n", + "> estimator does not produce: the dry run showed an **18.63 GB disk requirement\n", + "> that is now zero** on the Dask path, and the run logged a per-gene Dask path\n", + "> that [LazyGeneCovariance](https://kompot.readthedocs.io/en/latest/resource_planning.html)\n", + "> deliberately no longer takes. A blank cell that says why is better than a\n", + "> number that lies. **Re-run the notebook to regenerate them for your own data.**\n", + ">\n", + "> Every *result* below is unaffected. The 0.9.0 fix is numerically identical \u2014\n", + "> bitwise so on the in-memory and Dask paths \u2014 so the comparison table, the\n", + "> volcano plots and the heatmap are still exactly the ones this notebook was\n", + "> written around.\n" ] }, { "cell_type": "code", - "execution_count": 6, + "execution_count": null, "id": "cell-069", "metadata": { "execution": { @@ -332,66 +380,7 @@ "shell.execute_reply": "2026-07-11T01:32:22.335987Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "================================================================================\n", - "RESOURCE USAGE PLAN\n", - "================================================================================\n", - "\n", - "System Resources:\n", - " Memory: 631.87 GB available (of 754.57 GB total)\n", - " Disk: 361.54 GB available at /tmp\n", - "\n", - "Total Requirements:\n", - " Memory: 1.29 GB (0% of available)\n", - " Disk: 18.63 GB (5% of available)\n", - "\n", - "Memory Allocations:\n", - " \u2022 Mellon precision matrix L (condition 1, 2,917/3,116 cells) (np.int64(2917), np.int64(2917)): 64.92 MB\n", - " \u2022 Mellon precision matrix L (condition 2, 2,917/3,116 cells) (np.int64(3116), np.int64(3116)): 74.08 MB\n", - " \u2022 Smoothed expression (condition 1) (8090, 50): 3.09 MB \u2192 adata.layers['kompot_de_Young_smoothed']\n", - " \u2022 Smoothed expression (condition 2) (8090, 50): 3.09 MB \u2192 adata.layers['kompot_de_Old_smoothed']\n", - " \u2022 Fold change (8090, 50): 3.09 MB \u2192 adata.layers['kompot_de_Young_to_Old_fold_change']\n", - " \u2022 Standard deviation (condition 1) (8090, 50): 3.09 MB \u2192 adata.layers['kompot_de_Young_std']\n", - " \u2022 Standard deviation (condition 2) (8090, 50): 3.09 MB \u2192 adata.layers['kompot_de_Old_std']\n", - " \u2022 Temporary matrices during predictions (batch_size=0) (8090, 5000) + (8090, 50): 311.70 MB\n", - " \u2022 Peak intermediate arrays during predictions (~25 arrays) 25\u00d7(8090, 50): 77.15 MB\n", - " \u2022 Function predictor covariances (per condition) (5000, 5000): 381.47 MB\n", - " \u2022 Combined covariance matrix (5000, 5000): 190.73 MB\n", - " \u2022 Cholesky decomposition (for Mahalanobis) (5000, 5000): 190.73 MB\n", - " \u2022 Per-sample imputations (3 samples) (3, 5000, 50): 11.44 MB\n", - "\n", - "Disk Storage:\n", - " \u2022 Sample covariances (per condition, 3 samples) (5000, 5000, 50): 18.63 GB\n", - "\n", - "Output Fields:\n", - " adata.layers:\n", - " - kompot_de_Young_smoothed [OVERWRITES run_id=0]\n", - " - kompot_de_Old_smoothed [OVERWRITES run_id=0]\n", - " - kompot_de_Young_to_Old_fold_change [OVERWRITES run_id=0]\n", - " - kompot_de_Young_std\n", - " - kompot_de_Old_std\n", - " adata.var:\n", - " - kompot_de_Young_to_Old_mahalanobis_sample_var\n", - " - kompot_de_Young_to_Old_mean_lfc [OVERWRITES run_id=0]\n", - "\n", - "Info:\n", - " \u2139 Cell batch_size (0) processes all 8090 cells at once. Consider reducing via gp=GPSettings(batch_size=500) to lower peak memory by 292.43 MB.\n", - " \u2139 Prediction creates ~25 intermediate arrays of shape (8,090, 50). These coexist at peak memory (77.15 MB) but are freed before completion.\n", - " \u2139 Disk arrays will be stored at system temp: /tmp. To use a different location, set storage=StorageSettings(disk_storage_dir='/path/to/disk') or export TMPDIR=/path/to/disk\n", - "\n", - "Warnings:\n", - " \u26a0 Results with result_key='kompot_de' already exist (run_id=0). Previous run: 2026-07-11T01:32:21.770524 comparing Young to Old (null_genes=2000). Fields that will be overwritten: var.kompot_de_Young_to_Old_mean_lfc, layers.kompot_de_Young_smoothed, layers.kompot_de_Old_smoothed, layers.kompot_de_Young_to_Old_fold_change\n", - "\n", - "================================================================================\n", - "STATUS: \u26a0 FEASIBLE WITH WARNINGS - Proceed with caution\n", - "================================================================================\n" - ] - } - ], + "outputs": [], "source": [ "c1, c2 = CONDITIONS\n", "mahal = f\"kompot_de_{c1}_to_{c2}_mahalanobis\"\n", @@ -422,7 +411,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "cell-071", "metadata": { "execution": { @@ -432,302 +421,7 @@ "shell.execute_reply": "2026-07-11T01:36:30.374795Z" } }, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,345] [INFO ] Differential expression results with result_key='kompot_de' already exist in the dataset. Previous run was at 2026-07-11T01:32:21.770524 comparing Young to Old. Fields that will be overwritten: var.kompot_de_Young_to_Old_mean_lfc, layers.kompot_de_Young_smoothed, layers.kompot_de_Old_smoothed, layers.kompot_de_Young_to_Old_fold_change. Note: Only fields NOT affected by sample variance (like mean_log_fold_change, smoothed data, fold_change) will be overwritten since they don't use the sample variance suffix. These results will likely be identical if other parameters haven't changed. This is a partial rerun with sample variance added to a previous analysis with matching parameters. Set overwrite=False to prevent overwriting or overwrite=True to silence this message.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,347] [INFO ] Condition 1 (Young): 2,917 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,348] [INFO ] Condition 2 (Old): 3,116 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,905] [INFO ] Using sample column 'Replicate' for sample variance estimation\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,907] [INFO ] Found 3 unique sample(s) in condition 1\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:22,919] [INFO ] Found 3 unique sample(s) in condition 2\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:25,014] [INFO ] Fitting expression estimator for condition 1...\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:34,248] [INFO ] Found 3 unique groups for variance estimation\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:34,251] [INFO ] 3 groups have sufficient cells (>= 2) for variance estimation\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:34,252] [INFO ] Training estimator for group 1 with 944 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:32:52,176] [INFO ] Training estimator for group 2 with 1,007 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:08,107] [INFO ] Training estimator for group 3 with 966 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:25,634] [INFO ] Fitting expression estimator for condition 2...\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:39,761] [INFO ] Found 3 unique groups for variance estimation\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:39,763] [INFO ] 3 groups have sufficient cells (>= 2) for variance estimation\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:39,764] [INFO ] Training estimator for group 1 with 1,057 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:33:56,588] [INFO ] Training estimator for group 2 with 1,019 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:12,639] [INFO ] Training estimator for group 3 with 1,040 cells\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:51,054] [INFO ] Disk space at /tmp/kompot_arrays_014f3181_cuguwm9t:\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:51,057] [INFO ] Total: 377.28 GB | Used: 15.74 GB | Free: 361.54 GB\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:51,058] [INFO ] Expected storage need: 11.18 GB\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:51,059] [INFO ] Initialized disk storage at /tmp/kompot_arrays_014f3181_cuguwm9t with dask support (namespace: ns_014f3181)\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:51,060] [INFO ] Using gene-by-gene disk storage for covariance matrix (shape=(5000, 5000, 50))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:52,923] [INFO ] Using Dask for parallel disk-backed covariance computation (all cores) (shape=(5000, 5000, 50))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,738] [INFO ] Created Dask array for covariance with shape (5000, 5000, 50)\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,741] [INFO ] Disk space at /tmp/kompot_arrays_9ecd22cb_pt9wuuyh:\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,741] [INFO ] Total: 377.28 GB | Used: 15.74 GB | Free: 361.54 GB\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,742] [INFO ] Expected storage need: 11.18 GB\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,742] [INFO ] Initialized disk storage at /tmp/kompot_arrays_9ecd22cb_pt9wuuyh with dask support (namespace: ns_9ecd22cb)\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:53,742] [INFO ] Using gene-by-gene disk storage for covariance matrix (shape=(5000, 5000, 50))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:54,990] [INFO ] Using Dask for parallel disk-backed covariance computation (all cores) (shape=(5000, 5000, 50))\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:34:56,309] [INFO ] Created Dask array for covariance with shape (5000, 5000, 50)\n" - ] - }, - { - "data": { - "application/vnd.jupyter.widget-view+json": { - "model_id": "8a0d86bad2ad4f4cbcf5e60d5bd74506", - "version_major": 2, - "version_minor": 0 - }, - "text/plain": [ - "Processing gene-specific covariance matrices: 0%| | 0/50 [00:00LIL->CSR to avoid slow per-gene modifications.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:36:30,134] [INFO ] This run will have `run_id=1`.\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:36:30,169] [INFO ] Removed temporary storage directory /tmp/kompot_arrays_014f3181_cuguwm9t\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "[2026-07-11 01:36:30,284] [INFO ] Removed temporary storage directory /tmp/kompot_arrays_9ecd22cb_pt9wuuyh\n" - ] - } - ], + "outputs": [], "source": [ "de_sample_var = kompot.de(\n", " adata,\n", @@ -1343,11 +1037,28 @@ "## Best Practices\n", "\n", "**Use sample variance when** you have \u22653 replicates per condition and want conservative,\n", - "reproducible calls. For DE specifically: always `dry_run` first, subset to top genes,\n", - "store arrays on disk, and set `null_genes=0` unless you need FDR.\n", + "reproducible calls. Effect sizes stay identical; only confidence changes, and the genes\n", + "that survive the correction are the robust, replicate-consistent signals.\n", + "\n", + "**Always run it as a second pass.** The levers, and what each one actually buys:\n", + "\n", + "| Lever | Memory | Per-gene compute | Practical setting |\n", + "|---|---|---|---|\n", + "| `genes=` | **linear** | **linear** | top ~200 from pass 1 |\n", + "| `StorageSettings(store_arrays_on_disk=True)` | removes the tensor term | \u2014 | set `disk_storage_dir` to real scratch |\n", + "| `GPSettings(n_landmarks=\u2026)` | **quadratic** | reduces it; measure rather than extrapolate | 1 000\u20132 000 when 5 000 is too slow |\n", + "| `FDRSettings(null_genes=0)` | removes 2 000 genes' worth | likewise | already the default once `sample_col` is set |\n", + "\n", + "Note what the table says about `genes`: it is the only row that appears in both cost\n", + "columns. Offloading solves the memory axis outright \u2014 a whole-transcriptome run goes\n", + "from 7 552 GiB to 101 GiB \u2014 and leaves the per-gene factorisation untouched, which is a\n", + "long job rather than an impossible one (~11.6 h at 5 000 landmarks, single-threaded).\n", + "So restricting the gene list is about spending your compute on the genes you will\n", + "actually read, not about clearing a wall.\n", "\n", - "Effect sizes stay identical; only confidence changes. Genes that survive the correction\n", - "are the robust, replicate-consistent signals." + "And price it before you run it: `dry_run=True` on the exact call you intend to make takes\n", + "a second. Full treatment, with measured plans:\n", + "[Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html).\n" ] }, { diff --git a/examples/04_differential_abundance.ipynb b/examples/04_differential_abundance.ipynb index c3d8202..5b1dab3 100644 --- a/examples/04_differential_abundance.ipynb +++ b/examples/04_differential_abundance.ipynb @@ -766,7 +766,7 @@ "\n", "Now read the `p` column, because that is where the count starts costing you:\n", "\n", - "- **Only 5 of 20 cell types clear `p < 0.05`, and the *largest* effect in the table is not one\n", + "- Only 5 of 20 cell types clear `p < 0.05`, **and the largest effect in the table is not one\n", " of them.** Naive CD8 T cells fall from 5.0% of the mature compartment to 0.05% — a ~90-fold\n", " depletion — at `p = 0.08`, and CD8 TEM cells rise ~25-fold at `p = 0.07`. (The second-largest\n", " effect, HSC, does clear it, at `p = 0.02`.) A per-cluster proportion crushes every cell into a\n", @@ -779,7 +779,7 @@ "\n", "`kompot.da` never sees the labels. It estimates the cell-state density of each condition\n", "over `DM_EigenVectors` and scores **every cell**.\n", - "See [`kompot.da`](https://kompot.readthedocs.io/en/latest/simplified.html#kompot.da)." + "See [kompot.da](https://kompot.readthedocs.io/en/latest/simplified.html#kompot.da)." ] }, { @@ -895,7 +895,7 @@ "source": [ "### Volcano plot\n", "\n", - "Effect size vs. significance, one point per cell ([`volcano_da`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.volcano_da)):" + "Effect size vs. significance, one point per cell ([volcano_da](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.volcano_da)):" ] }, { @@ -960,7 +960,7 @@ "metadata": {}, "source": [ "Show only the significantly changed cells, split into enriched vs. depleted with `mgroups`\n", - "([`kompot.plot.embedding`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.embedding)):" + "([kompot.plot.embedding](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.embedding)):" ] }, { @@ -1018,7 +1018,7 @@ "### Summary by cell type\n", "\n", "The per-cell calls can of course still be *summarized* over an annotation — that is a\n", - "read-out, not an input ([`direction_barplot`](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.direction_barplot)):" + "read-out, not an input ([direction_barplot](https://kompot.readthedocs.io/en/latest/plotting.html#kompot.plot.direction_barplot)):" ] }, { diff --git a/examples/05_smooth_expression.ipynb b/examples/05_smooth_expression.ipynb index dcb4c9d..2f3e86f 100644 --- a/examples/05_smooth_expression.ipynb +++ b/examples/05_smooth_expression.ipynb @@ -14,7 +14,7 @@ "themselves scatter around it (the empirical variance, driven by the gene). Section 2 separates\n", "them.\n", "\n", - "This is the **same machinery `kompot.de` uses internally**: differential expression fits one\n", + "This is the **same machinery** that `kompot.de` uses internally: differential expression fits one\n", "of these functions per condition and contrasts them. Using the primitive directly unlocks\n", "two things DE wraps up: reading the **uncertainties**, and\n", "**training on some cells then predicting on others** (imputation, modality transfer,\n", @@ -24,7 +24,7 @@ "\n", "**Dive deeper:**\n", "[Getting Started / DE](./01_getting_started.ipynb) ·\n", - "[Advanced DE + `ls_factor`](./02_differential_expression_detailed.ipynb) ·\n", + "[Advanced DE and ls_factor](./02_differential_expression_detailed.ipynb) ·\n", "[Sample variance](./03_sample_variance.ipynb) ·\n", "[Paper](https://www.biorxiv.org/content/10.1101/2025.06.03.657769)" ] @@ -739,7 +739,7 @@ "\n", "The model is *fit* on the cells you give it but can be *evaluated* at any coordinates, so it\n", "doubles as a fit/predict surface. The bare-array estimator underneath `smooth_expression` is\n", - "[`ExpressionModel`](https://kompot.readthedocs.io/en/latest/api.html):\n", + "[ExpressionModel](https://kompot.readthedocs.io/en/latest/api.html):\n", "\n", "- `ExpressionModel().fit(X, y)` — learn the function from cell-state coordinates\n", " `X` `(n_cells, n_dims)` and values `y` `(n_cells, n_features)`.\n", @@ -1030,7 +1030,7 @@ "## Where to go deeper\n", "\n", "- **[Getting Started / DE](./01_getting_started.ipynb)** — the front door; `smooth_expression` is the primitive underneath.\n", - "- **[Advanced DE + `ls_factor`](./02_differential_expression_detailed.ipynb)** — the length-scale trade-off in the DE setting.\n", + "- **[Advanced DE and ls_factor](./02_differential_expression_detailed.ipynb)** — the length-scale trade-off in the DE setting.\n", "- **[DE with sample variance](./03_sample_variance.ipynb)** — per-gene aleatoric and replicate uncertainty (this tutorial shows the GP/epistemic term).\n", "- **[Differential abundance](./04_differential_abundance.ipynb)** — the same GP geometry applied to cell-state frequency.\n", "- **[Paper](https://www.biorxiv.org/content/10.1101/2025.06.03.657769)** ·\n", diff --git a/kompot/_provenance.py b/kompot/_provenance.py index 06f32cc..39e8525 100644 --- a/kompot/_provenance.py +++ b/kompot/_provenance.py @@ -65,35 +65,83 @@ def _find_git_dir(start: str) -> Optional[str]: current = parent -def _resolve_sha(git_dir: str) -> Optional[str]: - """Resolve HEAD to a full sha by reading git's on-disk refs.""" - head = _read_text(os.path.join(git_dir, "HEAD")) - if not head: +def _find_kompot_git_dir(package_dir: str) -> Optional[str]: + """The git dir whose work-tree root is the package's parent, or ``None``. + + ``_find_git_dir`` walks up to *any* enclosing ``.git``. A Kompot source + tree vendored into an unrelated repository, nested below its root, would + then be stamped with that repository's HEAD -- a wrong sha, which is + worse than none. In Kompot's own checkout the parent of the package + directory is the repository root, so the ``.git`` entry (a directory, or + a file for a worktree or submodule) must sit exactly there. + + Residual: a ``kompot/`` package copied to the *root* of another + repository passes this check, since by position it is indistinguishable + from Kompot's own checkout. + """ + repo_root = os.path.dirname(os.path.abspath(package_dir)) + if not os.path.exists(os.path.join(repo_root, ".git")): return None + return _find_git_dir(repo_root) - if not head.startswith("ref:"): - # Detached HEAD already holds the sha. - return head if _looks_like_sha(head) else None - - ref = head[len("ref:") :].strip() - # Loose ref. - loose = _read_text(os.path.join(git_dir, *ref.split("/"))) - if loose and _looks_like_sha(loose): - return loose +def _common_dir(git_dir: str) -> str: + """The directory holding the refs shared by every work tree of a repository. - # Packed ref. - packed = _read_text(os.path.join(git_dir, "packed-refs")) - if packed: + A linked worktree's git dir (``
/.git/worktrees/``) holds only + per-worktree state such as ``HEAD``; branch refs and ``packed-refs`` live + in the main repository's git dir, which a ``commondir`` file names, + usually relatively (settylab/kompot#28, #33). A plain repository has no + ``commondir`` and is its own common dir. + """ + common = _read_text(os.path.join(git_dir, "commondir")) + if not common: + return git_dir + if not os.path.isabs(common): + common = os.path.join(git_dir, common) + return os.path.normpath(common) + + +def _lookup_ref(ref: str, search_dirs) -> Optional[str]: + """Resolve *ref* to a sha or a further ``ref:`` target, loose before packed.""" + for directory in search_dirs: + loose = _read_text(os.path.join(directory, *ref.split("/"))) + if loose: + return loose + for directory in search_dirs: + packed = _read_text(os.path.join(directory, "packed-refs")) + if not packed: + continue for line in packed.splitlines(): if not line or line.startswith(("#", "^")): continue sha, _, name = line.partition(" ") - if name.strip() == ref and _looks_like_sha(sha): + if name.strip() == ref: return sha return None +def _resolve_sha(git_dir: str) -> Optional[str]: + """Resolve HEAD to a full sha by reading git's on-disk refs. + + ``HEAD`` is read from *git_dir*, which is per-worktree, while refs are + looked up there and then in the common dir, where a linked worktree's + branch refs actually live. + """ + value = _read_text(os.path.join(git_dir, "HEAD")) + common = _common_dir(git_dir) + search_dirs = [git_dir] if common == git_dir else [git_dir, common] + + # A symbolic ref may point at another symbolic ref; bound the chain. + for _ in range(5): + if not value: + return None + if not value.startswith("ref:"): + return value if _looks_like_sha(value) else None + value = _lookup_ref(value[len("ref:") :].strip(), search_dirs) + return None + + def _looks_like_sha(value: str) -> bool: return len(value) == 40 and all(c in "0123456789abcdef" for c in value.lower()) @@ -149,9 +197,19 @@ def _resolve() -> Dict[str, Any]: provenance["kompot_editable"] = False if editable is None else editable return provenance - git_dir = _find_git_dir(package_dir) + git_dir = _find_kompot_git_dir(package_dir) if git_dir is not None: provenance["kompot_git_sha"] = _resolve_sha(git_dir) + if provenance["kompot_git_sha"] is None: + # Inside a work tree a missing sha is a defect in resolution, not + # an expected absence, and every run stamped from here on carries + # it; say so once, at resolution time, rather than months later. + logger.warning( + "Kompot is running from a git checkout (%s) but its commit sha " + "could not be resolved; run history will record " + "kompot_git_sha=None.", + git_dir, + ) if editable is not None: provenance["kompot_editable"] = editable diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index 74eb2c3..802a65c 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -81,6 +81,39 @@ def de( equivalent to omitting it entirely. Extra ``**function_kwargs`` are forwarded to mellon's :class:`~mellon.FunctionEstimator`. + .. warning:: + + Supplying ``sample_col`` **multiplies the cost by the gene count.** + Sample variance replaces the single shared posterior covariance with + one ``(n_landmarks, n_landmarks)`` covariance matrix *per gene, per + condition*. Two costs follow, and they respond to different levers: + + * **memory**, ``2 * n_landmarks**2 * n_genes * 8`` bytes, roughly + 0.37 GiB per gene at the default ``n_landmarks=5000``. + ``StorageSettings(store_arrays_on_disk=True)`` keeps these out of + memory entirely, and ``n_landmarks`` shrinks them quadratically; + * **compute**, one Cholesky factorisation *per gene* instead of one in + total. ``GPSettings.batch_size`` does not bound that loop, and + ``store_arrays_on_disk`` does not touch it; lowering ``n_landmarks`` + does reduce it (measured ~0.016 s/gene at 500 against ~2.1 s at + 5 000, single-threaded), as does analysing fewer genes. + + Run it as a **second pass** over a restricted gene list:: + + kompot.de(adata, "condition", "Young", "Old") # pass 1, all genes + + mahal = "kompot_de_Young_to_Old_mahalanobis" + top = adata.var.sort_values(mahal, ascending=False).head(200).index + + kompot.de(adata, "condition", "Young", "Old", # pass 2 + sample_col="donor_id", genes=top, + gp=kompot.GPSettings(n_landmarks=2000)) + + Both costs are linear in ``genes``; memory is additionally + **quadratic** in ``n_landmarks``. Price any configuration first with + ``dry_run=True``. Full treatment: + https://kompot.readthedocs.io/en/latest/resource_planning.html + Parameters ---------- adata : AnnData @@ -96,7 +129,9 @@ def de( genes : list of str, optional Subset of genes to analyse. sample_col : str, optional - Column with biological-replicate labels. + Column with biological-replicate labels. Supplying it enables + **sample variance**, which is by far the most expensive option in + this function: read the warning above before using it. gp : GPSettings, optional GP model parameters (sigma, ls, n_landmarks, etc.). ``param_scheme`` sets which cells the shared length scale is @@ -118,7 +153,11 @@ def de( dry_run : bool, optional If True, estimate resource requirements and print a report instead of running the analysis. Returns a - :class:`~kompot.resource_estimation.ResourcePlan`. + :class:`~kompot.resource_estimation.ResourcePlan` carrying + ``total_memory_required``, ``total_disk_required``, ``is_feasible`` + and a per-array ``requirements`` list. The estimate resolves + ``FDRSettings.null_genes`` exactly as the run would, including the + ``"auto"`` default, so the plan prices the run it describes. **function_kwargs Forwarded to :class:`~mellon.FunctionEstimator`. @@ -228,6 +267,20 @@ def de( stacklevel=2, ) + # ---- 0. Resolve defaults ---- + # Resolved BEFORE the dry-run branch: the estimate must price the run that + # would actually happen. Resolving it afterwards left the estimator reading + # the literal string "auto", which matches neither its int nor its list + # branch, so it silently counted zero null genes (settylab/kompot#25). + if null_genes == "auto": + if sample_col is not None: + null_genes = 0 + logger.info( + "Defaulting null_genes=0 (FDR disabled) because sample_col is provided." + ) + else: + null_genes = 2000 + # ---- dry run ---- if dry_run: from ..resource_estimation import estimate_differential_expression_resources @@ -271,16 +324,6 @@ def de( print(plan.format_report()) return plan - # ---- 0. Resolve defaults ---- - if null_genes == "auto": - if sample_col is not None: - null_genes = 0 - logger.info( - "Defaulting null_genes=0 (FDR disabled) because sample_col is provided." - ) - else: - null_genes = 2000 - has_external_null = ( ext_null_mahalanobis is not None or ext_null_expression is not None ) diff --git a/kompot/cli/da.py b/kompot/cli/da.py index ce2b3d9..c5fff1c 100644 --- a/kompot/cli/da.py +++ b/kompot/cli/da.py @@ -9,7 +9,12 @@ from ..anndata import da from ..settings import GPSettings, DAThresholdSettings, StorageSettings, OutputSettings -from .utils import load_config, merge_args_with_config, validate_anndata_path +from .utils import ( + load_config, + merge_args_with_config, + validate_anndata_path, + write_output, +) from .compute_config import configure_compute @@ -367,15 +372,7 @@ def run_da(args): output_path = Path(args.output) logger.info(f"Saving results to {output_path}") - if str(output_path).endswith(".h5ad"): - adata.write_h5ad(output_path) - elif str(output_path).endswith(".zarr"): - adata.write_zarr(output_path) - else: - logger.error( - f"Unsupported output format: {output_path.suffix}. Use .h5ad or .zarr" - ) - sys.exit(1) + write_output(adata, output_path) # Save table output if specified if args.table_output: diff --git a/kompot/cli/de.py b/kompot/cli/de.py index 5ee951e..21088b8 100644 --- a/kompot/cli/de.py +++ b/kompot/cli/de.py @@ -16,7 +16,12 @@ StorageSettings, OutputSettings, ) -from .utils import load_config, merge_args_with_config, validate_anndata_path +from .utils import ( + load_config, + merge_args_with_config, + validate_anndata_path, + write_output, +) from .compute_config import configure_compute @@ -123,7 +128,15 @@ def add_de_parser(subparsers) -> argparse.ArgumentParser: parser.add_argument( "--sample-col", type=str, - help="Column in adata.obs with sample labels for sample variance estimation", + help=( + "Column in adata.obs with sample labels for sample variance " + "estimation. EXPENSIVE: gives every gene its own " + "(n_landmarks, n_landmarks) covariance matrix and its own Cholesky " + "factorisation -- ~0.37 GiB and ~2 s per gene at " + "--n-landmarks 5000. Run it as a second pass over a top-gene list " + "(config key 'genes'), set store_arrays_on_disk to keep the " + "tensors out of memory, and check with --dry-run first." + ), ) parser.add_argument( @@ -433,15 +446,7 @@ def run_de(args): output_path = Path(args.output) logger.info(f"Saving results to {output_path}") - if str(output_path).endswith(".h5ad"): - adata.write_h5ad(output_path) - elif str(output_path).endswith(".zarr"): - adata.write_zarr(output_path) - else: - logger.error( - f"Unsupported output format: {output_path.suffix}. Use .h5ad or .zarr" - ) - sys.exit(1) + write_output(adata, output_path) # Save table output if specified if args.table_output: diff --git a/kompot/cli/dm.py b/kompot/cli/dm.py index 6712247..8bae94f 100644 --- a/kompot/cli/dm.py +++ b/kompot/cli/dm.py @@ -7,7 +7,12 @@ import anndata as ad -from .utils import load_config, merge_args_with_config, validate_anndata_path +from .utils import ( + load_config, + merge_args_with_config, + validate_anndata_path, + write_output, +) logger = logging.getLogger("kompot.cli") @@ -214,14 +219,6 @@ def run_dm(args): output_path = Path(args.output) logger.info(f"Saving results to {output_path}") - if str(output_path).endswith(".h5ad"): - adata.write_h5ad(output_path) - elif str(output_path).endswith(".zarr"): - adata.write_zarr(output_path) - else: - logger.error( - f"Unsupported output format: {output_path.suffix}. Use .h5ad or .zarr" - ) - sys.exit(1) + write_output(adata, output_path) logger.info("Diffusion maps computation completed successfully") diff --git a/kompot/cli/smooth.py b/kompot/cli/smooth.py index 7e1d481..17402ec 100644 --- a/kompot/cli/smooth.py +++ b/kompot/cli/smooth.py @@ -9,7 +9,12 @@ from ..anndata.smooth import smooth_expression from ..settings import GPSettings, StorageSettings, OutputSettings -from .utils import load_config, merge_args_with_config, validate_anndata_path +from .utils import ( + load_config, + merge_args_with_config, + validate_anndata_path, + write_output, +) from .compute_config import configure_compute @@ -328,15 +333,7 @@ def run_smooth(args): output_path = Path(args.output) logger.info(f"Saving results to {output_path}") - if str(output_path).endswith(".h5ad"): - adata.write_h5ad(output_path) - elif str(output_path).endswith(".zarr"): - adata.write_zarr(output_path) - else: - logger.error( - f"Unsupported output format: {output_path.suffix}. Use .h5ad or .zarr" - ) - sys.exit(1) + write_output(adata, output_path) # Save table output if specified if args.table_output: diff --git a/kompot/cli/templates/de_config_template.yaml b/kompot/cli/templates/de_config_template.yaml index 6697c4a..2d4a5b7 100644 --- a/kompot/cli/templates/de_config_template.yaml +++ b/kompot/cli/templates/de_config_template.yaml @@ -23,6 +23,14 @@ landmarks: null # Pre-computed landmarks array (overrides n_landm store_landmarks: false # Whether to store landmarks in output # Sample variance (for biological replicates): +# COST: setting sample_col gives EVERY gene its own (n_landmarks, n_landmarks) +# covariance matrix per condition -- 2 * n_landmarks^2 * n_genes * 8 bytes, +# about 0.37 GiB per gene at n_landmarks=5000 -- plus one Cholesky +# factorisation per gene. Run it as a SECOND PASS over a restricted `genes` +# list (below), set store_arrays_on_disk to keep the tensors out of memory, +# lower n_landmarks (memory is quadratic in it), and check with --dry-run +# first. +# https://kompot.readthedocs.io/en/latest/resource_planning.html sample_col: null # Column in adata.obs with sample IDs allow_single_condition_variance: false # Allow variance with samples in only one condition @@ -48,9 +56,12 @@ param_scheme: null # Where the shared length scale is estimated from eps: 1.0e-8 # Numerical stability constant # Memory management: -store_arrays_on_disk: null # Store large arrays on disk (null = auto) -disk_storage_dir: null # Directory for disk storage -max_memory_ratio: 0.8 # Max memory usage ratio +store_arrays_on_disk: null # Consume per-gene covariances one gene at a + # time (null = true iff disk_storage_dir is set) +disk_storage_dir: null # Directory for disk storage; create it first -- + # --dry-run errors on a missing path + # (null = system temp, honours TMPDIR) +max_memory_ratio: 0.8 # Warning threshold only; does not switch modes # Output options: store_additional_stats: false # Store extra statistics @@ -59,7 +70,9 @@ progress: true # Show progress bars overwrite: null # Overwrite existing results (null = warn) # Advanced filtering: -genes: null # List of gene names to analyze (null = all) +genes: null # List of gene names to analyze (null = all). + # The lever for a sample_col run: restrict to + # the top ~200 genes of a first pass. min_cells: 2 # Minimum cells per gene min_percentage: null # Minimum percentage of cells expressing diff --git a/kompot/cli/utils.py b/kompot/cli/utils.py index 947989f..4b38936 100644 --- a/kompot/cli/utils.py +++ b/kompot/cli/utils.py @@ -4,12 +4,48 @@ import json from pathlib import Path from typing import Dict, Any +import contextlib import logging +import sys logger = logging.getLogger("kompot.cli") +def _nullable_strings_allowed(): + """Opt in to writing pandas nullable-string arrays, where anndata gates it. + + anndata 0.11 and 0.12 refuse to write a ``pd.arrays.StringArray`` (which + pandas 3 produces for ordinary string columns) unless + ``anndata.settings.allow_write_nullable_strings`` is set, so every CLI + write failed under their defaults (settylab/kompot#23). The opt-in is + scoped to the write rather than set at import, and is a no-op on anndata + versions without the setting. + """ + import anndata + + settings = getattr(anndata, "settings", None) + if settings is None or not hasattr(settings, "allow_write_nullable_strings"): + return contextlib.nullcontext() + return settings.override(allow_write_nullable_strings=True) + + +def write_output(adata, output_path) -> None: + """Write *adata* to an ``.h5ad`` or ``.zarr`` path, exiting on any other suffix.""" + output_path = Path(output_path) + if str(output_path).endswith(".h5ad"): + with _nullable_strings_allowed(): + adata.write_h5ad(output_path) + elif str(output_path).endswith(".zarr"): + with _nullable_strings_allowed(): + adata.write_zarr(output_path) + else: + logger.error( + f"Unsupported output format: {output_path.suffix}. Use .h5ad or .zarr" + ) + sys.exit(1) + + def load_config(config_path: str) -> Dict[str, Any]: """ Load configuration from YAML or JSON file. diff --git a/kompot/differential/differential_expression.py b/kompot/differential/differential_expression.py index 86c0cd7..e321382 100644 --- a/kompot/differential/differential_expression.py +++ b/kompot/differential/differential_expression.py @@ -1,14 +1,12 @@ """Differential expression analysis.""" import numpy as np -import jax from scipy.stats import chi2 as scipy_chi2 from typing import Optional, Dict, Any import logging from mellon.parameters import compute_landmarks -from tqdm.auto import tqdm -from ..utils import compute_mahalanobis_distances +from ..utils import compute_mahalanobis_distances, LazyGeneCovariance from ..batch_utils import apply_batched, is_jax_memory_error # noqa: F401 from ..settings import PARAM_SCHEMES, resolve_param_scheme from .expression_model import ExpressionModel @@ -846,121 +844,71 @@ def compute_mahalanobis_distances( gene_specific_covariance = None if self.use_sample_variance: - # Add empirical adjustments from sample variance - - # Create functions for computing sample variance + # Add the sample-variance covariance to the posterior covariance. + # + # When the predictors return 3-D tensors there is one + # (n_points, n_points) covariance PER GENE per condition. Summing + # them eagerly, and then adding combined_cov gene by gene with + # __setitem__, materialised a third dense tensor of that shape -- + # even when the inputs were disk-backed, which is precisely what + # store_arrays_on_disk exists to avoid, and which rebuilt the Dask + # graph once per gene (settylab/kompot#26). + # + # LazyGeneCovariance keeps references to the terms instead and + # assembles one gene's matrix at a time inside + # compute_mahalanobis_distances, so peak memory is bounded by a + # single (n_points, n_points) matrix rather than by the tensor. + variance_terms = [] if self.variance_predictor1 is not None: try: # Important: use diag=False to get full covariance matrix - variance1 = self.variance_predictor1( - variance_points, diag=False, progress=progress - ) - if self.variance_predictor2 is not None: - variance2 = self.variance_predictor2( + variance_terms.append( + self.variance_predictor1( variance_points, diag=False, progress=progress ) - # Add the covariance matrices for complete variance representation - combined_variance = variance1 + variance2 - del variance1, variance2 - - # Check if we have gene-specific covariance matrices (shape has 3 dimensions) - if len(combined_variance.shape) == 3: - # We have per-gene covariance matrices with shape (points, points, genes) - # Need to add combined_cov to each gene's covariance slice - gene_specific_covariance = combined_variance - # Check if combined_variance is a JAX array, if not, ensure combined_cov is numpy array - if not isinstance(combined_variance, jax.Array): - combined_cov_to_add = np.asarray(combined_cov) - else: - combined_cov_to_add = combined_cov - for g in tqdm( - range(combined_variance.shape[2]), - desc="Processing gene-specific covariance matrices", - disable=not progress, - ): - gene_specific_covariance[:, :, g] = ( - combined_variance[:, :, g] + combined_cov_to_add - ) - logger.debug( - f"Using gene-specific covariance matrices with shape {gene_specific_covariance.shape}" - ) - else: - # Add the sample variance to the combined covariance from function predictors - combined_cov += combined_variance - logger.debug( - "Added sample variance covariance matrix to function predictor covariance" - ) - else: - # Only add variance1 if variance2 is not available - if len(variance1.shape) == 3: - # We have per-gene covariance matrices - # Need to add combined_cov to each gene's covariance slice - gene_specific_covariance = variance1 - # Check if variance1 is a JAX array, if not, ensure combined_cov is numpy array - if not isinstance(variance1, jax.Array): - combined_cov_to_add = np.asarray(combined_cov) - else: - combined_cov_to_add = combined_cov - for g in tqdm( - range(variance1.shape[2]), - desc="Processing gene-specific covariance matrices (variance1)", - disable=not progress, - ): - gene_specific_covariance[:, :, g] = ( - variance1[:, :, g] + combined_cov_to_add - ) - logger.debug( - f"Using gene-specific covariance matrices from variance1 with shape {gene_specific_covariance.shape}" - ) - else: - combined_cov += variance1 - logger.debug( - "Added variance1 covariance matrix to function predictor covariance" - ) - del variance1 + ) except Exception as e: error_msg = f"Error computing sample variance from variance_predictor1: {e}." logger.error(error_msg) raise RuntimeError(error_msg) from e - elif self.variance_predictor2 is not None: + + if self.variance_predictor2 is not None: try: - # Important: use diag=False to get full covariance matrix - variance2 = self.variance_predictor2( - variance_points, diag=False, progress=progress - ) - # Check if we have gene-specific covariance matrices - if len(variance2.shape) == 3: - # We have per-gene covariance matrices - # Need to add combined_cov to each gene's covariance slice - gene_specific_covariance = variance2 - # Check if variance2 is a JAX array, if not, ensure combined_cov is numpy array - if not isinstance(variance2, jax.Array): - combined_cov_to_add = np.asarray(combined_cov) - else: - combined_cov_to_add = combined_cov - for g in tqdm( - range(variance2.shape[2]), - desc="Processing gene-specific covariance matrices (variance2)", - disable=not progress, - ): - gene_specific_covariance[:, :, g] = ( - variance2[:, :, g] + combined_cov_to_add - ) - logger.debug( - f"Using gene-specific covariance matrices from variance2 with shape {gene_specific_covariance.shape}" - ) - else: - # Add variance2 to the combined covariance - combined_cov += variance2 - logger.debug( - "Added variance2 covariance matrix to function predictor covariance" + variance_terms.append( + self.variance_predictor2( + variance_points, diag=False, progress=progress ) - del variance2 + ) except Exception as e: error_msg = f"Error computing sample variance from variance_predictor2: {e}." logger.error(error_msg) raise RuntimeError(error_msg) from e + if variance_terms: + gene_specific_terms = [v for v in variance_terms if len(v.shape) == 3] + shared_terms = [v for v in variance_terms if len(v.shape) != 3] + + for shared in shared_terms: + # Not gene-specific: fold straight into the shared matrix. + combined_cov = combined_cov + shared + logger.debug( + "Added a shared sample-variance covariance matrix to the " + "function predictor covariance" + ) + + if gene_specific_terms: + gene_specific_covariance = LazyGeneCovariance( + gene_specific_terms, base=np.asarray(combined_cov) + ) + logger.debug( + f"Using lazily assembled gene-specific covariance matrices " + f"with shape {gene_specific_covariance.shape} " + f"({len(gene_specific_terms)} sample-variance term(s) + " + f"the shared posterior covariance, one gene materialized " + f"at a time)" + ) + del variance_terms, gene_specific_terms, shared_terms + # Compute empirical variance at variance_points if enabled empirical_diag_var = None if ( diff --git a/kompot/differential/sample_variance_estimator.py b/kompot/differential/sample_variance_estimator.py index b948699..e760d80 100644 --- a/kompot/differential/sample_variance_estimator.py +++ b/kompot/differential/sample_variance_estimator.py @@ -4,6 +4,7 @@ import jax import jax.numpy as jnp import logging +import warnings from typing import Dict, Optional from mellon import FunctionEstimator, DensityEstimator from tqdm.auto import tqdm @@ -32,6 +33,21 @@ class SampleVarianceEstimator: predictions. Bessel's correction is applied to the variance calculation to ensure unbiased estimation, especially important when the number of samples is small. + .. warning:: + + For gene expression (``estimator_type='function'``) with more than one + gene, ``predict(..., diag=False)`` returns a **covariance matrix per + gene**, of shape ``(n_points, n_points, n_genes)`` and dtype float64. + In :func:`kompot.de` the evaluation points are the landmarks, so the + array is ``n_landmarks**2 * n_genes * 8`` bytes per condition: about + 190 MiB per gene at the default 5 000 landmarks. These are never + summed into a dense array -- :class:`kompot.utils.LazyGeneCovariance` + assembles one gene at a time -- and ``store_arrays_on_disk`` keeps them + out of memory altogether. What no storage mode removes is the Cholesky + factorisation per gene, which is why sample variance is prescribed as a + second pass over a restricted gene list. See + https://kompot.readthedocs.io/en/latest/resource_planning.html + Attributes ---------- group_predictors : Dict @@ -74,16 +90,26 @@ def __init__( store_arrays_on_disk will be set to True. If store_arrays_on_disk is False and this is provided, a warning will be logged and disk storage will not be used. dask_num_workers : int, optional - Number of parallel Dask workers to use for covariance computation. If None (default), - Dask uses all available CPU cores for maximum speed. Set to a smaller number (e.g., 2-4) - to limit CPU utilization at the cost of slower computation. Only applies when Dask is - available and disk storage is enabled. Note: This primarily controls CPU usage, not memory, - since the main memory footprint comes from the shared centered data array rather than - per-worker allocations. + Deprecated and ignored; passing a value emits a ``FutureWarning``. + The Dask covariance tensor is returned lazily and, since 0.9.0, + materialised one gene at a time, so there is no worker pool for + this setting to bound. It previously wrote a Dask config key the + scheduler never reads (settylab/kompot#31). To bound Dask's + threaded scheduler for your own computations, use + ``with dask.config.set(num_workers=N): ...``. """ self.eps = eps self.jit_compile = jit_compile self.estimator_type = estimator_type + if dask_num_workers is not None: + warnings.warn( + "SampleVarianceEstimator(dask_num_workers=...) has no effect and is " + "deprecated: the Dask covariance tensor is evaluated one gene at a " + "time, so there is no worker pool to bound. To bound Dask's threaded " + "scheduler, use `with dask.config.set(num_workers=N): ...`.", + FutureWarning, + stacklevel=2, + ) self.dask_num_workers = dask_num_workers # Determine store_arrays_on_disk based on disk_storage_dir if not explicitly set @@ -456,13 +482,8 @@ def compute_variance_from_predictions(X, predictions_list): if use_disk_storage: # Disk-backed version - use parallel computation with Dask if available if DASK_AVAILABLE: - worker_info = ( - f" with {self.dask_num_workers} workers" - if self.dask_num_workers - else " (all cores)" - ) logger.info( - f"Using Dask for parallel disk-backed covariance computation{worker_info} (shape={covariance_shape})" + f"Using Dask for lazy disk-backed covariance computation (shape={covariance_shape})" ) # Convert JAX to numpy once to avoid repeated conversions @@ -491,17 +512,11 @@ def compute_gene_cov(gene_data, n_groups): ) gene_arrays.append(gene_array) - # Stack along gene axis - # Note: Dask uses lazy evaluation - actual computation happens when array values are accessed - # Set global pool size if num_workers is specified - # This affects the threaded scheduler used by Dask delayed/array operations - if self.dask_num_workers is not None: - dask.config.set(pool=dask.config.get("pool", {}).copy()) - dask.config.set({"pool.num-workers": self.dask_num_workers}) - logger.info( - f"Configured Dask to use {self.dask_num_workers} workers (limits parallelism and memory)" - ) - + # Stack along gene axis. Dask evaluates lazily: nothing is + # computed until a consumer indexes a gene slice. No scheduler + # configuration is set here -- it would not govern that later + # evaluation, and a global dask.config.set would leak into the + # caller's process (settylab/kompot#31). dask_covariance = da.stack(gene_arrays, axis=2) # Clean up diff --git a/kompot/resource_estimation.py b/kompot/resource_estimation.py index e955010..f28f763 100644 --- a/kompot/resource_estimation.py +++ b/kompot/resource_estimation.py @@ -603,12 +603,31 @@ def estimate_differential_expression_resources( null_genes_param = kwargs.get("null_genes", 2000) # Default is 2000 compute_mahalanobis = kwargs.get("compute_mahalanobis", True) + # Refuse a null_genes we cannot interpret rather than silently counting + # zero. A string reaching here means the caller passed the unresolved + # "auto" sentinel, which used to fall through both branches below and + # under-report the plan by 2000 genes (settylab/kompot#25). + if isinstance(null_genes_param, str): + raise ValueError( + f"null_genes={null_genes_param!r} must be resolved to an int, a " + f"list of ints, or None before estimating resources. " + f"kompot.de() resolves 'auto' before the dry run; pass an " + f"explicit value when calling the estimator directly." + ) + n_null_genes = 0 if null_genes_param is not None and null_genes_param != 0 and compute_mahalanobis: - if isinstance(null_genes_param, int): - n_null_genes = null_genes_param - elif isinstance(null_genes_param, list): + if isinstance(null_genes_param, (int, np.integer)) and not isinstance( + null_genes_param, bool + ): + n_null_genes = int(null_genes_param) + elif isinstance(null_genes_param, (list, tuple)): n_null_genes = len(null_genes_param) + else: + raise TypeError( + f"null_genes must be an int, a list of ints, or None " + f"(got {type(null_genes_param).__name__})." + ) # Total genes processed includes null genes n_total_genes = n_genes + n_null_genes @@ -957,30 +976,60 @@ def estimate_differential_expression_resources( sv_cov_shape = (n_landmarks, n_landmarks, n_total_genes) sv_cov_size = estimate_array_size(sv_cov_shape) - resource_type = "disk" if store_arrays_on_disk else "memory" - plan.add_requirement( - f"Sample covariances (per condition, {n_samples} samples)", - sv_cov_size * 2, # variance1 and variance2 - both stored - resource_type, - shape=sv_cov_shape, - ) - - # Note: combined_variance (variance1 + variance2) behavior differs by storage type: - # - Disk storage: Creates a lazy Dask computation graph (no additional disk space) - # - Memory storage: Creates a third in-memory tensor (3x total memory) + # Where the two per-condition tensors live, and whether they are + # materialised at all, depends on the storage mode: + # + # store_arrays_on_disk=False two dense in-memory tensors + # True, dask installed lazy Dask graphs; NOTHING is + # written, each gene is evaluated on + # demand + # True, no dask two memory-mapped .npy files + # + # Their SUM is never materialised: compute_mahalanobis_distances + # assembles one (n_landmarks, n_landmarks) matrix per gene through + # kompot.utils.LazyGeneCovariance (settylab/kompot#26). Before that + # fix the plan charged a third tensor here, and charged disk for + # the dask path that writes nothing. if not store_arrays_on_disk: - # In memory: variance1, variance2, and combined_variance all exist as separate arrays plan.add_requirement( - "Combined sample covariances (temporary)", - sv_cov_size, + f"Sample covariances (per condition, {n_samples} samples)", + sv_cov_size * 2, # variance1 and variance2 "memory", shape=sv_cov_shape, ) - total_sv = sv_cov_size * 3 # variance1 + variance2 + combined plan.warnings.append( - f"Sample variance covariance tensors ({human_readable_size(total_sv)}) " - f"will be stored in memory. Consider storage=StorageSettings(store_arrays_on_disk=True) for large datasets." + f"Sample variance covariance tensors " + f"({human_readable_size(sv_cov_size * 2)}) will be held in " + f"memory. Restricting `genes` (linear) or lowering " + f"gp=GPSettings(n_landmarks=...) (quadratic) shrinks them; " + f"storage=StorageSettings(store_arrays_on_disk=True) keeps " + f"them out of memory entirely." + ) + elif DASK_AVAILABLE: + plan.info.append( + f"Sample covariances are Dask-backed: each gene's " + f"{n_landmarks}x{n_landmarks} matrix is evaluated on demand and " + f"the {human_readable_size(sv_cov_size * 2)} tensor is never " + f"materialised, in memory or on disk." ) + else: + plan.add_requirement( + f"Sample covariances (per condition, {n_samples} samples)", + sv_cov_size * 2, # variance1 and variance2, memory-mapped + "disk", + shape=sv_cov_shape, + ) + + # Per-gene working set: whichever storage mode is in force, the + # Mahalanobis step materialises one (n_landmarks, n_landmarks) + # matrix per sample-variance term plus the shared covariance. + per_gene_working = estimate_array_size((n_landmarks, n_landmarks)) * 3 + plan.add_requirement( + "Per-gene covariance working set (one gene at a time)", + per_gene_working, + "memory", + shape=(n_landmarks, n_landmarks), + ) if store_arrays_on_disk: # Add note about disk location @@ -995,8 +1044,13 @@ def estimate_differential_expression_resources( if not DASK_AVAILABLE: plan.warnings.append( - "Disk storage requested but dask is not installed. " - "Install dask for 50x faster computation: pip install 'dask[array]'" + "Disk storage requested but dask is not installed, so the " + "covariance tensors are written as memory-mapped .npy " + "files. Those pages stay resident as page cache, which a " + "cgroup (Slurm --mem, a container) charges against your " + "limit just like anonymous memory. Install dask to " + "evaluate the tensors lazily instead, which writes " + "nothing at all: pip install 'dask[array]'" ) # Add batch_size information diff --git a/kompot/settings.py b/kompot/settings.py index baf5496..a2a8628 100644 --- a/kompot/settings.py +++ b/kompot/settings.py @@ -65,13 +65,22 @@ class GPSettings: ls_factor : float Multiplier applied to the automatically inferred length scale. n_landmarks : int, optional - Number of landmarks for the Nystrom approximation. + Number of landmarks for the Nystrom approximation. Under sample + variance (``sample_col`` set on :func:`kompot.de`) the dominant + allocation is ``2 * n_landmarks**2 * n_genes * 8`` bytes, so memory is + **quadratic** in this value: halving it quarters the covariance + footprint. It also cuts the per-gene Cholesky factorisation, though by + an amount worth measuring rather than extrapolating. See + https://kompot.readthedocs.io/en/latest/resource_planning.html landmarks : np.ndarray, optional Pre-computed landmark coordinates. use_empirical_variance : bool Estimate per-gene heteroscedastic noise from GP residuals. batch_size : int, optional - Number of cells processed at once during prediction. + Number of cells processed at once during prediction. Also bounds the + gene batches of the *shared*-covariance Mahalanobis computation. It + does **not** bound the per-gene covariance loop used under sample + variance, which processes one gene at a time regardless. eps : float Small constant for numerical stability. jit_compile : bool @@ -328,11 +337,25 @@ class StorageSettings: store_additional_stats : bool Store extra columns (p-values, tail FDR, PTP, z-scores). store_arrays_on_disk : bool, optional - Use disk-backed arrays for large intermediate matrices. + Keep the per-gene sample-variance covariance tensors out of memory by + consuming them one gene at a time. With ``dask`` installed each + tensor is a lazy Dask graph and nothing is written to disk; without + ``dask`` each condition's tensor is written as a memory-mapped + ``.npy`` under ``disk_storage_dir``, which is slower because a gene + slice of that layout is strided. Defaults to ``None``, meaning *on if + and only if* ``disk_storage_dir`` is set; nothing enables it + automatically in response to memory pressure. See + https://kompot.readthedocs.io/en/latest/resource_planning.html disk_storage_dir : str, optional - Directory for disk-backed arrays. + Directory for disk-backed arrays. A real run creates it if missing, + but ``dry_run=True`` raises ``FileNotFoundError`` on a path that does + not exist, so create it before you plan against it. When unset, the + system temporary directory is used (honouring ``TMPDIR``), which on a + shared cluster is often small or RAM-backed. max_memory_ratio : float - Fraction of RAM before triggering disk storage. + Fraction of available RAM above which the resource estimator escalates + its warnings. It does **not** switch storage modes and does not cap + allocation. """ result_key: Optional[str] = None diff --git a/kompot/utils.py b/kompot/utils.py index a870224..c255fa4 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -41,6 +41,95 @@ def _sanitize_name(name): logger = logging.getLogger("kompot") + +class LazyGeneCovariance: + """Per-gene covariance tensor assembled one gene at a time. + + Sample variance produces a ``(n_points, n_points, n_genes)`` covariance + tensor per condition, and the Mahalanobis step needs their sum plus the + shared posterior covariance. Adding those up eagerly materialises a dense + tensor of that shape in memory, which defeats ``store_arrays_on_disk`` and + costs ``n_points**2 * n_genes * 8`` bytes even when every input is + disk-backed (settylab/kompot#26). + + This view holds references to the terms and materialises a single + ``(n_points, n_points)`` matrix only when a gene slice is requested, so + peak memory is bounded by one gene rather than by the whole tensor. It is + deliberately minimal: :func:`compute_mahalanobis_distances` only ever asks + for ``view[:, :, g]``. + + Parameters + ---------- + terms : sequence + 3-D array-likes of identical shape ``(n_points, n_points, n_genes)``. + NumPy arrays, ``np.memmap``, Dask arrays and JAX arrays all work. + base : np.ndarray, optional + 2-D ``(n_points, n_points)`` matrix added to every gene slice, i.e. + the shared posterior covariance. + """ + + def __init__(self, terms, base=None): + terms = [t for t in terms if t is not None] + if not terms: + raise ValueError("LazyGeneCovariance needs at least one term.") + shapes = {tuple(t.shape) for t in terms} + if len(shapes) != 1: + raise ValueError( + f"LazyGeneCovariance terms must share a shape (got {sorted(shapes)})." + ) + self._terms = terms + self._base = None if base is None else np.asarray(base) + self.shape = tuple(terms[0].shape) + if len(self.shape) != 3: + raise ValueError( + f"LazyGeneCovariance expects 3-D terms (got shape {self.shape})." + ) + if self._base is not None and self._base.shape != self.shape[:2]: + raise ValueError( + f"base has shape {self._base.shape}, expected {self.shape[:2]}." + ) + self.ndim = 3 + self.dtype = np.result_type(*[t.dtype for t in terms]) + + @property + def dask_backed(self) -> bool: + """True when any term is a Dask array, so slices are computed lazily.""" + from .memory_utils import DASK_AVAILABLE + + if not DASK_AVAILABLE: + return False + try: + import dask.array as da + except ImportError: # pragma: no cover - guarded by DASK_AVAILABLE + return False + return any(isinstance(t, da.Array) for t in self._terms) + + def gene(self, g: int) -> np.ndarray: + """Materialise the ``(n_points, n_points)`` covariance for gene ``g``.""" + out = np.asarray(self._terms[0][:, :, g], dtype=np.float64) + # Copy once so the accumulation below never writes into a caller's + # array (a NumPy basic slice is a view). + out = out.copy() + for term in self._terms[1:]: + out += np.asarray(term[:, :, g], dtype=np.float64) + if self._base is not None: + out += self._base + return out + + def __getitem__(self, key): + if ( + isinstance(key, tuple) + and len(key) == 3 + and key[0] == slice(None) + and key[1] == slice(None) + and isinstance(key[2], (int, np.integer)) + ): + return self.gene(int(key[2])) + raise TypeError( + "LazyGeneCovariance supports only [:, :, gene_index] indexing; " + f"got {key!r}." + ) + # Define standard colors for consistent use throughout the package KOMPOT_COLORS = { # Direction colors for differential abundance @@ -148,7 +237,32 @@ def compute_mahalanobis_distances( from tqdm.auto import tqdm from .memory_utils import DASK_AVAILABLE - # Check if covariance is a Dask array + # Check if covariance is a Dask array. + # + # REACHABILITY: since 0.9.0 nothing inside Kompot reaches this branch in + # production. Enumerating every in-repo caller: utils.py's singular + # `compute_mahalanobis_distance` passes a 2-D covariance; + # differential_expression.py passes a LazyGeneCovariance (3-D, but not a + # da.Array, and it deliberately does not take this branch -- see below) or + # a 2-D `combined_cov`. The differential-abundance path does not call this + # function at all. + # + # It is TESTED but production-unreachable, not dead: + # tests/test_mahalanobis_approaches.py builds a bare 3-D dask array with + # DiskStorage.as_dask_array and passes it straight in, asserting it agrees + # with the gene-specific path. An external caller can do the same. Deleting + # a covered path on a reading of the call graph needs more evidence than + # the call graph provides, and the cost of keeping it is this comment. + # + # A LazyGeneCovariance deliberately does NOT take this branch even when its + # terms are Dask-backed. The Dask branch below submits every gene as a + # delayed task at once, and a lazy view materialises its own slices inside + # each task, so the two nest: peak memory becomes (concurrent tasks) x + # (per-gene matrices) instead of one gene's worth. The sequential loop + # still gets Dask parallelism inside each slice computation. Measured at + # 600 landmarks x 150 genes: 1 886 MiB peak anonymous through the delayed + # path against 1 202 MiB sequential, for the same wall clock + # (settylab/kompot#26). is_dask = False if DASK_AVAILABLE: import dask.array as da @@ -188,13 +302,16 @@ def compute_mahalanobis_distances( # Create a custom function that can be mapped over gene dimensions def compute_gene_mahalanobis(g): - # Extract the difference vector and covariance matrix for this gene + # Extract the difference vector and covariance matrix for this gene. + # Only this gene's (n_points, n_points) matrix is ever materialized; + # a LazyGeneCovariance sums its terms here rather than up front. gene_diff = diffs[g] - gene_cov = covariance[:, :, g] - - # Convert to numpy (materializes Dask slices when using disk-backed storage) gene_diff_np = np.array(gene_diff) - gene_cov_np = np.array(gene_cov) + if isinstance(covariance, LazyGeneCovariance): + gene_cov_np = covariance.gene(g) + else: + # np.array materializes a Dask or memmap slice + gene_cov_np = np.array(covariance[:, :, g]) # Add per-gene empirical (diagonal) variance if provided if diagonal_variance is not None: @@ -260,8 +377,11 @@ def compute_gene_mahalanobis(g): return mahalanobis_distances - # For JAX arrays, proceed with the original approach - cov = jnp.array(covariance) + # Sequential per-gene path. Note there is deliberately no + # `jnp.array(covariance)` here: compute_gene_mahalanobis reads + # `covariance` directly, so materializing the whole tensor was a dead + # full-size copy of the array we are trying not to hold in memory + # (settylab/kompot#26). mahalanobis_distances = np.zeros(n_genes) # Process each gene separately to save memory, with progress bar @@ -601,6 +721,119 @@ def find_optimal_resolution( return resolution, best_partition +# Byte budget for one float64 working array in ``_exact_nearest_indices``. +_EXACT_NN_CHUNK_BYTES = 64 * 1024 * 1024 + + +def _exact_nearest_indices( + X: np.ndarray, queries: np.ndarray, chunk_size: Optional[int] = None +) -> np.ndarray: + """ + Exact Euclidean nearest row of ``X`` for each row of ``queries``. + + A brute-force search, chunked over the rows of ``X`` so that memory stays + bounded: no array larger than ``chunk_size x max(n_features, n_queries)`` + float64 values is formed. It is meant for few queries against many points, + such as snapping a few hundred cluster centroids to the nearest cell. + + The nearest row is defined by the squared distance ``((X[i] - q) ** 2).sum()`` + computed in float64. Exact ties, for example duplicate rows, go to the + lowest row index, so the result is deterministic. + + Parameters + ---------- + X : np.ndarray + Points of shape (n_samples, n_features). + queries : np.ndarray + Query points of shape (n_queries, n_features). + chunk_size : int or None, optional + Rows of ``X`` processed per step. By default about 64 MiB per working + array. + + Returns + ------- + np.ndarray + Integer array of shape (n_queries,) with the index of the nearest row + of ``X`` for each query. + """ + queries = np.asarray(queries, dtype=np.float64) + if queries.ndim != 2: + raise ValueError(f"queries must be 2-D, got shape {queries.shape}") + n_samples = X.shape[0] + n_queries, n_features = queries.shape + if X.ndim != 2 or X.shape[1] != n_features: + raise ValueError( + f"X has shape {X.shape}, incompatible with queries of shape {queries.shape}" + ) + if n_samples == 0: + raise ValueError("X has no rows to search") + if n_queries == 0: + return np.zeros(0, dtype=np.intp) + if chunk_size is None: + chunk_size = _EXACT_NN_CHUNK_BYTES // (8 * max(n_features, n_queries, 1)) + chunk_size = max(1, int(chunk_size)) + + # The search ranks rows by |x|^2 - 2 x.q, which differs from the squared + # distance by the per-query constant |q|^2. It is a matrix product and so + # fast, but it rounds differently from the direct sum, so it only + # shortlists candidates: every row whose score is within twice the rounding + # bound of the running minimum. The final minimum is never above the + # running one, so every row that can tie or beat it is on the shortlist. + q_sq = np.einsum("ij,ij->i", queries, queries) + q_norm = np.sqrt(q_sq) + eps = np.finfo(np.float64).eps + queries_t = np.ascontiguousarray(queries.T) + running_min = np.full(n_queries, np.inf) + max_x_norm = 0.0 + cand_rows, cand_cols, cand_scores = [], [], [] + for start in range(0, n_samples, chunk_size): + chunk = np.asarray(X[start : start + chunk_size], dtype=np.float64) + x_sq = np.einsum("ij,ij->i", chunk, chunk) + max_x_norm = max(max_x_norm, float(np.sqrt(x_sq.max()))) + score = chunk @ queries_t + score *= -2.0 + score += x_sq[:, None] + np.minimum(running_min, score.min(axis=0), out=running_min) + # Rounding bound on the score and on the direct sum it stands in for, + # with a 4x margin: (d + 3) eps (|x| + |q|)^2 each. The largest |x| seen + # so far only grows, so the bound applied later is never smaller. + bound = 4.0 * (n_features + 3) * eps * (max_x_norm + q_norm) ** 2 + rows, cols = np.nonzero(score <= (running_min + 2.0 * bound)[None, :]) + cand_rows.append(rows + start) + cand_cols.append(cols) + cand_scores.append(score[rows, cols]) + + bound = 4.0 * (n_features + 3) * eps * (max_x_norm + q_norm) ** 2 + rows = np.concatenate(cand_rows) + cols = np.concatenate(cand_cols) + scores = np.concatenate(cand_scores) + keep = scores <= (running_min + 2.0 * bound)[cols] + rows, cols = rows[keep], cols[keep] + + # Rank the shortlist by the direct sum; np.argmin returns the first of + # equal values and the rows are in ascending order, so exact ties go to + # the lowest index. + best_idx = np.full(n_queries, -1, dtype=np.intp) + order = np.lexsort((rows, cols)) + rows, cols = rows[order], cols[order] + bounds = np.searchsorted(cols, np.arange(n_queries + 1)) + for j in range(n_queries): + cand = rows[bounds[j] : bounds[j + 1]] + if cand.size == 0: + continue + points = np.asarray(X[cand], dtype=np.float64) + exact = ((points - queries[j]) ** 2).sum(axis=1) + best_idx[j] = cand[int(np.argmin(exact))] + + if np.any(best_idx < 0): + # Unreachable unless the input holds NaN or inf. + raise ValueError( + "Could not find a nearest point for every query; check X and the " + "queries for NaN or infinite values." + ) + return best_idx + + def find_landmarks( X: np.ndarray, n_clusters: int = 200, @@ -608,6 +841,7 @@ def find_landmarks( tol: float = 0.1, max_iter: int = 10, random_state: Optional[int] = None, + exact_snap: bool = True, ) -> Tuple[np.ndarray, np.ndarray]: """ Identify landmark points representing clusters in the dataset. @@ -635,6 +869,15 @@ def find_landmarks( landmark indices and coordinates). The default (``None``) preserves the historical non-deterministic behavior, so existing callers are unaffected. + 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, + ties to the lowest index), found by a chunked brute-force search whose + memory is bounded. With False the snap uses the approximate + 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. Returns ------- @@ -665,9 +908,12 @@ def find_landmarks( # Compute centroids centroids = np.array([X[clusters == c].mean(axis=0) for c in cluster_ids]) - # Find the nearest data point to each centroid - landmark_indices, _ = index.query(centroids, k=1) - landmark_indices = landmark_indices.flatten() + # Snap each centroid to a data point + if exact_snap: + landmark_indices = _exact_nearest_indices(X, centroids) + else: + landmark_indices, _ = index.query(centroids, k=1) + landmark_indices = landmark_indices.flatten() landmarks = X[landmark_indices] logger.info( diff --git a/kompot/version.py b/kompot/version.py index 028b8b0..ca25270 100644 --- a/kompot/version.py +++ b/kompot/version.py @@ -1,3 +1,3 @@ """Version information.""" -__version__ = "0.8.0" +__version__ = "0.9.0" diff --git a/pyproject.toml b/pyproject.toml index fc8d333..643c611 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -49,7 +49,7 @@ ignore = ["E203", "W503"] [project] name = "kompot" -version = "0.8.0" +version = "0.9.0" description = "Differential abundance and gene expression analysis using Mahalanobis distance with JAX backend" readme = "README.md" authors = [ diff --git a/tests/test_cli.py b/tests/test_cli.py index eb3a6fa..c8e3c0c 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -11,6 +11,10 @@ import subprocess import sys +# Fixtures write through the CLI's own writer: under anndata 0.11/0.12 and +# pandas 3 a plain write_h5ad of any AnnData fails (settylab/kompot#23). +from kompot.cli.utils import write_output + @pytest.fixture def sample_adata(): @@ -115,7 +119,7 @@ def test_validate_anndata_path(self, sample_adata, temp_dir): # Test with existing file adata_file = Path(temp_dir) / "test.h5ad" - sample_adata.write_h5ad(adata_file) + write_output(sample_adata, adata_file) validated = validate_anndata_path(str(adata_file)) assert validated.exists() @@ -132,7 +136,7 @@ def test_de_basic_cli_args(self, sample_adata, temp_dir): """Test DE with basic CLI arguments.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI cmd = [ @@ -178,7 +182,7 @@ def test_de_with_config_file(self, sample_adata, temp_dir): output_file = Path(temp_dir) / "output.h5ad" config_file = Path(temp_dir) / "config.yaml" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Create config config = { @@ -219,7 +223,7 @@ def test_de_cli_overrides_config(self, sample_adata, temp_dir): output_file = Path(temp_dir) / "output.h5ad" config_file = Path(temp_dir) / "config.yaml" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Create config with different condition1 config = { @@ -260,7 +264,7 @@ def test_de_missing_required_params(self, sample_adata, temp_dir): """Test error when required parameters are missing.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI without required params cmd = [ @@ -286,7 +290,7 @@ def test_de_table_output_csv(self, sample_adata, temp_dir): """Test DE with table output to CSV.""" input_file = Path(temp_dir) / "input.h5ad" table_file = Path(temp_dir) / "results.csv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI with table output only (no AnnData output) cmd = [ @@ -330,7 +334,7 @@ def test_de_table_output_tsv(self, sample_adata, temp_dir): """Test DE with table output to TSV.""" input_file = Path(temp_dir) / "input.h5ad" table_file = Path(temp_dir) / "results.tsv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -372,7 +376,7 @@ def test_de_both_outputs(self, sample_adata, temp_dir): input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" table_file = Path(temp_dir) / "results.csv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -409,7 +413,7 @@ def test_de_both_outputs(self, sample_adata, temp_dir): def test_de_no_output_error(self, sample_adata, temp_dir): """Test error when neither output nor table-output is specified.""" input_file = Path(temp_dir) / "input.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -442,7 +446,7 @@ def test_da_basic_cli_args(self, sample_adata, temp_dir): """Test DA with basic CLI arguments.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI cmd = [ @@ -482,7 +486,7 @@ def test_da_with_config_file(self, sample_adata, temp_dir): output_file = Path(temp_dir) / "output.h5ad" config_file = Path(temp_dir) / "config.yaml" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Create config config = { @@ -518,7 +522,7 @@ def test_da_with_thresholds(self, sample_adata, temp_dir): """Test DA with custom thresholds.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI with custom thresholds cmd = [ @@ -551,7 +555,7 @@ def test_da_table_output_csv(self, sample_adata, temp_dir): """Test DA with table output to CSV.""" input_file = Path(temp_dir) / "input.h5ad" table_file = Path(temp_dir) / "results.csv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI with table output only (no AnnData output) cmd = [ @@ -589,7 +593,7 @@ def test_da_table_output_tsv(self, sample_adata, temp_dir): """Test DA with table output to TSV.""" input_file = Path(temp_dir) / "input.h5ad" table_file = Path(temp_dir) / "results.tsv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -623,7 +627,7 @@ def test_da_table_output_tsv(self, sample_adata, temp_dir): def test_da_no_output_error(self, sample_adata, temp_dir): """Test error when neither output nor table-output is specified.""" input_file = Path(temp_dir) / "input.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -663,7 +667,7 @@ def test_dm_basic_cli_args(self, sample_adata, temp_dir): input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI cmd = [ @@ -708,7 +712,7 @@ def test_dm_with_config_file(self, sample_adata, temp_dir): output_file = Path(temp_dir) / "output.h5ad" config_file = Path(temp_dir) / "config.yaml" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Create config config = {"pca_key": "X_pca", "n_components": 10, "knn": 30} @@ -741,7 +745,7 @@ def test_dm_missing_pca(self, sample_adata, temp_dir): # Remove PCA from sample data del sample_adata.obsm["X_pca"] - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) # Run CLI without PCA in data cmd = [ @@ -772,7 +776,7 @@ def test_dm_without_palantir_says_so(self, sample_adata, temp_dir): else: pytest.skip("palantir is installed") input_file = Path(temp_dir) / "input.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [sys.executable, "-m", "kompot.cli", "dm", str(input_file), "-o", str(Path(temp_dir) / "output.h5ad")] result = subprocess.run(cmd, capture_output=True, text=True) diff --git a/tests/test_cli_compute_config.py b/tests/test_cli_compute_config.py index 2026532..fe872de 100644 --- a/tests/test_cli_compute_config.py +++ b/tests/test_cli_compute_config.py @@ -13,6 +13,10 @@ import pandas as pd from anndata import AnnData +# Fixtures write through the CLI's own writer: under anndata 0.11/0.12 and +# pandas 3 a plain write_h5ad of any AnnData fails (settylab/kompot#23). +from kompot.cli.utils import write_output + @pytest.fixture def sample_adata_for_cli(): @@ -223,7 +227,7 @@ def test_run_de_missing_output(self, sample_adata_for_cli, tmp_path, capsys): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args without output args = argparse.Namespace( @@ -307,7 +311,7 @@ def test_run_de_missing_required_params(self, sample_adata_for_cli, tmp_path): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args without required parameters args = argparse.Namespace( @@ -361,7 +365,7 @@ def test_run_de_with_config_file(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -414,7 +418,7 @@ def test_run_de_table_output_csv(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -478,7 +482,7 @@ def test_run_de_table_output_tsv(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -534,7 +538,7 @@ def test_run_de_unsupported_output_format( # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args with unsupported output format args = argparse.Namespace( @@ -588,7 +592,7 @@ def test_run_de_unsupported_table_format( # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args with unsupported table format args = argparse.Namespace( @@ -645,7 +649,7 @@ def test_dry_run_outputs_json(self, sample_adata_for_cli, tmp_path, monkeypatch, import argparse input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) args = argparse.Namespace( input=str(input_file), @@ -706,7 +710,7 @@ def test_dry_run_infeasible_exits_1(self, sample_adata_for_cli, tmp_path, monkey import argparse input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) args = argparse.Namespace( input=str(input_file), @@ -762,7 +766,7 @@ def test_dry_run_skips_output_validation(self, sample_adata_for_cli, tmp_path, m import argparse input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) args = argparse.Namespace( input=str(input_file), @@ -867,7 +871,7 @@ def test_run_da_missing_output(self, sample_adata_for_cli, tmp_path): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args without output args = argparse.Namespace( @@ -945,7 +949,7 @@ def test_run_da_missing_required_params(self, sample_adata_for_cli, tmp_path): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args without required parameters args = argparse.Namespace( @@ -996,7 +1000,7 @@ def test_run_da_with_config_file(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -1046,7 +1050,7 @@ def test_run_da_table_output_csv(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -1107,7 +1111,7 @@ def test_run_da_table_output_tsv(self, sample_adata_for_cli, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args args = argparse.Namespace( @@ -1160,7 +1164,7 @@ def test_run_da_unsupported_output_format( # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args with unsupported output format args = argparse.Namespace( @@ -1211,7 +1215,7 @@ def test_run_da_unsupported_table_format( # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_for_cli.write_h5ad(input_file) + write_output(sample_adata_for_cli, input_file) # Create args with unsupported table format args = argparse.Namespace( diff --git a/tests/test_cli_dm_coverage.py b/tests/test_cli_dm_coverage.py index 23f6c18..2cc276f 100644 --- a/tests/test_cli_dm_coverage.py +++ b/tests/test_cli_dm_coverage.py @@ -11,6 +11,10 @@ import pandas as pd from anndata import AnnData +# Fixtures write through the CLI's own writer: under anndata 0.11/0.12 and +# pandas 3 a plain write_h5ad of any AnnData fails (settylab/kompot#23). +from kompot.cli.utils import write_output + @pytest.fixture def sample_adata_with_pca(): @@ -160,7 +164,7 @@ def test_run_dm_missing_palantir( # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) # Mock palantir import to raise ImportError import builtins @@ -210,7 +214,7 @@ def test_run_dm_missing_pca_key(self, sample_adata_no_pca, tmp_path): # Save sample data (no PCA) input_file = tmp_path / "input.h5ad" - sample_adata_no_pca.write_h5ad(input_file) + write_output(sample_adata_no_pca, input_file) # Create args args = argparse.Namespace( @@ -290,7 +294,7 @@ def test_run_dm_unsupported_output_format(self, sample_adata_with_pca, tmp_path) # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) # Mock palantir.utils.run_diffusion_maps try: @@ -329,7 +333,7 @@ def test_run_dm_with_config_file(self, sample_adata_with_pca, tmp_path): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) # Create args args = argparse.Namespace( @@ -380,7 +384,7 @@ def test_run_dm_h5ad_success(self, sample_adata_with_pca, tmp_path, monkeypatch) # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) output_file = tmp_path / "output.h5ad" # Create args @@ -440,7 +444,7 @@ def test_run_dm_zarr_output(self, sample_adata_with_pca, tmp_path, monkeypatch): # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) output_file = tmp_path / "output.zarr" # Create args @@ -498,7 +502,7 @@ def test_dm_argument_conversion(self, sample_adata_with_pca, tmp_path, monkeypat # Save sample data input_file = tmp_path / "input.h5ad" - sample_adata_with_pca.write_h5ad(input_file) + write_output(sample_adata_with_pca, input_file) # Create args with hyphenated names (as they come from CLI) args = argparse.Namespace( diff --git a/tests/test_cli_smooth.py b/tests/test_cli_smooth.py index 110ad84..b5b2730 100644 --- a/tests/test_cli_smooth.py +++ b/tests/test_cli_smooth.py @@ -12,6 +12,10 @@ import yaml from anndata import AnnData +# Fixtures write through the CLI's own writer: under anndata 0.11/0.12 and +# pandas 3 a plain write_h5ad of any AnnData fails (settylab/kompot#23). +from kompot.cli.utils import write_output + @pytest.fixture def sample_adata(): @@ -42,7 +46,7 @@ def test_smooth_basic(self, sample_adata, temp_dir): """Test basic smooth with h5ad output.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -75,7 +79,7 @@ def test_smooth_with_condition(self, sample_adata, temp_dir): """Test smooth with condition selection.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -108,7 +112,7 @@ def test_smooth_table_output(self, sample_adata, temp_dir): """Test smooth with table output.""" input_file = Path(temp_dir) / "input.h5ad" table_file = Path(temp_dir) / "results.csv" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -139,7 +143,7 @@ def test_smooth_with_config(self, sample_adata, temp_dir): input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" config_file = Path(temp_dir) / "config.yaml" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) config = { "obsm_key": "X_pca", @@ -169,7 +173,7 @@ def test_smooth_with_config(self, sample_adata, temp_dir): def test_smooth_no_output_fails(self, sample_adata, temp_dir): """Test that missing output args fails.""" input_file = Path(temp_dir) / "input.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -188,7 +192,7 @@ def test_smooth_with_empirical_variance(self, sample_adata, temp_dir): """Test smooth with --use-empirical-variance flag.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, @@ -218,7 +222,7 @@ def test_smooth_gene_subset(self, sample_adata, temp_dir): """Test smooth with gene subsetting.""" input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - sample_adata.write_h5ad(input_file) + write_output(sample_adata, input_file) cmd = [ sys.executable, diff --git a/tests/test_cli_write_output.py b/tests/test_cli_write_output.py new file mode 100644 index 0000000..0681014 --- /dev/null +++ b/tests/test_cli_write_output.py @@ -0,0 +1,66 @@ +"""The CLI must be able to write its output under every supported anndata. + +settylab/kompot#23: anndata 0.11 and 0.12 refuse to write a pandas +``StringArray`` unless ``anndata.settings.allow_write_nullable_strings`` is +set, and pandas 3 produces one for every ordinary string column, so every +``kompot de/da/smooth/dm`` write failed under those defaults. An explicit +``dtype="string"`` column produces the same array under pandas 2 as well, so the +test does not depend on pandas 3 to build its input. (It has been run on anndata +0.12 only with pandas 3, and on anndata 0.10 with pandas 2, where it skips.) +""" + +import anndata +import numpy as np +import pandas as pd +import pytest + +from kompot.cli.utils import write_output + + +def _adata_with_nullable_strings(): + adata = anndata.AnnData(X=np.ones((4, 3), dtype=np.float32)) + adata.obs["label"] = pd.array(["a", "b", None, "d"], dtype="string") + adata.var["symbol"] = pd.array(["G1", "G2", "G3"], dtype="string") + return adata + + +def _requires_the_setting(): + """anndata < 0.11 has no opt-in and cannot write a StringArray at all. + + Under pandas 2 those versions produce object columns, never a StringArray, + so #23 does not arise there; the opt-in is a no-op on them by design. + """ + settings = getattr(anndata, "settings", None) + if settings is None or not hasattr(settings, "allow_write_nullable_strings"): + pytest.skip("anndata has no allow_write_nullable_strings setting (< 0.11)") + + +@pytest.mark.parametrize("suffix", [".h5ad", ".zarr"]) +def test_write_output_writes_nullable_string_columns(tmp_path, suffix): + _requires_the_setting() + if suffix == ".zarr": + pytest.importorskip("zarr") + adata = _adata_with_nullable_strings() + out = tmp_path / f"out{suffix}" + + write_output(adata, out) + + back = anndata.read_h5ad(out) if suffix == ".h5ad" else anndata.read_zarr(out) + assert list(back.obs["label"].astype(object).where(back.obs["label"].notna(), None)) == [ + "a", "b", None, "d" + ] + assert list(back.var["symbol"].astype(str)) == ["G1", "G2", "G3"] + + +def test_write_output_does_not_leak_the_setting(tmp_path): + """The opt-in is scoped to the write, not left switched on for the caller.""" + _requires_the_setting() + settings = anndata.settings + before = settings.allow_write_nullable_strings + write_output(_adata_with_nullable_strings(), tmp_path / "out.h5ad") + assert settings.allow_write_nullable_strings == before + + +def test_write_output_rejects_unknown_suffix(tmp_path): + with pytest.raises(SystemExit): + write_output(_adata_with_nullable_strings(), tmp_path / "out.csv") diff --git a/tests/test_dask_num_workers.py b/tests/test_dask_num_workers.py new file mode 100644 index 0000000..5c83897 --- /dev/null +++ b/tests/test_dask_num_workers.py @@ -0,0 +1,65 @@ +"""`SampleVarianceEstimator(dask_num_workers=...)` must not claim a limit it does not apply. + +settylab/kompot#31: the argument wrote the Dask config key ``pool.num-workers``, +which no scheduler reads (the threaded scheduler reads ``num_workers``), mutated +Dask's *global* config as a side effect, and logged "Configured Dask to use N +workers" regardless. The covariance tensor it governed is returned lazily and is +materialised one gene at a time, so there is no worker pool for the setting to +bound. It is now deprecated and ignored, and says so. +""" + +import logging +import warnings + +import numpy as np +import pytest + +from kompot.differential.sample_variance_estimator import ( + DASK_AVAILABLE, + SampleVarianceEstimator, +) + + +def _data(): + rng = np.random.default_rng(0) + X = rng.standard_normal((40, 3)) + Y = rng.standard_normal((40, 2)) + groups = np.repeat([0, 1, 2], [14, 13, 13]) + return X, Y, groups + + +def test_dask_num_workers_warns_that_it_has_no_effect(): + with pytest.warns(FutureWarning, match="no effect"): + SampleVarianceEstimator(dask_num_workers=2) + + +def test_default_does_not_warn(): + with warnings.catch_warnings(): + warnings.simplefilter("error", FutureWarning) + SampleVarianceEstimator() + + +@pytest.mark.skipif(not DASK_AVAILABLE, reason="dask not installed") +def test_disk_backed_predict_leaves_dask_config_alone_and_claims_nothing(tmp_path, caplog): + import dask + + X, Y, groups = _data() + before = dict(dask.config.config) + with pytest.warns(FutureWarning): + sve = SampleVarianceEstimator( + store_arrays_on_disk=True, + disk_storage_dir=str(tmp_path), + jit_compile=False, + dask_num_workers=2, + ) + sve.fit(X, Y, groups) + with caplog.at_level(logging.INFO, logger="kompot"): + cov = sve.predict(X[:8], diag=False) + + assert isinstance(cov, dask.array.Array) + assert np.isfinite(np.asarray(cov[:, :, 0])).all() + # No global side effect: neither the key the old code wrote nor any other. + assert dask.config.get("pool.num-workers", None) is None + assert dict(dask.config.config) == before + # And no log line asserting a limit that was never applied. + assert not any("workers" in r.getMessage() for r in caplog.records) diff --git a/tests/test_docs_markup_leak_check.py b/tests/test_docs_markup_leak_check.py new file mode 100644 index 0000000..e23ea86 --- /dev/null +++ b/tests/test_docs_markup_leak_check.py @@ -0,0 +1,90 @@ +"""The docs build must fail when RST markup survives into the HTML (settylab/kompot#30). + +Nested inline markup is invalid reStructuredText; docutils renders the inner +markers as literal text and ``sphinx-build`` exits 0. The check lives in +``docs/source/markup_leak_check.py`` and runs at ``build-finished``. +""" + +import sys +from pathlib import Path + +import pytest + +DOCS_SOURCE = Path(__file__).resolve().parents[1] / "docs" / "source" +sys.path.insert(0, str(DOCS_SOURCE)) +from markup_leak_check import find_leaks # noqa: E402 + +# The three shapes reported on #30, as docutils actually renders them. +NESTED = [ + "

Supplying ``sample_col`` multiplies the cost by the gene count.

", + "

The ``dask`` path wins on every instrument, and that is it.

", + "

The ``dask``-less path is a different bargain ...

", +] + + +@pytest.mark.parametrize("html", NESTED, ids=["sample_col", "dask", "dask-less"]) +def test_nested_literal_in_bold_is_a_leak(html): + assert [kind for kind, _ in find_leaks(html)] == ["double-backtick literal"] + + +def test_code_in_link_text_is_a_leak(): + """Markdown [`x`](url) through pandoc: the literal is a tag, only `...`__ is text.""" + html = ( + '

See `kompot.de' + ' <' + "https://example.org/x>`__.

" + ) + assert [kind for kind, _ in find_leaks(html)] == ["hyperlink"] + + +def test_a_rendered_link_is_not_a_hyperlink_leak(): + html = '

See kompot.de.

' + assert find_leaks(html) == [] + + +def test_role_in_a_code_block_comment_is_a_leak(): + html = '
# see :ref:`dry-run` for the plan
' + assert [kind for kind, _ in find_leaks(html)] == ["role"] + + +def test_directive_rendered_as_text_is_a_leak(): + html = "

.. warning::\nthis never became an admonition

" + assert [kind for kind, _ in find_leaks(html)] == ["directive"] + + +def test_correctly_rendered_markup_is_clean(): + html = ( + '

The ' + 'dask path wins; see ' + 'x.

' + "" + ) + assert find_leaks(html) == [] + + +def test_one_hit_per_literal_not_per_backtick_pair(): + assert len(find_leaks("

a ``x`` b ``y`` c

")) == 2 + + +def test_sphinx_build_fails_on_nested_markup(tmp_path): + """End to end: a real build of the #30 shape must exit non-zero, and a clean one zero.""" + pytest.importorskip("sphinx") + from sphinx.cmd.build import build_main + + def build(body, name): + src = tmp_path / name / "src" + src.mkdir(parents=True) + (src / "conf.py").write_text( + "import sys\n" + f"sys.path.insert(0, {str(DOCS_SOURCE)!r})\n" + "from markup_leak_check import check_after_build\n" + "def setup(app):\n" + " app.connect('build-finished', check_after_build)\n" + ) + (src / "index.rst").write_text("Title\n=====\n\n" + body + "\n") + return build_main(["-q", "-b", "html", str(src), str(tmp_path / name / "out")]) + + assert build("**The dask path wins** with ``dask``.", "clean") == 0 + assert build("**The ``dask`` path wins on every instrument**", "nested") != 0 + # What pandoc emits for Markdown [`kompot.de`](url). + assert build("See ```kompot.de`` `__.", "code_in_link") != 0 diff --git a/tests/test_landmark_snap.py b/tests/test_landmark_snap.py new file mode 100644 index 0000000..92ca06b --- /dev/null +++ b/tests/test_landmark_snap.py @@ -0,0 +1,206 @@ +"""The landmark snap in ``find_landmarks`` returns the cell exactly nearest each centroid. + +``find_landmarks`` clusters the cells and snaps each cluster centroid to a cell. +Through 0.8.0 the snap queried the approximate pynndescent index built for the +graph, which on large, high-dimensional data often returns a nearby cell +instead of the nearest one. The snap is now an exact, chunked brute-force +search with a defined tie-break. +""" + +import numpy as np +import pytest + +import kompot.utils as kutils +from kompot.utils import find_landmarks + + +def _brute_force_nearest(X, queries): + """Reference: direct float64 squared distances, first index on ties.""" + X = np.asarray(X, dtype=np.float64) + return np.array( + [int(np.argmin(((X - q) ** 2).sum(axis=1))) for q in np.asarray(queries)] + ) + + +class _FixedPartition: + def __init__(self, membership): + self.membership = list(membership) + + +def _clustered(n, d, k, seed): + rng = np.random.default_rng(seed) + centres = rng.normal(size=(k, d)) * 3.0 + labels = rng.integers(0, k, n) + X = centres[labels] + rng.normal(size=(n, d)) + return X, labels + + +class TestFindLandmarksExactSnap: + def test_landmarks_are_exact_nearest_to_centroids(self, monkeypatch): + """Each landmark is the cell exactly nearest its cluster centroid. + + 20 000 cells in 50 dimensions with 100 clusters: here the approximate + index query misses the exact nearest cell for roughly one centroid in + eight (13 of 100 with pynndescent 0.6.0), so this fails on the + approximate snap. The clustering is fixed to the generating labels so + that the centroids are known. + """ + X, labels = _clustered(20_000, 50, 100, seed=0) + monkeypatch.setattr( + kutils, + "find_optimal_resolution", + lambda *a, **k: (1.0, _FixedPartition(labels)), + ) + + landmarks, indices = find_landmarks(X, n_clusters=100, random_state=0) + + cluster_ids = np.unique(labels) + centroids = np.array([X[labels == c].mean(axis=0) for c in cluster_ids]) + expected = _brute_force_nearest(X, centroids) + missed = int(np.sum(indices != expected)) + assert missed == 0, f"{missed} of {len(expected)} landmarks are not the nearest cell" + np.testing.assert_array_equal(landmarks, X[expected]) + + def test_real_clustering_path_snaps_exactly(self, monkeypatch): + """Through the real Leiden step, each landmark is the nearest cell to its centroid.""" + captured = {} + real = kutils.find_optimal_resolution + + def spy(*args, **kwargs): + resolution, partition = real(*args, **kwargs) + captured["membership"] = np.array(partition.membership) + return resolution, partition + + monkeypatch.setattr(kutils, "find_optimal_resolution", spy) + X, _ = _clustered(3_000, 20, 25, seed=3) + _, indices = find_landmarks(X, n_clusters=25, random_state=0) + + membership = captured["membership"] + centroids = np.array( + [X[membership == c].mean(axis=0) for c in np.unique(membership)] + ) + np.testing.assert_array_equal(indices, _brute_force_nearest(X, centroids)) + + def test_approximate_snap_is_an_explicit_opt_out(self, monkeypatch): + """``exact_snap=False`` keeps the pre-0.9.0 index query and skips the exact search.""" + + def boom(*args, **kwargs): + raise AssertionError("exact search must not run when exact_snap=False") + + monkeypatch.setattr(kutils, "_exact_nearest_indices", boom) + 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 + ) + assert len(indices) == landmarks.shape[0] + np.testing.assert_array_equal(landmarks, X[indices]) + + def test_seeded_landmarks_are_deterministic(self): + X, _ = _clustered(2_000, 10, 15, seed=4) + first = find_landmarks(X, n_clusters=15, random_state=0) + second = find_landmarks(X, n_clusters=15, random_state=0) + np.testing.assert_array_equal(first[1], second[1]) + np.testing.assert_array_equal(first[0], second[0]) + + +class TestExactNearestIndices: + @pytest.mark.parametrize("dtype", [np.float64, np.float32]) + @pytest.mark.parametrize("chunk_size", [None, 1, 7, 512]) + def test_matches_brute_force(self, dtype, chunk_size): + rng = np.random.default_rng(1) + X = rng.normal(size=(2_000, 12)).astype(dtype) + queries = rng.normal(size=(40, 12)) * 0.3 + result = kutils._exact_nearest_indices(X, queries, chunk_size=chunk_size) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + assert result.dtype.kind == "i" + + def test_ties_go_to_the_lowest_index(self): + """Duplicate rows are equally near; the lowest index wins, across chunks too.""" + rng = np.random.default_rng(2) + base = rng.normal(size=(50, 6)) + X = np.concatenate([base, base, base]) # row i == row i+50 == row i+100 + perm = rng.permutation(len(X)) + X = X[perm] + queries = base[:20] + 1e-3 * rng.normal(size=(20, 6)) + + for chunk_size in (None, 1, 13, 64): + result = kutils._exact_nearest_indices(X, queries, chunk_size=chunk_size) + for j, q in enumerate(queries): + d2 = ((X - q) ** 2).sum(axis=1) + tied = np.flatnonzero(d2 == d2.min()) + assert len(tied) == 3 + assert result[j] == tied.min() + + def test_exact_ties_on_a_grid(self): + """Integer coordinates make many distances exactly equal.""" + rng = np.random.default_rng(3) + X = rng.integers(0, 3, size=(4_000, 3)).astype(np.float64) + queries = rng.integers(0, 3, size=(30, 3)).astype(np.float64) + 0.5 + for chunk_size in (None, 37): + result = kutils._exact_nearest_indices(X, queries, chunk_size=chunk_size) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + + def test_deterministic_and_chunk_invariant(self): + rng = np.random.default_rng(5) + X = rng.normal(size=(5_000, 30)) + queries = rng.normal(size=(25, 30)) * 0.2 + reference = kutils._exact_nearest_indices(X, queries) + for chunk_size in (None, 3, 100, 4_999, 5_000, 10_000): + np.testing.assert_array_equal( + kutils._exact_nearest_indices(X, queries, chunk_size=chunk_size), + reference, + ) + + def test_far_from_origin(self): + """Large coordinates relative to the spread stress the rounding bound.""" + rng = np.random.default_rng(6) + X = 1e4 + rng.normal(size=(3_000, 40)) + queries = X[rng.choice(3_000, 20, replace=False)] + 0.05 * rng.normal( + size=(20, 40) + ) + result = kutils._exact_nearest_indices(X, queries, chunk_size=128) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + + def test_near_duplicate_cells(self): + """Cells closer together than the matrix-product score can resolve. + + The fast score ranks these copies by rounding noise alone, so the result + matches brute force only if the shortlist keeps every near-tie and the + shortlist is then ranked by the direct sum. + """ + rng = np.random.default_rng(0) + base = rng.normal(size=(500, 20)) * 50 + X = np.concatenate( + [base, base + 1e-12 * rng.normal(size=base.shape), base * (1 + 1e-15)] + ) + X = X[rng.permutation(len(X))] + queries = base[:40] + 1e-10 * rng.normal(size=(40, 20)) + result = kutils._exact_nearest_indices(X, queries, chunk_size=97) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + + def test_query_at_the_midpoint_of_two_cells(self): + rng = np.random.default_rng(7) + X = rng.normal(size=(300, 10)) + pairs = rng.choice(300, (40, 2), replace=False) + queries = (X[pairs[:, 0]] + X[pairs[:, 1]]) / 2 + result = kutils._exact_nearest_indices(X, queries, chunk_size=11) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + + def test_tiny_spread_far_from_origin(self): + """Score cancellation: a spread of 1e-3 at an offset of 1e7.""" + rng = np.random.default_rng(8) + X = 1e7 + 1e-3 * rng.normal(size=(3_000, 64)) + queries = X[:20] + 1e-5 * rng.normal(size=(20, 64)) + result = kutils._exact_nearest_indices(X, queries, chunk_size=500) + np.testing.assert_array_equal(result, _brute_force_nearest(X, queries)) + + def test_rejects_bad_input(self): + X = np.zeros((10, 3)) + with pytest.raises(ValueError): + kutils._exact_nearest_indices(X, np.zeros((2, 4))) + with pytest.raises(ValueError): + kutils._exact_nearest_indices(np.zeros((0, 3)), np.zeros((2, 3))) + with pytest.raises(ValueError): + kutils._exact_nearest_indices(X, np.full((1, 3), np.nan)) + assert kutils._exact_nearest_indices(X, np.zeros((0, 3))).shape == (0,) diff --git a/tests/test_provenance.py b/tests/test_provenance.py index 493305e..0bc0084 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -82,6 +82,160 @@ def test_resolve_sha_returns_none_when_unresolvable(tmp_path): assert _resolve_sha(str(git_dir)) is None +# -------------------------------------------------------------------------- +# linked worktrees (settylab/kompot#28, #33) +# -------------------------------------------------------------------------- + + +def _make_linked_worktree_git_dir(root, commondir="../..", loose=SHA, packed=None): + """Mirror git's layout: the worktree's git dir holds HEAD, the main one holds refs.""" + main = _make_git_dir(root, loose=None, packed=packed) + if loose is not None: + (main / "refs" / "heads" / "feature").write_text(loose + "\n") + wt = main / "worktrees" / "wt" + wt.mkdir(parents=True) + (wt / "HEAD").write_text("ref: refs/heads/feature\n") + (wt / "commondir").write_text((str(main) if commondir is None else commondir) + "\n") + return wt + + +def test_resolve_sha_linked_worktree_on_a_branch(tmp_path): + """Branch refs live in the common dir, not in the per-worktree git dir.""" + wt = _make_linked_worktree_git_dir(tmp_path) + assert _resolve_sha(str(wt)) == SHA + + +def test_resolve_sha_linked_worktree_absolute_commondir(tmp_path): + wt = _make_linked_worktree_git_dir(tmp_path, commondir=None) + assert _resolve_sha(str(wt)) == SHA + + +def test_resolve_sha_linked_worktree_packed_refs_in_commondir(tmp_path): + """After `git gc` the branch ref is packed, and packed-refs is shared too.""" + wt = _make_linked_worktree_git_dir( + tmp_path, loose=None, packed=f"{SHA} refs/heads/feature\n" + ) + assert _resolve_sha(str(wt)) == SHA + + +def _git(cwd, *args): + import subprocess + + return subprocess.run( + ["git", "-c", "user.name=kompot-test", "-c", "user.email=test@example.invalid", + "-c", "commit.gpgsign=false", *args], + cwd=str(cwd), check=True, capture_output=True, text=True, + ).stdout.strip() + + +@pytest.mark.parametrize("pack", [False, True], ids=["loose", "packed"]) +@pytest.mark.parametrize("shape", ["clone", "branch-worktree", "detached-worktree"]) +def test_resolve_sha_matches_git_rev_parse(tmp_path, shape, pack): + """End to end against real git: every checkout shape resolves, loose or packed. + + `branch-worktree` is the shape that returned None before #28/#33 were fixed; + `detached-worktree` passed all along because its HEAD holds the sha. + """ + import shutil + + if shutil.which("git") is None: + pytest.skip("git binary not available") + repo = tmp_path / "repo" + repo.mkdir() + _git(repo, "init", "-q") + (repo / "f.txt").write_text("x\n") + _git(repo, "add", "f.txt") + _git(repo, "commit", "-q", "-m", "init") + + checkout = repo + if shape == "branch-worktree": + checkout = tmp_path / "wt" + _git(repo, "worktree", "add", "-b", "some/branch", str(checkout)) + (checkout / "g.txt").write_text("y\n") + _git(checkout, "add", "g.txt") + _git(checkout, "commit", "-q", "-m", "on the branch") + elif shape == "detached-worktree": + checkout = tmp_path / "wt" + _git(repo, "worktree", "add", "--detach", str(checkout)) + if pack: + _git(repo, "pack-refs", "--all", "--prune") + + pkg = checkout / "kompot" + pkg.mkdir() + git_dir = _find_git_dir(str(pkg)) + assert git_dir is not None + assert _resolve_sha(git_dir) == _git(checkout, "rev-parse", "HEAD") + + +@pytest.mark.parametrize("shape", ["clone", "branch-worktree"]) +def test_vendored_tree_does_not_borrow_an_unrelated_repo_sha(tmp_path, monkeypatch, shape): + """A kompot tree nested below another repository's root must not stamp its HEAD. + + The enclosing repo resolves perfectly well -- which is the danger: the + stamp would be a confident, wrong sha. Only a `.git` at the package's + parent counts. (A copy at another repository's root is not caught; by + position it looks like Kompot's own checkout.) + """ + import shutil + + if shutil.which("git") is None: + pytest.skip("git binary not available") + repo = tmp_path / "other" + repo.mkdir() + _git(repo, "init", "-q") + (repo / "f.txt").write_text("x\n") + _git(repo, "add", "f.txt") + _git(repo, "commit", "-q", "-m", "init") + checkout = repo + if shape == "branch-worktree": + checkout = tmp_path / "wt" + _git(repo, "worktree", "add", "-b", "some/branch", str(checkout)) + pkg = checkout / "vendor" / "lib" / "kompot" + pkg.mkdir(parents=True) + monkeypatch.setattr(_provenance, "__file__", str(pkg / "_provenance.py")) + monkeypatch.setattr(_provenance, "_pep610_editable", lambda: None) + + # The walk-up alone would find and resolve the unrelated repo... + assert _resolve_sha(_find_git_dir(str(pkg))) == _git(checkout, "rev-parse", "HEAD") + # ...and provenance must refuse it. + assert _provenance._resolve()["kompot_git_sha"] is None + + +def test_own_checkout_root_is_accepted(tmp_path, monkeypatch): + _make_git_dir(tmp_path) + pkg = tmp_path / "kompot" + pkg.mkdir() + monkeypatch.setattr(_provenance, "__file__", str(pkg / "_provenance.py")) + monkeypatch.setattr(_provenance, "_pep610_editable", lambda: True) + assert _provenance._resolve()["kompot_git_sha"] == SHA + + +def test_unresolvable_sha_in_a_work_tree_is_logged(tmp_path, monkeypatch, caplog): + """A None sha inside a checkout is a resolution failure and must be visible.""" + _make_git_dir(tmp_path, loose=None) + pkg = tmp_path / "kompot" + pkg.mkdir() + monkeypatch.setattr(_provenance, "__file__", str(pkg / "_provenance.py")) + monkeypatch.setattr(_provenance, "_pep610_editable", lambda: True) + + with caplog.at_level("WARNING", logger="kompot"): + result = _provenance._resolve() + assert result["kompot_git_sha"] is None + assert any("could not be resolved" in r.getMessage() for r in caplog.records) + + +def test_resolved_sha_logs_nothing(tmp_path, monkeypatch, caplog): + _make_git_dir(tmp_path) + pkg = tmp_path / "kompot" + pkg.mkdir() + monkeypatch.setattr(_provenance, "__file__", str(pkg / "_provenance.py")) + monkeypatch.setattr(_provenance, "_pep610_editable", lambda: True) + + with caplog.at_level("WARNING", logger="kompot"): + assert _provenance._resolve()["kompot_git_sha"] == SHA + assert not caplog.records + + def test_looks_like_sha_rejects_abbreviations(): assert _looks_like_sha(SHA) assert not _looks_like_sha(SHA[:7]) @@ -286,8 +440,10 @@ def test_source_checkout_resolves_sha_and_editable(): resolve and the wheel-degradation contract applies instead. """ package_dir = os.path.dirname(os.path.abspath(_provenance.__file__)) - if _is_installed_tree(package_dir) or _find_git_dir(package_dir) is None: - pytest.skip("kompot is not running from a source checkout") + # A .git at the package's parent, not merely an enclosing work tree: a tree + # exported or vendored below another repository's root must skip here. + if _is_installed_tree(package_dir) or _provenance._find_kompot_git_dir(package_dir) is None: + pytest.skip("kompot is not running from its own source checkout") p = get_provenance() assert p["kompot_version"] diff --git a/tests/test_resource_estimation.py b/tests/test_resource_estimation.py index d2d7c1f..e15134c 100644 --- a/tests/test_resource_estimation.py +++ b/tests/test_resource_estimation.py @@ -196,7 +196,16 @@ def test_dry_run_with_sample_variance(): def test_dry_run_with_disk_storage(): - """Test dry run with disk storage enabled.""" + """Disk is charged on the path that writes, and only on that path. + + Re-derived for 0.9.0 (settylab/kompot#26). This used to assert a bare + ``len(disk_reqs) > 0``, which passed for the wrong reason: the estimator + charged disk unconditionally, including for the Dask path, which writes + nothing to ``disk_storage_dir`` at all. The expected byte count below comes + from the allocation -- two ``(n_landmarks, n_landmarks, n_genes)`` float64 + tensors, one per condition -- not from observing a run. + """ + from kompot import resource_estimation from kompot.resource_estimation import dry_run_differential_expression import tempfile @@ -210,23 +219,51 @@ def test_dry_run_with_disk_storage(): } adata = ad.AnnData(X=X, obs=obs_data) + n_landmarks = 20 + expected_disk = 2 * n_landmarks * n_landmarks * n_genes * 8 + + def plan_with(dask_available, tmpdir): + original = resource_estimation.DASK_AVAILABLE + resource_estimation.DASK_AVAILABLE = dask_available + try: + return dry_run_differential_expression( + adata, + condition1="A", + condition2="B", + groupby="condition", + use_sample_variance=True, + sample_column="sample", + store_arrays_on_disk=True, + disk_storage_dir=tmpdir, + n_landmarks=n_landmarks, + null_genes=0, + verbose=False, + ) + finally: + resource_estimation.DASK_AVAILABLE = original with tempfile.TemporaryDirectory() as tmpdir: - plan = dry_run_differential_expression( - adata, - condition1="A", - condition2="B", - groupby="condition", - use_sample_variance=True, - sample_column="sample", - store_arrays_on_disk=True, - disk_storage_dir=tmpdir, - verbose=False, - ) - - # Should have disk requirements - disk_reqs = [r for r in plan.requirements if r.resource_type == "disk"] - assert len(disk_reqs) > 0 + # No dask: the tensors are written as memory-mapped .npy files. + without_dask = plan_with(False, tmpdir) + disk_reqs = [r for r in without_dask.requirements if r.resource_type == "disk"] + assert len(disk_reqs) == 1 + assert disk_reqs[0].size_bytes == expected_disk + assert without_dask.total_disk_required == expected_disk + + # With dask: lazily evaluated one gene at a time, nothing written. + with_dask = plan_with(True, tmpdir) + assert with_dask.total_disk_required == 0 + assert [r for r in with_dask.requirements if r.resource_type == "disk"] == [] + + # Either way the tensors are not charged to memory. + for plan in (without_dask, with_dask): + in_memory = [ + r + for r in plan.requirements + if r.name.startswith("Sample covariances") + and r.resource_type == "memory" + ] + assert in_memory == [] def test_format_report(): diff --git a/tests/test_sample_variance_cost_regressions.py b/tests/test_sample_variance_cost_regressions.py new file mode 100644 index 0000000..1610890 --- /dev/null +++ b/tests/test_sample_variance_cost_regressions.py @@ -0,0 +1,489 @@ +"""Regression guards for the sample-variance cost defects fixed in 0.9.0. + +Three distinct failures, all of which were silent — the code returned a +plausible answer at rc 0 and only a measurement disagreed: + +* ``settylab/kompot#25`` — ``kompot.de(dry_run=True)`` priced a run with zero + null genes while the run it was pricing would add 2 000. +* ``settylab/kompot#26`` — the per-gene covariance tensors were summed into a + third dense tensor before use, so ``store_arrays_on_disk=True`` did not keep + them out of memory, and ``compute_mahalanobis_distances`` materialised the + whole tensor again into a variable it never read. +* The estimator charged that third tensor, and charged disk for the Dask path + that writes nothing. +""" + +import io +import contextlib +import warnings + +import anndata as ad +import numpy as np +import pandas as pd +import pytest + +import kompot +from kompot.utils import LazyGeneCovariance + + +def _adata(n_cells=240, n_genes=12, n_samples=3, n_dims=6, seed=0): + rng = np.random.default_rng(seed) + obs = pd.DataFrame( + { + "condition": np.where(np.arange(n_cells) < n_cells // 2, "A", "B"), + "donor": [f"d{i % n_samples}" for i in range(n_cells)], + }, + index=[f"c{i}" for i in range(n_cells)], + ) + adata = ad.AnnData( + X=rng.normal(size=(n_cells, n_genes)), + obs=obs, + var=pd.DataFrame(index=[f"g{i}" for i in range(n_genes)]), + ) + adata.obsm["DM_EigenVectors"] = rng.normal(size=(n_cells, n_dims)) + return adata + + +def _plan(adata, **kwargs): + with contextlib.redirect_stdout(io.StringIO()): + return kompot.de( + adata, + groupby="condition", + condition1="A", + condition2="B", + obsm_key="DM_EigenVectors", + dry_run=True, + **kwargs, + ) + + +def _named(plan, needle): + return [r for r in plan.requirements if needle in r.name] + + +# -------------------------------------------------------------------------- +# settylab/kompot#25 — the dry run must price the run it is pricing +# -------------------------------------------------------------------------- + + +def test_dry_run_resolves_auto_null_genes(): + """A default dry run must account for the 2 000 null genes the run adds. + + ``null_genes="auto"`` used to reach the estimator unresolved, where it + matched neither the int nor the list branch and silently counted zero. + """ + adata = _adata() + auto = _plan(adata) # FDRSettings default -> "auto" + explicit = _plan(adata, fdr=kompot.FDRSettings(null_genes=2000)) + + assert auto.total_memory_required == explicit.total_memory_required + fold_change = _named(auto, "Fold change")[0] + assert fold_change.shape == (adata.n_obs, adata.n_vars + 2000) + + +def test_dry_run_auto_null_genes_is_zero_with_sample_col(): + """With ``sample_col`` set, ``"auto"`` resolves to 0, so no null genes.""" + adata = _adata() + auto = _plan(adata, sample_col="donor") + explicit = _plan(adata, sample_col="donor", fdr=kompot.FDRSettings(null_genes=0)) + + assert auto.total_memory_required == explicit.total_memory_required + assert _named(auto, "Fold change")[0].shape == (adata.n_obs, adata.n_vars) + + +def test_estimator_refuses_unresolved_null_genes(): + """The estimator must refuse a sentinel it cannot price, not count zero.""" + from kompot.resource_estimation import estimate_differential_expression_resources + + adata = _adata() + with pytest.raises(ValueError, match="must be resolved"): + estimate_differential_expression_resources( + adata, "A", "B", "condition", null_genes="auto" + ) + + +# -------------------------------------------------------------------------- +# settylab/kompot#26 — the per-gene tensor is assembled one gene at a time +# -------------------------------------------------------------------------- + + +def test_lazy_gene_covariance_matches_eager_sum(): + """The view must be numerically identical to the dense sum it replaces.""" + rng = np.random.default_rng(1) + n_points, n_genes = 5, 4 + a = rng.normal(size=(n_points, n_points, n_genes)) + b = rng.normal(size=(n_points, n_points, n_genes)) + base = rng.normal(size=(n_points, n_points)) + + view = LazyGeneCovariance([a, b], base=base) + eager = a + b + base[:, :, None] + + assert view.shape == (n_points, n_points, n_genes) + assert view.ndim == 3 + for g in range(n_genes): + np.testing.assert_allclose(view[:, :, g], eager[:, :, g], rtol=0, atol=0) + + +def test_lazy_gene_covariance_does_not_mutate_its_terms(): + """A NumPy basic slice is a view; the accumulation must copy first.""" + rng = np.random.default_rng(2) + a = rng.normal(size=(4, 4, 3)) + base = rng.normal(size=(4, 4)) + before = a.copy() + + view = LazyGeneCovariance([a], base=base) + for g in range(3): + view.gene(g) + view.gene(g) # twice: a second call must not accumulate again + + np.testing.assert_array_equal(a, before) + np.testing.assert_allclose(view.gene(0), before[:, :, 0] + base) + + +def test_lazy_gene_covariance_rejects_bad_input(): + with pytest.raises(ValueError, match="at least one term"): + LazyGeneCovariance([]) + with pytest.raises(ValueError, match="share a shape"): + LazyGeneCovariance([np.zeros((3, 3, 2)), np.zeros((3, 3, 4))]) + with pytest.raises(ValueError, match="3-D"): + LazyGeneCovariance([np.zeros((3, 3))]) + with pytest.raises(ValueError, match="base has shape"): + LazyGeneCovariance([np.zeros((3, 3, 2))], base=np.zeros((4, 4))) + with pytest.raises(TypeError, match="only"): + LazyGeneCovariance([np.zeros((3, 3, 2))])[0] + + +def test_mahalanobis_accepts_lazy_view_and_matches_dense(): + """The consumer must give the same distances for a view and a dense tensor.""" + from kompot.utils import compute_mahalanobis_distances + + rng = np.random.default_rng(3) + n_points, n_genes = 8, 5 + base = np.eye(n_points) * 2.0 + + def spd(): + m = rng.normal(size=(n_points, n_points)) + return m @ m.T / n_points + np.eye(n_points) + + a = np.stack([spd() for _ in range(n_genes)], axis=2) + b = np.stack([spd() for _ in range(n_genes)], axis=2) + diffs = rng.normal(size=(n_genes, n_points)) + + dense = a + b + base[:, :, None] + view = LazyGeneCovariance([a, b], base=base) + + from_dense = compute_mahalanobis_distances( + diffs, dense, jit_compile=False, progress=False + ) + from_view = compute_mahalanobis_distances( + diffs, view, jit_compile=False, progress=False + ) + np.testing.assert_allclose(from_view, from_dense, rtol=1e-12, atol=0) + + +def test_de_with_sample_variance_produces_finite_scores(): + """End-to-end: the lazy path still yields a usable sample-variance column.""" + adata = _adata(n_cells=200, n_genes=8, n_samples=3) + with warnings.catch_warnings(), contextlib.redirect_stdout(io.StringIO()): + warnings.simplefilter("ignore") + kompot.de( + adata, + groupby="condition", + condition1="A", + condition2="B", + obsm_key="DM_EigenVectors", + sample_col="donor", + gp=kompot.GPSettings(n_landmarks=40, random_state=0), + fdr=kompot.FDRSettings(null_genes=0), + output=kompot.OutputSettings(progress=False), + ) + col = "kompot_de_A_to_B_mahalanobis_sample_var" + assert col in adata.var + assert np.isfinite(adata.var[col].to_numpy()).all() + + +# -------------------------------------------------------------------------- +# The plan must describe the code that now runs +# -------------------------------------------------------------------------- + + +def test_plan_no_longer_charges_a_third_covariance_tensor(): + """In-memory mode holds variance1 + variance2, and nothing else. + + The expected figure is re-derived from the allocation, not copied from a + previous run: two (n_landmarks, n_landmarks, n_genes) float64 tensors. + """ + adata = _adata(n_cells=200, n_genes=10, n_samples=3) + plan = _plan( + adata, + sample_col="donor", + gp=kompot.GPSettings(n_landmarks=40), + storage=kompot.StorageSettings(store_arrays_on_disk=False), + ) + + assert _named(plan, "Combined sample covariances") == [] + sv = _named(plan, "Sample covariances") + assert len(sv) == 1 + assert sv[0].resource_type == "memory" + assert sv[0].size_bytes == 2 * 40 * 40 * adata.n_vars * 8 + + +def test_plan_disk_column_matches_the_selected_backend(): + """Disk is charged only on the path that actually writes files. + + With dask the tensor is a lazy graph and nothing reaches + ``disk_storage_dir``; without dask it is written as memory-mapped ``.npy``. + """ + from kompot import resource_estimation + + adata = _adata(n_cells=200, n_genes=10, n_samples=3) + kwargs = dict( + sample_col="donor", + gp=kompot.GPSettings(n_landmarks=40), + storage=kompot.StorageSettings(store_arrays_on_disk=True), + ) + expected = 2 * 40 * 40 * adata.n_vars * 8 + + original = resource_estimation.DASK_AVAILABLE + try: + resource_estimation.DASK_AVAILABLE = True + with_dask = _plan(adata, **kwargs) + resource_estimation.DASK_AVAILABLE = False + without_dask = _plan(adata, **kwargs) + finally: + resource_estimation.DASK_AVAILABLE = original + + assert with_dask.total_disk_required == 0 + assert _named(with_dask, "Sample covariances") == [] + + disk_reqs = [r for r in without_dask.requirements if r.resource_type == "disk"] + assert len(disk_reqs) == 1 + assert disk_reqs[0].size_bytes == expected + assert without_dask.total_disk_required == expected + + +def test_plan_charges_the_per_gene_working_set(): + """Whatever the backend, one gene's matrices are held at a time.""" + adata = _adata(n_cells=200, n_genes=10, n_samples=3) + for on_disk in (False, True): + plan = _plan( + adata, + sample_col="donor", + gp=kompot.GPSettings(n_landmarks=40), + storage=kompot.StorageSettings(store_arrays_on_disk=on_disk), + ) + working = _named(plan, "Per-gene covariance working set") + assert len(working) == 1, on_disk + assert working[0].resource_type == "memory" + # two sample-variance terms + the shared posterior covariance + assert working[0].size_bytes == 3 * 40 * 40 * 8 + + +# -------------------------------------------------------------------------- +# The restructured branches in DifferentialExpression.compute_mahalanobis +# -------------------------------------------------------------------------- + + +def _spd(rng, n): + m = rng.normal(size=(n, n)) + return m @ m.T / n + np.eye(n) + + +@pytest.mark.parametrize("n_terms", [1, 2]) +def test_lazy_view_matches_dense_for_one_or_two_predictors(n_terms): + """Both predictors, or only one, must give the dense arithmetic's answer. + + The fix replaced a nested if/else over (variance_predictor1, + variance_predictor2) x (gene-specific, shared) with a flat list of terms. + The single-predictor arms are reachable through ``ModelSettings`` and are + easy to leave untested. + """ + from kompot.utils import compute_mahalanobis_distances + + rng = np.random.default_rng(11) + n_points, n_genes = 7, 4 + base = _spd(rng, n_points) + terms = [ + np.stack([_spd(rng, n_points) for _ in range(n_genes)], axis=2) + for _ in range(n_terms) + ] + diffs = rng.normal(size=(n_genes, n_points)) + + dense = sum(terms) + base[:, :, None] + view = LazyGeneCovariance(terms, base=base) + + from_view = compute_mahalanobis_distances( + diffs, view, jit_compile=False, progress=False + ) + from_dense = compute_mahalanobis_distances( + diffs, dense, jit_compile=False, progress=False + ) + np.testing.assert_allclose(from_view, from_dense, rtol=1e-12, atol=0) + + +def test_lazy_view_refuses_a_shared_term(): + """A 2-D sample variance is folded into ``base`` by the caller. + + It must never be handed to the view as a term; if it is, that is a bug in + the caller and the view should say so rather than broadcast silently. + """ + rng = np.random.default_rng(12) + gene_specific = np.stack([_spd(rng, 5) for _ in range(3)], axis=2) + shared = _spd(rng, 5) + + with pytest.raises(ValueError, match="share a shape"): + LazyGeneCovariance([gene_specific, shared]) + + +# -------------------------------------------------------------------------- +# Potency: guards for the ALLOCATION property, not just the numbers +# +# The first version of this file asserted numerical equivalence and input +# validation, and an eagerly-materialising implementation satisfies both. All +# 14 cases passed against a mutant that rebuilt the dense +# (n_points, n_points, n_genes) sum in __init__ -- settylab/kompot#26's exact +# defect. Base absence could not stand in either: the pre-existing suite passed +# on 19ee1a1, which HAD that behaviour. +# +# A guard for an allocation property has to be able to fail on an allocation +# mutation, so these two assert HOW the terms are read rather than what comes +# out. +# -------------------------------------------------------------------------- + + +class _RecordingTerm: + """A 3-D array that records how it is read. + + Distinguishes per-gene slicing (``[:, :, g]``) from any read that touches + the array as a whole -- ``np.asarray``, ``+``, or a slice spanning genes. + An eager implementation cannot avoid the second kind; a lazy one must never + perform it. + """ + + def __init__(self, arr): + self._a = arr + self.shape = arr.shape + self.dtype = arr.dtype + self.ndim = arr.ndim + self.gene_reads = [] + self.bulk_reads = 0 + + def _is_gene_key(self, key): + return ( + isinstance(key, tuple) + and len(key) == 3 + and key[0] == slice(None) + and key[1] == slice(None) + and isinstance(key[2], (int, np.integer)) + ) + + def __getitem__(self, key): + if self._is_gene_key(key): + self.gene_reads.append(int(key[2])) + else: + self.bulk_reads += 1 + return self._a[key] + + def __array__(self, dtype=None): + self.bulk_reads += 1 + arr = self._a + return arr if dtype is None else arr.astype(dtype) + + def __add__(self, other): + self.bulk_reads += 1 + return self._a + (other._a if isinstance(other, _RecordingTerm) else other) + + __radd__ = __add__ + + +def test_mahalanobis_reads_the_terms_one_gene_at_a_time(): + """POTENCY GUARD for #26: no whole-tensor read may occur. + + Fails on the eager-dense mutant, which reaches the terms through ``+`` + rather than through ``[:, :, g]``. Verified red against that mutation -- + see the module docstring. + """ + from kompot.utils import compute_mahalanobis_distances + + rng = np.random.default_rng(31) + n_points, n_genes = 6, 5 + + def spd(): + m = rng.normal(size=(n_points, n_points)) + return m @ m.T / n_points + np.eye(n_points) + + raw = [ + np.stack([spd() for _ in range(n_genes)], axis=2), + np.stack([spd() for _ in range(n_genes)], axis=2), + ] + terms = [_RecordingTerm(t) for t in raw] + view = LazyGeneCovariance(terms, base=np.eye(n_points)) + diffs = rng.normal(size=(n_genes, n_points)) + + compute_mahalanobis_distances(diffs, view, jit_compile=False, progress=False) + + for i, term in enumerate(terms): + assert term.bulk_reads == 0, ( + f"term {i} was read as a whole {term.bulk_reads} time(s); the " + f"covariance must only ever be reached one gene at a time" + ) + assert sorted(term.gene_reads) == list(range(n_genes)), ( + f"term {i} gene reads were {sorted(term.gene_reads)}, " + f"expected each of {list(range(n_genes))} exactly once" + ) + + +def test_peak_memory_stays_bounded_by_one_gene(): + """POTENCY GUARD for #26, measured rather than structural. + + Builds a tensor far larger than one gene and asserts that walking it + through the Mahalanobis step grows anonymous memory by roughly one gene's + matrix rather than by the tensor. The eager mutant allocates the whole + thing and blows the bound. + """ + from kompot.utils import compute_mahalanobis_distances + + def anon_bytes(): + try: + with open("/proc/self/smaps_rollup") as fh: + for line in fh: + key, _, value = line.partition(":") + if key == "Anonymous": + return int(value.split()[0]) * 1024 + except OSError: + return None + return None + + if anon_bytes() is None: + pytest.skip("/proc/self/smaps_rollup unavailable") + + rng = np.random.default_rng(32) + n_points, n_genes = 260, 120 # one gene 0.52 MiB, tensor 62 MiB + one_gene = n_points * n_points * 8 + base = np.eye(n_points) + spd = rng.normal(size=(n_points, n_points)) + spd = spd @ spd.T / n_points + np.eye(n_points) + terms = [np.repeat(spd[:, :, None], n_genes, axis=2) for _ in range(2)] + diffs = rng.normal(size=(n_genes, n_points)) + + # The window MUST open before the constructor. An eager implementation + # materialises in __init__ -- that is #26's actual shape -- so a window + # opened after it measures only incidental churn, whose size depends on + # allocator state left by sibling tests. Measured: with the window opened + # after construction this guard passed against the eager mutant in every + # whole-file run and failed only when run alone, i.e. it was green in the + # configuration CI uses. + before = anon_bytes() + view = LazyGeneCovariance(terms, base=base) + compute_mahalanobis_distances(diffs, view, jit_compile=False, progress=False) + grew = anon_bytes() - before + + tensor = one_gene * n_genes + # Generous: 12 gene-matrices of headroom for interpreter churn, while still + # an order of magnitude below the 120-gene tensor the mutant materialises. + assert grew < 12 * one_gene, ( + f"anonymous memory grew {grew / 2**20:.1f} MiB walking a " + f"{tensor / 2**20:.1f} MiB tensor; one gene is " + f"{one_gene / 2**20:.2f} MiB, so the covariance is not being " + f"assembled a gene at a time" + ) diff --git a/tests/test_sample_variance_graph_scaling.py b/tests/test_sample_variance_graph_scaling.py new file mode 100644 index 0000000..17cf280 --- /dev/null +++ b/tests/test_sample_variance_graph_scaling.py @@ -0,0 +1,219 @@ +"""Regression tests for the gene-specific (sample-variance) covariance path. + +These guard a SCALING property, not a single timing. The defect they cover was not +a constant factor: with a Dask-backed covariance tensor the per-gene +``__setitem__`` loop in ``DifferentialExpression.compute_mahalanobis_distances`` +bound its result to the same object as its input, so every assignment layered onto +the graph the previous iteration had just built. Task count grew QUADRATICALLY in +``n_genes`` (measured: 1,804 tasks after one iteration, 994,740 after 120), which +is why runs of a few hundred genes did not finish while the same arithmetic done +directly takes ~16 ms/gene. + +Task count is used rather than wall time on purpose: it is deterministic, so these +tests do not flake on a loaded CI machine. +""" + +import numpy as np +import pytest + +from kompot.memory_utils import DASK_AVAILABLE +from kompot.utils import compute_mahalanobis_distances + +N_POINTS = 24 +N_GROUPS = 4 + + +def _shared_cov(n_points=N_POINTS, seed=0): + rng = np.random.default_rng(seed) + a = rng.normal(size=(n_points, 6)) * 0.3 + return a @ a.T + np.eye(n_points) + + +def _centred(n_genes, seed, n_points=N_POINTS): + """(n_points, n_groups, n_genes), as SampleVarianceEstimator.predict builds.""" + rng = np.random.default_rng(seed) + p = rng.normal(size=(N_GROUPS, n_points, n_genes)) * 0.1 + p = p - p.mean(axis=0, keepdims=True) + return np.moveaxis(p, 1, 0) + + +def _dask_gene_covariance(n_genes, seed, n_points=N_POINTS): + """Rebuilds the Dask array shape produced by the disk-backed estimator path.""" + import dask + import dask.array as da + + centred = _centred(n_genes, seed, n_points) + + @dask.delayed + def gene_cov(gene_data, n_groups): + return np.dot(gene_data, gene_data.T) / (n_groups - 1) + + return da.stack( + [ + da.from_delayed( + gene_cov(centred[:, :, g], N_GROUPS), + shape=(n_points, n_points), + dtype=np.float64, + ) + for g in range(n_genes) + ], + axis=2, + ) + + +def _numpy_gene_covariance(n_genes, seed, n_points=N_POINTS): + centred = _centred(n_genes, seed, n_points) + out = np.empty((n_points, n_points, n_genes)) + for g in range(n_genes): + gc = centred[:, :, g] + out[:, :, g] = gc @ gc.T / (N_GROUPS - 1) + return out + + +def _reference_distances(n_genes, diffs, shared, seed_pair=(1, 2)): + """Straight per-gene computation, independent of the code under test.""" + from scipy.linalg import solve_triangular + + c1 = _centred(n_genes, seed_pair[0]) + c2 = _centred(n_genes, seed_pair[1]) + out = np.empty(n_genes) + for g in range(n_genes): + a, b = c1[:, :, g], c2[:, :, g] + cov = (a @ a.T + b @ b.T) / (N_GROUPS - 1) + shared + np.eye(N_POINTS) * 1e-8 + chol = np.linalg.cholesky(cov) + out[g] = np.sqrt(np.sum(solve_triangular(chol, diffs[g], lower=True) ** 2)) + return out + + +def _task_count(arr): + """Graph size of the object handed to the Mahalanobis step, or None if it has no graph. + + None is a PASS, not a skip: an object that carries no task graph (a lazy view + that materialises one gene at a time) cannot exhibit the graph blowup this + module guards against. Returning None rather than raising is what lets the + same test express the property across implementations that differ in KIND, + not just in degree. + """ + graph = getattr(arr, "__dask_graph__", None) + if graph is None: + return None + return len(graph()) + + +class _FakeFunctionPredictor: + """Minimal stand-in for a fitted mellon predictor.""" + + def __init__(self, landmarks, n_genes, seed): + self.landmarks = landmarks + rng = np.random.default_rng(seed) + self._pred = rng.normal(size=(len(landmarks), n_genes)) * 0.2 + a = rng.normal(size=(len(landmarks), 6)) * 0.3 + self._cov = a @ a.T + np.eye(len(landmarks)) + + def __call__(self, X): + return self._pred + + def covariance(self, X, diag=False): + return self._cov + + +def _make_variance_predictor(n_genes, seed, n_points=N_POINTS): + def predict(X, diag=False, progress=False): + assert not diag, "the DE path must request the full covariance" + return _dask_gene_covariance(n_genes, seed, n_points) + + return predict + + +def _covariance_handed_to_mahalanobis(n_genes, monkeypatch): + """Run the real DE path and return the covariance object it passes downstream. + + Intercepting at this boundary is what makes this test version independent: it + measures the array the shipped code actually builds, whatever route it took to + build it. + """ + from kompot.differential import differential_expression as de_mod + + landmarks = np.linspace(0, 1, N_POINTS).reshape(-1, 1) + captured = {} + + def spy(diff_values, covariance, **kwargs): + captured["covariance"] = covariance + return np.zeros(np.shape(diff_values)[0]) + + monkeypatch.setattr(de_mod, "compute_mahalanobis_distances", spy) + + de = de_mod.DifferentialExpression( + use_sample_variance=True, + function_predictor1=_FakeFunctionPredictor(landmarks, n_genes, 11), + function_predictor2=_FakeFunctionPredictor(landmarks, n_genes, 12), + variance_predictor1=_make_variance_predictor(n_genes, 1), + variance_predictor2=_make_variance_predictor(n_genes, 2), + ) + de.compute_mahalanobis_distances( + X=landmarks, landmarks_override=landmarks, progress=False + ) + return captured["covariance"] + + +@pytest.mark.skipif(not DASK_AVAILABLE, reason="dask is not installed") +def test_gene_specific_covariance_graph_scales_linearly(monkeypatch): + """Task count must grow at most ~linearly with n_genes, not quadratically. + + This is the regression guard for the defect. It drives the real + ``DifferentialExpression.compute_mahalanobis_distances`` and measures the graph + of the covariance tensor that method hands downstream, so it exercises whatever + the shipped code does rather than a particular helper. + + Before the fix the per-gene ``__setitem__`` loop aliased its own input, so each + assignment layered onto the graph the previous iteration built and the task + count grew quadratically. The bound below is deliberately generous - 4x genes + may cost up to 8x tasks - so it tolerates a change in how dask chunks a stack + while still failing on quadratic growth, which at these sizes is ~16x. + """ + small, large = 20, 80 # 4x genes + + tasks = {} + for n_genes in (small, large): + cov = _covariance_handed_to_mahalanobis(n_genes, monkeypatch) + assert cov.shape == (N_POINTS, N_POINTS, n_genes) + tasks[n_genes] = _task_count(cov) + + if tasks[small] is None or tasks[large] is None: + # A lazy view with no task graph: the blowup is not merely bounded here, + # it is unrepresentable. Nothing left to assert. + return + + growth = tasks[large] / tasks[small] + assert growth <= 8.0, ( + "gene-specific covariance graph grew superlinearly: " + f"{tasks[small]} tasks at {small} genes vs {tasks[large]} at {large} " + f"({growth:.1f}x for 4x the genes; linear is 4x, quadratic is 16x)" + ) + + +@pytest.mark.parametrize("backend", ["numpy", "dask"]) +def test_gene_specific_mahalanobis_matches_reference(backend): + """The optimised path must return exactly what a direct computation returns.""" + if backend == "dask" and not DASK_AVAILABLE: + pytest.skip("dask is not installed") + + n_genes = 16 + shared = _shared_cov() + rng = np.random.default_rng(7) + diffs = rng.normal(size=(n_genes, N_POINTS)) * 0.2 + + build = _dask_gene_covariance if backend == "dask" else _numpy_gene_covariance + # Built without the helper so this checks the MATHS on any build. + cov = build(n_genes, seed=1) + build(n_genes, seed=2) + shared[:, :, None] + + got = compute_mahalanobis_distances( + diff_values=diffs, covariance=cov, batch_size=None, + jit_compile=False, eps=1e-8, progress=False, + ) + expected = _reference_distances(n_genes, diffs, shared) + + assert np.all(np.isfinite(got)) + np.testing.assert_allclose(got, expected, rtol=1e-8, atol=1e-10) + +