Repository: theislab/cellrank (language: python)

Issue:
Refactor kernels
<!-- Describe the feature request in the block below: -->
TODOs:
- [x] differentiate between directional and non-directional
- [x] make inversion not an inplace operation
- [x] don't compute `CytoTRACE` score on initialization
- [x] remove adaptive scaling (if reintroduced, will be added in a better form)
- [x] use elementwise multiplication when multiplying 2 kernels, not matmul
- [x] move `_tmap_as_tmat` from #664 
- [x] refactor `RandomWalk` class to include plotting
- [x] refactor `compute_projection`
- [x] add a better way of handling optional connectivities (`ConnectivityMixin`?)
- [x] add common decorator to check whether `.transition_matrix` is present
- [x] remove condition number check?
- [x] add docs/comments for some internal code
- [x] improve public docstrings



Discussion / hints:

Fix patch:
diff --git a/.bumpversion.cfg b/.bumpversion.cfg
index 94e029182..2069fd375 100644
--- a/.bumpversion.cfg
+++ b/.bumpversion.cfg
@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 1.5.0
+current_version = 1.5.1
 commit = True
 tag = True
 files = setup.py cellrank/__init__.py
diff --git a/README.rst b/README.rst
index 7ca561a15..8ddc82f20 100644
--- a/README.rst
+++ b/README.rst
@@ -23,6 +23,10 @@ fate-informed `circular embeddings`_, to name just a few.
 CellRank scales to large cell numbers, is fully compatible with `scanpy`_ and `scvelo`_ and is easy to use.
 For **installation instructions**, **documentation** and **tutorials**, visit `cellrank.org`_.
 
+Manuscript
+^^^^^^^^^^
+Please check out our manuscript `Lange et al. (2022)`_ in **Nature Methods** to learn more.
+
 Getting started with CellRank
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 If you're new to CellRank, make sure to go though the `basic tutorial`_ which introduces you to CellRank's high-level
@@ -30,10 +34,6 @@ API. Most biological systems require a bit more control, so be sure to check out
 which allows to unlock the full power of CellRank. If you want to see individual functions in action, visit our
 `gallery`_.
 
-Manuscript
-^^^^^^^^^^
-Please see our `preprint`_ on **bioRxiv** to learn more.
-
 CellRank's key applications
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 - compute initial & terminal as well as intermediate `macrostates`_ of your biological system
@@ -140,7 +140,7 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
 .. _Theislab: https://www.helmholtz-muenchen.de/icb/research/groups/theis-lab/overview/index.html
 .. _Peerlab: https://www.mskcc.org/research/ski/labs/dana-pe-er
 .. _`tweet`: https://twitter.com/MariusLange8
-.. _preprint: https://www.biorxiv.org/content/10.1101/2020.10.19.345983v2
+.. _Lange et al. (2022): https://www.nature.com/articles/s41592-021-01346-6
 .. _PageRank: https://en.wikipedia.org/wiki/PageRank#cite_note-1
 
 .. _gallery: https://cellrank.readthedocs.io/en/stable/auto_examples/index.html
diff --git a/cellrank/__init__.py b/cellrank/__init__.py
index f37679061..cf4095bb2 100644
--- a/cellrank/__init__.py
+++ b/cellrank/__init__.py
@@ -14,7 +14,7 @@
 
 __author__ = ", ".join(["Marius Lange", "Michal Klein"])
 __maintainer__ = ", ".join(["Marius Lange", "Michal Klein"])
-__version__ = "1.5.0"
+__version__ = "1.5.1"
 __email__ = "info@cellrank.org"
 
 try:
diff --git a/cellrank/_key.py b/cellrank/_key.py
index 82270abb3..53fb968d2 100644
--- a/cellrank/_key.py
+++ b/cellrank/_key.py
@@ -15,15 +15,15 @@ class Key:
     """Class which manages keys in :class:`anndata.AnnData`."""
 
     @classmethod
-    def backward(cls, bwd: bool) -> str:
+    def backward(cls, bwd: Optional[bool]) -> str:
         return "bwd" if bwd else "fwd"
 
     @classmethod
-    def where(cls, bwd: bool) -> str:
+    def where(cls, bwd: Optional[bool]) -> str:
         return "from" if bwd else "to"
 
     @classmethod
-    def initial(cls, bwd: bool) -> str:
+    def initial(cls, bwd: Optional[bool]) -> str:
         return "initial" if bwd else "terminal"
 
     @classmethod
@@ -36,15 +36,15 @@ def probs(cls, key: str) -> str:
             return f"{key}_probs"
 
         @classmethod
-        def macrostates(cls, bwd: bool) -> str:
+        def macrostates(cls, bwd: Optional[bool]) -> str:
             return f"macrostates_{Key.backward(bwd)}"
 
         @classmethod
-        def term_states(cls, bwd: bool) -> str:
+        def term_states(cls, bwd: Optional[bool]) -> str:
             return f"{Key.initial(bwd)}_states"
 
         @classmethod
-        def priming_degree(cls, bwd: bool) -> str:
+        def priming_degree(cls, bwd: Optional[bool]) -> str:
             return f"priming_degree_{Key.backward(bwd)}"
 
     class obsm:
@@ -53,33 +53,33 @@ def memberships(cls, key: str) -> str:
             return f"{key}_memberships"
 
         @classmethod
-        def schur_vectors(cls, bwd: bool) -> str:
+        def schur_vectors(cls, bwd: Optional[bool]) -> str:
             return f"schur_vectors_{Key.backward(bwd)}"
 
         @classmethod
-        def macrostates(cls, bwd: bool) -> str:
+        def macrostates(cls, bwd: Optional[bool]) -> str:
             return f"macrostates_{Key.backward(bwd)}"
 
         @classmethod
-        def abs_probs(cls, bwd: bool) -> str:
-            return ("from" if bwd else "to") + "_" + Key.obs.term_states(bwd)
+        def abs_probs(cls, bwd: Optional[bool]) -> str:
+            return Key.where(bwd) + "_" + Key.obs.term_states(bwd)
 
         @classmethod
-        def abs_times(cls, bwd: bool) -> str:
+        def abs_times(cls, bwd: Optional[bool]) -> str:
             return f"absorption_times_{Key.backward(bwd)}"
 
     class varm:
         @classmethod
-        def lineage_drivers(cls, bwd: bool):
-            return ("initial" if bwd else "terminal") + "_lineage_drivers"
+        def lineage_drivers(cls, bwd: Optional[bool]):
+            return Key.initial(bwd) + "_lineage_drivers"
 
     class uns:
         @classmethod
-        def kernel(cls, bwd: bool, key: Optional[str] = None) -> str:
+        def kernel(cls, bwd: Optional[bool], key: Optional[str] = None) -> str:
             return key if key is not None else f"T_{Key.backward(bwd)}"
 
         @classmethod
-        def estimator(cls, bwd: bool, key: Optional[str] = None) -> str:
+        def estimator(cls, bwd: Optional[bool], key: Optional[str] = None) -> str:
             return key if key is not None else f"{Key.backward(bwd)}_estimator"
 
         @classmethod
@@ -91,13 +91,13 @@ def colors(cls, key: str) -> str:
             return f"{key}_colors"
 
         @classmethod
-        def eigen(cls, bwd: bool) -> str:
+        def eigen(cls, bwd: Optional[bool]) -> str:
             return f"eigendecomposition_{Key.backward(bwd)}"
 
         @classmethod
-        def schur_matrix(cls, bwd: bool) -> str:
+        def schur_matrix(cls, bwd: Optional[bool]) -> str:
             return f"schur_matrix_{Key.backward(bwd)}"
 
         @classmethod
-        def coarse(cls, bwd: bool) -> str:
+        def coarse(cls, bwd: Optional[bool]) -> str:
             return f"coarse_{Key.backward(bwd)}"
diff --git a/cellrank/datasets/_datasets.py b/cellrank/datasets/_datasets.py
index 685a61f55..c9eeed2e7 100644
--- a/cellrank/datasets/_datasets.py
+++ b/cellrank/datasets/_datasets.py
@@ -11,8 +11,17 @@
 from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d, inject_docs
 
+__all__ = (
+    "pancreas",
+    "pancreas_preprocessed",
+    "lung",
+    "reprogramming_morris",
+    "reprogramming_schiebinger",
+    "zebrafish",
+)
 
-class ReprogrammingSubset(ModeEnum):  # noqa: D101
+
+class ReprogrammingSubset(ModeEnum):
     FULL = auto()
     K48 = "48k"
     K85 = "85k"
@@ -47,9 +56,8 @@ class ReprogrammingSubset(ModeEnum):  # noqa: D101
 
 
 def _load_dataset_from_url(
-    fpath: Union[os.PathLike, str], url: str, expected_shape: Tuple[int, int], **kwargs
+    fpath: Union[str, Path], url: str, expected_shape: Tuple[int, int], **kwargs: Any
 ) -> AnnData:
-
     fpath = str(fpath)
     if not fpath.endswith(".h5ad"):
         fpath += ".h5ad"
@@ -101,7 +109,7 @@ def pancreas(
     ----------
     path
         Path where to save the dataset.
-    **kwargs
+    kwargs
         Keyword arguments for :func:`scanpy.read`.
 
     Returns
diff --git a/cellrank/external/estimators/_import_error_estim.py b/cellrank/external/estimators/_import_error_estim.py
index 028fee767..1ac27e8a5 100644
--- a/cellrank/external/estimators/_import_error_estim.py
+++ b/cellrank/external/estimators/_import_error_estim.py
@@ -1,8 +1,11 @@
 from typing import Any
 
+from anndata import AnnData
 from cellrank.tl.estimators import BaseEstimator
 from cellrank.external._error_mixin import ImportErrorMixin
 
+__all__ = ("ErroredEstimator",)
+
 
 class ErroredEstimator(ImportErrorMixin, BaseEstimator):
     """
@@ -11,10 +14,18 @@ class ErroredEstimator(ImportErrorMixin, BaseEstimator):
     Subclasses can modify the message by overriding `__import_error_message__`.
     """
 
-    __import_error_message__ = "Unable to import the estimator."
+    __import_error_message__ = "Unable to import external estimator."
 
     def _fit_terminal_states(self, *args: Any, **kwargs: Any) -> None:
         raise NotImplementedError
 
-    def compute_terminal_states(self, *args: Any, **kwargs: Any) -> None:  # noqa: D102
+    def fit(self, *args: Any, **kwargs: Any) -> "ErroredEstimator":
+        """Not implemented."""
+        raise NotImplementedError
+
+    def predict(self, *args: Any, **kwargs: Any) -> None:
+        """Not implemented."""
+        raise NotImplementedError
+
+    def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
         raise NotImplementedError
diff --git a/cellrank/external/kernels/_import_error_kernel.py b/cellrank/external/kernels/_import_error_kernel.py
index 528b2cb13..03ac28df1 100644
--- a/cellrank/external/kernels/_import_error_kernel.py
+++ b/cellrank/external/kernels/_import_error_kernel.py
@@ -1,10 +1,12 @@
-from typing import Any
+from typing import Any, Optional
 
-from cellrank.tl.kernels import Kernel
 from cellrank.external._error_mixin import ImportErrorMixin
-from cellrank.tl.kernels._base_kernel import KernelExpression
+from cellrank.tl.kernels._base_kernel import Kernel
 
+__all__ = ("ErroredKernel",)
 
+
+# can't subclass UnidirectionalKernel since StationaryOTKernel have to use it
 class ErroredKernel(ImportErrorMixin, Kernel):
     """
     Utility kernel class which always throw :class:`ImportError` when instantiated.
@@ -12,12 +14,17 @@ class ErroredKernel(ImportErrorMixin, Kernel):
     Subclasses can modify the message by overriding `__import_error_message__`.
     """
 
-    __import_error_message__ = "Unable to import the kernel."
+    __import_error_message__ = "Unable to import external kernel."
 
-    def compute_transition_matrix(  # noqa: D102
-        self, *args: Any, **kwargs: Any
-    ) -> KernelExpression:
+    def compute_transition_matrix(self, *args: Any, **kwargs: Any) -> "ErroredKernel":
+        """Not implemented."""
         raise NotImplementedError
 
-    def copy(self) -> KernelExpression:  # noqa: D102
+    def copy(self, deep: bool = False) -> "ErroredKernel":
+        """Not implemented."""
         raise NotImplementedError
+
+    @property
+    def backward(self) -> Optional[bool]:
+        """None."""
+        return None
diff --git a/cellrank/external/kernels/_statot_kernel.py b/cellrank/external/kernels/_statot_kernel.py
index 45fb6b2a5..f458c7039 100644
--- a/cellrank/external/kernels/_statot_kernel.py
+++ b/cellrank/external/kernels/_statot_kernel.py
@@ -4,6 +4,7 @@
 from anndata import AnnData
 from cellrank._key import Key
 from cellrank.ul._docs import d
+from cellrank.tl.kernels._mixins import UnidirectionalMixin
 
 import numpy as np
 import pandas as pd
@@ -19,8 +20,11 @@
     _error = e
 
 
+__all__ = ("StationaryOTKernel",)
+
+
 @d.dedent
-class StationaryOTKernel(OTKernel_, error=_error):
+class StationaryOTKernel(UnidirectionalMixin, OTKernel_, error=_error):
     """
     Stationary optimal transport kernel from :cite:`zhang:21`.
 
@@ -31,12 +35,12 @@ class StationaryOTKernel(OTKernel_, error=_error):
     %(adata)s
     terminal_states
         Key in :attr:`anndata.AnnData.obs` or a categorical :class:`pandas.Series` where non-`NaN` values
-        mark terminal states.
+        mark the terminal states.
         If `None`, terminal states are assumed to be present in :attr:`anndata.AnnData.obs` ``['terminal_states']``.
     g
         Key in :attr:`anndata.AnnData.obs` containing relative growth rates for cells or the array itself.
     kwargs
-        Additional keyword arguments.
+        Keyword arguments for the parent class.
     """
 
     __import_error_message__ = "Unable to import the kernel. Please install `statOT` first as `pip install statot POT`."
@@ -58,14 +62,6 @@ def __init__(
         except Exception as e:  # noqa: B902
             raise RuntimeError("Unable to initialize the kernel.") from e
 
-    def _read_from_adata(
-        self,
-        conn_key: Optional[str] = "connectivities",
-        read_conn: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        super()._read_from_adata(conn_key=conn_key, read_conn=False, **kwargs)
-
     def compute_transition_matrix(
         self,
         eps: float,
@@ -92,7 +88,7 @@ def compute_transition_matrix(
         basis
             Key in :attr:`anndata.AnnData.obsm` where the basis is stored.
         cost_norm_method
-            Cost normalization method to use. Use "mean" to ensure `mean(C) = 1` or refer to
+            Cost normalization method to use. Use "mean" to ensure ``mean(C) = 1`` or refer to
             :func:`ot.utils.cost_normalization` for more information.
         method
             Choice of regularization. Valid options are:
@@ -100,7 +96,6 @@ def compute_transition_matrix(
                 - `'ent'` - entropy.
                 - `'quad'` - L2-norm.
                 - `'unbal'` - unbalanced transport (not yet implemented).
-
         tol
             Relative tolerance for OT solver convergence.
         thresh
@@ -116,7 +111,7 @@ def compute_transition_matrix(
 
         Returns
         -------
-        Self and makes :attr:`transition_matrix` available.
+        Self and updates :attr:`transition_matrix` and :attr:`params`.
         """
         if method not in ("ent", "quad", "unbal"):
             raise ValueError(
diff --git a/cellrank/external/kernels/_wot_kernel.py b/cellrank/external/kernels/_wot_kernel.py
index 24329449f..622313152 100644
--- a/cellrank/external/kernels/_wot_kernel.py
+++ b/cellrank/external/kernels/_wot_kernel.py
@@ -10,7 +10,8 @@
 from cellrank.ul._docs import d
 from cellrank.tl._utils import _maybe_subset_hvgs
 from cellrank.external.kernels._utils import MarkerGenes
-from cellrank.tl.kernels._exp_time_kernel import LastTimePoint
+from cellrank.tl.kernels.utils._tmat_flow import Numeric_t
+from cellrank.tl.kernels._transport_map_kernel import Pair_t, LastTimePoint
 
 import numpy as np
 import pandas as pd
@@ -27,13 +28,18 @@
     wot = None
 
 
-class nstr(str):  # used for params+cache (precomputed cost matrices)
+__all__ = ("WOTKernel",)
+
+
+# TODO(michalk8): remove me
+class nstr(str):  # used for params + cache (precomputed cost matrices)
     """String class that is not equal to any other string."""
 
     def __eq__(self, other: str) -> bool:
         return False
 
 
+# TODO(michalk8): refactor me
 @d.dedent
 class WOTKernel(Kernel, error=_error):
     """
@@ -49,7 +55,8 @@ class WOTKernel(Kernel, error=_error):
     time_key
         Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
         The experimental time can be of either of a numeric or an ordered categorical type.
-    %(cond_num)s
+    kwargs
+        Keyword arguments for the parent class.
 
     Examples
     --------
@@ -93,14 +100,12 @@ def __init__(
         adata: AnnData,
         backward: bool = False,
         time_key: str = "exp_time",
-        compute_cond_num: bool = False,
         **kwargs: Any,
     ):
         super().__init__(
             adata,
             backward=backward,
             time_key=time_key,
-            compute_cond_num=compute_cond_num,
             **kwargs,
         )
 
@@ -112,14 +117,6 @@ def __init__(
         self.adata.obs[self._time_key] = self.experimental_time
         self._growth_rates = None
 
-    def _read_from_adata(
-        self,
-        conn_key: Optional[str] = "connectivities",
-        read_conn: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        super()._read_from_adata(conn_key=conn_key, read_conn=False, **kwargs)
-
     def compute_initial_growth_rates(
         self,
         proliferation_key: Optional[str] = None,
@@ -255,9 +252,7 @@ def delta(
     @d.dedent
     def compute_transition_matrix(
         self,
-        cost_matrices: Optional[
-            Union[str, Mapping[Tuple[float, float], np.ndarray]]
-        ] = None,
+        cost_matrices: Optional[Union[str, Mapping[Pair_t, np.ndarray]]] = None,
         lambda1: float = 1,
         lambda2: float = 50,
         epsilon: float = 0.05,
@@ -311,7 +306,7 @@ def compute_transition_matrix(
 
         Returns
         -------
-        Self and makes available the following attributes:
+        Self and updates the following attributes:
 
             - :attr:`transition_matrix` - transition matrix.
             - :attr:`transport_maps` - transport maps between consecutive time points.
@@ -319,7 +314,7 @@ def compute_transition_matrix(
 
         Also modifies :attr:`anndata.AnnData.obs` with the following key:
 
-            - `'estimated_growth_rates'` - the estimated final growth rates.
+            - `'estimated_growth_rates'` - the final estimated growth rates.
 
         Notes
         -----
@@ -367,11 +362,7 @@ def compute_transition_matrix(
         tmap = self._restich_tmaps(tmap, last_time_point, conn_kwargs=conn_kwargs)
         self._growth_rates = tmap.obs
 
-        self._compute_transition_matrix(
-            matrix=tmap.X,
-            density_normalize=False,
-            check_irreducibility=False,
-        )
+        self.transition_matrix = tmap.X
         if threshold:
             self._threshold_transition_matrix(threshold)
         self.adata.obs["estimated_growth_rates"] = self.growth_rates[f"g{growth_iters}"]
@@ -380,7 +371,7 @@ def compute_transition_matrix(
 
         return self
 
-    def _compute_tmap(self, t1: Any, t2: Any, **kwargs: Any) -> AnnData:
+    def _compute_tmap(self, t1: Numeric_t, t2: Numeric_t, **kwargs: Any) -> AnnData:
         ot_model: wot.ot.OTModel = kwargs.pop("ot_model")
         cost_matrix: Optional[Union[str, np.ndarray]] = kwargs.pop("cost_matrix")
 
@@ -392,7 +383,7 @@ def _compute_tmap(self, t1: Any, t2: Any, **kwargs: Any) -> AnnData:
         if tmap is None:
             raise TypeError(
                 f"Unable to compute transport map for time pair `{(t1, t2)}`. "
-                f"Please ensure `adata.obs[{self._time_key!r}]` has the correct dtype (float)."
+                f"Please ensure that `adata.obs[{self._time_key!r}]` has the correct dtype (float)."
             )
         tmap.X = tmap.X.astype(np.float64)
         nans = int(np.sum(~np.isfinite(tmap.X)))
@@ -406,9 +397,7 @@ def _compute_tmap(self, t1: Any, t2: Any, **kwargs: Any) -> AnnData:
     def _compute_tmaps(
         self,
         adata: AnnData,
-        cost_matrices: Optional[
-            Mapping[Tuple[float, float], Optional[np.ndarray]]
-        ] = None,
+        cost_matrices: Optional[Mapping[Pair_t, Optional[np.ndarray]]] = None,
         solver: Literal["fixed_iters", "duality_gap"] = "duality_gap",
         growth_rate_field: Optional[str] = None,
         **kwargs: Any,
@@ -451,10 +440,8 @@ def _compute_tmaps(
     def _generate_cost_matrices(
         self,
         adata: AnnData,
-        cost_matrices: Optional[
-            Union[str, Mapping[Tuple[float, float], np.ndarray]]
-        ] = None,
-    ) -> Tuple[Mapping[Tuple[float, float], Optional[np.ndarray]], str]:
+        cost_matrices: Optional[Union[str, Mapping[Pair_t, np.ndarray]]] = None,
+    ) -> Tuple[Mapping[Pair_t, Optional[np.ndarray]], str]:
         timepoints = self.experimental_time.cat.categories
         timepoints = list(zip(timepoints[:-1], timepoints[1:]))
 
@@ -530,7 +517,7 @@ def growth_rates(self) -> Optional[pd.DataFrame]:
         return self._growth_rates
 
     def __invert__(self) -> "WOTKernel":
-        super().__invert__()
-        # because WOT reads from `adata`
-        self.adata.obs[self._time_key] = self.experimental_time
-        return self
+        wk = super().__invert__()
+        # needed for WOT
+        wk.adata.obs[wk._time_key] = wk.experimental_time
+        return wk
diff --git a/cellrank/external/models/_import_error_model.py b/cellrank/external/models/_import_error_model.py
index f4029af2f..0662f4b41 100644
--- a/cellrank/external/models/_import_error_model.py
+++ b/cellrank/external/models/_import_error_model.py
@@ -1,10 +1,12 @@
-from typing import Optional
+from typing import Any, Optional
 
 from cellrank.ul.models import BaseModel
 from cellrank.external._error_mixin import ImportErrorMixin
 
 import numpy as np
 
+__all__ = ("ErroredModel",)
+
 
 class ErroredModel(ImportErrorMixin, BaseModel):
     """
@@ -13,29 +15,33 @@ class ErroredModel(ImportErrorMixin, BaseModel):
     Subclasses can modify the message by overriding `__import_error_message__`.
     """
 
-    __import_error_message__ = "Unable to import the model."
+    __import_error_message__ = "Unable to import external model."
 
-    def fit(  # noqa: D102
+    def fit(
         self,
         x: Optional[np.ndarray] = None,
         y: Optional[np.ndarray] = None,
         w: Optional[np.ndarray] = None,
-        **kwargs,
-    ) -> "BaseModel":
+        **kwargs: Any,
+    ) -> "ErroredModel":
+        """Not implemented."""
         raise NotImplementedError
 
-    def predict(  # noqa: D102
+    def predict(
         self,
         x_test: Optional[np.ndarray] = None,
         key_added: Optional[str] = "_x_test",
-        **kwargs,
+        **kwargs: Any,
     ) -> np.ndarray:
+        """Not implemented."""
         raise NotImplementedError
 
-    def confidence_interval(  # noqa: D102
-        self, x_test: Optional[np.ndarray] = None, **kwargs
+    def confidence_interval(
+        self, x_test: Optional[np.ndarray] = None, **kwargs: Any
     ) -> np.ndarray:
+        """Not implemented."""
         raise NotImplementedError
 
-    def copy(self) -> BaseModel:  # noqa: D102
+    def copy(self) -> "ErroredModel":
+        """Not implemented."""
         raise NotImplementedError
diff --git a/cellrank/kernels.py b/cellrank/kernels.py
index 2cc5cbf5f..9560bf1f8 100644
--- a/cellrank/kernels.py
+++ b/cellrank/kernels.py
@@ -1,19 +1,10 @@
+from cellrank.tl.kernels.utils import _similarity as similarity
+from cellrank.tl.kernels.utils import _pseudotime_scheme as pseudotime_scheme
 from cellrank.tl.kernels._base_kernel import Kernel
-from cellrank.tl.kernels._exp_time_kernel import ExperimentalTimeKernel
 from cellrank.tl.kernels._velocity_kernel import VelocityKernel
 from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEKernel
-from cellrank.tl.kernels._velocity_schemes import (
-    CosineScheme,
-    DotProductScheme,
-    CorrelationScheme,
-    SimilaritySchemeABC,
-)
 from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
 from cellrank.tl.kernels._precomputed_kernel import PrecomputedKernel
-from cellrank.tl.kernels._pseudotime_schemes import (
-    ThresholdSchemeABC,
-    HardThresholdScheme,
-    SoftThresholdScheme,
-)
 from cellrank.tl.kernels._connectivity_kernel import ConnectivityKernel
 from cellrank.tl.kernels._transport_map_kernel import TransportMapKernel
+from cellrank.tl.kernels._experimental_time_kernel import ExperimentalTimeKernel
diff --git a/cellrank/tl/_init_term_states.py b/cellrank/tl/_init_term_states.py
index 2c30c86d1..35fcc044b 100644
--- a/cellrank/tl/_init_term_states.py
+++ b/cellrank/tl/_init_term_states.py
@@ -13,8 +13,8 @@
 from cellrank.tl.kernels import PrecomputedKernel
 from cellrank.tl.estimators import GPCCA, CFLARE, TermStatesEstimator
 from cellrank.tl._transition_matrix import transition_matrix
-from cellrank.tl.kernels._velocity_kernel import BackwardMode, VelocityMode
 from cellrank.tl.estimators._base_estimator import BaseEstimator
+from cellrank.tl.kernels.utils._velocity_model import BackwardMode, VelocityModel
 
 _docstring = """\
 Find {direction} states of a dynamic process of single cells based on RNA velocity :cite:`manno:18`.
@@ -73,6 +73,7 @@ def _fit(
     estim: TermStatesEstimator,
     n_lineages: Optional[int] = None,
     keys: Optional[Sequence[str]] = None,
+    method: Optional[str] = None,
     cluster_key: Optional[str] = None,
     compute_absorption_probabilities: bool = True,
     **kwargs: Any,
@@ -86,7 +87,7 @@ def _fit(
             use=n_lineages,
             cluster_key=cluster_key,
             n_clusters_kmeans=n_lineages,
-            method=kwargs.pop("method", "kmeans"),
+            method=method or "kmeans",
             **kwargs,
         )
     elif isinstance(estim, GPCCA):
@@ -96,7 +97,7 @@ def _fit(
                 n_lineages = estim.eigendecomposition["eigengap"] + 1
 
         if n_lineages > 1:
-            estim.compute_schur(n_lineages, method=kwargs.pop("method", "krylov"))
+            estim.compute_schur(n_lineages, method=method or "krylov")
 
         try:
             estim.compute_macrostates(
@@ -130,17 +131,18 @@ def _initial_terminal(
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
     backward: bool = False,
-    mode: str = VelocityMode.DETERMINISTIC,
+    model: str = VelocityModel.DETERMINISTIC,
     backward_mode: str = BackwardMode.TRANSPOSE,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
+    method: Optional[str] = None,
     key: Optional[str] = None,
     force_recompute: bool = False,
     show_plots: bool = False,
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
-    **kwargs,
+    **kwargs: Any,
 ) -> Optional[Union[AnnData, BaseEstimator]]:
     _check_estimator_type(estimator)
 
@@ -149,14 +151,14 @@ def _initial_terminal(
     try:
         if force_recompute:
             raise KeyError("Forcing transition matrix recomputation.")
-        kernel = PrecomputedKernel(key, adata=adata, backward=backward)
+        kernel = PrecomputedKernel(adata, obsp_key=key, backward=backward)
         logg.info("Using precomputed transition matrix")
     except KeyError:
         # compute kernel object
         kernel = transition_matrix(
             adata,
             backward=backward,
-            mode=mode,
+            model=model,
             backward_mode=backward_mode,
             **kwargs,
         )
@@ -175,6 +177,7 @@ def _initial_terminal(
         n_lineages=n_states,
         cluster_key=cluster_key,
         compute_absorption_probabilities=False,
+        method=method,
         **fit_kwargs,
     )
 
@@ -196,7 +199,7 @@ def _initial_terminal(
 
 
 @_deprecate(version="2.0")
-@inject_docs(m=VelocityMode, b=BackwardMode)
+@inject_docs(m=VelocityModel, b=BackwardMode)
 @d.dedent
 @inject_docs(
     __doc__=_docstring.format(
@@ -208,7 +211,7 @@ def _initial_terminal(
 def initial_states(  # noqa: D103
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
-    mode: str = VelocityMode.DETERMINISTIC,
+    model: str = VelocityModel.DETERMINISTIC,
     backward_mode: str = BackwardMode.TRANSPOSE,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
@@ -218,13 +221,13 @@ def initial_states(  # noqa: D103
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
-    **kwargs,
+    **kwargs: Any,
 ) -> Optional[Union[AnnData, BaseEstimator]]:
 
     return _initial_terminal(
         adata,
         estimator=estimator,
-        mode=mode,
+        model=model,
         backward_mode=backward_mode,
         backward=True,
         n_states=n_states,
@@ -240,7 +243,7 @@ def initial_states(  # noqa: D103
 
 
 @_deprecate(version="2.0")
-@inject_docs(m=VelocityMode, b=BackwardMode)
+@inject_docs(m=VelocityModel, b=BackwardMode)
 @d.dedent
 @inject_docs(
     __doc__=_docstring.format(
@@ -250,7 +253,7 @@ def initial_states(  # noqa: D103
 def terminal_states(  # noqa: D103
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
-    mode: str = VelocityMode.DETERMINISTIC,
+    model: str = VelocityModel.DETERMINISTIC,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
@@ -259,13 +262,13 @@ def terminal_states(  # noqa: D103
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
-    **kwargs,
+    **kwargs: Any,
 ) -> Optional[Union[AnnData, BaseEstimator]]:
 
     return _initial_terminal(
         adata,
         estimator=estimator,
-        mode=mode,
+        model=model,
         backward=False,
         n_states=n_states,
         cluster_key=cluster_key,
diff --git a/cellrank/tl/_mixins/_anndata.py b/cellrank/tl/_mixins/_anndata.py
index a727bfa9e..0e69f92bd 100644
--- a/cellrank/tl/_mixins/_anndata.py
+++ b/cellrank/tl/_mixins/_anndata.py
@@ -58,7 +58,7 @@ def from_adata(cls, adata: AnnData, **kwargs: Any) -> "AnnDataMixin":
         ----------
         %(adata)s
         kwargs
-            Additional keyword arguments.
+            Additional keyword arguments for :meth:`__init__`.
 
         Returns
         -------
diff --git a/cellrank/tl/_transition_matrix.py b/cellrank/tl/_transition_matrix.py
index 1566077a6..65b9d4b46 100644
--- a/cellrank/tl/_transition_matrix.py
+++ b/cellrank/tl/_transition_matrix.py
@@ -7,12 +7,12 @@
 from cellrank.tl._utils import _deprecate
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
 from cellrank.tl.kernels._base_kernel import KernelExpression
-from cellrank.tl.kernels._velocity_kernel import BackwardMode, VelocityMode
-from cellrank.tl.kernels._velocity_schemes import Scheme
+from cellrank.tl.kernels.utils._similarity import Similarity
+from cellrank.tl.kernels.utils._velocity_model import BackwardMode, VelocityModel
 
 
 @_deprecate(version="2.0")
-@inject_docs(m=VelocityMode, b=BackwardMode, s=Scheme)  # don't swap the order
+@inject_docs(m=VelocityModel, b=BackwardMode, s=Similarity)  # don't swap the order
 @d.dedent
 def transition_matrix(
     adata: AnnData,
@@ -21,13 +21,13 @@ def transition_matrix(
     xkey: str = "Ms",
     conn_key: str = "connectivities",
     gene_subset: Optional[Iterable] = None,
-    mode: Literal[
-        "deterministic", "stochastic", "sampling", "monte_carlo"
-    ] = VelocityMode.DETERMINISTIC,
+    model: Literal[
+        "deterministic", "stochastic", "monte_carlo"
+    ] = VelocityModel.DETERMINISTIC,
     backward_mode: Literal["transpose", "negate"] = BackwardMode.TRANSPOSE,
-    scheme: Union[
+    similarity: Union[
         Literal["dot_product", "cosine", "correlation"], Callable
-    ] = Scheme.CORRELATION,
+    ] = Similarity.CORRELATION,
     softmax_scale: Optional[float] = None,
     weight_connectivities: float = 0.2,
     density_normalize: bool = True,
@@ -84,18 +84,18 @@ def compute_velocity_kernel() -> VelocityKernel:
             conn_key=conn_key,
         ).compute_transition_matrix(
             softmax_scale=softmax_scale,
-            mode=mode,
+            model=model,
             backward_mode=backward_mode,
-            scheme=scheme,
+            similarity=similarity,
             **kwargs,
         )
 
     if 0 < weight_connectivities < 1:
         vk = compute_velocity_kernel()
         logg.info(f"Using a connectivity kernel with weight `{weight_connectivities}`")
-        ck = ConnectivityKernel(
-            adata, backward=backward, conn_key=conn_key
-        ).compute_transition_matrix(density_normalize=density_normalize)
+        ck = ConnectivityKernel(adata, conn_key=conn_key).compute_transition_matrix(
+            density_normalize=density_normalize
+        )
         final = (
             (1 - weight_connectivities) * vk + weight_connectivities * ck
         ).compute_transition_matrix()
@@ -104,7 +104,6 @@ def compute_velocity_kernel() -> VelocityKernel:
     elif weight_connectivities == 1:
         final = ConnectivityKernel(
             adata,
-            backward=backward,
             conn_key=conn_key,
         ).compute_transition_matrix(density_normalize=density_normalize)
     else:
diff --git a/cellrank/tl/_utils.py b/cellrank/tl/_utils.py
index a2fcd6b8c..67e346022 100644
--- a/cellrank/tl/_utils.py
+++ b/cellrank/tl/_utils.py
@@ -726,24 +726,14 @@ def _irreducible(d: Union[spmatrix, np.ndarray]) -> bool:
 
     import networkx as nx
 
-    start = logg.debug("Checking the transition matrix for irreducibility")
-
     G = nx.DiGraph(d) if not isinstance(d, nx.DiGraph) else d
-
     try:
         it = iter(nx.strongly_connected_components(G))
         _ = next(it)
         _ = next(it)
-        is_irreducible = False
+        return False
     except StopIteration:
-        is_irreducible = True
-
-    if not is_irreducible:
-        logg.warning("Transition matrix is not irreducible", time=start)
-    else:
-        logg.debug("Transition matrix is irreducible", time=start)
-
-    return is_irreducible
+        return True
 
 
 def _symmetric(
diff --git a/cellrank/tl/estimators/_base_estimator.py b/cellrank/tl/estimators/_base_estimator.py
index 1ca958dbd..5cb8cb243 100644
--- a/cellrank/tl/estimators/_base_estimator.py
+++ b/cellrank/tl/estimators/_base_estimator.py
@@ -13,7 +13,7 @@
 from cellrank.tl._mixins import IOMixin, KernelMixin, AnnDataMixin
 from cellrank.tl.kernels import PrecomputedKernel
 from cellrank.tl._lineage import Lineage
-from cellrank.tl.kernels._base_kernel import Kernel, KernelExpression
+from cellrank.tl.kernels._base_kernel import KernelExpression
 
 import numpy as np
 import pandas as pd
@@ -31,52 +31,33 @@ class BaseEstimator(IOMixin, KernelMixin, AnnDataMixin, ABC):
 
     Parameters
     ----------
-    obj
-        Can be one of the following:
-
-            - :class:`cellrank.tl.kernels.Kernel` - kernel object.
-            - :class:`anndata.AnnData` - annotated data object containing transition matrix in
-              :attr:`anndata.AnnData.obsp`.
-            - :class:`numpy.ndarray` - row-normalized sparse transition matrix.
-            - :class:`scipy.sparse.spmatrix` - row-normalized sparse transition matrix.
-    obsp_key
-        Key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
-        Only used when ``obj`` is an :class:`anndata.AnnData` object.
+    object
+        Can be one of the following types:
+
+            - :class:`anndata.AnnData` - annotated data object.
+            - :class:`scipy.sparse.spmatrix`, :class:`numpy.ndarray` - row-normalized transition matrix.
+            - :class:`cellrank.tl.kernels.KernelExpression` - kernel expression.
+            - :class:`str` - key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
+              ``adata`` must be provided in this case.
+            - :class:`bool` - directionality of the transition matrix that will be used to infer its storage location.
+              If `None`, the directionality will be determined automatically. ``adata`` must be provided in this case.
+    kwargs
+        Keyword arguments for :class:`cellrank.tl.kernels.PrecomputedKernel`.
     """
 
     def __init__(
         self,
-        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
-        obsp_key: Optional[str] = None,
+        object: Union[str, bool, np.ndarray, spmatrix, AnnData, KernelExpression],
+        **kwargs: Any,
     ):
-        if isinstance(obj, Kernel):
-            if obj._transition_matrix is None:
-                raise ValueError(
+        if isinstance(object, KernelExpression):
+            if object.transition_matrix is None:
+                raise RuntimeError(
                     "Compute transition matrix first as `.compute_transition_matrix()`."
                 )
-            kernel = obj
-        elif isinstance(obj, KernelExpression):
-            # this will fail if not all kernels have transition matrix computed
-            kernel = obj.compute_transition_matrix()
-        elif isinstance(obj, (np.ndarray, spmatrix)):
-            kernel = PrecomputedKernel(obj)
-        elif isinstance(obj, AnnData):
-            if obsp_key is None:
-                raise ValueError(
-                    "Specify `obsp_key=...` when supplying an `AnnData` object."
-                )
-            elif obsp_key not in obj.obsp:
-                raise KeyError(
-                    f"Unable to find transition matrix in `adata.obsp[{obsp_key!r}]`."
-                )
-            kernel = PrecomputedKernel(obsp_key, adata=obj)
         else:
-            raise TypeError(
-                f"Expected an object of type `KernelExpression`, `numpy.ndarray`, `scipy.sparse.spmatrix` "
-                f"or `anndata.AnnData`, got `{type(obj).__name__}`."
-            )
-
-        super().__init__(kernel=kernel)
+            object = PrecomputedKernel(object, copy=False, **kwargs)
+        super().__init__(kernel=object)
 
         self._params: Dict[str, Any] = {}
         self._shadow_adata = AnnData(
diff --git a/cellrank/tl/estimators/mixins/_absorption_probabilities.py b/cellrank/tl/estimators/mixins/_absorption_probabilities.py
index 322301939..ad504ca4a 100644
--- a/cellrank/tl/estimators/mixins/_absorption_probabilities.py
+++ b/cellrank/tl/estimators/mixins/_absorption_probabilities.py
@@ -26,7 +26,7 @@
 
 import numpy as np
 import pandas as pd
-from scipy.sparse import issparse, spmatrix
+from scipy.sparse import issparse, spmatrix, csr_matrix
 from pandas.api.types import infer_dtype, is_categorical_dtype
 
 
@@ -280,11 +280,17 @@ def compute_absorption_probabilities(
         s = self.transition_matrix[trans_indices, :][:, rec_indices]
 
         # take individual solutions and piece them together to get absorption probabilities towards the classes
-        # fmt: off
-        macro_ix_helper = np.cumsum([0] + [len(indices) for indices in lookup_dict.values()])
-        s = np.concatenate([s[:, np.arange(a, b)].sum(axis=1) for a, b in _pairwise(macro_ix_helper)], axis=1)
-        # fmt: on
-
+        macro_ix_helper = np.cumsum(
+            [0] + [len(indices) for indices in lookup_dict.values()]
+        )
+        # `s` can be sparse or dense, ensure the correct shape
+        s = np.concatenate(
+            [
+                s[:, np.arange(a, b)].sum(axis=1).reshape(-1, 1)
+                for a, b in _pairwise(macro_ix_helper)
+            ],
+            axis=1,
+        )
         abs_probs = self._compute_absorption_probabilities(
             q,
             s,
@@ -298,9 +304,7 @@ def compute_absorption_probabilities(
             show_progress_bar=show_progress_bar,
             preconditioner=preconditioner,
         )
-        abs_probs = Lineage(abs_probs, names=keys, colors=colors)
 
-        abs_times = None
         if time_to_absorption is not None:
             lineages = _normalize_abs_times(keys, time_to_absorption=time_to_absorption)
             abs_times = _calculate_lineage_absorption_time_means(
@@ -318,12 +322,17 @@ def compute_absorption_probabilities(
                 preconditioner=preconditioner,
                 index=self.adata.obs_names,
             )
+        else:
+            abs_times = None
 
         params = self._create_params(
             remove=["use_petsc", "n_jobs", "backend", "show_progress_bar"]
         )
         self._write_absorption_probabilities(
-            abs_probs, abs_times, params=params, time=start
+            Lineage(abs_probs, names=keys, colors=colors),
+            abs_times,
+            params=params,
+            time=start,
         )
 
     @d.dedent
diff --git a/cellrank/tl/estimators/terminal_states/_gpcca.py b/cellrank/tl/estimators/terminal_states/_gpcca.py
index eb10527e0..33047aac9 100644
--- a/cellrank/tl/estimators/terminal_states/_gpcca.py
+++ b/cellrank/tl/estimators/terminal_states/_gpcca.py
@@ -74,11 +74,10 @@ class GPCCA(TermStatesEstimator, LinDriversMixin, SchurMixin, EigenMixin):
 
     def __init__(
         self,
-        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
-        obsp_key: Optional[str] = None,
+        object: Union[str, bool, np.ndarray, spmatrix, AnnData, KernelExpression],
         **kwargs: Any,
     ):
-        super().__init__(obj=obj, obsp_key=obsp_key, **kwargs)
+        super().__init__(object=object, **kwargs)
 
         self._coarse_init_dist: Optional[pd.Series] = None
         self._coarse_stat_dist: Optional[pd.Series] = None
@@ -427,7 +426,9 @@ def fit(
         if isinstance(n_states, int) and n_states == 1:
             self.compute_eigendecomposition()
 
-        self.compute_macrostates(n_states=n_states, cluster_key=cluster_key, **kwargs)
+        self.compute_macrostates(
+            n_states=n_states, cluster_key=cluster_key, n_cells=n_cells, **kwargs
+        )
 
         return self
 
diff --git a/cellrank/tl/estimators/terminal_states/_term_states_estimator.py b/cellrank/tl/estimators/terminal_states/_term_states_estimator.py
index 47b05d574..93af644ac 100644
--- a/cellrank/tl/estimators/terminal_states/_term_states_estimator.py
+++ b/cellrank/tl/estimators/terminal_states/_term_states_estimator.py
@@ -37,11 +37,10 @@ class TermStatesEstimator(BaseEstimator, ABC):
 
     def __init__(
         self,
-        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
-        obsp_key: Optional[str] = None,
+        object: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
         **kwargs: Any,
     ):
-        super().__init__(obj=obj, obsp_key=obsp_key, **kwargs)
+        super().__init__(object=object, **kwargs)
 
         self._term_states: Optional[pd.Series] = None
         self._term_states_probs: Optional[pd.Series] = None
diff --git a/cellrank/tl/kernels/__init__.py b/cellrank/tl/kernels/__init__.py
index 2cc5cbf5f..9fd962067 100644
--- a/cellrank/tl/kernels/__init__.py
+++ b/cellrank/tl/kernels/__init__.py
@@ -1,19 +1,9 @@
+from cellrank.tl.kernels import utils
 from cellrank.tl.kernels._base_kernel import Kernel
-from cellrank.tl.kernels._exp_time_kernel import ExperimentalTimeKernel
 from cellrank.tl.kernels._velocity_kernel import VelocityKernel
 from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEKernel
-from cellrank.tl.kernels._velocity_schemes import (
-    CosineScheme,
-    DotProductScheme,
-    CorrelationScheme,
-    SimilaritySchemeABC,
-)
 from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
 from cellrank.tl.kernels._precomputed_kernel import PrecomputedKernel
-from cellrank.tl.kernels._pseudotime_schemes import (
-    ThresholdSchemeABC,
-    HardThresholdScheme,
-    SoftThresholdScheme,
-)
 from cellrank.tl.kernels._connectivity_kernel import ConnectivityKernel
 from cellrank.tl.kernels._transport_map_kernel import TransportMapKernel
+from cellrank.tl.kernels._experimental_time_kernel import ExperimentalTimeKernel
diff --git a/cellrank/tl/kernels/_base_kernel.py b/cellrank/tl/kernels/_base_kernel.py
index e93f2fb7e..9ded7a82c 100644
--- a/cellrank/tl/kernels/_base_kernel.py
+++ b/cellrank/tl/kernels/_base_kernel.py
@@ -1,178 +1,52 @@
-"""Kernel module."""
-from typing import (
-    Any,
-    Dict,
-    List,
-    Type,
-    Tuple,
-    Union,
-    Callable,
-    Iterable,
-    Optional,
-    Sequence,
-)
-
-import warnings
+from typing import Any, Dict, List, Tuple, Union, Callable, Optional, Sequence
+
 from abc import ABC, abstractmethod
-from copy import copy
+from copy import deepcopy
 from pathlib import Path
-from functools import reduce
 
-import scvelo as scv
 from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import (
-    save_fig,
-    _connected,
-    _normalize,
-    _symmetric,
-    _get_neighs,
-    _irreducible,
-)
-from scvelo.plotting.utils import default_size, plot_outline
-from cellrank.tl._mixins._io import IOMixin
-from cellrank.tl.kernels._utils import _get_basis, _filter_kwargs
-from cellrank.tl.kernels._tmat_flow import FlowPlotter
-from cellrank.tl.kernels._random_walk import RandomWalk
+from cellrank.tl._utils import save_fig, _normalize
+from cellrank.tl._mixins import IOMixin
+from cellrank.tl.kernels.utils import RandomWalk, FlowPlotter
+from cellrank.tl.kernels._utils import require_tmat
+from cellrank.tl.kernels._mixins import BidirectionalMixin, UnidirectionalMixin
+from cellrank.tl.kernels.utils._projection import Projector
+from cellrank.tl.kernels.utils._random_walk import Indices_t
 
 import numpy as np
-from scipy.sparse import spdiags, issparse, spmatrix, csr_matrix, isspmatrix_csr
-from pandas.api.types import infer_dtype, is_numeric_dtype
-from pandas.core.dtypes.common import is_categorical_dtype
+from scipy.sparse import issparse, spmatrix, csr_matrix, isspmatrix_csr
 
 import matplotlib.pyplot as plt
-from matplotlib.colors import LinearSegmentedColormap, to_hex
-from matplotlib.collections import LineCollection
-
-_ERROR_DIRECTION_MSG = "Can only combine kernels that have the same direction."
-_ERROR_EMPTY_CACHE_MSG = (
-    "Fatal error: tried to used cached values, but the cache was empty."
-)
-_ERROR_CONF_ADAPT = (
-    "Confidence adaptive operator is only supported for kernels, found type `{!r}`."
-)
-_ERROR_VAR_NOT_FOUND = "Variances not found in kernel `{!r}`."
+from matplotlib.colors import LinearSegmentedColormap
 
-_LOG_USING_CACHE = "Using cached transition matrix"
+__all__ = ("UnidirectionalKernel", "BidirectionalKernel")
 
-_RTOL = 1e-12
-_n_dec = 2
-_dtype = np.float64
-_cond_num_tolerance = 1e-15
-Indices_t = Optional[
-    Union[Sequence[str], Dict[str, Union[str, Sequence[str], Tuple[float, float]]]]
-]
+Tmat_t = Union[np.ndarray, spmatrix]
 
 
 class KernelExpression(IOMixin, ABC):
-    """Base class for all kernels and kernel expressions."""
-
     def __init__(
         self,
-        op_name: Optional[str] = None,
-        backward: bool = False,
-        compute_cond_num: bool = False,
+        parent: Optional["KernelExpression"] = None,
+        **kwargs: Any,
     ):
-        self._op_name = op_name
+        super().__init__()
+        self._parent = parent
+        self._normalize = parent is None
         self._transition_matrix = None
-        self._backward = backward
-        self._compute_cond_num = compute_cond_num
-        self._cond_num = None
-        self._params = {}
-        self._normalize = True
-        self._parent = None
-
-    def __init_subclass__(cls, **kwargs: Any):
-        super().__init_subclass__()
-
-    @property
-    def condition_number(self) -> Optional[int]:
-        """Condition number of the transition matrix."""
-        return self._cond_num
-
-    @property
-    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
-        """
-        Return row-normalized transition matrix.
-
-        If not present, it is computed iff all underlying kernels have been initialized.
-        """
-
-        if self._parent is None and self._transition_matrix is None:
-            self.compute_transition_matrix()
-
-        return self._transition_matrix
-
-    @property
-    def backward(self) -> bool:
-        """Direction of the process."""
-        return self._backward
-
-    @property
-    @abstractmethod
-    def adata(self) -> AnnData:
-        """Annotated data object."""
-
-    @adata.setter
-    @abstractmethod
-    def adata(self, value: AnnData) -> None:
-        pass
-
-    @property
-    @abstractmethod
-    def shape(self) -> Tuple[int, int]:
-        """`(n_cells, n_cells)`."""
-
-    @property
-    def params(self) -> Dict[str, Any]:
-        """Parameters which are used to compute the transition matrix."""
-        if len(self.kernels) == 1:
-            return self._params
-        # we need some identifier
-        return {f"{repr(k)}:{i}": k.params for i, k in enumerate(self.kernels)}
-
-    def _format_params(self):
-        return ", ".join(
-            f"{k}={round(v, _n_dec) if isinstance(v, float) else v}"
-            for k, v in self.params.items()
-        )
-
-    @transition_matrix.setter
-    def transition_matrix(self, value: Union[np.ndarray, spmatrix]) -> None:
-        """
-        Set a new value of the transition matrix.
-
-        Parameters
-        ----------
-        value
-            The new transition matrix. If the expression has no parent, the matrix is normalized, if needed.
-
-        Returns
-        -------
-        None
-            Nothing, just updates the :attr:`transition_matrix` and optionally normalizes it.
-        """
-        # fmt: off
-        should_norm = ~np.isclose(np.asarray(value.sum(1)).squeeze(), 1.0, rtol=_RTOL).all()
-        # fmt: on
-
-        if self._parent is None:
-            value = _normalize(value) if should_norm else value
-        else:
-            # it's AND, not OR, because of combinations
-            value = _normalize(value) if self._normalize and should_norm else value
+        self._params: Dict[str, Any] = {}
 
-        if not np.all(np.isfinite(value.data if issparse(value) else value)):
-            raise ValueError("Not all values are finite.")
-        self._transition_matrix = value
+    def __init_subclass__(cls, **_: Any) -> None:
+        super().__init_subclass__()
 
     @abstractmethod
     def compute_transition_matrix(
         self, *args: Any, **kwargs: Any
     ) -> "KernelExpression":
         """
-        Compute a transition matrix.
+        Compute transition matrix.
 
         Parameters
         ----------
@@ -183,183 +57,126 @@ def compute_transition_matrix(
 
         Returns
         -------
-        :class:`cellrank.tl.kernels.KernelExpression`
-            Self.
-        """
-
-    @d.get_sections(base="write_to_adata", sections=["Parameters"])
-    @inject_docs()  # get rid of {{}}
-    @d.dedent
-    def write_to_adata(self, key: Optional[str] = None) -> None:
+        Modifies :attr:`transition_matrix` and returns self.
         """
-        Write the transition matrix and parameters used for computation to the underlying :attr:`adata` object.
 
-        Parameters
-        ----------
-        key
-            Key used when writing transition matrix to :attr:`adata`. If `None`, determine the key automatically.
+    @abstractmethod
+    def copy(self, *, deep: bool = False) -> "KernelExpression":
+        """Return a copy of itself. The underlying :attr:`adata` object is not copied."""
 
-        Returns
-        -------
-        None
-            %(write_to_adata)s
-        """
-        from cellrank._key import Key
+    @property
+    @abstractmethod
+    def adata(self) -> AnnData:
+        """Annotated data object."""
 
-        if self._transition_matrix is None:
-            raise ValueError(
-                "Compute transition matrix first as `.compute_transition_matrix()`."
-            )
+    @adata.setter
+    @abstractmethod
+    def adata(self, value: Optional[AnnData]) -> None:
+        pass
 
-        key = Key.uns.kernel(self.backward, key=key)
-        # retain the embedding info
-        self.adata.uns[f"{key}_params"] = {
-            **self.adata.uns.get(f"{key}_params", {}),
-            **{"params": self.params},
-        }
-        self.adata.obsp[key] = self.transition_matrix
+    @property
+    @abstractmethod
+    def kernels(self) -> Tuple["KernelExpression", ...]:
+        """Underlying base kernels."""
 
+    @property
     @abstractmethod
-    def copy(self) -> "KernelExpression":
-        """Return a copy of itself. Note that the underlying :attr:`adata` object is not copied."""
-
-    def _maybe_compute_cond_num(self) -> None:
-        """Optionally compute condition number."""
-        if self._compute_cond_num and self._cond_num is None:
-            logg.debug("Computing condition number")
-            self._cond_num = np.linalg.cond(
-                self._transition_matrix.toarray()
-                if issparse(self._transition_matrix)
-                else self._transition_matrix
-            )
-            if self._cond_num > _cond_num_tolerance:
-                logg.warning(
-                    f"Transition matrix may be ill-conditioned, its condition number is `{self._cond_num:.2e}`"
-                )
-            else:
-                logg.info(f"Condition number is `{self._cond_num:.2e}`")
+    def shape(self) -> Tuple[int, int]:
+        """``(n_cells, n_cells)``."""
 
-    def _reuse_cache(
-        self, expected_params: Dict[str, Any], *, time: Optional[Any] = None
-    ) -> bool:
-        try:
-            if expected_params == self._params:
-                assert self.transition_matrix is not None, _ERROR_EMPTY_CACHE_MSG
-                logg.debug(_LOG_USING_CACHE)
-                logg.info("    Finish", time=time)
-                return True
-            return False
-        except AssertionError:
-            raise
-        except Exception as e:  # noqa: B902
-            logg.error(f"Unable to load the cache, reason: `{e}`")
-            return False
-        finally:
-            self._params = expected_params
+    @property
+    @abstractmethod
+    def backward(self) -> Optional[bool]:
+        """Direction of the process."""
 
     @abstractmethod
-    def _get_kernels(self) -> Iterable["Kernel"]:
+    def __getitem__(self, ix: int) -> "KernelExpression":
         pass
 
-    @property
-    def kernels(self) -> List["Kernel"]:
-        """Get the kernels of the kernel expression, except for constants."""
-        return list(set(self._get_kernels()))
+    @abstractmethod
+    def __len__(self) -> int:
+        pass
 
-    def compute_projection(
+    @d.get_full_description(base="plot_single_flow")
+    @d.get_sections(base="plot_single_flow", sections=["Parameters", "Returns"])
+    @d.dedent
+    @require_tmat
+    def plot_single_flow(
         self,
-        basis: str = "umap",
-        key_added: Optional[str] = None,
-        copy: bool = False,
-    ) -> Optional[np.ndarray]:
+        cluster: str,
+        cluster_key: str,
+        time_key: str,
+        clusters: Optional[Sequence[Any]] = None,
+        time_points: Optional[Sequence[Union[int, float]]] = None,
+        min_flow: float = 0,
+        remove_empty_clusters: bool = True,
+        ascending: Optional[bool] = False,
+        legend_loc: Optional[str] = "upper right out",
+        alpha: Optional[float] = 0.8,
+        xticks_step_size: Optional[int] = 1,
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        show: bool = True,
+    ) -> Optional[plt.Axes]:
         """
-        Compute a projection of the transition matrix in the embedding.
-
-        Projections can only be calculated for kNN based kernels. The projected matrix
-        can be then visualized as::
-
-            scvelo.pl.velocity_embedding(adata, vkey='T_fwd', basis='umap')
+        Visualize outgoing flow from a cluster of cells :cite:`mittnenzweig:21`.
 
         Parameters
         ----------
-        basis
-            Basis in :attr:`anndata.AnnData.obsm` for which to compute the projection.
-        key_added
-            If not `None` and ``copy = False``, save the result to :attr:`anndata.AnnData.obsm` ``['{key_added}']``.
-            Otherwise, save the result to `'T_fwd_{basis}'` or `T_bwd_{basis}`, depending on the direction.
-        copy
-            Whether to return the projection or modify :attr:`adata` inplace.
+        cluster
+            Cluster for which to visualize outgoing flow.
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` where clustering is stored.
+        time_key
+            Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
+        clusters
+            Visualize flow only for these clusters. If `None`, use all clusters.
+        time_points
+            Visualize flow only for these time points. If `None`, use all time points.
+        %(flow.parameters)s
+        %(plotting)s
+        show
+            If `False`, return :class:`matplotlib.pyplot.Axes`.
 
         Returns
         -------
-        If ``copy=True``, the projection array of shape `(n_cells, n_components)`.
-        Otherwise, it modifies :attr:`anndata.AnnData.obsm` with a key based on ``key_added``.
-        """
-        # modified from: https://github.com/theislab/scvelo/blob/master/scvelo/tools/velocity_embedding.py
-        from cellrank._key import Key
-        from scvelo.tools.velocity_embedding import quiver_autoscale
-
-        if self._transition_matrix is None:
-            raise RuntimeError(
-                "Compute transition matrix first as `.compute_transition_matrix()`."
-            )
-
-        for kernel in self.kernels:
-            if kernel._conn is None:
-                raise AttributeError(
-                    f"{kernel!r} is not a kNN based kernel. The embedding projection "
-                    "only works for kNN based kernels."
-                )
-
-        start = logg.info(f"Projecting transition matrix onto `{basis}`")
-        emb = _get_basis(self.adata, basis)
-        T_emb = np.empty_like(emb)
-
-        conn = self.kernels[0]._conn
-
-        with warnings.catch_warnings():
-            warnings.simplefilter("ignore")
-            for row_id, row in enumerate(self.transition_matrix):
-                conn_idxs = conn[row_id, :].indices
-
-                dX = emb[conn_idxs] - emb[row_id, None]
-
-                if np.any(np.isnan(dX)):
-                    T_emb[row_id, :] = np.nan
-                else:
-                    probs = row[:, conn_idxs]
-                    if issparse(probs):
-                        probs = probs.A.squeeze()
-
-                    dX /= np.linalg.norm(dX, axis=1)[:, None]
-                    dX = np.nan_to_num(dX)
-                    T_emb[row_id, :] = probs.dot(dX) - dX.sum(0) / dX.shape[0]
-
-        T_emb /= 3 * quiver_autoscale(np.nan_to_num(emb), T_emb)
+        The axes object, if ``show = False``.
+        %(just_plots)s
 
-        if copy:
-            return T_emb
+        Notes
+        -----
+        This function is a Python reimplementation of the following
+        `original R function <https://github.com/tanaylab/embflow/blob/main/scripts/generate_paper_figures/plot_vein.r>`_
+        with some minor stylistic differences.
+        This function will not recreate the results from :cite:`mittnenzweig:21`, because there, the Metacell model
+        :cite:`baran:19` was used to compute the flow, whereas here the transition matrix is used.
+        """  # noqa: E501
+        fp = FlowPlotter(self.adata, self.transition_matrix, cluster_key, time_key)
+        fp = fp.prepare(cluster, clusters, time_points)
 
-        key = Key.uns.kernel(self.backward, key=key_added)
-        ukey = f"{key}_params"
+        ax = fp.plot(
+            min_flow=min_flow,
+            remove_empty_clusters=remove_empty_clusters,
+            ascending=ascending,
+            alpha=alpha,
+            xticks_step_size=xticks_step_size,
+            legend_loc=legend_loc,
+            figsize=figsize,
+            dpi=dpi,
+        )
 
-        embs = self.adata.uns.get(ukey, {}).get("embeddings", [])
-        if basis not in embs:
-            embs = list(embs) + [basis]
-            self.adata.uns[ukey] = self.adata.uns.get(ukey, {})
-            self.adata.uns[ukey]["embeddings"] = embs
+        if save is not None:
+            save_fig(ax.figure, save)
 
-        key = key + "_" + basis
-        logg.info(
-            f"Adding `adata.obsm[{key!r}]`\n    Finish",
-            time=start,
-        )
-        self.adata.obsm[key] = T_emb
+        if not show:
+            return ax
 
     @d.dedent
+    @require_tmat
     def plot_random_walks(
         self,
-        n_sims: int,
+        n_sims: int = 100,
         max_iter: Union[int, float] = 0.25,
         seed: Optional[int] = None,
         successive_hits: int = 0,
@@ -421,61 +238,9 @@ def plot_random_walks(
         %(just_plots)s
         For each random walk, the first/last cell is marked by the start/end colors of ``cmap``.
         """
-
-        def create_ixs(ixs: Indices_t, *, kind: str) -> Optional[np.ndarray]:
-            if ixs is None:
-                return None
-            if isinstance(ixs, dict):
-                # fmt: off
-                if len(ixs) != 1:
-                    raise ValueError(f"Expected to find only 1 cluster key, found `{len(ixs)}`.")
-                key = next(iter(ixs.keys()))
-                if key not in self.adata.obs:
-                    raise KeyError(f"Unable to find data in `adata.obs[{key!r}]`.")
-
-                vals = self.adata.obs[key]
-                if is_categorical_dtype(vals):
-                    ixs = np.where(np.isin(vals, ixs[key]))[0]
-                elif is_numeric_dtype(vals):
-                    if len(ixs[key]) != 2:
-                        raise ValueError(f"Expected range to be of length `2`, found `{len(ixs[key])}`")
-                    minn, maxx = sorted(ixs[key])
-                    ixs = np.where((vals >= minn) & (vals <= maxx))[0]
-                else:
-                    raise TypeError(f"Expected `adata.obs[{key!r}]` to be numeric or categorical, "
-                                    f"found `{infer_dtype(vals)}`.")
-                # fmt: on
-            elif isinstance(ixs, str):
-                ixs = np.where(self.adata.obs_names == ixs)[0]
-            else:
-                ixs = np.where(np.isin(self.adata.obs_names, ixs))[0]
-
-            if not len(ixs):
-                logg.warning(f"No {kind} indices have been selected, using `None`")
-                return None
-
-            return ixs
-
-        if self._transition_matrix is None:
-            raise RuntimeError(
-                "Compute transition matrix first as `.compute_transition_matrix()`."
-            )
-        emb = _get_basis(self.adata, basis)
-
-        if isinstance(cmap, str):
-            cmap = plt.get_cmap(cmap)
-        if not isinstance(cmap, LinearSegmentedColormap):
-            if not hasattr(cmap, "colors"):
-                raise AttributeError(
-                    "Unable to create a colormap, `cmap` does not have attribute `colors`."
-                )
-            cmap = LinearSegmentedColormap.from_list(
-                "random_walk", colors=cmap.colors, N=max_iter
-            )
-
-        start_ixs = create_ixs(start_ixs, kind="start")
-        stop_ixs = create_ixs(stop_ixs, kind="stop")
-        rw = RandomWalk(self.transition_matrix, start_ixs=start_ixs, stop_ixs=stop_ixs)
+        rw = RandomWalk(
+            self.adata, self.transition_matrix, start_ixs=start_ixs, stop_ixs=stop_ixs
+        )
         sims = rw.simulate_many(
             n_sims=n_sims,
             max_iter=max_iter,
@@ -486,329 +251,252 @@ def create_ixs(ixs: Indices_t, *, kind: str) -> Optional[np.ndarray]:
             show_progress_bar=show_progress_bar,
         )
 
-        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
-        scv.pl.scatter(self.adata, basis=basis, show=False, ax=ax, **kwargs)
-
-        logg.info("Plotting random walks")
-        for sim in sims:
-            x = emb[sim][:, 0]
-            y = emb[sim][:, 1]
-            points = np.array([x, y]).T.reshape(-1, 1, 2)
-            segments = np.concatenate([points[:-1], points[1:]], axis=1)
-            n_seg = len(segments)
-
-            lc = LineCollection(
-                segments,
-                linewidths=linewidth,
-                colors=[cmap(float(i) / n_seg) for i in range(n_seg)],
-                alpha=linealpha,
-                zorder=2,
-            )
-            ax.add_collection(lc)
-
-        for ix in [0, -1]:
-            ixs = [sim[ix] for sim in sims]
-            plot_outline(
-                x=emb[ixs][:, 0],
-                y=emb[ixs][:, 1],
-                outline_color=("black", to_hex(cmap(float(abs(ix))))),
-                kwargs={
-                    "s": kwargs.get("s", default_size(self.adata)) * 1.1,
-                    "alpha": 0.9,
-                },
-                ax=ax,
-                zorder=4,
-            )
-
-        if ixs_legend_loc not in (None, "none"):
-            from cellrank.pl._utils import _position_legend
-
-            h1 = ax.scatter([], [], color=cmap(0.0), label="start")
-            h2 = ax.scatter([], [], color=cmap(1.0), label="stop")
-            legend = ax.get_legend()
-            if legend is not None:
-                ax.add_artist(legend)
-            _position_legend(ax, legend_loc=ixs_legend_loc, handles=[h1, h2])
-
-        if save is not None:
-            save_fig(fig, save)
+        rw.plot(
+            sims,
+            basis=basis,
+            cmap=cmap,
+            linewidth=linewidth,
+            linealpha=linealpha,
+            ixs_legend_loc=ixs_legend_loc,
+            figsize=figsize,
+            dpi=dpi,
+            save=save,
+            **kwargs,
+        )
 
-    @d.get_full_description(base="plot_single_flow")
-    @d.get_sections(base="plot_single_flow", sections=["Parameters", "Returns"])
-    @d.dedent
-    def plot_single_flow(
+    @require_tmat
+    def plot_projection(
         self,
-        cluster: str,
-        cluster_key: str,
-        time_key: str,
-        clusters: Optional[Sequence[Any]] = None,
-        time_points: Optional[Sequence[Union[int, float]]] = None,
-        min_flow: float = 0,
-        remove_empty_clusters: bool = True,
-        ascending: Optional[bool] = False,
-        legend_loc: Optional[str] = "upper right out",
-        alpha: Optional[float] = 0.8,
-        xticks_step_size: Optional[int] = 1,
-        figsize: Optional[Tuple[float, float]] = None,
-        dpi: Optional[int] = None,
-        save: Optional[Union[str, Path]] = None,
-        show: bool = True,
-    ) -> Optional[plt.Axes]:
+        basis: str = "umap",
+        key_added: Optional[str] = None,
+        recompute: bool = False,
+        stream: bool = True,
+        **kwargs: Any,
+    ) -> None:
         """
-        Visualize outgoing flow from a cluster of cells :cite:`mittnenzweig:21`.
+        Plot :attr:`transition_matrix` as a stream or a grid plot.
 
         Parameters
         ----------
-        cluster
-            Cluster for which to visualize outgoing flow.
-        cluster_key
-            Key in :attr:`anndata.AnnData.obs` where clustering is stored.
-        time_key
-            Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
-        clusters
-            Visualize flow only for these clusters. If `None`, use all clusters.
-        time_points
-            Visualize flow only for these time points. If `None`, use all time points.
-        %(flow.parameters)s
-        %(plotting)s
-        show
-            If `False`, return :class:`matplotlib.pyplot.Axes`.
+        basis
+            Key in :attr:`anndata.AnnData.obsm` containing the basis onto which to project.
+        key_added
+            If not `None`, save the result to :attr:`anndata.AnnData.obsm` ``['{key_added}']``.
+            Otherwise, save the result to `'T_fwd_{basis}'` or `T_bwd_{basis}`, depending on the direction.
+        recompute
+            Whether to recompute the projection if it already exists.
+        stream
+            If ``True``, use :func:`scvelo.pl.velocity_embedding_stream`.
+            Otherwise, use :func:`scvelo.pl.velocity_embedding_grid`.
+        kwargs
+            Keyword argument for the plotting function.
 
         Returns
         -------
-        The axes object, if ``show = False``.
-        %(just_plots)s
-
-        Notes
-        -----
-        This function is a Python reimplementation of the following
-        `original R function <https://github.com/tanaylab/embflow/blob/main/scripts/generate_paper_figures/plot_vein.r>`_
-        with some minor stylistic differences.
-        This function will not recreate the results from :cite:`mittnenzweig:21`, because there, the Metacell model
-        :cite:`baran:19` was used to compute the flow, whereas here the transition matrix is used.
-        """  # noqa: E501
-        if self._transition_matrix is None:
-            raise RuntimeError(
-                "Compute transition matrix first as `.compute_transition_matrix()`."
-            )
-
-        fp = FlowPlotter(self.adata, self.transition_matrix, cluster_key, time_key)
-        fp = fp.prepare(cluster, clusters, time_points)
-
-        ax = fp.plot(
-            min_flow=min_flow,
-            remove_empty_clusters=remove_empty_clusters,
-            ascending=ascending,
-            alpha=alpha,
-            xticks_step_size=xticks_step_size,
-            legend_loc=legend_loc,
-            figsize=figsize,
-            dpi=dpi,
-        )
-
-        if save is not None:
-            save_fig(ax.figure, save)
-
-        if not show:
-            return ax
-
-    def __xor__(self, other: "KernelExpression") -> "KernelExpression":
-        return self.__rxor__(other)
-
-    def __rxor__(self, other: "KernelExpression") -> "KernelExpression":
-        def convert(obj):
-            if _is_adaptive_type(obj):
-                if obj._mat_scaler is None:
-                    raise ValueError(_ERROR_VAR_NOT_FOUND.format(obj))
-                return KernelMul(
-                    [
-                        ConstantMatrix(
-                            obj.adata, 1, obj._mat_scaler, backward=obj.backward
-                        ),
-                        obj,
-                    ]
-                )
-            if _is_adaptive_type(_is_bin_mult(obj, return_constant=False)):
-                e, c = _get_expr_and_constant(obj)
-                if e._mat_scaler is None:
-                    raise ValueError(_ERROR_VAR_NOT_FOUND.format(e))
-                return KernelMul(
-                    [ConstantMatrix(e.adata, c, e._mat_scaler, backward=e.backward), e]
-                )
-
-            return obj
-
-        if (
-            not _is_adaptive_type(self)
-            and not isinstance(self, KernelAdaptiveAdd)
-            and not _is_adaptive_type(_is_bin_mult(self, return_constant=False))
-        ):
-            raise TypeError(_ERROR_CONF_ADAPT.format(self.__class__.__name__))
-        if (
-            not _is_adaptive_type(other)
-            and not isinstance(other, KernelAdaptiveAdd)
-            and not _is_adaptive_type(_is_bin_mult(other, return_constant=False))
-        ):
-            raise TypeError(_ERROR_CONF_ADAPT.format(other.__class__.__name__))
-
-        s = convert(self)
-        o = convert(other)
-
-        if isinstance(s, KernelAdaptiveAdd):
-            exprs = list(s) + (list(o) if isinstance(o, KernelAdaptiveAdd) else [o])
-            # (c1 * x + c2 * y + ...) + (c3 * z) => (c1 * x + c2 * y + c3 + ... + * z)
-            if all(_is_bin_mult(k, ConstantMatrix) for k in exprs):
-                return KernelAdaptiveAdd(exprs)
-
-        # same but reverse
-        if isinstance(o, KernelAdaptiveAdd):
-            exprs = (list(s) if isinstance(s, KernelAdaptiveAdd) else [s]) + list(o)
-            if all(_is_bin_mult(k, ConstantMatrix) for k in exprs):
-                return KernelAdaptiveAdd(exprs)
-
-        return KernelAdaptiveAdd([s, o])
+        Nothing, just plots and modifies :attr:`anndata.AnnData.obsm` with a key based on ``key_added``.
+        """
+        proj = Projector(self, basis=basis)
+        proj.project(key_added=key_added, recompute=recompute)
+        proj.plot(stream=stream, **kwargs)
 
     def __add__(self, other: "KernelExpression") -> "KernelExpression":
         return self.__radd__(other)
 
     def __radd__(self, other: "KernelExpression") -> "KernelExpression":
+        def same_level_add(k1: "KernelExpression", k2: "KernelExpression") -> bool:
+            if not (isinstance(k1, KernelAdd) and isinstance(k2, KernelMul)):
+                return False
+
+            for kexpr in k1:
+                if not isinstance(kexpr, KernelMul):
+                    return False
+                if not kexpr._bin_consts:
+                    return False
+
+            return True
+
         if not isinstance(other, KernelExpression):
-            raise TypeError(
-                f"Expected type `KernelExpression`, found `{other.__class__.__name__}`."
-            )
+            return NotImplemented
 
-        s = self * 1 if isinstance(self, Kernel) else self
-        o = other * 1 if isinstance(other, Kernel) else other
-
-        if isinstance(s, KernelSimpleAdd):
-            exprs = list(s) + [o]
-            # (c1 * x + c2 * y + ...) + (c3 * z) => (c1 * x + c2 * y + c3 + ... + * z)
-            if all(_is_bin_mult(k) for k in exprs):
-                return KernelSimpleAdd(exprs)
-
-        # same but reverse
-        if isinstance(o, KernelSimpleAdd):
-            exprs = [s] + list(o)
-            if all(_is_bin_mult(k) for k in exprs):
-                return KernelSimpleAdd(exprs)
-
-        # (c1 + c2) => c3
-        if isinstance(s, Constant) and isinstance(o, Constant):
-            assert s.backward == o.backward, _ERROR_DIRECTION_MSG
-            return Constant(
-                s.adata, s.transition_matrix + o.transition_matrix, backward=s.backward
-            )
+        s = self * 1.0 if isinstance(self, Kernel) else self
+        o = other * 1.0 if isinstance(other, Kernel) else other
 
-        ss = s * 1 if not isinstance(s, KernelMul) else s
-        oo = o * 1 if not isinstance(o, KernelMul) else o
+        # (c1 * x + c2 * y + ...) + (c3 * z) => (c1 * x + c2 * y + c3 + ... + * z)
+        if same_level_add(s, o):
+            return KernelAdd(*tuple(s) + (o,))
+        if same_level_add(o, s):
+            return KernelAdd(*tuple(o) + (s,))
 
-        return KernelSimpleAdd([ss, oo])
+        # add virtual constant
+        if not isinstance(s, KernelMul):
+            s = s * 1.0
+        if not isinstance(o, KernelMul):
+            o = o * 1.0
 
-    def __rmul__(
-        self, other: Union[int, float, "KernelExpression"]
-    ) -> "KernelExpression":
-        return self.__mul__(other)
+        return KernelAdd(s, o)
 
     def __mul__(
         self, other: Union[float, int, "KernelExpression"]
     ) -> "KernelExpression":
+        return self.__rmul__(other)
+
+    def __rmul__(
+        self, other: Union[int, float, "KernelExpression"]
+    ) -> "KernelExpression":
+        def same_level_mul(k1: "KernelExpression", k2: "KernelExpression") -> bool:
+            return (
+                isinstance(k1, KernelMul)
+                and isinstance(k2, Constant)
+                and k1._bin_consts
+            )
 
-        if isinstance(other, (int, float)):
-            other = Constant(self.adata, other, backward=self.backward)
+        if isinstance(other, (int, float, np.integer, np.floating)):
+            other = Constant(self.adata, other)
 
         if not isinstance(other, KernelExpression):
-            raise TypeError(
-                f"Expected type `KernelExpression`, found `{other.__class__.__name__}`."
-            )
+            return NotImplemented
 
-        # (c1 * c2) => c3
-        if isinstance(self, Constant) and isinstance(
-            other, Constant
-        ):  # small optimization
-            assert self.backward == other.backward, _ERROR_DIRECTION_MSG
-            return Constant(
-                self.adata,
-                self.transition_matrix * other.transition_matrix,
-                backward=self.backward,
-            )
+        # fmt: off
+        s = self if isinstance(self, (KernelMul, Constant)) else KernelMul(Constant(self.adata, 1.0), self)
+        o = other if isinstance(other, (KernelMul, Constant)) else KernelMul(Constant(other.adata, 1.0), other)
+        # fmt: on
 
-        s = (
-            self
-            if isinstance(self, (KernelMul, Constant))
-            else KernelMul([Constant(self.adata, 1, backward=self.backward), self])
-        )
-        o = (
-            other
-            if isinstance(other, (KernelMul, Constant))
-            else KernelMul([Constant(other.adata, 1, backward=other.backward), other])
+        # at this point, only KernelMul and Constant is possible (two constants are not possible)
+        # (c1 * k) * c2 => (c1 * c2) * k
+        if same_level_mul(s, o):
+            c, expr = s._split_const
+            return KernelMul(c * o, expr)
+        if same_level_mul(o, s):
+            c, expr = o._split_const
+            return KernelMul(c * s, expr)
+
+        return KernelMul(s, o)
+
+    @d.get_sections(base="write_to_adata", sections=["Parameters"])
+    @inject_docs()  # gets rid of {{}} in %(write_to_adata)s
+    @d.dedent
+    @require_tmat
+    def write_to_adata(self, key: Optional[str] = None, copy: bool = False) -> None:
+        """
+        Write the transition matrix and parameters used for computation to the underlying :attr:`adata` object.
+
+        Parameters
+        ----------
+        key
+            Key used when writing transition matrix to :attr:`adata`.
+            If `None`, the key will be determined automatically.
+
+        Returns
+        -------
+        %(write_to_adata)s
+        """
+        from cellrank._key import Key
+
+        if self.adata is None:
+            raise ValueError("Underlying annotated data object is not set.")
+
+        key = Key.uns.kernel(self.backward, key=key)
+        # retain the embedding info
+        self.adata.uns[f"{key}_params"] = {
+            **self.adata.uns.get(f"{key}_params", {}),
+            **{"params": self.params},
+        }
+        self.adata.obsp[key] = (
+            self.transition_matrix.copy() if copy else self.transition_matrix
         )
 
-        cs, co = _is_bin_mult(s), _is_bin_mult(o)
-        if cs and isinstance(o, Constant):
-            cs._transition_matrix *= o.transition_matrix
-            return s
+    @property
+    def transition_matrix(self) -> Union[np.ndarray, csr_matrix]:
+        """Row-normalized transition matrix."""
+        if self._parent is None and self._transition_matrix is None:
+            self.compute_transition_matrix()
+        return self._transition_matrix
 
-        if co and isinstance(s, Constant):
-            co._transition_matrix *= s.transition_matrix
-            return o
+    @transition_matrix.setter
+    def transition_matrix(self, matrix: Union[np.ndarray, spmatrix]) -> None:
+        """
+        Set the transition matrix.
 
-        return KernelMul([s, o])
+        Parameters
+        ----------
+        matrix
+            Transition matrix. The matrix is row-normalized if necessary.
 
-    def __invert__(self: "KernelExpression") -> "KernelExpression":
-        # mustn't return a copy because transition matrix
-        self._transition_matrix = None
-        self._params = {}
-        self._backward = not self.backward
+        Returns
+        -------
+        Nothing, just updates the :attr:`transition_matrix` and optionally normalizes it.
+        """
+        # fmt: off
+        if matrix.shape != self.shape:
+            raise ValueError(
+                f"Expected matrix to be of shape `{self.shape}`, found `{matrix.shape}`."
+            )
 
-        return self
+        def should_norm(mat: Union[np.ndarray, spmatrix]) -> bool:
+            return not np.all(np.isclose(np.asarray(mat.sum(1)).squeeze(), 1.0, rtol=1e-12))
 
-    def __copy__(self) -> "KernelExpression":
-        return self.copy()
+        if issparse(matrix) and not isspmatrix_csr(matrix):
+            matrix = csr_matrix(matrix)
+        matrix = matrix.astype(np.float64, copy=False)
+
+        force_normalize = (self._parent is None or self._normalize) and should_norm(matrix)
+        if force_normalize:
+            if np.any((matrix.data if issparse(matrix) else matrix) < 0):
+                raise ValueError("Unable to normalize matrix with negative values.")
+            matrix = _normalize(matrix)
+            if should_norm(matrix):  # some rows are all 0s/contain invalid values
+                n_inv = np.sum(~np.isclose(np.asarray(matrix.sum(1)).squeeze(), 1.0, rtol=1e-12))
+                raise ValueError(f"Transition matrix is not row stochastic, {n_inv} rows do not sum to 1.")
+        # fmt: on
 
-    def __deepcopy__(self, memodict={}) -> "KernelExpression":  # noqa
-        res = self.copy()
-        if self._parent is None:
-            res.adata = self.adata.copy()
-        memodict[id(self)] = res
+        self._transition_matrix = matrix
 
-        return res
+    @property
+    def params(self) -> Dict[str, Any]:
+        """Parameters which are used to compute the transition matrix."""
+        if len(self.kernels) == 1:
+            return self._params
+        return {f"{k!r}:{i}": k.params for i, k in enumerate(self.kernels)}
 
+    def _reuse_cache(
+        self, expected_params: Dict[str, Any], *, time: Optional[Any] = None
+    ) -> bool:
+        # fmt: off
+        try:
+            if expected_params == self._params:
+                assert self.transition_matrix is not None
+                logg.debug("Using cached transition matrix")
+                logg.info("    Finish", time=time)
+                return True
+            return False
+        except AssertionError:
+            logg.warning("Transition matrix does not exist for the given parameters. Recomputing")
+            return False
+        except Exception as e:  # noqa: B902
+            # e.g. the dict is not comparable
+            logg.warning(f"Expected and actually parameters are not comparable, reason `{e}`. Recomputing")
+            expected_params = {}  # clear the params
+            return False
+        finally:
+            self._params = expected_params
+        # fmt: on
 
-class UnaryKernelExpression(KernelExpression, ABC):
-    """Base class for unary kernel expressions, such as kernels or constants."""
 
-    def __init__(
-        self,
-        adata,
-        backward: bool = False,
-        op_name: Optional[str] = None,
-        compute_cond_num: bool = False,
-        **kwargs,
-    ):
-        super().__init__(op_name, backward=backward, compute_cond_num=compute_cond_num)
-        assert (
-            op_name is None
-        ), "Unary kernel does not support any kind operation associated with it."
+class Kernel(KernelExpression, ABC):
+    def __init__(self, adata: AnnData, **kwargs: Any):
+        super().__init__(**kwargs)
         self._adata = adata
         self._n_obs = adata.n_obs
+        self._read_from_adata(**kwargs)
+
+    def _read_from_adata(self, **kwargs: Any) -> None:
+        pass
 
     @property
-    @d.dedent
     def adata(self) -> AnnData:
-        """
-        Annotated data object.
-
-        Returns
-        -------
-        %(adata_ret)s
-        """
+        """Annotated data object."""
         return self._adata
 
-    @property
-    def shape(self) -> Tuple[int, int]:
-        """`(n_cells, n_cells)`."""
-        return self._n_obs, self._n_obs
-
     @adata.setter
     def adata(self, adata: Optional[AnnData]) -> None:
         if adata is None:
@@ -816,579 +504,342 @@ def adata(self, adata: Optional[AnnData]) -> None:
             return
         if not isinstance(adata, AnnData):
             raise TypeError(
-                f"Expected argument of type `anndata.AnnData`, found `{type(adata).__name__!r}`."
+                f"Expected `adata` to be of type `AnnData`, found `{type(adata).__name__}`."
             )
         shape = (adata.n_obs, adata.n_obs)
         if self.shape != shape:
             raise ValueError(
-                f"Expected the new object to have same shape as previous object `{self.shape}`, "
-                f"found `{shape}`."
+                f"Expected new `AnnData` object to have same shape as the previous `{self.shape}`, found `{shape}`."
             )
         self._adata = adata
 
-    def __repr__(self):
-        return f"{'~' if self.backward and self._parent is None else ''}<{self.__class__.__name__}>"
-
-    def __str__(self):
-        params_fmt = self._format_params()
-        if params_fmt:
-            return (
-                f"{'~' if self.backward and self._parent is None else ''}"
-                f"<{self.__class__.__name__}[{params_fmt}]>"
-            )
-        return repr(self)
-
-
-class NaryKernelExpression(KernelExpression, ABC):
-    """Base class for n-ary kernel expressions."""
-
-    def __init__(self, kexprs: List[KernelExpression], op_name: Optional[str] = None):
-        assert len(kexprs), "No kernel expressions specified."
-
-        backward = kexprs[0].backward
-        assert all(k.backward == backward for k in kexprs), _ERROR_DIRECTION_MSG
+    def copy(self, *, deep: bool = False) -> "Kernel":
+        with self._remove_adata:
+            k = deepcopy(self)
+        k.adata = self.adata.copy() if deep else self.adata
+        return k
+
+    def _copy_ignore(self, *attrs: str) -> "Kernel":
+        # prevent copying attributes that are not necessary, e.g. during inversion
+        sentinel, attrs = object(), set(attrs)
+        objects = [
+            (attr, obj)
+            for attr, obj in ((attr, getattr(self, attr, sentinel)) for attr in attrs)
+            if obj is not sentinel
+        ]
+        try:
+            for attr, _ in objects:
+                setattr(self, attr, None)
+            return self.copy(deep=False)
+        finally:
+            for attr, obj in objects:
+                setattr(self, attr, obj)
 
-        # use OR instead of AND
-        super().__init__(
-            op_name,
-            backward=backward,
-            compute_cond_num=any(k._compute_cond_num for k in kexprs),
+    def _format_params(self) -> str:
+        return ", ".join(
+            f"{k}={round(v, 3) if isinstance(v, float) else v!r}"
+            for k, v in self.params.items()
         )
 
-        # copies of constants are necessary because of the recalculation
-        self._kexprs = [copy(k) if isinstance(k, Constant) else k for k in kexprs]
-
-        for kexprs in self._kexprs:
-            kexprs._parent = self
-
     @property
-    def shape(self) -> Tuple[int, int]:
-        """`(n_cells, n_cells)`."""
-        return self.kernels[0].shape
-
-    def _maybe_recalculate_constants(self, type_: Type):
-        if type_ == Constant:
-            accessor = "transition_matrix"
-        elif type_ == ConstantMatrix:
-            accessor = "_value"
-        else:
-            raise RuntimeError(
-                f"Unable to determine accessor for type `{type(type_).__name__}`."
-            )
-
-        constants = [_is_bin_mult(k, type_) for k in self]
-        if all(c is not None for c in constants):
-            assert all(
-                isinstance(getattr(c, accessor), (int, float)) for c in constants
-            )
-            total = sum(getattr(c, accessor) for c in constants) + 0.0
-            for c in constants:
-                c._recalculate(getattr(c, accessor) / total)
-
-            for kexpr in self._kexprs:  # don't normalize  (c * x)
-                kexpr._normalize = False
+    def transition_matrix(self) -> Union[np.ndarray, csr_matrix]:
+        """Row-normalized transition matrix."""
+        return self._transition_matrix
 
-    def _get_kernels(self) -> Iterable["Kernel"]:
-        for k in self:
-            if isinstance(k, Kernel) and not isinstance(k, (Constant, ConstantMatrix)):
-                yield k
-            elif isinstance(k, NaryKernelExpression):
-                yield from k._get_kernels()
+    @transition_matrix.setter
+    def transition_matrix(self, matrix: Any) -> None:
+        KernelExpression.transition_matrix.fset(self, matrix)
 
     @property
-    @d.dedent
-    def adata(self) -> AnnData:
-        """
-        Annotated data object.
-
-        Returns
-        -------
-        %(adata_ret)s
-        """
-        # we can do this because Constant requires adata as well
-        return self._kexprs[0].adata
+    def kernels(self) -> Tuple["KernelExpression", ...]:
+        return (self,)
 
-    @adata.setter
-    def adata(self, adata: Optional[AnnData]) -> None:
-        for kexpr in self._kexprs:
-            kexpr.adata = adata
+    @property
+    def shape(self) -> Tuple[int, int]:
+        """``(n_cells, n_cells)``."""
+        return self._n_obs, self._n_obs
 
-    def __invert__(self) -> "NaryKernelExpression":
-        super().__invert__()
-        self._kexprs = [~kexpr for kexpr in self]
+    def __getitem__(self, ix: int) -> "Kernel":
+        if ix != 0:
+            raise IndexError(ix)
         return self
 
     def __len__(self) -> int:
-        return len(self._kexprs)
-
-    def __getitem__(self, item) -> "KernelExpression":
-        return self._kexprs[item]
+        return 1
 
     def __repr__(self) -> str:
-        return (
-            f"{'~' if self.backward and self._parent is None else ''}("
-            + f" {self._op_name} ".join(repr(kexpr) for kexpr in self._kexprs)
-            + ")"
-        )
+        return f"{'~' if self.backward and self._parent is None else ''}{self.__class__.__name__}"
 
     def __str__(self) -> str:
-        return (
-            f"{'~' if self.backward and self._parent is None else ''}("
-            + f" {self._op_name} ".join(str(kexpr) for kexpr in self._kexprs)
-            + ")"
-        )
-
-
-@d.dedent
-class Kernel(UnaryKernelExpression, ABC):
-    """
-    A base class from which all kernels are derived.
-
-    These kernels read from a given AnnData object, usually the KNN graph and additional variables, to compute a
-    weighted, directed graph. Every kernel object has a direction. The kernels defined in the derived classes are not
-    strictly kernels in the mathematical sense because they often only take one input argument - however, they build
-    on other functions which have computed a similarity based on two input arguments. The role of the kernels defined
-    here is to add directionality to these symmetric similarity relations or to transform them.
-
-    Parameters
-    ----------
-    %(adata)s
-    %(backward)s
-    %(cond_num)s
-    check_connectivity
-        Check whether the underlying KNN graph is connected.
-    kwargs
-        Keyword arguments which can specify key to be read from :attr:`adata` object.
-    """
-
-    def __init__(
-        self,
-        adata: AnnData,
-        backward: bool = False,
-        compute_cond_num: bool = False,
-        check_connectivity: bool = False,
-        **kwargs: Any,
-    ):
-        super().__init__(
-            adata, backward, op_name=None, compute_cond_num=compute_cond_num, **kwargs
-        )
-        self._read_from_adata(check_connectivity=check_connectivity, **kwargs)
-
-    # TODO: move to a mixin class
-    def _read_from_adata(
-        self,
-        conn_key: Optional[str] = "connectivities",
-        read_conn: bool = True,
-        **kwargs: Any,
-    ) -> None:
-        """
-        Import the base-KNN graph and optionally check for symmetry and connectivity.
-
-        Parameters
-        ----------
-        conn_key
-            Key in :attr:`anndata.AnnData.uns` where connectivities are stored.
-        read_conn
-            Whether to read connectivities or set them to `None`. Useful when not exposing density normalization or
-            when KNN connectivities are not used to compute the transition matrix.
-        kwargs
-            Additional keyword arguments.
-
-        Returns
-        -------
-        Nothing, just sets :attr:`_conn`.
-        """
-        if not read_conn:
-            self._conn = None
-            return
-
-        self._conn = _get_neighs(
-            self.adata, mode="connectivities", key=conn_key
-        ).astype(_dtype)
-
-        check_connectivity = kwargs.pop("check_connectivity", False)
-        if check_connectivity:
-            start = logg.debug("Checking the KNN graph for connectedness")
-            if not _connected(self._conn):
-                logg.warning("KNN graph is not connected", time=start)
-            else:
-                logg.debug("KNN graph is connected", time=start)
-
-        start = logg.debug("Checking the KNN graph for symmetry")
-        if not _symmetric(self._conn):
-            logg.warning("KNN graph is not symmetric", time=start)
-        else:
-            logg.debug("KNN graph is symmetric", time=start)
-
-    def _density_normalize(
-        self, other: Union[np.ndarray, spmatrix]
-    ) -> Union[np.ndarray, spmatrix]:
-        """
-        Density normalization by the underlying KNN graph.
-
-        Parameters
-        ----------
-        other
-            Matrix to normalize.
-
-        Returns
-        -------
-        :class:`np.ndarray` or :class:`scipy.sparse.spmatrix`
-            Density normalized transition matrix.
-
-        Raises
-        ------
-        ValueError
-            If KNN connectivities are not set.
-        """
-        if self._conn is None:
-            raise ValueError(
-                "Unable to density normalize the transition matrix "
-                "because KNN connectivities are not set."
-            )
-
-        logg.debug("Density-normalizing the transition matrix")
-
-        q = np.asarray(self._conn.sum(axis=0))
-        Q = spdiags(1.0 / q, 0, other.shape[0], other.shape[0])
-
-        return Q @ other @ Q
-
-    def _get_kernels(self) -> Iterable["Kernel"]:
-        yield self
-
-    def _compute_transition_matrix(
-        self,
-        matrix: Union[np.ndarray, spmatrix],
-        density_normalize: bool = True,
-        check_irreducibility: bool = False,
-    ):
-        if matrix.shape[0] != matrix.shape[1]:
-            raise ValueError(f"Expected a square matrix, found `{matrix.shape}`.")
-        if matrix.shape[0] != self.adata.n_obs:
-            raise ValueError(
-                f"Expected matrix to be of shape `{(self.adata.n_obs, self.adata.n_obs)}`, "
-                f"found `{matrix.shape}`."
+        params_fmt = self._format_params()
+        if params_fmt:
+            return (
+                f"{'~' if self.backward and self._parent is None else ''}"
+                f"{self.__class__.__name__}[{params_fmt}]"
             )
+        return repr(self)
 
-        matrix = matrix.astype(_dtype)
-        if issparse(matrix) and not isspmatrix_csr(matrix):
-            matrix = csr_matrix(matrix)
 
-        # density correction based on node degrees in the KNN graph
-        if density_normalize:
-            matrix = self._density_normalize(matrix)
+class UnidirectionalKernel(UnidirectionalMixin, Kernel, ABC):
+    """Kernel with no directionality."""
 
-        # check for zero-rows
-        problematic_indices = np.where(np.array(matrix.sum(1)).flatten() == 0)[0]
-        if len(problematic_indices):
-            logg.warning(
-                f"Detected `{len(problematic_indices)}` absorbing states in the transition matrix. "
-                f"This matrix won't be irreducible"
-            )
-            matrix[problematic_indices, problematic_indices] = 1.0
 
-        if check_irreducibility:
-            _irreducible(matrix)
+class BidirectionalKernel(BidirectionalMixin, Kernel, ABC):
+    """Kernel with either forward or backward direction that can be inverted using :meth:`__invert__`."""
 
-        # setting this property automatically row-normalizes
-        self.transition_matrix = matrix
-        self._maybe_compute_cond_num()
 
+class Constant(UnidirectionalKernel):
+    def __init__(self, adata: AnnData, value: Union[int, float]):
+        super().__init__(adata)
+        self.transition_matrix = value
 
-@d.dedent
-class Constant(Kernel):
-    """
-    Kernel representing a multiplication by a constant number.
+    def compute_transition_matrix(self, value: Union[int, float]) -> "Constant":
+        self.transition_matrix = value
+        return self
 
-    Parameters
-    ----------
-    %(adata)s
-    value
-        Constant value by which to multiply. Must be a positive number.
-    %(backward)s
-    """
+    @property
+    def transition_matrix(self) -> Union[int, float]:
+        return self._transition_matrix
 
-    def __init__(
-        self, adata: AnnData, value: Union[int, float], backward: bool = False
-    ):
-        super().__init__(adata, backward=backward)
-        if not isinstance(value, (int, float)):
+    @transition_matrix.setter
+    def transition_matrix(self, value: Union[int, float]) -> None:
+        if not isinstance(value, (int, float, np.integer, np.floating)):
             raise TypeError(
                 f"Value must be a `float` or `int`, found `{type(value).__name__}`."
             )
         if value <= 0:
-            raise ValueError(f"Expected the constant to be positive, found `{value}`.")
-        self._recalculate(value)
+            raise ValueError(f"Expected the scalar to be positive, found `{value}`.")
 
-    def _recalculate(self, value) -> None:
         self._transition_matrix = value
         self._params = {"value": value}
 
-    def _read_from_adata(self, **kwargs: Any) -> None:
-        pass
+    def copy(self, *, deep: bool = False) -> "Constant":
+        return Constant(self.adata, self.transition_matrix)
 
-    def compute_transition_matrix(self, *args, **kwargs) -> "Constant":
-        """Return self."""
-        return self
+    # fmt: off
+    def __radd__(self, other: Union[int, float, "KernelExpression"]) -> "Constant":
+        if isinstance(other, (int, float, np.integer, np.floating)):
+            other = Constant(self.adata, other)
+        if isinstance(other, Constant):
+            if self.shape != other.shape:
+                raise ValueError(f"Expected kernel shape to be `{self.shape}`, found `{other.shape}`.")
+            return Constant(self.adata, self.transition_matrix + other.transition_matrix)
 
-    @d.dedent
-    def copy(self) -> "Constant":
-        """%(copy)s"""  # noqa: D400, D401
-        return Constant(self.adata, self.transition_matrix, self.backward)
+        return super().__radd__(other)
 
-    def __invert__(self) -> "Constant":
-        # do not call parent's invert, since it removes the transition matrix
-        self._backward = not self.backward
-        return self
+    def __rmul__(self, other: Union[int, float, "KernelExpression"]) -> "Constant":
+        if isinstance(other, (int, float, np.integer, np.floating)):
+            other = Constant(self.adata, other)
+        if isinstance(other, Constant):
+            if self.shape != other.shape:
+                raise ValueError(f"Expected kernel shape to be `{self.shape}`, found `{other.shape}`.")
+            return Constant(self.adata, self.transition_matrix * other.transition_matrix)
+
+        return super().__rmul__(other)
+    # fmt: on
 
     def __repr__(self) -> str:
-        return repr(round(self.transition_matrix, _n_dec))
+        return repr(round(self.transition_matrix, 3))
 
     def __str__(self) -> str:
-        return str(round(self.transition_matrix, _n_dec))
-
+        return str(round(self.transition_matrix, 3))
 
-class ConstantMatrix(Kernel):
-    """Kernel representing multiplication by a constant matrix."""
 
+class NaryKernelExpression(BidirectionalMixin, KernelExpression):
     def __init__(
-        self,
-        adata: AnnData,
-        value: Union[int, float],
-        variances: Union[np.ndarray, spmatrix],
-        backward: bool = False,
+        self, *kexprs: KernelExpression, parent: Optional[KernelExpression] = None
     ):
-        super().__init__(adata, backward=backward)
-        conn_shape = self._conn.shape
-        if variances.shape != conn_shape:
-            raise ValueError(
-                f"Expected variances of shape `{conn_shape}`, found `{variances.shape}`."
-            )
-        if not isinstance(value, (int, float)):
-            raise TypeError(
-                f"Value must be on `float` or `int`, found `{type(value).__name__!r}`."
-            )
-        if value <= 0:
-            raise ValueError(f"Expected the constant to be positive, found `{value}`.")
-
-        self._value = value
-        self._mat_scaler = (
-            csr_matrix(variances) if not issparse(variances) else variances
-        )
-        self._recalculate(value)
-
-    @d.dedent
-    def copy(self) -> "ConstantMatrix":
-        """%(copy)s"""  # noqa: D400, D401
-        return ConstantMatrix(
-            self.adata, self._value, copy(self._mat_scaler), self.backward
-        )
-
-    def _recalculate(self, value) -> None:
-        self._value = value
-        self._params = {"value": value}
-        self._transition_matrix = value * self._mat_scaler
-
-    def compute_transition_matrix(self, *args, **kwargs) -> "ConstantMatrix":
-        """Return self."""
-        return self
-
-    def __invert__(self) -> "ConstantMatrix":
-        # do not call parent's invert, since it removes the transition matrix
-        self._backward = not self.backward
-        return self
-
-    def __repr__(self) -> str:
-        return repr(round(self._value, _n_dec))
-
-    def __str__(self) -> str:
-        return str(round(self._value, _n_dec))
-
+        super().__init__(parent=parent)
+        self._validate(kexprs)
 
-class SimpleNaryExpression(NaryKernelExpression):
-    """Base class for n-ary operations."""
+    @abstractmethod
+    def _combine_transition_matrices(self, t1: Tmat_t, t2: Tmat_t) -> Tmat_t:
+        pass
 
-    def __init__(self, kexprs: List[KernelExpression], op_name: str, fn: Callable):
-        super().__init__(kexprs, op_name=op_name)
-        self._fn = fn
+    @property
+    @abstractmethod
+    def _initial_value(self) -> Union[int, float, Tmat_t]:
+        pass
 
-    def compute_transition_matrix(self, *args, **kwargs) -> "SimpleNaryExpression":
-        """Compute and combine the transition matrices."""
-        # must be done before, because the underlying expression don't have to be normed
-        if isinstance(self, KernelSimpleAdd):
-            self._maybe_recalculate_constants(Constant)
-        elif isinstance(self, KernelAdaptiveAdd):
-            self._maybe_recalculate_constants(ConstantMatrix)
+    @property
+    @abstractmethod
+    def _combiner(self) -> str:
+        pass
 
+    def compute_transition_matrix(self) -> "NaryKernelExpression":
         for kexpr in self:
-            if kexpr._transition_matrix is None:
+            if kexpr.transition_matrix is None:
                 if isinstance(kexpr, Kernel):
                     raise RuntimeError(
-                        f"Kernel `{kexpr}` is uninitialized. "
-                        f"Compute its transition matrix first as `.compute_transition_matrix()`."
+                        f"`{kexpr}` is uninitialized. Compute its "
+                        f"transition matrix first as `.compute_transition_matrix()`."
                     )
                 kexpr.compute_transition_matrix()
-            elif isinstance(kexpr, Kernel):
-                logg.debug(_LOG_USING_CACHE)
 
-        self.transition_matrix = csr_matrix(
-            self._fn([kexpr.transition_matrix for kexpr in self])
-        )
-
-        # only the top level expression and kernels will have condition number computed
-        if self._parent is None:
-            self._maybe_compute_cond_num()
+        tmat = self._initial_value
+        for kexpr in self:
+            tmat = self._combine_transition_matrices(tmat, kexpr.transition_matrix)
+        self.transition_matrix = tmat
 
         return self
 
-    @d.dedent
-    def copy(self) -> "SimpleNaryExpression":
-        """%(copy)s"""  # noqa: D400, D401
-        constructor = type(self)
-        kwargs = {"op_name": self._op_name, "fn": self._fn}
-
-        # preserve the type information so that combination can properly work
-        # we test for this
-        sne = constructor(
-            [copy(k) for k in self], **_filter_kwargs(constructor, **kwargs)
-        )
-
-        # TODO: copying's a bit buggy - the parent stays the same
-        # we could disallow it for inner expressions
-        sne._transition_matrix = copy(self._transition_matrix)
-        sne._condition_number = self._cond_num
-        sne._normalize = self._normalize
-
-        return sne
-
-
-class KernelAdd(SimpleNaryExpression):
-    """Base class that represents the addition of :class:`KernelExpression`."""
-
-    def __init__(self, kexprs: List[KernelExpression], op_name: str):
-        super().__init__(kexprs, op_name=op_name, fn=Reductor(np.add, 0))
+    def _validate(self, kexprs: Sequence[KernelExpression]) -> None:
+        if not len(kexprs):
+            raise ValueError("No kernels to combined.")
 
+        shapes = {kexpr.shape for kexpr in kexprs}
+        if len(shapes) > 1:
+            raise ValueError(
+                f"Expected all kernels to have the same shapes, found `{sorted(shapes)}`."
+            )
 
-class KernelSimpleAdd(KernelAdd):
-    """Addition of :class:`KernelExpression`."""
-
-    def __init__(self, kexprs: List[KernelExpression]):
-        super().__init__(kexprs, op_name="+")
-
-
-class KernelAdaptiveAdd(KernelAdd):
-    """Adaptive addition of :class:`KernelExpression`."""
-
-    def __init__(self, kexprs: List[KernelExpression]):
-        super().__init__(kexprs, op_name="^")
+        directions = {kexpr.backward for kexpr in kexprs}
+        if True in directions and False in directions:
+            raise ValueError("Unable to combine both forward and backward kernels.")
 
+        self._backward = (
+            True if True in directions else False if False in directions else None
+        )
+        self._kexprs = kexprs
+        for kexpr in self:
+            kexpr._parent = self
 
-class KernelMul(SimpleNaryExpression):
-    """Multiplication of :class:`KernelExpression`."""
+    @property
+    def adata(self) -> AnnData:
+        """Annotated data object."""
+        return self[0].adata
 
-    def __init__(self, kexprs: List[KernelExpression]):
-        super().__init__(kexprs, op_name="*", fn=Reductor(np.multiply, 1))
+    @adata.setter
+    def adata(self, adata: Optional[AnnData]) -> None:
+        # allow resetting (use for temp. pickling without adata)
+        for kexpr in self:
+            kexpr.adata = adata
 
+    @property
+    def kernels(self) -> Tuple["KernelExpression", ...]:
+        """Underlying unique basic kernels."""
+        kernels = []
+        for kexpr in self:
+            if isinstance(kexpr, Kernel) and not isinstance(kexpr, Constant):
+                kernels.append(kexpr)
+            elif isinstance(kexpr, NaryKernelExpression):  # recurse
+                kernels.extend(kexpr.kernels)
 
-class Reductor:
-    """
-    Class that reduces an iterable.
+        # return only unique kernels
+        return tuple(set(kernels))
 
-    We don't define a decorator because of pickling.
+    def copy(self, *, deep: bool = False) -> "KernelExpression":
+        kexprs = (k.copy(deep=deep) for k in self)
+        return type(self)(*kexprs, parent=self._parent)
 
-    Parameters
-    ----------
-    func
-        Reduction function.
-    initial
-        Initial value for :func:`reduce`.
-    """
+    @property
+    def shape(self) -> Tuple[int, int]:
+        """``(n_cells, n_cells)``."""
+        # all kernels have the same shape
+        return self[0].shape
 
-    def __init__(self, func: Callable, initial: Union[int, float]):
-        self._func = func
-        self._initial = initial
+    def _format(self, formatter: Callable[[KernelExpression], str]) -> str:
+        return (
+            f"{'~' if self.backward and self._parent is None else ''}("
+            + f" {self._combiner} ".join(formatter(kexpr) for kexpr in self)
+            + ")"
+        )
 
-    def __call__(self, seq: Iterable):  # noqa
-        return reduce(self._func, seq, self._initial)
+    def __getitem__(self, ix: int) -> "KernelExpression":
+        return self._kexprs[ix]
 
+    def __len__(self) -> int:
+        return len(self._kexprs)
 
-def _get_expr_and_constant(k: KernelMul) -> Tuple[KernelExpression, Union[int, float]]:
-    """
-    Get the value of a constant in binary multiplication.
+    def __invert__(self) -> "KernelExpression":
+        if self.backward is None:
+            return self.copy()
+        kexprs = tuple(
+            ~k if isinstance(k, BidirectionalMixin) else k.copy() for k in self
+        )
+        kexpr = type(self)(*kexprs, parent=self._parent)
+        kexpr._transition_matrix = None
+        return kexpr
 
-    Parameters
-    ----------
-    k
-        Binary multiplication involving a constant and a kernel.
+    def __repr__(self) -> str:
+        return self._format(repr)
 
-    Returns
-    -------
-    :class:`KernelExpression`, int or float
-        The expression which is being multiplied and the value of the constant.
-    """
+    def __str__(self) -> str:
+        return self._format(str)
 
-    if not isinstance(k, KernelMul):
-        raise TypeError(
-            f"Expected expression to be of type `KernelMul`, found `{type(k).__name__}`."
-        )
-    if len(k) != 2:
-        raise ValueError(
-            f"Expected expression to be binary, found `{len(k)}` subexpressions."
-        )
-    e1, e2 = k[0], k[1]
-
-    if isinstance(e1, Constant):
-        return e2, e1.transition_matrix
-    elif isinstance(e2, Constant):
-        return e1, e2.transition_matrix
-    else:
-        raise ValueError(
-            "Expected one of the subexpressions to be `Constant`, found "
-            f"`{type(e1).__name__}` and `{type(e2).__name__}`."
-        )
 
+class KernelAdd(NaryKernelExpression):
+    def compute_transition_matrix(self) -> "KernelAdd":
+        self._maybe_recalculate_constants()
+        return super().compute_transition_matrix()
 
-def _is_bin_mult(
-    k: KernelExpression,
-    const_type: Union[Type, Tuple[Type]] = Constant,
-    return_constant: bool = True,
-) -> Optional[KernelExpression]:
-    """
-    Check if an expression is a binary multiplication.
+    def _combine_transition_matrices(self, t1: Tmat_t, t2: Tmat_t) -> Tmat_t:
+        return t1 + t2
 
-    Parameters
-    ----------
-    k
-        Kernel expression to check.
-    const_type
-        Type of the constant.
-    return_constant
-        Whether to return the constant in the expression or the multiplied expression.
+    def _maybe_recalculate_constants(self) -> None:
+        """Normalize constants to sum to 1."""
+        constants = self._bin_consts
+        if constants:
+            total = sum((c.transition_matrix for c in constants), 0.0)
+            for c in constants:
+                c.transition_matrix = c.transition_matrix / total
+            for kexpr in self:  # don't normalize  (c * x)
+                kexpr._normalize = False
 
-    Returns
-    -------
-    None
-        If the expression is not a binary multiplication.
-    :class:`cellrank.tl.kernels.KernelExpression`
-        Depending on ``return_constant``, it either returns the constant multiplier or the expression being multiplied.
-    """
+    @property
+    def _bin_consts(self) -> List[KernelExpression]:
+        """Return constant expressions for each binary multiplication children with at least 1 constant."""
+        return [c for k in self if isinstance(k, KernelMul) for c in k._bin_consts]
 
-    if not isinstance(k, KernelMul):
-        return None
+    @property
+    def _combiner(self) -> str:
+        return "+"
 
-    if len(k) != 2:
-        return None
+    @property
+    def _initial_value(self) -> Union[int, float, Tmat_t]:
+        return 0.0
+
+
+class KernelMul(NaryKernelExpression):
+    def _combine_transition_matrices(self, t1: Tmat_t, t2: Tmat_t) -> Tmat_t:
+        if issparse(t1):
+            return t1.multiply(t2)
+        if issparse(t2):
+            return t2.multiply(t1)
+        return t1 * t2
+
+    def _format(self, formatter: Callable[[KernelExpression], str]) -> str:
+        fmt = super()._format(formatter)
+        if fmt[0] == "~" or not self._bin_consts:
+            return fmt
+        if fmt.startswith("("):
+            fmt = fmt[1:]
+        if fmt.endswith(")"):
+            fmt = fmt[:-1]
+        return fmt
 
-    lhs, rhs = k[0], k[1]
+    @property
+    def _bin_consts(self) -> List[KernelExpression]:
+        """Return all constants if this expression contains only 2 subexpressions."""
+        if len(self) != 2:
+            return []
 
-    if isinstance(lhs, const_type):
-        return lhs if return_constant else rhs
-    if isinstance(rhs, const_type):
-        return rhs if return_constant else lhs
+        return [k for k in self if isinstance(k, Constant)]
 
-    return None
+    @property
+    def _split_const(self) -> Tuple[Optional[Constant], Optional[KernelExpression]]:
+        """Return a constant and the other expression, iff this expression is of length 2 and contains a constant."""
+        if not self._bin_consts:
+            return None, None
+        k1, k2 = self
+        if isinstance(k1, Constant):
+            return k1, k2
+        return k2, k1
 
+    @property
+    def _combiner(self) -> str:
+        return "*"
 
-def _is_adaptive_type(k: KernelExpression) -> bool:
-    return isinstance(k, Kernel) and not isinstance(k, (Constant, ConstantMatrix))
+    @property
+    def _initial_value(self) -> Union[int, float, Tmat_t]:
+        return 1.0
diff --git a/cellrank/tl/kernels/_connectivity_kernel.py b/cellrank/tl/kernels/_connectivity_kernel.py
index e0edec166..bb2c6a5f3 100644
--- a/cellrank/tl/kernels/_connectivity_kernel.py
+++ b/cellrank/tl/kernels/_connectivity_kernel.py
@@ -1,13 +1,14 @@
-from copy import copy
-
 from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d
-from cellrank.tl.kernels import Kernel
+from cellrank.tl.kernels._mixins import ConnectivityMixin
+from cellrank.tl.kernels._base_kernel import UnidirectionalKernel
+
+__all__ = ("ConnectivityKernel",)
 
 
 @d.dedent
-class ConnectivityKernel(Kernel):
+class ConnectivityKernel(ConnectivityMixin, UnidirectionalKernel):
     """
     Kernel which computes transition probabilities based on similarities among cells.
 
@@ -18,37 +19,30 @@ class ConnectivityKernel(Kernel):
 
     The resulting transition matrix is symmetric and thus cannot be used to learn about the direction of the biological
     process. To include this direction, consider combining with a velocity-derived transition matrix via
-    :class:`cellrank.tl.kernels.VelocityKernel`.
+    :class:`cellrank.kernels.VelocityKernel`.
 
     %(density_correction)s
 
     Parameters
     ----------
     %(adata)s
-    %(backward)s
     conn_key
-        Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix describing cell-cell similarity.
-    %(cond_num)s
+        Key in :attr:`anndata.AnnData.obsp` where connectivity matrix describing cell-cell similarity is stored.
     check_connectivity
-        Check whether the underlying KNN graph is connected.
+        Check whether the underlying kNN graph is connected.
     """
 
     def __init__(
         self,
         adata: AnnData,
-        backward: bool = False,
         conn_key: str = "connectivities",
-        compute_cond_num: bool = False,
         check_connectivity: bool = False,
     ):
         super().__init__(
             adata,
-            backward=backward,
-            compute_cond_num=compute_cond_num,
-            check_connectivity=check_connectivity,
             conn_key=conn_key,
+            check_connectivity=check_connectivity,
         )
-        self._key = conn_key
 
     def compute_transition_matrix(
         self, density_normalize: bool = True
@@ -56,37 +50,26 @@ def compute_transition_matrix(
         """
         Compute transition matrix based on transcriptomic similarity.
 
-        Uses symmetric, weighted KNN graph to compute symmetric transition matrix. The connectivities are computed
+        Uses symmetric, weighted kNN graph to compute symmetric transition matrix. The connectivities are computed
         using :func:`scanpy.pp.neighbors`. Depending on the parameters used there, they can be UMAP connectivities or
         gaussian-kernel-based connectivities with adaptive kernel width.
 
         Parameters
         ----------
         density_normalize
-            Whether or not to use the underlying KNN graph for density normalization.
+            Whether or not to use the underlying kNN graph for density normalization.
 
         Returns
         -------
-        Self and updated :attr:`transition_matrix`.
+        Self and updates :attr:`transition_matrix` and :attr:`params`.
         """
-
         # fmt: off
-        start = logg.info(f"Computing transition matrix based on `adata.obsp[{self._key!r}]`")
-
-        if self._reuse_cache({"dnorm": density_normalize, "key": self._key}, time=start):
+        start = logg.info(f"Computing transition matrix based on `adata.obsp[{self._conn_key!r}]`")
+        if self._reuse_cache({"dnorm": density_normalize, "key": self._conn_key}, time=start):
             return self
 
-        self._compute_transition_matrix(matrix=self._conn.copy(), density_normalize=density_normalize)
+        self.transition_matrix = self._density_normalize(self._conn) if density_normalize else self._conn
         logg.info("    Finish", time=start)
         # fmt: on
 
         return self
-
-    def copy(self) -> "ConnectivityKernel":
-        """Return a copy of self."""
-        ck = ConnectivityKernel(self.adata, backward=self.backward)
-        ck._params = copy(self.params)
-        ck._cond_num = self.condition_number
-        ck._transition_matrix = copy(self._transition_matrix)
-
-        return ck
diff --git a/cellrank/tl/kernels/_cytotrace_kernel.py b/cellrank/tl/kernels/_cytotrace_kernel.py
index 20c444db5..8c38f2152 100644
--- a/cellrank/tl/kernels/_cytotrace_kernel.py
+++ b/cellrank/tl/kernels/_cytotrace_kernel.py
@@ -1,4 +1,4 @@
-from typing import Any, Optional
+from typing import Any, List, Tuple, Optional
 from typing_extensions import Literal
 
 from enum import auto
@@ -8,15 +8,18 @@
 from cellrank._key import Key
 from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import _correlation_test, _correlation_test_helper
+from cellrank.tl._utils import _correlation_test
 from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
 
 import numpy as np
+import pandas as pd
 from scipy.stats import gmean, hmean
 from scipy.sparse import issparse
 
+__all__ = ("CytoTRACEKernel",)
 
-class CytoTRACEAggregation(ModeEnum):  # noqa: D101
+
+class CytoTRACEAggregation(ModeEnum):
     MEAN = auto()
     MEDIAN = auto()
     GMEAN = auto()
@@ -33,7 +36,7 @@ class CytoTRACEKernel(PseudotimeKernel):
     CytoTRACE can be used to estimate cellular plasticity and in turn, a pseudotemporal ordering of cells from more
     plastic to less plastic states. It relies on the assumption that differentiated cells express, on average,
     less genes than naive cells.
-    This kernel internally uses the :class:`cellrank.tl.kernels.PseudotimeKernel` to direct the KNN graph
+    This kernel internally uses the :class:`cellrank.kernels.PseudotimeKernel` to direct the kNN graph
     on the basis of the CytoTRACE-derived pseudotime.
 
     %(density_correction)s
@@ -42,12 +45,8 @@ class CytoTRACEKernel(PseudotimeKernel):
     ----------
     %(adata)s
     %(backward)s
-    %(cytotrace.parameters)s
-    %(cond_num)s
-    check_connectivity
-        Check whether the underlying KNN graph is connected.
     kwargs
-        Keyword arguments for :class:`cellrank.tl.kernels.PseudotimeKernel`.
+        Keyword arguments for :class:`cellrank.kernels.PseudotimeKernel`.
 
     Example
     -------
@@ -64,63 +63,47 @@ class CytoTRACEKernel(PseudotimeKernel):
         sc.pp.log1p(adata)
         sc.pp.highly_variable_genes(adata)
 
-        # CytoTRACE by default uses imputed data - a simple way to compute KNN-imputed data is to use scVelo's moments
+        # CytoTRACE by default uses imputed data - a simple way to compute kNN-imputed data is to use scVelo's moments
         # function. However, note that this function expects `spliced` counts because it's designed for RNA velocity,
         # so we're using a simple hack here:
         if 'spliced' not in adata.layers or 'unspliced' not in adata.layers:
             adata.layers['spliced'] = adata.X
             adata.layers['unspliced'] = adata.X
 
-        # compute KNN-imputation using scVelo's moments function
+        # compute kNN-imputation using scVelo's moments function
         scv.pp.moments(adata)
 
         # import and initialize the CytoTRACE kernel, compute transition matrix - done!
-        from cellrank.tl.kernels import CytoTRACEKernel
-        ctk = CytoTRACEKernel(adata).compute_transition_matrix()
+        from cellrank.kernels import CytoTRACEKernel
+        ctk = CytoTRACEKernel(adata).compute_cytotrace().compute_transition_matrix()
     """
 
     def __init__(
         self,
         adata: AnnData,
         backward: bool = False,
-        layer: Optional[str] = "Ms",
-        aggregation: Literal[
-            "mean", "median", "hmean", "gmean"
-        ] = CytoTRACEAggregation.MEAN,
-        use_raw: bool = False,
-        n_top_genes: int = 200,
-        compute_cond_num: bool = False,
-        check_connectivity: bool = False,
         **kwargs: Any,
     ):
+        kwargs.pop("time_key", None)
         super().__init__(
             adata,
             backward=backward,
             time_key=Key.cytotrace("pseudotime"),
-            compute_cond_num=compute_cond_num,
-            check_connectivity=check_connectivity,
-            layer=layer,
-            aggregation=aggregation,
-            n_top_genes=n_top_genes,
-            use_raw=use_raw,
             **kwargs,
         )
-        # TODO(michlak8): needed by the PT kernel?
-        self._time_key = Key.cytotrace("pseudotime")
 
     def _read_from_adata(
         self,
-        time_key: str,
-        layer: Optional[str] = "Ms",
-        aggregation: Literal[
-            "mean", "median", "hmean", "gmean"
-        ] = CytoTRACEAggregation.MEAN,
-        use_raw: bool = True,
+        time_key: str = Key.cytotrace("pseudotime"),
         **kwargs: Any,
     ) -> None:
-        self.compute_cytotrace(layer=layer, aggregation=aggregation, use_raw=use_raw)
-
-        super()._read_from_adata(time_key=time_key, **kwargs)
+        try:
+            super()._read_from_adata(time_key=time_key, **kwargs)
+        except KeyError as e:
+            self._pseudotime: Optional[np.ndarray] = None
+            if "Unable to find pseudotime" not in str(e):
+                raise
+            super(PseudotimeKernel, self)._read_from_adata(**kwargs)
 
     @d.get_sections(base="cytotrace", sections=["Parameters"])
     @inject_docs(ct=CytoTRACEAggregation)
@@ -131,8 +114,8 @@ def compute_cytotrace(
             "mean", "median", "hmean", "gmean"
         ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = False,
-        n_top_genes: int = 200,
-    ) -> None:
+        n_genes: int = 200,
+    ) -> "CytoTRACEKernel":
         """
         Re-implementation of the CytoTRACE algorithm :cite:`gulati:20` to estimate cellular plasticity.
 
@@ -156,8 +139,8 @@ def compute_cytotrace(
         use_raw
             Whether to use the :attr:`anndata.AnnData.raw` to compute the number of genes expressed per cell
             (#genes/cell) and the correlation of gene expression across cells with #genes/cell.
-        n_top_genes
-            Number of genes used to compute the CytoTRACE score.
+        n_genes
+            Number of top positively correlated genes to compute the CytoTRACE score.
 
         Returns
         -------
@@ -171,7 +154,7 @@ def compute_cytotrace(
 
             - `'ct_gene_corr'` - the correlation as specified above.
             - `'ct_correlates'` - indication of the genes used to compute the CytoTRACE score, i.e. the ones that
-              correlated best with `'num_exp_genes'`.
+              correlated positively with `'ct_num_exp_genes'`.
 
         Notes
         -----
@@ -182,10 +165,6 @@ def compute_cytotrace(
 
         aggregation = CytoTRACEAggregation(aggregation)
 
-        if n_top_genes <= 0:
-            raise ValueError(
-                f"Expected `n_top_genes` to be `> 0`, found `{n_top_genes}`."
-            )
         if use_raw and self.adata.raw is None:
             logg.warning("`adata.raw` is `None`. Setting `use_raw=False`")
             use_raw = False
@@ -211,53 +190,29 @@ def compute_cytotrace(
             gene_names=adata_mraw.var_names,
         )["gene_corr"]
 
-        # fmt: off
-        top_genes = [g for g in gene_corr.sort_values(ascending=False).index if g in self.adata.var_names]
-        top_genes = top_genes[:n_top_genes]
-        null_genes = int(gene_corr.loc[top_genes].isnull().sum())
-        if null_genes:
-            raise ValueError(f"Top `{len(top_genes)}` correlated genes contain `{null_genes}` NaN values.")
-        if len(top_genes) != n_top_genes:
-            logg.warning(f"Unable to get requested top correlated `{n_top_genes}`. "
-                         f"Using top `{len(top_genes)}` genes")
-
-        logg.debug(
-            f"Aggregating imputed gene expression using "
-            f"`{aggregation}` top `{len(top_genes)}` in layer `{layer}`"
+        cytotrace_score, pos_top_genes = self._compute_score(
+            gene_corr,
+            layer=layer,
+            aggregation=aggregation,
+            ascending=False,
+            n_genes=n_genes,
         )
-        # compute mean/median over top n genes, aggregate over genes
-        if layer == "X":
-            imputed_exp = self.adata[:, top_genes].X
-        else:
-            imputed_exp = self.adata[:, top_genes].layers[layer]
-        if issparse(imputed_exp):
-            imputed_exp = imputed_exp.A
-
-        if aggregation == CytoTRACEAggregation.MEAN:
-            cytotrace_score = np.mean(imputed_exp, axis=1)
-        elif aggregation == CytoTRACEAggregation.MEDIAN:
-            cytotrace_score = np.median(imputed_exp, axis=1)
-        elif aggregation == CytoTRACEAggregation.GMEAN:
-            cytotrace_score = gmean(imputed_exp, axis=1)
-        elif aggregation == CytoTRACEAggregation.HMEAN:
-            cytotrace_score = hmean(imputed_exp, axis=1)
-        else:
-            raise NotImplementedError(f"Aggregation method `{aggregation}` is not yet implemented.")
-        # fmt: on
+        cytotrace_score -= cytotrace_score.min()
+        cytotrace_score /= cytotrace_score.max()
 
+        self.adata.obs[Key.cytotrace("score")] = cytotrace_score
+        self.adata.obs[Key.cytotrace("pseudotime")] = self._pseudotime = (
+            1 - cytotrace_score
+        )
         self.adata.obs[Key.cytotrace("num_exp_genes")] = num_exp_genes
         self.adata.var[Key.cytotrace("gene_corr")] = gene_corr
         self.adata.var[Key.cytotrace("correlates")] = False
-        self.adata.var.loc[top_genes, Key.cytotrace("correlates")] = True
-        cytotrace_score -= np.min(cytotrace_score)
-        cytotrace_score /= np.max(cytotrace_score)
-        self.adata.obs[Key.cytotrace("score")] = cytotrace_score
-        self.adata.obs[Key.cytotrace("pseudotime")] = 1 - cytotrace_score
+        self.adata.var.loc[pos_top_genes, Key.cytotrace("correlates")] = True
         self.adata.uns[Key.cytotrace("params")] = {
             "aggregation": aggregation,
             "layer": layer,
             "use_raw": use_raw,
-            "n_top_genes": len(top_genes),
+            "n_genes": len(pos_top_genes),
         }
 
         logg.info(
@@ -270,3 +225,70 @@ def compute_cytotrace(
             f"    Finish",
             time=start,
         )
+
+        return self
+
+    def _compute_score(
+        self,
+        gene_corr: pd.Series,
+        *,
+        layer: Optional[str] = None,
+        aggregation: CytoTRACEAggregation.MEAN,
+        ascending: bool = False,
+        n_genes: int = 200,
+    ) -> Tuple[np.ndarray, List[str]]:
+        from sklearn.utils.sparsefuncs import csc_median_axis_0
+
+        # fmt: off
+        modifier = "negatively" if ascending else "positively"
+        if n_genes <= 0:
+            raise ValueError(f"Expected number of {modifier} correlated genes to be positive, found `{n_genes}`.")
+        top_genes = [g for g in gene_corr.sort_values(ascending=ascending).index if g in self.adata.var_names]
+        top_genes = top_genes[:n_genes]
+        invalid_genes = int(gene_corr.loc[top_genes].isnull().sum())
+        # fmt: on
+
+        if invalid_genes:
+            raise ValueError(
+                f"Top `{len(top_genes)}` {modifier} correlated genes contain `{invalid_genes}` NaN values."
+            )
+        if not len(top_genes):
+            raise ValueError("No genes have been selected.")
+
+        if len(top_genes) != n_genes:
+            logg.warning(
+                f"Unable to get requested top {modifier} correlated `{n_genes}`. "
+                f"Using top `{len(top_genes)}` genes"
+            )
+
+        # fmt: off
+        if layer == "X":
+            imputed_exp = self.adata[:, top_genes].X
+        else:
+            imputed_exp = self.adata[:, top_genes].layers[layer]
+        if issparse(imputed_exp) and aggregation not in (CytoTRACEAggregation.MEAN, CytoTRACEAggregation.MEDIAN):
+            imputed_exp = imputed_exp.A
+
+        if aggregation == CytoTRACEAggregation.MEAN:
+            cytotrace_score = np.asarray(imputed_exp.mean(axis=1)).reshape((-1,))
+        elif aggregation == CytoTRACEAggregation.MEDIAN:
+            if issparse(imputed_exp):
+                cytotrace_score = np.asarray(csc_median_axis_0(imputed_exp.T.tocsc())).reshape((-1,))
+            else:
+                cytotrace_score = np.median(imputed_exp, axis=1)
+        elif aggregation == CytoTRACEAggregation.GMEAN:
+            cytotrace_score = gmean(imputed_exp, axis=1)
+        elif aggregation == CytoTRACEAggregation.HMEAN:
+            cytotrace_score = hmean(imputed_exp, axis=1)
+        else:
+            raise NotImplementedError(f"Aggregation method `{aggregation}` is not yet implemented.")
+        # fmt: on
+        return cytotrace_score, top_genes
+
+    def __invert__(self) -> "CytoTRACEKernel":
+        ck = self._copy_ignore("_transition_matrix")
+        if ck.pseudotime is not None:  # already computed
+            ck._pseudotime = np.max(ck.pseudotime) - ck.pseudotime
+        ck._backward = not self.backward
+        ck._params = {}
+        return ck
diff --git a/cellrank/tl/kernels/_exp_time_kernel.py b/cellrank/tl/kernels/_experimental_time_kernel.py
similarity index 58%
rename from cellrank/tl/kernels/_exp_time_kernel.py
rename to cellrank/tl/kernels/_experimental_time_kernel.py
index e6c1452dc..4849f24ae 100644
--- a/cellrank/tl/kernels/_exp_time_kernel.py
+++ b/cellrank/tl/kernels/_experimental_time_kernel.py
@@ -1,13 +1,11 @@
 from typing import Any, Optional
 
 from abc import ABC
-from copy import copy
 
 from anndata import AnnData
-from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d
-from cellrank.tl.kernels import Kernel
-from cellrank.tl.kernels._utils import _ensure_numeric_ordered
+from cellrank.tl.kernels._utils import require_tmat, _ensure_numeric_ordered
+from cellrank.tl.kernels._base_kernel import BidirectionalKernel
 
 import numpy as np
 import pandas as pd
@@ -15,15 +13,11 @@
 from matplotlib.colors import Normalize, to_hex
 from matplotlib.pyplot import get_cmap
 
-
-class LastTimePoint(ModeEnum):  # noqa: D101
-    UNIFORM = "uniform"
-    DIAGONAL = "diagonal"
-    CONNECTIVITIES = "connectivities"
+__all__ = ("ExperimentalTimeKernel",)
 
 
 @d.dedent
-class ExperimentalTimeKernel(Kernel, ABC):
+class ExperimentalTimeKernel(BidirectionalKernel, ABC):
     """
     Kernel base class which computes directed transition probabilities based on experimental time.
 
@@ -36,7 +30,8 @@ class ExperimentalTimeKernel(Kernel, ABC):
     time_key
         Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
         The experimental time can be of either of a numeric or an ordered categorical type.
-    %(cond_num)s
+    kwargs
+        Keyword arguments for the parent class.
     """
 
     def __init__(
@@ -44,34 +39,33 @@ def __init__(
         adata: AnnData,
         backward: bool = False,
         time_key: str = "exp_time",
-        compute_cond_num: bool = False,
         **kwargs: Any,
     ):
         super().__init__(
             adata,
             backward=backward,
             time_key=time_key,
-            compute_cond_num=compute_cond_num,
-            check_connectivity=False,
             **kwargs,
         )
-        self._time_key = time_key
 
-    def _read_from_adata(self, **kwargs: Any) -> None:
+    def _read_from_adata(
+        self, time_key: str = "exp_time", cmap: str = "gnuplot", **kwargs: Any
+    ) -> None:
         super()._read_from_adata(**kwargs)
 
-        time_key = kwargs.pop("time_key", "exp_time")
+        self._time_key = time_key
         self._exp_time = _ensure_numeric_ordered(self.adata, time_key)
         self.adata.obs[time_key] = self.experimental_time.values
 
         # fmt: off
-        cmap = get_cmap(kwargs.pop("cmap", "gnuplot"))
+        cmap = get_cmap(cmap)
         cats = self.experimental_time.cat.categories
         norm = Normalize(vmin=cats.min(), vmax=cats.max())
         self.adata.uns[f"{time_key}_colors"] = np.array([to_hex(c) for c in cmap(norm(cats))])
         # fmt: on
 
     @d.dedent
+    @require_tmat
     def plot_single_flow(
         self,
         cluster: str,
@@ -100,24 +94,13 @@ def experimental_time(self) -> pd.Series:
         """Experimental time."""
         return self._exp_time
 
-    def copy(self) -> "ExperimentalTimeKernel":
-        """Return a copy of self."""
-        pk = type(self)(self.adata, backward=self.backward, time_key=self._time_key)
-        pk._exp_time = copy(self.experimental_time)
-        pk._params = copy(self._params)
-        pk._cond_num = self.condition_number
-        pk._transition_matrix = copy(self._transition_matrix)
-
-        return pk
-
-    def __invert__(self) -> "ExperimentalTimeKernel":
-        super().__invert__()
-        if len(self.experimental_time.cat.categories) > 1:
-            minn, maxx = self.experimental_time.min(), self.experimental_time.max()
-            self._exp_time = pd.Series(
-                pd.Categorical(
-                    maxx - np.array(self.experimental_time) + minn, ordered=True
-                ),
-                index=self.experimental_time.index,
-            )
-        return self
+    def __invert__(self, *args: str) -> "ExperimentalTimeKernel":
+        etk = self._copy_ignore(*(("_transition_matrix",) + args))
+        minn, maxx = etk.experimental_time.min(), etk.experimental_time.max()
+        etk._exp_time = pd.Series(
+            pd.Categorical(maxx - np.array(etk.experimental_time) + minn, ordered=True),
+            index=etk.experimental_time.index,
+        )
+        etk._backward = not self.backward
+        etk._params = {}
+        return etk
diff --git a/cellrank/tl/kernels/_mixins.py b/cellrank/tl/kernels/_mixins.py
new file mode 100644
index 000000000..e3050d1bb
--- /dev/null
+++ b/cellrank/tl/kernels/_mixins.py
@@ -0,0 +1,87 @@
+from typing import Any, Union
+
+from abc import ABC, abstractmethod
+
+from cellrank import logging as logg
+from cellrank.tl._utils import _connected, _symmetric
+from cellrank.ul._utils import _get_neighs
+
+import numpy as np
+from scipy.sparse import spdiags, spmatrix, csr_matrix
+
+__all__ = ("ConnectivityMixin", "UnidirectionalMixin", "BidirectionalMixin")
+
+
+class ConnectivityMixin:
+    """Mixin class that reads kNN connectivities and allows for density normalization."""
+
+    def _read_from_adata(
+        self,
+        conn_key: str = "connectivities",
+        check_connectivity: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        super()._read_from_adata(**kwargs)
+        # fmt: off
+        self._conn_key = conn_key
+        self._conn = _get_neighs(self.adata, mode="connectivities", key=conn_key)
+        self._conn = csr_matrix(self._conn).astype(np.float64, copy=False)
+        # fmt: on
+        if check_connectivity:
+            if not _connected(self._conn):
+                logg.warning("kNN graph is not connected")
+
+        if not _symmetric(self._conn):
+            logg.warning("kNN graph is not symmetric")
+
+    def _density_normalize(
+        self, matrix: Union[np.ndarray, spmatrix]
+    ) -> Union[np.ndarray, spmatrix]:
+        """
+        Density normalization by the underlying kNN graph.
+
+        Parameters
+        ----------
+        matrix
+            Matrix to normalize.
+
+        Returns
+        -------
+        Density normalized matrix.
+        """
+        logg.debug("Density normalizing the transition matrix")
+
+        q = np.asarray(self._conn.sum(axis=0)).squeeze()
+        Q = spdiags(1.0 / q, 0, matrix.shape[0], matrix.shape[0])
+
+        return Q @ matrix @ Q
+
+
+class UnidirectionalMixin:
+    """Mixin specifying that its kernel doesn't have a direction."""
+
+    @property
+    def backward(self) -> None:
+        """None."""
+        return None
+
+
+class BidirectionalMixin(ABC):
+    """Mixin specifying that its kernel has forward or backward directions."""
+
+    def __init__(self, *args: Any, backward: bool = False, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        if not isinstance(backward, bool):
+            raise TypeError(
+                f"Expected `backward` to be `bool`, found `{type(backward).__name__}`."
+            )
+        self._backward = backward
+
+    @abstractmethod
+    def __invert__(self) -> "BidirectionalMixin":
+        pass
+
+    @property
+    def backward(self) -> bool:
+        """Direction of the process."""
+        return self._backward
diff --git a/cellrank/tl/kernels/_precomputed_kernel.py b/cellrank/tl/kernels/_precomputed_kernel.py
index a381c43dd..b0840d3fb 100644
--- a/cellrank/tl/kernels/_precomputed_kernel.py
+++ b/cellrank/tl/kernels/_precomputed_kernel.py
@@ -1,142 +1,134 @@
-from typing import Any, Dict, Union, Optional
-
-from copy import copy
+from typing import Any, Union, Optional
 
 from anndata import AnnData
 from cellrank import logging as logg
 from cellrank._key import Key
 from cellrank.ul._docs import d
 from cellrank.ul._utils import _read_graph_data
-from cellrank.tl.kernels._base_kernel import _RTOL, Kernel, KernelExpression
+from cellrank.tl.kernels._base_kernel import KernelExpression, UnidirectionalKernel
 
 import numpy as np
 from scipy.sparse import spmatrix, csr_matrix
 
+__all__ = ("PrecomputedKernel",)
+
 
 @d.dedent
-class PrecomputedKernel(Kernel):
+class PrecomputedKernel(UnidirectionalKernel):
     """
     Kernel which contains a precomputed transition matrix.
 
     Parameters
     ----------
-    transition_matrix
-        Row-normalized transition matrix or a key in :attr:`anndata.AnnData.obsp`.
-        or a :class:`cellrank.tl.kernels.KernelExpression` with a precomputed transition matrix.
-        If `None`, try to determine the key based on ``backward``.
+    object
+        Can be one of the following types:
+
+            - :class:`anndata.AnnData` - annotated data object.
+            - :class:`scipy.sparse.spmatrix`, :class:`numpy.ndarray` - row-normalized transition matrix.
+            - :class:`cellrank.tl.kernels.KernelExpression` - kernel expression.
+            - :class:`str` - key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
+              ``adata`` must be provided in this case.
+            - :class:`bool` - directionality of the transition matrix that will be used to infer its storage location.
+              If `None`, the directionality will be determined automatically. ``adata`` must be provided in this case.
     %(adata)s
-        If `None`, a temporary placeholder object is created.
-    %(backward)s
-    %(cond_num)s
-    kwargs
-        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
+        Must be provided when ``object`` is :class:`str` or :class:`bool`.
+    obsp_key
+        Key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
+        If `None`, it will be determined automatically. Only used when ``object`` is :class:`anndata.AnnData`.
+    copy
+        Whether or not to copy the stored transition matrix.
+    backward
+        Hint whether this is a forward, backward or a unidirectional kernel. Only used when ``object`` is
+        :class:`anndata.AnnData`.
     """
 
+    _SENTINEL = object()
+
     def __init__(
         self,
-        transition_matrix: Optional[
-            Union[np.ndarray, spmatrix, KernelExpression, str]
-        ] = None,
+        object: Union[str, bool, np.ndarray, spmatrix, AnnData, KernelExpression],
         adata: Optional[AnnData] = None,
-        backward: bool = False,
-        compute_cond_num: bool = False,
+        obsp_key: Optional[str] = None,
         **kwargs: Any,
     ):
-        origin = "'array'"
-        params: Dict[str, Any] = {}
-
-        if transition_matrix is None:
-            transition_matrix = Key.uns.kernel(backward)
-            logg.info(f"Accessing `adata.obsp[{transition_matrix!r}]`")
-
-        if isinstance(transition_matrix, str):
-            if adata is None:
-                raise ValueError(
-                    "When `transition_matrix` specifies a key to `adata.obsp`, `adata` cannot be None."
-                )
-            origin = f"adata.obsp[{transition_matrix!r}]"
-            backward = Key.uns.kernel(bwd=True) == transition_matrix
-            transition_matrix = _read_graph_data(adata, transition_matrix)
-        elif isinstance(transition_matrix, KernelExpression):
-            if transition_matrix._transition_matrix is None:
-                raise ValueError(
-                    "Compute transition matrix first as `.compute_transition_matrix()`."
-                )
-            if adata is not None and adata is not transition_matrix.adata:
-                logg.warning(
-                    "Ignoring supplied `adata` object because it differs from the kernel's `adata` object."
-                )
-
-            # use `str` rather than `repr` because it captures the parameters
-            origin = str(transition_matrix).strip("~<>")
-            params = transition_matrix.params.copy()
-            backward = transition_matrix.backward
-            adata = transition_matrix.adata
-            transition_matrix = transition_matrix.transition_matrix
-
-        if not isinstance(transition_matrix, (np.ndarray, spmatrix)):
+        if isinstance(object, AnnData):
+            self._from_adata(object, obsp_key=obsp_key, **kwargs)
+        elif isinstance(object, KernelExpression):
+            self._from_kernel(object, copy=kwargs.get("copy", False))
+        elif isinstance(object, (np.ndarray, spmatrix)):
+            self._from_matrix(object, adata=adata, **kwargs)
+        elif isinstance(adata, AnnData):
+            if isinstance(object, str):
+                self._from_adata(adata, obsp_key=object, **kwargs)
+            elif object is None or isinstance(object, bool):
+                kwargs["backward"] = object
+                self._from_adata(adata, **kwargs)
+            else:
+                raise ValueError("Unable to interpret the data.")
+        else:
             raise TypeError(
-                f"Expected transition matrix to be of type `numpy.ndarray` or `scipy.sparse.spmatrix`, "
-                f"found `{type(transition_matrix).__name__}`."
+                f"Expected `object` to be either `str`, `bool`, `numpy.ndarray`, "
+                f"`scipy.sparse.spmatrix`, `AnnData` or `KernelExpression`, "
+                f"found `{type(object).__name__}`"
             )
 
-        if transition_matrix.shape[0] != transition_matrix.shape[1]:
-            raise ValueError(
-                f"Expected transition matrix to be square, found `{transition_matrix.shape}`."
-            )
-
-        if not np.allclose(np.sum(transition_matrix, axis=1), 1.0, rtol=_RTOL):
-            raise ValueError("Not a valid transition matrix, not all rows sum to 1.")
-
-        if adata is None:
-            logg.warning("Creating empty `AnnData` object")
-            adata = AnnData(
-                csr_matrix((transition_matrix.shape[0], 1), dtype=np.float32)
+    def _from_adata(
+        self,
+        adata: AnnData,
+        obsp_key: Optional[str] = None,
+        backward: Optional[bool] = _SENTINEL,
+        copy: bool = False,
+    ) -> None:
+        if obsp_key is None:
+            obsp_key = Key.uns.kernel(backward)
+        tmat = _read_graph_data(adata, obsp_key)
+        if backward is PrecomputedKernel._SENTINEL:
+            # not ideal, since None/False share the same key, we prefer False
+            if obsp_key == Key.uns.kernel(bwd=False):
+                backward = False
+            elif obsp_key == Key.uns.kernel(bwd=True):
+                backward = True
+            else:
+                backward = None
+        self._from_matrix(tmat, adata=adata, backward=backward, copy=copy)
+        self.params["origin"] = f"adata.obsp[{obsp_key!r}]"
+
+    def _from_kernel(self, kernel: KernelExpression, copy: bool = False) -> None:
+        if kernel.transition_matrix is None:
+            raise RuntimeError(
+                "Compute transition matrix first as `.compute_transition_matrix()`."
             )
-
-        super().__init__(
-            adata, backward=backward, compute_cond_num=compute_cond_num, **kwargs
+        self._from_matrix(
+            kernel.transition_matrix,
+            backward=kernel.backward,
+            adata=kernel.adata,
+            copy=copy,
         )
+        self._params = kernel.params.copy()
+        self.params["origin"] = repr(kernel)
 
-        self._transition_matrix = csr_matrix(transition_matrix)
-        self._maybe_compute_cond_num()
-        self._params = params
-        self._origin = origin
-
-    def _read_from_adata(self, **kwargs: Any) -> None:
-        self._conn = None
-
-    @d.dedent
-    def copy(self) -> "PrecomputedKernel":
-        """%(copy)s"""  # noqa
-        pk = PrecomputedKernel(
-            copy(self.transition_matrix),
-            adata=self.adata,
-            backward=self.backward,
-            compute_cond_num=False,
-        )
-        pk._cond_num = self.condition_number
-        pk._origin = self._origin
-        pk._params = self._params.copy()
-
-        return pk
-
-    def compute_transition_matrix(
-        self, *args: Any, **kwargs: Any
-    ) -> "PrecomputedKernel":
-        """Return self."""
-        return self
-
-    def __invert__(self) -> "PrecomputedKernel":
-        # do not call parent's invert, since it removes the transition matrix
-        self._backward = not self.backward
+    def _from_matrix(
+        self,
+        matrix: Union[np.ndarray, spmatrix],
+        adata: Optional[AnnData] = None,
+        backward: Optional[bool] = None,
+        copy: bool = False,
+    ) -> None:
+        # fmt: off
+        if adata is None:
+            logg.warning(f"Creating empty `AnnData` object of shape `{matrix.shape[0], 1}`")
+            adata = AnnData(csr_matrix((matrix.shape[0], 1), dtype=np.float64))
+        super().__init__(adata)
+        self._backward: Optional[bool] = backward
+        self.transition_matrix = matrix.copy() if copy else matrix
+        self.params['origin'] = "array"
+        # fmt: on
+
+    def compute_transition_matrix(self, *_: Any, **__: Any) -> "PrecomputedKernel":
+        """Do nothing and return self."""
         return self
 
-    def __repr__(self):
-        return (
-            f"{'~' if self.backward and self._parent is None else ''}"
-            f"<{self.__class__.__name__}[origin={self._origin}]>"
-        )
-
-    def __str__(self):
-        return repr(self)
+    @property
+    def backward(self) -> Optional[bool]:
+        """Direction of the process."""
+        return self._backward
diff --git a/cellrank/tl/kernels/_pseudotime_kernel.py b/cellrank/tl/kernels/_pseudotime_kernel.py
index 7b6e1e70e..0db7cf930 100644
--- a/cellrank/tl/kernels/_pseudotime_kernel.py
+++ b/cellrank/tl/kernels/_pseudotime_kernel.py
@@ -1,17 +1,16 @@
 from typing import Any, Union, Callable, Optional
 from typing_extensions import Literal
 
-from copy import copy
 from enum import auto
 
 from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum, Backend_t
 from cellrank.ul._docs import d
-from cellrank.tl._utils import _connected
-from cellrank.tl.kernels import Kernel
-from cellrank.tl.kernels._base_kernel import _dtype
-from cellrank.tl.kernels._pseudotime_schemes import (
+from cellrank.tl._utils import _connected, _irreducible
+from cellrank.tl.kernels._mixins import ConnectivityMixin
+from cellrank.tl.kernels._base_kernel import BidirectionalKernel
+from cellrank.tl.kernels.utils._pseudotime_scheme import (
     ThresholdSchemeABC,
     HardThresholdScheme,
     SoftThresholdScheme,
@@ -20,30 +19,31 @@
 
 import numpy as np
 
+__all__ = ("PseudotimeKernel",)
 
-class ThresholdScheme(ModeEnum):  # noqa: D101
+
+class ThresholdScheme(ModeEnum):
     SOFT = auto()
     HARD = auto()
 
 
 @d.dedent
-class PseudotimeKernel(Kernel):
+class PseudotimeKernel(ConnectivityMixin, BidirectionalKernel):
     """
     Kernel which computes directed transition probabilities based on a KNN graph and pseudotime.
 
     The KNN graph contains information about the (undirected) connectivities among cells, reflecting their similarity.
     Pseudotime can be used to either remove edges that point against the direction of increasing pseudotime
-    :cite:`setty:19`, or to downweight them :cite:`stassen:21`.
+    :cite:`setty:19` or to downweight them :cite:`stassen:21`.
 
     Parameters
     ----------
     %(adata)s
     %(backward)s
     time_key
-        Key in :attr:`adata` ``.obs`` where the pseudotime is stored.
-    %(cond_num)s
+        Key in :attr:`anndata.AnnData.obs` where the pseudotime is stored.
     kwargs
-        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
+        Keyword arguments for the parent class.
     """
 
     def __init__(
@@ -51,37 +51,34 @@ def __init__(
         adata: AnnData,
         backward: bool = False,
         time_key: str = "dpt_pseudotime",
-        compute_cond_num: bool = False,
-        check_connectivity: bool = False,
         **kwargs: Any,
     ):
         super().__init__(
             adata,
             backward=backward,
             time_key=time_key,
-            compute_cond_num=compute_cond_num,
-            check_connectivity=check_connectivity,
             **kwargs,
         )
-        self._time_key = time_key
 
-    def _read_from_adata(self, time_key: str, **kwargs: Any) -> None:
+    def _read_from_adata(self, time_key: str = "dpt_pseudotime", **kwargs: Any) -> None:
         super()._read_from_adata(**kwargs)
-
+        # fmt: off
+        self._time_key = time_key
         if time_key not in self.adata.obs:
-            raise KeyError(f"Could not find time key in `adata.obs[{time_key!r}]`.")
-
-        self._pseudotime = np.array(self.adata.obs[time_key]).astype(_dtype)
-        if self.backward:
-            self._pseudotime = np.max(self.pseudotime) - self.pseudotime
+            raise KeyError(f"Unable to find pseudotime in `adata.obs[{time_key!r}]`.")
 
-        if np.any(np.isnan(self._pseudotime)):
+        self._pseudotime = np.array(self.adata.obs[time_key]).astype(np.float64, copy=True)
+        if np.any(np.isnan(self.pseudotime)):
             raise ValueError("Encountered NaN values in pseudotime.")
+        # fmt: on
 
     @d.dedent
     def compute_transition_matrix(
         self,
-        threshold_scheme: Union[Literal["soft", "hard"], Callable] = "hard",
+        threshold_scheme: Union[
+            Literal["soft", "hard"],
+            Callable[[float, np.ndarray, np.ndarray], np.ndarray],
+        ] = "hard",
         frac_to_keep: float = 0.3,
         b: float = 10.0,
         nu: float = 0.5,
@@ -94,7 +91,7 @@ def compute_transition_matrix(
         """
         Compute transition matrix based on KNN graph and pseudotemporal ordering.
 
-        Depending on the choice of the `thresholding_scheme`, this is based on ideas by either *Palantir*
+        Depending on the choice of the ``threshold_scheme``, it is based on ideas by either *Palantir*
         :cite:`setty:19` or *VIA* :cite:`stassen:21`.
 
         Parameters
@@ -110,11 +107,11 @@ def compute_transition_matrix(
                   of increasing pseudotime. Essentially, the further "behind" a query cell is in pseudotime with respect
                   to the current reference cell, the more penalized will be its graph-connectivity.
                 - :class:`callable` - any function conforming to the signature of
-                  :func:`cellrank.tl.kernels.ThresholdSchemeABC.__call__`.
+                  :func:`cellrank.kernels.utils.ThresholdSchemeABC.__call__`.
         frac_to_keep
-            The `frac_to_keep` * number of the closest neighbors (according to graph connectivities) are kept, no matter
-            whether they lie in the pseudotemporal past or future. This is done to ensure that the graph remains
-            connected. Only used when ``threshold_scheme = 'hard'``. Needs to fall within the interval `[0, 1]`.
+            Fraction of the closest neighbors (according to graph connectivities) are kept, no matter whether they lie
+            in the pseudotemporal past or future. This is done to ensure that the graph remains connected.
+            Only used when ``threshold_scheme = 'hard'``. Needs to fall within the interval `[0, 1]`.
         %(soft_scheme_kernel)s
         check_irreducibility
             Optional check for irreducibility of the final transition matrix.
@@ -124,9 +121,12 @@ def compute_transition_matrix(
 
         Returns
         -------
-        Self and updated :attr:`transition_matrix`.
+        Self and updates :attr:`transition_matrix` and :attr:`params`.
         """
-        start = logg.info(f"Computing transition matrix based on `{self._time_key}`")
+        if self.pseudotime is None:
+            raise ValueError("Compute pseudotime first.")  # CytoTraceKernel
+
+        start = logg.info("Computing transition matrix based on pseudotime`")
         if isinstance(threshold_scheme, str):
             threshold_scheme = ThresholdScheme(threshold_scheme)
             if threshold_scheme == ThresholdScheme.SOFT:
@@ -166,34 +166,22 @@ def compute_transition_matrix(
         # make sure the biased graph is still connected
         if not _connected(biased_conn):
             logg.warning("Biased KNN graph is disconnected")
+        if check_irreducibility and not _irreducible(biased_conn):
+            logg.warning("Biased KNN graph is not irreducible")
 
-        self._compute_transition_matrix(
-            matrix=biased_conn,
-            density_normalize=False,
-            check_irreducibility=check_irreducibility,
-        )
+        self.transition_matrix = biased_conn
         logg.info("    Finish", time=start)
 
         return self
 
     @property
-    def pseudotime(self) -> np.array:
+    def pseudotime(self) -> Optional[np.array]:
         """Pseudotemporal ordering of cells."""
         return self._pseudotime
 
-    def copy(self) -> "PseudotimeKernel":
-        """Return a copy of self."""
-        pk = PseudotimeKernel(
-            self.adata, backward=self.backward, time_key=self._time_key
-        )
-        pk._pseudotime = copy(self.pseudotime)
-        pk._params = copy(self._params)
-        pk._cond_num = self.condition_number
-        pk._transition_matrix = copy(self._transition_matrix)
-
-        return pk
-
     def __invert__(self) -> "PseudotimeKernel":
-        super().__invert__()
-        self._pseudotime = np.max(self.pseudotime) - self.pseudotime
-        return self
+        pk = self._copy_ignore("_transition_matrix")
+        pk._pseudotime = np.max(pk.pseudotime) - pk.pseudotime
+        pk._backward = not self.backward
+        pk._params = {}
+        return pk
diff --git a/cellrank/tl/kernels/_random_walk.py b/cellrank/tl/kernels/_random_walk.py
deleted file mode 100644
index 865977436..000000000
--- a/cellrank/tl/kernels/_random_walk.py
+++ /dev/null
@@ -1,190 +0,0 @@
-from typing import Any, List, Union, Optional, Sequence
-
-from itertools import chain
-
-from cellrank import logging as logg
-from cellrank.ul._docs import d
-from cellrank.ul._parallelize import parallelize
-
-import numpy as np
-from scipy.sparse import issparse, spmatrix
-
-
-class RandomWalk:
-    """
-    Class that simulates a random walk on a Markov chain.
-
-    Parameters
-    ----------
-    transition_matrix
-        Row-stochastic transition matrix.
-    start_ixs
-        Indices from which to uniformly sample the starting points. If `None`, use all points.
-    stop_ixs
-        Indices which when hit, the random walk is terminated.
-    """
-
-    def __init__(
-        self,
-        transition_matrix: Union[np.ndarray, spmatrix],
-        start_ixs: Optional[Sequence[int]] = None,
-        stop_ixs: Optional[Sequence[int]] = None,
-    ):
-        if transition_matrix.ndim != 2 or (
-            transition_matrix.shape[0] != transition_matrix.shape[1]
-        ):
-            raise ValueError("Array is not a square matrix.")
-        if not np.allclose(transition_matrix.sum(1), 1.0):
-            raise ValueError("Transition matrix is not row-stochastic.")
-
-        self._tmat = transition_matrix
-        self._ixs = np.arange(self._tmat.shape[0])
-        self._is_sparse = issparse(self._tmat)
-        self._stop_ixs = set([] if stop_ixs is None or not len(stop_ixs) else stop_ixs)
-        self._starting_dist = (
-            np.ones_like(self._ixs)
-            if start_ixs is None
-            else np.isin(self._ixs, start_ixs)
-        )
-        if np.sum(self._starting_dist) == 0:
-            raise ValueError("No starting indices have been selected.")
-
-        self._starting_dist = self._starting_dist.astype(np.float64) / np.sum(
-            self._starting_dist
-        )
-
-    def _should_stop(self, ix: int) -> bool:
-        return ix in self._stop_ixs
-
-    def _sample(self, ix: int, *, rs: np.random.RandomState) -> int:
-        return rs.choice(
-            self._ixs,
-            p=self._tmat[ix].A.squeeze() if self._is_sparse else self._tmat[ix],
-        )
-
-    def _max_iter(self, max_iter: Union[int, float]) -> int:
-        if isinstance(max_iter, float):
-            max_iter = int(np.ceil(max_iter * len(self._ixs)))
-        if max_iter <= 1:
-            raise ValueError(
-                f"Expected number of iterations to be > 1, found `{max_iter}`."
-            )
-        return max_iter
-
-    @d.get_sections(base="rw_sim", sections=["Parameters"])
-    def simulate_one(
-        self,
-        max_iter: Union[int, float] = 0.25,
-        seed: Optional[int] = None,
-        successive_hits: int = 0,
-    ) -> np.ndarray:
-        """
-        Simulate one random walk.
-
-        Parameters
-        ----------
-        max_iter
-            Maximum number of steps of a random walk. If a :class:`float`, it can be specified
-            as a fraction of the number of cells.
-        seed
-            Random seed.
-        successive_hits
-            Number of successive hits in the ``stop_ixs`` required to stop prematurely.
-
-        Returns
-        -------
-        Array of shape ``(max_iter + 1,)`` of states that have been visited. If ``stop_ixs`` was specified, the array
-        may have smaller shape.
-        """
-        max_iter = self._max_iter(max_iter)
-        if successive_hits < 0:
-            raise ValueError(
-                f"Expected number of successive hits to be positive, found `{successive_hits}`."
-            )
-
-        rs = np.random.RandomState(seed)
-        ix = rs.choice(self._ixs, p=self._starting_dist)
-        sim, cnt = [ix], -1
-
-        for _ in range(max_iter):
-            ix = self._sample(ix, rs=rs)
-            sim.append(ix)
-            cnt = (cnt + 1) if self._should_stop(ix) else -1
-            if cnt >= successive_hits:
-                break
-
-        return np.array(sim)
-
-    def _simulate_many(
-        self,
-        sims: np.ndarray,
-        max_iter: Union[int, float] = 0.25,
-        seed: Optional[int] = None,
-        successive_hits: int = 0,
-        queue: Optional[Any] = None,
-    ) -> List[np.ndarray]:
-        res = []
-        for s in sims:
-            sim = self.simulate_one(
-                max_iter=max_iter,
-                seed=None if seed is None else seed + s,
-                successive_hits=successive_hits,
-            )
-            res.append(sim)
-            if queue is not None:
-                queue.put(1)
-
-        if queue is not None:
-            queue.put(None)
-
-        return res
-
-    @d.dedent
-    def simulate_many(
-        self,
-        n_sims: int,
-        max_iter: Union[int, float] = 0.25,
-        seed: Optional[int] = None,
-        successive_hits: int = 0,
-        n_jobs: Optional[int] = None,
-        backend: str = "loky",
-        show_progress_bar: bool = True,
-    ) -> List[np.ndarray]:
-        """
-        Simulate many random walks.
-
-        Parameters
-        ----------
-        n_sims
-            Number of random walks to simulate.
-        %(rw_sim.params)s
-        %(parallel)s
-
-        Returns
-        -------
-        List of arrays of shape ``(max_iter + 1,)`` of states that have been visited. If ``stop_ixs`` was specified,
-        the arrays may have smaller shape.
-        """
-        if n_sims <= 0:
-            raise ValueError(
-                f"Expected number of simulations to be positive, found `{n_sims}`."
-            )
-        max_iter = self._max_iter(max_iter)
-        start = logg.info(
-            f"Simulating `{n_sims}` random walks of maximum length `{max_iter}`"
-        )
-
-        simss = parallelize(
-            self._simulate_many,
-            collection=np.arange(n_sims),
-            n_jobs=n_jobs,
-            backend=backend,
-            show_progress_bar=show_progress_bar,
-            as_array=False,
-            unit="sim",
-        )(max_iter=max_iter, seed=seed, successive_hits=successive_hits)
-        simss = list(chain.from_iterable(simss))
-
-        logg.info("    Finish", time=start)
-
-        return simss
diff --git a/cellrank/tl/kernels/_transport_map_kernel.py b/cellrank/tl/kernels/_transport_map_kernel.py
index cd1a622b3..461acc2f2 100644
--- a/cellrank/tl/kernels/_transport_map_kernel.py
+++ b/cellrank/tl/kernels/_transport_map_kernel.py
@@ -3,30 +3,42 @@
 
 from abc import ABC, abstractmethod
 from types import MappingProxyType
+from contextlib import contextmanager
 
 import scanpy as sc
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d, inject_docs
-from cellrank.tl.kernels import ExperimentalTimeKernel
-from cellrank.tl.kernels._base_kernel import KernelExpression
-from cellrank.tl.kernels._exp_time_kernel import LastTimePoint
+from cellrank.tl.kernels._experimental_time_kernel import ExperimentalTimeKernel
 
 import numpy as np
 import pandas as pd
 from scipy.sparse import bmat, spdiags, spmatrix
 
+__all__ = ("TransportMapKernel",)
+
+
+class LastTimePoint(ModeEnum):
+    UNIFORM = "uniform"
+    DIAGONAL = "diagonal"
+    CONNECTIVITIES = "connectivities"
+
+
+Numeric_t = Union[int, float]
+Pair_t = Tuple[Numeric_t, Numeric_t]
+
 
 class TransportMapKernel(ExperimentalTimeKernel, ABC):
-    """Kernel base class which computes transition matrix based on transport maps for consecutive time pairs."""
+    """Kernel base class which computes transition matrix based on transport maps for consecutive time point pairs."""
 
     def __init__(self, *args: Any, **kwargs: Any):
         super().__init__(*args, **kwargs)
-        self._tmaps: Optional[Dict[Tuple[Any, Any], AnnData]] = None
+        self._tmaps: Optional[Dict[Pair_t, AnnData]] = None
 
     @abstractmethod
     def _compute_tmap(
-        self, t1: Any, t2: Any, **kwargs: Any
+        self, t1: Numeric_t, t2: Numeric_t, **kwargs: Any
     ) -> Union[np.ndarray, spmatrix, AnnData]:
         """
         Compute transport matrix for a time point pair.
@@ -46,57 +58,17 @@ def _compute_tmap(
         and :attr:`anndata.AnnData.var_names` correspond to subsets of observations from :attr:`adata`.
         """
 
-    @d.get_sections(base="tmk_thresh", sections=["Parameters"])
-    def _threshold_transition_matrix(
-        self, threshold: Union[float, Literal["auto"]]
-    ) -> None:
-        """
-        Remove small non-zero values from :attr:`transition_matrix`.
-
-        Parameters
-        ----------
-        threshold
-            How to remove small non-zero values from the transition matrix. Valid options are:
-
-                - `'auto'` - find the maximum threshold value which will not remove every non-zero value from any row.
-                - :class:`float` - value in `[0, 100]` corresponding to a percentage of non-zeros to remove.
-                  Rows where all values are removed will have uniform distribution.
-                - `None` - do not threshold.
-
-        Returns
-        -------
-        Nothing, just updates :attr:`transition_matrix`.
-        """
-        tmat = self.transition_matrix
-        if threshold == "auto":
-            threshold = min(np.max(tmat[i].data) for i in range(tmat.shape[0]))
-            logg.info(f"Using `threshold={threshold}`")
-            tmat.data[tmat.data < threshold] = 0.0
-        else:
-            if not (0 <= threshold <= 100):
-                raise ValueError(
-                    f"Expected `threshold to be in `[0, 100]`, found `{threshold}`.`"
-                )
-            threshold = np.percentile(tmat.data, threshold)
-            logg.info(f"Using `threshold={threshold}`")
-            tmat.data[tmat.data <= threshold] = 0.0
-
-        tmat.eliminate_zeros()
-        self._compute_transition_matrix(
-            matrix=tmat,
-            density_normalize=False,
-            check_irreducibility=False,
-        )
-
     @d.get_sections(base="tmk_tmat", sections=["Parameters"])
     @d.dedent
     def compute_transition_matrix(
         self,
         threshold: Optional[Union[float, Literal["auto"]]] = "auto",
-        last_time_point: LastTimePoint = LastTimePoint.DIAGONAL,
+        last_time_point: Literal[
+            "uniform", "diagonal", "connectivities"
+        ] = LastTimePoint.DIAGONAL,
         conn_kwargs: Mapping[str, Any] = MappingProxyType({}),
         **kwargs: Any,
-    ) -> KernelExpression:
+    ) -> "TransportMapKernel":
         """
         Compute transition matrix using transport maps.
 
@@ -112,11 +84,11 @@ def compute_transition_matrix(
                   derived from the last time point subset of :attr:`adata`.
         conn_kwargs
             Keyword arguments for :func:`scanpy.pp.neighbors` when using ``last_time_point = {ltp.CONNECTIVITIES!r}``.
-            Can have `'density_normalize'` for :meth:`cellrank.tl.kernels.ConnectivityKernel.compute_transition_matrix`.
+            Can have `'density_normalize'` for :meth:`cellrank.kernels.ConnectivityKernel.compute_transition_matrix`.
 
         Returns
         -------
-        Self and updated :attr:`transition_matrix`.
+        Self and updates :attr:`transition_matrix` and :attr:`params`.
         """
         ckwargs = dict(kwargs)
         ckwargs["threshold"] = threshold
@@ -127,19 +99,16 @@ def compute_transition_matrix(
         timepoints = self.experimental_time.cat.categories
         timepoints = list(zip(timepoints[:-1], timepoints[1:]))
 
-        # fmt: off
+        self._tmaps = {
+            (t1, t2): self._tmat_to_adata(t1, t2, self._compute_tmap(t1, t2, **kwargs))
+            for t1, t2 in timepoints
+        }
         tmap = self._restich_tmaps(
-            {(t1, t2): self._tmat_to_adata(t1, t2, self._compute_tmap(t1, t2, **kwargs)) for t1, t2 in timepoints},
+            self._tmaps,
             last_time_point=last_time_point,
             conn_kwargs=conn_kwargs,
         )
-        # fmt: on
-
-        self._compute_transition_matrix(
-            matrix=tmap.X,
-            density_normalize=False,
-            check_irreducibility=False,
-        )
+        self.transition_matrix = tmap.X
         if threshold:
             self._threshold_transition_matrix(threshold)
 
@@ -149,7 +118,7 @@ def compute_transition_matrix(
     @inject_docs(ltp=LastTimePoint)
     def _restich_tmaps(
         self,
-        tmaps: Mapping[Tuple[Any, Any], AnnData],
+        tmaps: Mapping[Pair_t, AnnData],
         last_time_point: LastTimePoint = LastTimePoint.DIAGONAL,
         conn_kwargs: Mapping[str, Any] = MappingProxyType({}),
     ) -> AnnData:
@@ -164,15 +133,17 @@ def _restich_tmaps(
 
         Returns
         -------
-        Merged transport maps into 1 :class:`anndata.AnnData` object.
+        Merged transport maps into one :class:`anndata.AnnData` object.
         """
         from cellrank.tl.kernels import ConnectivityKernel
 
+        last_time_point = LastTimePoint(last_time_point)
         tmaps = self._validate_tmaps(tmaps)
 
         conn_kwargs = dict(conn_kwargs)
         conn_kwargs["copy"] = False
         _ = conn_kwargs.pop("key_added", None)
+        # same default as in `ConnectivityKernel`
         density_normalize = conn_kwargs.pop("density_normalize", True)
 
         blocks = [[None] * (len(tmaps) + 1) for _ in range(len(tmaps) + 1)]
@@ -227,9 +198,9 @@ def _restich_tmaps(
 
     def _validate_tmaps(
         self,
-        tmaps: Mapping[Tuple[Any, Any], AnnData],
+        tmaps: Dict[Pair_t, AnnData],
         allow_reorder: bool = True,
-    ) -> Mapping[Tuple[Any, Any], AnnData]:
+    ) -> Mapping[Pair_t, AnnData]:
         """
         Validate that transport maps conform to various invariants.
 
@@ -293,15 +264,68 @@ def _validate_tmaps(
             )
         except AssertionError as e:
             raise KeyError(
-                "Observations from transport maps don't match"
+                "Observations from transport maps don't match "
                 "the observations from the underlying `AnnData` object."
             ) from e
 
         return tmaps
 
+    @d.get_sections(base="tmk_thresh", sections=["Parameters"])
+    def _threshold_transition_matrix(
+        self,
+        threshold: Union[float, Literal["auto"]],
+        normalize: bool = True,
+    ) -> None:
+        """
+        Remove small non-zero values from :attr:`transition_matrix`.
+
+        Parameters
+        ----------
+        threshold
+            How to remove small non-zero values from the transition matrix. Valid options are:
+
+                - `'auto'` - find the maximum threshold value which will not remove every non-zero value from any row.
+                - :class:`float` - value in `[0, 100]` corresponding to a percentage of non-zeros to remove.
+
+            Rows where all values are removed will have uniform distribution and a warning will be issued.
+
+        Returns
+        -------
+        Nothing, just updates :attr:`transition_matrix`.
+        """
+        tmat = self.transition_matrix
+        if threshold == "auto":
+            threshold = min(np.max(tmat[i].data) for i in range(tmat.shape[0]))
+        else:
+            if not (0 <= threshold <= 100):
+                raise ValueError(
+                    f"Expected `threshold` to be in `[0, 100]`, found `{threshold}`.`"
+                )
+            threshold = np.percentile(tmat.data, threshold)
+        logg.info(f"Using `threshold={threshold}`")
+
+        tmat.data[tmat.data < threshold] = 0.0
+        zeros_mask = np.where(np.asarray(tmat.sum(1)).squeeze() == 0)[0]
+        if len(zeros_mask):
+            logg.warning(
+                f"After thresholding, `{len(zeros_mask)}` row(s) are forced to be uniform"
+            )
+            for ix in zeros_mask:
+                start, end = tmat.indptr[ix], tmat.indptr[ix + 1]
+                size = end - start
+                tmat.data[start:end] = np.ones(size, dtype=tmat.dtype) / size
+        tmat.eliminate_zeros()
+
+        if normalize:
+            self.transition_matrix = tmat
+        else:
+            # e.g. when in `_tmap_as_tmat`
+            self._transition_matrix = tmat
+
     def _tmat_to_adata(
-        self, t1: Any, t2: Any, tmat: Union[np.ndarray, spmatrix, AnnData]
+        self, t1: Numeric_t, t2: Numeric_t, tmat: Union[np.ndarray, spmatrix, AnnData]
     ) -> AnnData:
+        """Convert transport map ``tmat`` to :class:`anndata.AnnData`."""
         if isinstance(tmat, AnnData):
             return tmat
 
@@ -311,7 +335,63 @@ def _tmat_to_adata(
 
         return tmat
 
+    @contextmanager
+    def _tmap_as_tmat(
+        self, threshold: Optional[Union[float, Literal["auto"]]] = None, **kwargs: Any
+    ) -> None:
+        """Temporarily set :attr:`transport_maps` as :attr:`transition_matrix`."""
+        if self.transport_maps is None:
+            raise RuntimeError(
+                "Compute transport maps first as `.compute_transition_matrix()`."
+            )
+
+        tmat = self._transition_matrix
+        try:
+            self._transition_matrix = self._restich_tmaps(
+                self.transport_maps, normalize=False, **kwargs
+            ).X
+            if threshold is not None:
+                self._threshold_transition_matrix(threshold, normalize=False)
+            yield
+        finally:
+            self._transition_matrix = tmat
+
+    @d.dedent
+    def plot_single_flow(
+        self,
+        cluster: str,
+        cluster_key: str,
+        time_key: Optional[str] = None,
+        use_transport_maps: bool = False,
+        threshold: Optional[Union[float, Literal["auto"]]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        %(plot_single_flow.full_desc)s
+
+        Parameters
+        ----------
+        %(plot_single_flow.parameters)s
+
+        Returns
+        -------
+        %(plot_single_flow.returns)s
+        """  # noqa: D400
+        if use_transport_maps:
+            with self._tmap_as_tmat(threshold):
+                return super().plot_single_flow(
+                    cluster, cluster_key, time_key=time_key, **kwargs
+                )
+        return super().plot_single_flow(
+            cluster, cluster_key, time_key=time_key, **kwargs
+        )
+
     @property
-    def transport_maps(self) -> Optional[Dict[Tuple[Any, Any], AnnData]]:
+    def transport_maps(self) -> Optional[Dict[Pair_t, AnnData]]:
         """Transport maps for consecutive time pairs."""
         return self._tmaps
+
+    def __invert__(self) -> "TransportMapKernel":
+        tk = super().__invert__("_tmaps")  # don't copy transport maps
+        tk._tmaps = None
+        return tk
diff --git a/cellrank/tl/kernels/_utils.py b/cellrank/tl/kernels/_utils.py
index 6e74970e2..40e0d9aec 100644
--- a/cellrank/tl/kernels/_utils.py
+++ b/cellrank/tl/kernels/_utils.py
@@ -1,5 +1,6 @@
-from typing import Tuple, Callable, Optional
+from typing import Any, Tuple, Callable, Optional
 
+import wrapt
 from inspect import signature
 
 from anndata import AnnData
@@ -219,7 +220,7 @@ def _reconstruct_one(
 
 
 @njit(**jit_kwargs)
-def _calculate_starts(indptr: np.ndarray, ixs) -> np.ndarray:
+def _calculate_starts(indptr: np.ndarray, ixs: np.ndarray) -> np.ndarray:
     """
     Get the position where to put the data.
 
@@ -295,3 +296,19 @@ def _ensure_numeric_ordered(adata: AnnData, key: str) -> pd.Series:
         raise ValueError(f"Expected to find at least `2` categories, found `{n_cats}`.")
 
     return exp_time
+
+
+@wrapt.decorator
+def require_tmat(
+    wrapped: Callable[..., Any],
+    instance: "KernelExpression",  # noqa: F821
+    args: Any,
+    kwargs: Any,
+) -> Any:
+    """Require that the transition matrix is computed before calling the wrapped function."""
+    # this can trigger combinations, but not individual kernels
+    if instance.transition_matrix is None:
+        raise RuntimeError(
+            "Compute transition matrix first as `.compute_transition_matrix()`."
+        )
+    return wrapped(*args, **kwargs)
diff --git a/cellrank/tl/kernels/_velocity_kernel.py b/cellrank/tl/kernels/_velocity_kernel.py
index 4aa952a95..9c4ff1ec7 100644
--- a/cellrank/tl/kernels/_velocity_kernel.py
+++ b/cellrank/tl/kernels/_velocity_kernel.py
@@ -1,48 +1,27 @@
-from typing import Any, Union, Callable, Iterable, Optional
+from typing import Any, Tuple, Union, Callable, Optional, Sequence
 from typing_extensions import Literal
 
-from copy import copy
-from enum import auto
-from math import fsum
+from abc import ABC
 
 from anndata import AnnData
 from cellrank import logging as logg
-from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum
+from cellrank.tl._enum import _DEFAULT_BACKEND, Backend_t
 from cellrank.ul._docs import d, inject_docs
-from cellrank.ul._utils import valuedispatch
-from cellrank.tl.kernels import Kernel
-from cellrank.ul._parallelize import parallelize
-from cellrank.tl.kernels._utils import (
-    prange,
-    np_mean,
-    _filter_kwargs,
-    _random_normal,
-    _reconstruct_one,
-    _calculate_starts,
-    _get_probs_for_zero_vec,
-)
+from cellrank.tl.kernels.utils import MonteCarlo, Stochastic, Deterministic
+from cellrank.tl.kernels._mixins import ConnectivityMixin
 from scvelo.preprocessing.moments import get_moments
-from cellrank.tl.kernels._base_kernel import _RTOL
-from cellrank.tl.kernels._velocity_schemes import Scheme, _get_scheme
+from cellrank.tl.kernels._base_kernel import BidirectionalKernel
+from cellrank.tl.kernels.utils._similarity import Similarity, SimilarityABC
+from cellrank.tl.kernels.utils._velocity_model import BackwardMode, VelocityModel
 
 import numpy as np
-from scipy.sparse import issparse, csr_matrix
+from scipy.sparse import issparse
 
-
-class VelocityMode(ModeEnum):  # noqa: D101
-    DETERMINISTIC = auto()
-    STOCHASTIC = auto()
-    SAMPLING = auto()
-    MONTE_CARLO = auto()
-
-
-class BackwardMode(ModeEnum):  # noqa: D101
-    TRANSPOSE = auto()
-    NEGATE = auto()
+__all__ = ("VelocityKernel",)
 
 
 @d.dedent
-class VelocityKernel(Kernel):
+class VelocityKernel(ConnectivityMixin, BidirectionalKernel, ABC):
     """
     Kernel which computes a transition matrix based on RNA velocity.
 
@@ -56,113 +35,74 @@ class VelocityKernel(Kernel):
     ----------
     %(adata)s
     %(backward)s
-    vkey
-        Key in :attr:`anndata.AnnData.layers` where velocities are stored.
     xkey
         Key in :attr:`anndata.AnnData.layers` where expected gene expression counts are stored.
+    vkey
+        Key in :attr:`anndata.AnnData.layers` where velocities are stored.
     gene_subset
         List of genes to be used to compute transition probabilities.
-        By default, genes from :attr:`anndata.AnnData.var` ``['velocity_genes']`` are used.
-    %(cond_num)s
-    check_connectivity
-        Check whether the underlying KNN graph is connected.
+        If not specified, genes from :attr:`anndata.AnnData.var` ``['{vkey}_genes']`` are used.
     kwargs
-        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
+        Keyword arguments for the parent class.
     """
 
     def __init__(
         self,
         adata: AnnData,
         backward: bool = False,
-        vkey: str = "velocity",
-        xkey: str = "Ms",
-        gene_subset: Optional[Iterable] = None,
-        compute_cond_num: bool = False,
-        check_connectivity: bool = False,
+        xkey: Optional[str] = "Ms",
+        vkey: Optional[str] = "velocity",
         **kwargs: Any,
     ):
         super().__init__(
             adata,
             backward=backward,
-            vkey=vkey,
             xkey=xkey,
-            gene_subset=gene_subset,
-            compute_cond_num=compute_cond_num,
-            check_connectivity=check_connectivity,
+            vkey=vkey,
             **kwargs,
         )
-        self._vkey = vkey  # for copy
-        self._xkey = xkey
-        self._gene_subset = gene_subset
-        self._logits = None
+        self._logits: Optional[np.ndarray] = None
 
-    def _read_from_adata(self, **kwargs: Any) -> None:
+    def _read_from_adata(
+        self,
+        xkey: Optional[str] = "Ms",
+        vkey: Optional[str] = "velocity",
+        gene_subset: Optional[Union[str, Sequence[str]]] = None,
+        **kwargs: Any,
+    ) -> None:
         super()._read_from_adata(**kwargs)
 
-        # check whether velocities have been computed
-        vkey = kwargs.pop("vkey", "velocity")
-        if vkey not in self.adata.layers.keys():
-            raise KeyError("Compute RNA velocity first as `scvelo.tl.velocity()`.")
-
-        # restrict genes to a subset, i.e. velocity genes or user provided list
-        gene_subset = kwargs.pop("gene_subset", None)
-        subset = np.ones(self.adata.n_vars, bool)
-        if gene_subset is not None:
-            var_names_subset = self.adata.var_names.isin(gene_subset)
-            subset &= var_names_subset if len(var_names_subset) > 0 else gene_subset
-        elif f"{vkey}_genes" in self.adata.var.keys():
-            subset &= np.array(self.adata.var[f"{vkey}_genes"].values, dtype=bool)
-
-        # choose data representation to use for transcriptomic displacements
-        xkey = kwargs.pop("xkey", "Ms")
-        xkey = xkey if xkey in self.adata.layers.keys() else "spliced"
-
-        # filter both the velocities and the gene expression profiles to the gene subset and densify the matrices
-        X = np.array(
-            self.adata.layers[xkey].A[:, subset]
-            if issparse(self.adata.layers[xkey])
-            else self.adata.layers[xkey][:, subset]
-        )
-        V = np.array(
-            self.adata.layers[vkey].A[:, subset]
-            if issparse(self.adata.layers[vkey])
-            else self.adata.layers[vkey][:, subset]
-        )
+        if gene_subset is None and f"{vkey}_genes" in self.adata.var:
+            gene_subset = self.adata.var[f"{vkey}_genes"]
 
-        # remove genes that have any Nan values (in both X and V)
-        nans = np.isnan(np.sum(V, axis=0))
-        if np.any(nans):
-            X = X[:, ~nans]
-            V = V[:, ~nans]
-
-        # check the velocity parameters
-        par_key = f"{vkey}_params"
-        velocity_params = self.adata.uns.get(par_key, None)
-        if velocity_params is None:
-            logg.debug(
-                f"Unable to load velocity parameters from `adata.uns[{par_key!r}]`"
-            )
+        self._xdata = self._extract_layer(xkey, subset=gene_subset)
+        self._vdata = self._extract_layer(vkey, subset=gene_subset)
 
-        # add to self
-        self._velocity = V.astype(np.float64)
-        self._gene_expression = X.astype(np.float64)
-        self._velocity_params = velocity_params
+        nans = np.isnan(np.sum(self._vdata, axis=0))
+        if np.any(nans):
+            self._xdata = self._xdata[:, ~nans]
+            self._vdata = self._vdata[:, ~nans]
+        # fmt: off
+        self._vexp = get_moments(self.adata, self._vdata, second_order=False).astype(np.float64, copy=False)
+        self._vvar = get_moments(self.adata, self._vdata, second_order=True).astype(np.float64, copy=False)
+        # fmt: on
 
-    @inject_docs(m=VelocityMode, b=BackwardMode, s=Scheme)  # don't swap the order
+    # TODO(michalk8): remove the docrep in 2.0
+    @inject_docs(m=VelocityModel, b=BackwardMode, s=Similarity)  # don't swap the order
     @d.dedent
     def compute_transition_matrix(
         self,
-        mode: Literal[
-            "deterministic", "stochastic", "sampling", "monte_carlo"
-        ] = VelocityMode.DETERMINISTIC,
+        model: Literal[
+            "deterministic", "stochastic", "monte_carlo"
+        ] = VelocityModel.DETERMINISTIC,
         backward_mode: Literal["transpose", "negate"] = BackwardMode.TRANSPOSE,
-        scheme: Union[
-            Literal["dot_product", "cosine", "correlation"], Callable
-        ] = Scheme.CORRELATION,
+        similarity: Union[
+            Literal["correlation", "cosine", "dot_product"],
+            Callable[[np.ndarray, np.ndarray, float], Tuple[np.ndarray, np.ndarray]],
+        ] = Similarity.CORRELATION,
         softmax_scale: Optional[float] = None,
         n_samples: int = 1000,
         seed: Optional[int] = None,
-        check_irreducibility: bool = False,
         **kwargs: Any,
     ) -> "VelocityKernel":
         """
@@ -178,389 +118,162 @@ def compute_transition_matrix(
         %(softmax_scale)s
         %(velocity_scheme)s
         n_samples
-            Number of bootstrap samples when ``mode = {m.MONTE_CARLO!r}``.
+            Number of samples when ``mode = {m.MONTE_CARLO!r}``.
         seed
-            Set the seed for random state when the method requires ``n_samples``.
-        check_irreducibility
-            Optional check for irreducibility of the final transition matrix.
+            Random seed when ``mode = {m.MONTE_CARLO!r}``.
         %(parallel)s
+        kwargs
+            Keyword arguments for the underlying ``model``.
 
         Returns
         -------
-        Self and updates the following fields:
-
-            - :attr:`transition_matrix`.
-            - :attr:`logits`.
+        Self and updates :attr:`transition_matrix`, :attr:`logits` and :attr:`params`.
         """
-        mode = VelocityMode(mode)
-        backward_mode = BackwardMode(backward_mode)
-
-        if isinstance(scheme, str):
-            scheme = _get_scheme(Scheme(scheme))
-
-        if not callable(scheme):
-            raise TypeError(
-                f"Expected `scheme` to be a function, found `{type(scheme)!r}`."
-            )
-
-        if self.backward and mode != VelocityMode.DETERMINISTIC:
-            logg.warning(
-                f"Mode `{mode!r}` is currently not supported for the backward process. "
-                f"Defaulting to mode `{VelocityMode.DETERMINISTIC!r}`"
-            )
-            mode = VelocityMode.DETERMINISTIC
-
-        if mode == VelocityMode.STOCHASTIC:
-            try:
-                if not hasattr(
-                    scheme, "hessian"
-                ):  # this can also raise in our velocity scheme definition
-                    raise NotImplementedError
-            except NotImplementedError:
-                logg.warning(
-                    f"Unable to detect a method for Hessian computation. If using predefined functions, consider "
-                    f"installing `jax` as `pip install jax jaxlib`\n"
-                    f"Defaulting to `mode={VelocityMode.MONTE_CARLO!r}` and `n_samples={n_samples}`"
-                )
-                mode = VelocityMode.MONTE_CARLO
-
-        start = logg.info(
-            f"Computing transition matrix based on logits using `{mode!r}` mode"
-        )
-
-        if seed is None:
-            seed = np.random.randint(0, 2 ** 16)
+        start = logg.info(f"Computing transition matrix using `{model!r}` model")
 
         # fmt: off
-        params = {"softmax_scale": softmax_scale, "mode": mode, "seed": seed, "scheme": str(scheme)}
+        params = {"model": model, "similarity": str(similarity), "softmax_scale": softmax_scale}
         if self.backward:
-            params["bwd_mode"] = backward_mode
+            params["bwd_mode"] = str(backward_mode)
+        if VelocityModel(model) == VelocityModel.MONTE_CARLO:
+            params["n_samples"] = n_samples
+            params["seed"] = seed
         if self._reuse_cache(params, time=start):
             return self
 
-        # compute first and second order moments to model the distribution of the velocity vector
-        np.random.seed(seed)
-        velocity_expectation = get_moments(self.adata, self._velocity, second_order=False).astype(np.float64)
-        velocity_variance = get_moments(self.adata, self._velocity, second_order=True).astype(np.float64)
-        # fmt: on
-
-        if mode == VelocityMode.MONTE_CARLO and n_samples == 1:
-            logg.debug("Setting mode to sampling because `n_samples=1`")
-            mode = VelocityMode.SAMPLING
-
-        backend = kwargs.pop("backend", _DEFAULT_BACKEND)
-        if mode != VelocityMode.STOCHASTIC and backend == "multiprocessing":
-            # this is because on jitting and pickling (cloudpickle, used by loky, handles it correctly)
-            logg.warning(
-                f"Multiprocessing backend is supported only for mode `{VelocityMode.STOCHASTIC!r}`. "
-                f"Defaulting to `{_DEFAULT_BACKEND!r}`"
-            )
-            backend = _DEFAULT_BACKEND
-
         if softmax_scale is None:
-            logg.info(
-                f"Estimating `softmax_scale` using `{VelocityMode.DETERMINISTIC!r}` mode"
-            )
-            _, cmat = _dispatch_computation(
-                VelocityMode.DETERMINISTIC,
-                scheme=scheme,
-                conn=self._conn,
-                expression=self._gene_expression,
-                velocity=self._velocity,
-                expectation=velocity_expectation,
-                variance=velocity_variance,
-                softmax_scale=1.0,
-                backward=self.backward,
-                backward_mode=backward_mode,
-                n_samples=n_samples,
-                seed=seed,
-                backend=backend,
-                **kwargs,
-            )
-            softmax_scale = 1.0 / np.median(np.abs(cmat.data))
+            softmax_scale = self._estimate_softmax_scale(backward_mode=backward_mode, similarity=similarity)
+            logg.info(f"Using `softmax_scale={softmax_scale:.4f}`")
             params["softmax_scale"] = softmax_scale
-            logg.info(f"Setting `softmax_scale={softmax_scale:.4f}`")
-
-        tmat, cmat = _dispatch_computation(
-            mode,
-            scheme=scheme,
-            conn=self._conn,
-            expression=self._gene_expression,
-            velocity=self._velocity,
-            expectation=velocity_expectation,
-            variance=velocity_variance,
-            softmax_scale=softmax_scale,
-            backward=self.backward,
+        # fmt: on
+
+        model = self._create_model(
+            model,
             backward_mode=backward_mode,
+            similarity=similarity,
+            softmax_scale=softmax_scale,
             n_samples=n_samples,
             seed=seed,
-            backend=backend,
-            **kwargs,
-        )
-        self._compute_transition_matrix(
-            tmat, density_normalize=False, check_irreducibility=check_irreducibility
         )
-        self._logits = cmat
+        if isinstance(model, Stochastic):
+            kwargs["backend"] = _DEFAULT_BACKEND
+        self.transition_matrix, self._logits = model(**kwargs)
 
         logg.info("    Finish", time=start)
 
         return self
 
-    @property
-    def logits(self) -> csr_matrix:
-        """Array of shape ``(n_cells, n_cells)`` containing the logits."""
-        return self._logits
-
-    @d.dedent
-    def copy(self) -> "VelocityKernel":
-        """%(copy)s"""  # noqa
-        vk = VelocityKernel(
-            self.adata,
-            backward=self.backward,
-            vkey=self._vkey,
-            xkey=self._xkey,
-            gene_subset=self._gene_subset,
-        )
-        vk._params = copy(self.params)
-        vk._cond_num = self.condition_number
-        vk._transition_matrix = copy(self._transition_matrix)
-        vk._logits = copy(self.logits)
-
-        return vk
-
-
-@valuedispatch
-def _dispatch_computation(mode, *_args, **_kwargs):
-    raise NotImplementedError(mode)
-
-
-def _run_in_parallel(fn: Callable, conn: csr_matrix, **kwargs) -> Any:
-    fname = fn.__name__
-    if fname == "_run_stochastic":
-        ixs = np.argsort(np.array((conn != 0).sum(1)).ravel())[::-1]
-    else:
-        ixs = np.arange(conn.shape[0])
-        np.random.shuffle(ixs)
-
-    unit = (
-        "sample" if (fname == "_run_mc") and kwargs.get("n_samples", 1) > 1 else "cell"
-    )
-    kwargs["indices"] = conn.indices
-    kwargs["indptr"] = conn.indptr
-
-    return parallelize(
-        fn,
-        ixs,
-        as_array=False,
-        extractor=lambda res: _reconstruct_one(np.concatenate(res, axis=-1), conn, ixs),
-        unit=unit,
-        **_filter_kwargs(parallelize, **kwargs),
-    )(**_filter_kwargs(fn, **kwargs))
-
-
-def _run_deterministic(
-    ixs: Union[prange, np.ndarray],
-    scheme: Callable,
-    indices: np.ndarray,
-    indptr: np.ndarray,
-    expression: np.ndarray,
-    velocity: np.ndarray,
-    backward: bool = False,
-    backward_mode: BackwardMode = BackwardMode.TRANSPOSE,
-    softmax_scale: float = 1,
-    queue=None,
-) -> np.ndarray:
-
-    starts = _calculate_starts(indptr, ixs)
-    probs_cors = np.empty((2, starts[-1]))
-
-    for i in prange(len(ixs)):
-        ix = ixs[i]
-        start, end = indptr[ix], indptr[ix + 1]
-
-        nbhs_ixs = indices[start:end]
-        n_neigh = len(nbhs_ixs)
-        assert n_neigh, f"Cell {ix} does not have any neighbors."
-
-        W = expression[nbhs_ixs, :] - expression[ix, :]
-
-        if not backward or (backward and backward_mode == BackwardMode.NEGATE):
-            v = np.expand_dims(velocity[ix], 0)
-
-            # for the negate backward mode
-            if backward:
-                v *= -1
-
-            if np.all(v == 0):
-                p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
-            else:
-                p, c = scheme(v, W, softmax_scale)
-        else:
-            # compute how likely all neighbors are to transition to this cell
-            V = velocity[nbhs_ixs, :]
-            p, c = scheme(V, -1 * W, softmax_scale)
-
-        probs_cors[0, starts[i] : starts[i] + n_neigh] = p
-        probs_cors[1, starts[i] : starts[i] + n_neigh] = c
-
-        if queue is not None:
-            queue.put(1)
-
-    if queue is not None:
-        queue.put(None)
-
-    return probs_cors
-
-
-def _run_mc(
-    ixs: Optional[np.ndarray],
-    scheme: Callable,
-    indices: np.ndarray,
-    indptr: np.ndarray,
-    expression: np.ndarray,
-    expectation: np.ndarray,
-    variance: np.ndarray,
-    n_samples: int = 1,
-    softmax_scale: float = 1,
-    seed: int = 0,
-    queue=None,
-) -> np.ndarray:
-
-    starts = _calculate_starts(indptr, ixs)
-    probs_cors = np.empty((2, starts[-1]))
-
-    for i in prange(len(ixs)):
-        ix = ixs[i]
-        start, end = indptr[ix], indptr[ix + 1]
+    def _create_model(
+        self,
+        model: Union[str, VelocityModel],
+        backward_mode: Literal["negate", "transpose"],
+        similarity: Union[str, Similarity, Callable],
+        n_samples: int = 1000,
+        seed: Optional[int] = None,
+        **kwargs: Any,
+    ):
+        model = VelocityModel(model)
+        backward_mode = BackwardMode(backward_mode) if self.backward else None
 
-        np.random.seed(seed + ix)
+        if isinstance(similarity, str):
+            similarity = SimilarityABC.create(Similarity(similarity))
+        elif not callable(similarity):
+            raise TypeError(
+                f"Expected `scheme` to be a function, found `{type(similarity).__name__}`."
+            )
 
-        nbhs_ixs = indices[start:end]
-        n_neigh = len(nbhs_ixs)
-        assert n_neigh, f"Cell {ix} does not have any neighbors."
+        if self.backward and model != VelocityModel.DETERMINISTIC:
+            logg.warning(
+                f"Mode `{model!r}` is currently not supported for backward process. "
+                f"Using to `mode={VelocityModel.DETERMINISTIC!r}`"
+            )
+            model = VelocityModel.DETERMINISTIC
 
-        # get the displacement matrix. Changing dimensions b/c varying numbers of neighbors slow down autograd
-        W = expression[nbhs_ixs, :] - expression[ix, :]
+        if model == VelocityModel.STOCHASTIC and not hasattr(similarity, "hessian"):
+            model = VelocityModel.MONTE_CARLO
+            logg.warning(
+                f"Unable to detect a method for Hessian computation. If using one of the "
+                f"predefined similarity functions, consider installing `jax` as "
+                f"`pip install jax jaxlib`. Using `mode={model!r}` and `n_samples={n_samples}`"
+            )
 
-        # much faster (1.8x than sampling only 1 if average)
-        samples = _random_normal(
-            expectation[ix],
-            variance[ix],
-            n_samples=n_samples,
-        )
+        if model == VelocityModel.DETERMINISTIC:
+            return Deterministic(
+                self._conn,
+                self._xdata,
+                self._vdata,
+                similarity=similarity,
+                backward_mode=backward_mode,
+                **kwargs,
+            )
+        if model == VelocityModel.STOCHASTIC:
+            return Stochastic(
+                self._conn,
+                self._xdata,
+                self._vexp,
+                self._vvar,
+                similarity=similarity,
+                backward_mode=backward_mode,
+                **kwargs,
+            )
+        if model == VelocityModel.MONTE_CARLO:
+            return MonteCarlo(
+                self._conn,
+                self._xdata,
+                self._vexp,
+                self._vvar,
+                similarity=similarity,
+                backward_mode=backward_mode,
+                n_samples=n_samples,
+                seed=seed,
+                **kwargs,
+            )
 
-        probs_cors_tmp = np.empty((2, n_samples, n_neigh))
-        for j in prange(n_samples):
-            prob, cor = scheme(np.atleast_2d(samples[j]), W, softmax_scale)
-            probs_cors_tmp[0, j, :] = prob
-            probs_cors_tmp[1, j, :] = cor
+        raise NotImplementedError(f"Model `{model}` is not yet implemented.")
 
-        probs_cors[0, starts[i] : starts[i] + n_neigh] = np_mean(
-            probs_cors_tmp[0], axis=0
-        )
-        probs_cors[1, starts[i] : starts[i] + n_neigh] = np_mean(
-            probs_cors_tmp[1], axis=0
+    def _estimate_softmax_scale(
+        self,
+        n_jobs: Optional[int] = None,
+        backend: Backend_t = _DEFAULT_BACKEND,
+        **kwargs,
+    ) -> float:
+        model = self._create_model(
+            VelocityModel.DETERMINISTIC, softmax_scale=1.0, **kwargs
         )
+        _, logits = model(n_jobs, backend)
+        return 1.0 / np.median(np.abs(logits.data))
 
-        if queue is not None:
-            queue.put(1)
-
-    if queue is not None:
-        queue.put(None)
-
-    return probs_cors
-
-
-def _run_stochastic(
-    ixs: np.ndarray,
-    scheme: Callable,
-    indices: np.ndarray,
-    indptr: np.ndarray,
-    expression: np.ndarray,
-    expectation: np.ndarray,
-    variance: np.ndarray,
-    softmax_scale: float = 1,
-    queue=None,
-) -> np.ndarray:
-    if not hasattr(scheme, "hessian"):
-        raise AttributeError("Velocity scheme doesn't have `hessian` attribute.")
-
-    starts = _calculate_starts(indptr, ixs)
-    probs_cors = np.empty((2, starts[-1]))
-    n_genes = expression.shape[1]
-
-    for i, ix in enumerate(ixs):
-        start, end = indptr[ix], indptr[ix + 1]
-
-        nbhs_ixs = indices[start:end]
-        n_neigh = len(nbhs_ixs)
-        assert n_neigh, f"Cell {ix} does not have any neighbors."
-
-        v = expectation[ix]
-
-        if np.all(v == 0):
-            p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
+    def _extract_layer(
+        self,
+        key: Optional[str] = None,
+        subset: Optional[Union[str, Sequence[str]]] = None,
+        dtype: np.dtype = np.float64,
+    ) -> np.ndarray:
+        if key in (None, "X"):
+            data = self.adata.X
+        elif key in self.adata.layers:
+            data = self.adata.layers[key]
         else:
-            # compute the Hessian tensor, and turn it into a matrix that has the diagonal elements in its rows
-            W = expression[nbhs_ixs, :] - expression[ix, :]
-
-            H = scheme.hessian(v, W, softmax_scale)
-            if H.shape == (n_neigh, n_genes):
-                H_diag = H
-            elif H.shape == (n_neigh, n_genes, n_genes):
-                H_diag = np.array([np.diag(h) for h in H])
-            else:
-                raise ValueError(
-                    f"Expected full Hessian matrix of shape `{(n_neigh, n_genes, n_genes)}` "
-                    f"or its diagonal of shape `{(n_neigh, n_genes)}`, found `{H.shape}`."
-                )
-
-            # compute zero order term
-            p_0, c = scheme(v[None, :], W, softmax_scale)
-
-            # compute second order term (note that the first order term cancels)
-            p_2 = 0.5 * H_diag.dot(variance[ix])
-
-            # combine both to give the second order Taylor approximation. Can sometimes be negative because we
-            # neglected higher order terms, so force it to be non-negative
-            p = np.clip(p_0 + p_2, a_min=0, a_max=1)
-
-            mask = np.isnan(p)
-            p[mask] = 0
-            c[mask] = 0
-
-            if np.all(p == 0):
-                p, c = _get_probs_for_zero_vec(len(nbhs_ixs))
-
-            sum_ = fsum(p)
-            if not np.isclose(sum_, 1.0, rtol=_RTOL):
-                p[~mask] = p[~mask] / sum_
-
-        probs_cors[0, starts[i] : starts[i] + n_neigh] = p
-        probs_cors[1, starts[i] : starts[i] + n_neigh] = c
-
-        if queue is not None:
-            queue.put(1)
-
-    if queue is not None:
-        queue.put(None)
-
-    return probs_cors
+            raise KeyError(f"Unable to find data in `adata.layers[{key}]`.")
 
+        if isinstance(subset, str):
+            subset = self.adata.var[subset]
+        subset = np.asarray(subset)
+        if np.issubdtype(subset.dtype, bool) and subset.shape == (data.shape[1],):
+            data = data[:, subset]
+        else:
+            data = data[:, np.isin(self.adata.var_names, subset)]
 
-@_dispatch_computation.register(VelocityMode.SAMPLING)
-def _(**kwargs):
-    # cant 't use partial - parallelize needs to know whether function has `py_func` attribute + doesn't have a __name__
-    kwargs["n_samples"] = 1
-    return _run_in_parallel(_run_mc, **kwargs)
+        data = data.astype(dtype, copy=False)
+        return data.toarray() if issparse(data) else data
 
+    @property
+    def logits(self) -> Optional[np.ndarray]:
+        """Array of shape ``(n_cells, n_cells)`` containing unnormalized transition matrix."""
+        return self._logits
 
-_dispatch_computation.register(VelocityMode.STOCHASTIC)(
-    lambda **kwargs: _run_in_parallel(_run_stochastic, **kwargs)
-)
-_dispatch_computation.register(VelocityMode.DETERMINISTIC)(
-    lambda **kwargs: _run_in_parallel(_run_deterministic, **kwargs)
-)
-_dispatch_computation.register(VelocityMode.MONTE_CARLO)(
-    lambda **kwargs: _run_in_parallel(_run_mc, **kwargs)
-)
+    def __invert__(self) -> "VelocityKernel":
+        dk = self._copy_ignore("_transition_matrix", "_logits")
+        dk._backward = not self.backward
+        dk._params = {}
+        return dk
diff --git a/cellrank/tl/kernels/utils/__init__.py b/cellrank/tl/kernels/utils/__init__.py
new file mode 100644
index 000000000..e38ad82cf
--- /dev/null
+++ b/cellrank/tl/kernels/utils/__init__.py
@@ -0,0 +1,19 @@
+from cellrank.tl.kernels.utils._tmat_flow import FlowPlotter
+from cellrank.tl.kernels.utils._similarity import (
+    Cosine,
+    DotProduct,
+    Correlation,
+    SimilarityABC,
+)
+from cellrank.tl.kernels.utils._random_walk import RandomWalk
+from cellrank.tl.kernels.utils._velocity_model import (
+    MonteCarlo,
+    Stochastic,
+    Deterministic,
+)
+from cellrank.tl.kernels.utils._pseudotime_scheme import (
+    ThresholdSchemeABC,
+    HardThresholdScheme,
+    SoftThresholdScheme,
+    CustomThresholdScheme,
+)
diff --git a/cellrank/tl/kernels/utils/_projection.py b/cellrank/tl/kernels/utils/_projection.py
new file mode 100644
index 000000000..f9c01670e
--- /dev/null
+++ b/cellrank/tl/kernels/utils/_projection.py
@@ -0,0 +1,135 @@
+from typing import Any, Optional
+
+import warnings
+
+import scvelo as scv
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.ul._docs import d
+from cellrank.tl.kernels._utils import _get_basis
+from scvelo.tools.velocity_embedding import quiver_autoscale
+
+import numpy as np
+from scipy.sparse import issparse
+
+
+class Projector:
+    """
+    Project transition matrix onto a low-dimensional embedding.
+
+    Should be used for visualization purposes.
+
+    Parameters
+    ----------
+    kexpr
+        Kernel that contains a transition matrix.
+    basis
+        Key in :attr:`anndata.AnnData.obsm` where the basis is stored.
+    """
+
+    def __init__(self, kexpr: "ExpernelExpression", basis: str = "umap"):  # noqa: F821
+        from cellrank.tl.kernels._mixins import ConnectivityMixin
+
+        for kernel in kexpr.kernels:
+            if not isinstance(kernel, ConnectivityMixin):
+                raise TypeError(
+                    f"{kernel!r} is not a kNN based kernel. The embedding projection "
+                    "only works for kNN based kernels."
+                )
+        self._kexpr = kexpr
+        self._basis = basis[2:] if basis.startswith("X_") else basis
+        self._key: Optional[str] = None
+
+    def project(
+        self,
+        key_added: Optional[str] = None,
+        recompute: bool = False,
+    ) -> None:
+        """
+        Project transition matrix onto an embedding.
+
+        Parameters
+        ----------
+        key_added
+            Key in :attr:`anndata.AnnData.obsm` where to store the projection.
+        recompute
+            Whether to recompute the projection if it already exists.
+
+        Returns
+        -------
+        Nothing, just updates :attr:`anndata.AnnData` with the projection and the parameters used for computation.
+        """
+        # modified from: https://github.com/theislab/scvelo/blob/master/scvelo/tools/velocity_embedding.py
+
+        self._key = Key.uns.kernel(self._kexpr.backward, key=key_added)
+        ukey = f"{self._key}_params"
+        key = self._key + "_" + self._basis
+
+        if not recompute and key in self._kexpr.adata.obsm:
+            logg.info(f"Using precomputed projection `adata.obsm[{key!r}]`")
+            return
+
+        start = logg.info(f"Projecting transition matrix onto `{self._basis}`")
+
+        emb = _get_basis(self._kexpr.adata, self._basis)
+        T_emb = np.empty_like(emb)
+        conn = self._kexpr.kernels[0]._conn
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            for row_id, row in enumerate(self._kexpr.transition_matrix):
+                conn_idxs = conn[row_id, :].indices
+
+                dX = emb[conn_idxs] - emb[row_id, None]
+
+                if np.any(np.isnan(dX)):
+                    T_emb[row_id, :] = np.nan
+                else:
+                    probs = row[:, conn_idxs]
+                    if issparse(probs):
+                        probs = probs.A.squeeze()
+
+                    dX /= np.linalg.norm(dX, axis=1)[:, None]
+                    dX = np.nan_to_num(dX)
+                    T_emb[row_id, :] = probs.dot(dX) - dX.sum(0) / dX.shape[0]
+
+        T_emb /= 3 * quiver_autoscale(np.nan_to_num(emb), T_emb)
+
+        embs = self._kexpr.adata.uns.get(ukey, {}).get("embeddings", [])
+        if self._basis not in embs:
+            embs = list(embs) + [self._basis]
+            self._kexpr.adata.uns[ukey] = self._kexpr.adata.uns.get(ukey, {})
+            self._kexpr.adata.uns[ukey]["embeddings"] = embs
+
+        logg.info(
+            f"Adding `adata.obsm[{key!r}]`\n    Finish",
+            time=start,
+        )
+        self._kexpr.adata.obsm[key] = T_emb
+
+    @d.dedent
+    def plot(self, *args: Any, stream: bool = True, **kwargs: Any) -> None:
+        """
+        Plot projected transition matrix in a embedding.
+
+        Parameters
+        ----------
+        args
+            Positional argument for the plotting function.
+        stream
+            If ``True``, use :func:`scvelo.pl.velocity_embedding_stream`.
+            Otherwise, use :func:`scvelo.pl.velocity_embedding_grid`.
+        kwargs
+            Keyword argument for the plotting function.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+        if stream:
+            return scv.pl.velocity_embedding_stream(
+                self._kexpr.adata, *args, basis=self._basis, vkey=self._key, **kwargs
+            )
+        return scv.pl.velocity_embedding_grid(
+            self._kexpr.adata, *args, basis=self._basis, vkey=self._key, **kwargs
+        )
diff --git a/cellrank/tl/kernels/_pseudotime_schemes.py b/cellrank/tl/kernels/utils/_pseudotime_scheme.py
similarity index 98%
rename from cellrank/tl/kernels/_pseudotime_schemes.py
rename to cellrank/tl/kernels/utils/_pseudotime_scheme.py
index 9d9efe337..2a77217fb 100644
--- a/cellrank/tl/kernels/_pseudotime_schemes.py
+++ b/cellrank/tl/kernels/utils/_pseudotime_scheme.py
@@ -8,6 +8,8 @@
 import numpy as np
 from scipy.sparse import csr_matrix
 
+__all__ = ("HardThresholdScheme", "SoftThresholdScheme", "CustomThresholdScheme")
+
 
 class ThresholdSchemeABC(ABC):
     """Base class for all connectivity biasing schemes."""
diff --git a/cellrank/tl/kernels/utils/_random_walk.py b/cellrank/tl/kernels/utils/_random_walk.py
new file mode 100644
index 000000000..6dea268a0
--- /dev/null
+++ b/cellrank/tl/kernels/utils/_random_walk.py
@@ -0,0 +1,375 @@
+from typing import Any, List, Tuple, Union, Mapping, Optional, Sequence
+from typing_extensions import Literal
+
+from pathlib import Path
+from itertools import chain
+
+import scvelo as scv
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import save_fig
+from cellrank.ul._parallelize import parallelize
+from cellrank.tl.kernels._utils import _get_basis
+
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+from pandas.api.types import infer_dtype, is_numeric_dtype, is_categorical_dtype
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap, to_hex
+from matplotlib.collections import LineCollection
+
+__all__ = ("RandomWalk",)
+Indices_t = Optional[
+    Union[Sequence[str], Mapping[str, Union[str, Sequence[str], Tuple[float, float]]]]
+]
+
+
+@d.dedent
+class RandomWalk:
+    """
+    Class that simulates a random walk on a Markov chain.
+
+    Parameters
+    ----------
+    %(adata)s
+    transition_matrix
+        Row-stochastic transition matrix.
+    start_ixs
+        Indices from which to uniformly sample the starting points. If `None`, use all points.
+    stop_ixs
+        Indices which when hit, the random walk is terminated.
+    """
+
+    def __init__(
+        self,
+        adata: AnnData,
+        transition_matrix: Union[np.ndarray, spmatrix],
+        start_ixs: Optional[Sequence[int]] = None,
+        stop_ixs: Optional[Sequence[int]] = None,
+    ):
+        if transition_matrix.ndim != 2 or (
+            transition_matrix.shape[0] != transition_matrix.shape[1]
+        ):
+            raise ValueError(
+                f"Expected transition matrix to be a square matrix, found `{transition_matrix.ndim}`."
+            )
+        if transition_matrix.shape[0] != adata.n_obs:
+            raise ValueError(
+                f"Expected transition matrix to be of shape `{adata.n_obs, adata.n_obs}`,"
+                f"found `{transition_matrix.shape}`."
+            )
+        if not np.allclose(transition_matrix.sum(1), 1.0):
+            raise ValueError("Transition matrix is not row-stochastic.")
+
+        self._adata = adata
+        self._tmat = transition_matrix
+        self._ixs = np.arange(self._tmat.shape[0])
+        self._is_sparse = issparse(self._tmat)
+
+        start_ixs = self._normalize_ixs(start_ixs, kind="start")
+        stop_ixs = self._normalize_ixs(stop_ixs, kind="stop")
+        self._stop_ixs = set([] if stop_ixs is None or not len(stop_ixs) else stop_ixs)
+        self._starting_dist = (
+            np.ones_like(self._ixs)
+            if start_ixs is None
+            else np.isin(self._ixs, start_ixs)
+        )
+        _sum = np.sum(self._starting_dist)
+        if _sum == 0:
+            raise ValueError("No starting indices have been selected.")
+
+        self._starting_dist = self._starting_dist.astype(transition_matrix.dtype) / _sum
+
+    @d.get_sections(base="rw_sim", sections=["Parameters"])
+    def simulate_one(
+        self,
+        max_iter: Union[int, float] = 0.25,
+        seed: Optional[int] = None,
+        successive_hits: int = 0,
+    ) -> np.ndarray:
+        """
+        Simulate one random walk.
+
+        Parameters
+        ----------
+        max_iter
+            Maximum number of steps of a random walk. If a :class:`float`, it can be specified
+            as a fraction of the number of cells.
+        seed
+            Random seed.
+        successive_hits
+            Number of successive hits in the ``stop_ixs`` required to stop prematurely.
+
+        Returns
+        -------
+        Array of shape ``(max_iter + 1,)`` of states that have been visited. If ``stop_ixs`` was specified, the array
+        may have smaller shape.
+        """
+        max_iter = self._max_iter(max_iter)
+        if successive_hits < 0:
+            raise ValueError(
+                f"Expected number of successive hits to be positive, found `{successive_hits}`."
+            )
+
+        rs = np.random.RandomState(seed)
+        ix = rs.choice(self._ixs, p=self._starting_dist)
+        sim, cnt = [ix], -1
+
+        for _ in range(max_iter):
+            ix = self._sample(ix, rs=rs)
+            sim.append(ix)
+            cnt = (cnt + 1) if self._should_stop(ix) else -1
+            if cnt >= successive_hits:
+                break
+
+        return np.array(sim)
+
+    def _simulate_many(
+        self,
+        sims: np.ndarray,
+        max_iter: Union[int, float] = 0.25,
+        seed: Optional[int] = None,
+        successive_hits: int = 0,
+        queue: Optional[Any] = None,
+    ) -> List[np.ndarray]:
+        res = []
+        for s in sims:
+            sim = self.simulate_one(
+                max_iter=max_iter,
+                seed=None if seed is None else seed + s,
+                successive_hits=successive_hits,
+            )
+            res.append(sim)
+            if queue is not None:
+                queue.put(1)
+
+        if queue is not None:
+            queue.put(None)
+
+        return res
+
+    @d.dedent
+    def simulate_many(
+        self,
+        n_sims: int,
+        max_iter: Union[int, float] = 0.25,
+        seed: Optional[int] = None,
+        successive_hits: int = 0,
+        n_jobs: Optional[int] = None,
+        backend: str = "loky",
+        show_progress_bar: bool = True,
+    ) -> List[np.ndarray]:
+        """
+        Simulate many random walks.
+
+        Parameters
+        ----------
+        n_sims
+            Number of random walks to simulate.
+        %(rw_sim.params)s
+        %(parallel)s
+
+        Returns
+        -------
+        List of arrays of shape ``(max_iter + 1,)`` of states that have been visited.
+        If ``stop_ixs`` was specified, the arrays may have smaller shape.
+        """
+        if n_sims <= 0:
+            raise ValueError(
+                f"Expected number of simulations to be positive, found `{n_sims}`."
+            )
+        max_iter = self._max_iter(max_iter)
+        start = logg.info(
+            f"Simulating `{n_sims}` random walks of maximum length `{max_iter}`"
+        )
+
+        simss = parallelize(
+            self._simulate_many,
+            collection=np.arange(n_sims),
+            n_jobs=n_jobs,
+            backend=backend,
+            show_progress_bar=show_progress_bar,
+            as_array=False,
+            unit="sim",
+        )(max_iter=max_iter, seed=seed, successive_hits=successive_hits)
+        simss = list(chain.from_iterable(simss))
+
+        logg.info("    Finish", time=start)
+
+        return simss
+
+    @d.dedent
+    def plot(
+        self,
+        sims: List[np.ndarray],
+        basis: str = "umap",
+        cmap: Union[str, LinearSegmentedColormap] = "gnuplot",
+        linewidth: float = 1.0,
+        linealpha: float = 0.3,
+        ixs_legend_loc: Optional[str] = None,
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Plot simulated random walks.
+
+        Parameters
+        ----------
+        sims
+            The simulated random walks.
+        basis
+            Basis used for plotting.
+        cmap
+            Colormap for the random walks.
+        linewidth
+            Line width for the random walks.
+        linealpha
+            Line alpha.
+        ixs_legend_loc
+            Position of the legend describing start- and endpoints.
+        %(plotting)s
+        kwargs
+            Keyword arguments for :func:`scvelo.pl.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+        emb = _get_basis(self._adata, basis)
+        if isinstance(cmap, str):
+            cmap = plt.get_cmap(cmap)
+        if not isinstance(cmap, LinearSegmentedColormap):
+            if not hasattr(cmap, "colors"):
+                raise AttributeError(
+                    "Unable to create a colormap, `cmap` does not have attribute `colors`."
+                )
+            cmap = LinearSegmentedColormap.from_list(
+                "random_walk",
+                colors=cmap.colors,
+                N=max(map(len, sims)),
+            )
+
+        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
+        scv.pl.scatter(self._adata, basis=basis, show=False, ax=ax, **kwargs)
+
+        logg.info("Plotting random walks")
+        for sim in sims:
+            x = emb[sim][:, 0]
+            y = emb[sim][:, 1]
+            points = np.array([x, y]).T.reshape(-1, 1, 2)
+            segments = np.concatenate([points[:-1], points[1:]], axis=1)
+            n_seg = len(segments)
+
+            lc = LineCollection(
+                segments,
+                linewidths=linewidth,
+                colors=[cmap(float(i) / n_seg) for i in range(n_seg)],
+                alpha=linealpha,
+                zorder=2,
+            )
+            ax.add_collection(lc)
+
+        for ix in [0, -1]:
+            ixs = [sim[ix] for sim in sims]
+            from scvelo.plotting.utils import default_size, plot_outline
+
+            plot_outline(
+                x=emb[ixs][:, 0],
+                y=emb[ixs][:, 1],
+                outline_color=("black", to_hex(cmap(float(abs(ix))))),
+                kwargs={
+                    "s": kwargs.get("s", default_size(self._adata)) * 1.1,
+                    "alpha": 0.9,
+                },
+                ax=ax,
+                zorder=4,
+            )
+
+        if ixs_legend_loc not in (None, "none"):
+            from cellrank.pl._utils import _position_legend
+
+            h1 = ax.scatter([], [], color=cmap(0.0), label="start")
+            h2 = ax.scatter([], [], color=cmap(1.0), label="stop")
+            legend = ax.get_legend()
+            if legend is not None:
+                ax.add_artist(legend)
+            _position_legend(ax, legend_loc=ixs_legend_loc, handles=[h1, h2])
+
+        if save is not None:
+            save_fig(fig, save)
+
+    def _normalize_ixs(
+        self, ixs: Indices_t, *, kind: Literal["start", "stop"]
+    ) -> Optional[np.ndarray]:
+        if ixs is None:
+            return None
+
+        if isinstance(ixs, dict):
+            # fmt: off
+            if len(ixs) != 1:
+                raise ValueError(f"Expected to find only 1 cluster key, found `{len(ixs)}`.")
+            key = next(iter(ixs.keys()))
+            if key not in self._adata.obs:
+                raise KeyError(f"Unable to find data in `adata.obs[{key!r}]`.")
+
+            vals = self._adata.obs[key]
+            if is_categorical_dtype(vals):
+                ixs = np.where(np.isin(vals, ixs[key]))[0]
+            elif is_numeric_dtype(vals):
+                if len(ixs[key]) != 2:
+                    raise ValueError(f"Expected range to be of length `2`, found `{len(ixs[key])}`")
+                minn, maxx = sorted(ixs[key])
+                ixs = np.where((vals >= minn) & (vals <= maxx))[0]
+            else:
+                raise TypeError(f"Expected `adata.obs[{key!r}]` to be numeric or categorical, "
+                                f"found `{infer_dtype(vals)}`.")
+            # fmt: on
+        elif isinstance(ixs, str):
+            ixs = np.where(self._adata.obs_names == ixs)[0]
+        elif isinstance(ixs[0], str):
+            ixs = np.where(np.isin(self._adata.obs_names, ixs))[0]
+        elif isinstance(ixs[0], bool):
+            if len(ixs) != self._adata.n_obs:
+                raise ValueError(
+                    f"Expected `bool` {kind} indices of length"
+                    f"`{self._adata.n_obs}`, found `{len(ixs)}`."
+                )
+            ixs = np.where(ixs)[0]
+        elif isinstance(ixs[0], int):
+            ixs = list(set(ixs))
+            if max(ixs) >= self._adata.n_obs:
+                raise IndexError(max(ixs))
+            if min(ixs) < -self._adata.n_obs:
+                raise IndexError(min(ixs))
+        else:
+            raise TypeError(
+                f"Expected {kind} indices to be either `dict` or a sequence of "
+                f"`int`, `str`, `bool`, found `{type(ixs).__nam__}`."
+            )
+
+        if not len(ixs):
+            raise ValueError(f"No {kind} indices have been selected.")
+
+        return ixs
+
+    def _should_stop(self, ix: int) -> bool:
+        return ix in self._stop_ixs
+
+    def _sample(self, ix: int, *, rs: np.random.RandomState) -> int:
+        return rs.choice(
+            self._ixs,
+            p=self._tmat[ix].A.squeeze() if self._is_sparse else self._tmat[ix],
+        )
+
+    def _max_iter(self, max_iter: Union[int, float]) -> int:
+        if isinstance(max_iter, float):
+            max_iter = int(np.ceil(max_iter * len(self._ixs)))
+        if max_iter <= 1:
+            raise ValueError(
+                f"Expected number of iterations to be > 1, found `{max_iter}`."
+            )
+        return max_iter
diff --git a/cellrank/tl/kernels/_velocity_schemes.py b/cellrank/tl/kernels/utils/_similarity.py
similarity index 92%
rename from cellrank/tl/kernels/_velocity_schemes.py
rename to cellrank/tl/kernels/utils/_similarity.py
index 06f83dcf6..514ccad98 100644
--- a/cellrank/tl/kernels/_velocity_schemes.py
+++ b/cellrank/tl/kernels/utils/_similarity.py
@@ -6,14 +6,15 @@
 
 from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d
-from cellrank.ul._utils import valuedispatch
 from cellrank.tl.kernels._utils import norm, np_mean, jit_kwargs
 
 import numpy as np
 from numba import njit
 
+__all__ = ("DotProduct", "Cosine", "Correlation")
 
-class Scheme(ModeEnum):  # noqa: D101
+
+class Similarity(ModeEnum):
     DOT_PRODUCT = auto()
     COSINE = auto()
     CORRELATION = auto()
@@ -46,7 +47,7 @@ def _predict_transition_probabilities_jax(
         softmax_scale: float = 1.0,
         center_mean: bool = True,
         scale_by_norm: bool = True,
-    ):
+    ) -> np.ndarray:
         if center_mean:
             # pearson correlation, otherwise cosine
             W -= W.mean(axis=1)[:, None]
@@ -140,7 +141,7 @@ def _predict_transition_probabilities_numpy(
     )
 
 
-class Hessian(ABC):  # noqa: D101
+class Hessian(ABC):
     @d.get_full_description(base="hessian")
     @d.get_sections(base="hessian", sections=["Parameters", "Returns"])
     @abstractmethod
@@ -175,7 +176,7 @@ def hessian(
         """
 
 
-class SimilaritySchemeABC(ABC):
+class SimilarityABC(ABC):
     """Base class for all similarity schemes."""
 
     @d.get_full_description(base="sim_scheme")
@@ -200,9 +201,19 @@ def __call__(
 
         Returns
         -------
-        The probability and logits arrays of shape ``(n_neighbors,)``.
+        The probability and unscaled logits arrays of shape ``(n_neighbors,)``.
         """
 
+    @staticmethod
+    def create(scheme: Similarity) -> "SimilarityABC":
+        if scheme == Similarity.CORRELATION:
+            return Correlation()
+        if scheme == Similarity.COSINE:
+            return Cosine()
+        if scheme == Similarity.DOT_PRODUCT:
+            return DotProduct()
+        raise NotImplementedError(scheme)
+
     def __repr__(self):
         return f"<{self.__class__.__name__}>"
 
@@ -210,7 +221,7 @@ def __str__(self):
         return repr(self)
 
 
-class SimilaritySchemeHessian(SimilaritySchemeABC, Hessian):
+class SimilarityHessian(SimilarityABC, Hessian):
     """
     Base class for all similarity schemes as defined in :cite:`li:20`.
 
@@ -272,7 +283,7 @@ def __getattribute__(self, name: str) -> Any:
         return super().__getattribute__(name)
 
 
-class DotProductScheme(SimilaritySchemeHessian):
+class DotProduct(SimilarityHessian):
     r"""
     Dot product scheme as defined in eq. (4.9) :cite:`li:20`.
 
@@ -286,7 +297,7 @@ def __init__(self):
         super().__init__(center_mean=False, scale_by_norm=False)
 
 
-class CosineScheme(SimilaritySchemeHessian):
+class Cosine(SimilarityHessian):
     r"""
     Cosine similarity scheme as defined in eq. (4.7) :cite:`li:20`.
 
@@ -300,7 +311,7 @@ def __init__(self):
         super().__init__(center_mean=False, scale_by_norm=True)
 
 
-class CorrelationScheme(SimilaritySchemeHessian):
+class Correlation(SimilarityHessian):
     r"""
     Pearson correlation scheme as defined in eq. (4.8) :cite:`li:20`.
 
@@ -312,23 +323,3 @@ class CorrelationScheme(SimilaritySchemeHessian):
 
     def __init__(self):
         super().__init__(center_mean=True, scale_by_norm=True)
-
-
-@valuedispatch
-def _get_scheme(scheme: Scheme, *_args, **_kwargs) -> SimilaritySchemeABC:
-    raise NotImplementedError(scheme)
-
-
-@_get_scheme.register(Scheme.DOT_PRODUCT)
-def _():
-    return DotProductScheme()
-
-
-@_get_scheme.register(Scheme.COSINE)
-def _():
-    return CosineScheme()
-
-
-@_get_scheme.register(Scheme.CORRELATION)
-def _():
-    return CorrelationScheme()
diff --git a/cellrank/tl/kernels/_tmat_flow.py b/cellrank/tl/kernels/utils/_tmat_flow.py
similarity index 99%
rename from cellrank/tl/kernels/_tmat_flow.py
rename to cellrank/tl/kernels/utils/_tmat_flow.py
index 94e48f181..f25db352d 100644
--- a/cellrank/tl/kernels/_tmat_flow.py
+++ b/cellrank/tl/kernels/utils/_tmat_flow.py
@@ -23,11 +23,13 @@
 from matplotlib.colors import to_rgb
 from matplotlib.collections import PolyCollection
 
+__all__ = ("FlowPlotter",)
+
 Numeric_t = Union[float, int]
 
 
 @dataclass(frozen=True)
-class Point:  # noqa: D101
+class Point:
     x: float
     xt: float
 
@@ -45,9 +47,9 @@ class FlowPlotter:
     tmat
         Matrix of shape ``(adata.n_obs, adata.n_obs)``.
     cluster_key
-        Key in :attr:`adata` ``.obs`` where clustering is stored.
+        Key in :attr:`anndata.AnnData.obs` where clustering is stored.
     time_key
-        Key in :attr:`adata` ``.obs`` where experimental time is stored.
+        Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
     """
 
     TIME_KEY = "time"
diff --git a/cellrank/tl/kernels/utils/_velocity_model.py b/cellrank/tl/kernels/utils/_velocity_model.py
new file mode 100644
index 000000000..576b24142
--- /dev/null
+++ b/cellrank/tl/kernels/utils/_velocity_model.py
@@ -0,0 +1,344 @@
+from typing import Any, Tuple, Union, Callable, Optional
+
+from abc import ABC, abstractmethod
+from enum import auto
+from math import fsum
+
+from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum, Backend_t
+from cellrank.ul._parallelize import parallelize
+from cellrank.tl.kernels._utils import _random_normal, _calculate_starts
+from cellrank.tl.kernels.utils._similarity import SimilarityABC
+
+import numpy as np
+from scipy.sparse import spmatrix, csr_matrix
+
+__all__ = ("Deterministic", "Stochastic", "MonteCarlo")
+
+
+class VelocityModel(ModeEnum):
+    DETERMINISTIC = auto()
+    STOCHASTIC = auto()
+    MONTE_CARLO = auto()
+
+
+class BackwardMode(ModeEnum):
+    TRANSPOSE = auto()
+    NEGATE = auto()
+
+
+class ModelABC(ABC):
+    def __init__(
+        self,
+        conn: spmatrix,
+        x: np.ndarray,
+        v: np.ndarray,
+        similarity: Union[
+            SimilarityABC,
+            Callable[[np.ndarray, np.ndarray, float], Tuple[np.ndarray, np.ndarray]],
+        ],
+        backward_mode: Optional[BackwardMode] = None,
+        softmax_scale: float = 1.0,
+        dtype: np.dtype = np.float64,
+    ):
+        self._conn = conn.astype(dtype, copy=False)
+        self._x = x.astype(dtype, copy=False)
+        self._v = v.astype(dtype, copy=False)
+        self._similarity = similarity
+        self._backward_mode = backward_mode
+        self._softmax_scale = softmax_scale
+        self._dtype = dtype
+
+    @abstractmethod
+    def _compute(
+        self, ix: int, neighs_ixs: np.ndarray, **kwargs: Any
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        pass
+
+    def __call__(
+        self,
+        n_jobs: Optional[int] = None,
+        backend: Backend_t = _DEFAULT_BACKEND,
+        show_progress_bar: bool = True,
+        **kwargs: Any,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        ixs = self._ixs
+        return parallelize(
+            self._compute_helper,
+            ixs,
+            n_jobs=n_jobs,
+            backend=backend,
+            show_progress_bar=show_progress_bar,
+            as_array=False,
+            extractor=lambda data: self._reconstruct_output(
+                np.concatenate(data, axis=-1), ixs
+            ),
+            unit=self._unit,
+        )(**kwargs)
+
+    def _compute_helper(
+        self, ixs: np.ndarray, queue=None, **kwargs
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        indptr, indices = self._conn.indptr, self._conn.indices
+        starts = _calculate_starts(indptr, ixs)
+        probs_logits = np.empty((2, starts[-1]), dtype=np.float64)
+
+        for i, ix in enumerate(ixs):
+            start, end = indptr[ix], indptr[ix + 1]
+            neigh_ixs = indices[start:end]
+            n_neigh = len(neigh_ixs)
+
+            ps, ls = self._compute(ix, neigh_ixs, **kwargs)
+            if np.shape(ps) != (n_neigh,):
+                raise ValueError(
+                    f"Expected row of shape `{(2, n_neigh)}`, found `{np.shape(ps)}`."
+                )
+
+            probs_logits[0, starts[i] : starts[i] + n_neigh] = ps
+            probs_logits[1, starts[i] : starts[i] + n_neigh] = ls
+            if queue is not None:
+                queue.put(1)
+
+        if queue is not None:
+            queue.put(None)
+
+        return probs_logits
+
+    def _reconstruct_output(
+        self,
+        data: np.ndarray,
+        ixs: Optional[np.ndarray] = None,
+    ) -> csr_matrix:
+        """
+        Transform :class:`numpy.ndarray` into :class:`scipy.sparse.csr_matrix`.
+
+        Parameters
+        ----------
+        data
+            Array of shape ``(2, nnz)``.
+        ixs
+            Indices that were used to sort the data.
+
+        Returns
+        -------
+        The reconstructed CSR matrix.
+        """
+
+        def reconstruct(data: np.ndarray) -> Tuple[csr_matrix, csr_matrix]:
+            data = csr_matrix(
+                (
+                    np.array(data, copy=True),
+                    np.array(conn.indices, copy=True),
+                    np.array(conn.indptr, copy=True),
+                )
+            )
+            if aixs is not None:
+                data = data[aixs]
+            data.eliminate_zeros()
+            return data
+
+        if data.shape != (2, self._conn.nnz):
+            raise ValueError(
+                f"Dimension or shape mismatch: `{data.shape}`, `{(2, self._conn.nnz)}`."
+            )
+
+        if ixs is None:
+            aixs, conn = None, self._conn
+        else:
+            aixs = np.argsort(ixs)
+            conn = self._conn[ixs]
+
+        probs, logits = reconstruct(data[0]), reconstruct(data[1])
+        close_to_1 = np.isclose(probs.sum(1), 1.0)
+        if not np.all(close_to_1):
+            raise ValueError(
+                f"Matrix is not row-stochastic, `{(~close_to_1).sum()}` do not sum to 1."
+            )
+
+        return probs, logits
+
+    @property
+    def _ixs(self) -> np.ndarray:
+        """Order in which to process rows."""
+        ixs = np.arange(self._conn.shape[0])
+        np.random.shuffle(ixs)
+        return ixs
+
+    def _uniform(self, n: int) -> Tuple[np.ndarray, np.ndarray]:
+        """Uninformative probability vector when velocity is 0."""
+        return np.ones(n, dtype=self._dtype) / n, np.zeros((n,), dtype=self._dtype)
+
+    @property
+    def _unit(self) -> str:
+        """Progress bar unit."""
+        return "cell"
+
+
+class Deterministic(ModelABC):
+    """Deterministic model."""
+
+    def _compute(self, ix: int, neigh_ixs: np.ndarray) -> Tuple[np.ndarray, np.ndarray]:
+        W = self._x[neigh_ixs, :] - self._x[ix, :]
+
+        if self._backward_mode not in (None, BackwardMode.NEGATE):
+            return self._similarity(self._v[neigh_ixs, :], -1 * W, self._softmax_scale)
+
+        v = self._v[ix]
+        if np.all(v == 0):
+            return self._uniform(len(neigh_ixs))
+
+        if self._backward_mode == BackwardMode.NEGATE:
+            v *= -1.0
+
+        return self._similarity(v[None, :], W, self._softmax_scale)
+
+
+class Stochastic(ModelABC):
+    """
+    Stochastic model.
+
+    Parameters
+    ----------
+    conn
+        The connectivity matrix of shape ``(n_cells, n_cells)``.
+    x
+        Array of shape ``(n_cells, n_features)``, such as imputed RNA expression used to calculate the displacement.
+    vmean
+        Array of shape ``(n_cells, n_features)`` containing the velocity mean.
+    vvar
+        Array of shape ``(n_cells, n_features)`` containing the velocity variance.
+    dtype
+        The data type to convert the arrays.
+    kwargs
+        Additional keyword arguments for the parent class.
+    """
+
+    def __init__(
+        self,
+        conn: spmatrix,
+        x: np.ndarray,
+        vmean: np.ndarray,
+        vvar: np.ndarray,
+        dtype: np.dtype = np.float64,
+        **kwargs: Any,
+    ):
+        super().__init__(conn, x, vmean, dtype=dtype, **kwargs)
+        if not hasattr(self._similarity, "hessian"):
+            raise AttributeError(
+                "SimilarityABC scheme doesn't have a `hessian` function."
+            )
+        self._var = vvar.astype(dtype, copy=False)
+
+    def _compute(
+        self,
+        ix: int,
+        nbhs_ixs: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        v = self._v[ix]
+        n_neigh, n_feat = len(nbhs_ixs), self._x.shape[1]
+
+        if np.all(v == 0):
+            return self._uniform(n_neigh)
+        # compute the Hessian tensor, and turn it into a matrix that has the diagonal elements in its rows
+        W = self._x[nbhs_ixs, :] - self._x[ix, :]
+
+        H = self._similarity.hessian(v, W, self._softmax_scale)
+        if H.shape == (n_neigh, n_feat):
+            H_diag = H
+        elif H.shape == (n_neigh, n_feat, n_feat):
+            H_diag = np.array([np.diag(h) for h in H])
+        else:
+            raise ValueError(
+                f"Expected full Hessian matrix of shape `{(n_neigh, n_feat, n_feat)}` "
+                f"or its diagonal of shape `{(n_neigh, n_feat)}`, found `{H.shape}`."
+            )
+
+        # compute zero order term
+        p_0, c = self._similarity(v[None, :], W, self._softmax_scale)
+
+        # compute second order term (note that the first order term cancels)
+        p_2 = 0.5 * H_diag.dot(self._var[ix])
+
+        # combine both to give the second order Taylor approximation.
+        # Can sometimes be negative because we neglected higher order terms, so force it to be non-negative
+        p = np.clip(p_0 + p_2, a_min=0, a_max=1)
+
+        nan_mask = np.isnan(p)
+        p[nan_mask] = 0
+        c[nan_mask] = 0
+
+        if np.all(p == 0):
+            return self._uniform(n_neigh)
+
+        sum_ = fsum(p)
+        if not np.isclose(sum_, 1.0, rtol=1e-12):
+            p[~nan_mask] = p[~nan_mask] / sum_
+
+        return p, c
+
+    @property
+    def _ixs(self) -> np.ndarray:
+        # less recompilation for JAX
+        return np.argsort((self._conn != 0).sum(1).A1)[::-1]
+
+
+class MonteCarlo(ModelABC):
+    """
+    Stochastic model.
+
+    Parameters
+    ----------
+    conn
+        The connectivity matrix of shape ``(n_cells, n_cells)``.
+    x
+        Array of shape ``(n_cells, n_features)``, such as imputed RNA expression used to calculate the displacement.
+    vmean
+        Array of shape ``(n_cells, n_features)`` containing the velocity mean.
+    vvar
+        Array of shape ``(n_cells, n_features)`` containing the velocity variance.
+    n_samples
+        Number of samples to generate from MVN distribution using ``exp`` and ``var``.
+    dtype
+        The data type used to convert the arrays.
+    kwargs
+        Additional keyword arguments for the parent class.
+    """
+
+    def __init__(
+        self,
+        conn: spmatrix,
+        x: np.ndarray,
+        vmean: np.ndarray,
+        vvar: np.ndarray,
+        n_samples: int = 1,
+        seed: Optional[int] = None,
+        dtype: np.dtype = np.float64,
+        **kwargs: Any,
+    ):
+        super().__init__(conn, x, vmean, dtype=dtype, **kwargs)
+        self._var = vvar.astype(dtype, copy=False)
+        self._n_samples = n_samples
+        np.random.seed(seed)
+
+    def _compute(
+        self,
+        ix: int,
+        nbhs_ixs: np.ndarray,
+    ) -> Tuple[np.ndarray, np.ndarray]:
+        n_neigh = len(nbhs_ixs)
+        W = self._x[nbhs_ixs, :] - self._x[ix, :]
+
+        # TODO(michalk8): don't store this?
+        samples = _random_normal(self._v[ix], self._var[ix], n_samples=self._n_samples)
+
+        probs = np.zeros((n_neigh,), dtype=np.float64)
+        logits = np.zeros((n_neigh,), dtype=np.float64)
+        for j in range(self._n_samples):
+            ps, ls = self._similarity(np.atleast_2d(samples[j]), W, self._softmax_scale)
+            probs += ps
+            logits += ls
+
+        return probs / self._n_samples, logits / self._n_samples
+
+    @property
+    def _unit(self) -> str:
+        return "sample"
diff --git a/cellrank/ul/_docs.py b/cellrank/ul/_docs.py
index 075542664..f8cf3207f 100644
--- a/cellrank/ul/_docs.py
+++ b/cellrank/ul/_docs.py
@@ -82,8 +82,7 @@
 
         - `{m.DETERMINISTIC!r}` - deterministic computation that doesn't propagate uncertainty.
         - `{m.MONTE_CARLO!r}` - Monte Carlo average of randomly sampled velocity vectors.
-        - `{m.STOCHASTIC!r}` - second order approximation, only available when :mod:`jax` is installed.
-        - `{m.SAMPLING!r}` - sample 1 transition matrix from the velocity distribution."""
+        - `{m.STOCHASTIC!r}` - second order approximation, only available when :mod:`jax` is installed."""
 _velocity_backward_mode = """\
 backward_mode
     Only matters if initialized as :attr:`backward` ``= True``.  Valid options are:
@@ -121,7 +120,7 @@
 Updates the :attr:`adata` with the following fields:
 
         - ``.obsp['{{key}}']`` - the transition matrix.
-        - ``.uns['{{key}}_params']`` - parameters used for calculation."""
+        - ``.uns['{{key}}_params']`` - parameters used for the calculation."""
 _en_cutoff_p_thresh = """\
 en_cutoff
     If ``cluster_key`` is given, this parameter determines when an approximate recurrent class will
@@ -137,14 +136,14 @@
     Basis to use when ``mode = 'embedding'``. If `None`, use `'umap'`."""
 _velocity_scheme = """\
 scheme
-    Similarity scheme between cells as described in :cite:`li:20`. Can be one of the following:
+    Similarity measure between cells as described in :cite:`li:20`. Can be one of the following:
 
-        - `{s.DOT_PRODUCT!r}` - :class:`cellrank.tl.kernels.DotProductScheme`.
-        - `{s.COSINE!r}` - :class:`cellrank.tl.kernels.CosineScheme`.
-        - `{s.CORRELATION!r}` - :class:`cellrank.tl.kernels.CorrelationScheme`.
+        - `{s.CORRELATION!r}` - :class:`cellrank.tl.kernels.utils.Correlation`.
+        - `{s.COSINE!r}` - :class:`cellrank.tl.kernels.utils.Cosine`.
+        - `{s.DOT_PRODUCT!r}` - :class:`cellrank.tl.kernels.utils.DotProduct`.
 
     Alternatively, any function can be passed as long as it follows the signature of
-    :meth:`cellrank.tl.kernels.SimilaritySchemeABC.__call__`."""
+    :meth:`cellrank.tl.kernels.utils.SimilarityABC.__call__`."""
 _cond_num = """\
 compute_cond_num
     Whether to compute condition number of the transition matrix. Note that this might be costly,
diff --git a/docs/source/classes.rst b/docs/source/classes.rst
index ccf656135..3ac1c334c 100644
--- a/docs/source/classes.rst
+++ b/docs/source/classes.rst
@@ -29,17 +29,17 @@ Velocity Kernel
 
 Cosine Similarity Scheme
 ++++++++++++++++++++++++
-.. autoclass:: cellrank.tl.kernels.CosineScheme
+.. autoclass:: cellrank.tl.kernels.utils.Cosine
     :members: __call__, hessian
 
 Correlation Scheme
 ++++++++++++++++++
-.. autoclass:: cellrank.tl.kernels.CorrelationScheme
+.. autoclass:: cellrank.tl.kernels.utils.Correlation
     :members: __call__, hessian
 
 Dot Product Scheme
 ++++++++++++++++++
-.. autoclass:: cellrank.tl.kernels.DotProductScheme
+.. autoclass:: cellrank.tl.kernels.utils.DotProduct
     :members: __call__, hessian
 
 Connectivity Kernel
@@ -56,13 +56,13 @@ Pseudotime Kernel
 
 Hard Threshold Scheme
 +++++++++++++++++++++
-.. autoclass:: cellrank.tl.kernels.HardThresholdScheme
+.. autoclass:: cellrank.tl.kernels.utils.HardThresholdScheme
     :members:
     :special-members: __call__
 
 Soft Threshold Scheme
 +++++++++++++++++++++
-.. autoclass:: cellrank.tl.kernels.SoftThresholdScheme
+.. autoclass:: cellrank.tl.kernels.utils.SoftThresholdScheme
     :members:
     :special-members: __call__
 
@@ -130,14 +130,14 @@ TransportMap Kernel
 
 Similarity Scheme
 -----------------
-.. autoclass:: cellrank.tl.kernels.SimilaritySchemeABC
+.. autoclass:: cellrank.tl.kernels.utils.SimilarityABC
     :members:
     :special-members: __call__
     :inherited-members:
 
 Threshold Scheme
 ----------------
-.. autoclass:: cellrank.tl.kernels.ThresholdSchemeABC
+.. autoclass:: cellrank.tl.kernels.utils.ThresholdSchemeABC
     :members:
     :special-members: __call__
     :inherited-members:
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 79c42491f..726bafd97 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -24,6 +24,10 @@ or fate-informed `circular embeddings`_, to name just a few.
 
 CellRank scales to large cell numbers, is fully compatible with `scanpy`_ and `scvelo`_ and is easy to use.
 
+Manuscript
+^^^^^^^^^^
+Please check out our manuscript :cite:`lange:22` in **Nature Methods** to learn more.
+
 Getting started with CellRank
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 If you're new to CellRank, make sure to go though the `basic tutorial`_ which introduces you to CellRank's high-level
@@ -34,10 +38,6 @@ which allows to unlock the full power of CellRank. If you want to see individual
 To use CellRank without RNA velocity information, check out the `beyond RNA velocity tutorial`_ as well as the
 `time-series tutorial`_.
 
-Manuscript
-^^^^^^^^^^
-Please see our `preprint`_ on **bioRxiv** to learn more.
-
 CellRank's key applications
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^
 - compute initial & terminal as well as intermediate `macrostates`_ of your biological system
@@ -154,7 +154,6 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
 .. _Theislab: https://www.helmholtz-muenchen.de/icb/research/groups/theis-lab/overview/index.html
 .. _Peerlab: https://www.mskcc.org/research/ski/labs/dana-pe-er
 .. _`tweet`: https://twitter.com/MariusLange8
-.. _preprint: https://www.biorxiv.org/content/10.1101/2020.10.19.345983v2
 .. _PageRank: https://en.wikipedia.org/wiki/PageRank#cite_note-1
 
 .. _gallery: https://cellrank.readthedocs.io/en/stable/auto_examples/index.html
diff --git a/docs/source/references.bib b/docs/source/references.bib
index 68d093331..d5bae7303 100644
--- a/docs/source/references.bib
+++ b/docs/source/references.bib
@@ -413,14 +413,15 @@ @article{tirosh:16:science
 	journal = {Science}
 }
 
-@article{lange:20,
+@article{lange:22,
     title={CellRank for directed single-cell fate mapping},
     author={Lange, Marius and Bergen, Volker and Klein, Michal and Setty, Manu and Reuter, Bernhard and Bakhti, Mostafa
-            and Lickert, Heiko and Ansari, Meshal and Schniering, Janine and Schiller, Herbert B and others},
-    journal={BioRxiv},
-    year={2020},
-    doi={10.1101/2020.10.19.345983},
-    publisher={Cold Spring Harbor Laboratory}
+            and Lickert, Heiko and Ansari, Meshal and Schniering, Janine and Schiller, Herbert B. and
+            Pe'er, Dana and Theis, Fabian J.},
+    journal={Nat. Methods},
+    year={2022},
+    doi={10.1038/s41592-021-01346-6},
+    publisher={Nature Publishing Group}
 }
 
 @article{nowotschin:19,
diff --git a/docs/source/release/notes-1.1.0.rst b/docs/source/release/notes-1.1.0.rst
index e91edd485..ece10d480 100644
--- a/docs/source/release/notes-1.1.0.rst
+++ b/docs/source/release/notes-1.1.0.rst
@@ -19,5 +19,4 @@ Features
   that the choice of metric can lead to slightly different results. Users can now also supply their own scheme as long
   as it follows the signature of :class:`cellrank.tl.kernels.SimilaritySchemeABC`.
 - :func:`cellrank.datasets.reprogramming` has been added to allow for easy reproducibility of the time & memory
-  benchmarking results in our `CellRank preprint <https://doi.org/10.1101/2020.10.19.345983>`_. This is a reprogramming
-  dataset from :cite:`morris:18`.
+  benchmarking results in our manuscript :cite:`lange:22`. This is a reprogramming dataset from :cite:`morris:18`.
diff --git a/docs/source/release/notes-1.5.1.rst b/docs/source/release/notes-1.5.1.rst
new file mode 100644
index 000000000..6135c19ec
--- /dev/null
+++ b/docs/source/release/notes-1.5.1.rst
@@ -0,0 +1,3 @@
+CellRank 1.5.1 (2022-01-13)
+===========================
+No significant changes.
diff --git a/docs/source/release/notes-dev.rst b/docs/source/release/notes-dev.rst
index d20577b32..4c17b66dd 100644
--- a/docs/source/release/notes-dev.rst
+++ b/docs/source/release/notes-dev.rst
@@ -4,7 +4,7 @@ CellRank dev (2022-02-10)
 Features
 --------
 
-- Update CytoTRACE calculation to check inf vals
+- Update CytoTRACE calculation to check for infinite values.
   `#761 <https://github.com/theislab/cellrank/pull/761>`__
 
 - Refactor :class:`cellrank.tl.kernels.TransportMapKernel` for easier definition of extensions.
diff --git a/examples/kernels/compute_kernel_tricks.py b/examples/kernels/compute_kernel_tricks.py
deleted file mode 100644
index 18d7d4c98..000000000
--- a/examples/kernels/compute_kernel_tricks.py
+++ /dev/null
@@ -1,97 +0,0 @@
-"""
-Kernel tricks
--------------
-
-This example shows some niche, but useful functionalities of :class:`cellrank.tl.kernels.Kernel`.
-
-CellRank is split into :mod:`cellrank.tl.kernels` and :mod:`cellrank.tl.estimators`. Kernels compute transition matrices
-based on some inputs, like RNA velocity :cite:`manno:18` :cite:`bergen:20`, while estimators perform inference based on
-a given kernel, e.g. they compute initial and terminal cells and fate probabilities.
-
-Here, will will dive a bit deeper into the how these kernel objects work.
-"""
-
-import cellrank as cr
-
-adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
-adata
-
-# %%
-# First, we create some kernels which will be used to compute the cell-to-cell transition matrix:
-#
-# - :class:`cellrank.tl.kernels.ConnectivityKernel` computes the transition matrix using the KNN graph from
-#   :func:`scanpy.pp.neighbors` :cite:`wolf:18`. Note that this kernel is by itself directionless and
-#   should be used in conjunction with e.g. :class:`cellrank.tl.kernels.VelocityKernel`.
-# - :class:`cellrank.tl.kernels.VelocityKernel` uses RNA velocity for the transition matrix computation :cite:`manno:18`
-#   :cite:`bergen:20`, but can also take into account uncertainty in RNA velocity.
-# - :class:`cellrank.tl.kernels.PseudotimeKernel` works similarly as in Palantir :cite:`setty:19` - it orients the edges
-#   of the KNN graph constructed in the expression space using the pseudotemporal ordering of cells, such as
-#   Diffusion Pseudotime (DPT) :cite:`haghverdi:16`.
-ck = cr.tl.kernels.ConnectivityKernel(adata)
-vk = cr.tl.kernels.VelocityKernel(adata)
-pk = cr.tl.kernels.PseudotimeKernel(adata)
-
-# %%
-# Kernels can be easily combined with each other. All of the constants within parentheses
-# (if they are needed) will be normalized to 1 automatically. This effect will be shown after the transition matrix
-# has been computed.
-10 * vk + 15 * ck
-
-# %%
-# We can also build more complex kernel expression. Note that multiple usage of the same kernel instance in the
-# expression caches it's transition matrix, so it's only computed once.
-k = ((vk + vk + 42 * vk) + ck * 2) + ck * (3 * vk + 4 * pk)
-k
-
-# %%
-# For complex expression like the one above, one can get the unique base kernels as follows.
-k.kernels
-
-# %%
-# Kernels can also be inverted (i.e. the ``backward`` attribute is set to the opposite value to the current one).
-# Note that this operation does *NOT* create a copy and modifies the kernel inplace, most importantly, the computed
-# transition matrix is removed. This makes it more easier to recompute the transition matrix of kernels,
-# since the handles point to the original objects.
-#
-# Note that in the 2nd :func:`print` statement, we access the private attribute - that's because accessing
-# :attr:`cellrank.tl.kernels.Kernel.transition_matrix` computes the transition matrix with default values.
-# This happens only with basic kernels and not the kernel expressions and only if they are not part of a
-# larger expression.
-ck.compute_transition_matrix()
-print(ck.transition_matrix is not None)
-
-inv_ck = ~ck
-print(ck._transition_matrix is not None)
-print(inv_ck is ck)
-
-# %%
-# We can also easily copy the kernel objects.
-vk.compute_transition_matrix(
-    mode="deterministic", softmax_scale=4, show_progress_bar=False
-)
-print(vk.transition_matrix is not None)
-
-inv_vk = ~(vk.copy())
-print(vk._transition_matrix is not None)
-print(vk is inv_vk)
-
-# %%
-# After the transition matrix is computed, we can inspect the parameters used in its computation.
-ck.compute_transition_matrix()
-ck.params
-
-# %%
-# The transition matrix can be written to :class:`anndata.AnnData` object. It is saved in the ``.obsp`` attribute.
-# This also writes the parameters to the ``.uns`` attribute.
-ck.write_to_adata(key="transition_matrix")
-adata.obsp["transition_matrix"]
-
-# %%
-# Precomputed kernels are useful if you have a transition matrix that was computed outside of CellRank that you would
-# like to analyze using one of our estimators. It’s essentially an interface to CellRank’s estimators.
-#
-# Below we supply a transition matrix saved in ``adata.obsp["transition_matrix"]``, but :class:`anndata.AnnData` object
-# is not required and we can just supply either :mod:`numpy` or :mod:`scipy.sparse` array. In that case, a minimal
-# empty :class:`anndata.AnnData` object is created, as shown below.
-pk = cr.tl.kernels.PrecomputedKernel(adata.obsp["transition_matrix"])
-pk.adata
diff --git a/examples/kernels/plot_projection.py b/examples/kernels/plot_projection.py
index 69caf3c7c..d05bb15c9 100644
--- a/examples/kernels/plot_projection.py
+++ b/examples/kernels/plot_projection.py
@@ -22,7 +22,7 @@
 # %%
 # Next, we compute the projection of the transition matrix onto an embedding, in this case UMAP.
 # The projected transition matrix can be found in ``adata.obsm['T_fwd_umap']``.
-k.compute_projection(basis="umap")
+k.plot_projection(basis="umap")
 
 # %%
 # The projected transition matrix can now be visualized using :mod:`scvelo`.
diff --git a/examples/plotting/plot_initial_states.py b/examples/plotting/plot_initial_states.py
index 23793e166..becf34f62 100644
--- a/examples/plotting/plot_initial_states.py
+++ b/examples/plotting/plot_initial_states.py
@@ -20,10 +20,10 @@
 cr.tl.initial_states(
     adata,
     cluster_key="clusters",
-    n_cells=30,
     softmax_scale=4,
     n_states=1,
     show_progress_bar=False,
+    fit_kwargs={"n_cells": 30},
 )
 
 # %%
diff --git a/examples/plotting/plot_lineage_drivers.py b/examples/plotting/plot_lineage_drivers.py
index 9e2716c67..cbc0d5a46 100644
--- a/examples/plotting/plot_lineage_drivers.py
+++ b/examples/plotting/plot_lineage_drivers.py
@@ -17,7 +17,6 @@
 cr.tl.terminal_states(
     adata,
     cluster_key="clusters",
-    n_cells=30,
     n_states=3,
     softmax_scale=4,
     show_progress_bar=False,
diff --git a/examples/plotting/plot_lineages.py b/examples/plotting/plot_lineages.py
index 25b4a63e8..c176210d0 100644
--- a/examples/plotting/plot_lineages.py
+++ b/examples/plotting/plot_lineages.py
@@ -23,7 +23,6 @@
 cr.tl.terminal_states(
     adata,
     cluster_key="clusters",
-    n_cells=30,
     n_states=3,
     softmax_scale=4,
     show_progress_bar=False,
diff --git a/examples/plotting/plot_terminal_states.py b/examples/plotting/plot_terminal_states.py
index f8a075027..ea5961157 100644
--- a/examples/plotting/plot_terminal_states.py
+++ b/examples/plotting/plot_terminal_states.py
@@ -20,10 +20,10 @@
 cr.tl.terminal_states(
     adata,
     cluster_key="clusters",
-    n_cells=30,
     softmax_scale=4,
     n_states=3,
     show_progress_bar=False,
+    fit_kwargs={"n_cells": 30},
 )
 
 # %%
diff --git a/requirements.txt b/requirements.txt
index 5958ef493..c4c424964 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -6,7 +6,7 @@ numba>=0.51.0
 numpy>=1.17.0
 pandas>=1.3.0,!=1.4.0
 pygam>=0.8.0
-pygpcca>=1.0.2
+pygpcca>=1.0.3
 scanpy>=1.7.2
 scikit-learn>=0.24.0
 scipy>=1.2.0
diff --git a/setup.py b/setup.py
index 0345c5272..41eb2a3ef 100644
--- a/setup.py
+++ b/setup.py
@@ -8,7 +8,7 @@
     __author__ = "Marius Lange, Michal Klein"
     __maintainer__ = "Marius Lange, Michal Klein"
     __email__ = "info@cellrank.org"
-    __version__ = "1.5.0"
+    __version__ = "1.5.1"
 
 
 if __name__ == "__main__":


Test patch:
diff --git a/tests/_ground_truth_figures/all_models_for_1_gene_failed.png b/tests/_ground_truth_figures/all_models_for_1_gene_failed.png
index 996b648ce..1284d6cec 100644
Binary files a/tests/_ground_truth_figures/all_models_for_1_gene_failed.png and b/tests/_ground_truth_figures/all_models_for_1_gene_failed.png differ
diff --git a/tests/_ground_truth_figures/all_models_for_1_gene_failed_same_plot.png b/tests/_ground_truth_figures/all_models_for_1_gene_failed_same_plot.png
index ba56e4256..b260340d4 100644
Binary files a/tests/_ground_truth_figures/all_models_for_1_gene_failed_same_plot.png and b/tests/_ground_truth_figures/all_models_for_1_gene_failed_same_plot.png differ
diff --git a/tests/_ground_truth_figures/all_models_for_1_lineage_failed.png b/tests/_ground_truth_figures/all_models_for_1_lineage_failed.png
index f49804a55..358b8f829 100644
Binary files a/tests/_ground_truth_figures/all_models_for_1_lineage_failed.png and b/tests/_ground_truth_figures/all_models_for_1_lineage_failed.png differ
diff --git a/tests/_ground_truth_figures/bar.png b/tests/_ground_truth_figures/bar.png
index 7afca4616..7969ac931 100644
Binary files a/tests/_ground_truth_figures/bar.png and b/tests/_ground_truth_figures/bar.png differ
diff --git a/tests/_ground_truth_figures/bar_bwd.png b/tests/_ground_truth_figures/bar_bwd.png
index 387156dcd..c28ce356c 100644
Binary files a/tests/_ground_truth_figures/bar_bwd.png and b/tests/_ground_truth_figures/bar_bwd.png differ
diff --git a/tests/_ground_truth_figures/bar_cluster_subset.png b/tests/_ground_truth_figures/bar_cluster_subset.png
index 4da34863a..912a94e3e 100644
Binary files a/tests/_ground_truth_figures/bar_cluster_subset.png and b/tests/_ground_truth_figures/bar_cluster_subset.png differ
diff --git a/tests/_ground_truth_figures/bar_lineage_subset.png b/tests/_ground_truth_figures/bar_lineage_subset.png
index c0a964792..844f268d2 100644
Binary files a/tests/_ground_truth_figures/bar_lineage_subset.png and b/tests/_ground_truth_figures/bar_lineage_subset.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage.png b/tests/_ground_truth_figures/cluster_lineage.png
index 1a5bbf532..d30bedc99 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage.png and b/tests/_ground_truth_figures/cluster_lineage.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png b/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png
index 73e93af20..434767853 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png and b/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_bwd.png b/tests/_ground_truth_figures/cluster_lineage_bwd.png
index 874886a8e..9d09464e2 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_bwd.png and b/tests/_ground_truth_figures/cluster_lineage_bwd.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates.png b/tests/_ground_truth_figures/cluster_lineage_covariates.png
index 3d5045bb3..b2f4786f8 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_covariates.png and b/tests/_ground_truth_figures/cluster_lineage_covariates.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png b/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png
index 4b528f0c6..d134f8e04 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png and b/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png b/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png
index 0cbb34f42..c2e376372 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png and b/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_data_key.png b/tests/_ground_truth_figures/cluster_lineage_data_key.png
index 320de1d33..675226d94 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_data_key.png and b/tests/_ground_truth_figures/cluster_lineage_data_key.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_gene_symbols.png b/tests/_ground_truth_figures/cluster_lineage_gene_symbols.png
index f1a7cbabe..dc91771b5 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_gene_symbols.png and b/tests/_ground_truth_figures/cluster_lineage_gene_symbols.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_leiden.png b/tests/_ground_truth_figures/cluster_lineage_leiden.png
index 1a5bbf532..d30bedc99 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_leiden.png and b/tests/_ground_truth_figures/cluster_lineage_leiden.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_no_norm.png b/tests/_ground_truth_figures/cluster_lineage_no_norm.png
index 21323f9e3..87ba3e9a8 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_no_norm.png and b/tests/_ground_truth_figures/cluster_lineage_no_norm.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_random_state.png b/tests/_ground_truth_figures/cluster_lineage_random_state.png
index 1a5bbf532..d30bedc99 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_random_state.png and b/tests/_ground_truth_figures/cluster_lineage_random_state.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_raw.png b/tests/_ground_truth_figures/cluster_lineage_raw.png
index 211574a2e..b47d16d7b 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_raw.png and b/tests/_ground_truth_figures/cluster_lineage_raw.png differ
diff --git a/tests/_ground_truth_figures/composition.png b/tests/_ground_truth_figures/composition.png
index b46791d90..e3656e813 100644
Binary files a/tests/_ground_truth_figures/composition.png and b/tests/_ground_truth_figures/composition.png differ
diff --git a/tests/_ground_truth_figures/composition_kwargs_autopct.png b/tests/_ground_truth_figures/composition_kwargs_autopct.png
index fb2b71fe2..6fd726e65 100644
Binary files a/tests/_ground_truth_figures/composition_kwargs_autopct.png and b/tests/_ground_truth_figures/composition_kwargs_autopct.png differ
diff --git a/tests/_ground_truth_figures/driver_corr.png b/tests/_ground_truth_figures/driver_corr.png
index c0a796dc3..990c4e53d 100644
Binary files a/tests/_ground_truth_figures/driver_corr.png and b/tests/_ground_truth_figures/driver_corr.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_adjust_text.png b/tests/_ground_truth_figures/driver_corr_adjust_text.png
index 95caa6f8b..4d8414b7a 100644
Binary files a/tests/_ground_truth_figures/driver_corr_adjust_text.png and b/tests/_ground_truth_figures/driver_corr_adjust_text.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_cmap.png b/tests/_ground_truth_figures/driver_corr_cmap.png
index cab76756c..a87ec56db 100644
Binary files a/tests/_ground_truth_figures/driver_corr_cmap.png and b/tests/_ground_truth_figures/driver_corr_cmap.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_color.png b/tests/_ground_truth_figures/driver_corr_color.png
index a6385b72a..b27f6e5f2 100644
Binary files a/tests/_ground_truth_figures/driver_corr_color.png and b/tests/_ground_truth_figures/driver_corr_color.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_fontsize.png b/tests/_ground_truth_figures/driver_corr_fontsize.png
index 384dd2b5b..22d962ab3 100644
Binary files a/tests/_ground_truth_figures/driver_corr_fontsize.png and b/tests/_ground_truth_figures/driver_corr_fontsize.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets.png b/tests/_ground_truth_figures/driver_corr_gene_sets.png
index 79c92609d..5935b545f 100644
Binary files a/tests/_ground_truth_figures/driver_corr_gene_sets.png and b/tests/_ground_truth_figures/driver_corr_gene_sets.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png
index f24953d5e..12df083ca 100644
Binary files a/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png and b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_legend_loc.png b/tests/_ground_truth_figures/driver_corr_legend_loc.png
index a16c49400..4e6e83b7b 100644
Binary files a/tests/_ground_truth_figures/driver_corr_legend_loc.png and b/tests/_ground_truth_figures/driver_corr_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_lenged_loc.png b/tests/_ground_truth_figures/driver_corr_lenged_loc.png
deleted file mode 100644
index 555b5744c..000000000
Binary files a/tests/_ground_truth_figures/driver_corr_lenged_loc.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/driver_corr_return_ax.png b/tests/_ground_truth_figures/driver_corr_return_ax.png
index 19beb887b..1e5c49568 100644
Binary files a/tests/_ground_truth_figures/driver_corr_return_ax.png and b/tests/_ground_truth_figures/driver_corr_return_ax.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_use_raw.png b/tests/_ground_truth_figures/driver_corr_use_raw.png
index dc9642800..71166be67 100644
Binary files a/tests/_ground_truth_figures/driver_corr_use_raw.png and b/tests/_ground_truth_figures/driver_corr_use_raw.png differ
diff --git a/tests/_ground_truth_figures/drivers_ascending.png b/tests/_ground_truth_figures/drivers_ascending.png
index 8fdbbe8bd..9a1d1b972 100644
Binary files a/tests/_ground_truth_figures/drivers_ascending.png and b/tests/_ground_truth_figures/drivers_ascending.png differ
diff --git a/tests/_ground_truth_figures/drivers_backward.png b/tests/_ground_truth_figures/drivers_backward.png
index b02731242..b46e96f92 100644
Binary files a/tests/_ground_truth_figures/drivers_backward.png and b/tests/_ground_truth_figures/drivers_backward.png differ
diff --git a/tests/_ground_truth_figures/drivers_cmap.png b/tests/_ground_truth_figures/drivers_cmap.png
index 332fde86c..0dc9de314 100644
Binary files a/tests/_ground_truth_figures/drivers_cmap.png and b/tests/_ground_truth_figures/drivers_cmap.png differ
diff --git a/tests/_ground_truth_figures/drivers_n_genes.png b/tests/_ground_truth_figures/drivers_n_genes.png
index 2a282bfae..8125e92a9 100644
Binary files a/tests/_ground_truth_figures/drivers_n_genes.png and b/tests/_ground_truth_figures/drivers_n_genes.png differ
diff --git a/tests/_ground_truth_figures/drivers_title_fmt.png b/tests/_ground_truth_figures/drivers_title_fmt.png
index 31ccd086f..6b7523405 100644
Binary files a/tests/_ground_truth_figures/drivers_title_fmt.png and b/tests/_ground_truth_figures/drivers_title_fmt.png differ
diff --git a/tests/_ground_truth_figures/failed_only_main_diagonal.png b/tests/_ground_truth_figures/failed_only_main_diagonal.png
index 421d8d56f..6e656ecb8 100644
Binary files a/tests/_ground_truth_figures/failed_only_main_diagonal.png and b/tests/_ground_truth_figures/failed_only_main_diagonal.png differ
diff --git a/tests/_ground_truth_figures/failed_only_off_diagonal.png b/tests/_ground_truth_figures/failed_only_off_diagonal.png
index a37f4e5ed..df43f17cf 100644
Binary files a/tests/_ground_truth_figures/failed_only_off_diagonal.png and b/tests/_ground_truth_figures/failed_only_off_diagonal.png differ
diff --git a/tests/_ground_truth_figures/fitted_cluster_fates.png b/tests/_ground_truth_figures/fitted_cluster_fates.png
index 69a6d2a6b..465aa5520 100644
Binary files a/tests/_ground_truth_figures/fitted_cluster_fates.png and b/tests/_ground_truth_figures/fitted_cluster_fates.png differ
diff --git a/tests/_ground_truth_figures/fitted_empty_model.png b/tests/_ground_truth_figures/fitted_empty_model.png
index fb385925c..7ca73cd5b 100644
Binary files a/tests/_ground_truth_figures/fitted_empty_model.png and b/tests/_ground_truth_figures/fitted_empty_model.png differ
diff --git a/tests/_ground_truth_figures/fitted_gene_trends.png b/tests/_ground_truth_figures/fitted_gene_trends.png
index 3eabc2a7d..d1278676d 100644
Binary files a/tests/_ground_truth_figures/fitted_gene_trends.png and b/tests/_ground_truth_figures/fitted_gene_trends.png differ
diff --git a/tests/_ground_truth_figures/fitted_heatmap/lineage_0-failed-diff.png b/tests/_ground_truth_figures/fitted_heatmap/lineage_0-failed-diff.png
new file mode 100644
index 000000000..e1994255f
Binary files /dev/null and b/tests/_ground_truth_figures/fitted_heatmap/lineage_0-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/fitted_heatmap/lineage_0.png b/tests/_ground_truth_figures/fitted_heatmap/lineage_0.png
index 7d8f5b5dd..fdce20921 100644
Binary files a/tests/_ground_truth_figures/fitted_heatmap/lineage_0.png and b/tests/_ground_truth_figures/fitted_heatmap/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/fitted_heatmap/lineage_1.png b/tests/_ground_truth_figures/fitted_heatmap/lineage_1.png
index 3ccaeedac..a2cf51f89 100644
Binary files a/tests/_ground_truth_figures/fitted_heatmap/lineage_1.png and b/tests/_ground_truth_figures/fitted_heatmap/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/fitted_heatmap/lineage_2.png b/tests/_ground_truth_figures/fitted_heatmap/lineage_2.png
index 987b2f8ab..69d29f299 100644
Binary files a/tests/_ground_truth_figures/fitted_heatmap/lineage_2.png and b/tests/_ground_truth_figures/fitted_heatmap/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/fitted_ignore_plot_smoothed_lineage.png b/tests/_ground_truth_figures/fitted_ignore_plot_smoothed_lineage.png
index 00d6e6501..652339956 100644
Binary files a/tests/_ground_truth_figures/fitted_ignore_plot_smoothed_lineage.png and b/tests/_ground_truth_figures/fitted_ignore_plot_smoothed_lineage.png differ
diff --git a/tests/_ground_truth_figures/fitted_model_cells_with_weights.png b/tests/_ground_truth_figures/fitted_model_cells_with_weights.png
index 535f9fa7a..fa74605f2 100644
Binary files a/tests/_ground_truth_figures/fitted_model_cells_with_weights.png and b/tests/_ground_truth_figures/fitted_model_cells_with_weights.png differ
diff --git a/tests/_ground_truth_figures/fitted_model_conf_int.png b/tests/_ground_truth_figures/fitted_model_conf_int.png
index d5d78c5c8..9e5a07b50 100644
Binary files a/tests/_ground_truth_figures/fitted_model_conf_int.png and b/tests/_ground_truth_figures/fitted_model_conf_int.png differ
diff --git a/tests/_ground_truth_figures/fitted_model_conf_int_no_conf_int_computed.png b/tests/_ground_truth_figures/fitted_model_conf_int_no_conf_int_computed.png
index 2c4136a4e..0820e9014 100644
Binary files a/tests/_ground_truth_figures/fitted_model_conf_int_no_conf_int_computed.png and b/tests/_ground_truth_figures/fitted_model_conf_int_no_conf_int_computed.png differ
diff --git a/tests/_ground_truth_figures/fitted_model_weights.png b/tests/_ground_truth_figures/fitted_model_weights.png
index 69084e3af..b333b40ac 100644
Binary files a/tests/_ground_truth_figures/fitted_model_weights.png and b/tests/_ground_truth_figures/fitted_model_weights.png differ
diff --git a/tests/_ground_truth_figures/flow_alpha.png b/tests/_ground_truth_figures/flow_alpha.png
index 0e6b91862..4b73f044c 100644
Binary files a/tests/_ground_truth_figures/flow_alpha.png and b/tests/_ground_truth_figures/flow_alpha.png differ
diff --git a/tests/_ground_truth_figures/flow_cluster_ascending.png b/tests/_ground_truth_figures/flow_cluster_ascending.png
index 8fc9cddd1..6d1a4d23f 100644
Binary files a/tests/_ground_truth_figures/flow_cluster_ascending.png and b/tests/_ground_truth_figures/flow_cluster_ascending.png differ
diff --git a/tests/_ground_truth_figures/flow_cluster_descending.png b/tests/_ground_truth_figures/flow_cluster_descending.png
index 6dfafdca8..e49251c74 100644
Binary files a/tests/_ground_truth_figures/flow_cluster_descending.png and b/tests/_ground_truth_figures/flow_cluster_descending.png differ
diff --git a/tests/_ground_truth_figures/flow_clusters_subset.png b/tests/_ground_truth_figures/flow_clusters_subset.png
index 38f20949d..ff5774bff 100644
Binary files a/tests/_ground_truth_figures/flow_clusters_subset.png and b/tests/_ground_truth_figures/flow_clusters_subset.png differ
diff --git a/tests/_ground_truth_figures/flow_explicit_cluster_order.png b/tests/_ground_truth_figures/flow_explicit_cluster_order.png
index 463213764..4dcedae1e 100644
Binary files a/tests/_ground_truth_figures/flow_explicit_cluster_order.png and b/tests/_ground_truth_figures/flow_explicit_cluster_order.png differ
diff --git a/tests/_ground_truth_figures/flow_legend_loc.png b/tests/_ground_truth_figures/flow_legend_loc.png
index 64cd1b056..36b8cb0fa 100644
Binary files a/tests/_ground_truth_figures/flow_legend_loc.png and b/tests/_ground_truth_figures/flow_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png b/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png
index 10317cd85..4e9492eab 100644
Binary files a/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png and b/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png b/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png
index 2cd45b681..ef3b8681a 100644
Binary files a/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png and b/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_no_xticks.png b/tests/_ground_truth_figures/flow_no_xticks.png
index 59321da57..cf79d276a 100644
Binary files a/tests/_ground_truth_figures/flow_no_xticks.png and b/tests/_ground_truth_figures/flow_no_xticks.png differ
diff --git a/tests/_ground_truth_figures/flow_return_ax.png b/tests/_ground_truth_figures/flow_return_ax.png
index 6dfafdca8..e49251c74 100644
Binary files a/tests/_ground_truth_figures/flow_return_ax.png and b/tests/_ground_truth_figures/flow_return_ax.png differ
diff --git a/tests/_ground_truth_figures/flow_source_clusters.png b/tests/_ground_truth_figures/flow_source_clusters.png
index ee24238d0..422cae394 100644
Binary files a/tests/_ground_truth_figures/flow_source_clusters.png and b/tests/_ground_truth_figures/flow_source_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_time_categories_too_close.png b/tests/_ground_truth_figures/flow_time_categories_too_close.png
index 8248fc56c..7500cfefc 100644
Binary files a/tests/_ground_truth_figures/flow_time_categories_too_close.png and b/tests/_ground_truth_figures/flow_time_categories_too_close.png differ
diff --git a/tests/_ground_truth_figures/gamr_ci_50.png b/tests/_ground_truth_figures/gamr_ci_50.png
index b60b9ec54..f81ebe8d5 100644
Binary files a/tests/_ground_truth_figures/gamr_ci_50.png and b/tests/_ground_truth_figures/gamr_ci_50.png differ
diff --git a/tests/_ground_truth_figures/gamr_default.png b/tests/_ground_truth_figures/gamr_default.png
index 9a978f569..02dfbc4c7 100644
Binary files a/tests/_ground_truth_figures/gamr_default.png and b/tests/_ground_truth_figures/gamr_default.png differ
diff --git a/tests/_ground_truth_figures/gamr_lineage_prob.png b/tests/_ground_truth_figures/gamr_lineage_prob.png
index 61a5bbe00..225c67399 100644
Binary files a/tests/_ground_truth_figures/gamr_lineage_prob.png and b/tests/_ground_truth_figures/gamr_lineage_prob.png differ
diff --git a/tests/_ground_truth_figures/gamr_no_cbar.png b/tests/_ground_truth_figures/gamr_no_cbar.png
index c6f15299e..31fa3615e 100644
Binary files a/tests/_ground_truth_figures/gamr_no_cbar.png and b/tests/_ground_truth_figures/gamr_no_cbar.png differ
diff --git a/tests/_ground_truth_figures/gamr_no_ci.png b/tests/_ground_truth_figures/gamr_no_ci.png
index 109564bba..02dfbc4c7 100644
Binary files a/tests/_ground_truth_figures/gamr_no_ci.png and b/tests/_ground_truth_figures/gamr_no_ci.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T.png b/tests/_ground_truth_figures/gpcca_coarse_T.png
index aa435f5c4..62f3c7a92 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T.png and b/tests/_ground_truth_figures/gpcca_coarse_T.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_cmap.png b/tests/_ground_truth_figures/gpcca_coarse_T_cmap.png
index 01da7e2eb..7ee5fbea7 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_cmap.png and b/tests/_ground_truth_figures/gpcca_coarse_T_cmap.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_init_dist.png b/tests/_ground_truth_figures/gpcca_coarse_T_init_dist.png
index fb50f23e1..951945a5b 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_init_dist.png and b/tests/_ground_truth_figures/gpcca_coarse_T_init_dist.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_no_annot.png b/tests/_ground_truth_figures/gpcca_coarse_T_no_annot.png
index c675cf453..2b0d5b72e 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_no_annot.png and b/tests/_ground_truth_figures/gpcca_coarse_T_no_annot.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_no_cbar.png b/tests/_ground_truth_figures/gpcca_coarse_T_no_cbar.png
index 40e1369f4..c763f3ade 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_no_cbar.png and b/tests/_ground_truth_figures/gpcca_coarse_T_no_cbar.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_no_order.png b/tests/_ground_truth_figures/gpcca_coarse_T_no_order.png
index 4e7e55129..5f46ad82c 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_no_order.png and b/tests/_ground_truth_figures/gpcca_coarse_T_no_order.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_stat_dist.png b/tests/_ground_truth_figures/gpcca_coarse_T_stat_dist.png
index 81865ff75..1c05a1910 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_stat_dist.png and b/tests/_ground_truth_figures/gpcca_coarse_T_stat_dist.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_stat_init_dist.png b/tests/_ground_truth_figures/gpcca_coarse_T_stat_init_dist.png
index 968d16b1b..90b2a398a 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_stat_init_dist.png and b/tests/_ground_truth_figures/gpcca_coarse_T_stat_init_dist.png differ
diff --git a/tests/_ground_truth_figures/gpcca_coarse_T_xtick_rot.png b/tests/_ground_truth_figures/gpcca_coarse_T_xtick_rot.png
index 4ab69c9de..588305daf 100644
Binary files a/tests/_ground_truth_figures/gpcca_coarse_T_xtick_rot.png and b/tests/_ground_truth_figures/gpcca_coarse_T_xtick_rot.png differ
diff --git a/tests/_ground_truth_figures/gpcca_complex_spectrum.png b/tests/_ground_truth_figures/gpcca_complex_spectrum.png
index 023d8bab3..b2d785195 100644
Binary files a/tests/_ground_truth_figures/gpcca_complex_spectrum.png and b/tests/_ground_truth_figures/gpcca_complex_spectrum.png differ
diff --git a/tests/_ground_truth_figures/gpcca_real_spectrum.png b/tests/_ground_truth_figures/gpcca_real_spectrum.png
index 14301732c..3ad51c86c 100644
Binary files a/tests/_ground_truth_figures/gpcca_real_spectrum.png and b/tests/_ground_truth_figures/gpcca_real_spectrum.png differ
diff --git a/tests/_ground_truth_figures/gpcca_real_spectrum_hide_eigengap.png b/tests/_ground_truth_figures/gpcca_real_spectrum_hide_eigengap.png
index b7ee4f9d6..a536b7453 100644
Binary files a/tests/_ground_truth_figures/gpcca_real_spectrum_hide_eigengap.png and b/tests/_ground_truth_figures/gpcca_real_spectrum_hide_eigengap.png differ
diff --git a/tests/_ground_truth_figures/gpcca_schur_matrix.png b/tests/_ground_truth_figures/gpcca_schur_matrix.png
index 20e3501d6..9c658fa9b 100644
Binary files a/tests/_ground_truth_figures/gpcca_schur_matrix.png and b/tests/_ground_truth_figures/gpcca_schur_matrix.png differ
diff --git a/tests/_ground_truth_figures/gpcca_schur_matrix_cmap.png b/tests/_ground_truth_figures/gpcca_schur_matrix_cmap.png
index 359e03753..2f9c33d09 100644
Binary files a/tests/_ground_truth_figures/gpcca_schur_matrix_cmap.png and b/tests/_ground_truth_figures/gpcca_schur_matrix_cmap.png differ
diff --git a/tests/_ground_truth_figures/gpcca_schur_matrix_title.png b/tests/_ground_truth_figures/gpcca_schur_matrix_title.png
index 27cc7662d..6bf9d6544 100644
Binary files a/tests/_ground_truth_figures/gpcca_schur_matrix_title.png and b/tests/_ground_truth_figures/gpcca_schur_matrix_title.png differ
diff --git a/tests/_ground_truth_figures/gpcca_spectrum_evals.png b/tests/_ground_truth_figures/gpcca_spectrum_evals.png
index 523789b99..0dc7560b7 100644
Binary files a/tests/_ground_truth_figures/gpcca_spectrum_evals.png and b/tests/_ground_truth_figures/gpcca_spectrum_evals.png differ
diff --git a/tests/_ground_truth_figures/gpcca_spectrum_evals_complex.png b/tests/_ground_truth_figures/gpcca_spectrum_evals_complex.png
index 64bf68a0d..9903286df 100644
Binary files a/tests/_ground_truth_figures/gpcca_spectrum_evals_complex.png and b/tests/_ground_truth_figures/gpcca_spectrum_evals_complex.png differ
diff --git a/tests/_ground_truth_figures/gpcca_spectrum_title.png b/tests/_ground_truth_figures/gpcca_spectrum_title.png
index 6021f663b..d7d1e77b0 100644
Binary files a/tests/_ground_truth_figures/gpcca_spectrum_title.png and b/tests/_ground_truth_figures/gpcca_spectrum_title.png differ
diff --git a/tests/_ground_truth_figures/graph.png b/tests/_ground_truth_figures/graph.png
index d143e2f29..42c407598 100644
Binary files a/tests/_ground_truth_figures/graph.png and b/tests/_ground_truth_figures/graph.png differ
diff --git a/tests/_ground_truth_figures/graph_bwd.png b/tests/_ground_truth_figures/graph_bwd.png
index 9d3d9f17d..13d102cf7 100644
Binary files a/tests/_ground_truth_figures/graph_bwd.png and b/tests/_ground_truth_figures/graph_bwd.png differ
diff --git a/tests/_ground_truth_figures/graph_categorical_key.png b/tests/_ground_truth_figures/graph_categorical_key.png
index 4c74becaf..f3d0ad708 100644
Binary files a/tests/_ground_truth_figures/graph_categorical_key.png and b/tests/_ground_truth_figures/graph_categorical_key.png differ
diff --git a/tests/_ground_truth_figures/graph_cmap.png b/tests/_ground_truth_figures/graph_cmap.png
index f1b3ab2e7..229bc6775 100644
Binary files a/tests/_ground_truth_figures/graph_cmap.png and b/tests/_ground_truth_figures/graph_cmap.png differ
diff --git a/tests/_ground_truth_figures/graph_curved_edges.png b/tests/_ground_truth_figures/graph_curved_edges.png
index e3f89fbb1..92ed68b5e 100644
Binary files a/tests/_ground_truth_figures/graph_curved_edges.png and b/tests/_ground_truth_figures/graph_curved_edges.png differ
diff --git a/tests/_ground_truth_figures/graph_dict_layout.png b/tests/_ground_truth_figures/graph_dict_layout.png
index 3ee2704a7..fe7594451 100644
Binary files a/tests/_ground_truth_figures/graph_dict_layout.png and b/tests/_ground_truth_figures/graph_dict_layout.png differ
diff --git a/tests/_ground_truth_figures/graph_edge_normalize.png b/tests/_ground_truth_figures/graph_edge_normalize.png
index 47f6ca1ab..93322c46b 100644
Binary files a/tests/_ground_truth_figures/graph_edge_normalize.png and b/tests/_ground_truth_figures/graph_edge_normalize.png differ
diff --git a/tests/_ground_truth_figures/graph_edge_reductions.png b/tests/_ground_truth_figures/graph_edge_reductions.png
index 6c5b82c18..414ecb537 100644
Binary files a/tests/_ground_truth_figures/graph_edge_reductions.png and b/tests/_ground_truth_figures/graph_edge_reductions.png differ
diff --git a/tests/_ground_truth_figures/graph_edge_reductions_restriction_incoming.png b/tests/_ground_truth_figures/graph_edge_reductions_restriction_incoming.png
index 3c863bdc6..21553b9de 100644
Binary files a/tests/_ground_truth_figures/graph_edge_reductions_restriction_incoming.png and b/tests/_ground_truth_figures/graph_edge_reductions_restriction_incoming.png differ
diff --git a/tests/_ground_truth_figures/graph_edge_reductions_restriction_outgoing.png b/tests/_ground_truth_figures/graph_edge_reductions_restriction_outgoing.png
index 13c3fddf9..0ca8e2e07 100644
Binary files a/tests/_ground_truth_figures/graph_edge_reductions_restriction_outgoing.png and b/tests/_ground_truth_figures/graph_edge_reductions_restriction_outgoing.png differ
diff --git a/tests/_ground_truth_figures/graph_edge_weight_scale.png b/tests/_ground_truth_figures/graph_edge_weight_scale.png
index 52be8cb99..534fca417 100644
Binary files a/tests/_ground_truth_figures/graph_edge_weight_scale.png and b/tests/_ground_truth_figures/graph_edge_weight_scale.png differ
diff --git a/tests/_ground_truth_figures/graph_filter_edges.png b/tests/_ground_truth_figures/graph_filter_edges.png
index a70dc226e..c2b78497a 100644
Binary files a/tests/_ground_truth_figures/graph_filter_edges.png and b/tests/_ground_truth_figures/graph_filter_edges.png differ
diff --git a/tests/_ground_truth_figures/graph_keys.png b/tests/_ground_truth_figures/graph_keys.png
index cd817c983..7830f2abd 100644
Binary files a/tests/_ground_truth_figures/graph_keys.png and b/tests/_ground_truth_figures/graph_keys.png differ
diff --git a/tests/_ground_truth_figures/graph_labels.png b/tests/_ground_truth_figures/graph_labels.png
index a16044d78..de2a7855f 100644
Binary files a/tests/_ground_truth_figures/graph_labels.png and b/tests/_ground_truth_figures/graph_labels.png differ
diff --git a/tests/_ground_truth_figures/graph_layout.png b/tests/_ground_truth_figures/graph_layout.png
index d143e2f29..42c407598 100644
Binary files a/tests/_ground_truth_figures/graph_layout.png and b/tests/_ground_truth_figures/graph_layout.png differ
diff --git a/tests/_ground_truth_figures/graph_networkx_layout.png b/tests/_ground_truth_figures/graph_networkx_layout.png
index 7bd01caa9..999495da9 100644
Binary files a/tests/_ground_truth_figures/graph_networkx_layout.png and b/tests/_ground_truth_figures/graph_networkx_layout.png differ
diff --git a/tests/_ground_truth_figures/graph_precomputed_layour_pca.png b/tests/_ground_truth_figures/graph_precomputed_layour_pca.png
index 3568e3907..7ed927c03 100644
Binary files a/tests/_ground_truth_figures/graph_precomputed_layour_pca.png and b/tests/_ground_truth_figures/graph_precomputed_layour_pca.png differ
diff --git a/tests/_ground_truth_figures/graph_show_arrows.png b/tests/_ground_truth_figures/graph_show_arrows.png
index d9d04fc0e..db06ab81e 100644
Binary files a/tests/_ground_truth_figures/graph_show_arrows.png and b/tests/_ground_truth_figures/graph_show_arrows.png differ
diff --git a/tests/_ground_truth_figures/graph_title.png b/tests/_ground_truth_figures/graph_title.png
index 56ec4594f..891616a62 100644
Binary files a/tests/_ground_truth_figures/graph_title.png and b/tests/_ground_truth_figures/graph_title.png differ
diff --git a/tests/_ground_truth_figures/graph_titles.png b/tests/_ground_truth_figures/graph_titles.png
index 50698b14f..79ce9b654 100644
Binary files a/tests/_ground_truth_figures/graph_titles.png and b/tests/_ground_truth_figures/graph_titles.png differ
diff --git a/tests/_ground_truth_figures/graph_top_n_edges_incoming.png b/tests/_ground_truth_figures/graph_top_n_edges_incoming.png
index abea0dde9..056c54023 100644
Binary files a/tests/_ground_truth_figures/graph_top_n_edges_incoming.png and b/tests/_ground_truth_figures/graph_top_n_edges_incoming.png differ
diff --git a/tests/_ground_truth_figures/graph_top_n_edges_outgoing.png b/tests/_ground_truth_figures/graph_top_n_edges_outgoing.png
index f4bc4ec93..587806b2b 100644
Binary files a/tests/_ground_truth_figures/graph_top_n_edges_outgoing.png and b/tests/_ground_truth_figures/graph_top_n_edges_outgoing.png differ
diff --git a/tests/_ground_truth_figures/heatmap_abs_probs_genes.png b/tests/_ground_truth_figures/heatmap_abs_probs_genes.png
index 4196d6863..3014e21b7 100644
Binary files a/tests/_ground_truth_figures/heatmap_abs_probs_genes.png and b/tests/_ground_truth_figures/heatmap_abs_probs_genes.png differ
diff --git a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_0.png b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_0.png
index 468c64d65..9c07fb5a2 100644
Binary files a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_0.png and b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_1.png b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_1.png
index 3f130b137..418321373 100644
Binary files a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_1.png and b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2-failed-diff.png b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2-failed-diff.png
new file mode 100644
index 000000000..3876cd3ec
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2.png b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2.png
index 1e66d80e9..0c46826b3 100644
Binary files a/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2.png and b/tests/_ground_truth_figures/heatmap_abs_probs_lineages/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_genes.png b/tests/_ground_truth_figures/heatmap_cluster_genes.png
index 3e865da18..e5a1a7a82 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_genes.png and b/tests/_ground_truth_figures/heatmap_cluster_genes.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_0.png b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_0.png
index 74a93ec23..03c219914 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_0.png and b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_1.png b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_1.png
index 27a9e9625..5f3f1b916 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_1.png and b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2-failed-diff.png b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2-failed-diff.png
new file mode 100644
index 000000000..78c455ff2
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2.png b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2.png
index a9e99d7d3..852c6aa2d 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2.png and b/tests/_ground_truth_figures/heatmap_cluster_key_abs_probs/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_0.png b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_0.png
index 23bfa694e..a7157c8f9 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_0.png and b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1-failed-diff.png
new file mode 100644
index 000000000..0c4ea4c89
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1.png b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1.png
index 6c99d665d..734b6422d 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1.png and b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_2.png b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_2.png
index 4aa633eeb..4f97de758 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_2.png and b/tests/_ground_truth_figures/heatmap_cluster_key_no_abs_probs/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cluster_no_scale.png b/tests/_ground_truth_figures/heatmap_cluster_no_scale.png
index 7e77b3a1f..c7fc5e4e8 100644
Binary files a/tests/_ground_truth_figures/heatmap_cluster_no_scale.png and b/tests/_ground_truth_figures/heatmap_cluster_no_scale.png differ
diff --git a/tests/_ground_truth_figures/heatmap_cmap.png b/tests/_ground_truth_figures/heatmap_cmap.png
index 242ff98a8..ce92d279b 100644
Binary files a/tests/_ground_truth_figures/heatmap_cmap.png and b/tests/_ground_truth_figures/heatmap_cmap.png differ
diff --git a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0-failed-diff.png b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0-failed-diff.png
new file mode 100644
index 000000000..aa9f78c81
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0.png b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0.png
index 30ccc293e..9ba16ff94 100644
Binary files a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0.png and b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_1.png b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_1.png
index b13d59179..8f99beab2 100644
Binary files a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_1.png and b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_2.png b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_2.png
index daeed2297..d962fc786 100644
Binary files a/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_2.png and b/tests/_ground_truth_figures/heatmap_gene_symbols/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_genes.png b/tests/_ground_truth_figures/heatmap_genes.png
index 96d2b907e..a30684839 100644
Binary files a/tests/_ground_truth_figures/heatmap_genes.png and b/tests/_ground_truth_figures/heatmap_genes.png differ
diff --git a/tests/_ground_truth_figures/heatmap_genes_1_gene_failed.png b/tests/_ground_truth_figures/heatmap_genes_1_gene_failed.png
index cf66c2be7..68792c2bb 100644
Binary files a/tests/_ground_truth_figures/heatmap_genes_1_gene_failed.png and b/tests/_ground_truth_figures/heatmap_genes_1_gene_failed.png differ
diff --git a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_0.png b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_0.png
index ae3992080..83796a20a 100644
Binary files a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_0.png and b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1-failed-diff.png
new file mode 100644
index 000000000..4a2e0537b
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1.png b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1.png
index 45dc787de..7c97f539e 100644
Binary files a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1.png and b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_2.png b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_2.png
index a35a74228..d8a871be0 100644
Binary files a/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_2.png and b/tests/_ground_truth_figures/heatmap_keep_gene_order/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_0.png b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_0.png
index ae3992080..83796a20a 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_0.png and b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1-failed-diff.png
new file mode 100644
index 000000000..207a685b3
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1.png b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1.png
index af77e2605..e4a7914bb 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1.png and b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_2.png b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_2.png
index 2ac6bca0d..b2372fdf2 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineage_height/lineage_2.png and b/tests/_ground_truth_figures/heatmap_lineage_height/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages/lineage_0.png b/tests/_ground_truth_figures/heatmap_lineages/lineage_0.png
index ae3992080..83796a20a 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages/lineage_0.png and b/tests/_ground_truth_figures/heatmap_lineages/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_lineages/lineage_1-failed-diff.png
new file mode 100644
index 000000000..207a685b3
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_lineages/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages/lineage_1.png b/tests/_ground_truth_figures/heatmap_lineages/lineage_1.png
index af77e2605..e4a7914bb 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages/lineage_1.png and b/tests/_ground_truth_figures/heatmap_lineages/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages/lineage_2.png b/tests/_ground_truth_figures/heatmap_lineages/lineage_2.png
index 2ac6bca0d..b2372fdf2 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages/lineage_2.png and b/tests/_ground_truth_figures/heatmap_lineages/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1-failed-diff.png
new file mode 100644
index 000000000..1d4ced5ce
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1.png b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1.png
index fdbe40729..7b5edd93a 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1.png and b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_2.png b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_2.png
index d2b2534fc..b2372fdf2 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_2.png and b/tests/_ground_truth_figures/heatmap_lineages_1_lineage_failed/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_0.png b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_0.png
index 346e2bdcb..9468d26a5 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_0.png and b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1-failed-diff.png
new file mode 100644
index 000000000..481216659
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1.png b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1.png
index 15db6a927..000963f3f 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1.png and b/tests/_ground_truth_figures/heatmap_lineages_bwd/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_0.png b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_0.png
index aea9ce929..7b75fa876 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_0.png and b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1-failed-diff.png
new file mode 100644
index 000000000..977e31c08
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1.png b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1.png
index 6f96d0ea0..d055dc778 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1.png and b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_2.png b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_2.png
index eee6d43fc..6c11a90d2 100644
Binary files a/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_2.png and b/tests/_ground_truth_figures/heatmap_lineages_raw/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_0.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_0.png
index 605a4573c..0d7372663 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_0.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_1.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_1.png
index 42a705845..8ade96814 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_1.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2-failed-diff.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2-failed-diff.png
new file mode 100644
index 000000000..b0150ed9a
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2.png
index 7f6f60603..f68561ef8 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_0.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_0.png
index ae3992080..83796a20a 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_0.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1-failed-diff.png
new file mode 100644
index 000000000..207a685b3
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1.png
index af77e2605..e4a7914bb 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_2.png b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_2.png
index 2ac6bca0d..b2372fdf2 100644
Binary files a/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_2.png and b/tests/_ground_truth_figures/heatmap_multiple_cluster_keys_show_all_genes/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_0.png b/tests/_ground_truth_figures/heatmap_n_jobs/lineage_0.png
deleted file mode 100644
index 605a4573c..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_0.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_1.png b/tests/_ground_truth_figures/heatmap_n_jobs/lineage_1.png
deleted file mode 100644
index 42a705845..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_1.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_2.png b/tests/_ground_truth_figures/heatmap_n_jobs/lineage_2.png
deleted file mode 100644
index 7f6f60603..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs/lineage_2.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_0.png b/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_0.png
deleted file mode 100644
index 605a4573c..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_0.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_1.png b/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_1.png
deleted file mode 100644
index 42a705845..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_1.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_2.png b/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_2.png
deleted file mode 100644
index 7f6f60603..000000000
Binary files a/tests/_ground_truth_figures/heatmap_n_jobs_multiprocessing/lineage_2.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cbar_genes.png b/tests/_ground_truth_figures/heatmap_no_cbar_genes.png
index f96f89da6..01387a599 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cbar_genes.png and b/tests/_ground_truth_figures/heatmap_no_cbar_genes.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0-failed-diff.png b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0-failed-diff.png
new file mode 100644
index 000000000..637984850
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0.png b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0.png
index cd32af1ec..38f82810f 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0.png and b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_1.png b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_1.png
index b2ceee687..c435c8269 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_1.png and b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_2.png b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_2.png
index a834afcf2..acdb0a9f1 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_2.png and b/tests/_ground_truth_figures/heatmap_no_cbar_lineages/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_0.png b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_0.png
index 96f0993ca..ffb01c4ee 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_0.png and b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1-failed-diff.png
new file mode 100644
index 000000000..40c132a77
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1.png b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1.png
index 469d8bda5..1a162ad8e 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1.png and b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_2.png b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_2.png
index 2abada4b7..f2b3132cd 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster/lineage_2.png and b/tests/_ground_truth_figures/heatmap_no_cluster/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_0.png b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_0.png
index ae3992080..83796a20a 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_0.png and b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1-failed-diff.png
new file mode 100644
index 000000000..207a685b3
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1.png b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1.png
index af77e2605..e4a7914bb 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1.png and b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_2.png b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_2.png
index 2ac6bca0d..b2372fdf2 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_2.png and b/tests/_ground_truth_figures/heatmap_no_cluster_genes/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_0.png b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_0.png
index 96f0993ca..ffb01c4ee 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_0.png and b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1-failed-diff.png
new file mode 100644
index 000000000..40c132a77
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1.png b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1.png
index 469d8bda5..1a162ad8e 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1.png and b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_2.png b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_2.png
index 2abada4b7..f2b3132cd 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_convolve/lineage_2.png and b/tests/_ground_truth_figures/heatmap_no_convolve/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_scale_genes.png b/tests/_ground_truth_figures/heatmap_no_scale_genes.png
index 585b9bd4d..e8323f871 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_scale_genes.png and b/tests/_ground_truth_figures/heatmap_no_scale_genes.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_0.png b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_0.png
index e1169ad41..c4831dfee 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_0.png and b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1-failed-diff.png
new file mode 100644
index 000000000..87c1a5c88
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1.png b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1.png
index f4a11ea71..9a916d42c 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1.png and b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_2.png b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_2.png
index 3f02f57f2..47d42fdff 100644
Binary files a/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_2.png and b/tests/_ground_truth_figures/heatmap_no_scale_lineages/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/heatmap_show_dendrogram.png b/tests/_ground_truth_figures/heatmap_show_dendrogram.png
index 3e865da18..e5a1a7a82 100644
Binary files a/tests/_ground_truth_figures/heatmap_show_dendrogram.png and b/tests/_ground_truth_figures/heatmap_show_dendrogram.png differ
diff --git a/tests/_ground_truth_figures/heatmap_time_range/lineage_0.png b/tests/_ground_truth_figures/heatmap_time_range/lineage_0.png
index 3c3a6ae46..96dcf7db5 100644
Binary files a/tests/_ground_truth_figures/heatmap_time_range/lineage_0.png and b/tests/_ground_truth_figures/heatmap_time_range/lineage_0.png differ
diff --git a/tests/_ground_truth_figures/heatmap_time_range/lineage_1-failed-diff.png b/tests/_ground_truth_figures/heatmap_time_range/lineage_1-failed-diff.png
new file mode 100644
index 000000000..c8193ec38
Binary files /dev/null and b/tests/_ground_truth_figures/heatmap_time_range/lineage_1-failed-diff.png differ
diff --git a/tests/_ground_truth_figures/heatmap_time_range/lineage_1.png b/tests/_ground_truth_figures/heatmap_time_range/lineage_1.png
index 04f4c2fc5..add3dd517 100644
Binary files a/tests/_ground_truth_figures/heatmap_time_range/lineage_1.png and b/tests/_ground_truth_figures/heatmap_time_range/lineage_1.png differ
diff --git a/tests/_ground_truth_figures/heatmap_time_range/lineage_2.png b/tests/_ground_truth_figures/heatmap_time_range/lineage_2.png
index 7862783c2..d377fe084 100644
Binary files a/tests/_ground_truth_figures/heatmap_time_range/lineage_2.png and b/tests/_ground_truth_figures/heatmap_time_range/lineage_2.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_basis.png b/tests/_ground_truth_figures/kernel_random_walk_basis.png
index e30a6ea56..3bc1e0943 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_basis.png and b/tests/_ground_truth_figures/kernel_random_walk_basis.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_cmap.png b/tests/_ground_truth_figures/kernel_random_walk_cmap.png
index 1bac74045..66cab3093 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_cmap.png and b/tests/_ground_truth_figures/kernel_random_walk_cmap.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_ixs_legend_loc.png b/tests/_ground_truth_figures/kernel_random_walk_ixs_legend_loc.png
index a4d7a072d..56f246edb 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_ixs_legend_loc.png and b/tests/_ground_truth_figures/kernel_random_walk_ixs_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_kwargs.png b/tests/_ground_truth_figures/kernel_random_walk_kwargs.png
index 10b6812dd..f021c2393 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_kwargs.png and b/tests/_ground_truth_figures/kernel_random_walk_kwargs.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_line_alpha.png b/tests/_ground_truth_figures/kernel_random_walk_line_alpha.png
index 7f511aa4f..f7bbd1e61 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_line_alpha.png and b/tests/_ground_truth_figures/kernel_random_walk_line_alpha.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_line_width.png b/tests/_ground_truth_figures/kernel_random_walk_line_width.png
index e8d18630d..317a4eff2 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_line_width.png and b/tests/_ground_truth_figures/kernel_random_walk_line_width.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_params.png b/tests/_ground_truth_figures/kernel_random_walk_params.png
index 243249944..680e6d4d1 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_params.png and b/tests/_ground_truth_figures/kernel_random_walk_params.png differ
diff --git a/tests/_ground_truth_figures/kernel_random_walk_start_ixs_range.png b/tests/_ground_truth_figures/kernel_random_walk_start_ixs_range.png
index 9a7ebbe55..a22470add 100644
Binary files a/tests/_ground_truth_figures/kernel_random_walk_start_ixs_range.png and b/tests/_ground_truth_figures/kernel_random_walk_start_ixs_range.png differ
diff --git a/tests/_ground_truth_figures/log_odds.png b/tests/_ground_truth_figures/log_odds.png
index 9d27f3074..0453246a5 100644
Binary files a/tests/_ground_truth_figures/log_odds.png and b/tests/_ground_truth_figures/log_odds.png differ
diff --git a/tests/_ground_truth_figures/log_odds_alpha.png b/tests/_ground_truth_figures/log_odds_alpha.png
index fc09919e2..119d2a122 100644
Binary files a/tests/_ground_truth_figures/log_odds_alpha.png and b/tests/_ground_truth_figures/log_odds_alpha.png differ
diff --git a/tests/_ground_truth_figures/log_odds_bwd.png b/tests/_ground_truth_figures/log_odds_bwd.png
index fd471d05f..a501112e6 100644
Binary files a/tests/_ground_truth_figures/log_odds_bwd.png and b/tests/_ground_truth_figures/log_odds_bwd.png differ
diff --git a/tests/_ground_truth_figures/log_odds_categorical_keys.png b/tests/_ground_truth_figures/log_odds_categorical_keys.png
index 37bec6768..49a374c04 100644
Binary files a/tests/_ground_truth_figures/log_odds_categorical_keys.png and b/tests/_ground_truth_figures/log_odds_categorical_keys.png differ
diff --git a/tests/_ground_truth_figures/log_odds_cmap.png b/tests/_ground_truth_figures/log_odds_cmap.png
index 72608562a..bac3477a0 100644
Binary files a/tests/_ground_truth_figures/log_odds_cmap.png and b/tests/_ground_truth_figures/log_odds_cmap.png differ
diff --git a/tests/_ground_truth_figures/log_odds_continuous_keys.png b/tests/_ground_truth_figures/log_odds_continuous_keys.png
index 592429fea..89d74cf73 100644
Binary files a/tests/_ground_truth_figures/log_odds_continuous_keys.png and b/tests/_ground_truth_figures/log_odds_continuous_keys.png differ
diff --git a/tests/_ground_truth_figures/log_odds_fontsize.png b/tests/_ground_truth_figures/log_odds_fontsize.png
index 3b08cb430..b2d9daba7 100644
Binary files a/tests/_ground_truth_figures/log_odds_fontsize.png and b/tests/_ground_truth_figures/log_odds_fontsize.png differ
diff --git a/tests/_ground_truth_figures/log_odds_jitter.png b/tests/_ground_truth_figures/log_odds_jitter.png
index 9d27f3074..0453246a5 100644
Binary files a/tests/_ground_truth_figures/log_odds_jitter.png and b/tests/_ground_truth_figures/log_odds_jitter.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs.png b/tests/_ground_truth_figures/log_odds_kwargs.png
index d0dee6b02..742b8b869 100644
Binary files a/tests/_ground_truth_figures/log_odds_kwargs.png and b/tests/_ground_truth_figures/log_odds_kwargs.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png b/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png
index f64218f89..814821fd3 100644
Binary files a/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png and b/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png b/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png
index fb93111b1..687639bfe 100644
Binary files a/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png and b/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png differ
diff --git a/tests/_ground_truth_figures/log_odds_layer.png b/tests/_ground_truth_figures/log_odds_layer.png
index c6d9cf32c..b5c86ee6b 100644
Binary files a/tests/_ground_truth_figures/log_odds_layer.png and b/tests/_ground_truth_figures/log_odds_layer.png differ
diff --git a/tests/_ground_truth_figures/log_odds_legend_loc.png b/tests/_ground_truth_figures/log_odds_legend_loc.png
index 428b2df49..6263ecb39 100644
Binary files a/tests/_ground_truth_figures/log_odds_legend_loc.png and b/tests/_ground_truth_figures/log_odds_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/log_odds_multiple_threshold.png b/tests/_ground_truth_figures/log_odds_multiple_threshold.png
index 9cd7b5eff..1acceeae1 100644
Binary files a/tests/_ground_truth_figures/log_odds_multiple_threshold.png and b/tests/_ground_truth_figures/log_odds_multiple_threshold.png differ
diff --git a/tests/_ground_truth_figures/log_odds_ncols.png b/tests/_ground_truth_figures/log_odds_ncols.png
index 3be2a680b..519c7c0c5 100644
Binary files a/tests/_ground_truth_figures/log_odds_ncols.png and b/tests/_ground_truth_figures/log_odds_ncols.png differ
diff --git a/tests/_ground_truth_figures/log_odds_rest.png b/tests/_ground_truth_figures/log_odds_rest.png
index 823083da6..921323318 100644
Binary files a/tests/_ground_truth_figures/log_odds_rest.png and b/tests/_ground_truth_figures/log_odds_rest.png differ
diff --git a/tests/_ground_truth_figures/log_odds_size.png b/tests/_ground_truth_figures/log_odds_size.png
index 4a27ab00c..456fa95cd 100644
Binary files a/tests/_ground_truth_figures/log_odds_size.png and b/tests/_ground_truth_figures/log_odds_size.png differ
diff --git a/tests/_ground_truth_figures/log_odds_threshold.png b/tests/_ground_truth_figures/log_odds_threshold.png
index 22b9dc129..4b675e3b4 100644
Binary files a/tests/_ground_truth_figures/log_odds_threshold.png and b/tests/_ground_truth_figures/log_odds_threshold.png differ
diff --git a/tests/_ground_truth_figures/log_odds_threshold_color.png b/tests/_ground_truth_figures/log_odds_threshold_color.png
index 9e2061305..3708caa09 100644
Binary files a/tests/_ground_truth_figures/log_odds_threshold_color.png and b/tests/_ground_truth_figures/log_odds_threshold_color.png differ
diff --git a/tests/_ground_truth_figures/log_odds_use_raw.png b/tests/_ground_truth_figures/log_odds_use_raw.png
index 9c4890661..787a3222e 100644
Binary files a/tests/_ground_truth_figures/log_odds_use_raw.png and b/tests/_ground_truth_figures/log_odds_use_raw.png differ
diff --git a/tests/_ground_truth_figures/log_odds_xticks_steps_size.png b/tests/_ground_truth_figures/log_odds_xticks_steps_size.png
index 9e5f03ac7..8132d8614 100644
Binary files a/tests/_ground_truth_figures/log_odds_xticks_steps_size.png and b/tests/_ground_truth_figures/log_odds_xticks_steps_size.png differ
diff --git a/tests/_ground_truth_figures/mc_complex_spectrum.png b/tests/_ground_truth_figures/mc_complex_spectrum.png
index debed89e1..fbca3d58b 100644
Binary files a/tests/_ground_truth_figures/mc_complex_spectrum.png and b/tests/_ground_truth_figures/mc_complex_spectrum.png differ
diff --git a/tests/_ground_truth_figures/mc_kwargs_linewidths.png b/tests/_ground_truth_figures/mc_kwargs_linewidths.png
index 71f2127ed..c19e5c671 100644
Binary files a/tests/_ground_truth_figures/mc_kwargs_linewidths.png and b/tests/_ground_truth_figures/mc_kwargs_linewidths.png differ
diff --git a/tests/_ground_truth_figures/mc_marker.png b/tests/_ground_truth_figures/mc_marker.png
index df973d713..546c8af86 100644
Binary files a/tests/_ground_truth_figures/mc_marker.png and b/tests/_ground_truth_figures/mc_marker.png differ
diff --git a/tests/_ground_truth_figures/mc_real_spectrum.png b/tests/_ground_truth_figures/mc_real_spectrum.png
index a375d2aeb..eb37e2f41 100644
Binary files a/tests/_ground_truth_figures/mc_real_spectrum.png and b/tests/_ground_truth_figures/mc_real_spectrum.png differ
diff --git a/tests/_ground_truth_figures/mc_real_spectrum_hide_eigengap.png b/tests/_ground_truth_figures/mc_real_spectrum_hide_eigengap.png
index 954c63b54..859f61f02 100644
Binary files a/tests/_ground_truth_figures/mc_real_spectrum_hide_eigengap.png and b/tests/_ground_truth_figures/mc_real_spectrum_hide_eigengap.png differ
diff --git a/tests/_ground_truth_figures/mc_real_spectrum_hide_xticks.png b/tests/_ground_truth_figures/mc_real_spectrum_hide_xticks.png
index a5965c919..15daa5fc3 100644
Binary files a/tests/_ground_truth_figures/mc_real_spectrum_hide_xticks.png and b/tests/_ground_truth_figures/mc_real_spectrum_hide_xticks.png differ
diff --git a/tests/_ground_truth_figures/mc_spectrum.png b/tests/_ground_truth_figures/mc_spectrum.png
index debed89e1..fbca3d58b 100644
Binary files a/tests/_ground_truth_figures/mc_spectrum.png and b/tests/_ground_truth_figures/mc_spectrum.png differ
diff --git a/tests/_ground_truth_figures/mc_spectrum_evals.png b/tests/_ground_truth_figures/mc_spectrum_evals.png
index 523789b99..0dc7560b7 100644
Binary files a/tests/_ground_truth_figures/mc_spectrum_evals.png and b/tests/_ground_truth_figures/mc_spectrum_evals.png differ
diff --git a/tests/_ground_truth_figures/mc_spectrum_evals_complex.png b/tests/_ground_truth_figures/mc_spectrum_evals_complex.png
index 64bf68a0d..9903286df 100644
Binary files a/tests/_ground_truth_figures/mc_spectrum_evals_complex.png and b/tests/_ground_truth_figures/mc_spectrum_evals_complex.png differ
diff --git a/tests/_ground_truth_figures/mc_spectrum_title.png b/tests/_ground_truth_figures/mc_spectrum_title.png
index 9d912a5ae..8b630e59c 100644
Binary files a/tests/_ground_truth_figures/mc_spectrum_title.png and b/tests/_ground_truth_figures/mc_spectrum_title.png differ
diff --git a/tests/_ground_truth_figures/mode_clustermap.png b/tests/_ground_truth_figures/mode_clustermap.png
index bcd9b5228..2724e6340 100644
Binary files a/tests/_ground_truth_figures/mode_clustermap.png and b/tests/_ground_truth_figures/mode_clustermap.png differ
diff --git a/tests/_ground_truth_figures/mode_clustermap_format.png b/tests/_ground_truth_figures/mode_clustermap_format.png
index 8686ab885..ba84abaf3 100644
Binary files a/tests/_ground_truth_figures/mode_clustermap_format.png and b/tests/_ground_truth_figures/mode_clustermap_format.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap.png b/tests/_ground_truth_figures/mode_heatmap.png
index fe29c72b0..29c82252e 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap.png and b/tests/_ground_truth_figures/mode_heatmap.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_clusters.png b/tests/_ground_truth_figures/mode_heatmap_clusters.png
index 5f7afdad0..a284cce29 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_clusters.png and b/tests/_ground_truth_figures/mode_heatmap_clusters.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_cmap.png b/tests/_ground_truth_figures/mode_heatmap_cmap.png
index a07d9778a..86b5a4634 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_cmap.png and b/tests/_ground_truth_figures/mode_heatmap_cmap.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_format.png b/tests/_ground_truth_figures/mode_heatmap_format.png
index 96fbb88f8..d132c57f9 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_format.png and b/tests/_ground_truth_figures/mode_heatmap_format.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_lineages.png b/tests/_ground_truth_figures/mode_heatmap_lineages.png
index 24d789bfe..64c1d605e 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_lineages.png and b/tests/_ground_truth_figures/mode_heatmap_lineages.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_title.png b/tests/_ground_truth_figures/mode_heatmap_title.png
index 331967ea1..079b38a17 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_title.png and b/tests/_ground_truth_figures/mode_heatmap_title.png differ
diff --git a/tests/_ground_truth_figures/mode_heatmap_xticks_rotation.png b/tests/_ground_truth_figures/mode_heatmap_xticks_rotation.png
index f4f586f4e..ec031744f 100644
Binary files a/tests/_ground_truth_figures/mode_heatmap_xticks_rotation.png and b/tests/_ground_truth_figures/mode_heatmap_xticks_rotation.png differ
diff --git a/tests/_ground_truth_figures/model_1_lineage.png b/tests/_ground_truth_figures/model_1_lineage.png
index a163a5a23..3aa2b4a5b 100644
Binary files a/tests/_ground_truth_figures/model_1_lineage.png and b/tests/_ground_truth_figures/model_1_lineage.png differ
diff --git a/tests/_ground_truth_figures/model_default.png b/tests/_ground_truth_figures/model_default.png
index 179d7ab12..18316f985 100644
Binary files a/tests/_ground_truth_figures/model_default.png and b/tests/_ground_truth_figures/model_default.png differ
diff --git a/tests/_ground_truth_figures/model_default_bwd.png b/tests/_ground_truth_figures/model_default_bwd.png
index 8f9ae46ab..a71cb181d 100644
Binary files a/tests/_ground_truth_figures/model_default_bwd.png and b/tests/_ground_truth_figures/model_default_bwd.png differ
diff --git a/tests/_ground_truth_figures/model_no_legend.png b/tests/_ground_truth_figures/model_no_legend.png
index 0d098fb70..9154ec1ef 100644
Binary files a/tests/_ground_truth_figures/model_no_legend.png and b/tests/_ground_truth_figures/model_no_legend.png differ
diff --git a/tests/_ground_truth_figures/model_no_lineage.png b/tests/_ground_truth_figures/model_no_lineage.png
index c0e440126..080ffb23e 100644
Binary files a/tests/_ground_truth_figures/model_no_lineage.png and b/tests/_ground_truth_figures/model_no_lineage.png differ
diff --git a/tests/_ground_truth_figures/model_no_lineage_show_lin_probs.png b/tests/_ground_truth_figures/model_no_lineage_show_lin_probs.png
index b2f2bcf32..d1353796c 100644
Binary files a/tests/_ground_truth_figures/model_no_lineage_show_lin_probs.png and b/tests/_ground_truth_figures/model_no_lineage_show_lin_probs.png differ
diff --git a/tests/_ground_truth_figures/model_obs_data_key.png b/tests/_ground_truth_figures/model_obs_data_key.png
index 7ea67f41b..1b007f521 100644
Binary files a/tests/_ground_truth_figures/model_obs_data_key.png and b/tests/_ground_truth_figures/model_obs_data_key.png differ
diff --git a/tests/_ground_truth_figures/model_show_lib_prob_cells_ci.png b/tests/_ground_truth_figures/model_show_lib_prob_cells_ci.png
deleted file mode 100644
index f0e50c5e2..000000000
Binary files a/tests/_ground_truth_figures/model_show_lib_prob_cells_ci.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/model_show_lin_prob_cells_ci.png b/tests/_ground_truth_figures/model_show_lin_prob_cells_ci.png
index 44f56f493..e65b51108 100644
Binary files a/tests/_ground_truth_figures/model_show_lin_prob_cells_ci.png and b/tests/_ground_truth_figures/model_show_lin_prob_cells_ci.png differ
diff --git a/tests/_ground_truth_figures/model_show_lin_prob_cells_lineage_ci.png b/tests/_ground_truth_figures/model_show_lin_prob_cells_lineage_ci.png
index 15a121629..e9c74f0a0 100644
Binary files a/tests/_ground_truth_figures/model_show_lin_prob_cells_lineage_ci.png and b/tests/_ground_truth_figures/model_show_lin_prob_cells_lineage_ci.png differ
diff --git a/tests/_ground_truth_figures/msc_default.png b/tests/_ground_truth_figures/msc_default.png
index 0981fbc9f..4e387b5aa 100644
Binary files a/tests/_ground_truth_figures/msc_default.png and b/tests/_ground_truth_figures/msc_default.png differ
diff --git a/tests/_ground_truth_figures/msc_labelrot.png b/tests/_ground_truth_figures/msc_labelrot.png
index d4417854d..514b8d34a 100644
Binary files a/tests/_ground_truth_figures/msc_labelrot.png and b/tests/_ground_truth_figures/msc_labelrot.png differ
diff --git a/tests/_ground_truth_figures/msc_legend_loc.png b/tests/_ground_truth_figures/msc_legend_loc.png
index 90771dab8..0feea8b98 100644
Binary files a/tests/_ground_truth_figures/msc_legend_loc.png and b/tests/_ground_truth_figures/msc_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/msc_title.png b/tests/_ground_truth_figures/msc_title.png
index ae7e99213..3935d922a 100644
Binary files a/tests/_ground_truth_figures/msc_title.png and b/tests/_ground_truth_figures/msc_title.png differ
diff --git a/tests/_ground_truth_figures/msc_width.png b/tests/_ground_truth_figures/msc_width.png
index 1051fb42d..2ad7b9e89 100644
Binary files a/tests/_ground_truth_figures/msc_width.png and b/tests/_ground_truth_figures/msc_width.png differ
diff --git a/tests/_ground_truth_figures/paga.png b/tests/_ground_truth_figures/paga.png
index c8023cca7..2daf4b58d 100644
Binary files a/tests/_ground_truth_figures/paga.png and b/tests/_ground_truth_figures/paga.png differ
diff --git a/tests/_ground_truth_figures/paga_lineage_subset.png b/tests/_ground_truth_figures/paga_lineage_subset.png
index 9c1cd69e8..0c65ea079 100644
Binary files a/tests/_ground_truth_figures/paga_lineage_subset.png and b/tests/_ground_truth_figures/paga_lineage_subset.png differ
diff --git a/tests/_ground_truth_figures/paga_pie.png b/tests/_ground_truth_figures/paga_pie.png
index 27dac8e04..aaab1bcb6 100644
Binary files a/tests/_ground_truth_figures/paga_pie.png and b/tests/_ground_truth_figures/paga_pie.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_embedding.png b/tests/_ground_truth_figures/paga_pie_embedding.png
index ce6c3a0ca..d5112529f 100644
Binary files a/tests/_ground_truth_figures/paga_pie_embedding.png and b/tests/_ground_truth_figures/paga_pie_embedding.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_legend_position.png b/tests/_ground_truth_figures/paga_pie_legend_position.png
index f39f71e3b..9640c4b7d 100644
Binary files a/tests/_ground_truth_figures/paga_pie_legend_position.png and b/tests/_ground_truth_figures/paga_pie_legend_position.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_legend_position_out.png b/tests/_ground_truth_figures/paga_pie_legend_position_out.png
index bba7834fb..6cbe2ea35 100644
Binary files a/tests/_ground_truth_figures/paga_pie_legend_position_out.png and b/tests/_ground_truth_figures/paga_pie_legend_position_out.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_legend_simple.png b/tests/_ground_truth_figures/paga_pie_legend_simple.png
index 3d83ac94c..55d04ec18 100644
Binary files a/tests/_ground_truth_figures/paga_pie_legend_simple.png and b/tests/_ground_truth_figures/paga_pie_legend_simple.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_no_legend.png b/tests/_ground_truth_figures/paga_pie_no_legend.png
index 668c12bdc..b2c8bf12a 100644
Binary files a/tests/_ground_truth_figures/paga_pie_no_legend.png and b/tests/_ground_truth_figures/paga_pie_no_legend.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_only_abs_prob.png b/tests/_ground_truth_figures/paga_pie_only_abs_prob.png
index ab19fd185..efeed9f84 100644
Binary files a/tests/_ground_truth_figures/paga_pie_only_abs_prob.png and b/tests/_ground_truth_figures/paga_pie_only_abs_prob.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_only_clusters.png b/tests/_ground_truth_figures/paga_pie_only_clusters.png
index 254b2bf41..0b2209c89 100644
Binary files a/tests/_ground_truth_figures/paga_pie_only_clusters.png and b/tests/_ground_truth_figures/paga_pie_only_clusters.png differ
diff --git a/tests/_ground_truth_figures/paga_pie_title.png b/tests/_ground_truth_figures/paga_pie_title.png
index 3820ff995..937dc7a3a 100644
Binary files a/tests/_ground_truth_figures/paga_pie_title.png and b/tests/_ground_truth_figures/paga_pie_title.png differ
diff --git a/tests/_ground_truth_figures/pie.png b/tests/_ground_truth_figures/pie.png
index 9c5cad17f..8c3804e71 100644
Binary files a/tests/_ground_truth_figures/pie.png and b/tests/_ground_truth_figures/pie.png differ
diff --git a/tests/_ground_truth_figures/pie_autopct_none.png b/tests/_ground_truth_figures/pie_autopct_none.png
index 6d7380093..409a7e249 100644
Binary files a/tests/_ground_truth_figures/pie_autopct_none.png and b/tests/_ground_truth_figures/pie_autopct_none.png differ
diff --git a/tests/_ground_truth_figures/pie_legend_kwargs.png b/tests/_ground_truth_figures/pie_legend_kwargs.png
index 8c38182ce..a636ff3fd 100644
Binary files a/tests/_ground_truth_figures/pie_legend_kwargs.png and b/tests/_ground_truth_figures/pie_legend_kwargs.png differ
diff --git a/tests/_ground_truth_figures/pie_legend_loc.png b/tests/_ground_truth_figures/pie_legend_loc.png
index 7beb32849..f43cf5f98 100644
Binary files a/tests/_ground_truth_figures/pie_legend_loc.png and b/tests/_ground_truth_figures/pie_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/pie_legend_loc_one.png b/tests/_ground_truth_figures/pie_legend_loc_one.png
index b67720ca5..61f8f3aa9 100644
Binary files a/tests/_ground_truth_figures/pie_legend_loc_one.png and b/tests/_ground_truth_figures/pie_legend_loc_one.png differ
diff --git a/tests/_ground_truth_figures/pie_reduction.png b/tests/_ground_truth_figures/pie_reduction.png
index df42c1028..821a6d149 100644
Binary files a/tests/_ground_truth_figures/pie_reduction.png and b/tests/_ground_truth_figures/pie_reduction.png differ
diff --git a/tests/_ground_truth_figures/pie_t.png b/tests/_ground_truth_figures/pie_t.png
index 9c5cad17f..8c3804e71 100644
Binary files a/tests/_ground_truth_figures/pie_t.png and b/tests/_ground_truth_figures/pie_t.png differ
diff --git a/tests/_ground_truth_figures/pie_title.png b/tests/_ground_truth_figures/pie_title.png
index 87636305b..7274a131a 100644
Binary files a/tests/_ground_truth_figures/pie_title.png and b/tests/_ground_truth_figures/pie_title.png differ
diff --git a/tests/_ground_truth_figures/proj_default_ordering.png b/tests/_ground_truth_figures/proj_default_ordering.png
index 271b3d153..f0e230059 100644
Binary files a/tests/_ground_truth_figures/proj_default_ordering.png and b/tests/_ground_truth_figures/proj_default_ordering.png differ
diff --git a/tests/_ground_truth_figures/proj_dont_normalize_by_mean.png b/tests/_ground_truth_figures/proj_dont_normalize_by_mean.png
index 8289957a9..deed4b767 100644
Binary files a/tests/_ground_truth_figures/proj_dont_normalize_by_mean.png and b/tests/_ground_truth_figures/proj_dont_normalize_by_mean.png differ
diff --git a/tests/_ground_truth_figures/proj_duplicate_keys.png b/tests/_ground_truth_figures/proj_duplicate_keys.png
index 16cd0c524..f0e230059 100644
Binary files a/tests/_ground_truth_figures/proj_duplicate_keys.png and b/tests/_ground_truth_figures/proj_duplicate_keys.png differ
diff --git a/tests/_ground_truth_figures/proj_extra_keys.png b/tests/_ground_truth_figures/proj_extra_keys.png
index c8a2ae9f5..89bfa234f 100644
Binary files a/tests/_ground_truth_figures/proj_extra_keys.png and b/tests/_ground_truth_figures/proj_extra_keys.png differ
diff --git a/tests/_ground_truth_figures/proj_hide_edges.png b/tests/_ground_truth_figures/proj_hide_edges.png
index 4ffc17539..e7a1ba57a 100644
Binary files a/tests/_ground_truth_figures/proj_hide_edges.png and b/tests/_ground_truth_figures/proj_hide_edges.png differ
diff --git a/tests/_ground_truth_figures/proj_key_added.png b/tests/_ground_truth_figures/proj_key_added.png
index 393822db6..6e4993d62 100644
Binary files a/tests/_ground_truth_figures/proj_key_added.png and b/tests/_ground_truth_figures/proj_key_added.png differ
diff --git a/tests/_ground_truth_figures/proj_labeldistance.png b/tests/_ground_truth_figures/proj_labeldistance.png
index 3963aec1b..61cbf2893 100644
Binary files a/tests/_ground_truth_figures/proj_labeldistance.png and b/tests/_ground_truth_figures/proj_labeldistance.png differ
diff --git a/tests/_ground_truth_figures/proj_labelrot.png b/tests/_ground_truth_figures/proj_labelrot.png
index 135c8175f..b5b9f1fdf 100644
Binary files a/tests/_ground_truth_figures/proj_labelrot.png and b/tests/_ground_truth_figures/proj_labelrot.png differ
diff --git a/tests/_ground_truth_figures/proj_ncols.png b/tests/_ground_truth_figures/proj_ncols.png
index 4965f8f82..baacb5f11 100644
Binary files a/tests/_ground_truth_figures/proj_ncols.png and b/tests/_ground_truth_figures/proj_ncols.png differ
diff --git a/tests/_ground_truth_figures/proj_no_cbar.png b/tests/_ground_truth_figures/proj_no_cbar.png
index e1d91ed44..55a81ef64 100644
Binary files a/tests/_ground_truth_figures/proj_no_cbar.png and b/tests/_ground_truth_figures/proj_no_cbar.png differ
diff --git a/tests/_ground_truth_figures/proj_scvelo_kwargs.png b/tests/_ground_truth_figures/proj_scvelo_kwargs.png
index 6b10a33bc..2607d3e31 100644
Binary files a/tests/_ground_truth_figures/proj_scvelo_kwargs.png and b/tests/_ground_truth_figures/proj_scvelo_kwargs.png differ
diff --git a/tests/_ground_truth_figures/proj_text_kwargs.png b/tests/_ground_truth_figures/proj_text_kwargs.png
index 6361fb440..723067ec2 100644
Binary files a/tests/_ground_truth_figures/proj_text_kwargs.png and b/tests/_ground_truth_figures/proj_text_kwargs.png differ
diff --git a/tests/_ground_truth_figures/proj_use_raw.png b/tests/_ground_truth_figures/proj_use_raw.png
index a9733ba65..fd9dd28ef 100644
Binary files a/tests/_ground_truth_figures/proj_use_raw.png and b/tests/_ground_truth_figures/proj_use_raw.png differ
diff --git a/tests/_ground_truth_figures/scvelo_connectivity_kernel_emb_stream.png b/tests/_ground_truth_figures/scvelo_connectivity_kernel_emb_stream.png
index 0d03fd011..7b50e5680 100644
Binary files a/tests/_ground_truth_figures/scvelo_connectivity_kernel_emb_stream.png and b/tests/_ground_truth_figures/scvelo_connectivity_kernel_emb_stream.png differ
diff --git a/tests/_ground_truth_figures/scvelo_final_states.png b/tests/_ground_truth_figures/scvelo_final_states.png
index 1e9f3d570..80874c9ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_final_states.png and b/tests/_ground_truth_figures/scvelo_final_states.png differ
diff --git a/tests/_ground_truth_figures/scvelo_final_states_clusters.png b/tests/_ground_truth_figures/scvelo_final_states_clusters.png
index 8d713ff6f..07c01fc87 100644
Binary files a/tests/_ground_truth_figures/scvelo_final_states_clusters.png and b/tests/_ground_truth_figures/scvelo_final_states_clusters.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_not_same.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_not_same.png
index 07565a45b..3996b4741 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_not_same.png and b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_not_same.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_clusters.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_clusters.png
index 50e882229..5df34c5ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_clusters.png and b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_clusters.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_no_clusters.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_no_clusters.png
index 50e882229..5df34c5ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_no_clusters.png and b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_cont_same_no_clusters.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png
index 7022de41c..93b946bdd 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png and b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png
index 04470b807..16ff03e7a 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png and b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states.png
index 371fb1cdd..3923880c6 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_cluster_key.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_cluster_key.png
index 371fb1cdd..3923880c6 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_cluster_key.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_cluster_key.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_cmap.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_cmap.png
index 662bcd7f4..e4cb6ddc9 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_cmap.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_cmap.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_discrete.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_discrete.png
index 681563288..20eea27c7 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_discrete.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_discrete.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_lineages.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_lineages.png
index bf628e0a5..8a738df64 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_lineages.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_lineages.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_no_same_plot.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_no_same_plot.png
index db56e6540..b9e3081bd 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_no_same_plot.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_no_same_plot.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_time.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_time.png
index 19e9c14e2..3a939c159 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_time.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_time.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_final_states_title.png b/tests/_ground_truth_figures/scvelo_gpcca_final_states_title.png
index 5d716566e..e50e6f2db 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_final_states_title.png and b/tests/_ground_truth_figures/scvelo_gpcca_final_states_title.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states.png
index 0240385f8..b971cfa64 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cluster_key.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cluster_key.png
index 0240385f8..b971cfa64 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cluster_key.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cluster_key.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cmap.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cmap.png
index baabe8c29..6c52ffa73 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cmap.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_cmap.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_discrete.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_discrete.png
index 0bc188108..58b0fd5bb 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_discrete.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_discrete.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_lineages.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_lineages.png
index 8d099d045..56f2961f8 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_lineages.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_lineages.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_no_same_plot.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_no_same_plot.png
index 6d4d286e8..a7d30ad61 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_no_same_plot.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_no_same_plot.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_time.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_time.png
index a404d943f..a06d35095 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_time.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_time.png differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_title.png b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_title.png
index fdf3aa904..e50e6f2db 100644
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_meta_states_title.png and b/tests/_ground_truth_figures/scvelo_gpcca_meta_states_title.png differ
diff --git a/tests/_ground_truth_figures/scvelo_initial_states_disc.png b/tests/_ground_truth_figures/scvelo_initial_states_disc.png
index 1405d11fd..5f1d9bba0 100644
Binary files a/tests/_ground_truth_figures/scvelo_initial_states_disc.png and b/tests/_ground_truth_figures/scvelo_initial_states_disc.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lin_probs.png b/tests/_ground_truth_figures/scvelo_lin_probs.png
index 141b7f083..abfc65f0f 100644
Binary files a/tests/_ground_truth_figures/scvelo_lin_probs.png and b/tests/_ground_truth_figures/scvelo_lin_probs.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lin_probs_clusters.png b/tests/_ground_truth_figures/scvelo_lin_probs_clusters.png
index 141b7f083..abfc65f0f 100644
Binary files a/tests/_ground_truth_figures/scvelo_lin_probs_clusters.png and b/tests/_ground_truth_figures/scvelo_lin_probs_clusters.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lin_probs_cmap.png b/tests/_ground_truth_figures/scvelo_lin_probs_cmap.png
index 141b7f083..abfc65f0f 100644
Binary files a/tests/_ground_truth_figures/scvelo_lin_probs_cmap.png and b/tests/_ground_truth_figures/scvelo_lin_probs_cmap.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lin_probs_lineages.png b/tests/_ground_truth_figures/scvelo_lin_probs_lineages.png
index 5625e7069..742885dd5 100644
Binary files a/tests/_ground_truth_figures/scvelo_lin_probs_lineages.png and b/tests/_ground_truth_figures/scvelo_lin_probs_lineages.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lin_probs_time.png b/tests/_ground_truth_figures/scvelo_lin_probs_time.png
index 470f6af5e..1d29bc3f7 100644
Binary files a/tests/_ground_truth_figures/scvelo_lin_probs_time.png and b/tests/_ground_truth_figures/scvelo_lin_probs_time.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lineages.png b/tests/_ground_truth_figures/scvelo_lineages.png
index 50e882229..5df34c5ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_lineages.png and b/tests/_ground_truth_figures/scvelo_lineages.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lineages_cmap.png b/tests/_ground_truth_figures/scvelo_lineages_cmap.png
index 50e882229..5df34c5ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_lineages_cmap.png and b/tests/_ground_truth_figures/scvelo_lineages_cmap.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lineages_subset.png b/tests/_ground_truth_figures/scvelo_lineages_subset.png
index 50e882229..5df34c5ee 100644
Binary files a/tests/_ground_truth_figures/scvelo_lineages_subset.png and b/tests/_ground_truth_figures/scvelo_lineages_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_lineages_time.png b/tests/_ground_truth_figures/scvelo_lineages_time.png
index 917992604..9403c7047 100644
Binary files a/tests/_ground_truth_figures/scvelo_lineages_time.png and b/tests/_ground_truth_figures/scvelo_lineages_time.png differ
diff --git a/tests/_ground_truth_figures/scvelo_pseudotime_kernel_hard_threshold_emb_stream.png b/tests/_ground_truth_figures/scvelo_pseudotime_kernel_hard_threshold_emb_stream.png
index cef09be5e..a648f911e 100644
Binary files a/tests/_ground_truth_figures/scvelo_pseudotime_kernel_hard_threshold_emb_stream.png and b/tests/_ground_truth_figures/scvelo_pseudotime_kernel_hard_threshold_emb_stream.png differ
diff --git a/tests/_ground_truth_figures/scvelo_pseudotime_kernel_soft_threshold_emb_stream.png b/tests/_ground_truth_figures/scvelo_pseudotime_kernel_soft_threshold_emb_stream.png
index 00a6b8b4e..f2ba18547 100644
Binary files a/tests/_ground_truth_figures/scvelo_pseudotime_kernel_soft_threshold_emb_stream.png and b/tests/_ground_truth_figures/scvelo_pseudotime_kernel_soft_threshold_emb_stream.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_cluster_key_discrete.png b/tests/_ground_truth_figures/scvelo_terminal_cluster_key_discrete.png
index 0d690fe57..f7e524192 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_cluster_key_discrete.png and b/tests/_ground_truth_figures/scvelo_terminal_cluster_key_discrete.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_cont_not_same_subset.png b/tests/_ground_truth_figures/scvelo_terminal_cont_not_same_subset.png
index bf628e0a5..8a738df64 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_cont_not_same_subset.png and b/tests/_ground_truth_figures/scvelo_terminal_cont_not_same_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_cont_same_subset.png b/tests/_ground_truth_figures/scvelo_terminal_cont_same_subset.png
index bf628e0a5..8a738df64 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_cont_same_subset.png and b/tests/_ground_truth_figures/scvelo_terminal_cont_same_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_diff_plot.png b/tests/_ground_truth_figures/scvelo_terminal_diff_plot.png
index db56e6540..b9e3081bd 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_diff_plot.png and b/tests/_ground_truth_figures/scvelo_terminal_diff_plot.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_diff_plot_titles.png b/tests/_ground_truth_figures/scvelo_terminal_diff_plot_titles.png
index 8ced1e61b..929e8a05d 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_diff_plot_titles.png and b/tests/_ground_truth_figures/scvelo_terminal_diff_plot_titles.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_disc_not_same_subset.png b/tests/_ground_truth_figures/scvelo_terminal_disc_not_same_subset.png
index 2331392df..7201fec18 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_disc_not_same_subset.png and b/tests/_ground_truth_figures/scvelo_terminal_disc_not_same_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_disc_same_subset.png b/tests/_ground_truth_figures/scvelo_terminal_disc_same_subset.png
index 2331392df..7201fec18 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_disc_same_subset.png and b/tests/_ground_truth_figures/scvelo_terminal_disc_same_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_states_cont.png b/tests/_ground_truth_figures/scvelo_terminal_states_cont.png
index deb5c3d2a..ae0914721 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_states_cont.png and b/tests/_ground_truth_figures/scvelo_terminal_states_cont.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_states_disc.png b/tests/_ground_truth_figures/scvelo_terminal_states_disc.png
index 89a308a3b..20eea27c7 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_states_disc.png and b/tests/_ground_truth_figures/scvelo_terminal_states_disc.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_time_mode.png b/tests/_ground_truth_figures/scvelo_terminal_time_mode.png
index 19e9c14e2..3a939c159 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_time_mode.png and b/tests/_ground_truth_figures/scvelo_terminal_time_mode.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_time_mode_clusters.png b/tests/_ground_truth_figures/scvelo_terminal_time_mode_clusters.png
index 19e9c14e2..3a939c159 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_time_mode_clusters.png and b/tests/_ground_truth_figures/scvelo_terminal_time_mode_clusters.png differ
diff --git a/tests/_ground_truth_figures/scvelo_terminal_time_mode_subset.png b/tests/_ground_truth_figures/scvelo_terminal_time_mode_subset.png
index 558a41489..9f48cfe89 100644
Binary files a/tests/_ground_truth_figures/scvelo_terminal_time_mode_subset.png and b/tests/_ground_truth_figures/scvelo_terminal_time_mode_subset.png differ
diff --git a/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png b/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png
index 53f8c2366..06867d984 100644
Binary files a/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png and b/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png differ
diff --git a/tests/_ground_truth_figures/scvelo_velocity_kernel_emb_stream.png b/tests/_ground_truth_figures/scvelo_velocity_kernel_emb_stream.png
index 2c92c8042..7e9ec0341 100644
Binary files a/tests/_ground_truth_figures/scvelo_velocity_kernel_emb_stream.png and b/tests/_ground_truth_figures/scvelo_velocity_kernel_emb_stream.png differ
diff --git a/tests/_ground_truth_figures/transpose.png b/tests/_ground_truth_figures/transpose.png
index 928cc09e9..eeff49f30 100644
Binary files a/tests/_ground_truth_figures/transpose.png and b/tests/_ground_truth_figures/transpose.png differ
diff --git a/tests/_ground_truth_figures/transpose_all_models_for_1_gene_failed.png b/tests/_ground_truth_figures/transpose_all_models_for_1_gene_failed.png
index 4c213b18a..df42b6ccf 100644
Binary files a/tests/_ground_truth_figures/transpose_all_models_for_1_gene_failed.png and b/tests/_ground_truth_figures/transpose_all_models_for_1_gene_failed.png differ
diff --git a/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed.png b/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed.png
index fdae19489..0d79e63c1 100644
Binary files a/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed.png and b/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed.png differ
diff --git a/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed_same_plot.png b/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed_same_plot.png
index 960500436..10d59c983 100644
Binary files a/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed_same_plot.png and b/tests/_ground_truth_figures/transpose_all_models_for_1_lineage_failed_same_plot.png differ
diff --git a/tests/_ground_truth_figures/transpose_failed_only_off_diagonal.png b/tests/_ground_truth_figures/transpose_failed_only_off_diagonal.png
index 44bfea289..7fde11353 100644
Binary files a/tests/_ground_truth_figures/transpose_failed_only_off_diagonal.png and b/tests/_ground_truth_figures/transpose_failed_only_off_diagonal.png differ
diff --git a/tests/_ground_truth_figures/transpose_same_plot.png b/tests/_ground_truth_figures/transpose_same_plot.png
index e031cf9d0..87d5691bd 100644
Binary files a/tests/_ground_truth_figures/transpose_same_plot.png and b/tests/_ground_truth_figures/transpose_same_plot.png differ
diff --git a/tests/_ground_truth_figures/trends.png b/tests/_ground_truth_figures/trends.png
index 209fa368b..02859eb5a 100644
Binary files a/tests/_ground_truth_figures/trends.png and b/tests/_ground_truth_figures/trends.png differ
diff --git a/tests/_ground_truth_figures/trends_abs_prob_cmap.png b/tests/_ground_truth_figures/trends_abs_prob_cmap.png
index c025ac205..a31213f86 100644
Binary files a/tests/_ground_truth_figures/trends_abs_prob_cmap.png and b/tests/_ground_truth_figures/trends_abs_prob_cmap.png differ
diff --git a/tests/_ground_truth_figures/trends_bwd.png b/tests/_ground_truth_figures/trends_bwd.png
index 5f0dffe6c..e2fccdba5 100644
Binary files a/tests/_ground_truth_figures/trends_bwd.png and b/tests/_ground_truth_figures/trends_bwd.png differ
diff --git a/tests/_ground_truth_figures/trends_cell_alpha.png b/tests/_ground_truth_figures/trends_cell_alpha.png
index b6e2aea02..74a0d3799 100644
Binary files a/tests/_ground_truth_figures/trends_cell_alpha.png and b/tests/_ground_truth_figures/trends_cell_alpha.png differ
diff --git a/tests/_ground_truth_figures/trends_conf_int.png b/tests/_ground_truth_figures/trends_conf_int.png
index aca784fb3..c25864f31 100644
Binary files a/tests/_ground_truth_figures/trends_conf_int.png and b/tests/_ground_truth_figures/trends_conf_int.png differ
diff --git a/tests/_ground_truth_figures/trends_gam_ci_100.png b/tests/_ground_truth_figures/trends_gam_ci_100.png
index 748641d29..4ad53ca2b 100644
Binary files a/tests/_ground_truth_figures/trends_gam_ci_100.png and b/tests/_ground_truth_figures/trends_gam_ci_100.png differ
diff --git a/tests/_ground_truth_figures/trends_gam_ci_20.png b/tests/_ground_truth_figures/trends_gam_ci_20.png
index 4cb67b2ac..7431a6347 100644
Binary files a/tests/_ground_truth_figures/trends_gam_ci_20.png and b/tests/_ground_truth_figures/trends_gam_ci_20.png differ
diff --git a/tests/_ground_truth_figures/trends_gam_ci_80.png b/tests/_ground_truth_figures/trends_gam_ci_80.png
deleted file mode 100644
index 4bbeb6817..000000000
Binary files a/tests/_ground_truth_figures/trends_gam_ci_80.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/trends_gene_as_title.png b/tests/_ground_truth_figures/trends_gene_as_title.png
index 1df09fbd0..bf74b1d73 100644
Binary files a/tests/_ground_truth_figures/trends_gene_as_title.png and b/tests/_ground_truth_figures/trends_gene_as_title.png differ
diff --git a/tests/_ground_truth_figures/trends_gene_legend_out.png b/tests/_ground_truth_figures/trends_gene_legend_out.png
index ee582ba21..c72453c5d 100644
Binary files a/tests/_ground_truth_figures/trends_gene_legend_out.png and b/tests/_ground_truth_figures/trends_gene_legend_out.png differ
diff --git a/tests/_ground_truth_figures/trends_gene_no_legend.png b/tests/_ground_truth_figures/trends_gene_no_legend.png
index 390ea8fe2..37a018ad4 100644
Binary files a/tests/_ground_truth_figures/trends_gene_no_legend.png and b/tests/_ground_truth_figures/trends_gene_no_legend.png differ
diff --git a/tests/_ground_truth_figures/trends_gene_symbols.png b/tests/_ground_truth_figures/trends_gene_symbols.png
index d84136315..6154b7481 100644
Binary files a/tests/_ground_truth_figures/trends_gene_symbols.png and b/tests/_ground_truth_figures/trends_gene_symbols.png differ
diff --git a/tests/_ground_truth_figures/trends_hide_cells.png b/tests/_ground_truth_figures/trends_hide_cells.png
index 47b276504..be6fb5591 100644
Binary files a/tests/_ground_truth_figures/trends_hide_cells.png and b/tests/_ground_truth_figures/trends_hide_cells.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_alpha.png b/tests/_ground_truth_figures/trends_lineage_alpha.png
index 96e98d307..bd8e0bd0d 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_alpha.png and b/tests/_ground_truth_figures/trends_lineage_alpha.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_cell_color.png b/tests/_ground_truth_figures/trends_lineage_cell_color.png
index 98735c304..bb982bcb0 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_cell_color.png and b/tests/_ground_truth_figures/trends_lineage_cell_color.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_cell_color_clusters.png b/tests/_ground_truth_figures/trends_lineage_cell_color_clusters.png
index 18c288483..1908f258f 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_cell_color_clusters.png and b/tests/_ground_truth_figures/trends_lineage_cell_color_clusters.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_cell_color_clusters_obs_legend_loc.png b/tests/_ground_truth_figures/trends_lineage_cell_color_clusters_obs_legend_loc.png
index ca84de339..0ce0a15d3 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_cell_color_clusters_obs_legend_loc.png and b/tests/_ground_truth_figures/trends_lineage_cell_color_clusters_obs_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_cell_color_gene.png b/tests/_ground_truth_figures/trends_lineage_cell_color_gene.png
index 73c4d3a1e..65a1d919e 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_cell_color_gene.png and b/tests/_ground_truth_figures/trends_lineage_cell_color_gene.png differ
diff --git a/tests/_ground_truth_figures/trends_lineage_cmap.png b/tests/_ground_truth_figures/trends_lineage_cmap.png
index 72adb6310..bcacae9de 100644
Binary files a/tests/_ground_truth_figures/trends_lineage_cmap.png and b/tests/_ground_truth_figures/trends_lineage_cmap.png differ
diff --git a/tests/_ground_truth_figures/trends_lw.png b/tests/_ground_truth_figures/trends_lw.png
index 7189ab800..d749fd4ec 100644
Binary files a/tests/_ground_truth_figures/trends_lw.png and b/tests/_ground_truth_figures/trends_lw.png differ
diff --git a/tests/_ground_truth_figures/trends_margins.png b/tests/_ground_truth_figures/trends_margins.png
index 494055a6f..acfe14123 100644
Binary files a/tests/_ground_truth_figures/trends_margins.png and b/tests/_ground_truth_figures/trends_margins.png differ
diff --git a/tests/_ground_truth_figures/trends_no_cbar.png b/tests/_ground_truth_figures/trends_no_cbar.png
index f1eb0e353..ff7111e7d 100644
Binary files a/tests/_ground_truth_figures/trends_no_cbar.png and b/tests/_ground_truth_figures/trends_no_cbar.png differ
diff --git a/tests/_ground_truth_figures/trends_perc.png b/tests/_ground_truth_figures/trends_perc.png
index 390ea8fe2..37a018ad4 100644
Binary files a/tests/_ground_truth_figures/trends_perc.png and b/tests/_ground_truth_figures/trends_perc.png differ
diff --git a/tests/_ground_truth_figures/trends_perc_per_lineage.png b/tests/_ground_truth_figures/trends_perc_per_lineage.png
index 987d79be2..065fac277 100644
Binary files a/tests/_ground_truth_figures/trends_perc_per_lineage.png and b/tests/_ground_truth_figures/trends_perc_per_lineage.png differ
diff --git a/tests/_ground_truth_figures/trends_raw.png b/tests/_ground_truth_figures/trends_raw.png
index 56c020119..757691de6 100644
Binary files a/tests/_ground_truth_figures/trends_raw.png and b/tests/_ground_truth_figures/trends_raw.png differ
diff --git a/tests/_ground_truth_figures/trends_same_plot.png b/tests/_ground_truth_figures/trends_same_plot.png
index f395c46e9..b52f4135c 100644
Binary files a/tests/_ground_truth_figures/trends_same_plot.png and b/tests/_ground_truth_figures/trends_same_plot.png differ
diff --git a/tests/_ground_truth_figures/trends_sharex.png b/tests/_ground_truth_figures/trends_sharex.png
index 390ea8fe2..37a018ad4 100644
Binary files a/tests/_ground_truth_figures/trends_sharex.png and b/tests/_ground_truth_figures/trends_sharex.png differ
diff --git a/tests/_ground_truth_figures/trends_sharey.png b/tests/_ground_truth_figures/trends_sharey.png
index 209fa368b..02859eb5a 100644
Binary files a/tests/_ground_truth_figures/trends_sharey.png and b/tests/_ground_truth_figures/trends_sharey.png differ
diff --git a/tests/_ground_truth_figures/trends_show_lineage_ci.png b/tests/_ground_truth_figures/trends_show_lineage_ci.png
index 8d828bee2..c4671bb77 100644
Binary files a/tests/_ground_truth_figures/trends_show_lineage_ci.png and b/tests/_ground_truth_figures/trends_show_lineage_ci.png differ
diff --git a/tests/_ground_truth_figures/trends_show_lineage_diff_plot.png b/tests/_ground_truth_figures/trends_show_lineage_diff_plot.png
index 96c126a14..b1a87d86b 100644
Binary files a/tests/_ground_truth_figures/trends_show_lineage_diff_plot.png and b/tests/_ground_truth_figures/trends_show_lineage_diff_plot.png differ
diff --git a/tests/_ground_truth_figures/trends_show_lineage_ignores_no_transpose.png b/tests/_ground_truth_figures/trends_show_lineage_ignores_no_transpose.png
index 2ec438a7c..60287e291 100644
Binary files a/tests/_ground_truth_figures/trends_show_lineage_ignores_no_transpose.png and b/tests/_ground_truth_figures/trends_show_lineage_ignores_no_transpose.png differ
diff --git a/tests/_ground_truth_figures/trends_show_lineage_same_plot.png b/tests/_ground_truth_figures/trends_show_lineage_same_plot.png
index 119a38815..88bf2417a 100644
Binary files a/tests/_ground_truth_figures/trends_show_lineage_same_plot.png and b/tests/_ground_truth_figures/trends_show_lineage_same_plot.png differ
diff --git a/tests/_ground_truth_figures/trends_size.png b/tests/_ground_truth_figures/trends_size.png
index e96a58bb8..9a1e4e36f 100644
Binary files a/tests/_ground_truth_figures/trends_size.png and b/tests/_ground_truth_figures/trends_size.png differ
diff --git a/tests/_ground_truth_figures/trends_suptitle.png b/tests/_ground_truth_figures/trends_suptitle.png
index b1cfe54fd..95d2cf063 100644
Binary files a/tests/_ground_truth_figures/trends_suptitle.png and b/tests/_ground_truth_figures/trends_suptitle.png differ
diff --git a/tests/_ground_truth_figures/trends_time_key.png b/tests/_ground_truth_figures/trends_time_key.png
index 21d23c5ca..8e80607cc 100644
Binary files a/tests/_ground_truth_figures/trends_time_key.png and b/tests/_ground_truth_figures/trends_time_key.png differ
diff --git a/tests/_ground_truth_figures/trends_time_key_del_latent_time.png b/tests/_ground_truth_figures/trends_time_key_del_latent_time.png
index 21d23c5ca..8e80607cc 100644
Binary files a/tests/_ground_truth_figures/trends_time_key_del_latent_time.png and b/tests/_ground_truth_figures/trends_time_key_del_latent_time.png differ
diff --git a/tests/_ground_truth_figures/trends_time_range.png b/tests/_ground_truth_figures/trends_time_range.png
index 1599a1c32..7df2e3d4c 100644
Binary files a/tests/_ground_truth_figures/trends_time_range.png and b/tests/_ground_truth_figures/trends_time_range.png differ
diff --git a/tests/_ground_truth_figures/violin.png b/tests/_ground_truth_figures/violin.png
index c311815ac..2ff04c8c4 100644
Binary files a/tests/_ground_truth_figures/violin.png and b/tests/_ground_truth_figures/violin.png differ
diff --git a/tests/_ground_truth_figures/violin_cluster_subset.png b/tests/_ground_truth_figures/violin_cluster_subset.png
index 28ca01037..2ff04c8c4 100644
Binary files a/tests/_ground_truth_figures/violin_cluster_subset.png and b/tests/_ground_truth_figures/violin_cluster_subset.png differ
diff --git a/tests/_ground_truth_figures/violin_lineage_subset.png b/tests/_ground_truth_figures/violin_lineage_subset.png
index f78cb44ba..3acf187a6 100644
Binary files a/tests/_ground_truth_figures/violin_lineage_subset.png and b/tests/_ground_truth_figures/violin_lineage_subset.png differ
diff --git a/tests/_ground_truth_figures/violin_no_cluster_key.png b/tests/_ground_truth_figures/violin_no_cluster_key.png
index 24a9fd506..39862640d 100644
Binary files a/tests/_ground_truth_figures/violin_no_cluster_key.png and b/tests/_ground_truth_figures/violin_no_cluster_key.png differ
diff --git a/tests/conftest.py b/tests/conftest.py
index 0624423ae..54cc0e501 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -58,7 +58,7 @@ def _create_cflare(*, backward: bool = False) -> Tuple[AnnData, CFLARE]:
     vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
         softmax_scale=4
     )
-    ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix()
+    ck = ConnectivityKernel(adata).compute_transition_matrix()
     final_kernel = 0.8 * vk + 0.2 * ck
 
     mc = CFLARE(final_kernel)
@@ -83,7 +83,7 @@ def _create_gpcca(*, backward: bool = False) -> Tuple[AnnData, GPCCA]:
     vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
         softmax_scale=4
     )
-    ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix()
+    ck = ConnectivityKernel(adata).compute_transition_matrix()
     final_kernel = 0.8 * vk + 0.2 * ck
 
     mc = GPCCA(final_kernel)
diff --git a/tests/test_external.py b/tests/test_external.py
index 7c54a623b..ae3522761 100644
--- a/tests/test_external.py
+++ b/tests/test_external.py
@@ -7,7 +7,7 @@
 from anndata import AnnData
 from cellrank.tl.kernels import ConnectivityKernel
 from cellrank.external.kernels._utils import MarkerGenes
-from cellrank.external.kernels._wot_kernel import LastTimePoint
+from cellrank.tl.kernels._transport_map_kernel import LastTimePoint
 
 import numpy as np
 import pandas as pd
@@ -47,22 +47,6 @@ def _statot_not_installed() -> bool:
 
 @statot_not_installed_skip
 class TestOTKernel:
-    def test_no_connectivities(self, adata_large: AnnData):
-        del adata_large.obsp["connectivities"]
-        terminal_states = np.full((adata_large.n_obs,), fill_value=np.nan, dtype=object)
-        ixs = np.where(adata_large.obs["clusters"] == "Granule immature")[0]
-        terminal_states[ixs] = "GI"
-
-        ok = cre.kernels.StationaryOTKernel(
-            adata_large,
-            terminal_states=pd.Series(terminal_states).astype("category"),
-            g=np.ones((adata_large.n_obs,), dtype=np.float64),
-        )
-        ok = ok.compute_transition_matrix(1, 0.001)
-
-        assert ok._conn is None
-        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
-
     def test_method_not_implemented(self, adata_large: AnnData):
         terminal_states = np.full((adata_large.n_obs,), fill_value=np.nan, dtype=object)
         ixs = np.where(adata_large.obs["clusters"] == "Granule immature")[0]
@@ -125,25 +109,14 @@ def test_compute_projection(
         else:
             combined_kernel = ok
 
-        expected_error = (
-            r"<StationaryOTKernel> is not a kNN based kernel. The embedding projection "
-            r"only works for kNN based kernels."
-        )
-        with pytest.raises(AttributeError, match=expected_error):
-            combined_kernel.compute_projection()
+        with pytest.raises(
+            TypeError, match=r"StationaryOTKernel is not a kNN based kernel"
+        ):
+            combined_kernel.plot_projection()
 
 
 @wot_not_installed_skip
 class TestWOTKernel:
-    def test_no_connectivities(self, adata_large: AnnData):
-        del adata_large.obsp["connectivities"]
-        ok = cre.kernels.WOTKernel(
-            adata_large, time_key="age(days)"
-        ).compute_transition_matrix()
-
-        assert ok._conn is None
-        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
-
     def test_invalid_solver_kwargs(self, adata_large: AnnData):
         ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
         with pytest.raises(TypeError, match="unexpected keyword argument 'foo'"):
@@ -159,6 +132,9 @@ def test_inversion_updates_adata(self, adata_large: AnnData):
 
         ok = ~ok
 
+        assert ok.transition_matrix is None
+        assert ok.transport_maps is None
+
         inverted_time = ok.experimental_time
         assert is_categorical_dtype(adata_large.obs[key])
         assert adata_large.obs[key].cat.ordered
@@ -292,6 +268,8 @@ def test_normal_run(self, adata_large: AnnData):
         assert isinstance(ok.params, dict)
         assert isinstance(ok.growth_rates, pd.DataFrame)
         assert isinstance(ok.transport_maps, dict)
+        for v in ok.transport_maps.values():
+            assert isinstance(v, AnnData)
         np.testing.assert_array_equal(adata_large.obs_names, ok.growth_rates.index)
         np.testing.assert_array_equal(ok.growth_rates.columns, ["g0", "g1"])
         assert isinstance(ok.transport_maps[12.0, 35.0], AnnData)
@@ -334,12 +312,8 @@ def test_compute_projection(
         else:
             combined_kernel = ok
 
-        expected_error = (
-            r"<WOTKernel> is not a kNN based kernel. The embedding projection only "
-            r"works for kNN based kernels."
-        )
-        with pytest.raises(AttributeError, match=expected_error):
-            combined_kernel.compute_projection()
+        with pytest.raises(TypeError, match=r"WOTKernel is not a kNN based kernel"):
+            combined_kernel.plot_projection()
 
     def test_wot_write(self, adata_large: AnnData, tmpdir):
         path = str(tmpdir / "wot.pickle")
diff --git a/tests/test_gpcca.py b/tests/test_gpcca.py
index 591f452f8..0fbe02644 100644
--- a/tests/test_gpcca.py
+++ b/tests/test_gpcca.py
@@ -223,7 +223,7 @@ def _fit_gpcca(adata, state: str, backward: bool = False) -> cr.estimators.GPCCA
     vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
         softmax_scale=4
     )
-    ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix()
+    ck = ConnectivityKernel(adata).compute_transition_matrix()
     terminal_kernel = 0.8 * vk + 0.2 * ck
 
     mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -505,7 +505,7 @@ def test_set_initial_states_from_forward(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -530,7 +530,7 @@ def test_compute_initial_states_from_forward_no_macro(self, adata_large: AnnData
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -543,10 +543,8 @@ def test_compute_initial_states_from_forward_no_macro(self, adata_large: AnnData
     def test_compute_initial_states_from_forward_too_many_states(
         self, adata_large: AnnData
     ):
-        vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
-            softmax_scale=4
-        )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -562,7 +560,7 @@ def test_compute_initial_states_from_forward_too_few_states(
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -578,7 +576,7 @@ def test_compute_initial_states_from_forward_no_stat_dist(
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
@@ -594,7 +592,7 @@ def test_compute_initial_states_from_forward_normal_run(self, adata_large: AnnDa
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata_large, backward=False).compute_transition_matrix()
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index bfbdf80a9..406250aef 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -27,13 +27,13 @@
     ConnectivityKernel,
     TransportMapKernel,
 )
+from cellrank.tl.kernels._mixins import ConnectivityMixin
 from cellrank.tl.kernels._base_kernel import (
     Kernel,
     Constant,
     KernelAdd,
     KernelMul,
-    _dtype,
-    _is_bin_mult,
+    UnidirectionalKernel,
 )
 from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEAggregation
 
@@ -47,7 +47,7 @@
 _rtol = 1e-6
 
 
-class CustomFunc(cr.tl.kernels.SimilaritySchemeABC):
+class CustomFunc(cr.tl.kernels.utils.SimilarityABC):
     def __call__(
         self, v: np.ndarray, D: np.ndarray, softmax_scale: float = 1.0
     ) -> Tuple[np.ndarray, np.ndarray]:
@@ -67,23 +67,21 @@ def hessian(
         return np.zeros((D.shape[0], v.shape[0], v.shape[0]))
 
 
-class CustomKernel(Kernel):
-    def compute_transition_matrix(
-        self, sparse: bool = False, dnorm: bool = False
-    ) -> "KernelExpression":
+class CustomKernel(UnidirectionalKernel):
+    def compute_transition_matrix(self, sparse: bool = False) -> "CustomKernel":
         if sparse:
             tmat = speye(self.adata.n_obs, dtype=np.float32)
         else:
             tmat = np.eye(self.adata.n_obs, dtype=np.float32)
 
-        self._compute_transition_matrix(tmat, density_normalize=dnorm)
+        self.transition_matrix = tmat
         return self
 
-    def copy(self) -> "KernelExpression":
+    def copy(self, deep: bool = False) -> "CustomKernel":
         return copy(self)
 
 
-class InvalidFuncProbs(cr.tl.kernels.SimilaritySchemeABC):
+class InvalidFuncProbs(cr.tl.kernels.utils.SimilarityABC):
     def __call__(
         self, v: np.ndarray, D: np.ndarray, _softmax_scale: float = 1.0
     ) -> Tuple[np.ndarray, np.ndarray]:
@@ -144,7 +142,9 @@ def _compute_tmap(
         tmat = sprandom(
             n, m, density=0.5, dtype=np.float64, format="csr", random_state=42
         )
+        tmat[:, 0] = 1e-3
         tmat.data[:] = np.abs(tmat.data)
+
         if dtype is spmatrix:
             return tmat
         if dtype is np.ndarray:
@@ -183,25 +183,25 @@ def test_not_none_transition_matrix_compute(self, adata: AnnData):
         assert ck.transition_matrix is not None
         assert pk.transition_matrix is not None
 
-    def test_not_none_transition_matrix_accessor(self, adata: AnnData):
+    def test_none_transition_matrix_accessor(self, adata: AnnData):
         vk = VelocityKernel(adata)
         ck = ConnectivityKernel(adata)
         pk = PseudotimeKernel(adata, time_key="latent_time")
 
-        assert vk.transition_matrix is not None
-        assert ck.transition_matrix is not None
-        assert pk.transition_matrix is not None
+        assert vk.transition_matrix is None
+        assert ck.transition_matrix is None
+        assert pk.transition_matrix is None
 
     def test_adding_hidden_constants(self, adata: AnnData):
         k = VelocityKernel(adata) + ConnectivityKernel(adata)
 
-        assert _is_bin_mult(k[0])
+        assert k[0]._bin_consts
         assert isinstance(k[0], KernelMul)
         assert isinstance(k[0][0], Constant)
         assert isinstance(k[0][1], VelocityKernel)
         assert k[0][0].transition_matrix == 1.0
 
-        assert _is_bin_mult(k[1])
+        assert k[1]._bin_consts
         assert isinstance(k[1], KernelMul)
         assert isinstance(k[1][0], Constant)
         assert isinstance(k[1][1], ConnectivityKernel)
@@ -220,10 +220,16 @@ def test_accessor_out_of_range(self, adata: AnnData):
     def test_parent(self, adata: AnnData):
         vk = VelocityKernel(adata)
         ck = ConnectivityKernel(adata)
-        k = vk + ck
+        k_old = vk + ck
+        k = k_old.copy()
+
+        assert vk._parent._parent is k_old
+        assert ck._parent._parent is k_old
+        assert k_old._parent is None
 
-        assert vk._parent._parent is k  # invisible constants
-        assert ck._parent._parent is k
+        k1, k2 = k.kernels
+        assert k1._parent._parent is k
+        assert k2._parent._parent is k
         assert k._parent is None
 
     def test_uninitialized_both(self, adata: AnnData):
@@ -263,44 +269,45 @@ def test_invalid_constant(self, adata: AnnData):
             _ = Constant(adata, None)
 
     def test_inversion(self, adata: AnnData):
-        c = ConnectivityKernel(adata, backward=False)
-        assert not c.backward
+        c = ConnectivityKernel(adata)
+        assert c.backward is None
 
-        nc = ~c
-        assert nc.backward
-
-    def test_inversion_inplace(self, adata: AnnData):
-        c = ConnectivityKernel(adata, backward=False)
+    def test_inversion_not_inplace(self, adata: AnnData):
+        c = VelocityKernel(adata, backward=False)
 
         assert not c.backward
-        _ = ~c
-        assert c.backward
+        k = ~c
+        assert k is not c
+        assert k.adata is c.adata
+        assert not c.backward
+        assert k.backward
 
     def test_inversion_propagation(self, adata: AnnData):
-        c = ConnectivityKernel(adata, backward=False)
+        c = ConnectivityKernel(adata)
         v = VelocityKernel(adata, backward=False)
         k = ~(c + v)
 
-        assert c.backward
-        assert v.backward
+        assert not c.backward
+        assert not v.backward
         assert k.backward
 
     def test_inversion_recalculation(self, adata: AnnData):
         c = ConnectivityKernel(adata).compute_transition_matrix()
         z = ~(c + c)
-        with pytest.raises(RuntimeError):
-            z.compute_transition_matrix()
 
-    def test_inversion_preservation_of_constants(self, adata: AnnData):
+        assert z._transition_matrix is None
+        # since Connectivities are directionless, this is allowed
+        assert z.transition_matrix is not None
+
+    def test_inversion_no_recalculation(self, adata: AnnData):
         c = ConnectivityKernel(adata).compute_transition_matrix()
-        a = (3 * c + 1 * c).compute_transition_matrix()
-        b = ~a
-        c.compute_transition_matrix()
+        v = VelocityKernel(adata).compute_transition_matrix(softmax_scale=1)
+        z = ~(v + c)
 
-        assert a[0][0].transition_matrix == 3 / 4
-        assert b[0][0].transition_matrix == 3 / 4
-        assert a[1][0].transition_matrix == 1 / 4
-        assert b[1][0].transition_matrix == 1 / 4
+        assert z._transition_matrix is None
+        with pytest.raises(RuntimeError, match=f"`{~v}` is uninitialized"):
+            # not allowed because VK is reset
+            assert z.transition_matrix is not None
 
     def test_addition_simple(self, adata: AnnData):
         k = VelocityKernel(adata) + ConnectivityKernel(adata)
@@ -309,20 +316,30 @@ def test_addition_simple(self, adata: AnnData):
 
     def test_multiplication_simple(self, adata: AnnData):
         k = 10 * VelocityKernel(adata)
-        c = _is_bin_mult(k)
+        c, _ = k._split_const
 
         assert isinstance(c, Constant)
         assert c.transition_matrix == 10
 
     def test_multiplication_simple_normalization(self, adata: AnnData):
         k = 10 * VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
-        c = _is_bin_mult(k)
+        c, _ = k._split_const
 
         assert c.transition_matrix == 10
 
+    def test_kernel_kernel_multiplication(self, adata: AnnData):
+        vk = VelocityKernel(adata).compute_transition_matrix()
+        ck = ConnectivityKernel(adata).compute_transition_matrix()
+
+        actual = (vk * ck).transition_matrix
+        expected = _normalize(vk.transition_matrix.A * ck.transition_matrix.A)
+
+        np.testing.assert_allclose(actual.A, expected)
+
     def test_constant(self, adata: AnnData):
         k = 9 * VelocityKernel(adata) + 1 * ConnectivityKernel(adata)
-        c1, c2 = _is_bin_mult(k[0]), _is_bin_mult(k[1])
+        c1, _ = k[0]._split_const
+        c2, _ = k[1]._split_const
 
         assert c1.transition_matrix == 9
         assert c2.transition_matrix == 1
@@ -333,7 +350,8 @@ def test_constant_normalize_2(self, adata: AnnData):
             + 1 * ConnectivityKernel(adata).compute_transition_matrix()
         )
         k.compute_transition_matrix()
-        c1, c2 = _is_bin_mult(k[0]), _is_bin_mult(k[1])
+        c1, _ = k[0]._split_const
+        c2, _ = k[1]._split_const
 
         assert c1.transition_matrix == 9 / 10
         assert c2.transition_matrix == 1 / 10
@@ -345,7 +363,9 @@ def test_constant_normalize_3(self, adata: AnnData):
             + ConnectivityKernel(adata).compute_transition_matrix()
         )
         k.compute_transition_matrix()
-        c1, c2, c3 = _is_bin_mult(k[0]), _is_bin_mult(k[1]), _is_bin_mult(k[2])
+        c1, _ = k[0]._split_const
+        c2, _ = k[1]._split_const
+        c3, _ = k[2]._split_const
 
         assert c1.transition_matrix == 1 / 3
         assert c2.transition_matrix == 1 / 3
@@ -357,7 +377,9 @@ def test_constant_wrong_parentheses(self, adata: AnnData):
             + ConnectivityKernel(adata).compute_transition_matrix()
         )
         k.compute_transition_matrix()
-        c1, c2, c3 = _is_bin_mult(k[0]), _is_bin_mult(k[1]), _is_bin_mult(k[2])
+        c1, _ = k[0]._split_const
+        c2, _ = k[1]._split_const
+        c3, _ = k[2]._split_const
 
         assert c1.transition_matrix == 1 / 3
         assert c2.transition_matrix == 1 / 3
@@ -369,97 +391,49 @@ def test_constant_correct_parentheses(self, adata: AnnData):
             + ConnectivityKernel(adata).compute_transition_matrix()
         )
         k.compute_transition_matrix()
-        c1, c2, c3 = (
-            _is_bin_mult(k[0]),
-            _is_bin_mult(k[1][1][0]),
-            _is_bin_mult(k[1][1][1]),
-        )
+        c1, _ = k[0]._split_const
+        c2, _ = k[1][1][0]._split_const
+        c3, _ = k[1][1][1]._split_const
 
         assert c1.transition_matrix == 1 / 2
         assert c2.transition_matrix == 1 / 2
         assert c3.transition_matrix == 1 / 2
 
-    def test_adaptive_kernel_constants(self, adata: AnnData):
-        ck1 = ConnectivityKernel(adata).compute_transition_matrix()
-        ck1._mat_scaler = np.random.normal(size=(adata.n_obs, adata.n_obs))
-
-        ck2 = ConnectivityKernel(adata).compute_transition_matrix()
-        ck2._mat_scaler = np.random.normal(size=(adata.n_obs, adata.n_obs))
-
-        k = (3 * ck1) ^ (1 * ck2)
-        k.compute_transition_matrix()
-
-        assert k[0][0]._value == 3 / 4
-        assert k[1][0]._value == 1 / 4
-
-    def test_adaptive_kernel_complex(self, adata: AnnData):
-        ck1 = ConnectivityKernel(adata).compute_transition_matrix()
-        ck1._mat_scaler = np.random.normal(size=(adata.n_obs, adata.n_obs))
-
-        ck2 = ConnectivityKernel(adata).compute_transition_matrix()
-        ck2._mat_scaler = np.random.normal(size=(adata.n_obs, adata.n_obs))
-
-        ck3 = ConnectivityKernel(adata).compute_transition_matrix()
-        ck3._mat_scaler = np.random.normal(size=(adata.n_obs, adata.n_obs))
-
-        k = 4 * ((3 * ck1) ^ (1 * ck2)) + 2 * ck3
-        k.compute_transition_matrix()
-
-        assert k[0][0].transition_matrix == 4 / 6
-        assert k[1][0].transition_matrix == 2 / 6
-        assert k[0][1][0][0]._value == 3 / 4
-        assert k[0][1][1][0]._value == 1 / 4
-
     def test_repr(self, adata: AnnData):
         rpr = repr(VelocityKernel(adata))
 
-        assert rpr == f"<{VelocityKernel.__name__}>"
+        assert rpr == f"{VelocityKernel.__name__}"
 
     def test_repr_inv(self, adata: AnnData):
         rpr = repr(~VelocityKernel(adata))
 
-        assert rpr == f"~<{VelocityKernel.__name__}>"
+        assert rpr == f"~{VelocityKernel.__name__}"
 
     def test_repr_inv_comb(self, adata: AnnData):
         rpr = repr(~(VelocityKernel(adata) + ConnectivityKernel(adata)))
 
         assert (
             rpr
-            == f"~((1 * <{VelocityKernel.__name__}>) + (1 * <{ConnectivityKernel.__name__}>))"
+            == f"~(1.0 * {VelocityKernel.__name__} + 1.0 * {ConnectivityKernel.__name__})"
         )
 
     def test_str_repr_equiv_no_transition_matrix(self, adata: AnnData):
         vk = VelocityKernel(adata)
-        string = str(vk)
-        rpr = repr(vk)
 
-        assert string == rpr
-        assert string == f"<{VelocityKernel.__name__}>"
+        assert repr(vk) == str(vk)
 
     def test_str(self, adata: AnnData):
         string = str(ConnectivityKernel(adata).compute_transition_matrix())
 
         assert (
-            string == f"<{ConnectivityKernel.__name__}[dnorm=True, key=connectivities]>"
-        )
-
-    def test_str_inv(self, adata: AnnData):
-        string = str(
-            ConnectivityKernel(adata, backward=True).compute_transition_matrix()
-        )
-
-        assert (
-            string
-            == f"~<{ConnectivityKernel.__name__}[dnorm=True, key=connectivities]>"
+            string == f"{ConnectivityKernel.__name__}[dnorm=True, key='connectivities']"
         )
 
     def test_combination_correct_parameters(self, adata: AnnData):
-        from cellrank.tl.kernels import CosineScheme
-
         k = VelocityKernel(adata).compute_transition_matrix(
             softmax_scale=4,
             seed=42,
-            scheme="cosine",
+            similarity="cosine",
         ) + (
             ConnectivityKernel(adata).compute_transition_matrix(density_normalize=False)
             + ConnectivityKernel(adata).compute_transition_matrix(
@@ -474,33 +448,12 @@ def test_combination_correct_parameters(self, adata: AnnData):
         assert {"dnorm": False, "key": "connectivities"} in k.params.values()
         assert {
             "softmax_scale": 4,
-            "mode": "deterministic",
-            "seed": 42,
-            "scheme": str(CosineScheme()),
+            "model": "deterministic",
+            "similarity": "cosine",
         } in k.params.values()
 
 
 class TestKernel:
-    def test_precomputed_not_array(self):
-        with pytest.raises(TypeError):
-            _ = PrecomputedKernel([[1, 0], [0, 1]])
-
-    def test_precomputed_not_square(self):
-        with pytest.raises(ValueError):
-            _ = PrecomputedKernel(np.random.normal(size=(10, 9)))
-
-    def test_precomputed_not_a_transition_matrix(self):
-        mat = random_transition_matrix(100)
-        mat[0, 0] = 0xDEADBEEF
-        with pytest.raises(ValueError):
-            _ = PrecomputedKernel(mat)
-
-    def test_precomputed_from_kernel_no_transition(self, adata: AnnData):
-        vk = VelocityKernel(adata)
-
-        with pytest.raises(ValueError):
-            PrecomputedKernel(vk)
-
     @pytest.mark.parametrize(
         "clazz",
         [
@@ -508,13 +461,14 @@ def test_precomputed_from_kernel_no_transition(self, adata: AnnData):
             VelocityKernel,
             PseudotimeKernel,
             CytoTRACEKernel,
-            PrecomputedKernel,
         ],
     )
     @pytest.mark.parametrize("key_added", [None, "foo"])
     def test_kernel_reads_correct_connectivities(
         self, adata: AnnData, key_added: Optional[str], clazz: type
     ):
+        if clazz == VelocityKernel and key_added == "foo":
+            pytest.skip("`get_moments` in scVelo doesn't support specifying key")
         del adata.uns["neighbors"]
         del adata.obsp["connectivities"]
         del adata.obsp["distances"]
@@ -535,98 +489,21 @@ def test_kernel_reads_correct_connectivities(
 
         k = clazz(**kwargs)
 
-        if isinstance(k, PrecomputedKernel):
-            assert k._conn is None
-        else:
+        if isinstance(k, ConnectivityMixin):
             np.testing.assert_array_equal(k._conn.A, conn.A)
+        else:
+            assert not hasattr(k, "_conn")
 
-    def test_precomputed_from_kernel(self, adata: AnnData):
-        vk = VelocityKernel(adata).compute_transition_matrix(
-            mode="deterministic",
-            softmax_scale=4,
-        )
-
-        pk = PrecomputedKernel(vk)
-        pk.write_to_adata()
-
-        assert pk.adata is vk.adata
-        assert pk._origin == str(vk).strip("~<>")
-        assert pk.params is not vk.params
-        assert pk.params == vk.params
-        assert pk.transition_matrix is not vk.transition_matrix
-        np.testing.assert_array_equal(pk.transition_matrix.A, vk.transition_matrix.A)
-
-    def test_precomputed_no_adata(self):
-        pk = PrecomputedKernel(random_transition_matrix(50))
-        pk.write_to_adata()
-
-        assert isinstance(pk.adata, AnnData)
-        assert pk._origin == "'array'"
-        assert pk.adata.shape == (50, 1)
-        assert pk.adata.obs.shape == (50, 0)
-        assert pk.adata.var.shape == (1, 0)
-        assert "T_fwd_params" in pk.adata.uns.keys()
-        assert pk.adata.uns["T_fwd_params"] == {"params": pk.params}
-        np.testing.assert_array_equal(
-            pk.adata.obsp["T_fwd"].toarray(), pk.transition_matrix.toarray()
-        )
-
-    def test_precomputed_different_adata(self, adata: AnnData):
-        vk = VelocityKernel(adata).compute_transition_matrix(
-            mode="deterministic", softmax_scale=4
-        )
-        bdata = adata.copy()
-
-        pk = PrecomputedKernel(vk, adata=bdata)
-
-        assert pk.adata is adata
-        assert pk.adata is vk.adata
-        assert pk.adata is not bdata
-
-    def test_precomputed_adata_origin(self, adata: AnnData):
-        vk = VelocityKernel(adata).compute_transition_matrix(
-            mode="deterministic", softmax_scale=4
-        )
-        vk.write_to_adata("foo")
-
-        pk = PrecomputedKernel("foo", adata=adata)
-
-        assert pk._origin == "adata.obsp['foo']"
-
-    def test_precomputed_adata(self, adata: AnnData):
-        pk = PrecomputedKernel(random_transition_matrix(adata.n_obs), adata=adata)
-
-        assert pk.adata is adata
-
-    def test_precomputed_transition_matrix(self, adata: AnnData):
-        mat = random_transition_matrix(adata.n_obs)
-        pk = PrecomputedKernel(mat)
-
-        np.testing.assert_array_equal(mat, pk.transition_matrix.toarray())
-
-    def test_precomputed_sum(self, adata: AnnData):
-        mat = random_transition_matrix(adata.n_obs)
-        pk = PrecomputedKernel(mat)
-        vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
-
-        expected = (0.5 * vk.transition_matrix) + (0.5 * pk.transition_matrix)
-        actual = (pk + vk).compute_transition_matrix()
-
-        np.testing.assert_array_almost_equal(
-            expected.toarray(), actual.transition_matrix.toarray()
-        )
-
-    @pytest.mark.parametrize("dnorm", [False, True])
     @pytest.mark.parametrize("sparse", [False, True])
-    def test_custom_preserves_type(self, adata: AnnData, sparse: bool, dnorm: bool):
-        c = CustomKernel(adata).compute_transition_matrix(sparse=sparse, dnorm=dnorm)
+    def test_custom_preserves_type(self, adata: AnnData, sparse: bool):
+        c = CustomKernel(adata).compute_transition_matrix(sparse=sparse)
 
         if sparse:
             assert isspmatrix_csr(c.transition_matrix)
         else:
             assert isinstance(c.transition_matrix, np.ndarray)
 
-        assert c.transition_matrix.dtype == _dtype
+        assert c.transition_matrix.dtype == np.float64
 
     def test_write_adata(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
@@ -648,13 +525,13 @@ def test_write_adata_key(self, adata: AnnData):
             adata.obsp["foo"].toarray(), vk.transition_matrix.toarray()
         )
 
-    @pytest.mark.parametrize("mode", ["deterministic", "stochastic"])
-    def test_vk_row_normalized(self, adata: AnnData, mode: str):
-        if mode == "stochastic":
+    @pytest.mark.parametrize("model", ["deterministic", "stochastic"])
+    def test_vk_row_normalized(self, adata: AnnData, model: str):
+        if model == "stochastic":
             pytest.importorskip("jax")
             pytest.importorskip("jaxlib")
         vk = VelocityKernel(adata)
-        vk.compute_transition_matrix(mode="stochastic", softmax_scale=4)
+        vk.compute_transition_matrix(model="stochastic", softmax_scale=4)
 
         np.testing.assert_allclose(vk.transition_matrix.sum(1), 1, rtol=_rtol)
 
@@ -695,20 +572,54 @@ def test_pseudotime_parallelize(self, adata: AnnData):
 
     def test_pseudotime_inverse(self, adata: AnnData):
         pk = PseudotimeKernel(adata, time_key="latent_time")
-        pt = pk.pseudotime.copy()
 
         pk_inv = ~pk
 
-        assert pk_inv is pk
+        assert pk_inv is not pk
         assert pk_inv.backward
-        np.testing.assert_allclose(pt, 1 - pk_inv.pseudotime)
+        np.testing.assert_allclose(pk.pseudotime, 1 - pk_inv.pseudotime)
+
+    @pytest.mark.parametrize("empty", [False, True])
+    def test_set_adata(self, adata: AnnData, empty: bool):
+        data = None if empty else adata.copy()
+        ck1 = ConnectivityKernel(adata).compute_transition_matrix()
+        ck2 = ConnectivityKernel(adata).compute_transition_matrix()
+        ck = ck1 + ck2
+
+        ck.adata = data
+
+        assert ck.adata is data
+        for k in ck:
+            assert k.adata is data
 
-    @pytest.mark.parametrize("mode", ["deterministic", "stochastic", "sampling"])
-    def test_manual_combination(self, adata: AnnData, mode: str):
-        if mode == "stochastic":
+    def test_set_adata_wrong_shape(self, adata: AnnData):
+        ck = ConnectivityKernel(adata)
+        with pytest.raises(
+            ValueError, match="Expected new `AnnData` .* to have same shape"
+        ):
+            ck.adata = adata[:2].copy()
+
+    def test_forward_backward_combination(self, adata: AnnData):
+        vkf = VelocityKernel(adata, backward=False)
+        vkb = VelocityKernel(adata, backward=True)
+        with pytest.raises(ValueError, match="Unable to combine both forward"):
+            _ = vkf + vkb
+
+    def test_unidirectional_combination(self, adata: AnnData):
+        ck1 = ConnectivityKernel(adata).compute_transition_matrix()
+        ck2 = ConnectivityKernel(adata).compute_transition_matrix()
+        ck = ck1 + ck2
+
+        assert ck.backward is None
+
+    @pytest.mark.parametrize("model", ["deterministic", "stochastic", "monte_carlo"])
+    def test_manual_combination(self, adata: AnnData, model: str):
+        if model == "stochastic":
             pytest.importorskip("jax")
             pytest.importorskip("jaxlib")
-        vk = VelocityKernel(adata).compute_transition_matrix(mode=mode, softmax_scale=4)
+        vk = VelocityKernel(adata).compute_transition_matrix(
+            model=model, softmax_scale=4, n_samples=1
+        )
         ck = ConnectivityKernel(adata).compute_transition_matrix()
 
         T_vk = vk.transition_matrix
@@ -742,34 +653,14 @@ def test_manual_combination_no_precomputed(self, adata: AnnData):
 
         np.testing.assert_allclose(T_comb_manual.A, T_comb_kernel.A, rtol=_rtol)
 
-    def test_manual_combination_backward(self, adata: AnnData):
-        backward, density_normalize = True, False
-
-        vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
-            softmax_scale=4
-        )
-
-        ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix(
-            density_normalize=density_normalize
-        )
-
-        T_vk = vk.transition_matrix
-        T_ck = ck.transition_matrix
-        T_comb_manual = 0.8 * T_vk + 0.2 * T_ck
-
-        comb_kernel = 0.8 * vk + 0.2 * ck
-        T_comb_kernel = comb_kernel.transition_matrix
-
-        np.testing.assert_allclose(T_comb_manual.A, T_comb_kernel.A, rtol=_rtol)
-
-    def test_manual_combination_backward_dense_norm(self, adata: AnnData):
-        backward, density_normalize = True, True
-
+    @pytest.mark.parametrize("density_normalize", [False, TransportMapKernel])
+    def test_manual_combination_backward(self, adata: AnnData, density_normalize):
+        backward = True
         vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
             softmax_scale=4
         )
 
-        ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix(
+        ck = ConnectivityKernel(adata).compute_transition_matrix(
             density_normalize=density_normalize
         )
 
@@ -782,8 +673,8 @@ def test_manual_combination_backward_dense_norm(self, adata: AnnData):
 
         np.testing.assert_allclose(T_comb_manual.A, T_comb_kernel.A, rtol=_rtol)
 
-    def compare_with_scanpy_density_normalize(self, adata: AnnData):
-        # check whether cellrank's transition matrix matches scanpy's
+    @pytest.mark.parametrize("density_normalize", [False, True])
+    def test_dnorm_scanpy(self, adata: AnnData, density_normalize: bool):
         density_normalize = True
         ck = ConnectivityKernel(adata).compute_transition_matrix(
             density_normalize=density_normalize
@@ -794,27 +685,8 @@ def compare_with_scanpy_density_normalize(self, adata: AnnData):
         neigh.compute_transitions(density_normalize=density_normalize)
         T_sc = neigh.transitions
 
-        # check whether these are the same while leaving them sparse
         assert T_sc.shape == T_cr.shape
-        assert len(T_sc.indices) == len(T_cr.indices)
-        assert np.allclose((T_cr - T_sc).data, 0)
-
-    def compare_with_scanpy(self, adata: AnnData):
-        # check whether cellrank's transition matrix matches scanpy's
-        density_normalize = False
-        ck = ConnectivityKernel(adata).compute_transition_matrix(
-            density_normalize=density_normalize
-        )
-        T_cr = ck.transition_matrix
-
-        neigh = Neighbors(adata)
-        neigh.compute_transitions(density_normalize=density_normalize)
-        T_sc = neigh.transitions
-
-        # check whether these are the same while leaving them sparse
-        assert T_sc.shape == T_cr.shape
-        assert len(T_sc.indices) == len(T_cr.indices)
-        assert np.allclose((T_cr - T_sc).data, 0)
+        np.testing.assert_allclose(T_cr.A, T_cr.A)
 
     def test_connectivities_key_kernel(self, adata: AnnData):
         key = "foobar"
@@ -825,9 +697,8 @@ def test_connectivities_key_kernel(self, adata: AnnData):
         T_cr = ck.transition_matrix
 
         assert key == ck.params["key"]
-        np.testing.assert_array_equal(T_cr, adata.obsp[key])
-
-        del adata.obsp[key]
+        assert T_cr is not adata.obsp[key]
+        np.testing.assert_array_equal(T_cr.A, adata.obsp[key])
 
 
 class TestKernelAddition:
@@ -873,100 +744,6 @@ def test_addition_3_kernels(self, adata: AnnData):
 
         np.testing.assert_allclose(k.transition_matrix.A, expected)
 
-    def test_addition_adaptive(self, adata: AnnData):
-        adata.obsp["velocity_variances"] = vv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        adata.obsp["connectivity_variances"] = cv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        vk, ck = create_kernels(
-            adata,
-            velocity_variances="velocity_variances",
-            connectivity_variances="connectivity_variances",
-        )
-
-        k = vk ^ ck
-        expected = _normalize(
-            0.5 * vv * vk.transition_matrix + 0.5 * cv * ck.transition_matrix
-        )
-
-        np.testing.assert_allclose(k.transition_matrix.A, expected)
-
-    def test_addition_adaptive_constants(self, adata: AnnData):
-        a, b = np.random.uniform(0, 10, 2)
-        s = a + b
-        adata.obsp["velocity_variances"] = vv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        adata.obsp["connectivity_variances"] = cv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        vk, ck = create_kernels(
-            adata,
-            velocity_variances="velocity_variances",
-            connectivity_variances="connectivity_variances",
-        )
-
-        k = a * vk ^ b * ck
-        expected = _normalize(
-            a / s * vv * vk.transition_matrix + b / s * cv * ck.transition_matrix
-        )
-
-        np.testing.assert_allclose(k.transition_matrix.A, expected)
-
-    def test_addition_adaptive_wrong_variances(self, adata: AnnData):
-        a, b = np.random.uniform(0, 10, 2)
-        s = a + b
-        adata.obsp["velocity_variances"] = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        adata.obsp["connectivity_variances"] = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        vk, ck = create_kernels(
-            adata,
-            velocity_variances="velocity_variances",
-            connectivity_variances="connectivity_variances",
-        )
-
-        k = a * vk ^ b * ck
-        expected = _normalize(
-            a / s * vk.transition_matrix + b / s * ck.transition_matrix
-        )
-
-        assert not np.allclose(k.transition_matrix.A, expected.A)
-
-    def test_addition_adaptive_4_kernels(self, adata: AnnData):
-        a, b, c, d = np.random.uniform(0, 10, 4)
-        s = a + b + c + d
-        adata.obsp["velocity_variances"] = vv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        adata.obsp["connectivity_variances"] = cv = np.random.random(
-            size=(adata.n_obs, adata.n_obs)
-        )
-        vk, ck = create_kernels(
-            adata,
-            velocity_variances="velocity_variances",
-            connectivity_variances="connectivity_variances",
-        )
-        vk1, ck1 = create_kernels(
-            adata,
-            velocity_variances="velocity_variances",
-            connectivity_variances="connectivity_variances",
-        )
-
-        k = a * vk ^ b * ck ^ c * vk1 ^ d * ck1
-        expected = _normalize(
-            a / s * vv * vk.transition_matrix
-            + b / s * cv * ck.transition_matrix
-            + c / s * vv * vk1.transition_matrix
-            + d / s * cv * ck1.transition_matrix
-        )
-
-        np.testing.assert_allclose(k.transition_matrix.A, expected)
-
 
 class TestKernelCopy:
     def test_copy_simple(self, adata: AnnData):
@@ -993,36 +770,13 @@ def test_copy_params(self, adata: AnnData):
         vk2 = vk1.copy()
 
         assert vk1.params == vk2.params
-
-    def test_copy_cond_num(self, adata: AnnData):
-        for KernelClass in [
-            VelocityKernel,
-            ConnectivityKernel,
-            PseudotimeKernel,
-            PrecomputedKernel,
-        ]:
-            if KernelClass is PrecomputedKernel:
-                k1 = KernelClass(
-                    random_transition_matrix(adata.n_obs), compute_cond_num=True
-                )
-            elif KernelClass is VelocityKernel:
-                k1 = KernelClass(
-                    adata, compute_cond_num=True
-                ).compute_transition_matrix(softmax_scale=4)
-            else:
-                k1 = KernelClass(
-                    adata, compute_cond_num=True
-                ).compute_transition_matrix()
-            k2 = k1.copy()
-
-            assert k1.condition_number == k2.condition_number
+        assert vk1.params is not vk2.params
 
     def test_copy_velocity_kernel(self, adata: AnnData):
         vk1 = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
         vk2 = vk1.copy()
 
         np.testing.assert_array_equal(vk1.transition_matrix.A, vk2.transition_matrix.A)
-        np.testing.assert_array_equal(vk1.logits.A, vk2.logits.A)
 
         assert vk1.params == vk2.params
         assert vk1.backward == vk2.backward
@@ -1048,10 +802,20 @@ def test_copy_works(self, adata: AnnData):
         ck2 = ck1.copy()
         ck1.compute_transition_matrix()
 
-        assert (
-            ck1._transition_matrix is not None
-        )  # calling the property would trigger the calculation
-        assert ck2._transition_matrix is None
+        assert ck1.transition_matrix is not None
+        assert ck2.transition_matrix is None
+
+    @pytest.mark.parametrize(
+        "ignored", (("_transition_matrix",), ("_params", "foobar"))
+    )
+    def test_copy_ignore(self, adata: AnnData, ignored: Tuple[str, ...]):
+        ck1 = ConnectivityKernel(adata).compute_transition_matrix()
+        ck2 = ck1._copy_ignore(*ignored)
+
+        for attr in ignored:
+            assert getattr(ck2, attr, None) is None
+        assert ck2._conn is not None
+        assert ck2._conn_key == ck1._conn_key
 
 
 class TestGeneral:
@@ -1086,115 +850,12 @@ def test_kernels_unique(self, adata: AnnData):
         assert len(v.kernels) == 1
         assert v.kernels[0] is vk
 
-    def test_no_comp_cond_num(self, adata: AnnData):
-        vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
-
-        assert vk.condition_number is None
-
-    def test_comp_cond_num(self, adata: AnnData):
-        vk = VelocityKernel(adata, compute_cond_num=True).compute_transition_matrix(
-            softmax_scale=4
-        )
-
-        assert isinstance(vk.condition_number, float)
-
-    def test_comp_cond_num_or_policy(self, adata: AnnData):
-        vk = VelocityKernel(adata, compute_cond_num=True).compute_transition_matrix(
-            softmax_scale=4
-        )
-        ck = ConnectivityKernel(
-            adata, compute_cond_num=False
-        ).compute_transition_matrix()
-        v = (vk + ck).compute_transition_matrix()
-
-        assert isinstance(vk.condition_number, float)
-        assert ck.condition_number is None
-        assert isinstance(v.condition_number, float)
-
-
-class TestTransitionProbabilities:
-    def test_pearson_correlations_fwd(self, adata: AnnData):
-        # test whether pearson correlations in cellrank match those from scvelo, forward case
-        backward = False
-
-        # compute pearson correlations using scvelo
-        velo_graph = adata.obsp["velocity_graph"] + adata.obsp["velocity_graph_neg"]
-
-        # compute pearson correlations using cellrank
-        vk = VelocityKernel(adata, backward=backward)
-        vk.compute_transition_matrix(mode="deterministic", softmax_scale=4)
-        pearson_correlations_cr = vk.logits
-
-        pc_r = velo_graph.copy()
-        pc_r.data = np.array(pearson_correlations_cr[(velo_graph != 0)]).squeeze()
-
-        assert np.max(np.abs((pc_r - velo_graph).data)) < _rtol
-
-    def test_pearson_correlations_bwd(self, adata: AnnData):
-        # test whether pearson correlations in cellrank match those from scvelo, backward case
-        backward = True
-
-        # compute pearson correlations using scvelo
-        velo_graph = (adata.obsp["velocity_graph"] + adata.obsp["velocity_graph_neg"]).T
-
-        # compute pearson correlations using cellrak
-        vk = VelocityKernel(adata, backward=backward)
-        vk.compute_transition_matrix(
-            mode="deterministic", backward_mode="transpose", softmax_scale=4
-        )
-        pearson_correlations_cr = vk.logits
-
-        pc_r = velo_graph.copy()
-        pc_r.data = np.array(pearson_correlations_cr[(velo_graph != 0)]).squeeze()
-
-        assert np.max(np.abs((pc_r - velo_graph.T).data)) < _rtol
-
-    def test_transition_probabilities_fwd(self, adata: AnnData):
-        # test whether transition probabilities in cellrank match those from scvelo, forward case
-        sigma_test = 3
-
-        # compute transition probabilities using cellrank
-        vk = VelocityKernel(adata)
-        vk.compute_transition_matrix(softmax_scale=sigma_test, mode="deterministic")
-        T_cr = vk.transition_matrix
-
-        pearson_correlation = vk.logits
-        T_exp = np.expm1(pearson_correlation * sigma_test)
-        T_exp.data += 1
-        T_exp = _normalize(T_exp)
-
-        np.testing.assert_allclose(T_exp.A, T_cr.A)  # don't use data, can be reordered
-
-    def test_transition_probabilities_bwd(self, adata: AnnData):
-        # test whether transition probabilities in cellrank match those from scvelo, backward case
-        sigma_test = 3
-
-        # compute transition probabilities using cellrank
-        vk = VelocityKernel(adata, backward=True)
-        vk.compute_transition_matrix(softmax_scale=sigma_test, mode="deterministic")
-        T_cr = vk.transition_matrix
-
-        pearson_correlation = vk.logits
-        T_exp = np.expm1(pearson_correlation * sigma_test)
-        T_exp.data += 1
-        T_exp = _normalize(T_exp)
-
-        np.testing.assert_allclose(T_exp.A, T_cr.A)  # don't use data, can be reordered
-
-    def test_estimate_softmax_scale(self, adata: AnnData):
-        vk = VelocityKernel(adata)
-        vk.compute_transition_matrix(
-            mode="deterministic", show_progress_bar=False, softmax_scale=None
-        )
-
-        assert isinstance(vk.params["softmax_scale"], float)
-
 
 class TestMonteCarlo:
     def test_mc_and_mc_fwd_1k(self, adata: AnnData):
         vk_mc = VelocityKernel(adata, backward=False)
         vk_mc.compute_transition_matrix(
-            mode="monte_carlo",
+            model="monte_carlo",
             show_progress_bar=False,
             n_samples=1000,
             n_jobs=4,
@@ -1203,7 +864,7 @@ def test_mc_and_mc_fwd_1k(self, adata: AnnData):
 
         vk_s = VelocityKernel(adata, backward=False)
         vk_s.compute_transition_matrix(
-            mode="monte_carlo",
+            model="monte_carlo",
             show_progress_bar=False,
             n_samples=1000,
             n_jobs=4,
@@ -1218,7 +879,7 @@ def test_mc_and_mc_fwd_1k(self, adata: AnnData):
     def test_monte_carlo_5k(self, adata: AnnData):
         vk_mc = VelocityKernel(adata, backward=False)
         vk_mc.compute_transition_matrix(
-            mode="monte_carlo",
+            model="monte_carlo",
             show_progress_bar=False,
             n_samples=5000,
             n_jobs=4,
@@ -1228,7 +889,7 @@ def test_monte_carlo_5k(self, adata: AnnData):
 
         vk_s = VelocityKernel(adata, backward=False)
         vk_s.compute_transition_matrix(
-            mode="monte_carlo",
+            model="monte_carlo",
             show_progress_bar=False,
             n_samples=5000,
             n_jobs=4,
@@ -1245,7 +906,7 @@ def test_monte_carlo_5k(self, adata: AnnData):
     def test_monte_carlo_and_stochastic(self, adata: AnnData):
         vk_mc = VelocityKernel(adata, backward=False)
         vk_mc.compute_transition_matrix(
-            mode="monte_carlo",
+            model="monte_carlo",
             show_progress_bar=False,
             n_samples=1000,
             n_jobs=4,
@@ -1254,7 +915,7 @@ def test_monte_carlo_and_stochastic(self, adata: AnnData):
 
         vk_s = VelocityKernel(adata, backward=False)
         vk_s.compute_transition_matrix(
-            mode="stochastic", show_progress_bar=False, n_jobs=4, softmax_scale=4
+            model="stochastic", show_progress_bar=False, n_jobs=4, softmax_scale=4
         )
 
         val = np.mean(
@@ -1267,40 +928,42 @@ class TestVelocityScheme:
     def test_invalid_string_key(self, adata: AnnData):
         vk = VelocityKernel(adata)
         with pytest.raises(ValueError):
-            vk.compute_transition_matrix(scheme="foobar")
+            vk.compute_transition_matrix(similarity="foobar")
 
     def test_not_callable(self, adata: AnnData):
         vk = VelocityKernel(adata)
         with pytest.raises(
             TypeError, match="Expected `scheme` to be a function, found"
         ):
-            vk.compute_transition_matrix(scheme=1311)
+            vk.compute_transition_matrix(similarity=1311)
 
     def test_custom_function_not_sum_to_1(self, adata: AnnData):
         vk = VelocityKernel(adata)
         with pytest.raises(ValueError, match=r"Matrix is not row-stochastic."):
-            vk.compute_transition_matrix(scheme=InvalidFuncProbs())
+            vk.compute_transition_matrix(similarity=InvalidFuncProbs())
 
     def test_custom_function_invalid_hessian(self, adata: AnnData):
         vk = VelocityKernel(adata)
         with pytest.raises(ValueError, match=r"Expected full Hessian matrix"):
             vk.compute_transition_matrix(
-                mode="stochastic", scheme=InvalidFuncHessianShape(), softmax_scale=4
+                model="stochastic",
+                similarity=InvalidFuncHessianShape(),
+                softmax_scale=4,
             )
 
     @pytest.mark.parametrize("backward", [True, False])
     def test_implementations_differ(self, adata: AnnData, backward: bool):
         vk_dot = VelocityKernel(adata, backward=backward)
         vk_dot.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme="dot_product"
+            model="deterministic", softmax_scale=4, similarity="dot_product"
         )
         vk_cos = VelocityKernel(adata, backward=backward)
         vk_cos.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme="cosine"
+            model="deterministic", softmax_scale=4, similarity="cosine"
         )
         vk_cor = VelocityKernel(adata, backward=backward)
         vk_cor.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme="correlation"
+            model="deterministic", softmax_scale=4, similarity="correlation"
         )
 
         np.testing.assert_allclose(vk_dot.transition_matrix.sum(1), 1.0)
@@ -1316,9 +979,9 @@ def test_implementations_differ(self, adata: AnnData, backward: bool):
         zip(
             ["dot_product", "cosine", "correlation"],
             [
-                cr.tl.kernels.DotProductScheme(),
-                cr.tl.kernels.CosineScheme(),
-                cr.tl.kernels.CorrelationScheme(),
+                cr.tl.kernels.utils.DotProduct(),
+                cr.tl.kernels.utils.Cosine(),
+                cr.tl.kernels.utils.Correlation(),
             ],
         ),
     )
@@ -1327,10 +990,10 @@ def test_function_and_string_key(self, adata: AnnData, key: str, fn: Callable):
         vk_fn = VelocityKernel(adata)
 
         vk_k.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme=key
+            model="deterministic", softmax_scale=4, similarity=key
         )
         vk_fn.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme=fn
+            model="deterministic", softmax_scale=4, similarity=fn
         )
 
         np.testing.assert_allclose(vk_k.transition_matrix.A, vk_fn.transition_matrix.A)
@@ -1339,43 +1002,43 @@ def test_function_and_string_key(self, adata: AnnData, key: str, fn: Callable):
     def test_custom_function(self, adata: AnnData, backward: bool):
         vk = VelocityKernel(adata, backward=backward)
         vk.compute_transition_matrix(
-            mode="deterministic", softmax_scale=4, scheme=CustomFuncHessian()
+            model="deterministic", softmax_scale=4, similarity=CustomFuncHessian()
         )
 
-        assert vk.params["scheme"] == str(CustomFuncHessian())
+        assert vk.params["similarity"] == str(CustomFuncHessian())
 
     def test_custom_function_stochastic_no_hessian(self, adata: AnnData):
         vk = VelocityKernel(adata)
         vk.compute_transition_matrix(
-            mode="stochastic", scheme=CustomFunc(), softmax_scale=4, n_samples=10
+            model="stochastic", similarity=CustomFunc(), softmax_scale=4, n_samples=10
         )
 
-        assert vk.params["mode"] == "monte_carlo"
-        assert vk.params["scheme"] == str(CustomFunc())
+        assert vk.params["model"] == "stochastic"
+        assert vk.params["similarity"] == str(CustomFunc())
 
 
 class TestComputeProjection:
     def test_no_transition_matrix(self, adata: AnnData):
         with pytest.raises(RuntimeError, match=r"Compute transition matrix first as"):
-            cr.tl.kernels.ConnectivityKernel(adata).compute_projection()
+            cr.tl.kernels.ConnectivityKernel(adata).plot_projection()
 
     def test_no_basis(self, adata: AnnData):
         ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
         with pytest.raises(KeyError, match=r"Unable to find a basis in"):
-            ck.compute_projection(basis="foo")
+            ck.plot_projection(basis="foo")
 
-    def test_basis_prefix(self, adata: AnnData):
+    def test_normal_run(self, adata: AnnData):
         ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
-        ck.compute_projection(basis="X_umap")
+        ck.plot_projection(basis="umap")
 
     @pytest.mark.parametrize("write_first", [True, False])
     def test_write_to_adata(self, adata: AnnData, write_first: bool):
         ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
         if write_first:
             ck.write_to_adata()
-            ck.compute_projection(basis="umap")
+            ck.plot_projection(basis="umap")
         else:
-            ck.compute_projection(basis="umap")
+            ck.plot_projection(basis="umap")
             ck.write_to_adata()
 
         assert adata.uns[Key.uns.kernel(ck.backward) + "_params"] == {
@@ -1386,7 +1049,7 @@ def test_write_to_adata(self, adata: AnnData, write_first: bool):
     @pytest.mark.parametrize("key_added", [None, "foo"])
     def test_key_added(self, adata: AnnData, key_added: Optional[str]):
         ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
-        ck.compute_projection(basis="umap", copy=False, key_added=key_added)
+        ck.plot_projection(basis="umap", key_added=key_added)
 
         key = Key.uns.kernel(ck.backward, key=key_added)
         ukey = f"{key}_params"
@@ -1395,29 +1058,11 @@ def test_key_added(self, adata: AnnData, key_added: Optional[str]):
         assert adata.uns[ukey] == {"embeddings": ["umap"]}
         np.testing.assert_array_equal(adata.obsm[key].shape, adata.obsm["X_umap"].shape)
 
-    @pytest.mark.parametrize("copy", [True, False])
-    def test_copy(self, adata: AnnData, copy: bool):
-        ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
-        res = ck.compute_projection(basis="umap", copy=copy)
-
-        if copy:
-            assert isinstance(res, np.ndarray)
-            np.testing.assert_array_equal(res.shape, adata.obsm["X_umap"].shape)
-        else:
-            assert res is None
-            key = Key.uns.kernel(ck.backward) + "_umap"
-            np.testing.assert_array_equal(
-                adata.obsm[key].shape, adata.obsm["X_umap"].shape
-            )
-
-    def test_nan_in_embedding(self, adata: AnnData):
-        adata.obsm["X_umap"][-1] = np.nan
+    def test_nan_in_embedding(self, adata_large: AnnData):
+        adata_large.obsm["X_umap"][0, :] = np.nan
 
-        ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
-        res = ck.compute_projection(basis="umap", copy=True)
-
-        assert not np.all(np.isnan(res))
-        assert np.all(np.isnan(res[-1, :]))
+        ck = cr.tl.kernels.ConnectivityKernel(adata_large).compute_transition_matrix()
+        ck.plot_projection(basis="umap")
 
 
 class TestPseudotimeKernelScheme:
@@ -1460,7 +1105,7 @@ def test_scheme(self, adata: AnnData, scheme: str):
             assert pk.params["frac_to_keep"] == 0.3
             assert "b" not in pk.params
             assert "nu" not in pk.params
-        elif scheme == "soft":
+        else:
             assert pk.params["b"] == 10
             assert pk.params["nu"] == 0.5
             assert "k" not in pk.params
@@ -1471,33 +1116,40 @@ class TestCytoTRACEKernel:
     def test_layer(self, adata: AnnData, layer: str):
         if layer == "foo":
             with pytest.raises(KeyError, match=layer):
-                _ = CytoTRACEKernel(adata, layer=layer)
+                _ = CytoTRACEKernel(adata).compute_cytotrace(layer=layer)
         else:
-            _ = CytoTRACEKernel(adata, layer=layer)
+            _ = CytoTRACEKernel(adata).compute_cytotrace(layer=layer)
             assert adata.uns[Key.cytotrace("params")]["layer"] == layer
 
     @pytest.mark.parametrize("agg", list(CytoTRACEAggregation))
     def test_aggregation(self, adata: AnnData, agg: CytoTRACEAggregation):
-        _ = CytoTRACEKernel(adata, aggregation=agg)
+        _ = CytoTRACEKernel(adata).compute_cytotrace(aggregation=agg)
         assert adata.uns[Key.cytotrace("params")]["aggregation"] == agg
 
     @pytest.mark.parametrize("use_raw", [False, True])
     def test_raw(self, adata: AnnData, use_raw: bool):
-        _ = CytoTRACEKernel(adata, use_raw=use_raw)
+        _ = CytoTRACEKernel(adata).compute_cytotrace(use_raw=use_raw)
         assert adata.uns[Key.cytotrace("params")]["use_raw"] == use_raw
 
-    def test_correct_class(self, adata: AnnData):
+    @pytest.mark.parametrize("copy", [False, True])
+    def test_correct_class(self, adata: AnnData, copy: bool):
         k = CytoTRACEKernel(adata)
-        assert isinstance(k, PseudotimeKernel)
+        if copy:
+            k = k.copy()
+
+        assert isinstance(k, CytoTRACEKernel)
         assert k._time_key == Key.cytotrace("pseudotime")
+        assert k.pseudotime is None
 
     def test_writes_params(self, adata: AnnData):
-        k = CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
+        k = CytoTRACEKernel(adata).compute_cytotrace(
+            use_raw=False, layer="X", aggregation="mean"
+        )
 
         assert adata.uns[Key.cytotrace("params")] == {
             "layer": "X",
             "aggregation": "mean",
-            "n_top_genes": 143,
+            "n_genes": 143,
             "use_raw": False,
         }
 
@@ -1518,34 +1170,56 @@ def test_writes_params(self, adata: AnnData):
 
     def test_raw_less_genes(self, adata: AnnData):
         adata.raw = adata.raw.to_adata()[:, :20]
-        _ = CytoTRACEKernel(adata, use_raw=True, n_top_genes=31)
+        _ = CytoTRACEKernel(adata).compute_cytotrace(use_raw=True, n_genes=31)
         assert adata.uns[Key.cytotrace("params")] == {
             "layer": "Ms",
             "aggregation": "mean",
-            "n_top_genes": 20,
+            "n_genes": 20,
             "use_raw": True,
         }
 
-    @pytest.mark.parametrize("n_top_genes", [0, 10, 300])
+    @pytest.mark.parametrize("n_genes", [0, 10, 300])
     @pytest.mark.parametrize("use_raw", [False, True])
-    def test_n_top_genes(self, adata: AnnData, use_raw: bool, n_top_genes: int):
-        n_genes = 50 if use_raw else 143
-        if n_top_genes < 0:
-            with pytest.raises(ValueError, match=r"Expected `n_top_genes`"):
-                _ = CytoTRACEKernel(adata, use_raw=use_raw, n_top_genes=n_top_genes)
+    def test_n_top_genes(self, adata: AnnData, use_raw: bool, n_genes: int):
+        n_genes = min(adata.raw.n_vars if use_raw else adata.n_vars, n_genes)
+        if n_genes <= 0:
+            with pytest.raises(ValueError, match=r"Expected .* genes to be positive"):
+                _ = CytoTRACEKernel(adata).compute_cytotrace(
+                    use_raw=use_raw, n_genes=n_genes
+                )
         else:
-            _ = CytoTRACEKernel(adata, use_raw=use_raw, n_top_genes=n_top_genes)
+            _ = CytoTRACEKernel(adata).compute_cytotrace(
+                use_raw=use_raw, n_genes=n_genes
+            )
             assert adata.var[Key.cytotrace("correlates")].sum() == n_genes
-            assert adata.uns[Key.cytotrace("params")]["n_top_genes"] == n_genes
+            assert adata.uns[Key.cytotrace("params")]["n_genes"] == n_genes
+
+    def test_rereads_pseudotime(self, adata: AnnData):
+        k1 = CytoTRACEKernel(adata).compute_cytotrace(
+            use_raw=False, layer="X", aggregation="mean"
+        )
+        k2 = CytoTRACEKernel(adata)
+
+        assert k1.pseudotime is not k2.pseudotime
+        np.testing.assert_array_equal(k1.pseudotime, k2.pseudotime)
+
+    def test_compute_transition_matrix_no_pt(self, adata: AnnData):
+        k = CytoTRACEKernel(adata)
+        with pytest.raises(ValueError, match=r"Compute pseudotime"):
+            k.compute_transition_matrix()
 
     def test_compute_transition_matrix(self, adata: AnnData):
-        k = CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
+        k = CytoTRACEKernel(adata).compute_cytotrace(
+            use_raw=False, layer="X", aggregation="mean"
+        )
         k.compute_transition_matrix()
 
         np.testing.assert_allclose(k.transition_matrix.sum(1), 1.0)
 
     def test_inversion(self, adata: AnnData):
-        k = ~CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
+        k = ~CytoTRACEKernel(adata).compute_cytotrace(
+            use_raw=False, layer="X", aggregation="mean"
+        )
 
         pt = adata.obs[Key.cytotrace("pseudotime")].values
         np.testing.assert_array_equal(np.max(pt) - pt, k.pseudotime)
@@ -1568,10 +1242,13 @@ def test_cache(self, adata: AnnData, reuse: bool):
         np.testing.assert_allclose(tmk.transition_matrix.sum(1), 1.0)
 
     @pytest.mark.parametrize("dtype", [spmatrix, np.ndarray, AnnData])
-    def test_returned_type(self, adata: AnnData, dtype: type):
+    def test_returned_dtype(self, adata: AnnData, dtype: type):
         tmk = DummyTMapKernel(adata).compute_transition_matrix(dtype=dtype)
 
         np.testing.assert_allclose(tmk.transition_matrix.sum(1), 1.0)
+        assert isinstance(tmk.transport_maps, dict)
+        for v in tmk.transport_maps.values():
+            assert isinstance(v, AnnData)
 
     @pytest.mark.parametrize("kind", ["source", "target"])
     def test_wrong_shape(self, adata: AnnData, kind: str):
@@ -1592,6 +1269,7 @@ def test_wrong_names(self, adata: AnnData, kind: str):
 
 class TestSingleFlow:
     def test_no_transition_matrix(self, kernel: Kernel):
+        kernel = kernel.kernels[0][0]
         kernel._transition_matrix = None
         with pytest.raises(RuntimeError, match=r"Compute transition matrix first as"):
             kernel.plot_single_flow("Astrocytes", "clusters", "age(days)")
@@ -1649,6 +1327,135 @@ def test_remove_empty_clusters_none_remain(self, kernel: Kernel):
             )
 
 
+class TestPrecomputedKernel:
+    def test_precomputed_not_array(self):
+        with pytest.raises(TypeError, match=r"Expected `object`"):
+            _ = PrecomputedKernel([[1, 0], [0, 1]])
+
+    def test_precomputed_not_square(self):
+        with pytest.raises(
+            ValueError, match=r"Expected matrix to be of shape `\(10, 10\)`"
+        ):
+            _ = PrecomputedKernel(np.random.normal(size=(10, 9)))
+
+    def test_precomputed_not_a_transition_matrix(self):
+        mat = random_transition_matrix(100)
+        mat[0, 0] = -1e-3
+        with pytest.raises(
+            ValueError, match=r"Unable to normalize matrix with negative values."
+        ):
+            _ = PrecomputedKernel(mat)
+
+    def test_precomputed_from_kernel_no_transition(self, adata: AnnData):
+        vk = VelocityKernel(adata)
+
+        with pytest.raises(RuntimeError, match="Compute transition matrix"):
+            PrecomputedKernel(vk)
+
+    def test_precomputed_from_kernel(self, adata: AnnData):
+        vk = VelocityKernel(adata).compute_transition_matrix(
+            model="deterministic",
+            softmax_scale=4,
+        )
+
+        pk = PrecomputedKernel(vk)
+        pk.write_to_adata()
+
+        assert pk.adata is vk.adata
+        assert pk.params is not vk.params
+        assert pk.params.pop("origin") == repr(vk)
+        assert pk.params == vk.params
+        assert pk.transition_matrix is vk.transition_matrix
+
+    def test_precomputed_no_adata(self):
+        pk = PrecomputedKernel(random_transition_matrix(50))
+        pk.write_to_adata()
+
+        assert isinstance(pk.adata, AnnData)
+        assert pk.params["origin"] == "array"
+        assert pk.adata.shape == (50, 1)
+        assert pk.adata.obs.shape == (50, 0)
+        assert pk.adata.var.shape == (1, 0)
+        assert "T_fwd_params" in pk.adata.uns.keys()
+        assert pk.adata.uns["T_fwd_params"] == {"params": pk.params}
+        np.testing.assert_array_equal(pk.adata.obsp["T_fwd"], pk.transition_matrix)
+
+    def test_precomputed_different_adata(self, adata: AnnData):
+        vk = VelocityKernel(adata).compute_transition_matrix(
+            model="deterministic", softmax_scale=4
+        )
+        bdata = adata.copy()
+
+        pk = PrecomputedKernel(vk, adata=bdata)
+
+        assert pk.adata is adata
+        assert pk.adata is vk.adata
+        assert pk.adata is not bdata
+
+    def test_precomputed_adata_origin(self, adata: AnnData):
+        vk = VelocityKernel(adata).compute_transition_matrix(
+            model="deterministic", softmax_scale=4
+        )
+        vk.write_to_adata("foo")
+
+        pk = PrecomputedKernel("foo", adata=adata)
+
+        assert pk.params["origin"] == "adata.obsp['foo']"
+
+    def test_precomputed_adata(self, adata: AnnData):
+        pk = PrecomputedKernel(random_transition_matrix(adata.n_obs), adata=adata)
+
+        assert pk.adata is adata
+
+    @pytest.mark.parametrize("backward", [False, True, None])
+    def test_precomputed_transition_matrix(self, adata: AnnData, backward: bool):
+        mat = random_transition_matrix(adata.n_obs)
+        pk = PrecomputedKernel(mat, backward=backward)
+
+        np.testing.assert_array_equal(mat, pk.transition_matrix)
+        assert pk.backward is backward
+
+    @pytest.mark.parametrize("backward", [False, True, None])
+    def test_precomputed_bool(self, adata: AnnData, backward: bool):
+        adata.obsp[Key.uns.kernel(backward)] = mat = random_transition_matrix(
+            adata.n_obs
+        )
+        pk = PrecomputedKernel(backward, adata=adata)
+
+        np.testing.assert_array_equal(mat, pk.transition_matrix)
+        assert pk.backward is backward
+
+    @pytest.mark.parametrize("backward", [False, True])
+    def test_precomputed_str(self, adata: AnnData, backward: bool):
+        key = Key.uns.kernel(backward)
+        adata.obsp[key] = mat = random_transition_matrix(adata.n_obs)
+        pk = PrecomputedKernel(key, adata=adata)
+
+        np.testing.assert_array_equal(mat, pk.transition_matrix)
+        # directionality inferred from key, will not work for `None` (same as key `False`)
+        assert pk.backward is backward
+
+    @pytest.mark.parametrize("backward", [False, True, None])
+    def test_precomputed_obsp(self, adata: AnnData, backward: bool):
+        key = "foo"
+        adata.obsp[key] = mat = random_transition_matrix(adata.n_obs)
+        pk = PrecomputedKernel(adata, obsp_key=key, backward=backward)
+
+        np.testing.assert_array_equal(mat, pk.transition_matrix)
+        # directionality inferred from key
+        assert pk.backward is backward
+
+    def test_precomputed_sum(self, adata: AnnData):
+        mat = random_transition_matrix(adata.n_obs)
+        pk = PrecomputedKernel(mat)
+        vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
+
+        expected = (0.5 * vk.transition_matrix) + (0.5 * pk.transition_matrix)
+        actual = (pk + vk).compute_transition_matrix()
+
+        np.testing.assert_array_almost_equal(expected, actual.transition_matrix)
+
+
 class TestKernelIO:
     @pytest.mark.parametrize("copy", [False, True])
     @pytest.mark.parametrize("write_adata", [False, True])
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index fbc02566f..f62f4cdd2 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -232,7 +232,7 @@ def test_bwd_pipeline_cflare(self, adata: AnnData):
         vk = VelocityKernel(adata, backward=True).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata, backward=True).compute_transition_matrix()
+        ck = ConnectivityKernel(adata).compute_transition_matrix()
         final_kernel = 0.8 * vk + 0.2 * ck
 
         estimator_bwd = cr.tl.estimators.CFLARE(final_kernel)
@@ -294,7 +294,7 @@ def test_bwd_pipeline_gpcca(self, adata: AnnData):
         vk = VelocityKernel(adata, backward=True).compute_transition_matrix(
             softmax_scale=4
         )
-        ck = ConnectivityKernel(adata, backward=True).compute_transition_matrix()
+        ck = ConnectivityKernel(adata).compute_transition_matrix()
         final_kernel = 0.8 * vk + 0.2 * ck
 
         estimator_bwd = cr.tl.estimators.GPCCA(final_kernel)
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 9cb6bb603..cd901edb0 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -2563,11 +2563,9 @@ def test_scvelo_gpcca_abs_probs_cont_not_same(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_transition_matrix_projection(self, mc: GPCCA, fpath: str):
-        mc.kernel.compute_projection(basis="umap")
-        scv.pl.velocity_embedding(
-            mc.kernel.adata,
-            vkey="T_fwd",
+        mc.kernel.plot_projection(
             basis="umap",
+            stream=False,
             arrow_length=6,
             arrow_size=6,
             dpi=DPI,
@@ -3819,8 +3817,7 @@ class TestProjectionEmbedding:
     def test_scvelo_connectivity_kernel_emb_stream(self, adata: AnnData, fpath: str):
         ck = ConnectivityKernel(adata)
         ck.compute_transition_matrix()
-        ck.compute_projection()
-        scv.pl.velocity_embedding_stream(adata, vkey="T_fwd", dpi=DPI, save=fpath)
+        ck.plot_projection(dpi=DPI, save=fpath)
 
     @compare()
     def test_scvelo_pseudotime_kernel_hard_threshold_emb_stream(
@@ -3828,8 +3825,7 @@ def test_scvelo_pseudotime_kernel_hard_threshold_emb_stream(
     ):
         ptk = PseudotimeKernel(adata)
         ptk.compute_transition_matrix(threshold_scheme="hard", frac_to_keep=0.3)
-        ptk.compute_projection()
-        scv.pl.velocity_embedding_stream(adata, vkey="T_fwd", dpi=DPI, save=fpath)
+        ptk.plot_projection(dpi=DPI, save=fpath)
 
     @compare()
     def test_scvelo_pseudotime_kernel_soft_threshold_emb_stream(
@@ -3837,12 +3833,10 @@ def test_scvelo_pseudotime_kernel_soft_threshold_emb_stream(
     ):
         ptk = PseudotimeKernel(adata)
         ptk.compute_transition_matrix(threshold_scheme="soft", frac_to_keep=0.3)
-        ptk.compute_projection()
-        scv.pl.velocity_embedding_stream(adata, vkey="T_fwd", dpi=DPI, save=fpath)
+        ptk.plot_projection(dpi=DPI, save=fpath)
 
     @compare()
     def test_scvelo_velocity_kernel_emb_stream(self, adata: AnnData, fpath: str):
         vk = VelocityKernel(adata)
         vk.compute_transition_matrix()
-        vk.compute_projection()
-        scv.pl.velocity_embedding_stream(adata, vkey="T_fwd", dpi=DPI, save=fpath)
+        vk.plot_projection(dpi=DPI, save=fpath)
diff --git a/tests/test_random_walk.py b/tests/test_random_walk.py
index d39760cd4..7873202ac 100644
--- a/tests/test_random_walk.py
+++ b/tests/test_random_walk.py
@@ -1,6 +1,7 @@
 import pytest
 
-from cellrank.tl.kernels._random_walk import RandomWalk
+from anndata import AnnData
+from cellrank.tl.kernels.utils._random_walk import RandomWalk
 
 import numpy as np
 
@@ -10,17 +11,16 @@ def test_matrix_not_row_stochastic(self, test_matrix_4: np.ndarray):
         with pytest.raises(
             ValueError, match=r"Transition matrix is not row-stochastic."
         ):
-            _ = RandomWalk(test_matrix_4)
+            _ = RandomWalk(AnnData(test_matrix_4), test_matrix_4)
 
     def test_starting_dist_does_not_sum_to_1(self, test_matrix_1: np.ndarray):
-        with pytest.raises(
-            ValueError, match=r"No starting indices have been selected."
-        ):
-            _ = RandomWalk(test_matrix_1, start_ixs=[1000])
+        adata = AnnData(test_matrix_1)
+        with pytest.raises(IndexError, match=r"1000"):
+            _ = RandomWalk(adata, test_matrix_1, start_ixs=[1000])
 
     @pytest.mark.parametrize("kind", ["simulations", "iterations", "hits"])
     def test_invalid_numbers(self, test_matrix_1: np.ndarray, kind: str):
-        rw = RandomWalk(test_matrix_1)
+        rw = RandomWalk(AnnData(test_matrix_1), test_matrix_1)
         with pytest.raises(ValueError, match=kind):
             if kind == "simulations":
                 rw.simulate_many(n_sims=0)
@@ -33,8 +33,8 @@ def test_invalid_numbers(self, test_matrix_1: np.ndarray, kind: str):
 
     @pytest.mark.parametrize("n_sims", [1, 10])
     def test_reproducibility(self, test_matrix_1: np.ndarray, n_sims: int):
-        rw1 = RandomWalk(test_matrix_1)
-        rw2 = RandomWalk(test_matrix_1)
+        rw1 = RandomWalk(AnnData(test_matrix_1), test_matrix_1)
+        rw2 = RandomWalk(AnnData(test_matrix_1), test_matrix_1)
 
         if n_sims == 1:
             r1, r2 = rw1.simulate_one(seed=42), rw2.simulate_one(seed=42)
@@ -44,19 +44,23 @@ def test_reproducibility(self, test_matrix_1: np.ndarray, n_sims: int):
         np.testing.assert_array_equal(r1, r2)
 
     def test_simulate_one(self, test_matrix_1: np.ndarray):
-        res = RandomWalk(test_matrix_1).simulate_one(max_iter=10)
+        res = RandomWalk(AnnData(test_matrix_1), test_matrix_1).simulate_one(
+            max_iter=10
+        )
 
         assert isinstance(res, np.ndarray)
         assert np.issubdtype(res.dtype, np.integer)
         np.testing.assert_array_equal(res.shape, (11,))
 
     def test_start_ixs(self, test_matrix_1: np.ndarray):
-        res = RandomWalk(test_matrix_1, start_ixs=[1]).simulate_one(max_iter=10)
+        adata = AnnData(test_matrix_1)
+        res = RandomWalk(adata, test_matrix_1, start_ixs=[1]).simulate_one(max_iter=10)
 
         assert res[0] == 1
 
     def test_stop_ixs(self, test_matrix_1: np.ndarray):
-        res = RandomWalk(test_matrix_1, stop_ixs=[1]).simulate_one(
+        adata = AnnData(test_matrix_1)
+        res = RandomWalk(adata, test_matrix_1, stop_ixs=[1]).simulate_one(
             max_iter=1000, seed=42
         )
 
@@ -65,9 +69,10 @@ def test_stop_ixs(self, test_matrix_1: np.ndarray):
         assert res[-2] != 1
 
     def test_successive_hits(self, test_matrix_1: np.ndarray):
+        adata = AnnData(test_matrix_1)
         test_matrix_1[0, 0] = 0.2
         test_matrix_1[0, 2] = 0.0
-        res = RandomWalk(test_matrix_1, stop_ixs=[0]).simulate_one(
+        res = RandomWalk(adata, test_matrix_1, stop_ixs=[0]).simulate_one(
             max_iter=1000, seed=42, successive_hits=1
         )
 
@@ -79,7 +84,7 @@ def test_successive_hits(self, test_matrix_1: np.ndarray):
     @pytest.mark.parametrize("n_jobs", [1, 2])
     @pytest.mark.parametrize("backend", ["threading", "loky"])
     def test_simulate_many(self, test_matrix_1: np.ndarray, backend: str, n_jobs: int):
-        res = RandomWalk(test_matrix_1).simulate_many(
+        res = RandomWalk(AnnData(test_matrix_1), test_matrix_1).simulate_many(
             10, max_iter=100, seed=42, backend=backend, n_jobs=n_jobs
         )
 
diff --git a/tests/test_tools.py b/tests/test_tools.py
index 27e789d9d..ec370e8d4 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -8,6 +8,7 @@
 
 import numpy as np
 import pandas as pd
+from scipy.sparse import isspmatrix_csr
 
 
 class TestLineages:
@@ -54,7 +55,7 @@ def test_normal_run(self, adata_cflare: AnnData, use_raw: bool):
             assert np.all(adata_cflare.varm[key][f"{name}_qval"] >= 0)
             assert np.all(adata_cflare.varm[key][f"{name}_qval"] <= 1.0)
 
-    def test_invalid_mode(self, adata_cflare: AnnData):
+    def test_invalid_method(self, adata_cflare: AnnData):
         cr.tl.lineages(adata_cflare)
         with pytest.raises(ValueError):
             cr.tl.lineage_drivers(adata_cflare, use_raw=False, method="foobar")
@@ -167,9 +168,9 @@ def test_invalid_weight(self, adata: AnnData):
         with pytest.raises(ValueError):
             cr.tl.initial_states(adata, weight_connectivities=10)
 
-    def test_invalid_invalid_mode(self, adata: AnnData):
+    def test_invalid_invalid_model(self, adata: AnnData):
         with pytest.raises(ValueError):
-            cr.tl.initial_states(adata, mode="foobar")
+            cr.tl.initial_states(adata, model="foobar")
 
     @pytest.mark.parametrize("force_recompute", [False, True])
     def test_force_recompute(self, adata: AnnData, force_recompute: bool):
@@ -223,8 +224,9 @@ def test_connectivities_conn_key(self, adata: AnnData):
         ck = cr.tl.transition_matrix(adata, weight_connectivities=1, conn_key=key)
 
         assert isinstance(ck, cr.tl.kernels.ConnectivityKernel)
+        assert isspmatrix_csr(ck.transition_matrix)
 
-        np.testing.assert_array_equal(ck.transition_matrix, adata.obsp[key])
+        np.testing.assert_array_equal(ck.transition_matrix.A, adata.obsp[key])
 
     def test_only_velocity(self, adata: AnnData):
         vk = cr.tl.transition_matrix(adata, weight_connectivities=0)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 4d2ea5cba..71e5ff532 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -39,7 +39,7 @@
     _np_apply_along_axis,
     _get_probs_for_zero_vec,
 )
-from cellrank.tl.kernels._velocity_schemes import (
+from cellrank.tl.kernels.utils._similarity import (
     _predict_transition_probabilities_jax,
     _predict_transition_probabilities_numpy,
 )
@@ -897,66 +897,18 @@ def test_one_feature(self):
 
 
 class TestKernelUtils:
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_mean_axis_0(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.mean(x, axis=0), np_mean(x, 0))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_mean_axis_1(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.mean(x, axis=1), np_mean(x, 1))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_max_axis_0(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.max(x, axis=0), np_max(x, 0))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_max_axis_1(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.max(x, axis=1), np_max(x, 1))
+    @pytest.mark.parametrize("fn", ["mean", "max", "sum", "norm"])
+    @pytest.mark.parametrize("axis", [0, 1])
+    def numba_function(self, fn: str, axis: int):
+        numpy_fn = getattr(np, fn)
+        numba_fn = globals()[f"np_{fn}"]
 
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_sum_axis_0(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.sum(x, axis=0), np_sum(x, 0))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_sum_axis_1(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.sum(x, axis=1), np_sum(x, 1))
+        x = np.random.RandomState(42).normal(size=(10, 10))
 
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_norm_axis_0(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
+        np.testing.assert_allclose(numpy_fn(x, axis=axis), numba_fn(x, axis))
 
-        np.testing.assert_allclose(np.linalg.norm(x, axis=0), norm(x, 0))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_numba_norm_axis_1(self, seed):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
-
-        np.testing.assert_allclose(np.linalg.norm(x, axis=1), norm(x, 1))
-
-    @pytest.mark.parametrize("seed", range(5))
-    def test_apply_along_axis(self, seed: int):
-        np.random.seed(seed)
-        x = np.random.normal(size=(10, 10))
+    def test_apply_along_axis(self):
+        x = np.random.RandomState(42).normal(size=(10, 10))
 
         def _create_numba_fn(fn):
             @njit
@@ -1010,7 +962,7 @@ def test_reconstruct_one(self, seed: int, shuffle: bool):
 
     @jax_not_installed_skip
     @pytest.mark.parametrize(
-        "seed,c,s",
+        "seed, c, s",
         zip(range(4), [True, True, False, False], [True, False, True, False]),
     )
     def test_numpy_and_jax(self, seed: int, c: bool, s: bool):