diff --git a/scenvi/utils.py b/scenvi/utils.py index 35b5e33..8b04f56 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,13 @@ 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 + # 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((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 +305,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 +320,29 @@ 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). 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 :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 +410,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) +