bootstrapping

Author: selmanozleyen

src/rapids_singlecell/pertpy_gpu/_distances.py

@@ -523,7 +523,7 @@ def _bootstrap_mode_precomputed(
             distance = self.metric_fct.from_precomputed(
                 bootstrap_sub_pwd, bootstrap_sub_idx, **kwargs
             )
-            distances.append(distance.get())
+            distances.append(distance)
 
         mean = np.mean(distances)
         variance = np.var(distances)

src/rapids_singlecell/pertpy_gpu/_distances_standalone.py

@@ -35,7 +35,7 @@ def _load_edistance_kernels():
 compute_group_distances_kernel = _load_edistance_kernels()
 
 
-def compute_d_other_gpu(
+def compute_pairwise_means_gpu(
     embedding: cp.ndarray, cat_offsets: cp.ndarray, cell_indices: cp.ndarray, k: int
 ) -> cp.ndarray:
     """
@@ -122,12 +122,177 @@ def compute_d_other_gpu(
     return pairwise_means
 
 
+def generate_bootstrap_indices(
+    cat_offsets: cp.ndarray,
+    k: int,
+    n_bootstrap: int = 100,
+    random_state: int = 0,
+) -> list[list[cp.ndarray]]:
+    """
+    Generate bootstrap indices for all groups and all bootstrap iterations.
+    This matches the CPU implementation's random sampling logic for reproducibility.
+    
+    Parameters
+    ----------
+    cat_offsets : cp.ndarray
+        Group start/end indices
+    k : int
+        Number of groups
+    n_bootstrap : int
+        Number of bootstrap samples
+    random_state : int
+        Random seed for reproducibility
+        
+    Returns
+    -------
+    bootstrap_indices : list[list[cp.ndarray]]
+        For each bootstrap iteration, list of indices arrays for each group
+        Shape: [n_bootstrap][k] where each element is cp.ndarray of group_size
+    """
+    import numpy as np
+    
+    # Use same RNG logic as CPU code
+    rng = np.random.default_rng(random_state)
+    
+    # Convert to numpy for CPU-based random generation
+    cat_offsets_np = cat_offsets.get()
+    
+    bootstrap_indices = []
+    
+    for bootstrap_iter in range(n_bootstrap):
+        group_indices = []
+        
+        for group_idx in range(k):
+            start_idx = cat_offsets_np[group_idx]
+            end_idx = cat_offsets_np[group_idx + 1]
+            group_size = end_idx - start_idx
+            
+            if group_size > 0:
+                # Generate bootstrap indices using same logic as CPU code
+                # rng.choice(a=X.shape[0], size=X.shape[0], replace=True)
+                bootstrap_group_indices = rng.choice(
+                    group_size, size=group_size, replace=True
+                )
+                # Convert to CuPy array
+                group_indices.append(cp.array(bootstrap_group_indices, dtype=cp.int32))
+            else:
+                # Empty group
+                group_indices.append(cp.array([], dtype=cp.int32))
+        
+        bootstrap_indices.append(group_indices)
+    
+    return bootstrap_indices
+
+
+def _bootstrap_sample_cells_from_indices(
+    *,
+    cat_offsets: cp.ndarray,
+    cell_indices: cp.ndarray,
+    k: int,
+    bootstrap_group_indices: list[cp.ndarray],
+) -> tuple[cp.ndarray, cp.ndarray]:
+    """
+    Bootstrap sample cells using pre-generated indices.
+    
+    Parameters
+    ----------
+    cat_offsets : cp.ndarray
+        Group start/end indices
+    cell_indices : cp.ndarray
+        Sorted cell indices by group
+    k : int
+        Number of groups
+    bootstrap_group_indices : list[cp.ndarray]
+        Pre-generated bootstrap indices for each group
+        
+    Returns
+    -------
+    new_cat_offsets, new_cell_indices : tuple[cp.ndarray, cp.ndarray]
+        New category structure with bootstrapped cells
+    """
+    new_cell_indices = []
+    new_cat_offsets = cp.zeros(k + 1, dtype=cp.int32)
+    
+    for group_idx in range(k):
+        start_idx = cat_offsets[group_idx]
+        end_idx = cat_offsets[group_idx + 1]
+        group_size = end_idx - start_idx
+        
+        if group_size > 0:
+            # Get original cell indices for this group
+            group_cells = cell_indices[start_idx:end_idx]
+            
+            # Use pre-generated bootstrap indices
+            bootstrap_indices = bootstrap_group_indices[group_idx]
+            bootstrap_cells = group_cells[bootstrap_indices]
+            
+            new_cell_indices.extend(bootstrap_cells.get().tolist())
+        
+        new_cat_offsets[group_idx + 1] = len(new_cell_indices)
+    
+    return new_cat_offsets, cp.array(new_cell_indices, dtype=cp.int32)
+
+
+def compute_pairwise_means_gpu_bootstrap(
+    embedding: cp.ndarray,
+    *,
+    cat_offsets: cp.ndarray,
+    cell_indices: cp.ndarray,
+    k: int,
+    n_bootstrap: int = 100,
+    random_state: int = 0,
+) -> tuple[cp.ndarray, cp.ndarray]:
+    """
+    Compute bootstrap statistics for between-group distances.
+    Uses CPU-compatible random generation for reproducibility.
+    
+    Returns:
+        means: [k, k] matrix of bootstrap means
+        variances: [k, k] matrix of bootstrap variances
+    """
+    # Generate all bootstrap indices upfront using CPU-compatible logic
+    bootstrap_indices = generate_bootstrap_indices(
+        cat_offsets, k, n_bootstrap, random_state
+    )
+    
+    bootstrap_results = []
+    
+    for bootstrap_iter in range(n_bootstrap):
+        # Use pre-generated indices for this bootstrap iteration
+        boot_cat_offsets, boot_cell_indices = _bootstrap_sample_cells_from_indices(
+            cat_offsets=cat_offsets,
+            cell_indices=cell_indices,
+            k=k,
+            bootstrap_group_indices=bootstrap_indices[bootstrap_iter],
+        )
+        
+        # Compute distances with bootstrapped samples
+        pairwise_means = compute_pairwise_means_gpu(
+            embedding=embedding,
+            cat_offsets=boot_cat_offsets,
+            cell_indices=boot_cell_indices,
+            k=k,
+        )
+        bootstrap_results.append(pairwise_means.get())
+    
+    # Compute statistics across bootstrap samples
+    bootstrap_stack = cp.array(bootstrap_results)  # [n_bootstrap, k, k]
+    means = cp.mean(bootstrap_stack, axis=0)
+    variances = cp.var(bootstrap_stack, axis=0)
+    
+    return means, variances
+
+
 def pairwise_edistance_gpu(
     adata: AnnData,
     groupby: str,
     *,
     obsm_key: str = "X_pca",
     groups: list[str] | None = None,
+    inplace: bool = False,
+    bootstrap: bool = False,
+    n_bootstrap: int = 100,
+    random_state: int = 0,
 ) -> pd.DataFrame:
     """
     GPU-accelerated pairwise edistance computation with decomposed components.
@@ -169,9 +334,117 @@ def pairwise_edistance_gpu(
     k = len(group_map)
     cat_offsets, cell_indices = _create_category_index_mapping(group_labels, k)
 
+    groups_list = (
+        list(original_groups.cat.categories.values) if groups is None else groups
+    )
+    if not bootstrap:
+        df = compute_pairwise_means_gpu_edistance(
+            embedding=embedding,
+            cat_offsets=cat_offsets,
+            cell_indices=cell_indices,
+            k=k,
+            groups_list=groups_list,
+            groupby=groupby,
+        )
+        if inplace:
+            adata.uns[f"{groupby}_pairwise_edistance"] = {
+                "distances": df,
+            }
+        return df
+
+    else:
+        df, df_var = compute_pairwise_means_gpu_edistance_bootstrap(
+            embedding=embedding,
+            cat_offsets=cat_offsets,
+            cell_indices=cell_indices,
+            k=k,
+            groups_list=groups_list,
+            groupby=groupby,
+            n_bootstrap=n_bootstrap,
+            random_state=random_state,
+        )
+
+        if inplace:
+            adata.uns[f"{groupby}_pairwise_edistance"] = {
+                "distances": df,
+                "distances_var": df_var,
+            }
+        return df, df_var
+
+
+def compute_pairwise_means_gpu_edistance_bootstrap(
+    embedding: cp.ndarray,
+    *,
+    cat_offsets: cp.ndarray,
+    cell_indices: cp.ndarray,
+    k: int,
+    groups_list: list[str],
+    groupby: str,
+    n_bootstrap: int = 100,
+    random_state: int = 0,
+) -> tuple[pd.DataFrame, pd.DataFrame]:
+    # Bootstrap computation
+    pairwise_means_boot, pairwise_vars_boot = compute_pairwise_means_gpu_bootstrap(
+        embedding=embedding,
+        cat_offsets=cat_offsets,
+        cell_indices=cell_indices,
+        k=k,
+        n_bootstrap=n_bootstrap,
+        random_state=random_state,
+    )
+
+    # 4. Compute final edistance for means and variances
+    edistance_means = cp.zeros((k, k), dtype=np.float32)
+    edistance_vars = cp.zeros((k, k), dtype=np.float32)
+
+    for a in range(k):
+        for b in range(a + 1, k):
+            # Bootstrap mean edistance
+            edistance_means[a, b] = (
+                2 * pairwise_means_boot[a, b]
+                - pairwise_means_boot[a, a]
+                - pairwise_means_boot[b, b]
+            )
+            edistance_means[b, a] = edistance_means[a, b]
+
+            # Bootstrap variance edistance (using delta method approximation)
+            # Var(2*X - Y - Z) = 4*Var(X) + Var(Y) + Var(Z) (assuming independence)
+            edistance_vars[a, b] = (
+                4 * pairwise_vars_boot[a, b]
+                + pairwise_vars_boot[a, a]
+                + pairwise_vars_boot[b, b]
+            )
+            edistance_vars[b, a] = edistance_vars[a, b]
+
+    # 5. Create output DataFrames
+
+    df_mean = pd.DataFrame(
+        edistance_means.get(), index=groups_list, columns=groups_list
+    )
+    df_mean.index.name = groupby
+    df_mean.columns.name = groupby
+    df_mean.name = "pairwise edistance"
+
+    df_var = pd.DataFrame(edistance_vars.get(), index=groups_list, columns=groups_list)
+    df_var.index.name = groupby
+    df_var.columns.name = groupby
+    df_var.name = "pairwise edistance variance"
+
+    return df_mean, df_var
+
+
+def compute_pairwise_means_gpu_edistance(
+    embedding: cp.ndarray,
+    *,
+    cat_offsets: cp.ndarray,
+    cell_indices: cp.ndarray,
+    k: int,
+    groups_list: list[str],
+    groupby: str,
+) -> pd.DataFrame:
     # 3. Compute decomposed components
     # d_itself = compute_d_itself_gpu(embedding, cat_offsets, cell_indices, k)
-    pairwise_means = compute_d_other_gpu(embedding, cat_offsets, cell_indices, k)
+    pairwise_means = compute_pairwise_means_gpu(embedding, cat_offsets, cell_indices, k)
 
     # 4. Compute final edistance: df[a,b] = 2*d_other[a,b] - d_itself[a] - d_itself[b]
     edistance_matrix = cp.zeros((k, k), dtype=np.float32)
@@ -183,17 +456,10 @@ def pairwise_edistance_gpu(
             edistance_matrix[b, a] = edistance_matrix[a, b]
 
     # 5. Create output DataFrame
-    groups_list = (
-        list(original_groups.cat.categories.values) if groups is None else groups
-    )
+
     df = pd.DataFrame(edistance_matrix.get(), index=groups_list, columns=groups_list)
     df.index.name = groupby
     df.columns.name = groupby
     df.name = "pairwise edistance"
 
-    # Store in adata
-    adata.uns[f"{groupby}_pairwise_edistance"] = {
-        "distances": df,
-    }
-
     return df

tmp_scripts/compute_edistance.py

@@ -2,6 +2,7 @@
 
 import os
 import time
+from pathlib import Path
 
 import anndata as ad
 import cupy as cp
@@ -24,9 +25,10 @@
 
     # homedir/data/adamson_2016_upr_epistasis
     save_dir = os.path.join(
-        os.path.expanduser("~"), "data", "adamson_2016_upr_epistasis_pca.h5ad"
+        os.path.expanduser("~"),
+        "data",
     )
-    adata = ad.read_h5ad(save_dir)
+    adata = ad.read_h5ad(Path(save_dir) / "adamson_2016_upr_epistasis_pca.h5ad")
     rsc.get.anndata_to_GPU(adata, convert_all=True)
     dist = Distance(obsm_key="X_pca", metric="edistance")
     start_time = time.time()

tmp_scripts/compute_edistance_cpu.py

@@ -2,20 +2,37 @@
 
 import os
 import time
+from argparse import ArgumentParser
+from pathlib import Path
 
 import anndata as ad
 from pertpy.tools import Distance
 
 if __name__ == "__main__":
+    parser = ArgumentParser()
+    parser.add_argument("--bootstrap", action="store_true")
+    args = parser.parse_args()
     obs_key = "perturbation"
+    bootstrap = args.bootstrap
 
     # homedir/data/adamson_2016_upr_epistasis
     save_dir = os.path.join(
-        os.path.expanduser("~"), "data", "adamson_2016_upr_epistasis_pca.h5ad"
+        os.path.expanduser("~"),
+        "data",
     )
-    adata = ad.read_h5ad(save_dir)
+    adata = ad.read_h5ad(Path(save_dir) / "adamson_2016_upr_epistasis_pca.h5ad")
     dist = Distance(obsm_key="X_pca", metric="edistance")
     start_time = time.time()
-    df = dist.pairwise(adata, groupby=obs_key)
+    if bootstrap:
+        df, df_var = dist.pairwise(
+            adata, groupby=obs_key, bootstrap=True, n_bootstrap=100
+        )
+    else:
+        df = dist.pairwise(adata, groupby=obs_key)
     end_time = time.time()
     print(f"Time taken: {end_time - start_time} seconds")
+    if bootstrap:
+        df_var.to_csv(Path(save_dir) / "df_cpu_bootstrap_var.csv")
+        df.to_csv(Path(save_dir) / "df_cpu_bootstrap.csv")
+    else:
+        df.to_csv(Path(save_dir) / "df_cpu.csv")