numeric stability and progtess for run_local_variability

README.md

@@ -51,6 +51,8 @@ Release Notes
  * fixes for edge cases discoverd through extended testing
  * implement `plot_trajectory` function to show trajectory on the umap
  * scale pseudotime to unit intervall in anndata
+ * optional progress bar with `progress=True` in `palantir.utils.run_local_variability`
+ * avoid NaN in local variablility output
 
  ### Version 1.3.1
  * implemented `palantir.plot.plot_stats` to plot arbitray cell-wise statistics as x-/y-positions.

src/palantir/utils.py

@@ -488,7 +488,7 @@ def _dot_helper_func(x, y):
     return x.dot(y)
 
 
-def _local_var_helper(expressions, distances):
+def _local_var_helper(expressions, distances, eps=1e-16):
     if hasattr(expressions, "todense"):
 
         def cast(x):
@@ -510,7 +510,7 @@ def cast(x):
         except ValueError:
             raise ValueError(f"This cell caused the error: {cell}")
         expr_distance = np.sqrt(np.sum(expr_deltas**2, axis=1, keepdims=True))
-        change_rate = expr_deltas / expr_distance
+        change_rate = expr_deltas / (expr_distance + eps)
         yield np.max(change_rate**2, axis=0)
 
 
@@ -519,6 +519,8 @@ def run_local_variability(
     expression_key: str = "MAGIC_imputed_data",
     distances_key: str = "distances",
     localvar_key: str = "local_variability",
+    progress: bool = False,
+    eps: float = 1e-16,
 ) -> np.ndarray:
     """
     Compute local gene variability scores for each cell.
@@ -539,6 +541,10 @@ def run_local_variability(
     localvar_key : str, optional
         Key under which the computed local variability matrix is stored in the layers of the AnnData object.
         Default is 'local_variability'.
+    progress : bool
+        Show progress bar. Requires tqdm to be installed. Default is False.
+    eps : float
+        A small value preventing devision by 0. Defaults to 1e-16.
 
     Returns
     -------
@@ -557,7 +563,17 @@ def run_local_variability(
         raise KeyError(f"'{distances_key}' not found in .obsp.")
     X_dists = ad.obsp[distances_key]
 
-    local_variability = np.stack(list(_local_var_helper(X, X_dists)))
+    local_var_generator = _local_var_helper(X, X_dists, eps=eps)
+    if progress is True:
+        try:
+            from fastprogress.fastprogress import progress_bar
+        except ModuleNotFoundError:
+            raise Exception(
+                "Showing the progress bar requires the python module `fastprogress` to be installed."
+            )
+        local_var_generator = progress_bar(local_var_generator, total=X.shape[0])
+
+    local_variability = np.stack(list(local_var_generator))
 
     ad.layers[localvar_key] = local_variability