Merge pull request #41 from quadbio/docs/improve_docstrings

Decrease docstring duplication

Author: Marius1311

pyproject.toml

@@ -24,6 +24,7 @@ classifiers = [
 dynamic = [ "version" ]
 dependencies = [
   "anndata>=0.11",
+  "docrep",
   "numpy",
   "packaging",
   "pandas",
@@ -56,6 +57,7 @@ optional-dependencies.doc = [
 ]
 optional-dependencies.test = [
   "coverage",
+  "faiss-cpu",
   "pytest",
   "squidpy>=1.6",
 ]

src/cellmapper/_docs.py

@@ -0,0 +1,109 @@
+"""Shared documentation for cellmapper."""
+
+from docrep import DocstringProcessor
+
+__all__ = ["d"]
+
+_t = """\
+t
+    Number of diffusion time steps. This parameter controls the degree of
+    smoothing applied by the diffusion operator. Larger values lead to more
+    smoothing."""
+
+_diffusion_method = """\
+diffusion_method
+    Method for computing powers of the mapping matrix (only valid in self-mapping mode). Options are "iterative" for
+    iterative matrix multiplication (inspired by MAGIC :cite:`van2018recovering`) and "spectral" for
+    eigendecomposition-based approach. """
+
+_prediction_postfix = """\
+prediction_postfix
+    Postfix to add to mapped variables to identify them as predictions."""
+
+_symmetrize = """\
+symmetrize
+    If True, create a symmetrize connectivity matrix. Only valid for square matrices (self-mapping).
+    If None (default), uses True for self-mapping and False for cross-mapping."""
+
+_self_edges = """\
+self_edges
+    Control self-edges (diagonal entries) for square matrices (self-mapping).
+    If None (default), uses False for self-mapping (scanpy style) and None for cross-mapping.
+    This controls whether or not the kernel used to compute the connectivities is supplied with self-edges.
+    It does not determine whether the final connectivity matrix has self edges. For example, the `umap`
+    kernel expectes self-edges, but does not produce them in the final connectivity matrix."""
+
+_knn_method = """\
+knn_method
+    Method for computing k-nearest neighbors. Options include:
+    - "sklearn": Scikit-learn's NearestNeighbors. See https://scikit-learn.org/stable/modules/generated/sklearn.neighbors.NearestNeighbors.html
+    - "pynndescent": Pynndescent's approximate nearest neighbors. See https://pynndescent.readthedocs.io/en/latest/
+    - "rapids": RAPIDS cuML's NearestNeighbors (GPU). See https://docs.rapids.ai/api/cuml/stable/api.html#cuml.neighbors.NearestNeighbors
+    - "faiss-cpu": Facebook AI Similarity Search (FAISS) on CPU. See https://faiss.ai/
+    - "faiss-gpu": Facebook AI Similarity Search (FAISS) on GPU. See https://faiss.ai/
+
+
+    All methods return exactly `n_neighbors` neighbors, including the reference cell itself (in self-mapping mode)."""
+
+_only_yx = """\
+only_yx
+    If True, only compute the xy neighbors. In self-mapping mode, this is
+    automatically set to True for efficiency since all neighbor matrices contain the same information.
+    This is faster, but not suitable for Jaccard or HNOCA methods in cross-mapping mode."""
+
+_kernel_method = """\
+kernel_method
+    Method to use for computing the mapping matrix. Options include:
+
+    - "jaccard": Jaccard similarity. Inspired by GLUE :cite:`cao2022multi`
+    - "gauss": Gaussian kernel with (global) bandwith equal to the mean distance.
+    - "scarches": scArches kernel. Inspired by scArches :cite:`lotfollahi2022mapping`
+    - "inverse_distance": Inverse distance kernel.
+    - "random": Random kernel, useful for testing.
+    - "hnoca": HNOCA kernel. Inspired by HNOCA-tools :cite:`he2024integrated`
+    - "equal": All neighbors are equally weighted (1/n_neighbors).
+    - "umap": UMAP fuzzy simplicial set connectivities. Only available for self-mapping with true k-NN graphs."""
+
+_comparison_method = """\
+comparison_method
+    Method to use for comparing the mapping results. Options include:
+
+    - "pearson": Pearson correlation coefficient.
+    - "spearman": Spearman rank correlation coefficient.
+    - "js": Jenson-Shanon divergence.
+    - "rmse": Root Mean Square Error."""
+
+_layer_key = """\
+layer_key
+    Key in `self.query.layers` to use as the original expression. Use "X" to use `self.query.X`."""
+
+
+_n_neighbors = """\
+n_neighbors
+    Number of nearest neighbors. This parameter controls the sparsity of the connectivity matrix. """
+
+_use_rep = """\
+use_rep
+    Data representation based on which to find nearest neighbors. If None, a fallback representation will be
+    computed automatically. """
+
+_knn_dist_metric = """\
+knn_dist_metric
+    Distance metric to use for nearest neighbors. See the knn algorithms documentation for details. """
+
+
+d = DocstringProcessor(
+    t=_t,
+    diffusion_method=_diffusion_method,
+    prediction_postfix=_prediction_postfix,
+    symmetrize=_symmetrize,
+    self_edges=_self_edges,
+    knn_method=_knn_method,
+    only_yx=_only_yx,
+    kernel_method=_kernel_method,
+    comparison_method=_comparison_method,
+    layer_key=_layer_key,
+    n_neighbors=_n_neighbors,
+    use_rep=_use_rep,
+    knn_dist_metric=_knn_dist_metric,
+)

src/cellmapper/check.py

@@ -61,15 +61,18 @@ def check(self) -> None:
     "https://docs.rapids.ai/install/.",
     cupy="To speed up k-NN search on GPU, you may install cuPy following the guide from "
     "https://docs.rapids.ai/install/.",
-    faiss="To speed up k-NN search on GPU, you may install faiss following the guide from "
+    faiss_cpu="To speed up k-NN search on CPU, you may install faiss following the guide from "
+    "https://github.com/facebookresearch/faiss/blob/main/INSTALL.md",
+    faiss_gpu="To speed up k-NN search on GPU, you may install faiss following the guide from "
     "https://github.com/facebookresearch/faiss/blob/main/INSTALL.md",
     pynndescent="To use fast approximate k-NN search, install pynndescent: pip install pynndescent",
 )
 
 CHECKERS = {
     "cuml": Checker("cuml", vmin=None, install_hint=INSTALL_HINTS.cuml),
     "cupy": Checker("cupy", vmin=None, install_hint=INSTALL_HINTS.cupy),
-    "faiss": Checker("faiss", package_name="faiss", vmin="1.7.0", install_hint=INSTALL_HINTS.faiss),
+    "faiss-cpu": Checker("faiss", package_name="faiss-cpu", vmin="1.7.0", install_hint=INSTALL_HINTS.faiss_cpu),
+    "faiss-gpu": Checker("faiss", package_name="faiss", vmin="1.7.0", install_hint=INSTALL_HINTS.faiss_gpu),
     "pynndescent": Checker("pynndescent", vmin=None, install_hint=INSTALL_HINTS.pynndescent),
 }
 

src/cellmapper/constants.py

@@ -6,8 +6,8 @@ class PackageConstants:
     SKLEARN_WARNING_CUTOFF: int = 50000
 
     # Default mapping methods
-    DEFAULT_SELF_MAPPING_METHOD: str = "umap"
-    DEFAULT_CROSS_MAPPING_METHOD: str = "gauss"
+    DEFAULT_SELF_MAPPING_KERNEL_METHOD: str = "umap"
+    DEFAULT_CROSS_MAPPING_KERNEL_METHOD: str = "gauss"
 
     # Kernel method categories
     JACCARD_BASED_KERNELS = {"jaccard", "hnoca"}

src/cellmapper/model/_knn_backend.py

@@ -50,15 +50,44 @@ def query(self, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
         return distances, indices
 
 
-class _FaissBackend(_KNNBackend):
+class _FaissCpuBackend(_KNNBackend):
     def __init__(
         self,
         n_neighbors: int,
         metric: str,
         random_state: int = 0,
         **kwargs: Any,
     ):
-        check_deps("faiss")
+        check_deps("faiss-cpu")
+        import faiss
+
+        self.faiss = faiss
+        self._index = None
+
+    def fit(self, data: np.ndarray) -> None:
+        dims = data.shape[1]
+        index = self.faiss.IndexFlatL2(dims)
+        # Ensure data is float32 and C-contiguous
+        data_f32 = np.ascontiguousarray(data.astype(np.float32))
+        index.add(data_f32)
+        self._index = index
+
+    def query(self, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
+        # Ensure points are float32 and C-contiguous
+        points_f32 = np.ascontiguousarray(points.astype(np.float32))
+        distances, indices = self._index.search(points_f32, k)
+        return distances, indices
+
+
+class _FaissGpuBackend(_KNNBackend):
+    def __init__(
+        self,
+        n_neighbors: int,
+        metric: str,
+        random_state: int = 0,
+        **kwargs: Any,
+    ):
+        check_deps("faiss-gpu")
         import faiss
 
         self.faiss = faiss
@@ -69,11 +98,15 @@ def fit(self, data: np.ndarray) -> None:
         dims = data.shape[1]
         flat = self.faiss.IndexFlatL2(dims)
         gpu_index = self.faiss.index_cpu_to_gpu(self.res, 0, flat)
-        gpu_index.add(data.astype(np.float32))
+        # Ensure data is float32 and C-contiguous
+        data_f32 = np.ascontiguousarray(data.astype(np.float32))
+        gpu_index.add(data_f32)
         self._index = gpu_index
 
     def query(self, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
-        distances, indices = self._index.search(points.astype(np.float32), k)
+        # Ensure points are float32 and C-contiguous
+        points_f32 = np.ascontiguousarray(points.astype(np.float32))
+        distances, indices = self._index.search(points_f32, k)
         return distances, indices
 
 
@@ -148,16 +181,17 @@ def query(self, points: np.ndarray, k: int) -> tuple[np.ndarray, np.ndarray]:
 
 _BACKENDS = {
     "sklearn": _SklearnBackend,
-    "faiss": _FaissBackend,
+    "faiss-cpu": _FaissCpuBackend,
+    "faiss-gpu": _FaissGpuBackend,
     "rapids": _RapidsBackend,
     "pynndescent": _PyNNDescentBackend,
 }
 
 
-def get_backend(method: str, n_neighbors: int, metric: str, random_state: int = 0, **kwargs: Any) -> _KNNBackend:
+def get_backend(knn_method: str, n_neighbors: int, metric: str, random_state: int = 0, **kwargs: Any) -> _KNNBackend:
     """Factory to get a configured KNN backend."""
     try:
-        backend_cls = _BACKENDS[method]
+        backend_cls = _BACKENDS[knn_method]
     except KeyError:
-        raise ValueError(f"Unknown method: {method}. Supported methods: {list(_BACKENDS)}") from KeyError
+        raise ValueError(f"Unknown method: {knn_method}. Supported methods: {list(_BACKENDS)}") from KeyError
     return backend_cls(n_neighbors=n_neighbors, metric=metric, random_state=random_state, **kwargs)