Fix reverse key lookup

In the "match_gene_segment" function in _clonotype_neighbors, a "keys
array" is created to map from indicies in the reverse lookup table
(these refer to rows in the clonotype table) to numeric indices of
the reverse lookup table.

The size of this keys array was defined as the size of the clonotype
table. Instead it should have been the size of the distance matrix.

For real-world cases, this hardly was an issue, since the number
of clonotypes is typically larger than the number of columns
in the distance matrix.

Author: grst

src/scirpy/ir_dist/_clonotype_neighbors.py

@@ -329,7 +329,7 @@ def AND_max(a, b):
             if distance_matrix_name != distance_matrix_name_reverse:
                 raise ValueError("Forward and reverse lookup tablese must be defined on the same distance matrices.")
             reverse_lookup_values = np.vstack(list(reverse.lookup.values()))
-            reverse_lookup_keys = np.zeros(reverse.size, dtype=np.int64)
+            reverse_lookup_keys = np.zeros(reverse.dist_mat_size, dtype=np.int64)
             reverse_lookup_keys[list(reverse.lookup.keys())] = np.arange(len(list(reverse.lookup.keys())))
             match_column_mask = sp.csr_matrix(
                 (np.empty(len(has_distance_mask.indices)), has_distance_mask.indices, has_distance_mask.indptr),
@@ -419,8 +419,8 @@ def match_gene_segment(
                 raise ValueError("Forward and reverse lookup tablese must be defined on the same distance matrices.")
             empty_row = np.array([np.zeros(reverse.size, dtype=bool)])
             reverse_lookup_values = np.vstack((*reverse.lookup.values(), empty_row))
-            reverse_lookup_keys = np.full(id_len, -1, dtype=np.int32)
             keys_array = np.fromiter(reverse.lookup.keys(), dtype=int, count=len(reverse.lookup))
+            reverse_lookup_keys = np.full(reverse.dist_mat_size, -1, dtype=np.int32)
             reverse_lookup_keys[keys_array] = np.arange(len(keys_array))
             gene_segment_mask = sp.csr_matrix(
                 (np.empty(len(has_distance_mask.indices)), has_distance_mask.indices, has_distance_mask.indptr),

src/scirpy/ir_dist/_util.py

@@ -93,7 +93,7 @@ def reduce_and(*args, chain_count):
 
 
 class ReverseLookupTable:
-    def __init__(self, dist_type: Literal["boolean", "numeric"], size: int):
+    def __init__(self, dist_type: Literal["boolean", "numeric"], size: int, dist_mat: np.ndarray | sp.csr_matrix):
         """Reverse lookup table holds a mask that indicates which objects
         are neighbors of an object with a given index `i`.
 
@@ -113,18 +113,29 @@ def __init__(self, dist_type: Literal["boolean", "numeric"], size: int):
             Either `boolean` or `numeric`
         size
             The size of the masks.
+        dist_mat
+            A pointer to the associated distance matrix. This makes it easier to find the
+            corresponding distance matrix from an instance of the ReverseLookupTable
         """
         if dist_type not in ["boolean", "numeric"]:
             raise ValueError("invalid dist_type")
         self.dist_type = dist_type
         self.size = size
         self.lookup: dict[Hashable, sp.coo_matrix] = {}
+        self.dist_mat = dist_mat
+
+    @property
+    def dist_mat_size(self) -> int:
+        assert self.dist_mat.shape is not None
+        assert self.dist_mat.shape[0] == self.dist_mat.shape[1]
+        return self.dist_mat.shape[0]
 
     @staticmethod
     def from_dict_of_indices(
         dict_of_indices: Mapping,
         dist_type: Literal["boolean", "numeric"],
         size: int,
+        dist_mat: np.ndarray | sp.csr_matrix,
     ):
         """Convert a dict of indices to a ReverseLookupTable of row masks.
 
@@ -137,7 +148,7 @@ def from_dict_of_indices(
         size
             The size of the masks
         """
-        rlt = ReverseLookupTable(dist_type, size)
+        rlt = ReverseLookupTable(dist_type, size, dist_mat)
 
         # convert into coo matrices (numeric distances) or numpy boolean arrays.
         for k, v in dict_of_indices.items():
@@ -414,4 +425,6 @@ def _build_reverse_lookup_table(
             except KeyError:
                 tmp_reverse_lookup[tmp_key] = [i]
 
-        return ReverseLookupTable.from_dict_of_indices(tmp_reverse_lookup, dist_type, self.n_cols)
+        return ReverseLookupTable.from_dict_of_indices(
+            tmp_reverse_lookup, dist_type, self.n_cols, self.distance_matrices[distance_matrix]
+        )