COVET: make preprocessing explicit and batch_size numerically inert

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, 

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