Optimize runtime and memory usage

.github/workflows/build.yml

@@ -2,40 +2,44 @@ name: Build and Test
 
 on:
   push:
+    branches: ["main", "ci_optimization"]
+  pull_request:
+    branches: ["main"]
 
 permissions:
-  contents: read # to fetch code (actions/checkout)
+  contents: read
 
 jobs:
-  build_wheels:
-    name: ${{ matrix.os }} Wheels
+  test:
+    name: ${{ matrix.os }} - Python ${{ matrix.python-version }}
     runs-on: ${{ matrix.os }}
     strategy:
-      # Ensure that a wheel builder finishes even if another fails
       fail-fast: false
       matrix:
-        os: [windows-latest, macos-latest]
+        os: [ubuntu-latest, windows-latest, macos-latest]
         python-version: ['3.8', '3.9', '3.10', '3.11', '3.12']
 
     steps:
       - name: Checkout ${{ github.repository }}
-        uses: actions/checkout@v3
-        with:
-          fetch-depth: 0
+        uses: actions/checkout@v4
 
       - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
-      - name: Dependencies
+      - name: Install package
         run:  pip install . 
 
       - name: Test
         run: |
-            python -c "\
+            python -c "
             import anndata
             from sketchKH import *
             adata = anndata.read_h5ad('data/nk_cell_preprocessed.h5ad')
             kh_indices, adata_subsample = sketch(adata, sample_set_key = 'FCS_File', gamma = 1, num_subsamples = 500, frequency_seed = 0, n_jobs = 4)
-            "
+
+            n_samples = adata.obs['FCS_File'].nunique()
+            assert len(kh_indices) == n_samples, f'Expected {n_samples} indices, Actual {len(kh_indices)}'
+            assert adata_subsample.shape == (10000, 43), f'Expected shape (10000, 43), Actual {adata_subsample.shape}'
+            "

README.md

@@ -4,13 +4,13 @@ Distribution-Informed Sketching with Kernel Herding
 ## Overview
 We provide a set of functions for distribution-aware sketching of multiple profiled single-cell samples via Kernel Herding. Our sketches select a small, representative set of cells from each profiled sample so that all major immune cell-types and their relative frequencies are well-represented. 
 * Please see our paper for more information (ACM-BCB 2022) : https://arxiv.org/abs/2207.00584
