Merge pull request #18 from stanfish06/downsample

update: downsample (sign)

Author: stanfish06

src/scloop/data/types.py

@@ -1,6 +1,18 @@
 # Copyright 2025 Zhiyuan Yu (Heemskerk's lab, University of Michigan)
-from typing import Literal
+from typing import Literal, Annotated, TypeAlias
+from pydantic import Field
 
 FeatureSelectionMethod = Literal["hvg", "delve", "none"]
 EmbeddingMethod = Literal["pca", "diffmap", "scvi"]
 EmbeddingNeighbors = Literal["pca", "scvi"]
+
+Index_t = Annotated[int, Field(ge=0)]
+# need at least 2 points to compute PH. Maybe also set an upper bound later as it is not feasible to compute PH on a lot of points
+IndexListDownSample: TypeAlias = Annotated[
+    list[Index_t],
+    Field(min_length=2, description="Downsampled indices for PH computation"),
+]
+SizeDownSample = Annotated[
+    int, Field(ge=2, description="Sample to this number of cells")
+]
+# TODO: make a type for boundary matrix. Restrict matrix size for efficient computation

src/scloop/preprocessing/__init__.py

@@ -1,2 +1,3 @@
 from .prepare import prepare_adata
+from .downsample import sample
 from .delve import delve_fs

src/scloop/preprocessing/downsample.py

@@ -1 +1,90 @@
-# TODO: use numba
+# Copyright 2025 Zhiyuan Yu (Heemskerk's lab, University of Michigan)
+from numba import jit
+from anndata import AnnData
+from ..data.types import IndexListDownSample, SizeDownSample, EmbeddingMethod
+from pydantic import validate_call
+import pandas as pd
+import numpy as np
+
+__all__ = ["sample"]
+
+
+@jit(nopython=True)
+def _sample_impl(
+    data: np.ndarray,
+    class_labels: np.ndarray,
+    classes: np.ndarray,
+    seed_indices: np.ndarray,
+    n: int,
+) -> np.ndarray:
+    # selected observation indices
+    indices = np.zeros(n, dtype=np.int64)
+    min_dists = np.full(len(data), np.inf)
+
+    num_seeds = len(seed_indices)
+    num_classes = len(classes)
+    indices[0] = seed_indices[0]
+
+    # for each newly added points, recompute and select out-of-bag hausdorff point
+    for i in range(1, n):
+        last_point = data[indices[i - 1]]
+        dists = np.sum((data - last_point) ** 2, axis=1)
+        min_dists = np.minimum(min_dists, dists)
+        min_dists[indices[:i]] = -1
+        if i >= num_seeds:
+            class_indicies = np.where(class_labels == classes[i % num_classes])[0]
+            next_idx = class_indicies[np.argmax(min_dists[class_indicies])]
+            # if this class is exhausted, fall back to reguler sampling
+            if min_dists[next_idx] == -1:
+                next_idx = np.argmax(min_dists)
+            indices[i] = next_idx
+        else:
+            indices[i] = seed_indices[i]
+
+    return indices
+
+
+@validate_call(config={"arbitrary_types_allowed": True})
+def sample(
+    adata: AnnData,
+    groupby: str | None,
+    embedding_method: EmbeddingMethod,
+    n: SizeDownSample,
+    random_state: int = 0,
+) -> IndexListDownSample:
+    """
+    Topology-preserving downsampling using greedy farthest-point sampling.
+
+    Args:
+        adata: AnnData object containing the data
+        groupby: column in adata.obs for class-balanced sampling, or None
+        embedding_method: which embedding to use from adata.obsm
+        n: number of points to sample
+        random_state: random seed for reproducibility
+
+    Returns:
+        list of indices into adata.obs for the downsampled points
+    """
+    assert n <= adata.shape[0]
+    assert f"X_{embedding_method}" in adata.obsm
+    downsample_embedding = adata.obsm[f"X_{embedding_method}"]
+    assert type(downsample_embedding) is np.ndarray
+
+    if groupby is None:
+        class_labels = np.zeros(adata.shape[0], dtype=np.int64)
+        classes = np.array([0])
+        seed_indices = np.array([np.random.randint(adata.shape[0])])
+    else:
+        assert type(adata.obs) is pd.DataFrame
+        assert groupby in adata.obs.columns
+        class_labels, classes = pd.factorize(adata.obs.loc[:, groupby])
+        classes = np.arange(len(classes), dtype=np.int64)
+        seed_indices = []
+        np.random.seed(random_state)
+        for c in classes:
+            seed_indices.append(np.random.choice(np.where(class_labels == c)[0]))
+        seed_indices = np.array(seed_indices)
+
+    return _sample_impl(
+        downsample_embedding, class_labels, classes, seed_indices, n
+    ).tolist()