update: downsample

Author: stanfish06

src/scloop/data/types.py

@@ -8,5 +8,11 @@
 
 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)]
+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/downsample.py

@@ -1,15 +1,33 @@
 # Copyright 2025 Zhiyuan Yu (Heemskerk's lab, University of Michigan)
 from numba import jit
 from anndata import AnnData
-from ..data.types import IndexListDownSample
+from ..data.types import IndexListDownSample, SizeDownSample, EmbeddingMethod
 from pydantic import validate_call
-from typing import TypeAlias
+from scipy.spatial.distance import directed_hausdorff
 import numpy as np
 
-@jit
-def _sample_impl(data: np.ndarray, n: int) -> np.ndarray:
-    return np.array([1, 1])
+@jit(nopython=True)
+def _sample_impl(data: np.ndarray, seed_idx: int, n: int) -> np.ndarray:
+    indices = np.zeros(n, dtype=np.int64)
+    indices[0] = seed_idx
+    min_dists = np.full(len(data), np.inf)
+    
+    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
+        indices[i] = np.argmax(min_dists)
+    
+    return indices.tolist()
 
 @validate_call(config={"arbitrary_types_allowed": True})
-def sample(adata: AnnData, n: int) -> IndexListDownSample:
-    return _sample_impl(np.array(), n).tolist()
+def sample(adata: AnnData, embedding_method: EmbeddingMethod, n: SizeDownSample) -> IndexListDownSample:
+    """
+    topology-aware downsampling
+    """
+    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
+    return _sample_impl(downsample_embedding, 0, n)