From b4d1c69efc09dcc3bcd7018b5f52a280f020b895 Mon Sep 17 00:00:00 2001 From: Dominik Date: Thu, 22 Oct 2026 20:41:48 -0700 Subject: [PATCH 1/4] smooth_expression: add cell_mask to mask cells out (NaN) instead of evaluating them smooth_expression fits on the selected cells but evaluates the model on every cell. Cells that have no place in the representation (e.g. removed by a quality gate before the diffusion map, with NaN coordinates) then get NaN by accident or an unrequested extrapolation. cell_mask (boolean obs column or array; True keeps a cell): masked-out cells are excluded from fitting, never evaluated (their obsm rows are never read), and NaN in every per-cell output (smoothed, std, obs/sample variance). With groupby/condition the fit uses condition AND mask; kept cells are evaluated as before. Kept cells with non-finite coordinates raise. cell_mask=None (default) is unchanged. Tests: 13 new in tests/test_anndata_smooth.py. --- CHANGELOG.md | 27 ++++++++ kompot/anndata/smooth.py | 89 +++++++++++++++++++++++--- tests/test_anndata_smooth.py | 119 +++++++++++++++++++++++++++++++++++ 3 files changed, 226 insertions(+), 9 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index ba78d57..8e61e69 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,33 @@ All notable changes to this project will be documented in this file. +## [Unreleased] + +### Added — `smooth_expression(..., cell_mask=...)`: mask cells out instead of extrapolating to them + +`smooth_expression` fits on the selected cells and evaluates the model on **every** +cell. That is right for conditions, but wrong for cells that have no place in the +cell-state representation at all: cells removed by a quality gate before the +diffusion map was computed typically carry NaN coordinates, and a smoothed value +for them is either NaN by accident or an extrapolation nobody asked for. + +`cell_mask` (a boolean `adata.obs` column name, or a boolean array; `True` keeps a +cell) makes the exclusion explicit: + + - masked-out cells are excluded from fitting and are **never evaluated**, so their + `obsm` coordinates are never read and may be NaN; + - every per-cell output (smoothed values, standard deviations, and the empirical + and sample variance layers when requested) is **NaN** for them; + - with `groupby`/`condition`, the model is fitted on cells in the condition *and* + in the mask, and evaluated on every kept cell; + - kept cells with non-finite coordinates raise a `ValueError` instead of + propagating NaN silently; + - the run record stores the mask (column name, or `"array"`) and the number of + cells masked out. + +The default, `cell_mask=None`, is unchanged. On kept cells the result equals +smoothing the kept subset on its own (tested). + ## [0.9.0] - 2026-09-15 ### Fixed — sample variance no longer costs more memory than it needs to diff --git a/kompot/anndata/smooth.py b/kompot/anndata/smooth.py index 8e07946..1987559 100644 --- a/kompot/anndata/smooth.py +++ b/kompot/anndata/smooth.py @@ -45,6 +45,7 @@ def smooth_expression( storage: "StorageSettings | None" = None, output: "OutputSettings | None" = None, model: "ModelSettings | None" = None, + cell_mask: Optional[Union[str, np.ndarray, pd.Series]] = None, **function_kwargs, ) -> Optional[Union[Dict[str, Any], Any]]: """Smooth gene expression for a single condition using GP regression. @@ -92,6 +93,17 @@ def smooth_expression( model : ModelSettings, optional Pre-fitted :class:`~kompot.ExpressionModel` to inject via ``model1``. When provided, skips internal fitting. + cell_mask : str or array-like of bool, optional + Cells to keep. Either the name of a boolean column in ``adata.obs`` + or a boolean array of length ``adata.n_obs``; ``True`` keeps a cell. + Cells with ``False`` are **masked out**: they are excluded from + fitting, they are never evaluated, their ``obsm`` coordinates are + never read (so they may be NaN, e.g. cells left out of a diffusion + map), and every per-cell output (smoothed values, standard + deviations and, when requested, the variance layers) is NaN for + them. Combined with ``groupby``/``condition``, the model is fitted + on cells in the condition *and* in the mask, and evaluated on every + cell in the mask. *None* (default) keeps every cell, as before. **function_kwargs Forwarded to :class:`mellon.FunctionEstimator`. @@ -143,6 +155,8 @@ def smooth_expression( if groupby is not None: validate_groupby(adata, groupby) + keep = _resolve_cell_mask(adata, cell_mask) + # --- Build cell mask --- if groupby is not None and condition is not None: mask = (adata.obs[groupby] == condition).values @@ -155,6 +169,15 @@ def smooth_expression( mask = np.ones(adata.n_obs, dtype=bool) logger.info(f"Smoothing all {mask.sum():,} cells") + if keep is not None: + mask = mask & keep + if mask.sum() == 0: + raise ValueError("No cell is both in the selected condition and kept by cell_mask.") + logger.info( + f"cell_mask: {int((~keep).sum()):,} cells masked out (NaN in every output); " + f"fitting on {mask.sum():,} cells" + ) + # --- Build field names --- cond_label = _sanitize_name(condition) if condition else "all" field_names = _smooth_field_names( @@ -189,6 +212,13 @@ def smooth_expression( # --- Extract data --- X_all = adata.obsm[obsm_key] X_train = X_all[mask] + # Masked-out cells are never evaluated, so their coordinates are never read. + X_eval = X_all if keep is None else X_all[keep] + if keep is not None and not np.all(np.isfinite(np.asarray(X_eval))): + raise ValueError( + f"adata.obsm['{obsm_key}'] has non-finite values in cells kept by " + "cell_mask; mask those cells out or fix the representation." + ) if layer is not None: if layer not in adata.layers: @@ -248,12 +278,22 @@ def smooth_expression( # --- Predict on all cells --- # The model is trained on the condition cells but evaluated everywhere, - # so every cell gets a smoothed value and uncertainty estimate. - imputed = np.asarray( - expr_model.predict(X_all, batch_size=batch_size, progress=progress) + # so every cell gets a smoothed value and uncertainty estimate. With a + # cell_mask, "everywhere" means every kept cell; masked-out rows are NaN. + def _expand_cells(values): + """Place kept-cell rows into a full (n_obs, ...) array, NaN elsewhere.""" + values = np.asarray(values) + if keep is None: + return values + out = np.full((adata.n_obs,) + values.shape[1:], np.nan, dtype=np.result_type(values, np.float32)) + out[keep] = values + return out + + imputed = _expand_cells( + expr_model.predict(X_eval, batch_size=batch_size, progress=progress) ) - std_vals = np.asarray( - expr_model.std(X_all, batch_size=batch_size, progress=progress) + std_vals = _expand_cells( + expr_model.std(X_eval, batch_size=batch_size, progress=progress) ) # Broadcast (n, 1) to (n, n_genes) when only GP covariance is present n_genes_fitted = imputed.shape[1] @@ -288,8 +328,8 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["std_key"]] = std_vals if use_empirical_variance: - obs_var = np.asarray( - expr_model.obs_variance(X_all, batch_size=batch_size, progress=progress) + obs_var = _expand_cells( + expr_model.obs_variance(X_eval, batch_size=batch_size, progress=progress) ) if genes is not None: adata.layers[field_names["obs_variance_key"]] = _expand_genes( @@ -299,9 +339,9 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["obs_variance_key"]] = obs_var if sample_col is not None and expr_model.has_sample_variance: - sam_var = np.asarray( + sam_var = _expand_cells( expr_model.sample_variance( - X_all, diag=True, batch_size=batch_size, progress=progress + X_eval, diag=True, batch_size=batch_size, progress=progress ) ) if isinstance(sam_var, np.ndarray): @@ -353,6 +393,9 @@ def _expand_genes(values, n_cells, n_genes): "sample_col": sample_col, "n_genes": len(selected_genes), "n_cells": int(mask.sum()), + "cell_mask": (None if cell_mask is None else + cell_mask if isinstance(cell_mask, str) else "array"), + "n_cells_masked_out": 0 if keep is None else int((~keep).sum()), "batch_size": batch_size, "result_key": result_key, }, @@ -381,6 +424,34 @@ def _expand_genes(values, n_cells, n_genes): return None +def _resolve_cell_mask(adata, cell_mask): + """Return a boolean array of kept cells, or None when no mask is given.""" + if cell_mask is None: + return None + if isinstance(cell_mask, str): + if cell_mask not in adata.obs: + raise ValueError(f"cell_mask column '{cell_mask}' not found in adata.obs.") + col = adata.obs[cell_mask] + if not pd.api.types.is_bool_dtype(col) or col.isna().any(): + raise ValueError( + f"cell_mask column '{cell_mask}' must be boolean without missing values " + f"(got dtype {col.dtype})." + ) + keep = col.to_numpy(dtype=bool) + else: + arr = np.asarray(cell_mask) + if arr.dtype != bool: + raise ValueError(f"cell_mask must be a boolean array (got dtype {arr.dtype}).") + keep = arr.reshape(-1) + if keep.shape != (adata.n_obs,): + raise ValueError( + f"cell_mask has {keep.shape[0]} entries but adata has {adata.n_obs} observations." + ) + if not keep.any(): + raise ValueError("cell_mask keeps no cell.") + return keep + + def compute_smoothed_expression( adata, groupby: Optional[str] = None, diff --git a/tests/test_anndata_smooth.py b/tests/test_anndata_smooth.py index f513014..0f58a0e 100644 --- a/tests/test_anndata_smooth.py +++ b/tests/test_anndata_smooth.py @@ -168,3 +168,122 @@ def test_deprecated_compute_smoothed_expression(self, adata): progress=False, ) assert "kompot_smooth_all_smoothed" in adata.layers + + +GP_SEEDED = GPSettings(n_landmarks=30, random_state=0) + + +class TestCellMask: + """cell_mask: masked-out cells are excluded from fitting, never evaluated, and NaN in every output.""" + + @staticmethod + def _keep(adata): + keep = np.ones(adata.n_obs, dtype=bool) + keep[[0, 7, 41, 79]] = False + return keep + + def test_masked_rows_nan_kept_rows_finite(self, adata): + keep = self._keep(adata) + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + ) + for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): + layer = adata.layers[key] + assert np.all(np.isnan(layer[~keep])) + assert np.all(np.isfinite(layer[keep])) + + def test_masked_coordinates_are_never_read(self, adata): + keep = self._keep(adata) + adata.obsm["DM_EigenVectors"][~keep] = np.nan # e.g. cells left out of the diffusion map + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + ) + layer = adata.layers["kompot_smooth_all_smoothed"] + assert np.all(np.isnan(layer[~keep])) and np.all(np.isfinite(layer[keep])) + + def test_equals_smoothing_the_kept_subset(self, adata): + keep = self._keep(adata) + sub = adata[keep].copy() + kompot.smooth_expression(sub, gp=GP_SEEDED, output=OutputSettings(progress=False)) + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + ) + for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): + np.testing.assert_allclose(adata.layers[key][keep], sub.layers[key], rtol=1e-6, atol=1e-8) + + def test_obs_column_and_condition(self, adata): + keep = self._keep(adata) + adata.obs["in_state_space"] = keep + kompot.smooth_expression( + adata, + groupby="condition", + condition="A", + cell_mask="in_state_space", + gp=GP_SEEDED, + output=OutputSettings(progress=False), + ) + layer = adata.layers["kompot_smooth_A_smoothed"] + in_a = (adata.obs["condition"] == "A").values + assert np.all(np.isnan(layer[~keep])) + # kept cells outside the condition are still evaluated, as without a mask + assert np.all(np.isfinite(layer[keep & ~in_a])) + info = adata.uns["kompot_smooth"]["last_run_info"] + assert "in_state_space" in str(info) + + def test_variance_layers_masked(self, adata): + keep = self._keep(adata) + kompot.smooth_expression( + adata, + groupby="condition", + condition="A", + sample_col="sample", + cell_mask=keep, + gp=GPSettings(n_landmarks=30, random_state=0, use_empirical_variance=True), + output=OutputSettings(progress=False), + ) + for key in ("kompot_smooth_A_obs_variance", "kompot_smooth_A_sample_variance"): + if key in adata.layers: + assert np.all(np.isnan(adata.layers[key][~keep])) + assert "kompot_smooth_A_obs_variance" in adata.layers + + def test_genes_subset_with_mask(self, adata): + keep = self._keep(adata) + kompot.smooth_expression( + adata, genes=["gene_0", "gene_1"], gp=GP_SEEDED, cell_mask=keep, + output=OutputSettings(progress=False), + ) + layer = adata.layers["kompot_smooth_all_smoothed"] + assert np.all(np.isnan(layer[~keep])) and np.all(np.isfinite(layer[keep][:, :2])) + assert np.all(np.isnan(layer[:, 2:])) + + def test_no_mask_unchanged(self, adata): + """cell_mask=None (default) must give the same result as before the option existed""" + a2 = adata.copy() + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) + kompot.smooth_expression(a2, gp=GP_SEEDED, cell_mask=None, output=OutputSettings(progress=False)) + np.testing.assert_array_equal(adata.layers["kompot_smooth_all_smoothed"], a2.layers["kompot_smooth_all_smoothed"]) + assert np.all(np.isfinite(adata.layers["kompot_smooth_all_smoothed"])) + + @pytest.mark.parametrize( + "bad, match", + [ + ("no_such_column", "not found"), + (np.ones(5, dtype=bool), "entries"), + (np.zeros(80, dtype=bool), "keeps no cell"), + (np.ones(80, dtype=int), "boolean"), + ], + ) + def test_invalid_masks_raise(self, adata, bad, match): + with pytest.raises(ValueError, match=match): + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask=bad, output=OutputSettings(progress=False)) + + def test_non_bool_obs_column_raises(self, adata): + adata.obs["flag"] = ["yes"] * 80 + with pytest.raises(ValueError, match="boolean"): + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="flag", output=OutputSettings(progress=False)) + + def test_nan_coordinates_inside_mask_raise(self, adata): + keep = self._keep(adata) + adata.obsm["DM_EigenVectors"][3] = np.nan # cell 3 is kept + with pytest.raises(ValueError, match="non-finite"): + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False)) From 4d1ddf7fe711fd46eb7280e816066ebadff0c88e Mon Sep 17 00:00:00 2001 From: Dominik Date: Thu, 22 Oct 2026 21:05:56 -0700 Subject: [PATCH 2/4] cell_mask: take an .obs column only, record it in the run info, expose it in the CLI Maintainer review on #36: - smooth_expression's cell_mask is the name of a boolean adata.obs column (no missing values); the unreleased array form is removed. A missing column, a non-boolean or partly missing column, an all-False column and a non-string argument are refused with a clear error. - Run info: params["cell_mask"] (column), params["n_cells_masked_out"], params["masked_cells"] (states that masked cells are excluded and NaN); the field descriptions of every output layer say "NaN for the N cells masked out by obs['']". - CLI: kompot smooth --cell-mask COLUMN, documented in --help; config template key cell_mask. - Tests: kept rows bitwise equal to groupby=, condition=True (same fit); nullable-boolean column bitwise equal to numpy bool; run-info fields; refusals; CLI --help, CLI round-trip bitwise equal to the Python API, CLI refusals. --- kompot/anndata/smooth.py | 82 ++++++----- kompot/cli/smooth.py | 16 ++ .../cli/templates/smooth_config_template.yaml | 3 + tests/test_anndata_smooth.py | 137 +++++++++++++----- tests/test_cli_smooth.py | 66 +++++++++ 5 files changed, 227 insertions(+), 77 deletions(-) diff --git a/kompot/anndata/smooth.py b/kompot/anndata/smooth.py index 1987559..ff920d2 100644 --- a/kompot/anndata/smooth.py +++ b/kompot/anndata/smooth.py @@ -45,7 +45,7 @@ def smooth_expression( storage: "StorageSettings | None" = None, output: "OutputSettings | None" = None, model: "ModelSettings | None" = None, - cell_mask: Optional[Union[str, np.ndarray, pd.Series]] = None, + cell_mask: Optional[str] = None, **function_kwargs, ) -> Optional[Union[Dict[str, Any], Any]]: """Smooth gene expression for a single condition using GP regression. @@ -93,17 +93,19 @@ def smooth_expression( model : ModelSettings, optional Pre-fitted :class:`~kompot.ExpressionModel` to inject via ``model1``. When provided, skips internal fitting. - cell_mask : str or array-like of bool, optional - Cells to keep. Either the name of a boolean column in ``adata.obs`` - or a boolean array of length ``adata.n_obs``; ``True`` keeps a cell. - Cells with ``False`` are **masked out**: they are excluded from - fitting, they are never evaluated, their ``obsm`` coordinates are - never read (so they may be NaN, e.g. cells left out of a diffusion - map), and every per-cell output (smoothed values, standard - deviations and, when requested, the variance layers) is NaN for - them. Combined with ``groupby``/``condition``, the model is fitted - on cells in the condition *and* in the mask, and evaluated on every - cell in the mask. *None* (default) keeps every cell, as before. + cell_mask : str, optional + Name of a **boolean** column in ``adata.obs`` (no missing values); + ``True`` keeps a cell. Cells with ``False`` are **masked out**: they + are excluded from fitting, they are never evaluated, their ``obsm`` + coordinates are never read (so they may be NaN, e.g. cells left out + of a diffusion map), and every per-cell output (smoothed values, + standard deviations and, when requested, the variance layers) is + NaN for them. Combined with ``groupby``/``condition``, the model is + fitted on cells in the condition *and* in the mask, and evaluated on + every cell in the mask. The run info records the column, the number + of masked cells and that they are NaN (``params["cell_mask"]``, + ``params["n_cells_masked_out"]``, ``params["masked_cells"]``). + *None* (default) keeps every cell, as before. **function_kwargs Forwarded to :class:`mellon.FunctionEstimator`. @@ -353,29 +355,33 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["sample_variance_key"]] = sam_var # --- Metadata --- + _mask_note = ( + "" if keep is None else + f"; NaN for the {int((~keep).sum()):,} cells masked out by obs['{cell_mask}']" + ) field_mapping = { field_names["smoothed_key"]: { "location": "layers", "type": "smoothed", - "description": f"GP-smoothed expression for {cond_label}", + "description": f"GP-smoothed expression for {cond_label}" + _mask_note, }, field_names["std_key"]: { "location": "layers", "type": "std", - "description": f"Total posterior standard deviation for {cond_label}", + "description": f"Total posterior standard deviation for {cond_label}" + _mask_note, }, } if use_empirical_variance: field_mapping[field_names["obs_variance_key"]] = { "location": "layers", "type": "obs_variance", - "description": f"Per-gene empirical (aleatoric) variance for {cond_label}", + "description": f"Per-gene empirical (aleatoric) variance for {cond_label}" + _mask_note, } if sample_col is not None and expr_model.has_sample_variance: field_mapping[field_names["sample_variance_key"]] = { "location": "layers", "type": "sample_variance", - "description": f"Biological-replicate variance for {cond_label}", + "description": f"Biological-replicate variance for {cond_label}" + _mask_note, } run_info = { @@ -393,9 +399,13 @@ def _expand_genes(values, n_cells, n_genes): "sample_col": sample_col, "n_genes": len(selected_genes), "n_cells": int(mask.sum()), - "cell_mask": (None if cell_mask is None else - cell_mask if isinstance(cell_mask, str) else "array"), + "cell_mask": cell_mask, "n_cells_masked_out": 0 if keep is None else int((~keep).sum()), + "masked_cells": ( + None if keep is None else + f"cells with adata.obs['{cell_mask}'] == False are excluded from " + "fitting, never evaluated, and NaN in every output layer of this run" + ), "batch_size": batch_size, "result_key": result_key, }, @@ -425,30 +435,28 @@ def _expand_genes(values, n_cells, n_genes): def _resolve_cell_mask(adata, cell_mask): - """Return a boolean array of kept cells, or None when no mask is given.""" + """Boolean array of kept cells from a boolean ``adata.obs`` column, or None when no mask is given.""" if cell_mask is None: return None - if isinstance(cell_mask, str): - if cell_mask not in adata.obs: - raise ValueError(f"cell_mask column '{cell_mask}' not found in adata.obs.") - col = adata.obs[cell_mask] - if not pd.api.types.is_bool_dtype(col) or col.isna().any(): - raise ValueError( - f"cell_mask column '{cell_mask}' must be boolean without missing values " - f"(got dtype {col.dtype})." - ) - keep = col.to_numpy(dtype=bool) - else: - arr = np.asarray(cell_mask) - if arr.dtype != bool: - raise ValueError(f"cell_mask must be a boolean array (got dtype {arr.dtype}).") - keep = arr.reshape(-1) - if keep.shape != (adata.n_obs,): + if not isinstance(cell_mask, str): + raise TypeError( + "cell_mask must be the name of a boolean column in adata.obs " + f"(got {type(cell_mask).__name__})." + ) + if cell_mask not in adata.obs: + raise ValueError( + f"cell_mask column '{cell_mask}' not found in adata.obs. " + f"Available columns: {list(adata.obs.columns)}" + ) + col = adata.obs[cell_mask] + if not pd.api.types.is_bool_dtype(col) or col.isna().any(): raise ValueError( - f"cell_mask has {keep.shape[0]} entries but adata has {adata.n_obs} observations." + f"cell_mask column '{cell_mask}' must be boolean (True = keep) without missing " + f"values; got dtype '{col.dtype}'." ) + keep = col.to_numpy(dtype=bool) if not keep.any(): - raise ValueError("cell_mask keeps no cell.") + raise ValueError(f"cell_mask column '{cell_mask}' keeps no cell (all False).") return keep diff --git a/kompot/cli/smooth.py b/kompot/cli/smooth.py index 17402ec..2172f45 100644 --- a/kompot/cli/smooth.py +++ b/kompot/cli/smooth.py @@ -79,6 +79,18 @@ def add_smooth_parser(subparsers) -> argparse.ArgumentParser: help="Which condition to smooth (requires --groupby). If omitted, all cells are used.", ) + parser.add_argument( + "--cell-mask", + type=str, + metavar="COLUMN", + help=( + "Boolean column in adata.obs; True keeps a cell. Cells with False are " + "masked out: excluded from fitting, never evaluated (their obsm " + "coordinates may be NaN), and NaN in every output layer. The run info " + "records the column and the number of masked cells. Default: keep all cells." + ), + ) + # Common optional parameters parser.add_argument( "--obsm-key", @@ -261,6 +273,8 @@ def run_smooth(args): logger.info(f" Condition: {params['condition']}") else: logger.info(" Using all cells") + if params.get("cell_mask"): + logger.info(f" Cell mask: obs['{params['cell_mask']}'] (False = masked out, NaN in outputs)") logger.info(f" ObsM key: {params.get('obsm_key', 'DM_EigenVectors')}") if params.get("layer"): logger.info(f" Layer: {params['layer']}") @@ -282,6 +296,7 @@ def run_smooth(args): layer = params.pop("layer", None) genes = params.pop("genes", None) sample_col = params.pop("sample_col", None) + cell_mask = params.pop("cell_mask", None) # Build Settings from remaining params gp_keys = { @@ -319,6 +334,7 @@ def run_smooth(args): layer=layer, genes=genes, sample_col=sample_col, + cell_mask=cell_mask, gp=gp, storage=storage, output=output, diff --git a/kompot/cli/templates/smooth_config_template.yaml b/kompot/cli/templates/smooth_config_template.yaml index a9cfd53..caf63d4 100644 --- a/kompot/cli/templates/smooth_config_template.yaml +++ b/kompot/cli/templates/smooth_config_template.yaml @@ -7,6 +7,9 @@ groupby: null # Column in adata.obs with condition labels condition: null # Which condition to smooth (requires groupby) +# Cell mask (optional): +cell_mask: null # Boolean column in adata.obs; False = masked out (excluded, NaN in outputs) + # Cell state representation: obsm_key: "DM_EigenVectors" # Key in adata.obsm (e.g., X_pca, DM_EigenVectors) diff --git a/tests/test_anndata_smooth.py b/tests/test_anndata_smooth.py index 0f58a0e..94653b0 100644 --- a/tests/test_anndata_smooth.py +++ b/tests/test_anndata_smooth.py @@ -171,21 +171,24 @@ def test_deprecated_compute_smoothed_expression(self, adata): GP_SEEDED = GPSettings(n_landmarks=30, random_state=0) +MASKED = [0, 7, 41, 79] -class TestCellMask: - """cell_mask: masked-out cells are excluded from fitting, never evaluated, and NaN in every output.""" +def _with_mask(adata, col="in_state_space"): + keep = np.ones(adata.n_obs, dtype=bool) + keep[MASKED] = False + adata.obs[col] = keep + return keep + - @staticmethod - def _keep(adata): - keep = np.ones(adata.n_obs, dtype=bool) - keep[[0, 7, 41, 79]] = False - return keep +class TestCellMask: + """cell_mask (a boolean adata.obs column): masked-out cells are excluded from fitting, never evaluated, and NaN in + every output; the run info records the column, the count and that they are NaN.""" def test_masked_rows_nan_kept_rows_finite(self, adata): - keep = self._keep(adata) + keep = _with_mask(adata) kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) ) for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): layer = adata.layers[key] @@ -193,27 +196,53 @@ def test_masked_rows_nan_kept_rows_finite(self, adata): assert np.all(np.isfinite(layer[keep])) def test_masked_coordinates_are_never_read(self, adata): - keep = self._keep(adata) + keep = _with_mask(adata) adata.obsm["DM_EigenVectors"][~keep] = np.nan # e.g. cells left out of the diffusion map kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) ) layer = adata.layers["kompot_smooth_all_smoothed"] assert np.all(np.isnan(layer[~keep])) and np.all(np.isfinite(layer[keep])) + def test_kept_rows_bitwise_equal_same_fit_via_groupby(self, adata): + """the mask selects the SAME training cells as groupby=, condition=True; on kept cells the results are + bitwise identical, and only the masked rows differ (NaN instead of an extrapolation)""" + keep = _with_mask(adata) + ref = adata.copy() + kompot.smooth_expression( + ref, groupby="in_state_space", condition=True, gp=GP_SEEDED, output=OutputSettings(progress=False) + ) + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) + ) + r = ref.layers["kompot_smooth_True_smoothed"] + m = adata.layers["kompot_smooth_all_smoothed"] + np.testing.assert_array_equal(m[keep], r[keep]) + assert np.all(np.isfinite(r[~keep])) and np.all(np.isnan(m[~keep])) + + def test_nullable_boolean_column_bitwise_equal(self, adata): + """a pandas nullable 'boolean' column without missing values gives bitwise the same result as a numpy bool one""" + _with_mask(adata) + a2 = adata.copy() + a2.obs["in_state_space"] = a2.obs["in_state_space"].astype("boolean") + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) + kompot.smooth_expression(a2, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) + np.testing.assert_array_equal( + adata.layers["kompot_smooth_all_smoothed"], a2.layers["kompot_smooth_all_smoothed"] + ) + def test_equals_smoothing_the_kept_subset(self, adata): - keep = self._keep(adata) + keep = _with_mask(adata) sub = adata[keep].copy() kompot.smooth_expression(sub, gp=GP_SEEDED, output=OutputSettings(progress=False)) kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False) + adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) ) for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): np.testing.assert_allclose(adata.layers[key][keep], sub.layers[key], rtol=1e-6, atol=1e-8) - def test_obs_column_and_condition(self, adata): - keep = self._keep(adata) - adata.obs["in_state_space"] = keep + def test_with_condition(self, adata): + keep = _with_mask(adata) kompot.smooth_expression( adata, groupby="condition", @@ -227,29 +256,49 @@ def test_obs_column_and_condition(self, adata): assert np.all(np.isnan(layer[~keep])) # kept cells outside the condition are still evaluated, as without a mask assert np.all(np.isfinite(layer[keep & ~in_a])) - info = adata.uns["kompot_smooth"]["last_run_info"] - assert "in_state_space" in str(info) + + def test_run_info_records_the_mask(self, adata): + from kompot.anndata.utils.json_utils import get_json_metadata + + _with_mask(adata) + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) + ) + info = get_json_metadata(adata, "kompot_smooth.last_run_info") + p = info["params"] + assert p["cell_mask"] == "in_state_space" + assert p["n_cells_masked_out"] == len(MASKED) + assert "NaN" in p["masked_cells"] and "in_state_space" in p["masked_cells"] + desc = info["field_mapping"]["kompot_smooth_all_smoothed"]["description"] + assert f"NaN for the {len(MASKED)} cells masked out by obs['in_state_space']" in desc + + def test_run_info_without_mask(self, adata): + from kompot.anndata.utils.json_utils import get_json_metadata + + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) + p = get_json_metadata(adata, "kompot_smooth.last_run_info")["params"] + assert p["cell_mask"] is None and p["n_cells_masked_out"] == 0 and p["masked_cells"] is None def test_variance_layers_masked(self, adata): - keep = self._keep(adata) + keep = _with_mask(adata) kompot.smooth_expression( adata, groupby="condition", condition="A", sample_col="sample", - cell_mask=keep, + cell_mask="in_state_space", gp=GPSettings(n_landmarks=30, random_state=0, use_empirical_variance=True), output=OutputSettings(progress=False), ) + assert "kompot_smooth_A_obs_variance" in adata.layers for key in ("kompot_smooth_A_obs_variance", "kompot_smooth_A_sample_variance"): if key in adata.layers: assert np.all(np.isnan(adata.layers[key][~keep])) - assert "kompot_smooth_A_obs_variance" in adata.layers def test_genes_subset_with_mask(self, adata): - keep = self._keep(adata) + keep = _with_mask(adata) kompot.smooth_expression( - adata, genes=["gene_0", "gene_1"], gp=GP_SEEDED, cell_mask=keep, + adata, genes=["gene_0", "gene_1"], gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False), ) layer = adata.layers["kompot_smooth_all_smoothed"] @@ -257,33 +306,41 @@ def test_genes_subset_with_mask(self, adata): assert np.all(np.isnan(layer[:, 2:])) def test_no_mask_unchanged(self, adata): - """cell_mask=None (default) must give the same result as before the option existed""" + """cell_mask=None (default) gives the same result as before the option existed""" a2 = adata.copy() kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) kompot.smooth_expression(a2, gp=GP_SEEDED, cell_mask=None, output=OutputSettings(progress=False)) np.testing.assert_array_equal(adata.layers["kompot_smooth_all_smoothed"], a2.layers["kompot_smooth_all_smoothed"]) assert np.all(np.isfinite(adata.layers["kompot_smooth_all_smoothed"])) + def test_missing_column_refused(self, adata): + with pytest.raises(ValueError, match="not found in adata.obs"): + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="no_such_column", output=OutputSettings(progress=False)) + @pytest.mark.parametrize( - "bad, match", - [ - ("no_such_column", "not found"), - (np.ones(5, dtype=bool), "entries"), - (np.zeros(80, dtype=bool), "keeps no cell"), - (np.ones(80, dtype=int), "boolean"), - ], + "values", + [["yes"] * 80, list(range(80)), [True] * 79 + [None]], + ids=["strings", "integers", "missing-value"], ) - def test_invalid_masks_raise(self, adata, bad, match): - with pytest.raises(ValueError, match=match): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask=bad, output=OutputSettings(progress=False)) - - def test_non_bool_obs_column_raises(self, adata): - adata.obs["flag"] = ["yes"] * 80 - with pytest.raises(ValueError, match="boolean"): + def test_non_boolean_column_refused(self, adata, values): + adata.obs["flag"] = pd.Series(values, index=adata.obs_names, dtype=object if None in values else None) + with pytest.raises(ValueError, match="must be boolean"): kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="flag", output=OutputSettings(progress=False)) + def test_all_false_refused(self, adata): + adata.obs["none_kept"] = np.zeros(adata.n_obs, dtype=bool) + with pytest.raises(ValueError, match="keeps no cell"): + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="none_kept", output=OutputSettings(progress=False)) + + def test_array_form_refused(self, adata): + """the anndata interface takes an obs column only""" + with pytest.raises(TypeError, match="name of a boolean column"): + kompot.smooth_expression( + adata, gp=GP_SEEDED, cell_mask=np.ones(adata.n_obs, dtype=bool), output=OutputSettings(progress=False) + ) + def test_nan_coordinates_inside_mask_raise(self, adata): - keep = self._keep(adata) + _with_mask(adata) adata.obsm["DM_EigenVectors"][3] = np.nan # cell 3 is kept with pytest.raises(ValueError, match="non-finite"): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask=keep, output=OutputSettings(progress=False)) + kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) diff --git a/tests/test_cli_smooth.py b/tests/test_cli_smooth.py index b5b2730..7422162 100644 --- a/tests/test_cli_smooth.py +++ b/tests/test_cli_smooth.py @@ -252,3 +252,69 @@ def test_smooth_gene_subset(self, sample_adata, temp_dir): # Gene_0 and Gene_1 should have values, others NaN assert not np.all(np.isnan(smoothed[:, 0])) assert np.all(np.isnan(smoothed[:, 2])) + + +class TestCLISmoothCellMask: + """--cell-mask COLUMN: documented in --help; round-trip equals the Python API bitwise; bad columns refused.""" + + @staticmethod + def _write(sample_adata, temp_dir): + keep = np.ones(sample_adata.n_obs, dtype=bool) + keep[[1, 30, 60]] = False + sample_adata.obs["in_state_space"] = keep + input_file = Path(temp_dir) / "input.h5ad" + write_output(sample_adata, input_file) + return input_file, keep + + @staticmethod + def _cmd(input_file, output_file, *extra): + return [ + sys.executable, "-m", "kompot.cli", "smooth", str(input_file), "-o", str(output_file), + "--obsm-key", "X_pca", "--n-landmarks", "20", "--random-state", "0", "--no-progress", *extra, + ] + + def test_help_documents_cell_mask(self): + r = subprocess.run([sys.executable, "-m", "kompot.cli", "smooth", "--help"], capture_output=True, text=True) + assert r.returncode == 0 + assert "--cell-mask" in r.stdout and "NaN" in r.stdout + + def test_cell_mask_round_trip(self, sample_adata, temp_dir): + import anndata as ad + import kompot + from kompot import GPSettings, OutputSettings + from kompot.anndata.utils.json_utils import get_json_metadata + + input_file, keep = self._write(sample_adata, temp_dir) + output_file = Path(temp_dir) / "output.h5ad" + r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "in_state_space"), + capture_output=True, text=True) + assert r.returncode == 0, f"CLI failed: {r.stderr}" + out = ad.read_h5ad(output_file) + cli = np.asarray(out.layers["kompot_smooth_all_smoothed"]) + assert np.all(np.isnan(cli[~keep])) and np.all(np.isfinite(cli[keep])) + p = get_json_metadata(out, "kompot_smooth.last_run_info")["params"] + assert p["cell_mask"] == "in_state_space" and p["n_cells_masked_out"] == 3 + + # the same call through the Python API, in-process: bitwise equal + api = ad.read_h5ad(input_file) + kompot.smooth_expression( + api, obsm_key="X_pca", cell_mask="in_state_space", + gp=GPSettings(n_landmarks=20, random_state=0), output=OutputSettings(progress=False), + ) + np.testing.assert_array_equal(cli, np.asarray(api.layers["kompot_smooth_all_smoothed"])) + + def test_cell_mask_missing_column_fails(self, sample_adata, temp_dir): + input_file, _ = self._write(sample_adata, temp_dir) + output_file = Path(temp_dir) / "output.h5ad" + r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "no_such_column"), + capture_output=True, text=True) + assert r.returncode != 0 and "not found in adata.obs" in r.stderr + assert not output_file.exists() + + def test_cell_mask_non_boolean_column_fails(self, sample_adata, temp_dir): + input_file, _ = self._write(sample_adata, temp_dir) + output_file = Path(temp_dir) / "output.h5ad" + r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "condition"), + capture_output=True, text=True) + assert r.returncode != 0 and "must be boolean" in r.stderr + assert not output_file.exists() From 4a79a5e0062f4bfab31ece7a3bfea8417971d6fe Mon Sep 17 00:00:00 2001 From: Dominik Date: Thu, 22 Oct 2026 21:06:17 -0700 Subject: [PATCH 3/4] CHANGELOG: cell_mask is an .obs column (array form removed); run-info fields; CLI option --- CHANGELOG.md | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 8e61e69..325af4e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -12,8 +12,8 @@ cell-state representation at all: cells removed by a quality gate before the diffusion map was computed typically carry NaN coordinates, and a smoothed value for them is either NaN by accident or an extrapolation nobody asked for. -`cell_mask` (a boolean `adata.obs` column name, or a boolean array; `True` keeps a -cell) makes the exclusion explicit: +`cell_mask` (the name of a boolean `adata.obs` column without missing values; `True` +keeps a cell; CLI: `kompot smooth --cell-mask COLUMN`) makes the exclusion explicit: - masked-out cells are excluded from fitting and are **never evaluated**, so their `obsm` coordinates are never read and may be NaN; @@ -23,8 +23,11 @@ cell) makes the exclusion explicit: in the mask, and evaluated on every kept cell; - kept cells with non-finite coordinates raise a `ValueError` instead of propagating NaN silently; - - the run record stores the mask (column name, or `"array"`) and the number of - cells masked out. + - the run info records the column (`params["cell_mask"]`), the number of cells + masked out (`params["n_cells_masked_out"]`) and that they are NaN + (`params["masked_cells"]`, and the description of every output layer); + - a missing, non-boolean or partly missing column, or one that keeps no cell, is + refused with a clear error. The default, `cell_mask=None`, is unchanged. On kept cells the result equals smoothing the kept subset on its own (tested). From fa26738f415fe51c9ea07a7572635fa6a07ba209 Mon Sep 17 00:00:00 2001 From: Dominik Date: Fri, 23 Oct 2026 00:00:49 -0700 Subject: [PATCH 4/4] smooth_expression: no mask argument; cells with non-finite coordinates are excluded from fitting and NaN Maintainer review on #36: prediction is per cell, so a cell mask is not needed. Measured: ExpressionModel.predict and .std are row-independent (mellon LandmarksConditional: mean = K(x, landmarks) @ weights; the variance is solved against the fixed landmark Cholesky factor). With NaN rows planted, those rows come back NaN and every other row is bitwise equal to evaluating without them, at batch sizes 7, 50 and 1000, with and without the variance terms. What failed on main is the FIT: cells with NaN coordinates reached landmark selection (KMeans), which refuses NaN. Now cells whose coordinates are not all finite are excluded from fitting only; they are still evaluated and come out NaN. The run info records n_cells_nonfinite and nonfinite_cells, and each output layer's description gives the NaN count. cell_mask, the CLI option and its config key are removed (reverted to main). Tests: non-finite rows NaN; other rows bitwise equal to the same fit; equal to smoothing the valid subset; partial-NaN and inf rows; all-finite control; run info; condition, variance layers and gene subset; no valid cell raises; CLI round trip bitwise equal to the Python API. --- CHANGELOG.md | 45 ++-- kompot/anndata/smooth.py | 134 ++++------- kompot/cli/smooth.py | 16 -- .../cli/templates/smooth_config_template.yaml | 3 - tests/test_anndata_smooth.py | 223 +++++++----------- tests/test_cli_smooth.py | 69 ++---- 6 files changed, 159 insertions(+), 331 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 325af4e..2d93e6d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,33 +4,24 @@ All notable changes to this project will be documented in this file. ## [Unreleased] -### Added — `smooth_expression(..., cell_mask=...)`: mask cells out instead of extrapolating to them - -`smooth_expression` fits on the selected cells and evaluates the model on **every** -cell. That is right for conditions, but wrong for cells that have no place in the -cell-state representation at all: cells removed by a quality gate before the -diffusion map was computed typically carry NaN coordinates, and a smoothed value -for them is either NaN by accident or an extrapolation nobody asked for. - -`cell_mask` (the name of a boolean `adata.obs` column without missing values; `True` -keeps a cell; CLI: `kompot smooth --cell-mask COLUMN`) makes the exclusion explicit: - - - masked-out cells are excluded from fitting and are **never evaluated**, so their - `obsm` coordinates are never read and may be NaN; - - every per-cell output (smoothed values, standard deviations, and the empirical - and sample variance layers when requested) is **NaN** for them; - - with `groupby`/`condition`, the model is fitted on cells in the condition *and* - in the mask, and evaluated on every kept cell; - - kept cells with non-finite coordinates raise a `ValueError` instead of - propagating NaN silently; - - the run info records the column (`params["cell_mask"]`), the number of cells - masked out (`params["n_cells_masked_out"]`) and that they are NaN - (`params["masked_cells"]`, and the description of every output layer); - - a missing, non-boolean or partly missing column, or one that keeps no cell, is - refused with a clear error. - -The default, `cell_mask=None`, is unchanged. On kept cells the result equals -smoothing the kept subset on its own (tested). +### Fixed — `smooth_expression` with non-finite cell-state coordinates + +Cells removed by a quality gate before the diffusion map was computed often carry +NaN coordinates in `adata.obsm`. `smooth_expression` used them for fitting, and +the fit failed: landmark selection (KMeans) refuses NaN input. + +Now cells whose coordinates are not all finite (NaN or inf) are excluded from +fitting. They are still evaluated, and they come out NaN in every output layer +(smoothed values, standard deviation, and the variance layers when requested). +No argument is needed: prediction is per cell (`K(x, landmarks) @ weights`; the +variance is solved against the fixed landmark Cholesky factor), so a NaN row stays +NaN and no other cell's value changes. Measured: with NaN rows planted, every +other row of `predict` and `std` is bitwise equal to evaluating without those +rows, at batch sizes 7, 50 and 1000. + +The run info records `params["n_cells_nonfinite"]` and `params["nonfinite_cells"]`, +and each output layer's description gives the number of NaN cells. With all +coordinates finite, nothing changes. ## [0.9.0] - 2026-09-15 diff --git a/kompot/anndata/smooth.py b/kompot/anndata/smooth.py index ff920d2..f93dd4f 100644 --- a/kompot/anndata/smooth.py +++ b/kompot/anndata/smooth.py @@ -45,7 +45,6 @@ def smooth_expression( storage: "StorageSettings | None" = None, output: "OutputSettings | None" = None, model: "ModelSettings | None" = None, - cell_mask: Optional[str] = None, **function_kwargs, ) -> Optional[Union[Dict[str, Any], Any]]: """Smooth gene expression for a single condition using GP regression. @@ -77,7 +76,11 @@ def smooth_expression( Which group in ``groupby`` to smooth. If *None* and ``groupby`` is *None*, all cells are used. obsm_key : str - Key in ``adata.obsm`` for cell-state coordinates. + Key in ``adata.obsm`` for cell-state coordinates. Cells whose + coordinates are not all finite (NaN or inf, e.g. cells left out of + the diffusion map) are excluded from fitting and get NaN in every + output layer. Prediction is per cell, so no other cell's value + changes. The run info records how many (``params["n_cells_nonfinite"]``). layer : str, optional Layer to use as expression input. *None* means ``adata.X``. genes : list of str, optional @@ -93,19 +96,6 @@ def smooth_expression( model : ModelSettings, optional Pre-fitted :class:`~kompot.ExpressionModel` to inject via ``model1``. When provided, skips internal fitting. - cell_mask : str, optional - Name of a **boolean** column in ``adata.obs`` (no missing values); - ``True`` keeps a cell. Cells with ``False`` are **masked out**: they - are excluded from fitting, they are never evaluated, their ``obsm`` - coordinates are never read (so they may be NaN, e.g. cells left out - of a diffusion map), and every per-cell output (smoothed values, - standard deviations and, when requested, the variance layers) is - NaN for them. Combined with ``groupby``/``condition``, the model is - fitted on cells in the condition *and* in the mask, and evaluated on - every cell in the mask. The run info records the column, the number - of masked cells and that they are NaN (``params["cell_mask"]``, - ``params["n_cells_masked_out"]``, ``params["masked_cells"]``). - *None* (default) keeps every cell, as before. **function_kwargs Forwarded to :class:`mellon.FunctionEstimator`. @@ -157,8 +147,6 @@ def smooth_expression( if groupby is not None: validate_groupby(adata, groupby) - keep = _resolve_cell_mask(adata, cell_mask) - # --- Build cell mask --- if groupby is not None and condition is not None: mask = (adata.obs[groupby] == condition).values @@ -171,15 +159,6 @@ def smooth_expression( mask = np.ones(adata.n_obs, dtype=bool) logger.info(f"Smoothing all {mask.sum():,} cells") - if keep is not None: - mask = mask & keep - if mask.sum() == 0: - raise ValueError("No cell is both in the selected condition and kept by cell_mask.") - logger.info( - f"cell_mask: {int((~keep).sum()):,} cells masked out (NaN in every output); " - f"fitting on {mask.sum():,} cells" - ) - # --- Build field names --- cond_label = _sanitize_name(condition) if condition else "all" field_names = _smooth_field_names( @@ -213,14 +192,24 @@ def smooth_expression( # --- Extract data --- X_all = adata.obsm[obsm_key] - X_train = X_all[mask] - # Masked-out cells are never evaluated, so their coordinates are never read. - X_eval = X_all if keep is None else X_all[keep] - if keep is not None and not np.all(np.isfinite(np.asarray(X_eval))): - raise ValueError( - f"adata.obsm['{obsm_key}'] has non-finite values in cells kept by " - "cell_mask; mask those cells out or fix the representation." + # Cells with non-finite coordinates have no place in the representation: + # exclude them from fitting (landmark KMeans and the GP refuse NaN). They + # are still evaluated below; prediction is per cell (K(x, landmarks) + # @ weights; the variance solves against the fixed landmark Cholesky), so + # their rows come out NaN and no other cell's value changes. + finite = np.isfinite(np.asarray(X_all, dtype=float)).all(axis=1) + n_nonfinite = int((~finite).sum()) + if n_nonfinite: + mask = mask & finite + if mask.sum() == 0: + raise ValueError( + f"No selected cell has finite coordinates in adata.obsm['{obsm_key}']." + ) + logger.info( + f"{n_nonfinite:,} cells have non-finite coordinates in adata.obsm['{obsm_key}']: " + f"excluded from fitting, NaN in every output; fitting on {mask.sum():,} cells" ) + X_train = X_all[mask] if layer is not None: if layer not in adata.layers: @@ -280,22 +269,12 @@ def smooth_expression( # --- Predict on all cells --- # The model is trained on the condition cells but evaluated everywhere, - # so every cell gets a smoothed value and uncertainty estimate. With a - # cell_mask, "everywhere" means every kept cell; masked-out rows are NaN. - def _expand_cells(values): - """Place kept-cell rows into a full (n_obs, ...) array, NaN elsewhere.""" - values = np.asarray(values) - if keep is None: - return values - out = np.full((adata.n_obs,) + values.shape[1:], np.nan, dtype=np.result_type(values, np.float32)) - out[keep] = values - return out - - imputed = _expand_cells( - expr_model.predict(X_eval, batch_size=batch_size, progress=progress) + # so every cell gets a smoothed value and uncertainty estimate. + imputed = np.asarray( + expr_model.predict(X_all, batch_size=batch_size, progress=progress) ) - std_vals = _expand_cells( - expr_model.std(X_eval, batch_size=batch_size, progress=progress) + std_vals = np.asarray( + expr_model.std(X_all, batch_size=batch_size, progress=progress) ) # Broadcast (n, 1) to (n, n_genes) when only GP covariance is present n_genes_fitted = imputed.shape[1] @@ -330,8 +309,8 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["std_key"]] = std_vals if use_empirical_variance: - obs_var = _expand_cells( - expr_model.obs_variance(X_eval, batch_size=batch_size, progress=progress) + obs_var = np.asarray( + expr_model.obs_variance(X_all, batch_size=batch_size, progress=progress) ) if genes is not None: adata.layers[field_names["obs_variance_key"]] = _expand_genes( @@ -341,9 +320,9 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["obs_variance_key"]] = obs_var if sample_col is not None and expr_model.has_sample_variance: - sam_var = _expand_cells( + sam_var = np.asarray( expr_model.sample_variance( - X_eval, diag=True, batch_size=batch_size, progress=progress + X_all, diag=True, batch_size=batch_size, progress=progress ) ) if isinstance(sam_var, np.ndarray): @@ -355,33 +334,33 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["sample_variance_key"]] = sam_var # --- Metadata --- - _mask_note = ( - "" if keep is None else - f"; NaN for the {int((~keep).sum()):,} cells masked out by obs['{cell_mask}']" + _nan_note = ( + "" if not n_nonfinite else + f"; NaN for the {n_nonfinite:,} cells with non-finite coordinates in obsm['{obsm_key}']" ) field_mapping = { field_names["smoothed_key"]: { "location": "layers", "type": "smoothed", - "description": f"GP-smoothed expression for {cond_label}" + _mask_note, + "description": f"GP-smoothed expression for {cond_label}" + _nan_note, }, field_names["std_key"]: { "location": "layers", "type": "std", - "description": f"Total posterior standard deviation for {cond_label}" + _mask_note, + "description": f"Total posterior standard deviation for {cond_label}" + _nan_note, }, } if use_empirical_variance: field_mapping[field_names["obs_variance_key"]] = { "location": "layers", "type": "obs_variance", - "description": f"Per-gene empirical (aleatoric) variance for {cond_label}" + _mask_note, + "description": f"Per-gene empirical (aleatoric) variance for {cond_label}" + _nan_note, } if sample_col is not None and expr_model.has_sample_variance: field_mapping[field_names["sample_variance_key"]] = { "location": "layers", "type": "sample_variance", - "description": f"Biological-replicate variance for {cond_label}" + _mask_note, + "description": f"Biological-replicate variance for {cond_label}" + _nan_note, } run_info = { @@ -399,12 +378,11 @@ def _expand_genes(values, n_cells, n_genes): "sample_col": sample_col, "n_genes": len(selected_genes), "n_cells": int(mask.sum()), - "cell_mask": cell_mask, - "n_cells_masked_out": 0 if keep is None else int((~keep).sum()), - "masked_cells": ( - None if keep is None else - f"cells with adata.obs['{cell_mask}'] == False are excluded from " - "fitting, never evaluated, and NaN in every output layer of this run" + "n_cells_nonfinite": n_nonfinite, + "nonfinite_cells": ( + None if not n_nonfinite else + f"cells with non-finite coordinates in adata.obsm['{obsm_key}'] are excluded " + "from fitting and NaN in every output layer of this run" ), "batch_size": batch_size, "result_key": result_key, @@ -434,32 +412,6 @@ def _expand_genes(values, n_cells, n_genes): return None -def _resolve_cell_mask(adata, cell_mask): - """Boolean array of kept cells from a boolean ``adata.obs`` column, or None when no mask is given.""" - if cell_mask is None: - return None - if not isinstance(cell_mask, str): - raise TypeError( - "cell_mask must be the name of a boolean column in adata.obs " - f"(got {type(cell_mask).__name__})." - ) - if cell_mask not in adata.obs: - raise ValueError( - f"cell_mask column '{cell_mask}' not found in adata.obs. " - f"Available columns: {list(adata.obs.columns)}" - ) - col = adata.obs[cell_mask] - if not pd.api.types.is_bool_dtype(col) or col.isna().any(): - raise ValueError( - f"cell_mask column '{cell_mask}' must be boolean (True = keep) without missing " - f"values; got dtype '{col.dtype}'." - ) - keep = col.to_numpy(dtype=bool) - if not keep.any(): - raise ValueError(f"cell_mask column '{cell_mask}' keeps no cell (all False).") - return keep - - def compute_smoothed_expression( adata, groupby: Optional[str] = None, diff --git a/kompot/cli/smooth.py b/kompot/cli/smooth.py index 2172f45..17402ec 100644 --- a/kompot/cli/smooth.py +++ b/kompot/cli/smooth.py @@ -79,18 +79,6 @@ def add_smooth_parser(subparsers) -> argparse.ArgumentParser: help="Which condition to smooth (requires --groupby). If omitted, all cells are used.", ) - parser.add_argument( - "--cell-mask", - type=str, - metavar="COLUMN", - help=( - "Boolean column in adata.obs; True keeps a cell. Cells with False are " - "masked out: excluded from fitting, never evaluated (their obsm " - "coordinates may be NaN), and NaN in every output layer. The run info " - "records the column and the number of masked cells. Default: keep all cells." - ), - ) - # Common optional parameters parser.add_argument( "--obsm-key", @@ -273,8 +261,6 @@ def run_smooth(args): logger.info(f" Condition: {params['condition']}") else: logger.info(" Using all cells") - if params.get("cell_mask"): - logger.info(f" Cell mask: obs['{params['cell_mask']}'] (False = masked out, NaN in outputs)") logger.info(f" ObsM key: {params.get('obsm_key', 'DM_EigenVectors')}") if params.get("layer"): logger.info(f" Layer: {params['layer']}") @@ -296,7 +282,6 @@ def run_smooth(args): layer = params.pop("layer", None) genes = params.pop("genes", None) sample_col = params.pop("sample_col", None) - cell_mask = params.pop("cell_mask", None) # Build Settings from remaining params gp_keys = { @@ -334,7 +319,6 @@ def run_smooth(args): layer=layer, genes=genes, sample_col=sample_col, - cell_mask=cell_mask, gp=gp, storage=storage, output=output, diff --git a/kompot/cli/templates/smooth_config_template.yaml b/kompot/cli/templates/smooth_config_template.yaml index caf63d4..a9cfd53 100644 --- a/kompot/cli/templates/smooth_config_template.yaml +++ b/kompot/cli/templates/smooth_config_template.yaml @@ -7,9 +7,6 @@ groupby: null # Column in adata.obs with condition labels condition: null # Which condition to smooth (requires groupby) -# Cell mask (optional): -cell_mask: null # Boolean column in adata.obs; False = masked out (excluded, NaN in outputs) - # Cell state representation: obsm_key: "DM_EigenVectors" # Key in adata.obsm (e.g., X_pca, DM_EigenVectors) diff --git a/tests/test_anndata_smooth.py b/tests/test_anndata_smooth.py index 94653b0..5a0ea2c 100644 --- a/tests/test_anndata_smooth.py +++ b/tests/test_anndata_smooth.py @@ -171,176 +171,115 @@ def test_deprecated_compute_smoothed_expression(self, adata): GP_SEEDED = GPSettings(n_landmarks=30, random_state=0) -MASKED = [0, 7, 41, 79] +BAD = [0, 7, 41, 79] -def _with_mask(adata, col="in_state_space"): - keep = np.ones(adata.n_obs, dtype=bool) - keep[MASKED] = False - adata.obs[col] = keep - return keep +def _invalidate(adata, rows=BAD): + """set the coordinates of `rows` to NaN; return the boolean array of cells with finite coordinates""" + adata.obsm["DM_EigenVectors"][rows] = np.nan + ok = np.ones(adata.n_obs, dtype=bool) + ok[rows] = False + return ok -class TestCellMask: - """cell_mask (a boolean adata.obs column): masked-out cells are excluded from fitting, never evaluated, and NaN in - every output; the run info records the column, the count and that they are NaN.""" +class TestNonFiniteCoordinates: + """Cells whose obsm coordinates are not all finite are excluded from fitting and NaN in every output; prediction + is per cell, so every other cell's values are unchanged. No mask argument is needed.""" - def test_masked_rows_nan_kept_rows_finite(self, adata): - keep = _with_mask(adata) - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) - ) + def test_nan_rows_nan_other_rows_finite(self, adata): + ok = _invalidate(adata) + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): layer = adata.layers[key] - assert np.all(np.isnan(layer[~keep])) - assert np.all(np.isfinite(layer[keep])) - - def test_masked_coordinates_are_never_read(self, adata): - keep = _with_mask(adata) - adata.obsm["DM_EigenVectors"][~keep] = np.nan # e.g. cells left out of the diffusion map - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) - ) - layer = adata.layers["kompot_smooth_all_smoothed"] - assert np.all(np.isnan(layer[~keep])) and np.all(np.isfinite(layer[keep])) + assert np.all(np.isnan(layer[~ok])) + assert np.all(np.isfinite(layer[ok])) - def test_kept_rows_bitwise_equal_same_fit_via_groupby(self, adata): - """the mask selects the SAME training cells as groupby=, condition=True; on kept cells the results are - bitwise identical, and only the masked rows differ (NaN instead of an extrapolation)""" - keep = _with_mask(adata) + def test_other_rows_bitwise_equal_same_fit(self, adata): + """reference: the SAME training cells (groupby on a finiteness column) with the invalid cells given arbitrary + finite coordinates. Valid cells are bitwise equal; only the invalid rows differ (NaN instead of a value).""" ref = adata.copy() - kompot.smooth_expression( - ref, groupby="in_state_space", condition=True, gp=GP_SEEDED, output=OutputSettings(progress=False) - ) - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) - ) - r = ref.layers["kompot_smooth_True_smoothed"] - m = adata.layers["kompot_smooth_all_smoothed"] - np.testing.assert_array_equal(m[keep], r[keep]) - assert np.all(np.isfinite(r[~keep])) and np.all(np.isnan(m[~keep])) - - def test_nullable_boolean_column_bitwise_equal(self, adata): - """a pandas nullable 'boolean' column without missing values gives bitwise the same result as a numpy bool one""" - _with_mask(adata) - a2 = adata.copy() - a2.obs["in_state_space"] = a2.obs["in_state_space"].astype("boolean") - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) - kompot.smooth_expression(a2, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) - np.testing.assert_array_equal( - adata.layers["kompot_smooth_all_smoothed"], a2.layers["kompot_smooth_all_smoothed"] - ) - - def test_equals_smoothing_the_kept_subset(self, adata): - keep = _with_mask(adata) - sub = adata[keep].copy() + ok = _invalidate(adata) + ref.obs["finite"] = ok + ref.obsm["DM_EigenVectors"][~ok] = 5.0 + kompot.smooth_expression(ref, groupby="finite", condition=True, gp=GP_SEEDED, + output=OutputSettings(progress=False)) + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) + for a, b in (("all_smoothed", "True_smoothed"), ("all_std", "True_std")): + m = adata.layers[f"kompot_smooth_{a}"] + r = ref.layers[f"kompot_smooth_{b}"] + np.testing.assert_array_equal(m[ok], r[ok]) + assert np.all(np.isnan(m[~ok])) and np.all(np.isfinite(r[~ok])) + + def test_equals_smoothing_the_valid_subset(self, adata): + ok = _invalidate(adata) + sub = adata[ok].copy() kompot.smooth_expression(sub, gp=GP_SEEDED, output=OutputSettings(progress=False)) - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) - ) + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): - np.testing.assert_allclose(adata.layers[key][keep], sub.layers[key], rtol=1e-6, atol=1e-8) + np.testing.assert_allclose(adata.layers[key][ok], sub.layers[key], rtol=1e-6, atol=1e-8) - def test_with_condition(self, adata): - keep = _with_mask(adata) - kompot.smooth_expression( - adata, - groupby="condition", - condition="A", - cell_mask="in_state_space", - gp=GP_SEEDED, - output=OutputSettings(progress=False), - ) - layer = adata.layers["kompot_smooth_A_smoothed"] - in_a = (adata.obs["condition"] == "A").values - assert np.all(np.isnan(layer[~keep])) - # kept cells outside the condition are still evaluated, as without a mask - assert np.all(np.isfinite(layer[keep & ~in_a])) + def test_partial_nan_and_inf_rows(self, adata): + X = adata.obsm["DM_EigenVectors"] + X[3, 2] = np.nan # one coordinate + X[11, 0] = np.inf + X[12, :] = -np.inf + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) + layer = adata.layers["kompot_smooth_all_smoothed"] + bad = [3, 11, 12] + assert np.all(np.isnan(layer[bad])) + assert np.all(np.isfinite(np.delete(layer, bad, axis=0))) - def test_run_info_records_the_mask(self, adata): + def test_all_finite_control(self, adata): from kompot.anndata.utils.json_utils import get_json_metadata - _with_mask(adata) - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False) - ) - info = get_json_metadata(adata, "kompot_smooth.last_run_info") - p = info["params"] - assert p["cell_mask"] == "in_state_space" - assert p["n_cells_masked_out"] == len(MASKED) - assert "NaN" in p["masked_cells"] and "in_state_space" in p["masked_cells"] - desc = info["field_mapping"]["kompot_smooth_all_smoothed"]["description"] - assert f"NaN for the {len(MASKED)} cells masked out by obs['in_state_space']" in desc + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) + assert np.all(np.isfinite(adata.layers["kompot_smooth_all_smoothed"])) + p = get_json_metadata(adata, "kompot_smooth.last_run_info")["params"] + assert p["n_cells_nonfinite"] == 0 and p["nonfinite_cells"] is None - def test_run_info_without_mask(self, adata): + def test_run_info_records_the_cells(self, adata): from kompot.anndata.utils.json_utils import get_json_metadata + _invalidate(adata) kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) - p = get_json_metadata(adata, "kompot_smooth.last_run_info")["params"] - assert p["cell_mask"] is None and p["n_cells_masked_out"] == 0 and p["masked_cells"] is None + info = get_json_metadata(adata, "kompot_smooth.last_run_info") + p = info["params"] + assert p["n_cells_nonfinite"] == len(BAD) and p["n_cells"] == adata.n_obs - len(BAD) + assert "NaN" in p["nonfinite_cells"] and "DM_EigenVectors" in p["nonfinite_cells"] + desc = info["field_mapping"]["kompot_smooth_all_smoothed"]["description"] + assert f"NaN for the {len(BAD)} cells with non-finite coordinates in obsm['DM_EigenVectors']" in desc - def test_variance_layers_masked(self, adata): - keep = _with_mask(adata) + def test_with_condition(self, adata): + ok = _invalidate(adata) kompot.smooth_expression( - adata, - groupby="condition", - condition="A", - sample_col="sample", - cell_mask="in_state_space", + adata, groupby="condition", condition="A", gp=GP_SEEDED, output=OutputSettings(progress=False) + ) + layer = adata.layers["kompot_smooth_A_smoothed"] + in_a = (adata.obs["condition"] == "A").values + assert np.all(np.isnan(layer[~ok])) + assert np.all(np.isfinite(layer[ok & ~in_a])) # valid cells outside the condition are still evaluated + + def test_variance_layers(self, adata): + ok = _invalidate(adata) + kompot.smooth_expression( + adata, groupby="condition", condition="A", sample_col="sample", gp=GPSettings(n_landmarks=30, random_state=0, use_empirical_variance=True), output=OutputSettings(progress=False), ) assert "kompot_smooth_A_obs_variance" in adata.layers for key in ("kompot_smooth_A_obs_variance", "kompot_smooth_A_sample_variance"): if key in adata.layers: - assert np.all(np.isnan(adata.layers[key][~keep])) + assert np.all(np.isnan(adata.layers[key][~ok])) + assert np.all(np.isfinite(adata.layers[key][ok])) - def test_genes_subset_with_mask(self, adata): - keep = _with_mask(adata) - kompot.smooth_expression( - adata, genes=["gene_0", "gene_1"], gp=GP_SEEDED, cell_mask="in_state_space", - output=OutputSettings(progress=False), - ) + def test_genes_subset(self, adata): + ok = _invalidate(adata) + kompot.smooth_expression(adata, genes=["gene_0", "gene_1"], gp=GP_SEEDED, output=OutputSettings(progress=False)) layer = adata.layers["kompot_smooth_all_smoothed"] - assert np.all(np.isnan(layer[~keep])) and np.all(np.isfinite(layer[keep][:, :2])) + assert np.all(np.isnan(layer[~ok])) and np.all(np.isfinite(layer[ok][:, :2])) assert np.all(np.isnan(layer[:, 2:])) - def test_no_mask_unchanged(self, adata): - """cell_mask=None (default) gives the same result as before the option existed""" - a2 = adata.copy() - kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) - kompot.smooth_expression(a2, gp=GP_SEEDED, cell_mask=None, output=OutputSettings(progress=False)) - np.testing.assert_array_equal(adata.layers["kompot_smooth_all_smoothed"], a2.layers["kompot_smooth_all_smoothed"]) - assert np.all(np.isfinite(adata.layers["kompot_smooth_all_smoothed"])) - - def test_missing_column_refused(self, adata): - with pytest.raises(ValueError, match="not found in adata.obs"): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="no_such_column", output=OutputSettings(progress=False)) - - @pytest.mark.parametrize( - "values", - [["yes"] * 80, list(range(80)), [True] * 79 + [None]], - ids=["strings", "integers", "missing-value"], - ) - def test_non_boolean_column_refused(self, adata, values): - adata.obs["flag"] = pd.Series(values, index=adata.obs_names, dtype=object if None in values else None) - with pytest.raises(ValueError, match="must be boolean"): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="flag", output=OutputSettings(progress=False)) - - def test_all_false_refused(self, adata): - adata.obs["none_kept"] = np.zeros(adata.n_obs, dtype=bool) - with pytest.raises(ValueError, match="keeps no cell"): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="none_kept", output=OutputSettings(progress=False)) - - def test_array_form_refused(self, adata): - """the anndata interface takes an obs column only""" - with pytest.raises(TypeError, match="name of a boolean column"): - kompot.smooth_expression( - adata, gp=GP_SEEDED, cell_mask=np.ones(adata.n_obs, dtype=bool), output=OutputSettings(progress=False) - ) - - def test_nan_coordinates_inside_mask_raise(self, adata): - _with_mask(adata) - adata.obsm["DM_EigenVectors"][3] = np.nan # cell 3 is kept - with pytest.raises(ValueError, match="non-finite"): - kompot.smooth_expression(adata, gp=GP_SEEDED, cell_mask="in_state_space", output=OutputSettings(progress=False)) + def test_no_valid_cell_raises(self, adata): + adata.obsm["DM_EigenVectors"][:] = np.nan + with pytest.raises(ValueError, match="finite coordinates"): + kompot.smooth_expression(adata, gp=GP_SEEDED, output=OutputSettings(progress=False)) diff --git a/tests/test_cli_smooth.py b/tests/test_cli_smooth.py index 7422162..0beb1e3 100644 --- a/tests/test_cli_smooth.py +++ b/tests/test_cli_smooth.py @@ -254,67 +254,32 @@ def test_smooth_gene_subset(self, sample_adata, temp_dir): assert np.all(np.isnan(smoothed[:, 2])) -class TestCLISmoothCellMask: - """--cell-mask COLUMN: documented in --help; round-trip equals the Python API bitwise; bad columns refused.""" - - @staticmethod - def _write(sample_adata, temp_dir): - keep = np.ones(sample_adata.n_obs, dtype=bool) - keep[[1, 30, 60]] = False - sample_adata.obs["in_state_space"] = keep - input_file = Path(temp_dir) / "input.h5ad" - write_output(sample_adata, input_file) - return input_file, keep - - @staticmethod - def _cmd(input_file, output_file, *extra): - return [ - sys.executable, "-m", "kompot.cli", "smooth", str(input_file), "-o", str(output_file), - "--obsm-key", "X_pca", "--n-landmarks", "20", "--random-state", "0", "--no-progress", *extra, - ] - - def test_help_documents_cell_mask(self): - r = subprocess.run([sys.executable, "-m", "kompot.cli", "smooth", "--help"], capture_output=True, text=True) - assert r.returncode == 0 - assert "--cell-mask" in r.stdout and "NaN" in r.stdout +class TestCLISmoothNonFinite: + """CLI round trip with NaN coordinates in obsm: those rows are NaN, recorded in the run info, and the result is + bitwise equal to the Python API.""" - def test_cell_mask_round_trip(self, sample_adata, temp_dir): + def test_round_trip(self, sample_adata, temp_dir): import anndata as ad import kompot from kompot import GPSettings, OutputSettings from kompot.anndata.utils.json_utils import get_json_metadata - input_file, keep = self._write(sample_adata, temp_dir) + bad = [1, 30, 60] + sample_adata.obsm["X_pca"][bad] = np.nan + input_file = Path(temp_dir) / "input.h5ad" output_file = Path(temp_dir) / "output.h5ad" - r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "in_state_space"), - capture_output=True, text=True) + write_output(sample_adata, input_file) + cmd = [sys.executable, "-m", "kompot.cli", "smooth", str(input_file), "-o", str(output_file), + "--obsm-key", "X_pca", "--n-landmarks", "20", "--random-state", "0", "--no-progress"] + r = subprocess.run(cmd, capture_output=True, text=True) assert r.returncode == 0, f"CLI failed: {r.stderr}" out = ad.read_h5ad(output_file) cli = np.asarray(out.layers["kompot_smooth_all_smoothed"]) - assert np.all(np.isnan(cli[~keep])) and np.all(np.isfinite(cli[keep])) - p = get_json_metadata(out, "kompot_smooth.last_run_info")["params"] - assert p["cell_mask"] == "in_state_space" and p["n_cells_masked_out"] == 3 - - # the same call through the Python API, in-process: bitwise equal + ok = np.ones(out.n_obs, dtype=bool) + ok[bad] = False + assert np.all(np.isnan(cli[~ok])) and np.all(np.isfinite(cli[ok])) + assert get_json_metadata(out, "kompot_smooth.last_run_info")["params"]["n_cells_nonfinite"] == 3 api = ad.read_h5ad(input_file) - kompot.smooth_expression( - api, obsm_key="X_pca", cell_mask="in_state_space", - gp=GPSettings(n_landmarks=20, random_state=0), output=OutputSettings(progress=False), - ) + kompot.smooth_expression(api, obsm_key="X_pca", gp=GPSettings(n_landmarks=20, random_state=0), + output=OutputSettings(progress=False)) np.testing.assert_array_equal(cli, np.asarray(api.layers["kompot_smooth_all_smoothed"])) - - def test_cell_mask_missing_column_fails(self, sample_adata, temp_dir): - input_file, _ = self._write(sample_adata, temp_dir) - output_file = Path(temp_dir) / "output.h5ad" - r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "no_such_column"), - capture_output=True, text=True) - assert r.returncode != 0 and "not found in adata.obs" in r.stderr - assert not output_file.exists() - - def test_cell_mask_non_boolean_column_fails(self, sample_adata, temp_dir): - input_file, _ = self._write(sample_adata, temp_dir) - output_file = Path(temp_dir) / "output.h5ad" - r = subprocess.run(self._cmd(input_file, output_file, "--cell-mask", "condition"), - capture_output=True, text=True) - assert r.returncode != 0 and "must be boolean" in r.stderr - assert not output_file.exists()