From 27f1fcfc759aed5b4af276bcbbce55e9f093ae7b Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 12:28:22 -0700 Subject: [PATCH 01/28] docs: make the cost of sample variance explicit and prescribe the two-pass workflow Supplying `sample_col` to `kompot.de()` is the one option whose cost is multiplied by the gene count, and nothing a user read on the way to that decision said so. Without it, every gene is scored against a single shared posterior covariance of shape (n_landmarks, n_landmarks); with it, every gene gets its own, so the dominant allocation becomes `3 * n_landmarks**2 * n_genes * 8` bytes -- about 0.56 GiB per gene at the default n_landmarks=5000, ~560 GiB for 1 000 genes, and terabytes for a whole transcriptome. Compute scales too: the Mahalanobis step then performs one Cholesky factorisation per gene, in a Python loop, instead of one in total, and GPSettings.batch_size does not bound that loop. New canonical guide docs/source/resource_planning.rst, linked from the other surfaces rather than duplicated into them. It derives the scaling from the allocation in the source, carries measured dry-run plans at realistic sizes, prescribes the two-pass workflow (cheap first pass over all genes, then sample variance restricted to the top ~1 000), ranks the levers -- `genes` linear, `n_landmarks` quadratic, `null_genes=0`, disk offload -- and shows how to price a run with `dry_run=True`. The same warning now appears where the decision is actually made: `kompot.de`'s docstring (so `help()` carries it), GPSettings.n_landmarks, StorageSettings, SampleVarianceEstimator, `kompot de --sample-col`, the DE config template, the README, and notebooks 02 and 03 (markdown cells only; outputs untouched). Two docs/code disagreements corrected rather than papered over: - StorageSettings.max_memory_ratio was documented as "Fraction of RAM before triggering disk storage". It triggers nothing. `store_arrays_on_disk=None` resolves to `disk_storage_dir is not None`, and the ratio only sets the threshold at which the estimator escalates warnings; SampleVarianceEstimator hardcodes 0.8 for its own check, so a user's value never reaches it. The DE config template's `store_arrays_on_disk: null # (null = auto)` implied the same non-existent automatic decision. - The CLI "Complete Analysis" example paired `--sample-col Sample` with `--n-landmarks 5000` over every gene: the single most expensive configuration the package can express (~11 TiB on a 20 000-gene dataset), presented as the recommended form. Replaced with the restricted second pass through a config file. Two measured behavioural defects are documented rather than fixed here, and filed separately: - #25 the dry run does not resolve `null_genes="auto"`, so a default-settings plan without `sample_col` omits the 2 000 null genes the run will add (measured 4.44x optimistic on a 4 000x500 input). - #26 `store_arrays_on_disk=True` does not reduce peak memory, and with dask installed nothing is written to disk at all. Measured at two scales with peak anonymous memory from /proc/self/smaps_rollup, against a plan predicting a 17x-28x drop. Numbers in the guide come from `kompot.de(..., dry_run=True)` on synthetic AnnData, reproducible with the snippets shown alongside them. Refs #27 --- CHANGELOG.md | 32 ++ README.md | 36 ++ docs/source/anndata.rst | 15 +- docs/source/cli.rst | 54 +- docs/source/index.rst | 13 +- docs/source/resource_planning.rst | 495 ++++++++++++++++++ docs/source/simplified.rst | 25 +- .../02_differential_expression_detailed.ipynb | 21 +- examples/03_sample_variance.ipynb | 61 ++- kompot/anndata/differential_expression.py | 40 +- kompot/cli/de.py | 8 +- kompot/cli/templates/de_config_template.yaml | 19 +- .../differential/sample_variance_estimator.py | 12 + kompot/settings.py | 34 +- 14 files changed, 824 insertions(+), 41 deletions(-) create mode 100644 docs/source/resource_planning.rst diff --git a/CHANGELOG.md b/CHANGELOG.md index a336fd0..80a7cbc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,38 @@ All notable changes to this project will be documented in this file. +## [Unreleased] + +### 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. Supplying + `sample_col` enables sample variance, which replaces the single shared posterior + covariance with one `(n_landmarks, n_landmarks)` matrix **per gene**, so the dominant + allocation is `3 x n_landmarks^2 x n_genes x 8` bytes: about 0.56 GiB per gene at the + default `n_landmarks=5000`, ~560 GiB for 1 000 genes, and terabytes for a whole + transcriptome. Compute scales too, since the Mahalanobis step then factorises once per + gene instead of once in total. The guide prescribes the two-pass workflow (a cheap + first pass over all genes, then sample variance restricted to the top ~1 000), ranks + the levers (`genes` linear, `n_landmarks` quadratic, `null_genes=0`, disk offload), + and shows how to price a run with `dry_run=True`. It carries measured dry-run plans + at realistic sizes. + - The same warning now appears where the decision is made: `kompot.de`'s docstring (so + `help()` carries it), `GPSettings.n_landmarks`, `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 via a config file. + - Corrected `StorageSettings.max_memory_ratio`, which was documented as "Fraction of RAM + before triggering disk storage". It triggers nothing: `store_arrays_on_disk=None` + 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. + - Documented two measured divergences between the dry run and the runs it prices: + the estimate does not resolve `null_genes="auto"` (#25), and `store_arrays_on_disk=True` + does not reduce peak memory, with nothing written to disk when `dask` is installed (#26). + ## [0.8.0] - 2026-07-28 ### Changed — statistics now match the manuscript diff --git a/README.md b/README.md index 241692d..fa58259 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,40 @@ 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**. The dominant allocation becomes +`3 x n_landmarks^2 x n_genes x 8` bytes: at the default 5 000 landmarks that is +roughly **0.56 GiB per gene**, so a whole transcriptome asks for terabytes. +Compute scales too, since the Mahalanobis step then factorises once per gene +instead of once in total. + +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(1000).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 +114,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..b02b0be 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 when ``null_genes`` is given as an explicit number - **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 36dc71a..37e9feb 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: about 0.56 GiB per gene at ``--n-landmarks 5000``, so a whole + transcriptome asks for terabytes. 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 ~1 000 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,22 @@ 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 +``3 x n_landmarks^2 x n_genes x 8`` bytes, so the levers that matter are the +**gene list** (linear) and ``--n-landmarks`` (quadratic). ``--batch-size`` +does not bound that term. Full treatment: +:doc:`Planning Memory and Disk `. Speed Optimization ^^^^^^^^^^^^^^^^^^ @@ -750,7 +788,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 +852,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/index.rst b/docs/source/index.rst index 0a5830e..187679f 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,21 @@ 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. At default settings that costs roughly + **0.56 GiB per gene**, so it must be run as a second pass over a restricted + gene list, never over the whole transcriptome. 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 +151,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/resource_planning.rst b/docs/source/resource_planning.rst new file mode 100644 index 0000000..f9d3ca0 --- /dev/null +++ b/docs/source/resource_planning.rst @@ -0,0 +1,495 @@ +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**. The cost of that step is multiplied by the number of genes +analysed, and at Kompot's default landmark count it reaches hundreds of +gigabytes for a routine gene set. + +This page explains where the 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 1 000; see :ref:`how-many-genes`) + and cut ``n_landmarks`` if the plan is still too large. + +.. 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(1000).index + + # ---- Pass 2: sample variance, restricted --------------------------- + kompot.de( + adata, "condition", "Young", "Old", + sample_col="donor_id", + genes=top_genes, # the lever that matters most + gp=kompot.GPSettings(n_landmarks=2000), # quadratic; see "Other levers" + storage=kompot.StorageSettings( + store_arrays_on_disk=True, # read "Disk offload" below first + 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 levers in that call +that genuinely shrink the allocation are ``genes`` and ``n_landmarks``; +``store_arrays_on_disk`` is discussed, with measurements, under +:ref:`disk-offload`. + +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 Kompot holds one per condition, plus their sum. So the dominant term is + +.. code-block:: text + + bytes = k x n_landmarks^2 x n_genes x 8 (float64) + + k = 3 variance1 + variance2 + their sum, all dense in memory + k = 2 what the planner charges to disk when store_arrays_on_disk=True, + on the assumption that the sum stays lazy + +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 a sample-variance + run costs about 0.56 GiB**. One thousand genes is roughly 560 GiB; ten + thousand is over five terabytes. Nothing else in a Kompot run behaves this + way, which is why the gene list on pass 2 is the number to decide + deliberately. + +Compute scales with genes as well, and 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. 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. + +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, 1 000 genes and 6 donors this is 0.45 GiB, negligible + beside the covariance tensors but not beside everything else. +* **Two extra layers**, ``__std``, each + ``n_cells x n_genes x 8`` B, written into ``adata.layers``. + + +Measured plans +-------------- + +The figures below are 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``. They 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: 14 20 20 32 + + * - genes + - no sample variance + - sample variance, RAM + - sample variance, disk-backed + * - 200 + - 3.1 GiB + - 115.0 GiB + - 3.2 GiB RAM + 74.5 GiB disk + * - 1 000 + - 6.4 GiB + - 566.0 GiB + - 7.2 GiB RAM + 372.5 GiB disk + * - 2 000 + - 10.6 GiB + - 1 130 GiB + - 12.1 GiB RAM + 745.1 GiB disk + * - 20 000 (all) + - 85.7 GiB + - 11 276 GiB + - 100.6 GiB RAM + 7 451 GiB disk + +The last row is the configuration to avoid: sample variance over a full +transcriptome at default settings asks for eleven terabytes. The +no-sample-variance column grows only with ``n_cells x n_genes``, which is why +pass 1 over all 20 000 genes is comfortable at 86 GiB. + +The fourth column is what the *planner* charges when +``store_arrays_on_disk=True``. Read it as the scratch footprint the run is +allowed to reach, not as a memory saving you can count on; see the warning +under :ref:`disk-offload`. + +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), + 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 + + base = dict(groupby="condition", condition1="Young", condition2="Old", + genes=top_genes, 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`. + +.. note:: + + ``kompot.resource_estimation.dry_run_differential_expression()`` is + deprecated. Use ``kompot.de(..., dry_run=True)``. + +.. _null-genes-cost: + +.. warning:: + + The dry run reads ``null_genes`` as you pass it and does not resolve the + ``"auto"`` default. With ``FDRSettings`` left at its default and no + ``sample_col``, the real run adds 2 000 null genes but the dry run counts + none, so the plan is optimistic by that margin. Pass + ``fdr=kompot.FDRSettings(null_genes=2000)`` to the dry run to see the true + figure. Runs *with* ``sample_col`` are unaffected, because ``"auto"`` + resolves to ``0`` there and the dry run also counts 0. Tracked in + `settylab/kompot#25 `_. + + Null genes are charged at the same per-gene rate as real ones. If you + override the default and request FDR alongside sample variance, each null + gene gets its own landmark covariance matrix too: 1 000 real genes plus + 2 000 null genes is a 3 000-gene tensor. + + +.. _disk-offload: + +Disk offload +------------ + +``StorageSettings(store_arrays_on_disk=True)`` is intended to change how the +per-gene covariance tensors are represented, so they are consumed one gene at +a time instead of built as three dense in-memory arrays. It is the flag the +dry run's warnings point you at, and the one the estimator prices into its +disk column. **Read the warning at the end of this section before you rely on +it**: measured runs do not show the memory saving. + +.. 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]'``) the tensor is a lazy + Dask graph, evaluated one gene at a time as the Mahalanobis step reaches it. + Measured: **no file is written at all**, so the dry run's disk figure is an + upper bound on scratch rather than a prediction of it. +* **Without** ``dask`` Kompot writes the full tensor for each condition as a + memory-mapped ``covariance_matrix.npy`` under ``disk_storage_dir``, and the + dry run's disk figure is then what actually lands there. This path is + sequential rather than parallel, and the dry run warns that ``dask`` is + missing. + +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 writes into 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 +warnings escalate; it does not switch storage modes. + +.. warning:: + + **Measured runs do not show the memory saving this flag promises.** + ``store_arrays_on_disk=True`` moves the tensor from the memory column of + the plan to the disk column, but peak memory of the actual ``kompot.de()`` + run is unchanged. Two synthetic configurations, 1 500 cells and 4 donors, + with peak *anonymous* memory sampled from ``/proc/self/smaps_rollup`` so + that page cache for a memory-mapped file is excluded: + + .. list-table:: 150 genes, 600 landmarks (tensor 412 MiB per condition) + :header-rows: 1 + :widths: 34 14 14 16 14 12 + + * - run + - plan RAM + - plan disk + - peak anon + - on disk + - wall + * - no sample variance + - 69 MiB + - 0 + - 962 MiB + - 0 + - 49 s + * - sample variance, in memory + - 1 313 MiB + - 0 + - 2 357 MiB + - 0 + - 87 s + * - on disk, with ``dask`` + - 77 MiB + - 824 MiB + - 2 781 MiB + - **0** + - 283 s + * - on disk, no ``dask`` + - 77 MiB + - 824 MiB + - 2 346 MiB + - 824 MiB + - 50 s + + .. list-table:: 120 genes, 900 landmarks (tensor 742 MiB per condition) + :header-rows: 1 + :widths: 34 14 14 16 14 12 + + * - run + - plan RAM + - plan disk + - peak anon + - on disk + - wall + * - no sample variance + - 73 MiB + - 0 + - 1 052 MiB + - 0 + - 55 s + * - sample variance, in memory + - 2 307 MiB + - 0 + - 3 505 MiB + - 0 + - 67 s + * - on disk, with ``dask`` + - 82 MiB + - 1 483 MiB + - 4 347 MiB + - **0** + - 259 s + * - on disk, no ``dask`` + - 82 MiB + - 1 483 MiB + - 3 562 MiB + - 1 483 MiB + - 113 s + + Two things to read off these. The peak-memory column is flat across the + three sample-variance rows at both scales, against a plan predicting a 17x + to 28x drop. And with ``dask`` installed **nothing reaches disk at all**: + the tensor is a lazy graph rather than a file, so the plan's disk figure + bounds scratch consumption rather than predicting it, and that path was + also the slowest of the four. Tracked in `settylab/kompot#26 + `_. + + Until that is resolved, treat **restricting genes** and **reducing + landmarks** as the levers that actually bound memory. They shrink the + allocation itself, and the dry run prices them correctly. + +Other levers +------------ + +``n_landmarks`` **is the strongest lever**, because the tensor is quadratic in +it. Same 1 000 genes, sample variance on disk, from the dry run: + +.. list-table:: + :header-rows: 1 + :widths: 20 20 20 + + * - ``n_landmarks`` + - RAM + - disk + * - 1 000 + - 6.1 GiB + - 14.9 GiB + * - 2 000 + - 6.3 GiB + - 59.6 GiB + * - 3 000 + - 6.5 GiB + - 134.1 GiB + * - 5 000 (default) + - 7.2 GiB + - 372.5 GiB + +Halving the landmark count quarters the covariance cost. Landmarks control +the resolution of the cell-state approximation, so this is a genuine accuracy +trade-off rather than a free saving, but 5 000 landmarks 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 during the *shared*-covariance Mahalanobis computation. +It does **not** bound the per-gene covariance loop, so it is a lever for pass 1 +and for peak prediction memory, not for the sample-variance tensor. + +``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? +------------------------- + +Cost is 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 1 000 genes is a recommendation chosen from the cost curve, not a +tuned or validated parameter.** It is the point at which the covariance term +is around 560 GiB at default landmarks, or 370 GiB of scratch, while still +covering the interesting tail of a Mahalanobis ranking. Past roughly 2 000 +genes the in-memory figure passes a terabyte, and the per-gene Cholesky loop +becomes the dominant runtime term because it is the only stage that grows with +the gene count while the rest of the run is essentially fixed. If you need +more coverage than that, reduce ``n_landmarks`` in the same breath rather than +raising the gene count alone. + +Price it rather than guessing. A :ref:`dry run ` gives the exact +figure for your dimensions in a second. + + +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..d1788e9 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -40,10 +40,21 @@ Differential Expression kompot.de( adata, "condition", "Young", "Old", sample_col="donor_id", - genes=top_genes, # e.g. top 200 from a previous run + genes=top_genes, # e.g. top 1 000 from a previous run 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, so the dominant + allocation becomes ``3 x n_landmarks^2 x n_genes x 8`` bytes: about + 0.56 GiB per gene at the default 5 000 landmarks, and roughly 560 GiB for + 1 000 genes. 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 +177,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 +189,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, what to compare, and the caveat about ``null_genes``: +:doc:`Planning Memory and Disk `. + Run Tracking ------------ diff --git a/examples/02_differential_expression_detailed.ipynb b/examples/02_differential_expression_detailed.ipynb index 38cbe92..0e70d3b 100644 --- a/examples/02_differential_expression_detailed.ipynb +++ b/examples/02_differential_expression_detailed.ipynb @@ -153,10 +153,18 @@ "\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", + " (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 side by side).\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 — about **0.56 GiB per gene** at the default 5 000 landmarks. It belongs\n", + "in a second pass over a 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 +1336,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", + "Note that the estimate reads `null_genes` as you pass it and does not resolve the\n", + "`\"auto\"` default, so the explicit `null_genes=2000` below is what makes the plan match\n", + "the run it is pricing.\n" ] }, { diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index ba48efb..278ed2e 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -15,7 +15,20 @@ "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.**\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", + "> so the dominant allocation becomes `3 \u00d7 n_landmarks\u00b2 \u00d7 n_genes \u00d7 8` bytes: roughly\n", + "> **0.56 GiB per gene** at the default `n_landmarks=5000`. One thousand genes is about\n", + "> 560 GiB; the whole transcriptome is terabytes. Compute grows too, because the\n", + "> Mahalanobis step then factorises once *per gene* instead of once in total.\n", + ">\n", + "> This notebook therefore does what you should always do: a **cheap first pass over all\n", + "> genes**, then a **sample-variance pass restricted to the top ones**. 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 +236,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,16 +322,25 @@ "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, so the run is priced\n", + "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 cost is *linear* in the gene count, so this is\n", + " the lever that decides whether the run is possible at all. Fifty genes keeps this\n", + " tutorial fast; on a real study, on the order of **1 000** is the working\n", + " recommendation;\n", + "2. **lower `n_landmarks`** if the plan is still too large \u2014 cost is *quadratic* in it, so\n", + " halving the landmark count quarters the covariance footprint;\n", + "3. **`store_arrays_on_disk=True`** (`StorageSettings`), which asks Kompot to consume the\n", + " per-gene tensors one gene at a time rather than as three dense in-memory arrays;\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" ] }, { @@ -1343,11 +1366,21 @@ "## 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, in order of how much each one buys:\n", + "\n", + "| Lever | Effect on the per-gene covariance tensor | Practical setting |\n", + "|---|---|---|\n", + "| `genes=` | **linear** \u2014 the whole point of pass 2 | top ~1 000 from pass 1 |\n", + "| `GPSettings(n_landmarks=\u2026)` | **quadratic** | 1 000\u20132 000 when 5 000 does not fit |\n", + "| `FDRSettings(null_genes=0)` | removes 2 000 genes' worth of tensor | already the default once `sample_col` is set |\n", + "| `StorageSettings(store_arrays_on_disk=True)` | consumed gene by gene | set `disk_storage_dir` to real scratch |\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 and tells you whether it fits. Full treatment, with measured plans:\n", + "[Planning Memory and Disk](https://kompot.readthedocs.io/en/latest/resource_planning.html).\n" ] }, { diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index 4315a0d..41b0dae 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -79,6 +79,33 @@ 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*, so the + dominant allocation becomes ``3 * n_landmarks**2 * n_genes * 8`` bytes: + roughly 0.56 GiB per gene at the default ``n_landmarks=5000``, about + 560 GiB for 1 000 genes, and terabytes for a whole transcriptome. + Compute grows too, because the Mahalanobis step then performs one + Cholesky factorisation per gene instead of one in total, and + ``GPSettings.batch_size`` does not bound that loop. + + 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(1000).index + + kompot.de(adata, "condition", "Young", "Old", # pass 2 + sample_col="donor_id", genes=top, + gp=kompot.GPSettings(n_landmarks=2000)) + + Cost is linear in ``genes`` and **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 @@ -94,7 +121,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.). fdr : FDRSettings, optional @@ -112,7 +141,14 @@ 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. Note that the estimate reads + ``FDRSettings.null_genes`` as given and does not resolve the + ``"auto"`` default, so a dry run without ``sample_col`` and with + ``null_genes="auto"`` omits the 2 000 null genes the real run would + add; pass ``null_genes=2000`` explicitly to see them + (https://github.com/settylab/kompot/issues/25). **function_kwargs Forwarded to :class:`~mellon.FunctionEstimator`. diff --git a/kompot/cli/de.py b/kompot/cli/de.py index 5e6f80c..1bf2cce 100644 --- a/kompot/cli/de.py +++ b/kompot/cli/de.py @@ -123,7 +123,13 @@ 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, ~0.56 GiB per gene " + "at --n-landmarks 5000. Run it as a second pass over a top-gene " + "list (config key 'genes') and check with --dry-run first." + ), ) parser.add_argument( diff --git a/kompot/cli/templates/de_config_template.yaml b/kompot/cli/templates/de_config_template.yaml index 80abd92..b27d353 100644 --- a/kompot/cli/templates/de_config_template.yaml +++ b/kompot/cli/templates/de_config_template.yaml @@ -23,6 +23,12 @@ 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 -- 3 * n_landmarks^2 * n_genes * 8 bytes, about 0.56 GiB +# per gene at n_landmarks=5000. Run it as a SECOND PASS over a restricted +# `genes` list (below), lower n_landmarks (cost 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 @@ -46,9 +52,12 @@ ls_factor: 10.0 # Length scale multiplication factor 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 @@ -57,7 +66,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 ~1000 genes of a first pass. min_cells: 2 # Minimum cells per gene min_percentage: null # Minimum percentage of cells expressing diff --git a/kompot/differential/sample_variance_estimator.py b/kompot/differential/sample_variance_estimator.py index b948699..31a9970 100644 --- a/kompot/differential/sample_variance_estimator.py +++ b/kompot/differential/sample_variance_estimator.py @@ -32,6 +32,18 @@ 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. This is the dominant + cost of any sample-variance run, and it 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 diff --git a/kompot/settings.py b/kompot/settings.py index 1f80ca1..a75cb8f 100644 --- a/kompot/settings.py +++ b/kompot/settings.py @@ -43,13 +43,21 @@ 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 ``3 * n_landmarks**2 * n_genes * 8`` bytes, so cost is + **quadratic** in this value: halving it quarters the covariance + footprint. 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 @@ -272,11 +280,27 @@ 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. + Represent the per-gene sample-variance covariance tensors so they are + consumed one gene at a time instead of as dense in-memory arrays. + With ``dask`` installed this 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``. Defaults to + ``None``, meaning *on if and only if* ``disk_storage_dir`` is set; + nothing enables it automatically in response to memory pressure. + Note that the dry run currently reports a memory saving from this flag + that the measured run does not deliver + (https://github.com/settylab/kompot/issues/26); 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 From 833c6bf5ab8887928aaef42423f5b2512457d70c Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 13:50:15 -0700 Subject: [PATCH 02/28] fix: assemble the sample-variance covariance one gene at a time (0.9.0) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three silent defects around the per-gene covariance tensor, plus the version bump that ships them. settylab/kompot#26 — store_arrays_on_disk did not bound memory ------------------------------------------------------------- With `sample_col` set, each condition produces a (n_landmarks, n_landmarks, n_genes) float64 covariance tensor. The code summed the two into a THIRD dense tensor and then added the shared posterior covariance into it gene by gene with __setitem__. With np.memmap inputs that sum materialises the whole thing in RAM, which is exactly the array store_arrays_on_disk exists to avoid; under Dask the per-gene __setitem__ rebuilt the graph once per gene. Separately, compute_mahalanobis_distances ran `cov = jnp.array(covariance)` in its gene-specific branch and never read `cov` — a second full-size copy for nothing. The sum is now assembled lazily by kompot.utils.LazyGeneCovariance, which keeps references to the terms and materialises one (n_landmarks, n_landmarks) matrix when the Mahalanobis step asks for a gene. A LazyGeneCovariance deliberately does not take the Dask delayed branch: that branch submits every gene at once and the view materialises its own slices inside each task, so the two nest and peak memory scales with concurrency rather than with one gene. Measured, 1 500 cells / 150 genes / 600 landmarks (412 MiB per tensor), peak ANONYMOUS memory from /proc/self/smaps_rollup, same machine and input, 19ee1a1 vs this tree: run before after written wall no sample variance 949 MiB 972 MiB 0 -> 0 33.0 -> 32.6 s sample variance, in memory 2 401 MiB 1 973 MiB 0 -> 0 45.7 -> 46.0 s store_arrays_on_disk, with dask 2 782 MiB 1 153 MiB 0 -> 0 231.8 -> 47.8 s store_arrays_on_disk, no dask 2 284 MiB 1 148 MiB 824 MiB 48.2 -> 47.7 s Disk-backed storage is now cheaper than in-memory storage (1 153 vs 1 973), where before it was dearer (2 782 vs 2 401), and the Dask path is 4.9x faster. Results are unchanged: distances differ only by floating-point summation order, max relative difference 1.7e-12, the same order as the pre-existing difference BETWEEN the in-memory, Dask and memmap paths on 19ee1a1 alone. settylab/kompot#25 — the dry run did not price the run it was pricing --------------------------------------------------------------------- de(dry_run=True) built its plan before null_genes="auto" was resolved, so the estimator received the literal string, matched neither its int nor its list branch, and counted zero null genes. Measured on 4 000 x 500: 0.485 GiB planned against 2.156 GiB for the run it described, 4.44x optimistic. Resolution moved above the dry-run branch, and the estimator now raises on a null_genes it cannot interpret instead of falling through to zero. Estimator accounting -------------------- The plan no longer charges a third tensor that is never built, charges disk only on the path that actually writes (nothing reaches disk_storage_dir when dask is installed), and charges the per-gene working set that replaced the dense sum. Verified against the measurement: on the no-dask path the plan's disk figure and the bytes observed under disk_storage_dir now agree exactly (824 MiB = 824 MiB); on the dask path the plan reports 0 and 0 bytes were observed. Tests ----- tests/test_sample_variance_cost_regressions.py, 11 cases. Every expected value is re-derived from the allocation (`2 * L**2 * n_genes * 8`, `3 * L**2 * 8`), not copied from a run. Verified to FAIL on 19ee1a1: the five that can run there fail, and the LazyGeneCovariance cases cannot even import. One existing test changed and it is the one predicted: test_resource_estimation.py::test_dry_run_with_disk_storage asserted a bare `len(disk_reqs) > 0`, which is false once disk is charged only on the path that writes. It is re-derived rather than loosened — see that commit. Version 0.8.0 -> 0.9.0 in kompot/version.py and pyproject.toml, with a changelog entry per issue. Fixes #25 Fixes #26 --- CHANGELOG.md | 127 ++++++-- kompot/anndata/differential_expression.py | 24 +- .../differential/differential_expression.py | 150 ++++------ kompot/resource_estimation.py | 96 ++++-- kompot/utils.py | 121 +++++++- kompot/version.py | 2 +- pyproject.toml | 2 +- .../test_sample_variance_cost_regressions.py | 279 ++++++++++++++++++ 8 files changed, 633 insertions(+), 168 deletions(-) create mode 100644 tests/test_sample_variance_cost_regressions.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 80a7cbc..6f8f870 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,37 +2,112 @@ 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), peak +**anonymous** memory from `/proc/self/smaps_rollup`, same machine and same +synthetic input on both trees: + +| run | 0.8.0 | 0.9.0 | bytes written, 0.8.0 → 0.9.0 | 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 148 MiB | 824 → 824 MiB | 48.2 s → 47.7 s | + +Disk-backed storage is now cheaper than in-memory storage, which it was not +before: 1 153 MiB against 1 973 MiB, where 0.8.0 made it *more* expensive +(2 782 against 2 401). The Dask path is also 4.9x faster. + +**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. ### 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. Supplying - `sample_col` enables sample variance, which replaces the single shared posterior - covariance with one `(n_landmarks, n_landmarks)` matrix **per gene**, so the dominant - allocation is `3 x n_landmarks^2 x n_genes x 8` bytes: about 0.56 GiB per gene at the - default `n_landmarks=5000`, ~560 GiB for 1 000 genes, and terabytes for a whole - transcriptome. Compute scales too, since the Mahalanobis step then factorises once per - gene instead of once in total. The guide prescribes the two-pass workflow (a cheap - first pass over all genes, then sample variance restricted to the top ~1 000), ranks - the levers (`genes` linear, `n_landmarks` quadratic, `null_genes=0`, disk offload), - and shows how to price a run with `dry_run=True`. It carries measured dry-run plans - at realistic sizes. - - The same warning now appears where the decision is made: `kompot.de`'s docstring (so - `help()` carries it), `GPSettings.n_landmarks`, `StorageSettings`, - `SampleVarianceEstimator`, `kompot de --sample-col`, the DE config template, the - README, and the tutorial notebooks. + 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 genes — ranks the levers, and shows how to price a run + before committing to it. + - 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 via a config file. - - Corrected `StorageSettings.max_memory_ratio`, which was documented as "Fraction of RAM - before triggering disk storage". It triggers nothing: `store_arrays_on_disk=None` - 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. - - Documented two measured divergences between the dry run and the runs it prices: - the estimate does not resolve `null_genes="auto"` (#25), and `store_arrays_on_disk=True` - does not reduce peak memory, with nothing written to disk when `dask` is installed (#26). + `--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. ## [0.8.0] - 2026-07-28 diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index 41b0dae..2c8f07b 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -252,6 +252,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 @@ -295,16 +309,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/differential/differential_expression.py b/kompot/differential/differential_expression.py index 6af2bee..85815a5 100644 --- a/kompot/differential/differential_expression.py +++ b/kompot/differential/differential_expression.py @@ -8,7 +8,7 @@ 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 .expression_model import ExpressionModel @@ -635,121 +635,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/resource_estimation.py b/kompot/resource_estimation.py index e955010..b72e479 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,11 @@ 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 " + "and computed sequentially. Install dask to evaluate them " + "lazily instead, which writes nothing and parallelises the " + "per-gene work: pip install 'dask[array]'" ) # Add batch_size information diff --git a/kompot/utils.py b/kompot/utils.py index a870224..9fa5951 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,17 @@ 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. + # + # 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 +287,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 +362,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 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_sample_variance_cost_regressions.py b/tests/test_sample_variance_cost_regressions.py new file mode 100644 index 0000000..656d38b --- /dev/null +++ b/tests/test_sample_variance_cost_regressions.py @@ -0,0 +1,279 @@ +"""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 From f7d7476caf6ed0e40a7a1fb3e9306b8f1d9d0bfb Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 14:11:31 -0700 Subject: [PATCH 03/28] docs: rewrite the sample-variance cost guidance from the post-fix measurement MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The docs commit on this branch described 0.8.0's behaviour: 0.56 GiB per gene, and store_arrays_on_disk carrying a warning that it did not deliver the memory saving it advertised. The fix commit changed both facts, so this rewrites the prose from the new measurement rather than patching the old sentences — carrying them forward half-edited would leave a page that is pessimistic in one paragraph and confident in the next about the same flag. What changed in the guidance ---------------------------- The headline is now 2 * n_landmarks^2 * n_genes * 8 bytes, about 0.37 GiB per gene at the default 5 000 landmarks, and disk offload is presented as a genuine remedy with no hedge. More substantively, the ARGUMENT for the two-pass workflow has moved. Memory used to be the reason to restrict genes on pass 2; store_arrays_on_disk now answers that (1 000 genes plans at 7.7 GiB, the whole transcriptome at 101 GiB). So the prescription now rests on the per-gene Cholesky factorisation, which is the one cost no storage mode touches and only a shorter gene list reduces. Measured: 0.52 s per gene at 500 landmarks rising to 8.5 s at 2 000, against a shared-covariance column flat in the gene count. Multiplied out at 2 000 landmarks: ~2.4 h for 1 000 genes, ~47 h for 20 000. The prescription survives and is better founded; a whole-transcriptome sample-variance run is no longer impossible, merely interminable, which is the worse failure because it looks like progress. `genes` is now visibly the only lever that appears in BOTH cost columns. The memmap cost is documented where the path is chosen ------------------------------------------------------ Without dask the tensors are memory-mapped, and a gene slice of a C-contiguous (n_points, n_points, n_genes) map is strided, so per-gene reads touch the whole file where the old code read it once sequentially: about 12% slower on a measured 900-landmark pair, 74 s against 66 s, for 57% less memory. That sentence sits in the bullet describing the no-dask path, so a reader meets it before waiting on it rather than after. Also in this commit ------------------- tests/test_resource_estimation.py::test_dry_run_with_disk_storage re-derived. It asserted a bare `len(disk_reqs) > 0`, which passed for the wrong reason: the estimator charged disk unconditionally, including for the dask path that writes nothing. The replacement asserts an exact byte count taken from the allocation (2 * L^2 * n_genes * 8) on the no-dask path, zero on the dask path, and that the tensors are not charged to memory in either — strictly stronger than what it replaced, not loosened. One published snippet was wrong and is corrected: the with/without comparison left null_genes at "auto", which resolves to 2 000 on the no-sample-variance arm and 0 on the other, so it was silently comparing different gene counts. It now pins null_genes=0 on both. Caught by executing it, which is also how the whole snippet set is checked. Verification ------------ sphinx-build exits 0 with 24 warnings, an identical set to origin/main. Every published snippet executed against a synthetic AnnData. Rendered HTML scanned across every page for raw reStructuredText leakage: zero. --- CHANGELOG.md | 14 +- README.md | 15 +- docs/source/cli.rst | 14 +- docs/source/index.rst | 7 +- docs/source/resource_planning.rst | 501 ++++++++++-------- docs/source/simplified.rst | 15 +- .../02_differential_expression_detailed.ipynb | 9 +- examples/03_sample_variance.ipynb | 61 ++- kompot/anndata/differential_expression.py | 22 +- kompot/cli/de.py | 8 +- kompot/cli/templates/de_config_template.yaml | 10 +- .../differential/sample_variance_estimator.py | 9 +- kompot/settings.py | 23 +- tests/test_resource_estimation.py | 69 ++- 14 files changed, 464 insertions(+), 313 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 6f8f870..3e7e16b 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -38,7 +38,19 @@ synthetic input on both trees: Disk-backed storage is now cheaper than in-memory storage, which it was not before: 1 153 MiB against 1 973 MiB, where 0.8.0 made it *more* expensive -(2 782 against 2 401). The Dask path is also 4.9x faster. +(2 782 against 2 401). Repeated at 900 landmarks / 120 genes: 2 809 MiB in +memory, 1 325 MiB with Dask and 1 341 MiB without, against 3 502 / 4 341 / +3 079 MiB on 0.8.0. + +One cost moved the other way, on the path that is not recommended. Without +`dask` the tensors are memory-mapped, and a gene slice of a C-contiguous +`(n_points, n_points, n_genes)` 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 on a clean 900-landmark pair, 74 s against 66 s — about 12% slower in +exchange for 57% less memory. Installing `dask` avoids it in both directions: +that path is 4.9x *faster* than 0.8.0 as well as lighter. 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, diff --git a/README.md b/README.md index fa58259..2bcc7dd 100644 --- a/README.md +++ b/README.md @@ -71,13 +71,16 @@ kompot.de( Passing `sample_col` turns on **sample variance**, which replaces the single shared posterior covariance with **one `(n_landmarks, n_landmarks)` covariance -matrix per gene**. The dominant allocation becomes -`3 x n_landmarks^2 x n_genes x 8` bytes: at the default 5 000 landmarks that is -roughly **0.56 GiB per gene**, so a whole transcriptome asks for terabytes. -Compute scales too, since the Mahalanobis step then factorises once per gene -instead of once in total. +matrix per gene, per condition**. Two costs follow, and they respond to +different levers: -Run it in two passes, and price the second one first: +- **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. + Nothing makes this cheaper except analysing fewer genes. + +So run it in two passes, and price the second one first: ```python # Pass 1 — all genes, no sample variance diff --git a/docs/source/cli.rst b/docs/source/cli.rst index 37e9feb..cd653d1 100644 --- a/docs/source/cli.rst +++ b/docs/source/cli.rst @@ -252,9 +252,9 @@ Example: Complete Analysis ``--sample-col`` is deliberately absent above. It turns on sample variance, which gives every gene its own ``(n_landmarks, n_landmarks)`` covariance - matrix: about 0.56 GiB per gene at ``--n-landmarks 5000``, so a whole - transcriptome asks for terabytes. Run it as a **second pass** over a - restricted gene list, which the CLI takes through the ``genes:`` key of a + matrix and its own Cholesky factorisation: about 0.37 GiB and several + seconds 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 @@ -733,9 +733,11 @@ For large datasets: # disk_storage_dir: "/scratch/kompot" # real scratch, not tmpfs With ``--sample-col`` the dominant allocation is -``3 x n_landmarks^2 x n_genes x 8`` bytes, so the levers that matter are the -**gene list** (linear) and ``--n-landmarks`` (quadratic). ``--batch-size`` -does not bound that term. Full treatment: +``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 nothing bounds the +per-gene factorisation except a shorter gene list. Full treatment: :doc:`Planning Memory and Disk `. Speed Optimization diff --git a/docs/source/index.rst b/docs/source/index.rst index 187679f..b12ffbb 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -137,9 +137,10 @@ 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. At default settings that costs roughly - **0.56 GiB per gene**, so it must be run as a second pass over a restricted - gene list, never over the whole transcriptome. Read + own landmark covariance matrix and its own Cholesky factorisation. At + default settings that is roughly **0.37 GiB and several seconds per gene**, + so it belongs in a second pass over a restricted gene list, never in a + whole-transcriptome run. Read :doc:`Planning Memory and Disk ` before your first sample-variance run. diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index f9d3ca0..004d8bc 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -4,11 +4,10 @@ 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**. The cost of that step is multiplied by the number of genes -analysed, and at Kompot's default landmark count it reaches hundreds of -gigabytes for a routine gene set. +per gene**. Both the memory and the compute of that step are multiplied by the +number of genes analysed. -This page explains where the cost comes from, prescribes the two-pass workflow +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. @@ -22,9 +21,10 @@ 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 1 000; see :ref:`how-many-genes`) - and cut ``n_landmarks`` if the plan is still too large. +2. **Pass 2, top genes only, with sample variance.** Restrict to the genes that + survived pass 1 (on the order of 1 000; 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 @@ -37,23 +37,28 @@ instead: mahal = "kompot_de_Young_to_Old_mahalanobis" top_genes = adata.var.sort_values(mahal, ascending=False).head(1000).index - # ---- Pass 2: sample variance, restricted --------------------------- + # ---- Pass 2: sample variance, restricted and offloaded ------------- kompot.de( adata, "condition", "Young", "Old", sample_col="donor_id", - genes=top_genes, # the lever that matters most - gp=kompot.GPSettings(n_landmarks=2000), # quadratic; see "Other levers" + 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, # read "Disk offload" below first + 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 levers in that call -that genuinely shrink the allocation are ``genes`` and ``n_landmarks``; -``store_arrays_on_disk`` is discussed, with measurements, under -:ref:`disk-offload`. +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 reason the split is not optional is that the two costs respond to different +levers. ``store_arrays_on_disk`` removes the memory term almost entirely +(:ref:`disk-offload`), and ``n_landmarks`` shrinks it quadratically — but +neither touches the **per-gene Cholesky factorisation**, which is linear in the +gene count and nothing makes it cheaper except analysing fewer genes. Restrict +the gene list and both costs fall together; offload without restricting and you +have merely moved from running out of memory to running out of time. 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: @@ -90,22 +95,20 @@ 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 +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 Kompot holds one per condition, plus their sum. So the dominant term is +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 = k x n_landmarks^2 x n_genes x 8 (float64) - - k = 3 variance1 + variance2 + their sum, all dense in memory - k = 2 what the planner charges to disk when store_arrays_on_disk=True, - on the assumption that the sum stays lazy + 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 @@ -114,72 +117,152 @@ 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 a sample-variance - run costs about 0.56 GiB**. One thousand genes is roughly 560 GiB; ten - thousand is over five terabytes. Nothing else in a Kompot run behaves this - way, which is why the gene list on pass 2 is the number to decide - deliberately. + At the default 5 000 landmarks, **every gene added to an in-memory + sample-variance run costs about 0.37 GiB**. One thousand genes is roughly + 370 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:: -Compute scales with genes as well, and 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. With per-gene -covariances it factorises **once per gene**, in a Python loop over genes + 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. +loop. Timed on one core-count of this machine, per gene: + +.. list-table:: + :header-rows: 1 + :widths: 22 26 26 26 + + * - ``n_landmarks`` + - shared, 20 genes + - per-gene, 20 genes + - per-gene, per gene + * - 500 + - 0.61 s + - 10.5 s + - 0.52 s + * - 1 000 + - 8.7 s + - 121.6 s + - 6.1 s + * - 1 500 + - 11.8 s + - 156.7 s + - 7.8 s + * - 2 000 + - 10.8 s + - 170.9 s + - 8.5 s + +Those are wall-clock seconds on one shared machine and will not transfer +exactly, but the **shape** does: the shared column is flat in the gene count +while the per-gene column is linear in it. Multiply it out, because a per-gene +figure is not something you can act on. At 2 000 landmarks and 8.5 s per gene: + +.. list-table:: + :header-rows: 1 + :widths: 30 35 35 + + * - genes in pass 2 + - factorisation time + - in practice + * - 1 000 + - ~2.4 hours + - an overnight job + * - 2 000 + - ~4.7 hours + - still an overnight job + * - 20 000 (all) + - ~47 hours + - two days of compute + +That, rather than memory, is what bounds a sample-variance run once the +tensors are offloaded, and it is the reason +:ref:`the two-pass workflow ` remains the prescription. A +whole-transcriptome sample-variance run is no longer impossible; it is merely +interminable, which is a worse failure because it looks like progress. 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, 1 000 genes and 6 donors this is 0.45 GiB, negligible - beside the covariance tensors but not beside everything else. + At 5 000 landmarks, 1 000 genes and 6 donors this is 0.45 GiB. * **Two extra layers**, ``__std``, each ``n_cells x n_genes x 8`` B, written into ``adata.layers``. - Measured plans -------------- -The figures below are 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``. They are totals for the whole plan, not just the -covariance term. (Kompot's report labels these binary magnitudes ``GB``; they -are the same numbers.) +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: 14 20 20 32 + :widths: 12 18 20 22 18 * - genes - no sample variance - - sample variance, RAM - - sample variance, disk-backed + - sample variance, in memory + - sample variance, ``store_arrays_on_disk`` + - disk used * - 200 - 3.1 GiB - - 115.0 GiB - - 3.2 GiB RAM + 74.5 GiB disk + - 78.3 GiB + - 3.8 GiB + - 0 * - 1 000 - 6.4 GiB - - 566.0 GiB - - 7.2 GiB RAM + 372.5 GiB disk + - 380.2 GiB + - 7.7 GiB + - 0 * - 2 000 - 10.6 GiB - - 1 130 GiB - - 12.1 GiB RAM + 745.1 GiB disk + - 757.7 GiB + - 12.6 GiB + - 0 * - 20 000 (all) - 85.7 GiB - - 11 276 GiB - - 100.6 GiB RAM + 7 451 GiB disk - -The last row is the configuration to avoid: sample variance over a full -transcriptome at default settings asks for eleven terabytes. The -no-sample-variance column grows only with ``n_cells x n_genes``, which is why -pass 1 over all 20 000 genes is comfortable at 86 GiB. - -The fourth column is what the *planner* charges when -``store_arrays_on_disk=True``. Read it as the scratch footprint the run is -allowed to reach, not as a memory saving you can count on; see the warning -under :ref:`disk-offload`. + - 7 552 GiB + - 101.2 GiB + - 0 + +Three things to read off this. + +**The in-memory column is the one that explodes.** 1 000 genes asks for +380 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 1 000 genes plans at 7.7 GiB, barely +above the 6.4 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. + +**Memory stops being the binding constraint, and compute takes over.** The +20 000-gene disk-backed plan at 101 GiB would fit on a large node; the 20 000 +per-gene Cholesky factorisations behind it would not fit in your week. That is +why :ref:`the two-pass workflow ` is the prescription even +now that the memory problem is solved. Reproduce any row on your own data by substituting your ``adata``: @@ -191,6 +274,7 @@ Reproduce any row on your own data by substituting your ``adata``: 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) @@ -234,8 +318,12 @@ 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, dry_run=True) + 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) @@ -284,12 +372,10 @@ stderr, and exits non-zero when the plan is infeasible. See :doc:`cli`. Disk offload ------------ -``StorageSettings(store_arrays_on_disk=True)`` is intended to change how the -per-gene covariance tensors are represented, so they are consumed one gene at -a time instead of built as three dense in-memory arrays. It is the flag the -dry run's warnings point you at, and the one the estimator prices into its -disk column. **Read the warning at the end of this section before you rely on -it**: measured runs do not show the memory saving. +``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 @@ -305,157 +391,138 @@ it**: measured runs do not show the memory saving. What it does mechanically depends on whether ``dask`` is installed: -* **With** ``dask`` (``pip install 'kompot[dask]'``) the tensor is a lazy - Dask graph, evaluated one gene at a time as the Mahalanobis step reaches it. - Measured: **no file is written at all**, so the dry run's disk figure is an - upper bound on scratch rather than a prediction of it. -* **Without** ``dask`` Kompot writes the full tensor for each condition as a - memory-mapped ``covariance_matrix.npy`` under ``disk_storage_dir``, and the - dry run's disk figure is then what actually lands there. This path is - sequential rather than parallel, and the dry run warns that ``dask`` is - missing. - -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** +* **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), peak **anonymous** memory sampled from ``/proc/self/smaps_rollup`` so +that page cache for a memory-mapped file cannot flatter the reading: + +.. list-table:: + :header-rows: 1 + :widths: 30 16 16 18 20 + + * - run + - peak memory + - bytes written + - wall clock + - vs. no sample variance + * - no sample variance + - 972 MiB + - 0 + - 33 s + - — + * - sample variance, in memory + - 1 973 MiB + - 0 + - 46 s + - +1 001 MiB + * - ``store_arrays_on_disk``, with ``dask`` + - 1 153 MiB + - **0** + - 48 s + - +181 MiB + * - ``store_arrays_on_disk``, no ``dask`` + - 1 148 MiB + - 824 MiB + - 48 s + - +176 MiB + +The offloaded runs cost about a sixth of the extra memory that the in-memory +run does, and the gap widens with the gene count, because the in-memory column +carries a term linear in genes and the offloaded ones do not. Repeated at 900 +landmarks and 120 genes (742 MiB per tensor), peak memory was 2 809 MiB +in memory, 1 325 MiB with ``dask`` and 1 341 MiB without, against 1 051 MiB +for a run with no sample variance at all. + +The wall-clock column is from a shared machine under load and should be read +as an order of magnitude, not a benchmark. One caveat it does show reliably: +the ``dask``-less path is **modestly slower than it was before 0.9.0** — 74 s +against 66 s on a clean 900-landmark pair — because 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. That is the price of not holding the tensor, it is +paid only when ``dask`` is absent, and installing ``dask`` avoids it in both +directions. + +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 writes into 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. +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 -warnings escalate; it does not switch storage modes. +memory pressure. ``max_memory_ratio`` only sets the threshold at which the +plan's warnings escalate; it does not switch storage modes. -.. warning:: +.. note:: - **Measured runs do not show the memory saving this flag promises.** - ``store_arrays_on_disk=True`` moves the tensor from the memory column of - the plan to the disk column, but peak memory of the actual ``kompot.de()`` - run is unchanged. Two synthetic configurations, 1 500 cells and 4 donors, - with peak *anonymous* memory sampled from ``/proc/self/smaps_rollup`` so - that page cache for a memory-mapped file is excluded: - - .. list-table:: 150 genes, 600 landmarks (tensor 412 MiB per condition) - :header-rows: 1 - :widths: 34 14 14 16 14 12 - - * - run - - plan RAM - - plan disk - - peak anon - - on disk - - wall - * - no sample variance - - 69 MiB - - 0 - - 962 MiB - - 0 - - 49 s - * - sample variance, in memory - - 1 313 MiB - - 0 - - 2 357 MiB - - 0 - - 87 s - * - on disk, with ``dask`` - - 77 MiB - - 824 MiB - - 2 781 MiB - - **0** - - 283 s - * - on disk, no ``dask`` - - 77 MiB - - 824 MiB - - 2 346 MiB - - 824 MiB - - 50 s - - .. list-table:: 120 genes, 900 landmarks (tensor 742 MiB per condition) - :header-rows: 1 - :widths: 34 14 14 16 14 12 - - * - run - - plan RAM - - plan disk - - peak anon - - on disk - - wall - * - no sample variance - - 73 MiB - - 0 - - 1 052 MiB - - 0 - - 55 s - * - sample variance, in memory - - 2 307 MiB - - 0 - - 3 505 MiB - - 0 - - 67 s - * - on disk, with ``dask`` - - 82 MiB - - 1 483 MiB - - 4 347 MiB - - **0** - - 259 s - * - on disk, no ``dask`` - - 82 MiB - - 1 483 MiB - - 3 562 MiB - - 1 483 MiB - - 113 s - - Two things to read off these. The peak-memory column is flat across the - three sample-variance rows at both scales, against a plan predicting a 17x - to 28x drop. And with ``dask`` installed **nothing reaches disk at all**: - the tensor is a lazy graph rather than a file, so the plan's disk figure - bounds scratch consumption rather than predicting it, and that path was - also the slowest of the four. Tracked in `settylab/kompot#26 - `_. - - Until that is resolved, treat **restricting genes** and **reducing - landmarks** as the levers that actually bound memory. They shrink the - allocation itself, and the dry run prices them correctly. + 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 ------------ -``n_landmarks`` **is the strongest lever**, because the tensor is quadratic in -it. Same 1 000 genes, sample variance on disk, from the dry run: +``n_landmarks`` **is the strongest memory lever**, because the covariance term +is quadratic in it. Same 1 000 genes, from the dry run: .. list-table:: :header-rows: 1 - :widths: 20 20 20 + :widths: 24 26 26 * - ``n_landmarks`` - - RAM - - disk + - sample variance, in memory + - with ``store_arrays_on_disk`` * - 1 000 + - 21.0 GiB - 6.1 GiB - - 14.9 GiB * - 2 000 - - 6.3 GiB - - 59.6 GiB + - 66.0 GiB + - 6.4 GiB * - 3 000 - - 6.5 GiB - - 134.1 GiB + - 140.8 GiB + - 6.7 GiB * - 5 000 (default) - - 7.2 GiB - - 372.5 GiB + - 380.2 GiB + - 7.7 GiB -Halving the landmark count quarters the covariance cost. Landmarks control -the resolution of the cell-state approximation, so this is a genuine accuracy -trade-off rather than a free saving, but 5 000 landmarks is rarely required to -resolve a covariance structure that is being summarised gene by gene. +Halving the landmark count quarters the covariance footprint. It also cuts the +per-gene factorisation, which is roughly cubic in ``n_landmarks``: measured +above, one gene costs 0.52 s at 500 landmarks and 8.5 s at 2 000. Landmarks +control the resolution of the cell-state approximation, so this is a genuine +accuracy trade-off rather than a free saving, but 5 000 landmarks 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 during the *shared*-covariance Mahalanobis computation. -It does **not** bound the per-gene covariance loop, so it is a lever for pass 1 -and for peak prediction memory, not for the sample-variance tensor. +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. @@ -465,23 +532,27 @@ next section. How many genes in pass 2? ------------------------- -Cost is 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 1 000 genes is a recommendation chosen from the cost curve, not a -tuned or validated parameter.** It is the point at which the covariance term -is around 560 GiB at default landmarks, or 370 GiB of scratch, while still -covering the interesting tail of a Mahalanobis ranking. Past roughly 2 000 -genes the in-memory figure passes a terabyte, and the per-gene Cholesky loop -becomes the dominant runtime term because it is the only stage that grows with -the gene count while the rest of the run is essentially fixed. If you need -more coverage than that, reduce ``n_landmarks`` in the same breath rather than -raising the gene count alone. - -Price it rather than guessing. A :ref:`dry run ` gives the exact -figure for your dimensions in a second. - +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 1 000 genes is a recommendation chosen from the cost curves, not a +tuned or validated parameter.** With ``store_arrays_on_disk=True`` memory is no +longer what decides it — 1 000 genes plans at 7.7 GiB, and even the whole +transcriptome plans at 101 GiB, which would fit on a large node. What decides +it is the per-gene Cholesky: at 2 000 landmarks and roughly 8.5 s per gene, +1 000 genes is about **2.4 hours** and 20 000 genes about **47 hours**. One +thousand genes covers the interesting tail of a Mahalanobis ranking at a cost +you can absorb overnight; the whole transcriptome costs two days for a result +whose last nine-tenths you were never going to read. + +If you need more coverage than that, cut ``n_landmarks`` in the same breath: +it reduces the per-gene cost as well as the memory, so it buys back the time +the extra genes spend. + +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. See also -------- diff --git a/docs/source/simplified.rst b/docs/source/simplified.rst index d1788e9..cc07c76 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -47,12 +47,15 @@ Differential Expression .. warning:: ``sample_col`` is the one expensive option. It gives every gene **its own** - ``(n_landmarks, n_landmarks)`` covariance matrix, so the dominant - allocation becomes ``3 x n_landmarks^2 x n_genes x 8`` bytes: about - 0.56 GiB per gene at the default 5 000 landmarks, and roughly 560 GiB for - 1 000 genes. 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 + ``(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; + nothing removes the per-gene factorisation except analysing fewer genes. + 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 diff --git a/examples/02_differential_expression_detailed.ipynb b/examples/02_differential_expression_detailed.ipynb index 0e70d3b..9e25f31 100644 --- a/examples/02_differential_expression_detailed.ipynb +++ b/examples/02_differential_expression_detailed.ipynb @@ -157,12 +157,15 @@ " `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).\n", + "- **`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 — about **0.56 GiB per gene** at the default 5 000 landmarks. It belongs\n", - "in a second pass over a restricted gene list, never in a whole-transcriptome run. See\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" ] diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index 278ed2e..8182766 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -17,16 +17,19 @@ "\n", "This dataset has **3 replicates per age group**.\n", "\n", - "> **Cost warning \u2014 sample variance is priced per gene.**\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", - "> so the dominant allocation becomes `3 \u00d7 n_landmarks\u00b2 \u00d7 n_genes \u00d7 8` bytes: roughly\n", - "> **0.56 GiB per gene** at the default `n_landmarks=5000`. One thousand genes is about\n", - "> 560 GiB; the whole transcriptome is terabytes. Compute grows too, because the\n", - "> Mahalanobis step then factorises once *per gene* instead of once in total.\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", - "> This notebook therefore does what you should always do: a **cheap first pass over all\n", - "> genes**, then a **sample-variance pass restricted to the top ones**. The levers and\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", + "> Nothing makes this cheaper except 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" ] @@ -324,17 +327,18 @@ "source": [ "## Pass 2 \u2014 DE With Sample Variance, Restricted to the Top Genes\n", "\n", - "Add `sample_col`. Kompot now builds one covariance matrix per gene, so the run is priced\n", - "per gene and we pull every lever available:\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 pass 1 \u2014 cost is *linear* in the gene count, so this is\n", - " the lever that decides whether the run is possible at all. Fifty genes keeps this\n", - " tutorial fast; on a real study, on the order of **1 000** is the working\n", - " recommendation;\n", - "2. **lower `n_landmarks`** if the plan is still too large \u2014 cost is *quadratic* in it, so\n", - " halving the landmark count quarters the covariance footprint;\n", - "3. **`store_arrays_on_disk=True`** (`StorageSettings`), which asks Kompot to consume the\n", - " per-gene tensors one gene at a time rather than as three dense in-memory arrays;\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 **1 000** 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", + " the per-gene factorisation roughly *cubic*, 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", @@ -1369,17 +1373,22 @@ "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, in order of how much each one buys:\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 ~1 000 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** | ~**cubic** | 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", - "| Lever | Effect on the per-gene covariance tensor | Practical setting |\n", - "|---|---|---|\n", - "| `genes=` | **linear** \u2014 the whole point of pass 2 | top ~1 000 from pass 1 |\n", - "| `GPSettings(n_landmarks=\u2026)` | **quadratic** | 1 000\u20132 000 when 5 000 does not fit |\n", - "| `FDRSettings(null_genes=0)` | removes 2 000 genes' worth of tensor | already the default once `sample_col` is set |\n", - "| `StorageSettings(store_arrays_on_disk=True)` | consumed gene by gene | set `disk_storage_dir` to real scratch |\n", + "Note what the table says about `genes`: it is the only row that appears in both cost\n", + "columns. Offloading to disk solves memory and leaves the factorisation untouched, so a\n", + "whole-transcriptome sample-variance run stops being impossible and becomes merely\n", + "interminable. Restrict the gene list.\n", "\n", "And price it before you run it: `dry_run=True` on the exact call you intend to make takes\n", - "a second and tells you whether it fits. Full treatment, with measured plans:\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/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index 2c8f07b..ba9f73d 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -83,13 +83,16 @@ def de( 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*, so the - dominant allocation becomes ``3 * n_landmarks**2 * n_genes * 8`` bytes: - roughly 0.56 GiB per gene at the default ``n_landmarks=5000``, about - 560 GiB for 1 000 genes, and terabytes for a whole transcriptome. - Compute grows too, because the Mahalanobis step then performs one - Cholesky factorisation per gene instead of one in total, and - ``GPSettings.batch_size`` does not bound that loop. + 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 + nothing makes it cheaper except analysing fewer genes. Run it as a **second pass** over a restricted gene list:: @@ -102,8 +105,9 @@ def de( sample_col="donor_id", genes=top, gp=kompot.GPSettings(n_landmarks=2000)) - Cost is linear in ``genes`` and **quadratic** in ``n_landmarks``. - Price any configuration first with ``dry_run=True``. Full treatment: + 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 diff --git a/kompot/cli/de.py b/kompot/cli/de.py index 1bf2cce..8edb713 100644 --- a/kompot/cli/de.py +++ b/kompot/cli/de.py @@ -126,9 +126,11 @@ def add_de_parser(subparsers) -> argparse.ArgumentParser: help=( "Column in adata.obs with sample labels for sample variance " "estimation. EXPENSIVE: gives every gene its own " - "(n_landmarks, n_landmarks) covariance matrix, ~0.56 GiB per gene " - "at --n-landmarks 5000. Run it as a second pass over a top-gene " - "list (config key 'genes') and check with --dry-run first." + "(n_landmarks, n_landmarks) covariance matrix and its own Cholesky " + "factorisation -- ~0.37 GiB and several seconds 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." ), ) diff --git a/kompot/cli/templates/de_config_template.yaml b/kompot/cli/templates/de_config_template.yaml index b27d353..e11e994 100644 --- a/kompot/cli/templates/de_config_template.yaml +++ b/kompot/cli/templates/de_config_template.yaml @@ -24,10 +24,12 @@ 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 -- 3 * n_landmarks^2 * n_genes * 8 bytes, about 0.56 GiB -# per gene at n_landmarks=5000. Run it as a SECOND PASS over a restricted -# `genes` list (below), lower n_landmarks (cost is quadratic in it), and check -# with --dry-run first. +# 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 diff --git a/kompot/differential/sample_variance_estimator.py b/kompot/differential/sample_variance_estimator.py index 31a9970..ac61924 100644 --- a/kompot/differential/sample_variance_estimator.py +++ b/kompot/differential/sample_variance_estimator.py @@ -39,9 +39,12 @@ class SampleVarianceEstimator: 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. This is the dominant - cost of any sample-variance run, and it is why sample variance is - prescribed as a second pass over a restricted gene list. See + 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 diff --git a/kompot/settings.py b/kompot/settings.py index a75cb8f..c2c3903 100644 --- a/kompot/settings.py +++ b/kompot/settings.py @@ -45,9 +45,10 @@ class GPSettings: n_landmarks : int, optional Number of landmarks for the Nystrom approximation. Under sample variance (``sample_col`` set on :func:`kompot.de`) the dominant - allocation is ``3 * n_landmarks**2 * n_genes * 8`` bytes, so cost is + allocation is ``2 * n_landmarks**2 * n_genes * 8`` bytes, so memory is **quadratic** in this value: halving it quarters the covariance - footprint. See + footprint. It also cuts the per-gene Cholesky factorisation, which is + roughly cubic in it. See https://kompot.readthedocs.io/en/latest/resource_planning.html landmarks : np.ndarray, optional Pre-computed landmark coordinates. @@ -280,16 +281,14 @@ class StorageSettings: store_additional_stats : bool Store extra columns (p-values, tail FDR, PTP, z-scores). store_arrays_on_disk : bool, optional - Represent the per-gene sample-variance covariance tensors so they are - consumed one gene at a time instead of as dense in-memory arrays. - With ``dask`` installed this 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``. Defaults to - ``None``, meaning *on if and only if* ``disk_storage_dir`` is set; - nothing enables it automatically in response to memory pressure. - Note that the dry run currently reports a memory saving from this flag - that the measured run does not deliver - (https://github.com/settylab/kompot/issues/26); see + 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. A real run creates it if missing, 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(): From 8d6174ee67c2d0e1fcf31ba99e13ae48d88bacd5 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 14:15:48 -0700 Subject: [PATCH 04/28] test: cover the single-predictor and shared-term branches of the rewrite The #26 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 the suite did not exercise them, so a restructuring bug there would have been invisible. Three cases: one predictor and two both give the dense arithmetic's answer bitwise (maxdiff 0.000e+00 on the probe these were lifted from), and a 2-D shared term handed to LazyGeneCovariance is refused rather than broadcast silently -- a shared sample variance is folded into the base by the caller, so its arrival as a term is a caller bug the view should be loud about. --- .../test_sample_variance_cost_regressions.py | 56 +++++++++++++++++++ 1 file changed, 56 insertions(+) diff --git a/tests/test_sample_variance_cost_regressions.py b/tests/test_sample_variance_cost_regressions.py index 656d38b..99c6315 100644 --- a/tests/test_sample_variance_cost_regressions.py +++ b/tests/test_sample_variance_cost_regressions.py @@ -277,3 +277,59 @@ def test_plan_charges_the_per_gene_working_set(): 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]) From 792b324e9a07b48b31d957203fa633d4b296ae48 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 14:20:20 -0700 Subject: [PATCH 05/28] docs: qualify the Cholesky timing table The shared-covariance column is not monotone in n_landmarks (8.7 s at 1 000, 11.8 at 1 500, 10.8 at 2 000) because at these sizes it is dominated by JAX compilation rather than by the solve. Presenting it as a clean measurement would invite a reader to derive a scaling from it that is not there. It is labelled as a fixed cost independent of the gene count, which is the only property the argument actually uses, and the per-gene column -- the one that is linear in genes and carries the two-pass prescription -- is named as the one to take seriously. --- docs/source/resource_planning.rst | 14 ++++++++++---- 1 file changed, 10 insertions(+), 4 deletions(-) diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index 004d8bc..e3b5097 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -142,7 +142,7 @@ 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. Timed on one core-count of this machine, per gene: +loop. Timed on one machine, 20 genes per measurement: .. list-table:: :header-rows: 1 @@ -170,9 +170,15 @@ loop. Timed on one core-count of this machine, per gene: - 8.5 s Those are wall-clock seconds on one shared machine and will not transfer -exactly, but the **shape** does: the shared column is flat in the gene count -while the per-gene column is linear in it. Multiply it out, because a per-gene -figure is not something you can act on. At 2 000 landmarks and 8.5 s per gene: +exactly. The shared column is not monotone in ``n_landmarks`` because at these +sizes it is dominated by JAX compilation rather than by the factorisation, so +read it as *"a fixed cost, independent of the gene count"* rather than as a +measurement of the solve. The per-gene column is the one to take seriously, and +what it shows is that it is **linear in the gene count** where the shared +column is flat. + +Multiply it out, because a per-gene figure is not something anyone can act on. +At 2 000 landmarks and 8.5 s per gene: .. list-table:: :header-rows: 1 From 6685622cfe354307ee672bbb1eee361d79995b15 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 14:28:41 -0700 Subject: [PATCH 06/28] docs: withdraw the cubic-scaling claim for n_landmarks Caught by re-reading the page against its own table. I had written that the per-gene Cholesky is 'roughly cubic in n_landmarks', which a Cholesky's O(n^3) flop count makes plausible and which the measurement does not support. 500 -> 2 000 landmarks is 4x and 0.52 -> 8.5 s per gene is 16x, i.e. an overall exponent of 2.0, not 3. Worse, the per-interval ratios are 11.7x, 1.3x, 1.1x (exponents 3.55, 0.62, 0.30), which is not a power law at all -- it looks like JAX warm-up at the first size and then a flat region, not a scaling. So no exponent is claimed. The four points are presented as four measured points in the range they cover, with an explicit instruction not to extrapolate a scaling from them and to time a single gene at your own landmark count if the number matters. The memory claim is untouched: quadratic there is exact arithmetic (2 * n_landmarks^2 * n_genes * 8), not a fit. Same correction in GPSettings.n_landmarks and in notebook 03's lever table. --- docs/source/resource_planning.rst | 28 ++++++++++++++++++---------- examples/03_sample_variance.ipynb | 6 +++--- kompot/settings.py | 4 ++-- 3 files changed, 23 insertions(+), 15 deletions(-) diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index e3b5097..426082d 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -516,13 +516,20 @@ is quadratic in it. Same 1 000 genes, from the dry run: - 380.2 GiB - 7.7 GiB -Halving the landmark count quarters the covariance footprint. It also cuts the -per-gene factorisation, which is roughly cubic in ``n_landmarks``: measured -above, one gene costs 0.52 s at 500 landmarks and 8.5 s at 2 000. Landmarks -control the resolution of the cell-state approximation, so this is a genuine -accuracy trade-off rather than a free saving, but 5 000 landmarks is rarely -required to resolve a covariance structure that is being summarised gene by -gene. +Halving the landmark count quarters the covariance footprint, which is exact +arithmetic rather than a measurement. It also cuts the per-gene factorisation, +though **by how much is not something these four points establish**: the timing +table above rises from 0.52 s to 8.5 s per gene between 500 and 2 000 +landmarks, but the per-interval ratios are 11.7x, 1.3x and 1.1x, which is not a +power law. A Cholesky is O(n^3) in flops and the measurement plainly is not +following that at these sizes, so take the table as four measured points in +the range it covers and do not extrapolate a scaling from it. Measure at your +own landmark count if the time matters. + +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 @@ -552,9 +559,10 @@ thousand genes covers the interesting tail of a Mahalanobis ranking at a cost you can absorb overnight; the whole transcriptome costs two days for a result whose last nine-tenths you were never going to read. -If you need more coverage than that, cut ``n_landmarks`` in the same breath: -it reduces the per-gene cost as well as the memory, so it buys back the time -the extra genes spend. +If you need more coverage than that, cut ``n_landmarks`` in the same breath. +It reduces the per-gene cost as well as the memory, so it buys back some of the +time the extra genes spend — how much, at your dimensions, is worth a single +timed gene rather than an extrapolation. 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 diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index 8182766..3e7479a 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -337,8 +337,8 @@ "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", - " the per-gene factorisation roughly *cubic*, so it buys on both axes;\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", @@ -1379,7 +1379,7 @@ "|---|---|---|---|\n", "| `genes=` | **linear** | **linear** | top ~1 000 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** | ~**cubic** | 1 000\u20132 000 when 5 000 is too slow |\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", diff --git a/kompot/settings.py b/kompot/settings.py index c2c3903..d11325d 100644 --- a/kompot/settings.py +++ b/kompot/settings.py @@ -47,8 +47,8 @@ class GPSettings: 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, which is - roughly cubic in it. See + 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. From a489772817dcf3dfa64c1a9d4ceed831c367469b Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 14:33:40 -0700 Subject: [PATCH 07/28] style: drop the jax and tqdm imports the rewrite orphaned pyflakes over the whole package, diffed against origin/main, reported exactly two findings that were mine: 'jax' and 'tqdm.auto.tqdm' imported but unused in differential/differential_expression.py. Both were used only by the eager per-gene assembly that LazyGeneCovariance replaced -- jax for the isinstance(..., jax.Array) checks that chose whether to convert combined_cov, tqdm for the three per-gene __setitem__ progress loops. Nothing else in the module references either. With them gone the package's pyflakes findings are an identical SET to origin/main's: no new ones, none removed. CI does not run flake8, so this is housekeeping rather than a gate -- but they are dead because of this branch and should not be left for someone else to wonder about. --- kompot/differential/differential_expression.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/kompot/differential/differential_expression.py b/kompot/differential/differential_expression.py index 85815a5..68ad81d 100644 --- a/kompot/differential/differential_expression.py +++ b/kompot/differential/differential_expression.py @@ -1,12 +1,10 @@ """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, LazyGeneCovariance from ..batch_utils import apply_batched, is_jax_memory_error # noqa: F401 From 03b4730d03370419655546a6ff5ecd050e4552ab Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 16:52:23 -0700 Subject: [PATCH 08/28] fix(docs,tests): remediate the skeptic pass on #29 (F1-F6) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit w242sk returned `check`/second-pass on PR #29 with six findings. All six are accepted and remediated; every number below is my own re-measurement, not the skeptic's, taken with a positive control printed. F1 — the Cholesky table measured thread contention, not Kompot -------------------------------------------------------------- The published table (8.5 s/gene at n_landmarks=2000) and the 2.4 h / 4.7 h / 47 h projection built on it were taken with unrestricted BLAS threads on a shared node at load 40-96. They measured oversubscription. The table refuted itself and I could have seen it without measuring anything: it claimed 170.9 s for a 20-gene factorisation while a complete 20-gene sample-variance run at the same n_landmarks finishes in ~48 s. A part cannot exceed its whole. Re-measured inside a real kompot.de run, OMP_NUM_THREADS=1, with the covariance object printed as a positive control (`LazyGeneCovariance (2000, 2000, 20)`): L=500 0.0160 s/gene L=2000 0.2462 s/gene L=1000 0.0616 s/gene L=5000 2.0837 s/gene against 6.1353 s/gene unrestricted at L=2000 on the same host, same commit, same run — 25x. Bare np.linalg.cholesky at n=2000: 0.125 s at 21.3 GFLOPS pinned vs 5.20 s at 0.5 GFLOPS unrestricted. A throughput figure that COLLAPSES as threads are added is contention, not work. Corrected: 20 000 genes is ~1.4 h at L=2000 and ~11.6 h at L=5000. Two conclusions drawn from the wrong number are withdrawn rather than patched. "Memory stops being the binding constraint, compute takes over" and "merely interminable" are gone from every surface. The prescription now rests on the MEMORY argument, which is exact arithmetic, and names the regime: held in memory a whole transcriptome is 7 552 GiB and out of reach; offloaded it is 101 GiB and a long job. So ~1 000 genes is a budget justified on value — pass 1 ranks the genes anyway, 1 000 costs a twentieth on both axes, and FDR is not calibrated for sample variance so the rest buy no testable calls. A new warning tells users to pin OMP_NUM_THREADS on a shared node, which is the product finding hiding inside F1: the per-gene loop is one multi-threaded LAPACK call on a matrix small relative to the core count. F2 — the memory metric could not see what the memmap path moved --------------------------------------------------------------- Reporting only `Anonymous` and justifying it as keeping page cache from "flattering the reading" was backwards: choosing Anonymous is what EXCLUDES the 823 MiB the dask-less path moves into file-backed residency. Four instruments, 600 landmarks / 150 genes / 4 donors / 4 threads, MiB: no SV anon 958 rss 1218 file 260 disk 0 in-mem anon 1975 rss 2238 file 263 disk 0 dask anon 1180 rss 1442 file 262 disk 0 nodask anon 1139 rss 2224 file 1085 disk 824 As a share of the in-memory extra: dask 22% anon / 22% Rss; nodask 18% / 99%. So the dask path — the recommended one — wins on every instrument, and now leads the section. The dask-less path is described as converting unreclaimable anonymous pages into reclaimable page cache: a win for surviving a squeeze, not a smaller footprint, with the cgroup caveat stated plainly (Slurm --mem and container limits charge page cache to the same budget) and an explicit note that this is a mechanism, not an OOM verdict. F3 — the regression suite had no potency ----------------------------------------- All 14 cases passed against a mutant that rebuilt the dense sum in __init__ — #26's exact defect — because they asserted numerical equivalence and input validation, which an eager implementation satisfies. Base absence could not stand in: the pre-existing suite passed on 19ee1a1, which HAD the defect. Two guards added that assert the ALLOCATION property: - a recording proxy distinguishing `[:, :, g]` from any whole-array read; - a measured guard that walks a 62 MiB tensor and asserts anonymous growth stays within 12 gene-matrices. Both verified RED against the same mutation ("anonymous memory grew 67.4 MiB walking a 61.9 MiB tensor") while the other 14 stayed green on it, which is what makes this a demonstration rather than an assertion. utils.py restored from a byte-checked backup, md5 verified. F4, F5, F6 ---------- Five shipped sites still asserted the pre-fix #25 behaviour in the PR that fixes it — the de() docstring that ships in help(), the guide's warning, anndata.rst, simplified.rst and notebook 02 cell 24. Measured ratio is now 1.000 and they say so. examples/03_sample_variance.ipynb cells 10 and 12 shipped a pre-fix plan as stored output (18.63 GB of disk that is now 0, and a per-gene Dask path LazyGeneCovariance no longer takes). Cleared with a note that says what was cleared and why: a blank that says why beats a number that lies. Not re-executed — that means running the real analysis on shared infrastructure at a configuration this PR prices in hours, and publishing fresh outputs from a lab dataset into a public repo; both are the operator's call. Every RESULT cell is untouched and still correct, since the fix is bitwise identical. resource_estimation.py no longer promises that dask "parallelises the per-gene work" — LazyGeneCovariance deliberately bypasses the delayed branch — and now names the cgroup charge instead. The is_dask branch in compute_mahalanobis_distances is DOCUMENTED, not deleted: it is unreachable from Kompot's own call graph, but `compute_mahalanobis_distance` is in `kompot.__all__` and forwards a caller-supplied covariance, so an external caller can still reach it. Unreachable internally is not dead. The plan's sample-variance memory delta is 13.3x optimistic (16.7 MiB charged against 222 MiB measured). Documented as "the plan is a floor, not a total", with the reason it is useful anyway: what it gets exactly right is the part that explodes. Verification ------------ Suite 27 failed / 2014 passed. The one failure outside the established set, test_smooth_infrastructure.py::TestRunInfoSmooth::test_auto_detect_smooth, is a load artefact from a run at load 85 in which an xdist worker crashed and was replaced: it passes 3/3 serially at load 46, its whole file passes 27/27, and it passes on origin/main run the same way. It is in `smooth`, which this PR does not touch. Docs build 24 warnings, an identical set to origin/main. Rendered-HTML leak scan zero — it caught nested inline RST markup twice more in the new prose, the third and fourth instances this session, which is filed as #30 because a missing lint is not a habit problem. --- docs/source/anndata.rst | 2 +- docs/source/index.rst | 4 +- docs/source/resource_planning.rst | 325 +++++++++------ docs/source/simplified.rst | 2 +- .../02_differential_expression_detailed.ipynb | 6 +- examples/03_sample_variance.ipynb | 389 ++---------------- kompot/anndata/differential_expression.py | 9 +- kompot/cli/de.py | 2 +- kompot/resource_estimation.py | 10 +- kompot/utils.py | 9 + .../test_sample_variance_cost_regressions.py | 147 +++++++ 11 files changed, 396 insertions(+), 509 deletions(-) diff --git a/docs/source/anndata.rst b/docs/source/anndata.rst index b02b0be..bd014d9 100644 --- a/docs/source/anndata.rst +++ b/docs/source/anndata.rst @@ -52,7 +52,7 @@ overwrites, and verify parameters. **Key features:** - **Memory and disk estimation**: Calculates expected resource usage for all intermediate arrays and final results -- **Null genes accounting**: Estimates resource inflation from null distribution genes when ``null_genes`` is given as an explicit number +- **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 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`` diff --git a/docs/source/index.rst b/docs/source/index.rst index b12ffbb..7b33fea 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -138,8 +138,8 @@ See :doc:`simplified` for the full API and all available settings. 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 several seconds per gene**, - so it belongs in a second pass over a restricted gene list, never in a + 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. diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index 426082d..a931b54 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -52,13 +52,17 @@ instead: 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 reason the split is not optional is that the two costs respond to different -levers. ``store_arrays_on_disk`` removes the memory term almost entirely -(:ref:`disk-offload`), and ``n_landmarks`` shrinks it quadratically — but -neither touches the **per-gene Cholesky factorisation**, which is linear in the -gene count and nothing makes it cheaper except analysing fewer genes. Restrict -the gene list and both costs fall together; offload without restricting and you -have merely moved from running out of memory to running out of time. +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 ``n_landmarks`` shrinks +it quadratically, but neither touches the **per-gene Cholesky factorisation**, +which is linear in the gene count and which only a shorter gene list reduces. +Restricting ``genes`` is the one lever that moves both. + +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: @@ -142,66 +146,63 @@ 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. Timed on one machine, 20 genes per measurement: +loop. + +Measured inside a real ``kompot.de`` sample-variance run, 20 genes, with +``OMP_NUM_THREADS=1``: .. list-table:: :header-rows: 1 - :widths: 22 26 26 26 + :widths: 24 26 24 26 * - ``n_landmarks`` - - shared, 20 genes - - per-gene, 20 genes - - per-gene, per gene + - per gene + - 1 000 genes + - 20 000 genes * - 500 - - 0.61 s - - 10.5 s - - 0.52 s + - 0.016 s + - 16 s + - 5 min * - 1 000 - - 8.7 s - - 121.6 s - - 6.1 s - * - 1 500 - - 11.8 s - - 156.7 s - - 7.8 s + - 0.062 s + - 1 min + - 21 min * - 2 000 - - 10.8 s - - 170.9 s - - 8.5 s + - 0.246 s + - 4 min + - 1.4 h + * - 5 000 (default) + - 2.08 s + - 35 min + - 11.6 h -Those are wall-clock seconds on one shared machine and will not transfer -exactly. The shared column is not monotone in ``n_landmarks`` because at these -sizes it is dominated by JAX compilation rather than by the factorisation, so -read it as *"a fixed cost, independent of the gene count"* rather than as a -measurement of the solve. The per-gene column is the one to take seriously, and -what it shows is that it is **linear in the gene count** where the shared -column is flat. +.. warning:: -Multiply it out, because a per-gene figure is not something anyone can act on. -At 2 000 landmarks and 8.5 s per gene: + **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``: -.. list-table:: - :header-rows: 1 - :widths: 30 35 35 + .. code-block:: text - * - genes in pass 2 - - factorisation time - - in practice - * - 1 000 - - ~2.4 hours - - an overnight job - * - 2 000 - - ~4.7 hours - - still an overnight job - * - 20 000 (all) - - ~47 hours - - two days of compute + 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. -That, rather than memory, is what bounds a sample-variance run once the -tensors are offloaded, and it is the reason -:ref:`the two-pass workflow ` remains the prescription. A -whole-transcriptome sample-variance run is no longer impossible; it is merely -interminable, which is a worse failure because it looks like progress. +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``: @@ -264,11 +265,13 @@ 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. -**Memory stops being the binding constraint, and compute takes over.** The -20 000-gene disk-backed plan at 101 GiB would fit on a large node; the 20 000 -per-gene Cholesky factorisations behind it would not fit in your week. That is -why :ref:`the two-pass workflow ` is the prescription even -now that the memory problem is solved. +**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``: @@ -349,6 +352,23 @@ 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 @@ -356,21 +376,19 @@ stderr, and exits non-zero when the plan is infeasible. See :doc:`cli`. .. _null-genes-cost: -.. warning:: +.. note:: - The dry run reads ``null_genes`` as you pass it and does not resolve the - ``"auto"`` default. With ``FDRSettings`` left at its default and no - ``sample_col``, the real run adds 2 000 null genes but the dry run counts - none, so the plan is optimistic by that margin. Pass - ``fdr=kompot.FDRSettings(null_genes=2000)`` to the dry run to see the true - figure. Runs *with* ``sample_col`` are unaffected, because ``"auto"`` - resolves to ``0`` there and the dry run also counts 0. Tracked in - `settylab/kompot#25 `_. + 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 — 1 000 real genes + plus 2 000 null genes is a 3 000-gene tensor. - Null genes are charged at the same per-gene rate as real ones. If you - override the default and request FDR alongside sample variance, each null - gene gets its own landmark covariance matrix too: 1 000 real genes plus - 2 000 null genes is a 3 000-gene tensor. + 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: @@ -416,55 +434,93 @@ What it does mechanically depends on whether ``dask`` is installed: installing ``dask`` removes it. Measured on 1 500 cells, 150 genes, 600 landmarks and 4 donors (412 MiB per -tensor), peak **anonymous** memory sampled from ``/proc/self/smaps_rollup`` so -that page cache for a memory-mapped file cannot flatter the reading: +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: 30 16 16 18 20 + :widths: 28 13 13 15 12 12 * - run - - peak memory - - bytes written - - wall clock - - vs. no sample variance + - anon + - Rss + - file-backed + - written + - wall * - no sample variance - - 972 MiB + - 958 MiB + - 1 218 MiB + - 260 MiB - 0 - - 33 s - - — + - 36 s * - sample variance, in memory - - 1 973 MiB + - 1 975 MiB + - 2 238 MiB + - 263 MiB - 0 - - 46 s - - +1 001 MiB + - 49 s * - ``store_arrays_on_disk``, with ``dask`` - - 1 153 MiB + - 1 180 MiB + - 1 442 MiB + - 262 MiB - **0** - - 48 s - - +181 MiB + - 50 s * - ``store_arrays_on_disk``, no ``dask`` - - 1 148 MiB + - 1 139 MiB + - **2 224 MiB** + - **1 085 MiB** - 824 MiB - - 48 s - - +176 MiB - -The offloaded runs cost about a sixth of the extra memory that the in-memory -run does, and the gap widens with the gene count, because the in-memory column -carries a term linear in genes and the offloaded ones do not. Repeated at 900 -landmarks and 120 genes (742 MiB per tensor), peak memory was 2 809 MiB -in memory, 1 325 MiB with ``dask`` and 1 341 MiB without, against 1 051 MiB -for a run with no sample variance at all. - -The wall-clock column is from a shared machine under load and should be read -as an order of magnitude, not a benchmark. One caveat it does show reliably: -the ``dask``-less path is **modestly slower than it was before 0.9.0** — 74 s -against 66 s on a clean 900-landmark pair — because 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. That is the price of not holding the tensor, it is -paid only when ``dask`` is absent, and installing ``dask`` avoids it in both -directions. + - 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`` + - 22% + - **22%** + * - no ``dask`` + - 18% + - **99%** + +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** @@ -517,14 +573,12 @@ is quadratic in it. Same 1 000 genes, from the dry run: - 7.7 GiB Halving the landmark count quarters the covariance footprint, which is exact -arithmetic rather than a measurement. It also cuts the per-gene factorisation, -though **by how much is not something these four points establish**: the timing -table above rises from 0.52 s to 8.5 s per gene between 500 and 2 000 -landmarks, but the per-interval ratios are 11.7x, 1.3x and 1.1x, which is not a -power law. A Cholesky is O(n^3) in flops and the measurement plainly is not -following that at these sizes, so take the table as four measured points in -the range it covers and do not extrapolate a scaling from it. Measure at your -own landmark count if the time matters. +arithmetic rather than a measurement. It cuts the per-gene factorisation +steeply too: pinned to one thread, 0.016 s per gene at 500 landmarks against +2.08 s at 5 000. Those four points are close to the cubic flop count a Cholesky +implies — 500 to 5 000 is 10x the landmarks for 130x the time — but they are +four timings on one machine, not a scaling law, and 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 @@ -550,23 +604,34 @@ than a statistical one, and Kompot supplies no default: ``genes=None`` means *all of them*, which is the one setting to avoid here. **Roughly 1 000 genes is a recommendation chosen from the cost curves, not a -tuned or validated parameter.** With ``store_arrays_on_disk=True`` memory is no -longer what decides it — 1 000 genes plans at 7.7 GiB, and even the whole -transcriptome plans at 101 GiB, which would fit on a large node. What decides -it is the per-gene Cholesky: at 2 000 landmarks and roughly 8.5 s per gene, -1 000 genes is about **2.4 hours** and 20 000 genes about **47 hours**. One -thousand genes covers the interesting tail of a Mahalanobis ranking at a cost -you can absorb overnight; the whole transcriptome costs two days for a result -whose last nine-tenths you were never going to read. - -If you need more coverage than that, cut ``n_landmarks`` in the same breath. -It reduces the per-gene cost as well as the memory, so it buys back some of the -time the extra genes spend — how much, at your dimensions, is worth a single -timed gene rather than an extrapolation. +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 ~1 000 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; +* 1 000 genes costs about a twentieth 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 000 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. +count by a single-gene timing at your landmark count, pinned. See also -------- diff --git a/docs/source/simplified.rst b/docs/source/simplified.rst index cc07c76..f7b22de 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -200,7 +200,7 @@ report. plan.total_disk_required # bytes plan.is_feasible -How to read the report, what to compare, and the caveat about ``null_genes``: +How to read the report and what to compare: :doc:`Planning Memory and Disk `. diff --git a/examples/02_differential_expression_detailed.ipynb b/examples/02_differential_expression_detailed.ipynb index 9e25f31..5e47400 100644 --- a/examples/02_differential_expression_detailed.ipynb +++ b/examples/02_differential_expression_detailed.ipynb @@ -1344,9 +1344,9 @@ "`requirements` breakdown, so you can price two configurations against each other before\n", "committing to either.\n", "\n", - "Note that the estimate reads `null_genes` as you pass it and does not resolve the\n", - "`\"auto\"` default, so the explicit `null_genes=2000` below is what makes the plan match\n", - "the run it is pricing.\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 3e7479a..14eaefd 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -347,9 +347,28 @@ "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": { @@ -359,66 +378,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", @@ -449,7 +409,7 @@ }, { "cell_type": "code", - "execution_count": 7, + "execution_count": null, "id": "cell-071", "metadata": { "execution": { @@ -459,302 +419,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", @@ -1383,9 +1048,11 @@ "| `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 to disk solves memory and leaves the factorisation untouched, so a\n", - "whole-transcriptome sample-variance run stops being impossible and becomes merely\n", - "interminable. Restrict the gene list.\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", "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", diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index ba9f73d..63ca02a 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -147,12 +147,9 @@ def de( instead of running the analysis. Returns a :class:`~kompot.resource_estimation.ResourcePlan` carrying ``total_memory_required``, ``total_disk_required``, ``is_feasible`` - and a per-array ``requirements`` list. Note that the estimate reads - ``FDRSettings.null_genes`` as given and does not resolve the - ``"auto"`` default, so a dry run without ``sample_col`` and with - ``null_genes="auto"`` omits the 2 000 null genes the real run would - add; pass ``null_genes=2000`` explicitly to see them - (https://github.com/settylab/kompot/issues/25). + 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`. diff --git a/kompot/cli/de.py b/kompot/cli/de.py index 8edb713..1251def 100644 --- a/kompot/cli/de.py +++ b/kompot/cli/de.py @@ -127,7 +127,7 @@ def add_de_parser(subparsers) -> argparse.ArgumentParser: "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 several seconds per gene at " + "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." diff --git a/kompot/resource_estimation.py b/kompot/resource_estimation.py index b72e479..f28f763 100644 --- a/kompot/resource_estimation.py +++ b/kompot/resource_estimation.py @@ -1045,10 +1045,12 @@ def estimate_differential_expression_resources( if not DASK_AVAILABLE: plan.warnings.append( "Disk storage requested but dask is not installed, so the " - "covariance tensors are written as memory-mapped .npy files " - "and computed sequentially. Install dask to evaluate them " - "lazily instead, which writes nothing and parallelises the " - "per-gene work: pip install 'dask[array]'" + "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/utils.py b/kompot/utils.py index 9fa5951..c8821a2 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -239,6 +239,15 @@ def compute_mahalanobis_distances( # Check if covariance is a Dask array. # + # REACHABILITY: since 0.9.0 nothing inside Kompot reaches this branch. The + # only gene-specific covariance the package constructs is a + # LazyGeneCovariance (see below), and the shared path passes a 2-D array. + # It is retained rather than deleted because `compute_mahalanobis_distance` + # is exported in `kompot.__all__` and forwards a caller-supplied covariance + # straight through, so an external caller can still hand this function a + # `dask.array.Array` directly. Deleting it would remove a documented, + # tested path that Kompot itself simply no longer uses. + # # 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 diff --git a/tests/test_sample_variance_cost_regressions.py b/tests/test_sample_variance_cost_regressions.py index 99c6315..ddf9f9f 100644 --- a/tests/test_sample_variance_cost_regressions.py +++ b/tests/test_sample_variance_cost_regressions.py @@ -333,3 +333,150 @@ def test_lazy_view_refuses_a_shared_term(): 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)) + + view = LazyGeneCovariance(terms, base=base) + before = anon_bytes() + 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" + ) From fcfa4068829c250ec0e237e62481461a3c565d20 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 16:59:15 -0700 Subject: [PATCH 09/28] fix: correct the F6 reachability reason and publish the dask share as a range Three corrections from the skeptic's answer on the reconciliation channel. None is a dispute; all three are places where my remediation was imprecise. F6 REASON WAS WRONG, CONCLUSION WAS RIGHT. My comment justified keeping the is_dask branch by saying `compute_mahalanobis_distance` is in `kompot.__all__` and forwards a caller-supplied covariance. But the exported name is the SINGULAR, which passes a 2-D covariance at utils.py:602 and so cannot reach a 3-D branch; the PLURAL is not exported at all. Verified both. The real reason is concrete and checkable: tests/test_mahalanobis_approaches.py builds a bare 3-D dask array with DiskStorage.as_dask_array and passes it straight into compute_mahalanobis_distances, asserting it agrees with the gene-specific path. So the branch is TESTED but production-unreachable, which is not the same as dead. The comment now enumerates every in-repo caller and cites the covering test, so the next reader does not re-derive it. DASK MEMORY SHARE IS NOW A RANGE. I published 22%/22%; the skeptic measured 17%/17% on two independent rigs. Neither is wrong: Dask's threaded scheduler sizes its pool from the CPU count independently of OMP_NUM_THREADS, so that arm's peak is set by a concurrency neither of us pinned. Published as 17-22% with the mechanism stated, rather than two significant figures that would not reproduce. The dask-less row agreed to within a point across both rigs and is unchanged. PINNING ALSO BUYS REPRODUCIBILITY, which is worth more than the speed. Two independent measurements of the same configuration at 5 000 landmarks, single-threaded, at 1-minute loads of 45.9 and 62.8 -- a 37% difference -- came out 1.4% apart (2.084 vs 2.113 s/gene). Unpinned, the same step moved by 25x. That is now the second half of the threading warning. --- docs/source/resource_planning.rst | 19 ++++++++++++++++--- kompot/utils.py | 22 ++++++++++++++-------- 2 files changed, 30 insertions(+), 11 deletions(-) diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index a931b54..b7d3790 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -199,6 +199,13 @@ Measured inside a real ``kompot.de`` sample-variance run, 20 genes, with 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 — came out at 2.084 and 2.113 s per gene, **1.4% apart**. + Unpinned, the same step moved by a factor of 25. Pin the threads and the + measurement stops being a property of the machine's mood. + 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 @@ -485,11 +492,17 @@ sample variance at all: - on ``Anonymous`` - on ``Rss`` * - with ``dask`` - - 22% - - **22%** + - 17-22% + - **17-22%** * - no ``dask`` - 18% - - **99%** + - **96-99%** + +The ``dask`` row is given as a range because that path's peak is set by a +concurrency neither figure pins: Dask's threaded scheduler sizes its pool from +the CPU count, independently of ``OMP_NUM_THREADS``. Two independent rigs on +this host, one pinned to four BLAS threads and one unpinned, bracket it at 17% +and 22%. The ``dask``-less row reproduced to within a point across both. 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 diff --git a/kompot/utils.py b/kompot/utils.py index c8821a2..fa12f6b 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -239,14 +239,20 @@ def compute_mahalanobis_distances( # Check if covariance is a Dask array. # - # REACHABILITY: since 0.9.0 nothing inside Kompot reaches this branch. The - # only gene-specific covariance the package constructs is a - # LazyGeneCovariance (see below), and the shared path passes a 2-D array. - # It is retained rather than deleted because `compute_mahalanobis_distance` - # is exported in `kompot.__all__` and forwards a caller-supplied covariance - # straight through, so an external caller can still hand this function a - # `dask.array.Array` directly. Deleting it would remove a documented, - # tested path that Kompot itself simply no longer uses. + # 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 From 85fcb5ef31eb6d557ce44a8e4c6297220dbe8219 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 17:21:08 -0700 Subject: [PATCH 10/28] docs: the dask memory share is run-to-run noise, not an unpinned concurrency I published 17-22% with a mechanism -- Dask's threaded scheduler sizing its pool from the CPU count independently of OMP_NUM_THREADS -- and the mechanism does not survive measurement. Pinning the pool at 1, at 4, and leaving it at the default 36 gave dask-arm peaks of 1173/1209, 1152/1163 and 1134/1138 MiB: all inside one spread, and the SINGLE-worker runs were if anything the highest. What actually explains the 17% vs 22% between two rigs is that the share is a ratio of two differences -- an extra of ~200 MiB between peaks of ~950 and ~1150 MiB -- so ordinary noise in either peak swings it hard. Across twelve runs here the no-sample-variance baseline spans 938-981 MiB and the dask peak spans 1134-1209 MiB, which brackets the share at 14-27% with nothing changing but the run. Two rigs landing on 17 and 22 is that spread, not a methodology difference, and neither figure was reproducible enough to print. So the table now gives one significant figure and states the spread, and says what the measurement does support, which is not close: about a fifth on Anonymous for both offloaded paths, against essentially all of it on Rss for the dask-less one. A factor of five on the same run is the finding; the second digit never was. Found while settling whether the concurrency could be pinned at all. It can, by the caller, through dask.config -- but kompot's own knob for it, SampleVarianceEstimator(dask_num_workers=N), sets the config key "pool.num-workers" where the threaded scheduler reads "num_workers", so it configures nothing and logs that it succeeded. Filed as #31; pre-existing and not fixed here. --- docs/source/resource_planning.rst | 32 +++++++++++++++++++++---------- 1 file changed, 22 insertions(+), 10 deletions(-) diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index b7d3790..5f51fe0 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -492,17 +492,29 @@ sample variance at all: - on ``Anonymous`` - on ``Rss`` * - with ``dask`` - - 17-22% - - **17-22%** + - ~1/5 + - **~1/5** * - no ``dask`` - - 18% - - **96-99%** - -The ``dask`` row is given as a range because that path's peak is set by a -concurrency neither figure pins: Dask's threaded scheduler sizes its pool from -the CPU count, independently of ``OMP_NUM_THREADS``. Two independent rigs on -this host, one pinned to four BLAS threads and one unpinned, bracket it at 17% -and 22%. The ``dask``-less row reproduced to within a point across both. + - ~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 From b30c601bd2e7b871798655b5c6173fd6af3ddac3 Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 18:56:12 -0700 Subject: [PATCH 11/28] fix: remediate the delta skeptic pass (D1-D5) Second pass by w242sk at 85fcb5e. Five findings, all accepted. D5 -- MY POTENCY GUARD WAS A FALSE GREEN IN CI'S CONFIGURATION -------------------------------------------------------------- I had asked the skeptic not to take "both guards went red" on my word. It found that test_peak_memory_stays_bounded_by_one_gene passes against the eager-dense mutant whenever the whole file runs, and fails only standalone. Reproduced here before fixing: whole-file x3 against the mutant, only the recording-proxy guard failed; `-k peak_memory` alone, the memory guard failed. Cause was one line. The window opened AFTER the constructor: view = LazyGeneCovariance(terms, base=base) before = anon_bytes() # <- too late An eager implementation materialises in __init__ -- that is #26's actual shape -- so the 62 MiB allocation happened entirely outside the measured window, and what the guard measured was incidental churn whose size depends on allocator state left by sibling tests. Window now opens before the constructor. Verified RED 3/3 whole-file against the same mutant (67.4 MiB against a 6.19 MiB bound, the order-of-magnitude margin the comment claims), and GREEN 3/3 restored, utils.py md5-verified against a pre-mutation backup. I verified that guard went red. I verified it in the configuration where it works. The rule: run a potency check in the configuration the suite actually runs in, not the one that isolates the case. D1 -- THE CUBIC CLAIM I WITHDREW CAME BACK IN MY OWN REMEDIATION ---------------------------------------------------------------- 6685622 withdrew "roughly cubic in n_landmarks" as unsupported. 03b4730 then wrote "close to the cubic flop count a Cholesky implies -- 500 to 5 000 is 10x the landmarks for 130x the time", which refutes itself: cubic is 1000x, not 130x. Measured exponent is 2.11 over the four pinned points, consistent interval to interval (1.94, 2.00, 2.33); cubic predicts 16 s/gene at L=5000 against 2.08 s measured, a 7.7x overshoot. Now stated as quadratic-ish WITH the mechanism: the step also carries O(n^2) work (per-gene materialisation, the triangular solve) that dominates at small n, and single-threaded dpotrf throughput climbs 8.2 -> 37.7 GFLOPS across the range, a 4.6x improvement that cancels much of the extra work. Noted that the exponent should drift toward 3 at larger n. The original withdrawal was correct because the UNPINNED data was not a power law at all, so no exponent was claimable. Pinning the threads is what made a scaling statement possible -- and the statement made was the wrong one. A correction can restore the error it corrected, once better evidence makes the claim sayable again. D2 -- de() AND GPSettings CONTRADICTED EACH OTHER IN THE SAME PACKAGE --------------------------------------------------------------------- Four sites said n_landmarks does not touch the per-gene factorisation and only a shorter gene list reduces it: resource_planning.rst's top-of-page summary (543 lines above its own table showing 0.016 s at 500 against 2.08 s at 5 000), simplified.rst, README.md, and de()'s docstring -- which ships in help() -- while settings.py:50, written in the original PR, already said the correct thing. Swept for the class rather than patching the three that were named. Like D1 this is a consequence of the F1 fix rather than a survivor of it: the clean measurement made the landmark effect on compute visible for the first time, and the older prose was never revisited. D3 -- THE 25x WAS A LEVEL CONTRAST PRESENTED AS A SPREAD --------------------------------------------------------- The reproducibility paragraph read as "unpinned measurements vary by 25x". The 25x is pinned against unrestricted at one moment: it says how far the level moves, not how much an unpinned measurement scatters. The two unrestricted measurements taken here agreed to within 6-10%. What makes them unusable is not scatter but being wrong by more than an order of magnitude. Prescription and level claim were both right; only the inference was unsupported. D4 -- THE CHANGELOG WAS NEVER TOUCHED BY THE ENTIRE REMEDIATION ---------------------------------------------------------------- `git diff --name-only a489772 85fcb5e` did not include it. At the tip it still reported peak anonymous memory as the sole instrument -- zero occurrences of Rss, file-backed or page cache in the whole file -- still said "disk-backed storage is now cheaper than in-memory storage" unqualified across two offloaded rows presented identically, and still credited the dask-less path with "57% less memory", the Anonymous-only figure that is about 1% on Rss. Both instruments are now in the table, the "cheaper than in-memory" claim is scoped to the dask path where it holds on every instrument, the dask-less path is described as converting anonymous pages to reclaimable page cache rather than as a smaller footprint, and the cgroup caveat is stated. This file becomes the GitHub release note; the docs page had been corrected and the release narrative had not. VERIFICATION ------------ Suite 26 failed / 2014 passed, no crashed workers, failing set BYTE-IDENTICAL to the established pre-change set (cmp clean). Docs 24 warnings, identical set to origin/main. Rendered-HTML leak scan zero. Every published snippet executed. --- CHANGELOG.md | 62 +++++++++++-------- README.md | 4 +- docs/source/resource_planning.rst | 51 +++++++++++---- docs/source/simplified.rst | 5 +- kompot/anndata/differential_expression.py | 4 +- .../test_sample_variance_cost_regressions.py | 9 ++- 6 files changed, 91 insertions(+), 44 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 3e7e16b..21381b3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -25,32 +25,42 @@ gene by gene with `__setitem__`. Three consequences, all silent: 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), peak -**anonymous** memory from `/proc/self/smaps_rollup`, same machine and same -synthetic input on both trees: - -| run | 0.8.0 | 0.9.0 | bytes written, 0.8.0 → 0.9.0 | 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 148 MiB | 824 → 824 MiB | 48.2 s → 47.7 s | - -Disk-backed storage is now cheaper than in-memory storage, which it was not -before: 1 153 MiB against 1 973 MiB, where 0.8.0 made it *more* expensive -(2 782 against 2 401). Repeated at 900 landmarks / 120 genes: 2 809 MiB in -memory, 1 325 MiB with Dask and 1 341 MiB without, against 3 502 / 4 341 / -3 079 MiB on 0.8.0. - -One cost moved the other way, on the path that is not recommended. Without -`dask` the tensors are memory-mapped, and a gene slice of a C-contiguous -`(n_points, n_points, n_genes)` 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 on a clean 900-landmark pair, 74 s against 66 s — about 12% slower in -exchange for 57% less memory. Installing `dask` avoids it in both directions: -that path is 4.9x *faster* than 0.8.0 as well as lighter. Wall-clock figures -here come from a shared machine under load and should be read as orders of -magnitude. + +Measured on 1 500 cells / 150 genes / 600 landmarks (412 MiB per tensor), same +machine and same synthetic input on both trees. **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. + +| run | 0.8.0 anon | 0.9.0 anon | 0.9.0 Rss | written | wall, 0.8.0 → 0.9.0 | +|---|---|---|---|---|---| +| no sample variance | 949 MiB | 972 MiB | 1 218 MiB | 0 → 0 | 33.0 s → 32.6 s | +| sample variance, in memory | 2 401 MiB | 1 973 MiB | 2 238 MiB | 0 → 0 | 45.7 s → 46.0 s | +| `store_arrays_on_disk`, with `dask` | 2 782 MiB | **1 153 MiB** | **1 442 MiB** | 0 → 0 | 231.8 s → 47.8 s | +| `store_arrays_on_disk`, no `dask` | 2 284 MiB | 1 139 MiB | **2 224 MiB** | 824 → 824 MiB | 48.2 s → 47.7 s | + +**With `dask` — the recommended path — offloading is now cheaper than holding +the tensors in memory, on every instrument**, where on 0.8.0 it was dearer. The +extra over a run with no sample variance is about a fifth of the in-memory +extra on both `Anonymous` and `Rss`, and that path is also 4.8x faster than +0.8.0 at matched thread counts. + +**Without `dask` the picture is different and should not be read as a smaller +footprint.** That path memory-maps the tensors, so the pages become resident +page cache rather than private heap: about a fifth of the in-memory extra on +`Anonymous`, but roughly **all of it on `Rss`**. 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, so that advantage may not exist +in the environment you are actually in. If you are under a hard cap, install +`dask`. + +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, diff --git a/README.md b/README.md index 2bcc7dd..9053f0b 100644 --- a/README.md +++ b/README.md @@ -78,7 +78,9 @@ different levers: 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. - Nothing makes this cheaper except analysing fewer genes. + `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: diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index 5f51fe0..6b0b94f 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -54,10 +54,12 @@ 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 ``n_landmarks`` shrinks -it quadratically, but neither touches the **per-gene Cholesky factorisation**, -which is linear in the gene count and which only a shorter gene list reduces. -Restricting ``genes`` is the one lever that moves both. +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 @@ -202,9 +204,16 @@ Measured inside a real ``kompot.de`` sample-variance run, 20 genes, with 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 — came out at 2.084 and 2.113 s per gene, **1.4% apart**. - Unpinned, the same step moved by a factor of 25. Pin the threads and the - measurement stops being a property of the machine's mood. + 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 @@ -571,6 +580,8 @@ plan's warnings escalate; it does not switch storage modes. the ``dask`` path was 4.9x slower as well (`settylab/kompot#26 `_). +.. _other-levers: + Other levers ------------ @@ -598,12 +609,26 @@ is quadratic in it. Same 1 000 genes, from the dry run: - 7.7 GiB Halving the landmark count quarters the covariance footprint, which is exact -arithmetic rather than a measurement. It cuts the per-gene factorisation -steeply too: pinned to one thread, 0.016 s per gene at 500 landmarks against -2.08 s at 5 000. Those four points are close to the cubic flop count a Cholesky -implies — 500 to 5 000 is 10x the landmarks for 130x the time — but they are -four timings on one machine, not a scaling law, and they move with your BLAS -build and your node's load. Take them as a floor and time your own. +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. Expect the +exponent to drift up toward 3 at larger n as that headroom runs out. + +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 diff --git a/docs/source/simplified.rst b/docs/source/simplified.rst index f7b22de..b476f52 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -51,8 +51,9 @@ Differential Expression 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; - nothing removes the per-gene factorisation except analysing fewer genes. + ``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 diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index 63ca02a..7b5cc68 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -92,7 +92,9 @@ def de( 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 - nothing makes it cheaper except analysing fewer genes. + ``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:: diff --git a/tests/test_sample_variance_cost_regressions.py b/tests/test_sample_variance_cost_regressions.py index ddf9f9f..1610890 100644 --- a/tests/test_sample_variance_cost_regressions.py +++ b/tests/test_sample_variance_cost_regressions.py @@ -466,8 +466,15 @@ def anon_bytes(): terms = [np.repeat(spd[:, :, None], n_genes, axis=2) for _ in range(2)] diffs = rng.normal(size=(n_genes, n_points)) - view = LazyGeneCovariance(terms, base=base) + # 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 From 8bfcae048eb6a82e9e765b790ed188312f83b82f Mon Sep 17 00:00:00 2001 From: Dominik Date: Tue, 15 Sep 2026 20:14:09 -0700 Subject: [PATCH 12/28] fix: remediate the third skeptic pass (R1-R3) Round three by w242sk at b30c601. Three items, all accepted. R1 -- THREE SITES MY SWEEP MISSED, AND THE INSTRUMENT THAT HID THEM -------------------------------------------------------------------- My D2 sweep reported clean. It was not. Three sites still carried retired claims: docs/source/cli.rst "about 0.37 GiB and several seconds per gene" (D1) docs/source/cli.rst "nothing bounds the per-gene factorisation except a shorter gene list" (D2) examples/03_sample_variance.ipynb cell 0 "Nothing makes this cheaper except analysing fewer genes." (D2) The last is verbatim the sentence the README fix replaced, sitting in the notebook's opening cost warning -- the first thing a reader sees. THE INSTRUMENT IS THE FINDING. Every one of those is LINE-WRAPPED, so a line-based grep cannot match it. cli.rst has "about 0.37 GiB and several" on line 255 and "seconds per gene" on line 256. My grep for the phrase returned zero against text that plainly contains it, and the zero ran in the flattering direction: it confirmed a remediation that was incomplete. The working form: git show "${ref}:${path}" | tr '\n' ' ' | tr -s ' ' | grep -iF 'phrase' Re-swept all 219 tracked text files with twelve retired phrases, including eight negative controls (only a shorter gene list reduces / neither touches the / binding constraint / interminable / 47 hours / 2.4 hours / 8.5 s / close to the cubic): ZERO hits. So the rest of the sweep was complete and these three were the whole remainder. AND THE TWO FILES FAILED FOR DIFFERENT REASONS. cli.rst was last touched by f7d7476, an original-PR commit; no remediation opened it, so the sweep never reached the file. But 03_sample_variance.ipynb was opened BY 03b4730 -- my own first remediation -- to clear its stale outputs for F5, and the retired claim in its opening cell was left in place. A file being open in the editor is not a file being swept. Also reconciled: cli.rst and kompot/cli/de.py disagreed about --sample-col, the --help text saying ~2 s per gene and its own documentation page saying several seconds. Both now read ~0.37 GiB and ~2 s per gene at --n-landmarks 5000. R2 -- THE CHANGELOG TABLE MIXED TWO RUNS ------------------------------------------ Adding an Rss column spliced values from the later four-BLAS-thread run into a table whose Anonymous column came from the original paired run, updating only the dask-less anon figure. Three of four anon values then disagreed with the docs page. Worse, the anon-to-Rss gap IS the file-backed quantity, so the dask row computed it across two runs: 1442 - 1153 = 289 MiB against the 262 MiB the docs state, with no file-backed column in the CHANGELOG for a reader to catch it with. Rss is dropped from the CHANGELOG entirely and the two-instrument treatment is cross-referenced to the docs page. A release note wants a self-consistent single-instrument before/after; unlike carrying the whole table, this shape cannot re-acquire the defect later. Rebuilding from the paired run also identified the spliced value: the dask-less 0.9.0 figure is 1170 MiB, from the same run as its 2284 baseline, not the 1139 taken from the four-thread run. Every figure in that table is now one paired run. R3 -- A THIRD UNMEASURED SCALING CLAIM, IN A DOCUMENT THAT HAS WITHDRAWN TWO ----------------------------------------------------------------------------- "Expect the exponent to drift up toward 3 at larger n as that headroom runs out" was not measured, and the measurement refutes it: 2.3, 2.8, 2.3 across 2000->5000->8000->11000, no drift, with single-threaded throughput still climbing at n=11 000 (37.8 -> 42.3 -> 53.1 GFLOPS). So "as that headroom runs out" has not begun to operate anywhere a user can reach. Scoped to the range measured, with the counter-evidence stated. VERIFICATION ------------ Suite 26 failed / 2014 passed, no crashed workers, failing set byte-identical to the established pre-change set (cmp clean). Docs 24 warnings, identical set to origin/main. Rendered-HTML leak scan zero. Wrap-insensitive sweep zero. --- CHANGELOG.md | 53 ++++++++++++++++--------------- docs/source/cli.rst | 10 +++--- docs/source/resource_planning.rst | 9 ++++-- examples/03_sample_variance.ipynb | 4 ++- 4 files changed, 44 insertions(+), 32 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 21381b3..d9373f2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -27,33 +27,36 @@ 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. **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. - -| run | 0.8.0 anon | 0.9.0 anon | 0.9.0 Rss | written | wall, 0.8.0 → 0.9.0 | -|---|---|---|---|---|---| -| no sample variance | 949 MiB | 972 MiB | 1 218 MiB | 0 → 0 | 33.0 s → 32.6 s | -| sample variance, in memory | 2 401 MiB | 1 973 MiB | 2 238 MiB | 0 → 0 | 45.7 s → 46.0 s | -| `store_arrays_on_disk`, with `dask` | 2 782 MiB | **1 153 MiB** | **1 442 MiB** | 0 → 0 | 231.8 s → 47.8 s | -| `store_arrays_on_disk`, no `dask` | 2 284 MiB | 1 139 MiB | **2 224 MiB** | 824 → 824 MiB | 48.2 s → 47.7 s | +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, on every instrument**, where on 0.8.0 it was dearer. The -extra over a run with no sample variance is about a fifth of the in-memory -extra on both `Anonymous` and `Rss`, and that path is also 4.8x faster than -0.8.0 at matched thread counts. - -**Without `dask` the picture is different and should not be read as a smaller -footprint.** That path memory-maps the tensors, so the pages become resident -page cache rather than private heap: about a fifth of the in-memory extra on -`Anonymous`, but roughly **all of it on `Rss`**. 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, so that advantage may not exist -in the environment you are actually in. If you are under a hard cap, install -`dask`. +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 diff --git a/docs/source/cli.rst b/docs/source/cli.rst index cd653d1..b21701b 100644 --- a/docs/source/cli.rst +++ b/docs/source/cli.rst @@ -252,8 +252,8 @@ Example: Complete Analysis ``--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 several - seconds per gene at ``--n-landmarks 5000``. Run it as a **second pass** over + 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: @@ -736,8 +736,10 @@ 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 nothing bounds the -per-gene factorisation except a shorter gene list. Full treatment: +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 diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index 6b0b94f..f208d64 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -624,8 +624,13 @@ 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. Expect the -exponent to drift up toward 3 at larger n as that headroom runs out. +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. diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index 14eaefd..a8cf71a 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -26,7 +26,9 @@ "> 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", - "> Nothing makes this cheaper except analysing fewer genes.\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", From c4c7a3fa6f94e1264011bdf6290e46313157cdbb Mon Sep 17 00:00:00 2001 From: Dominik Date: Mon, 21 Sep 2026 16:25:52 -0700 Subject: [PATCH 13/28] test: guard the covariance graph against superlinear growth in n_genes An independent reproduction of the defect this branch fixes, arrived at from a stalled user run rather than from this branch, added as a regression guard. It drives the real DifferentialExpression.compute_mahalanobis_distances and asserts on the object that method hands downstream, so it exercises whatever the shipped code builds rather than a particular helper. It counts TASKS, not wall time, so it cannot flake on a loaded CI machine. On origin/main it fails: gene-specific covariance graph grew superlinearly: 7790 tasks at 20 genes vs 314360 at 80 (40.4x for 4x the genes; linear is 4x, quadratic is 16x) On this branch it passes unmodified, because a LazyGeneCovariance carries no task graph at all: the guard finds nothing to measure and returns. The blowup is not bounded here, it is unrepresentable. The first version of this test asserted the REPRESENTATION rather than the property and errored on this branch with AttributeError: 'LazyGeneCovariance' object has no attribute '__dask_graph__' which was the probe's assumption, not this branch's behaviour. _task_count now returns None for an object with no graph and the guard treats that as a pass, so the same test expresses the property across implementations that differ in KIND rather than in degree. The accompanying equivalence test checks the gene-specific path against an independent per-gene reference on both the numpy and dask backends, and passes on origin/main too - the old code computed the same numbers, only slowly. --- tests/test_sample_variance_graph_scaling.py | 219 ++++++++++++++++++++ 1 file changed, 219 insertions(+) create mode 100644 tests/test_sample_variance_graph_scaling.py 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) + + From 04c857443867d70c10ed3d509fe2a58a78392d7a Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 15:58:49 -0700 Subject: [PATCH 14/28] fix(provenance): resolve the sha in a linked worktree on a branch (#28, #33) Branch refs and packed-refs live in the common dir named by the worktree's commondir file; resolution looked only in the per-worktree git dir and returned None. Look up refs in both, follow symbolic refs a bounded number of steps, and log a warning when a sha cannot be resolved inside a checkout instead of stamping None silently. Fixes #28 Fixes #33 --- kompot/_provenance.py | 80 ++++++++++++++++++++-------- tests/test_provenance.py | 111 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 170 insertions(+), 21 deletions(-) diff --git a/kompot/_provenance.py b/kompot/_provenance.py index 06f32cc..fe04bd6 100644 --- a/kompot/_provenance.py +++ b/kompot/_provenance.py @@ -65,35 +65,63 @@ 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: - return None - - 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 - - # Packed ref. - packed = _read_text(os.path.join(git_dir, "packed-refs")) - if packed: +def _common_dir(git_dir: str) -> str: + """The directory holding the refs shared by every work tree of a repository. + + 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()) @@ -152,6 +180,16 @@ def _resolve() -> Dict[str, Any]: git_dir = _find_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/tests/test_provenance.py b/tests/test_provenance.py index 493305e..775a1e7 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -82,6 +82,117 @@ 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") + + +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]) From 74e62c6e042b4ee79f144d67b408f2dfa9884b54 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 15:58:49 -0700 Subject: [PATCH 15/28] fix(cli): write output under anndata 0.11/0.12 with nullable strings (#23) anndata 0.11 and 0.12 refuse to write a pandas StringArray unless allow_write_nullable_strings is set, and pandas 3 produces one for every string column, so every CLI write failed. One write_output helper now opts in for the duration of the write only; a no-op where the setting is absent. Fixes #23 --- kompot/cli/da.py | 17 ++++----- kompot/cli/de.py | 17 ++++----- kompot/cli/dm.py | 17 ++++----- kompot/cli/smooth.py | 17 ++++----- kompot/cli/utils.py | 36 +++++++++++++++++++ tests/test_cli_write_output.py | 65 ++++++++++++++++++++++++++++++++++ 6 files changed, 129 insertions(+), 40 deletions(-) create mode 100644 tests/test_cli_write_output.py 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 a0aae0e..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 @@ -441,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/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/tests/test_cli_write_output.py b/tests/test_cli_write_output.py new file mode 100644 index 0000000..833f866 --- /dev/null +++ b/tests/test_cli_write_output.py @@ -0,0 +1,65 @@ +"""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, which +is what makes this test meaningful on every pandas the package supports. +""" + +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") From e6ddfc7e2071812265d6123b968fa23757cac51f Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 15:58:50 -0700 Subject: [PATCH 16/28] fix: deprecate the no-op SampleVarianceEstimator(dask_num_workers) (#31) It set pool.num-workers, which no scheduler reads, mutated Dask's global config and logged a limit that was never applied. The covariance tensor is evaluated one gene at a time, so there is no pool to bound. Warn instead, and touch no configuration. Fixes #31 --- .../differential/sample_variance_estimator.py | 46 ++++++------- tests/test_dask_num_workers.py | 65 +++++++++++++++++++ 2 files changed, 88 insertions(+), 23 deletions(-) create mode 100644 tests/test_dask_num_workers.py diff --git a/kompot/differential/sample_variance_estimator.py b/kompot/differential/sample_variance_estimator.py index ac61924..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 @@ -89,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 @@ -471,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 @@ -506,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/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) From da7979c4e466f3c1d17ea27ef2d222de04658b48 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 16:03:35 -0700 Subject: [PATCH 17/28] docs: fail the build when RST markup survives into the HTML (#30) Nested inline markup renders as literal backticks and sphinx exits 0. A build-finished check scans every rendered page for unrendered literals, roles and directives and raises. It found ten live instances, all in notebook markdown where pandoc turns code-inside-bold or code-inside-link into nested RST; those are rewritten. Notebook 03 also carries the ~200-gene recommendation from the next commit. Fixes #30 --- docs/source/conf.py | 6 ++ docs/source/markup_leak_check.py | 101 ++++++++++++++++++ .../02_differential_expression_detailed.ipynb | 8 +- examples/03_sample_variance.ipynb | 10 +- examples/04_differential_abundance.ipynb | 2 +- examples/05_smooth_expression.ipynb | 6 +- tests/test_docs_markup_leak_check.py | 73 +++++++++++++ 7 files changed, 193 insertions(+), 13 deletions(-) create mode 100644 docs/source/markup_leak_check.py create mode 100644 tests/test_docs_markup_leak_check.py 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/markup_leak_check.py b/docs/source/markup_leak_check.py new file mode 100644 index 0000000..11a4545 --- /dev/null +++ b/docs/source/markup_leak_check.py @@ -0,0 +1,101 @@ +"""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 ``. +""" + +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]*``)?"), + "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/examples/02_differential_expression_detailed.ipynb b/examples/02_differential_expression_detailed.ipynb index 5e47400..a8808ab 100644 --- a/examples/02_differential_expression_detailed.ipynb +++ b/examples/02_differential_expression_detailed.ipynb @@ -151,17 +151,17 @@ "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, 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", + "- `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\n", + "- `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", + "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", diff --git a/examples/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index a8cf71a..d6d5a31 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -335,17 +335,17 @@ "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 **1 000** is the working recommendation;\n", - "2. **`store_arrays_on_disk=True`** (`StorageSettings`) \u2014 Kompot then assembles one gene's\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", + "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", - "**Always check the plan first with `dry_run=True`**, which reports memory, disk, and the\n", + "**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" ] }, @@ -1044,7 +1044,7 @@ "\n", "| Lever | Memory | Per-gene compute | Practical setting |\n", "|---|---|---|---|\n", - "| `genes=` | **linear** | **linear** | top ~1 000 from pass 1 |\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", diff --git a/examples/04_differential_abundance.ipynb b/examples/04_differential_abundance.ipynb index c3d8202..e720d5f 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", diff --git a/examples/05_smooth_expression.ipynb b/examples/05_smooth_expression.ipynb index dcb4c9d..5c8f4b0 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)" ] @@ -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/tests/test_docs_markup_leak_check.py b/tests/test_docs_markup_leak_check.py new file mode 100644 index 0000000..fba7e0e --- /dev/null +++ b/tests/test_docs_markup_leak_check.py @@ -0,0 +1,73 @@ +"""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) +def test_nested_literal_in_bold_is_a_leak(html): + assert [kind for kind, _ in find_leaks(html)] == ["double-backtick literal"] + + +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 From 8221782c64e62749d3368d9cfc9646f005a46308 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 16:03:36 -0700 Subject: [PATCH 18/28] docs: recommend ~200 genes, not ~1 000, for the sample-variance pass Per the maintainer's review on #29. Every suggestion of a gene budget for the second pass now says ~200, and the figures stated alongside it are recomputed for 200: the timing column is the measured per-gene time x 200, and the n_landmarks table is re-run from dry_run=True on the same synthetic 20 000 x 20 000 input, whose 200- and 1 000-gene rows reproduce the published plan table exactly. --- README.md | 2 +- docs/source/cli.rst | 2 +- docs/source/resource_planning.rst | 59 ++++++++++---------- docs/source/simplified.rst | 2 +- kompot/anndata/differential_expression.py | 2 +- kompot/cli/templates/de_config_template.yaml | 2 +- 6 files changed, 35 insertions(+), 34 deletions(-) diff --git a/README.md b/README.md index 9053f0b..5412696 100644 --- a/README.md +++ b/README.md @@ -89,7 +89,7 @@ So run it in two passes, and price the second one first: kompot.de(adata, "condition", "Young", "Old") mahal = "kompot_de_Young_to_Old_mahalanobis" -top_genes = adata.var.sort_values(mahal, ascending=False).head(1000).index +top_genes = adata.var.sort_values(mahal, ascending=False).head(200).index # Pass 2 — sample variance, restricted to the top genes plan = kompot.de( diff --git a/docs/source/cli.rst b/docs/source/cli.rst index 42a8867..5e0f7d3 100644 --- a/docs/source/cli.rst +++ b/docs/source/cli.rst @@ -261,7 +261,7 @@ Example: Complete Analysis # sample_variance.yaml sample_col: "Sample" - genes: ["GATA1", "KLF1", "..."] # top ~1 000 from the first pass + 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 diff --git a/docs/source/resource_planning.rst b/docs/source/resource_planning.rst index f208d64..e79a299 100644 --- a/docs/source/resource_planning.rst +++ b/docs/source/resource_planning.rst @@ -22,7 +22,7 @@ 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 1 000; see :ref:`how-many-genes`), offload + 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. @@ -35,7 +35,7 @@ instead: # ---- Pick the genes worth refining --------------------------------- mahal = "kompot_de_Young_to_Old_mahalanobis" - top_genes = adata.var.sort_values(mahal, ascending=False).head(1000).index + top_genes = adata.var.sort_values(mahal, ascending=False).head(200).index # ---- Pass 2: sample variance, restricted and offloaded ------------- kompot.de( @@ -124,8 +124,8 @@ covariance matrix is ``5000^2 x 8 B = 190 MiB``, so: :class: warning At the default 5 000 landmarks, **every gene added to an in-memory - sample-variance run costs about 0.37 GiB**. One thousand genes is roughly - 370 GiB; the whole transcriptome is over seven terabytes. Nothing else in a + 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 @@ -159,23 +159,23 @@ Measured inside a real ``kompot.de`` sample-variance run, 20 genes, with * - ``n_landmarks`` - per gene - - 1 000 genes + - 200 genes - 20 000 genes * - 500 - 0.016 s - - 16 s + - 3 s - 5 min * - 1 000 - 0.062 s - - 1 min + - 12 s - 21 min * - 2 000 - 0.246 s - - 4 min + - 49 s - 1.4 h * - 5 000 (default) - 2.08 s - - 35 min + - 7 min - 11.6 h .. warning:: @@ -224,7 +224,7 @@ 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, 1 000 genes and 6 donors this is 0.45 GiB. + 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``. @@ -269,12 +269,12 @@ not just the covariance term. (Kompot's report labels these binary magnitudes Three things to read off this. -**The in-memory column is the one that explodes.** 1 000 genes asks for -380 GiB, and the whole transcriptome for seven and a half terabytes, because +**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 1 000 genes plans at 7.7 GiB, barely -above the 6.4 GiB of a run with no sample variance at all, because the tensors +**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 @@ -398,8 +398,9 @@ stderr, and exits non-zero when the plan is infeasible. See :doc:`cli`. 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 — 1 000 real genes - plus 2 000 null genes is a 3 000-gene tensor. + 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 @@ -586,7 +587,7 @@ Other levers ------------ ``n_landmarks`` **is the strongest memory lever**, because the covariance term -is quadratic in it. Same 1 000 genes, from the dry run: +is quadratic in it. Same 200 genes, from the dry run: .. list-table:: :header-rows: 1 @@ -596,17 +597,17 @@ is quadratic in it. Same 1 000 genes, from the dry run: - sample variance, in memory - with ``store_arrays_on_disk`` * - 1 000 - - 21.0 GiB - - 6.1 GiB + - 5.4 GiB + - 2.5 GiB * - 2 000 - - 66.0 GiB - - 6.4 GiB + - 14.6 GiB + - 2.6 GiB * - 3 000 - - 140.8 GiB - - 6.7 GiB + - 29.7 GiB + - 2.9 GiB * - 5 000 (default) - - 380.2 GiB - - 7.7 GiB + - 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: @@ -658,7 +659,7 @@ 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 1 000 genes is a recommendation chosen from the cost curves, not a +**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: @@ -671,14 +672,14 @@ magnitude: 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 ~1 000 genes is still the right default, for reasons +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; -* 1 000 genes costs about a twentieth of the whole transcriptome on both axes, +* 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 000 genes buy no additional + 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 diff --git a/docs/source/simplified.rst b/docs/source/simplified.rst index b476f52..fb4979f 100644 --- a/docs/source/simplified.rst +++ b/docs/source/simplified.rst @@ -40,7 +40,7 @@ Differential Expression kompot.de( adata, "condition", "Young", "Old", sample_col="donor_id", - genes=top_genes, # e.g. top 1 000 from a previous run + genes=top_genes, # e.g. top 200 from a previous run fdr=kompot.FDRSettings(null_genes=0), ) diff --git a/kompot/anndata/differential_expression.py b/kompot/anndata/differential_expression.py index d74a1a2..802a65c 100644 --- a/kompot/anndata/differential_expression.py +++ b/kompot/anndata/differential_expression.py @@ -103,7 +103,7 @@ def de( 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(1000).index + 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, diff --git a/kompot/cli/templates/de_config_template.yaml b/kompot/cli/templates/de_config_template.yaml index 531c83b..2a7da4e 100644 --- a/kompot/cli/templates/de_config_template.yaml +++ b/kompot/cli/templates/de_config_template.yaml @@ -72,7 +72,7 @@ overwrite: null # Overwrite existing results (null = warn) # Advanced filtering: genes: null # List of gene names to analyze (null = all). # The lever for a sample_col run: restrict to - # the top ~1000 genes of a first pass. + # the top ~200 genes of a first pass. min_cells: 2 # Minimum cells per gene min_percentage: null # Minimum percentage of cells expressing From 16b26438e861b8715fabc526c4ac4e88d04160ae Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 16:03:44 -0700 Subject: [PATCH 19/28] docs(changelog): 0.9.0 entries for #23, #28/#33, #30, #31 and the 200-gene budget --- CHANGELOG.md | 40 ++++++++++++++++++++++++++++++++++------ 1 file changed, 34 insertions(+), 6 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 5592dd9..2c6ca2f 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -114,8 +114,9 @@ Following the fix above, `estimate_differential_expression_resources`: 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 genes — ranks the levers, and shows how to price a run - before committing to it. + 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`, @@ -134,7 +135,7 @@ Following the fix above, `estimate_differential_expression_resources`: Fixes settylab/kompot#27. -### Added +### 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 @@ -145,7 +146,7 @@ Fixes settylab/kompot#27. - `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. @@ -158,13 +159,40 @@ Fixes settylab/kompot#27. 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. + +### 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. + Its first run found ten such constructs in the tutorial notebooks, where pandoc turns code inside + bold or inside link text into nested RST; they are rewritten. ## [0.8.0] - 2026-07-28 From 151f9b865215b2ad58fd0ea78d82b272470c5a71 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 16:39:17 -0700 Subject: [PATCH 20/28] test(cli): write CLI fixtures through write_output (#23) Under anndata 0.11/0.12 with pandas 3 even the obs index is a StringArray, so the fixtures' own write_h5ad failed before the CLI ran: that, not the CLI, is where the 49 failures reported on #23 were raised. Writing the inputs through the CLI's writer keeps the global setting untouched, so the CLI tests still exercise the CLI's own opt-in. --- tests/test_cli.py | 42 +++++++++++++++++--------------- tests/test_cli_compute_config.py | 38 ++++++++++++++++------------- tests/test_cli_dm_coverage.py | 18 ++++++++------ tests/test_cli_smooth.py | 18 ++++++++------ 4 files changed, 66 insertions(+), 50 deletions(-) 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, From a308f94c695e987a6af3b8cb355dd6823c209803 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 17:22:34 -0700 Subject: [PATCH 21/28] test(docs): readable parametrize ids for the #30 leak cases --- tests/test_docs_markup_leak_check.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_docs_markup_leak_check.py b/tests/test_docs_markup_leak_check.py index fba7e0e..addb93e 100644 --- a/tests/test_docs_markup_leak_check.py +++ b/tests/test_docs_markup_leak_check.py @@ -22,7 +22,7 @@ ] -@pytest.mark.parametrize("html", NESTED) +@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"] From 533ab0fc8d543e3ee0189def0c2b5c25e29cf5e9 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 18:16:24 -0700 Subject: [PATCH 22/28] fix(#30, #28): catch code-in-link leaks; ignore a repository that is not kompot's own Skeptic review of a308f94. #30: Markdown [`x`](url) becomes nested RST that renders as `x `__; the inner literal is a tag, so tag-stripping hid it from the double-backtick pattern and the check passed while 14 such links were live (notebooks 01 x8, 03 x1, 04 x4, 05 x1). Add a pattern for the leaked anonymous-hyperlink tail and rewrite the 14 links. Measured on the base build with the final checker: 27 constructs in 5 pages, all fixed. #28: _find_git_dir accepts any enclosing .git, so a kompot tree vendored into an unrelated repository stamped that repository's HEAD -- a wrong sha, worse than None. Provenance now accepts a git dir only when its .git sits at the package's parent, kompot's own repository root. Also: drop a trailing space in the DE config template, and state what the nullable-string write test has and has not been run against. --- docs/source/markup_leak_check.py | 8 ++++ examples/01_getting_started.ipynb | 16 ++++---- examples/03_sample_variance.ipynb | 2 +- examples/04_differential_abundance.ipynb | 8 ++-- examples/05_smooth_expression.ipynb | 2 +- kompot/_provenance.py | 18 ++++++++- kompot/cli/templates/de_config_template.yaml | 2 +- tests/test_cli_write_output.py | 5 ++- tests/test_docs_markup_leak_check.py | 17 ++++++++ tests/test_provenance.py | 42 ++++++++++++++++++++ 10 files changed, 102 insertions(+), 18 deletions(-) diff --git a/docs/source/markup_leak_check.py b/docs/source/markup_leak_check.py index 11a4545..c14e3da 100644 --- a/docs/source/markup_leak_check.py +++ b/docs/source/markup_leak_check.py @@ -10,6 +10,10 @@ 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 @@ -21,6 +25,10 @@ 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|" 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/03_sample_variance.ipynb b/examples/03_sample_variance.ipynb index d6d5a31..e3eeaec 100644 --- a/examples/03_sample_variance.ipynb +++ b/examples/03_sample_variance.ipynb @@ -358,7 +358,7 @@ "> 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", + "> 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", diff --git a/examples/04_differential_abundance.ipynb b/examples/04_differential_abundance.ipynb index e720d5f..5b1dab3 100644 --- a/examples/04_differential_abundance.ipynb +++ b/examples/04_differential_abundance.ipynb @@ -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 5c8f4b0..2f3e86f 100644 --- a/examples/05_smooth_expression.ipynb +++ b/examples/05_smooth_expression.ipynb @@ -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", diff --git a/kompot/_provenance.py b/kompot/_provenance.py index fe04bd6..9feefe6 100644 --- a/kompot/_provenance.py +++ b/kompot/_provenance.py @@ -65,6 +65,22 @@ def _find_git_dir(start: str) -> Optional[str]: current = parent +def _find_kompot_git_dir(package_dir: str) -> Optional[str]: + """The git dir of *Kompot's own* checkout, or ``None``. + + ``_find_git_dir`` walks up to *any* enclosing ``.git``. A Kompot source + tree vendored or copied into an unrelated repository would then be + stamped with that repository's HEAD -- a wrong sha, which is worse than + none. Kompot's repository root is the parent of the package directory, + so the ``.git`` entry (a directory, or a file for a worktree or + submodule) must sit exactly there. + """ + 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) + + def _common_dir(git_dir: str) -> str: """The directory holding the refs shared by every work tree of a repository. @@ -177,7 +193,7 @@ 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: diff --git a/kompot/cli/templates/de_config_template.yaml b/kompot/cli/templates/de_config_template.yaml index 2a7da4e..2d4a5b7 100644 --- a/kompot/cli/templates/de_config_template.yaml +++ b/kompot/cli/templates/de_config_template.yaml @@ -72,7 +72,7 @@ overwrite: null # Overwrite existing results (null = warn) # Advanced filtering: 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. + # 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/tests/test_cli_write_output.py b/tests/test_cli_write_output.py index 833f866..0681014 100644 --- a/tests/test_cli_write_output.py +++ b/tests/test_cli_write_output.py @@ -4,8 +4,9 @@ ``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, which -is what makes this test meaningful on every pandas the package supports. +``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 diff --git a/tests/test_docs_markup_leak_check.py b/tests/test_docs_markup_leak_check.py index addb93e..e23ea86 100644 --- a/tests/test_docs_markup_leak_check.py +++ b/tests/test_docs_markup_leak_check.py @@ -27,6 +27,21 @@ 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"] @@ -71,3 +86,5 @@ def build(body, name): 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_provenance.py b/tests/test_provenance.py index 775a1e7..99e7d25 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -167,6 +167,48 @@ def test_resolve_sha_matches_git_rev_parse(tmp_path, shape, pack): 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 copied into another repository must not stamp that repository's 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 (Kompot's own repository root) counts. + """ + 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) From da038c532ba448aea32b0b93d75039001a990ea4 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 18:16:24 -0700 Subject: [PATCH 23/28] docs(changelog): correct the #30 counts; note the vendored-tree provenance guard --- CHANGELOG.md | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 2c6ca2f..ea90de7 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -171,6 +171,8 @@ Fixes settylab/kompot#27. 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. + Provenance also ignores a repository that is not Kompot's own: a Kompot source tree vendored + into another repository no longer stamps that repository's commit as `kompot_git_sha`. ### Deprecated @@ -191,8 +193,10 @@ Fixes settylab/kompot#27. (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. - Its first run found ten such constructs in the tutorial notebooks, where pandoc turns code inside - bold or inside link text into nested RST; they are rewritten. + 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 From eac37dffdf55a487f204769f5a9d71ec70432a11 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 18:50:21 -0700 Subject: [PATCH 24/28] test(provenance): the source-checkout test requires kompot's own checkout Its skip guard used the unguarded walk-up, so a kompot tree exported or vendored inside another work tree ran the test and, before the root guard, passed it by stamping the enclosing repository's sha: the defect the guard fixes, read as a pass. It now skips unless kompot has its own .git. --- tests/test_provenance.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/tests/test_provenance.py b/tests/test_provenance.py index 99e7d25..9869c07 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -439,8 +439,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") + # Kompot's OWN checkout, not merely an enclosing work tree: a tree exported + # or vendored inside another repository must skip here, not borrow its sha. + 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"] From 9cf061f375212e5844dba10c6513be7eef2ea1ee Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 25 Sep 2026 20:16:26 -0700 Subject: [PATCH 25/28] docs(#28): narrow the vendored-tree guard's claim to trees nested below a root A kompot/ package copied to the root of another repository still stamps that repository's sha: by position it cannot be told apart from kompot's own checkout. Say so in the CHANGELOG and the docstrings. Text only. --- CHANGELOG.md | 5 +++-- kompot/_provenance.py | 16 ++++++++++------ tests/test_provenance.py | 9 +++++---- 3 files changed, 18 insertions(+), 12 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ea90de7..4c12684 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -171,8 +171,9 @@ Fixes settylab/kompot#27. 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. - Provenance also ignores a repository that is not Kompot's own: a Kompot source tree vendored - into another repository no longer stamps that repository's commit as `kompot_git_sha`. + 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 diff --git a/kompot/_provenance.py b/kompot/_provenance.py index 9feefe6..39e8525 100644 --- a/kompot/_provenance.py +++ b/kompot/_provenance.py @@ -66,14 +66,18 @@ def _find_git_dir(start: str) -> Optional[str]: def _find_kompot_git_dir(package_dir: str) -> Optional[str]: - """The git dir of *Kompot's own* checkout, or ``None``. + """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 or copied into an unrelated repository would then be - stamped with that repository's HEAD -- a wrong sha, which is worse than - none. Kompot's repository root is the parent of the package directory, - so the ``.git`` entry (a directory, or a file for a worktree or - submodule) must sit exactly there. + 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")): diff --git a/tests/test_provenance.py b/tests/test_provenance.py index 9869c07..0bc0084 100644 --- a/tests/test_provenance.py +++ b/tests/test_provenance.py @@ -169,11 +169,12 @@ def test_resolve_sha_matches_git_rev_parse(tmp_path, shape, pack): @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 copied into another repository must not stamp that repository's HEAD. + """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 (Kompot's own repository root) counts. + parent counts. (A copy at another repository's root is not caught; by + position it looks like Kompot's own checkout.) """ import shutil @@ -439,8 +440,8 @@ 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__)) - # Kompot's OWN checkout, not merely an enclosing work tree: a tree exported - # or vendored inside another repository must skip here, not borrow its sha. + # 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") From 301139faac07c1c1fbc8393ce76838e8ef928577 Mon Sep 17 00:00:00 2001 From: Dominik Date: Sun, 4 Oct 2026 18:52:17 -0700 Subject: [PATCH 26/28] fix(#34): snap landmarks to the exact nearest cell find_landmarks snapped each cluster centroid to a cell with the approximate pynndescent index built for the graph. A centroid usually sits off the data, where that search is least reliable: at ~1M cells in 100 dimensions, 27% of centroids snapped to a cell that was not the nearest. The snap is now an exact Euclidean search, chunked over the cells (about 64 MiB per working array). A matrix-product score shortlists every cell within a rounding bound of the minimum, and the shortlist is ranked by the direct float64 sum, so the answer equals a brute-force argmin; exact ties go to the lowest index. About 2-4 s for 150 centroids against 1M x 100. exact_snap=False keeps the old index query, for reproducing 0.8.0 landmarks. --- CHANGELOG.md | 24 +++++ kompot/utils.py | 132 ++++++++++++++++++++++++++- tests/test_landmark_snap.py | 173 ++++++++++++++++++++++++++++++++++++ 3 files changed, 326 insertions(+), 3 deletions(-) create mode 100644 tests/test_landmark_snap.py diff --git a/CHANGELOG.md b/CHANGELOG.md index 4c12684..253a964 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -142,6 +142,30 @@ Fixes settylab/kompot#27. `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: at ~1M cells in 100 dimensions, 27% of +centroids snapped to a cell that was not the nearest (median distance ratio +1.04, max 1.40). On ~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). Exact ties, such as +duplicate cells, go to the lowest cell index, so the result is deterministic. +It costs about 2–4 s for 150 centroids against 1M cells in 100 dimensions. + +**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"`. diff --git a/kompot/utils.py b/kompot/utils.py index fa12f6b..4b3e135 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -721,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, @@ -728,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. @@ -755,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 (at ~1M cells in 100 dimensions, 27% of + centroids). Set False only to reproduce landmarks from earlier versions. Returns ------- @@ -785,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/tests/test_landmark_snap.py b/tests/test_landmark_snap.py new file mode 100644 index 0000000..46e63ec --- /dev/null +++ b/tests/test_landmark_snap.py @@ -0,0 +1,173 @@ +"""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_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,) From 729cbc1ca031109f1d1a5cdeaeec2319e2737a45 Mon Sep 17 00:00:00 2001 From: Dominik Date: Sun, 4 Oct 2026 18:57:06 -0700 Subject: [PATCH 27/28] docs(changelog, #34): landmark snap cost measured on a dedicated 8-core node (1.2 s at 1M x 100) --- CHANGELOG.md | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 253a964..1dcf09d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -154,7 +154,8 @@ centroids snapped to a cell that was not the nearest (median distance ratio The snap is now an exact Euclidean search, chunked over the cells so that memory stays bounded (about 64 MiB per working array). Exact ties, such as duplicate cells, go to the lowest cell index, so the result is deterministic. -It costs about 2–4 s for 150 centroids against 1M cells in 100 dimensions. +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 From 28542012cf3f25ba6ff8e0139d304d47d79393d7 Mon Sep 17 00:00:00 2001 From: Dominik Date: Sun, 4 Oct 2026 19:57:36 -0700 Subject: [PATCH 28/28] test, docs(#34): guard the rounding bound; cite reproducible snap figures Review follow-up. Three helper tests (near-duplicate cells, a query at the midpoint of two cells, a tiny spread far from the origin) now fail if the shortlist slack is removed, if the in-loop filter drops it, or if the shortlist is ranked by the matrix-product score instead of the direct sum. Those three mutants previously passed every test. The changelog and the exact_snap docstring now lead with the synthetic 1M x 100 measurement (36 of 150 centroids, 24%), and label the 27% figure as one real dataset. The changelog notes that the near-tie shortlist is outside the 64 MiB working-array budget. --- CHANGELOG.md | 14 +++++++++----- kompot/utils.py | 4 ++-- tests/test_landmark_snap.py | 33 +++++++++++++++++++++++++++++++++ 3 files changed, 44 insertions(+), 7 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1dcf09d..ba78d57 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -146,13 +146,17 @@ Fixes settylab/kompot#27. `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: at ~1M cells in 100 dimensions, 27% of -centroids snapped to a cell that was not the nearest (median distance ratio -1.04, max 1.40). On ~75k cells in 19 dimensions it was at most one centroid in -~130. +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). Exact ties, such as +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. diff --git a/kompot/utils.py b/kompot/utils.py index 4b3e135..c255fa4 100644 --- a/kompot/utils.py +++ b/kompot/utils.py @@ -876,8 +876,8 @@ def find_landmarks( 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 (at ~1M cells in 100 dimensions, 27% of - centroids). Set False only to reproduce landmarks from earlier versions. + 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 ------- diff --git a/tests/test_landmark_snap.py b/tests/test_landmark_snap.py index 46e63ec..92ca06b 100644 --- a/tests/test_landmark_snap.py +++ b/tests/test_landmark_snap.py @@ -162,6 +162,39 @@ def test_far_from_origin(self): 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):