-* Updated : March 17, 2026
+* Updated : June 23, 2026
 
 ![Sketching via KH Overview](https://github.com/CompCy-lab/SketchKH/blob/main/sketch_overview.png?raw=True)
 
 ## Installation
 Dependencies
-* Python >= 3.6, anndata >= 0.7.6, numpy >= 1.22.4, scipy >= 1.7.1, numba, joblib, tqdm_joblib
+* Python >= 3.8, anndata >= 0.7.6, numpy >= 1.22.4, scipy >= 1.7.1, numba, joblib, tqdm_joblib
 
 You can install the package with `pip` by,
 ```

pyproject.toml

@@ -4,10 +4,10 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "sketchKH"
-version = "0.1.3"
+version = "0.1.4"
 description = "Distribution-based sketching of single-cell samples"
 readme = "README.md"
-requires-python = ">=3.6"
+requires-python = ">=3.8"
 authors = [
     { name = "CompCy Lab", email = "compcylab@gmail.com" }
 ]

sketchKH/kh.py

@@ -1,13 +1,12 @@
 import numpy as np
-from tqdm import tqdm
 import anndata
 from typing import Union
-import scipy
 import logging
 from pandas.api.types import is_numeric_dtype
 from numba import njit
 from joblib import Parallel, delayed
 from tqdm_joblib import tqdm_joblib
+from scipy import sparse
 
 def random_feats(X: np.ndarray,
                 gamma: Union[int, float] = 1,
@@ -29,17 +28,14 @@ def random_feats(X: np.ndarray,
     ----------
     """
     scale = 1 / gamma
-
-    if (frequency_seed is not None):
-        np.random.seed(frequency_seed)
-        W = np.random.normal(scale = scale, size = (X.shape[1], 1000))
-    else:
-        W = np.random.normal(scale = scale, size = (X.shape[1], 1000))
+    rng = np.random.RandomState(frequency_seed)
+    W = rng.normal(scale = scale, size = (X.shape[1], 1000))
 
     XW = np.dot(X, W)
-    sin_XW = np.sin(XW)
-    cos_XW = np.cos(XW)
-    phi = np.concatenate((cos_XW, sin_XW), axis=1)
+
+    phi = np.empty((X.shape[0], 2000))
+    np.cos(XW, out = phi[:, :1000])
+    np.sin(XW, out = phi[:, 1000:])
 
     return phi
 
@@ -72,13 +68,12 @@ def kernel_herding(phi: np.ndarray, num_subsamples: int):
 
     w_0 = np.copy(w_t)
     for subsample_idx in range(num_subsamples): #find argmax
+        scores = phi @ w_t
         max_score = -1e20
         new_ind = -1
         for cell_idx in range(num_cells):
             if selected_mask[cell_idx] == 0:
-                score = 0.0
-                for feature_idx in range(num_features):
-                    score += phi[cell_idx, feature_idx] * w_t[feature_idx]
+                score = scores[cell_idx]
                 if score > max_score:
                     max_score = score
                     new_ind = cell_idx
@@ -93,7 +88,7 @@ def kernel_herding(phi: np.ndarray, num_subsamples: int):
     return kh_indices
 
 def _parse_input(adata: anndata.AnnData):
-    """accesses and parses data from adata object
+    """Accesses and parses data from adata object
 
     Parameters
     adata: anndata.AnnData
@@ -106,16 +101,12 @@ def _parse_input(adata: anndata.AnnData):
         array of data (dimensions = cells x features)
     ----------
     """
-    try:
-        if isinstance(adata, anndata.AnnData):
-            X = adata.X.copy()
-        if isinstance(X, scipy.sparse.csr_matrix):
-            X = np.asarray(X.todense())
-        if is_numeric_dtype(adata.obs_names):
-            logging.warning('Converting cell IDs to strings.')
-            adata.obs_names = adata.obs_names.astype('str')
-    except NameError:
-        pass
+    X = adata.X
+    if sparse.issparse(X):
+        X = X.toarray()
+    if is_numeric_dtype(adata.obs_names):
+        logging.warning('Converting cell IDs to strings.')
+        adata.obs_names = adata.obs_names.astype('str')
 
     return X
 
@@ -144,7 +135,7 @@ def kernel_herding_main(sample_set_ind,
         indices of subsampled cells within the sample-set
     ----------
     """
-    X = X[sample_set_ind, :]
+    X = X[sample_set_ind, :] #will create a copy
     phi = random_feats(X, gamma = gamma, frequency_seed = frequency_seed)
     kh_indices = kernel_herding(phi, num_subsamples)
 
@@ -156,8 +147,9 @@ def sketch(adata,
             gamma: Union[int, float] = 1,
             frequency_seed: int = None,
             num_subsamples: int = 500,
-            n_jobs: int = -1):
-    """constructs a sketch using kernel herding and random Fourier frequency features
+            backend: str = 'threading',
+            n_jobs: int = 1):
+    """Constructs a sketch using kernel herding and random Fourier frequency features
 
     Parameters
     adata: anndata.Anndata
@@ -175,7 +167,9 @@ def sketch(adata,
         random state parameter     
     num_samples: int (default = None)
         number of cells to subsample per sample-set
-    n_jobs: int (default = -1)
+    backend: str (default = 'threading')
+        backend for parallelization
+    n_jobs: int (default = 1)
         number of tasks
     ----------
 
@@ -209,8 +203,8 @@ def sketch(adata,
     def process_set(i, inds):
         return kernel_herding_main(sample_set_ind = inds, X = X, gamma = gamma, frequency_seed = frequency_seed, num_subsamples = num_subsamples)
 
-    with tqdm_joblib(tqdm(desc="Performing subsampling", total = n_sample_sets)):
-        kh_indices = Parallel(n_jobs=n_jobs)(delayed(process_set)(i, inds) for i, inds in enumerate(sample_set_inds))
+    with tqdm_joblib(total=n_sample_sets, desc='Performing subsampling'):
+        kh_indices = Parallel(n_jobs=n_jobs, backend=backend)(delayed(process_set)(i, inds) for i, inds in enumerate(sample_set_inds))
 
     subsampled_cell_indices = [sample_set_inds[i][kh_indices[i]] for i in range(n_sample_sets)]
     subsampled_cell_indices = np.concatenate(subsampled_cell_indices)