diff --git a/CHANGELOG.md b/CHANGELOG.md index ba78d57..2d93e6d 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,27 @@ All notable changes to this project will be documented in this file. +## [Unreleased] + +### 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 ### 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..f93dd4f 100644 --- a/kompot/anndata/smooth.py +++ b/kompot/anndata/smooth.py @@ -76,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 @@ -188,6 +192,23 @@ def smooth_expression( # --- Extract data --- X_all = adata.obsm[obsm_key] + # 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: @@ -313,29 +334,33 @@ def _expand_genes(values, n_cells, n_genes): adata.layers[field_names["sample_variance_key"]] = sam_var # --- Metadata --- + _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}", + "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}", + "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}", + "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}", + "description": f"Biological-replicate variance for {cond_label}" + _nan_note, } run_info = { @@ -353,6 +378,12 @@ def _expand_genes(values, n_cells, n_genes): "sample_col": sample_col, "n_genes": len(selected_genes), "n_cells": int(mask.sum()), + "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, }, diff --git a/tests/test_anndata_smooth.py b/tests/test_anndata_smooth.py index f513014..5a0ea2c 100644 --- a/tests/test_anndata_smooth.py +++ b/tests/test_anndata_smooth.py @@ -168,3 +168,118 @@ 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) +BAD = [0, 7, 41, 79] + + +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 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_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[~ok])) + assert np.all(np.isfinite(layer[ok])) + + 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() + 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, output=OutputSettings(progress=False)) + for key in ("kompot_smooth_all_smoothed", "kompot_smooth_all_std"): + np.testing.assert_allclose(adata.layers[key][ok], sub.layers[key], rtol=1e-6, atol=1e-8) + + 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_all_finite_control(self, adata): + from kompot.anndata.utils.json_utils import get_json_metadata + + 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_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)) + 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_with_condition(self, adata): + ok = _invalidate(adata) + kompot.smooth_expression( + 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][~ok])) + assert np.all(np.isfinite(adata.layers[key][ok])) + + 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[~ok])) and np.all(np.isfinite(layer[ok][:, :2])) + assert np.all(np.isnan(layer[:, 2:])) + + 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 b5b2730..0beb1e3 100644 --- a/tests/test_cli_smooth.py +++ b/tests/test_cli_smooth.py @@ -252,3 +252,34 @@ 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 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_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 + + 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" + 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"]) + 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", 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"]))