Allow including self in pre-computed distance matrices

Author: Marius1311

src/cellmapper/cellmapper.py

@@ -539,34 +539,43 @@ def fit(
 
         return self
 
-    def load_precomputed_distances(self, distances_key: str = "distances") -> None:
+    def load_precomputed_distances(self, distances_key: str = "distances", include_self: bool | None = None) -> None:
         """
-        Load a pre-computed distance matrix from AnnData.obsp.
+        Load precomputed distances from the AnnData object.
+
+        This method is only available in self-mapping mode.
 
         Parameters
         ----------
         distances_key
-            Key in adata.obsp where the distance matrix is stored
+            Key in adata.obsp where the precomputed distances are stored.
+        include_self
+            If True, include self as a neighbor (even if not present in the distance matrix).
+            If False, exclude self connections (even if present in the distance matrix).
+            If None (default), preserve the original behavior of the distance matrix.
 
         Returns
         -------
         None
 
         Notes
         -----
-        This method can only be used in self-mapping mode (when CellMapper was
-        initialized with query=None).
+        Updates the following attributes:
+
+        - ``knn``: Neighbors object constructed from the precomputed distances.
         """
         if not self._is_self_mapping:
-            raise ValueError("Pre-computed distances can only be used in self-mapping mode")
+            raise ValueError("load_precomputed_distances is only available in self-mapping mode.")
 
-        if distances_key not in self.query.obsp:
-            raise KeyError(f"Distance matrix '{distances_key}' not found in query.obsp")
+        # Access the precomputed distances
+        distances_matrix = self.query.obsp[distances_key]
 
-        self.knn = Neighbors.from_distances(self.query.obsp[distances_key])
+        # Create a neighbors object using the factory method
+        self.knn = Neighbors.from_distances(distances_matrix, include_self=include_self)
 
         logger.info(
-            "Loaded pre-computed distance matrix from query.obsp['%s'] with %d cells",
+            "Loaded precomputed distances from '%s' with %d cells and %d neighbors per cell.",
             distances_key,
-            self.query.n_obs,
+            distances_matrix.shape[0],
+            self.knn.xx.n_neighbors,
         )

src/cellmapper/knn.py

@@ -292,22 +292,26 @@ def __init__(self, xrep: np.ndarray, yrep: np.ndarray | None = None):
         self._is_self_mapping = yrep is None
 
     @classmethod
-    def from_distances(cls, distances_matrix: "csr_matrix") -> "Neighbors":
+    def from_distances(cls, distances_matrix: "csr_matrix", include_self: bool | None = None) -> "Neighbors":
         """
         Create a Neighbors object from a pre-computed distances matrix.
 
         Parameters
         ----------
         distances_matrix
             Sparse distance matrix, typically from adata.obsp['distances']
+        include_self
+            If True, include self as a neighbor (cells are their own neighbors).
+            If False, exclude self connections, even if present in the distance matrix.
+            If None (default), preserve the original behavior of the distance matrix.
 
         Returns
         -------
         Neighbors
             A new Neighbors object with pre-computed neighbor information
         """
         # Extract indices and distances from the sparse matrix
-        indices, distances = extract_neighbors_from_distances(distances_matrix)
+        indices, distances = extract_neighbors_from_distances(distances_matrix, include_self=include_self)
 
         # Create a minimal Neighbors object for self-mapping
         n_cells = distances_matrix.shape[0]
@@ -323,6 +327,9 @@ def from_distances(cls, distances_matrix: "csr_matrix") -> "Neighbors":
         neighbors.xy = neighbors_result
         neighbors.yx = neighbors_result
 
+        # Mark as self-mapping
+        neighbors._is_self_mapping = True
+
         logger.info("Created Neighbors object from distances matrix with %d cells", n_cells)
 
         return neighbors

src/cellmapper/utils.py

@@ -123,14 +123,20 @@ def create_imputed_anndata(
     return imputed_adata
 
 
-def extract_neighbors_from_distances(distances_matrix: "csr_matrix") -> tuple[np.ndarray, np.ndarray]:
+def extract_neighbors_from_distances(
+    distances_matrix: "csr_matrix", include_self: bool | None = None
+) -> tuple[np.ndarray, np.ndarray]:
     """
     Extract neighbor indices and distances from a sparse distance matrix.
 
     Parameters
     ----------
     distances_matrix
         Sparse matrix of distances, typically from adata.obsp['distances']
+    include_self
+        If True, include self as a neighbor (even if not present in the distance matrix).
+        If False, exclude self connections (even if present in the distance matrix).
+        If None (default), preserve the original behavior of the distance matrix.
 
     Returns
     -------
@@ -139,44 +145,64 @@ def extract_neighbors_from_distances(distances_matrix: "csr_matrix") -> tuple[np
     """
     # Check that the input is a sparse matrix
     if not issparse(distances_matrix):
-        raise TypeError("Distances matrix must be sparse")
+        raise TypeError("Distances matrix must be a sparse matrix")
+
+    # Verify that the matrix is square
+    if distances_matrix.shape[0] != distances_matrix.shape[1]:
+        raise ValueError(f"Square distance matrix required (got {distances_matrix.shape})")
 
     n_cells = distances_matrix.shape[0]
 
-    # Get the number of neighbors per cell
-    n_neighbors_per_cell = np.diff(distances_matrix.indptr)
-    max_n_neighbors = n_neighbors_per_cell.max()
-    min_n_neighbors = n_neighbors_per_cell.min()
-
-    # Check if all cells have the same number of neighbors
-    if max_n_neighbors != min_n_neighbors:
-        logger.warning(
-            "Variable neighborhood sizes detected: min=%d, max=%d neighbors per cell. "
-            "Some cells may have fewer neighbors than others, which could affect results.",
-            min_n_neighbors,
-            max_n_neighbors,
-        )
+    # Ensure the matrix is CSR format for efficient row-based operations
+    distances_matrix = distances_matrix.tocsr()
+
+    # First pass: determine the max number of neighbors after including/excluding self
+    max_n_neighbors = 0
+    for i in range(n_cells):
+        start, end = distances_matrix.indptr[i], distances_matrix.indptr[i + 1]
+        cell_indices = distances_matrix.indices[start:end]
 
-    # Pre-allocate arrays for indices and distances
-    # Use -1 as a sentinel value for missing neighbors (better than 0 which is a valid index)
+        # Calculate how many neighbors this cell will have after applying include_self
+        n_neighbors = len(cell_indices)
+        if include_self is True and i not in cell_indices:
+            n_neighbors += 1  # Will add self
+        elif include_self is False and i in cell_indices:
+            n_neighbors -= 1  # Will remove self
+
+        max_n_neighbors = max(max_n_neighbors, n_neighbors)
+
+    # Pre-allocate arrays for indices and distances with the correct size
     indices = np.full((n_cells, max_n_neighbors), -1, dtype=np.int64)
     distances = np.full((n_cells, max_n_neighbors), np.inf, dtype=np.float64)
 
-    # Extract indices and distances for each cell
+    # Second pass: extract and process neighbor data
     for i in range(n_cells):
         # Get start and end indices for this cell in the sparse matrix
         start, end = distances_matrix.indptr[i], distances_matrix.indptr[i + 1]
 
-        # Number of neighbors for this cell
-        n_neighbors = end - start
+        # Get neighbor indices and distances
+        cell_indices = distances_matrix.indices[start:end]
+        cell_distances = distances_matrix.data[start:end]
+
+        # Filter self-connection if requested and present
+        if include_self is False and i in cell_indices:
+            # Find the index of self in the neighbors
+            self_idx = np.where(cell_indices == i)[0]
+            if len(self_idx) > 0:
+                # Remove self from indices and distances
+                mask = cell_indices != i
+                cell_indices = cell_indices[mask]
+                cell_distances = cell_distances[mask]
+        # If include_self is True and self is not in the neighbors, add it
+        elif include_self is True and i not in cell_indices:
+            # Add self with distance 0
+            cell_indices = np.append(cell_indices, i)
+            cell_distances = np.append(cell_distances, 0.0)
+
+        # Number of neighbors after potential filtering
+        n_neighbors = len(cell_indices)
 
         if n_neighbors > 0:
-            # Get neighbor indices
-            cell_indices = distances_matrix.indices[start:end]
-
-            # Get distances
-            cell_distances = distances_matrix.data[start:end]
-
             # Sort by distance if they aren't already sorted
             if not np.all(np.diff(cell_distances) >= 0):
                 sort_idx = np.argsort(cell_distances)

tests/test_cellmapper.py

@@ -415,3 +415,45 @@ def test_load_squidpy_distances(self, adata_spatial, squidpy_params):
 
         assert "leiden_pred" in cm.query.obs
         assert "leiden_conf" in cm.query.obs
+
+    @pytest.mark.parametrize("include_self", [True, False])
+    def test_load_distances_with_include_self(self, adata_spatial):
+        """Test loading precomputed distances with and without self-connections."""
+
+        # Compute neighbors with scanpy
+        sc.pp.neighbors(adata_spatial, n_neighbors=10, use_rep="X_pca")
+
+        # Initialize CellMapper in self-mapping mode
+        cm_with_self = CellMapper(adata_spatial)
+        cm_without_self = CellMapper(adata_spatial)
+
+        # Load precomputed distances with different include_self settings
+        cm_with_self.load_precomputed_distances(distances_key="distances", include_self=True)
+        cm_without_self.load_precomputed_distances(distances_key="distances", include_self=False)
+
+        # Verify that neighbors were loaded with or without self
+        assert cm_with_self.knn is not None
+        assert cm_without_self.knn is not None
+
+        # Check that with include_self=True, each cell has itself as a neighbor
+        for i in range(min(10, cm_with_self.knn.xx.n_samples)):  # Check first 10 cells
+            assert i in cm_with_self.knn.xx.indices[i]
+
+        # Check that with include_self=False, no cell has itself as a neighbor
+        for i in range(min(10, cm_without_self.knn.xx.n_samples)):  # Check first 10 cells
+            assert i not in cm_without_self.knn.xx.indices[i]
+
+        # Both should work with the rest of the pipeline
+        cm_with_self.compute_mappping_matrix(method="gaussian")
+        cm_without_self.compute_mappping_matrix(method="gaussian")
+
+        # Compute label transfer for both
+        cm_with_self.transfer_labels(obs_keys="leiden", prediction_postfix="with_self")
+        cm_without_self.transfer_labels(obs_keys="leiden", prediction_postfix="without_self")
+
+        # Both should have created prediction columns
+        assert "leiden_with_self" in adata_spatial.obs
+        assert "leiden_without_self" in adata_spatial.obs
+
+        # The results should be different (excluding self changes the neighborhood)
+        assert not adata_spatial.obs["leiden_with_self"].equals(adata_spatial.obs["leiden_without_self"])

tests/test_neighbors.py

@@ -76,7 +76,9 @@ def test_from_distances_factory_method(self, small_data):
         # Test with the neighbors results
         nn_results = neighbors.xx
         assert nn_results.n_samples == n_samples
-        assert nn_results.n_neighbors == n_samples
+        assert (
+            nn_results.n_neighbors + 1 == n_samples
+        )  # the distance to self is exactly 0, so it won't be included as a neighbor
 
         # Get adjacency matrix and verify it reflects the original distances
         adj_matrix = nn_results.knn_graph_distances
@@ -154,9 +156,10 @@ def test_from_distances_different_kernels(self, kernel):
                     distances_data[i, j] = abs(i - j)  # Simple metric: difference in indices
 
         distances = csr_matrix(distances_data)
+        print(distances)
 
         # Create Neighbors object
-        neighbors = Neighbors.from_distances(distances)
+        neighbors = Neighbors.from_distances(distances, include_self=True)
 
         # Compute connectivities with different kernels
         connectivities = neighbors.xx.knn_graph_connectivities(kernel=kernel)
@@ -167,8 +170,7 @@ def test_from_distances_different_kernels(self, kernel):
         assert np.all(connectivities.data > 0)
         # Diagonal should be large values (except for random kernel)
         if kernel != "random":
-            diag_indices = list(range(n_samples))
-            diag_values = connectivities[diag_indices, diag_indices].toarray().flatten()
+            diag_values = connectivities.diagonal()
             # Diagonal elements should typically be the largest for each row
             for i in range(n_samples):
                 row = connectivities.getrow(i).toarray().flatten()

tests/test_utils.py

@@ -86,3 +86,30 @@ def test_infinite_values(self):
         # Check that infinite values are excluded
         for i in range(4):
             assert np.all(np.isfinite(nbr_distances[i]))
+
+    def test_include_self_parameter(self):
+        """Test the include_self parameter to control self-connections."""
+        # Create a distance matrix with self-connections (diagonal = 0)
+        distances_data = np.array(
+            [[0.0, 1.0, 2.0, 3.0], [1.0, 0.0, 4.0, 5.0], [2.0, 4.0, 0.0, 6.0], [3.0, 5.0, 6.0, 0.0]]
+        )
+        distances = csr_matrix(distances_data)
+
+        # Test with include_self=True (default)
+        indices_with_self, distances_with_self = extract_neighbors_from_distances(distances, include_self=True)
+
+        # Test with include_self=False
+        indices_without_self, distances_without_self = extract_neighbors_from_distances(distances, include_self=False)
+
+        # With include_self=True, diagonal elements should be included
+        # (self should be the first neighbor because distance is 0)
+        for i in range(4):
+            assert indices_with_self[i, 0] == i
+            assert distances_with_self[i, 0] == 0.0
+
+        # With include_self=False, diagonal elements should be excluded
+        for i in range(4):
+            # The self-index should not be in the neighbors
+            assert i not in indices_without_self[i]
+            # Check that no zero distances are present
+            assert np.all(distances_without_self[i] > 0)