[settylab-dotto-bot]

DO NOT MERGE — for the maintainer's review.

What

smooth_expression now handles cells whose cell-state coordinates are not all finite (NaN or inf, e.g. cells left out of a diffusion map). They are excluded from fitting and come out NaN in every output layer. No argument is needed. Every other cell's values are unchanged.

The run info records params["n_cells_nonfinite"] and params["nonfinite_cells"], and each output layer's description gives the NaN count. With all coordinates finite, nothing changes.

Why

On main (915b411), a NaN coordinate makes the fit fail:

• kompot/anndata/smooth.py L190-191 put every selected cell into the training set;
• kompot/differential/expression_model.py L233-239 run landmark selection (mellon compute_landmarks → k_means), which refuses NaN.

Prediction, by contrast, is per cell (mellon LandmarksConditional: mean K(x, landmarks) @ weights; the variance is solved against the fixed landmark Cholesky factor). So excluding these cells from the fit is the whole fix, and their rows come out NaN on their own. Measured: every other row of predict and std is bitwise equal to evaluating without the NaN rows, at batch sizes 7, 50 and 1000.

(An earlier revision of this PR added a cell_mask argument; it was removed on review.)

Tests

• tests/test_anndata_smooth.py::TestNonFiniteCoordinates, 10 tests, including a bitwise comparison against the same fit.
• tests/test_cli_smooth.py::TestCLISmoothNonFinite: a CLI round trip bitwise equal to the Python API.
• All synthetic.
• Locally, the two smoothing test files: 31 passed. The full suite runs in this repo's CI.

[katosh]

@settylab-dotto-bot the anndata interface should take an .obs column as mask and document this appropriately in the run info. The CLI should expose the option.

[settylab-dotto-bot]

Done, on this PR (commits 4d1ddf7 and 4a79a5e; still DO NOT MERGE).

.obs column as the mask. smooth_expression(..., cell_mask="<column>") takes the name of a boolean adata.obs column; True keeps a cell, and missing values are not allowed. The array form that this PR first added is removed. It was never released, and no lower-level function takes a mask instead: ExpressionModel is unchanged, and the masking lives in the anndata interface only. Refused with a clear error:

• a missing column (the error lists the available columns);
• a non-boolean or partly missing column;
• an all-False column;
• a non-string argument (TypeError).

Run info. uns[result_key]["last_run_info"]["params"] now records:

• cell_mask: the column name;
• n_cells_masked_out: the number of masked cells;
• masked_cells: "cells with adata.obs[''] == False are excluded from fitting, never evaluated, and NaN in every output layer of this run".

Each output layer's description in field_mapping ends with "; NaN for the N cells masked out by obs['']". Without a mask the three fields are None, 0 and None.

CLI. kompot smooth --cell-mask COLUMN, documented in --help. The config template has a cell_mask: key, and the log names the mask.

Tests (synthetic fixtures; 22 for the mask: 18 anndata + 4 CLI):

• bitwise: kept rows equal groupby=<column>, condition=True, which is the same fit; only the masked rows differ, NaN instead of an extrapolation;
• bitwise: a nullable boolean column gives the same result as a numpy bool column;
• refused: missing column; string, integer or partly missing column; all-False column; array argument;
• run info: all three fields and the layer descriptions; and their values without a mask;
• CLI: --help documents the option;
• CLI round-trip: NaN rows and run info are written, and the result is bitwise equal to the Python API called in-process. A missing column and a non-boolean column exit non-zero and write no output.

Full suite, this repo's CI on head 4a79a5e, 2,261 tests each, 0 failed:

• Python 3.10: 2,223 passed, 38 skipped;
• Python 3.11: 2,225 passed, 36 skipped;
• Python 3.12: 2,225 passed, 36 skipped.

The PR description is updated to match.

[katosh]

@settylab-dotto-bot reconsider this. The predict and std functions should predict all cells independently. So, cells with invalid coordinates should have NaN values returned. Is .predict and .std pure JAX? We should not need to pass a cell mask at all in order to ignore cells with NaN coordinates.

[settylab-dotto-bot]

Answer: you are right. No mask is needed, and cell_mask is gone. predict and std evaluate each cell independently. What breaks on main is the fit: cells with NaN coordinates reach landmark selection. This PR now excludes cells with non-finite coordinates from fitting only; they are still evaluated and come out NaN. Head fa26738f, still DO NOT MERGE.

Is it pure JAX and per cell? Yes, for the paths smooth_expression uses (mellon 1.7.2, as installed: kompot main's pyproject.toml L77 pins mellon>=1.7.1, which allows it. LandmarksConditional in mellon/conditional.py):

• mean, L651-658: mu + K(x_new, landmarks) @ weights, one row per cell;
• variance, L694-707: solve_triangular(L, K(landmarks, x_new)) against the fixed landmark Cholesky factor L, then a per-column sum. There is no factorisation or normalisation over the test points.
• kompot's ExpressionModel.predict/std batch the cells (apply_batched) and sum per-cell terms.
• mellon's input check (validate_array) converts to float and does not reject NaN.

Measured (synthetic: 400 cells, 5 dims, 6 genes, 40 landmarks; NaN coordinates planted in 7 cells; fit on the valid cells):

check	batch 7	batch 50	batch 1000
the 7 rows are NaN, exactly those, for predict and std	✓	✓	✓
every other row bitwise equal to evaluating without the 7 rows, predict and std	✓	✓	✓
control without NaN: all finite; repeat bitwise; valid rows equal full evaluation	✓	✓	✓
a row with ONE NaN coordinate comes back fully NaN	✓	✓	✓

All of it holds with and without the variance terms (use_empirical_variance=True with sample_indices). inf and -inf coordinates also come back NaN.

What fails on main (915b411), with 2 NaN rows in DM_EigenVectors:

• ValueError: Input X contains NaN. KMeans does not accept missing values …
• Path: kompot/anndata/smooth.py L190-191 put every selected cell into the training set (X_train = X_all[mask]); kompot/differential/expression_model.py L233-239 then call mellon's compute_landmarks, which runs k_means.

The change (kompot/anndata/smooth.py, +36/−5 against main):

• cells whose obsm coordinates are not all finite are dropped from the training mask;
• prediction is unchanged and still runs on all cells, so those rows come out NaN in every output layer;
• if no selected cell is finite, a ValueError is raised;
• run info gains params["n_cells_nonfinite"] and params["nonfinite_cells"] (the sentence above), and each output layer's description in field_mapping gives the NaN count;
• cell_mask, --cell-mask and the config key are removed; the CLI and template are back to main. The CLI needs nothing new: NaN coordinates in the input file behave the same way.

Tests (synthetic):

• TestNonFiniteCoordinates, 10 tests:
  • NaN rows are NaN and the rest finite;
  • bitwise: valid rows equal the same fit with the invalid cells given arbitrary finite coordinates;
  • equal to smoothing the valid subset alone;
  • partial-NaN and ±inf rows;
  • all-finite control;
  • run info;
  • condition, variance layers and gene subset;
  • no valid cell raises.
• TestCLISmoothNonFinite: a CLI round trip with NaN coordinates, bitwise equal to the Python API.
• tests/test_anndata_smooth.py + tests/test_cli_smooth.py: 31 passed locally. Full suite on CI at fa26738f (https://github.com/settylab/kompot/actions/runs/39915067924), 2,250 tests per Python, 0 failed: Python 3.10 2,212 passed / 38 skipped; 3.11 2,214 / 36; 3.12 2,214 / 36.