Fix: improve run_pca and fix Python 3.9 compatibility

- Extract _slice_pca to module level for testability
- Fix run_pca to properly preserve existing X_pca when using custom pca_key
- Capture results before modifying obsm to avoid accessing deleted keys
- Add future annotations import for Python 3.9 compatibility in tests
- Update release notes with PR #176 optimizations

Author: katosh

README.md

@@ -76,6 +76,14 @@ Release Notes
  ### Next Release
  * Avoid warning about coloring in `plot_trajectories`
  * Fix: forward `eigvec_key` to fallback method in `early_cell()`
+ * Performance optimizations (PR #176):
+   - Faster max-min waypoint sampling: O(N·K²) → O(N·K)
+   - PCA slicing optimization to avoid redundant recomputation
+   - Sparse LU solver for absorption probabilities
+   - Vectorized perspective matrix and Dijkstra shortest paths
+   - Vectorized Markov chain construction and terminal state detection
+   - Optional sklearn backend for diffusion kernel (`kernel_backend` parameter)
+ * Fix: `run_pca` now correctly preserves existing `X_pca` when using custom `pca_key`
 
  ### Version 1.4.2
  * Fix "lightgray" error in `plot_trend`

src/palantir/utils.py

@@ -26,6 +26,29 @@ class CellNotFoundException(Exception):
     pass
 
 
+def _slice_pca(ad: AnnData, n_comps: int) -> None:
+    """Slice PCA results to fewer components without recomputing.
+
+    Parameters
+    ----------
+    ad : AnnData
+        AnnData object with PCA results in obsm["X_pca"], uns["pca"], and varm["PCs"].
+    n_comps : int
+        Number of components to keep.
+    """
+    if "X_pca" in ad.obsm:
+        ad.obsm["X_pca"] = ad.obsm["X_pca"][:, :n_comps]
+    if "pca" in ad.uns:
+        for key in ("variance", "variance_ratio"):
+            if key in ad.uns["pca"]:
+                ad.uns["pca"][key] = np.asarray(ad.uns["pca"][key])[:n_comps]
+        params = ad.uns["pca"].get("params")
+        if isinstance(params, dict) and "n_comps" in params:
+            params["n_comps"] = n_comps
+    if "PCs" in ad.varm:
+        ad.varm["PCs"] = ad.varm["PCs"][:, :n_comps]
+
+
 def run_pca(
     data: Union[pd.DataFrame, AnnData],
     n_components: int = 300,
@@ -53,25 +76,14 @@ def run_pca(
         Tuple of PCA projections of the data and the explained variance.
         If AnnData is passed as data, the results are also written to the input object and None is returned.
     """
-    def _slice_pca(ad: AnnData, n_comps: int) -> None:
-        if "X_pca" in ad.obsm:
-            ad.obsm["X_pca"] = ad.obsm["X_pca"][:, :n_comps]
-        if "pca" in ad.uns:
-            for key in ("variance", "variance_ratio"):
-                if key in ad.uns["pca"]:
-                    ad.uns["pca"][key] = np.asarray(ad.uns["pca"][key])[:n_comps]
-            params = ad.uns["pca"].get("params")
-            if isinstance(params, dict) and "n_comps" in params:
-                params["n_comps"] = n_comps
-        if "PCs" in ad.varm:
-            ad.varm["PCs"] = ad.varm["PCs"][:, :n_comps]
-
     if isinstance(data, pd.DataFrame):
         ad = AnnData(data.values)
+        old_pca = None
     else:
         ad = data
-        if pca_key != "X_pca":
-            old_pca = ad.obsm.get("X_pca", None)
+        # Preserve existing X_pca if user wants results in a different key
+        if pca_key != "X_pca" and "X_pca" in ad.obsm:
+            old_pca = ad.obsm["X_pca"].copy()
         else:
             old_pca = None
 
@@ -91,15 +103,19 @@ def _slice_pca(ad: AnnData, n_comps: int) -> None:
         elif n_comps > l_n_comps:
             sc.pp.pca(ad, n_comps=n_comps, mask_var="highly_variable", zero_center=False)
 
-    if isinstance(data, AnnData):
+    # Capture results before potentially modifying obsm
+    pca_projections = pd.DataFrame(ad.obsm["X_pca"], index=ad.obs_names)
+    variance_ratio = ad.uns["pca"]["variance_ratio"]
+
+    if isinstance(data, AnnData) and pca_key != "X_pca":
         data.obsm[pca_key] = ad.obsm["X_pca"]
-        if pca_key != "X_pca":
-            del data.obsm["X_pca"]
-        elif old_pca is not None:
+        # Restore original X_pca if it existed, otherwise remove the temporary one
+        if old_pca is not None:
             data.obsm["X_pca"] = old_pca
+        else:
+            del data.obsm["X_pca"]
 
-    pca_projections = pd.DataFrame(ad.obsm[pca_key], index=ad.obs_names)
-    return pca_projections, ad.uns["pca"]["variance_ratio"]
+    return pca_projections, variance_ratio
 
 
 def run_low_density_variability(

tests/test_core_equivalence.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix, csgraph, find