Merge pull request #22 from devsystemslab/feat/typing

Add a few type hints, remove some unused inputs

Author: Marius1311

README.md

@@ -4,13 +4,15 @@
 [![PyPI][badge-pypi]][pypi]
 [![Docs Build][badge-docs]][docs]
 [![Downloads][badge-downloads]][downloads]
+[![Coverage][badge-coverage]][coverage]
 
 
 [badge-tests]: https://github.com/devsystemslab/hnoca-tools/actions/workflows/test.yaml/badge.svg
 [badge-pre-commit]: https://results.pre-commit.ci/badge/github/devsystemslab/hnoca-tools/main.svg
 [badge-pypi]: https://img.shields.io/pypi/v/hnoca.svg
 [badge-docs]: https://github.com/devsystemslab/hnoca-tools/actions/workflows/build-site.yaml/badge.svg
 [badge-downloads]: https://static.pepy.tech/badge/hnoca
+[badge-coverage]: https://codecov.io/gh/devsystemslab/hnoca-tools/branch/main/graph/badge.svg
 
 # Human Neural Organoid Cell Atlas Toolbox
 #### 🛠️ The Swiss Army Knive of the Single Cell Cartographer
@@ -129,3 +131,4 @@ de_df = stats.test_de_paired(
 [pypi]: https://pypi.org/project/hnoca/
 [docs]: https://devsystemslab.github.io/HNOCA-tools/
 [downloads]: https://pepy.tech/project/hnoca
+[coverage]: https://codecov.io/gh/devsystemslab/hnoca-tools

src/hnoca/mapping/mapper.py

@@ -6,6 +6,7 @@
 import anndata as ad
 import cloudpickle
 import numpy as np
+import scanpy as sc
 
 from hnoca.utils import check_deps
 
@@ -93,7 +94,7 @@ def _train_scanvi(self, query_adata, retrain="partial", **kwargs):
 
         self.query_model = vae_q
 
-    def _train_scvi(self, query_adata, retrain="partial", **kwargs):
+    def _train_scvi(self, query_adata: sc.AnnData, retrain: str = "partial", **kwargs):
         """Train a new scvi model on the query data."""
         unfrozen = retrain == "full"
         self.scvi.model.SCVI.prepare_query_anndata(query_adata, self.ref_model)
@@ -102,7 +103,7 @@ def _train_scvi(self, query_adata, retrain="partial", **kwargs):
 
         self.query_model = vae_q
 
-    def _train_scpoli(self, query_adata, retrain="partial", labeled_indices=None, **kwargs):
+    def _train_scpoli(self, query_adata: sc.AnnData, retrain: str = "partial", labeled_indices=None, **kwargs):
         """Train a new scpoli model on the query data"""
         freeze = retrain != "full"  # noqa F841
         labeled_indices = [] if labeled_indices is None else labeled_indices

src/hnoca/mapping/matching.py

@@ -1,11 +1,12 @@
 import anndata
 import numpy as np
+import scanpy as sc
 from scipy import sparse
 
 
 def get_matched_transcriptome(  # noqa D103
-    adata,
-    adata_ref,
+    adata: sc.AnnData,
+    adata_ref: sc.AnnData,
     wknn,
     rescale_factor=1,
 ):

src/hnoca/mapping/utils.py

@@ -1,12 +1,13 @@
 import anndata as ad
 import numpy as np
 import pandas as pd
+import scanpy as sc
 from scipy import sparse
 
 from hnoca._logging import logger
 
 
-def prepare_features(query_adata, ref_model):
+def prepare_features(query_adata: sc.AnnData, ref_model):
     """Prepare the features of the query dataset to match the reference dataset."""
     ref_features = ref_model.adata.var_names
     query_features = query_adata.var_names

src/hnoca/mapping/wknn.py

@@ -3,6 +3,7 @@
 
 import numpy as np
 import pandas as pd
+import scanpy as sc
 import tqdm
 from pynndescent import NNDescent
 from scipy import sparse
@@ -52,8 +53,6 @@ def build_nn(  # noqa: D103
     ref,
     query=None,
     k=100,
-    weight: Literal["unweighted", "dist", "gaussian_kernel"] = "unweighted",
-    sigma=None,
     use_rapids: bool = False,
 ):
     if query is None:
@@ -108,9 +107,9 @@ def random_walk_with_restart(init, transition_prob, alpha=0.5, num_rounds=100):
 
 
 def get_wknn(
-    ref,
-    query,
-    ref2=None,
+    ref: np.ndarray,
+    query: np.ndarray,
+    ref2: np.ndarray | None = None,
     k: int = 100,
     query2ref: bool = True,
     ref2query: bool = False,
@@ -123,23 +122,23 @@ def get_wknn(
 
     Parameters
     ----------
-    ref : np.ndarray
+    ref
         The reference representation to build ref-query neighbor graph
-    query : np.ndarray
+    query
         The query representation to build ref-query neighbor graph
-    ref2 : np.ndarray
+    ref2
         The reference representation to build ref-ref neighbor graph
-    k : int
+    k
         Number of neighbors per cell
-    query2ref : bool
+    query2ref
         Consider query-to-ref neighbors
-    ref2query : bool
+    ref2query
         Consider ref-to-query neighbors
-    weighting_scheme : str
+    weighting_scheme
         How to weight edges in the ref-query neighbor graph
-    top_n : int
+    top_n
         The number of top neighbors to consider
-    return_adjs : bool
+    return_adjs
         Whether to return the adjacency matrices of ref-query, query-ref, ref-ref, and ref-ref for weighting
     """
     adj_q2r = build_nn(ref=ref, query=query, k=k)
@@ -186,26 +185,26 @@ def get_wknn(
         return wknn
 
 
-def estimate_presence_score(  # noqa: D103
-    ref_adata,
-    query_adata,
+def estimate_presence_score(
+    ref_adata: sc.AnnData,
+    query_adata: sc.AnnData,
     wknn=None,
-    use_rep_ref_wknn="X_latent",
-    use_rep_query_wknn="X_latent",
-    k_wknn=100,
-    query2ref_wknn=True,
-    ref2query_wknn=False,
-    weighting_scheme_wknn="jaccard_square",
+    use_rep_ref_wknn: str = "X_latent",
+    use_rep_query_wknn: str = "X_latent",
+    k_wknn: int = 100,
+    query2ref_wknn: bool = True,
+    ref2query_wknn: bool = False,
+    weighting_scheme_wknn: str = "jaccard_square",
     ref_trans_prop=None,
     use_rep_ref_trans_prop=None,
-    k_ref_trans_prop=50,
-    symm_ref_trans_prop=True,
+    k_ref_trans_prop: int = 50,
     split_by=None,
-    do_random_walk=True,
-    alpha_random_walk=0.1,
-    num_rounds_random_walk=100,
+    do_random_walk: bool = True,
+    alpha_random_walk: float = 0.1,
+    num_rounds_random_walk: int = 100,
     log=True,
 ):
+    """Estimate presence score of query cells in reference dataset."""
     if wknn is None:
         ref = ref_adata.obsm[use_rep_ref_wknn]
         query = query_adata.obsm[use_rep_query_wknn]
@@ -271,7 +270,7 @@ def estimate_presence_score(  # noqa: D103
     }
 
 
-def transfer_labels(ref_adata, query_adata, wknn, label_key="celltype"):
+def transfer_labels(ref_adata: sc.AnnData, query_adata: sc.AnnData, wknn, label_key: str = "celltype"):
     """Transfer labels from reference to query data."""
     scores = pd.DataFrame(
         wknn @ pd.get_dummies(ref_adata.obs[label_key]),