More data type fixes and improve speed of generating random pairs.

Author: newmana

stlearn/tl/cci/analysis.py

@@ -259,7 +259,7 @@ def run(
         Whether print dialogue to user during run-time.
     Returns
     --------
-    adata: None
+    None
     Relevant information stored:
         adata.uns['lr_summary']
             Summary of significant spots detected per LR,
@@ -398,7 +398,7 @@ def adj_pvals(
     -------
     adata: AnnData
         Adjusts all of the LR results; warning, does not adjust
-        celltype-celltype results from running ran st.tl.run_cci downstream.
+        celltype-celltype results from running st.tl.run_cci downstream.
     """
     if "lr_summary" not in adata.uns:
         raise Exception("Need to run st.tl.cci.run first.")

stlearn/tl/cci/base.py

@@ -42,13 +42,13 @@ def calc_distance(adata: AnnData, distance: float | None) -> float:
 
 def get_lrs_scores(
     adata: AnnData,
-    lrs: np.ndarray,
+    lrs: npt.NDArray[np.str_],
     neighbours: np.ndarray,
     het_vals: np.ndarray,
     min_expr: float,
     filter_pairs: bool = True,
     spot_indices: npt.NDArray[np.int32] | None = None,
-):
+) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.str_]]:
     """Gets the scores for the indicated set of LR pairs & the heterogeneity values.
     Parameters
     ----------
@@ -69,7 +69,12 @@ def get_lrs_scores(
         Subset of spots to score, given as their integer row positions.
     Returns
     -------
-    lrs: np.ndarray   lr pairs from the database in format ['L1_R1', 'LN_RN']
+    lr_scores: npt.NDArray[np.float64]
+        Shape (n_scored_spots, n_pairs). LR score for each scored spot (rows in
+        the order of spot_indices) and each ligand-receptor pair (columns).
+    new_lrs: npt.NDArray[np.str_]
+        Shape (n_pairs,). Ligand-receptor pair labels in 'L_R' format
+        (e.g. 'L1_R1'), column-aligned with lr_scores.
     """
     if spot_indices is None:
         spot_indices = np.arange(len(adata), dtype=np.int32)

stlearn/tl/cci/perm_utils.py

@@ -4,6 +4,9 @@
 from numba.typed import List
 from scipy.spatial.distance import canberra
 from sklearn.preprocessing import MinMaxScaler
+from numba.core import types
+import numpy.typing as npt
+from numba.typed import Dict, List
 
 from .base import get_lrs_scores
 
@@ -264,23 +267,54 @@ def get_similar_genes_fast(
     return similar_genes
 
 
-@njit
-def gen_rand_pairs(genes1: np.ndarray, genes2: np.ndarray, n_pairs: int, seed: int):
-    """Generates random pairs of genes."""
+def gen_rand_pairs(
+    genes1: npt.NDArray[np.str_],
+    genes2: npt.NDArray[np.str_],
+    n_pairs: int,
+    seed: int,
+) -> npt.NDArray[np.str_]:
+    """Generate unique random gene pairs for building background LR scores.
+
+    Each pair is formed by drawing one gene from genes1 and one from genes2,
+    formatted as 'gene1_gene2'. Self-pairs (the same gene on both sides) and
+    duplicate pairs are rejected, so every returned pair is unique and its two
+    genes are distinct.
+
+    Parameters
+    ----------
+    genes1: npt.NDArray[np.str_]
+        Candidate genes for the first (ligand) position of each pair.
+    genes2: npt.NDArray[np.str_]
+        Candidate genes for the second (receptor) position of each pair.
+    n_pairs: int
+        Number of unique pairs to generate.
+    seed: int
+        Seed for the random generator, for reproducible pair selection.
+    """
+    n_possible = len(genes1) * len(genes2) - np.intersect1d(genes1, genes2).size
+    if n_pairs > n_possible:
+        raise ValueError(
+            f"Requested {n_pairs} unique pairs but only {n_possible} are "
+            f"possible from {len(genes1)}×{len(genes2)} genes."
+        )
+    return np.array(list(_gen_rand_pairs(genes1, genes2, n_pairs, seed)))
+
 
+@njit
+def _gen_rand_pairs(genes1, genes2, n_pairs, seed):
     np.random.seed(seed)  # noqa: NPY002 (numba requires legacy API)
     rand_pairs = List()
+    seen = Dict.empty(types.unicode_type, types.boolean)  # O(1) membership
     for _j in range(0, n_pairs):
         l_rand = np.random.choice(genes1, 1)[0]  # noqa: NPY002
         r_rand = np.random.choice(genes2, 1)[0]  # noqa: NPY002
         rand_pair = "_".join([l_rand, r_rand])
-        while rand_pair in rand_pairs or l_rand == r_rand:
+        while rand_pair in seen or l_rand == r_rand:        # was: in rand_pairs
             l_rand = np.random.choice(genes1, 1)[0]  # noqa: NPY002
             r_rand = np.random.choice(genes2, 1)[0]  # noqa: NPY002
             rand_pair = "_".join([l_rand, r_rand])
-
         rand_pairs.append(rand_pair)
-
+        seen[rand_pair] = True
     return rand_pairs
 
 

stlearn/tl/cci/permutation.py

@@ -13,7 +13,7 @@
 
 def perform_spot_testing(
     adata: AnnData,
-    lr_scores: np.ndarray,
+    lr_scores: npt.NDArray[np.float64],
     lrs: npt.NDArray[np.str_],
     n_pairs: int,
     neighbours: List,
@@ -145,14 +145,15 @@ def perform_spot_testing(
             else:  # Fitting NB per LR
                 lr_j_scores = lr_scores[spot_indices, lr_j]
                 bg_ = background.ravel()
-                bg_wScore = np.array(list(lr_j_scores) + list(bg_))
+                bg_w_score = np.array(list(lr_j_scores) + list(bg_))
 
                 # 1) rounding discretisation
                 # First multiple to get minimum value to be one before rounding #
-                bg_1 = bg_wScore * (1 / min(bg_wScore[bg_wScore != 0]))
-                bg_1 = np.round(bg_1)
-                lr_j_scores_1 = bg_1[0 : len(lr_j_scores)]
-                bg_1 = bg_1[len(lr_j_scores) : len(bg_1)]
+                scale = 1.0 / min(bg_w_score[bg_w_score != 0])
+                scaled = np.round(bg_w_score * scale)
+                n_obs = len(lr_j_scores)
+                lr_j_scores_1 = scaled[:n_obs]
+                bg_1 = scaled[n_obs:]
 
                 # Getting the pvalue from negative binomial approach
                 round_pvals, _, _, _ = get_stats(