smooth_expression: cell_mask to mask cells out (NaN) instead of evaluating them [DO NOT MERGE]

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

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,
             },

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))

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"]))