From 54424b40c80029af30d5876a1ad27a7fd049eef1 Mon Sep 17 00:00:00 2001 From: Marius1311 Date: Tue, 4 Aug 2026 15:09:59 +0200 Subject: [PATCH 1/2] COVET: make preprocessing explicit and batch_size inert Three things in compute_covet surprise callers who want COVET on data they have already prepared, which the README advertises as a supported use. 1. `X` is log-transformed a second time. The only test for "already log-transformed" is `spatial_data.X.min() < 0`. Log-normalized data is non-negative, so it is read as raw counts and gets `log(x + 1)` applied on top. Measured on log-normalized input, this moves COVET by 152% relative -- same genes, same graph, same estimator. Adds `log_transform=None|True|False`. `None` keeps the current heuristic, so existing results stay reproducible, but warns and says how to override. `use_obsm`/`use_layer` keep meaning "take this at face value". The upstream test suite itself triggers the warning, in test_covet_with_batches. 2. A missing `batch_key` silently pools the kNN across samples. `if batch_key not in spatial_data.obs.columns: batch_key = -1` cannot tell a typo from an absent default. Passing `batch_key="smaple"` builds one graph over all sections and returns niches that straddle samples, with nothing to indicate it. Now only the default `"batch"` falls back; anything else raises and lists the available columns. `batch_key=-1` is unaffected. 3. `batch_size` changes the numbers. The accumulator is a fixed float32 buffer while the unbatched branch takes its dtype from `np.matmul` on the centered data. Passing `batch_size` therefore accumulated in single precision before the regularization term (derived from this stack) and before the matrix square root. Both outputs are cast to float32 on return, so the dtype difference is invisible, but the values differ: ~7e-8 relative on COVET and ~5e-6 absolute on COVET_SQRT. A memory knob should be numerically inert. The buffer now matches the unbatched dtype. Adds tests/test_covet_preprocessing.py covering all three. Ten of its eleven tests fail against the current implementation. Co-Authored-By: Claude Opus 5 (1M context) --- scenvi/utils.py | 85 ++++++++++++++++------ tests/test_covet_preprocessing.py | 113 ++++++++++++++++++++++++++++++ 2 files changed, 178 insertions(+), 20 deletions(-) create mode 100644 tests/test_covet_preprocessing.py diff --git a/scenvi/utils.py b/scenvi/utils.py index 35b5e33..a0facae 100644 --- a/scenvi/utils.py +++ b/scenvi/utils.py @@ -1,3 +1,5 @@ +import warnings + import jax import jax.numpy as jnp import numpy as np @@ -211,10 +213,17 @@ def calculate_covariance_matrices(spatial_data, kNN, exp_data, spatial_key="spat # Get the global mean for each feature global_mean = exp_data.mean(axis=0) - # Initialize the output covariance matrices + # Initialize the output covariance matrices. The dtype has to match what the + # unbatched branch below produces (np.matmul on the centered data). With a fixed + # float32 buffer, passing batch_size instead accumulated in single precision -- + # before the regularization term, which is derived from this stack, and before + # the matrix square root -- so a pure memory knob moved the numbers. n_cells = exp_data.shape[0] n_features = exp_data.shape[1] - CovMats = np.zeros((n_cells, n_features, n_features), dtype=np.float32) + CovMats = np.zeros( + (n_cells, n_features, n_features), + dtype=np.promote_types(exp_data.dtype, np.float32), + ) # Process in batches if requested if batch_size is None or batch_size >= n_cells: @@ -300,8 +309,8 @@ def niche_cell_type( return cell_type_niche def compute_covet( - spatial_data, k=8, g=64, genes=None, spatial_key="spatial", batch_key="batch", - batch_size=None, use_obsm=None, use_layer=None + spatial_data, k=8, g=64, genes=None, spatial_key="spatial", batch_key="batch", + batch_size=None, use_obsm=None, use_layer=None, log_transform=None ): """ Compute niche covariance matrices for spatial data, run with scenvi.compute_covet @@ -315,18 +324,31 @@ def compute_covet( :param batch_size: (int) Number of cells/spots to process at once for large datasets (default None) :param use_obsm: (str) obsm key to use for COVET calculation instead of gene expression (e.g. 'X_pca', 'X_dc') (default None) :param use_layer: (str) layer to use for COVET calculation instead of log-transformed X (e.g. 'log', 'log1p') (default None) - + :param log_transform: (bool) whether to apply log(x + 1) to the selected expression data. + None (default) keeps the historical behaviour: data taken from `X` is log-transformed + unless it contains negative values, and data taken from `use_obsm`/`use_layer` is used + as is. Pass True or False to decide explicitly. + :return COVET: niche covariance matrices :return COVET_SQRT: matrix square-root of niche covariance matrices for approximate OT :return CovGenes: list of genes selected for COVET representation (or feature names if using obsm) """ genes = [] if genes is None else genes - - # Handle batch key - if batch_key not in spatial_data.obs.columns: - batch_key = -1 - + + # Handle batch key. Falling back to a single pooled kNN is only safe for the + # default: an explicitly requested batch_key that is missing is a mistake, and + # silently building niches that straddle samples gives no sign of it. + if batch_key != -1 and batch_key not in spatial_data.obs.columns: + if batch_key == "batch": + batch_key = -1 + else: + raise ValueError( + f"batch_key '{batch_key}' is not a column of spatial_data.obs. " + f"Available columns: {list(spatial_data.obs.columns)}. " + "Pass batch_key=-1 to build a single kNN graph across all cells." + ) + # Determine data source: obsm, layer, or X if use_obsm is not None: if use_obsm not in spatial_data.obsm: @@ -394,17 +416,40 @@ def compute_covet( if use_layer not in spatial_data.layers: raise ValueError(f"Layer '{use_layer}' not found in spatial_data.layers") print(f"Using expression data from layer '{use_layer}'") - exp_data = spatial_data[:, CovGenes].layers[use_layer].toarray() if scipy.sparse.issparse(spatial_data.layers[use_layer]) else spatial_data[:, CovGenes].layers[use_layer] + exp_data = spatial_data[:, CovGenes].layers[use_layer] else: - # Default: log-transform X if needed - if spatial_data.X.min() < 0: - # Data is already log-transformed - print("Using expression data from X (appears to be log-transformed)") - exp_data = spatial_data[:, CovGenes].X.toarray() if scipy.sparse.issparse(spatial_data.X) else spatial_data[:, CovGenes].X - else: - print("Log-transforming expression data from X") - exp_data = np.log(spatial_data[:, CovGenes].X.toarray() + 1) if scipy.sparse.issparse(spatial_data.X) else np.log(spatial_data[:, CovGenes].X + 1) - + print("Using expression data from X") + exp_data = spatial_data[:, CovGenes].X + + if scipy.sparse.issparse(exp_data): + exp_data = exp_data.toarray() + + # Decide on the log transform once, now that the data source is settled. + if log_transform is None: + if use_obsm is not None or use_layer is not None: + # Handed over explicitly, so take it at face value. + log_transform = False + else: + # Historical heuristic. It reads non-negative data as raw counts, but + # log-normalized data is non-negative too, so it silently log-transforms + # a second time -- which changes COVET substantially. Warn rather than + # change the default, so existing results stay reproducible. + log_transform = spatial_data.X.min() >= 0 + if log_transform: + warnings.warn( + "compute_covet is applying log(x + 1) to spatial_data.X because it " + "contains no negative values. If X is already normalized, this " + "log-transforms it a second time and substantially changes COVET. " + "Pass log_transform=False to use X as is, or log_transform=True to " + "silence this warning.", + UserWarning, + stacklevel=2, + ) + + if log_transform: + print("Log-transforming expression data") + exp_data = np.log(exp_data + 1) + # Calculate covariance matrices with batch processing COVET = calculate_covariance_matrices( spatial_data, k, exp_data, spatial_key=spatial_key, diff --git a/tests/test_covet_preprocessing.py b/tests/test_covet_preprocessing.py new file mode 100644 index 0000000..d258247 --- /dev/null +++ b/tests/test_covet_preprocessing.py @@ -0,0 +1,113 @@ +"""compute_covet's preprocessing decisions are explicit and consistent.""" + +import warnings + +import anndata +import numpy as np +import pytest + +import scenvi + +K = 6 + + +@pytest.fixture +def counts_data(): + """Raw counts: non-negative, so the historical heuristic log-transforms them.""" + rng = np.random.default_rng(0) + return anndata.AnnData( + X=rng.poisson(5.0, size=(64, 8)).astype(np.float64), + obsm={"spatial": rng.normal(size=(64, 2))}, + ) + + +@pytest.fixture +def batched_data(): + rng = np.random.default_rng(1) + adata = anndata.AnnData( + X=rng.poisson(5.0, size=(64, 8)).astype(np.float64), + obsm={"spatial": rng.normal(size=(64, 2))}, + ) + adata.obs["sample"] = ["a"] * 32 + ["b"] * 32 + return adata + + +class TestLogTransform: + def test_heuristic_warns_when_it_log_transforms(self, counts_data): + with pytest.warns(UserWarning, match="log_transform=False"): + scenvi.compute_covet(counts_data, k=K, g=8, batch_key=-1) + + @pytest.mark.parametrize("log_transform", [True, False]) + def test_explicit_choice_does_not_warn(self, counts_data, log_transform): + with warnings.catch_warnings(): + warnings.simplefilter("error", UserWarning) + scenvi.compute_covet(counts_data, k=K, g=8, batch_key=-1, log_transform=log_transform) + + def test_log_transform_false_matches_pre_logged_input(self, counts_data): + """log_transform=False must use X verbatim, i.e. agree with logging it ourselves.""" + expected, _, _ = scenvi.compute_covet(counts_data, k=K, g=8, batch_key=-1, log_transform=True) + + logged = counts_data.copy() + logged.X = np.log(counts_data.X + 1) + actual, _, _ = scenvi.compute_covet(logged, k=K, g=8, batch_key=-1, log_transform=False) + + np.testing.assert_allclose(actual, expected, rtol=1e-10, atol=1e-12) + + def test_default_is_unchanged_for_existing_users(self, counts_data): + """The heuristic still fires; this PR warns about it, it does not alter it.""" + with pytest.warns(UserWarning): + default, _, _ = scenvi.compute_covet(counts_data, k=K, g=8, batch_key=-1) + explicit, _, _ = scenvi.compute_covet(counts_data, k=K, g=8, batch_key=-1, log_transform=True) + + np.testing.assert_allclose(default, explicit, rtol=1e-10, atol=1e-12) + + def test_use_obsm_is_taken_at_face_value(self, counts_data): + counts_data.obsm["feat"] = np.asarray(counts_data.X[:, :4], dtype=np.float64) + covet, _, _ = scenvi.compute_covet(counts_data, k=K, use_obsm="feat", batch_key=-1) + + logged = counts_data.copy() + logged.obsm["feat"] = np.log(counts_data.obsm["feat"] + 1) + covet_logged, _, _ = scenvi.compute_covet(logged, k=K, use_obsm="feat", batch_key=-1) + + assert not np.allclose(covet, covet_logged) + + +class TestBatchKey: + def test_missing_explicit_batch_key_raises(self, batched_data): + with pytest.raises(ValueError, match="not a column of spatial_data.obs"): + scenvi.compute_covet(batched_data, k=K, g=8, batch_key="smaple", log_transform=False) + + def test_default_batch_key_still_falls_back(self, counts_data): + """`batch_key='batch'` with no such column keeps meaning "no batches".""" + covet, _, _ = scenvi.compute_covet(counts_data, k=K, g=8, log_transform=False) + assert covet.shape == (counts_data.n_obs, 8, 8) + + def test_explicit_minus_one_is_accepted(self, batched_data): + covet, _, _ = scenvi.compute_covet(batched_data, k=K, g=8, batch_key=-1, log_transform=False) + assert covet.shape == (batched_data.n_obs, 8, 8) + + def test_batches_change_the_result(self, batched_data): + pooled, _, _ = scenvi.compute_covet(batched_data, k=K, g=8, batch_key=-1, log_transform=False) + per_batch, _, _ = scenvi.compute_covet(batched_data, k=K, g=8, batch_key="sample", log_transform=False) + assert not np.allclose(pooled, per_batch) + + +class TestBatchSize: + def test_batch_size_does_not_change_the_result(self, counts_data): + """batch_size is a memory knob and must be numerically inert. + + Both results are cast to float32 on return, so the dtype difference is + invisible -- but the batched path used to accumulate in float32 before the + regularization term and the square root were computed, which moved COVET by + ~7e-8 relative and COVET_SQRT by ~6e-6. + """ + whole, sqrt_whole, _ = scenvi.compute_covet( + counts_data, k=K, g=8, batch_key=-1, log_transform=False + ) + chunked, sqrt_chunked, _ = scenvi.compute_covet( + counts_data, k=K, g=8, batch_key=-1, log_transform=False, batch_size=16 + ) + + np.testing.assert_allclose(chunked, whole, rtol=1e-12, atol=1e-14) + np.testing.assert_allclose(sqrt_chunked, sqrt_whole, rtol=1e-10, atol=1e-12) + From 2a88e6de352a6b251c1293ac7b53f0c445a5f197 Mon Sep 17 00:00:00 2001 From: Marius1311 Date: Wed, 5 Aug 2026 10:28:59 +0200 Subject: [PATCH 2/2] Address review: trim the dtype comment and log_transform docstring Both per review on #92. The dtype comment retold the whole bug; one line naming the invariant (batch_size must not move the numbers) is enough -- the reasoning that motivated it is in the commit that introduced it. The log_transform docstring restated what the code comment at the heuristic already says, so it is cut back to what the parameter does. Co-Authored-By: Claude Opus 5 (1M context) --- scenvi/utils.py | 12 +++--------- 1 file changed, 3 insertions(+), 9 deletions(-) diff --git a/scenvi/utils.py b/scenvi/utils.py index a0facae..8b04f56 100644 --- a/scenvi/utils.py +++ b/scenvi/utils.py @@ -213,11 +213,7 @@ def calculate_covariance_matrices(spatial_data, kNN, exp_data, spatial_key="spat # Get the global mean for each feature global_mean = exp_data.mean(axis=0) - # Initialize the output covariance matrices. The dtype has to match what the - # unbatched branch below produces (np.matmul on the centered data). With a fixed - # float32 buffer, passing batch_size instead accumulated in single precision -- - # before the regularization term, which is derived from this stack, and before - # the matrix square root -- so a pure memory knob moved the numbers. + # Match the unbatched branch's dtype so that batch_size stays numerically inert. n_cells = exp_data.shape[0] n_features = exp_data.shape[1] CovMats = np.zeros( @@ -324,10 +320,8 @@ def compute_covet( :param batch_size: (int) Number of cells/spots to process at once for large datasets (default None) :param use_obsm: (str) obsm key to use for COVET calculation instead of gene expression (e.g. 'X_pca', 'X_dc') (default None) :param use_layer: (str) layer to use for COVET calculation instead of log-transformed X (e.g. 'log', 'log1p') (default None) - :param log_transform: (bool) whether to apply log(x + 1) to the selected expression data. - None (default) keeps the historical behaviour: data taken from `X` is log-transformed - unless it contains negative values, and data taken from `use_obsm`/`use_layer` is used - as is. Pass True or False to decide explicitly. + :param log_transform: (bool) whether to apply log(x + 1). None (default) estimates whether + the data is already log-transformed, True/False decide explicitly (default None) :return COVET: niche covariance matrices :return COVET_SQRT: matrix square-root of niche covariance matrices for approximate OT