Repository: theislab/cellrank (language: python)

Issue:
Change branch structure, point readthedocs to master, introduce stable
<!-- What kind of feature would you like to request? -->
- [ ] Additional function parameters, change functionality or change defaults?
- [ ] New estimator in `cr.tl.estimators`?
- [ ] New kernel in `cr.tl.kernels`?
- [ ] New gene trend model in `cr.ul.models`?
- [ ] New plotting function in `cr.pl`?
- [ ] Other?

Master will always be pointing to the latest release, dev will contain new developments, documentation on readthedocs will point to master. 



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 0b85609e1..ecb7dd81e 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -6,11 +6,13 @@ on:
     push:
         branches:
         -   master
+        -   dev
         tags:
         -   v*
     pull_request:
         branches:
         -   master
+        -   dev
 
 jobs:
     init:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 7782c8a34..70cd08852 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -54,7 +54,7 @@ repos:
     -   id: autoflake
         args: [--in-place, --remove-unused-variable, --ignore-init-module-imports]
 -   repo: https://gitlab.com/pycqa/flake8
-    rev: 3.8.4
+    rev: 3.9.0
     hooks:
     -   id: flake8
         additional_dependencies: [flake8-tidy-imports, flake8-docstrings, flake8-comprehensions, flake8-bugbear, flake8-blind-except]
@@ -74,7 +74,7 @@ repos:
     -   id: blacken-docs
         additional_dependencies: [black==20.8b1]
 -   repo: https://github.com/asottile/pyupgrade
-    rev: v2.10.0
+    rev: v2.10.1
     hooks:
     -   id: pyupgrade
         args: [--py3-plus, --py37-plus]
diff --git a/.rstcheck.cfg b/.rstcheck.cfg
index 6a28b16e4..63b9ec02c 100644
--- a/.rstcheck.cfg
+++ b/.rstcheck.cfg
@@ -1,3 +1,3 @@
 [rstcheck]
-ignore_messages=(Unknown target name:.*|No (directive|role) entry for "(auto)?(class|method|property|function|func|mod)" in module "docutils\.parsers\.rst\.languages\.en"\.|Duplicate implicit target name: "kernels"\.)
+ignore_messages=(Unknown target name:.*|No (directive|role) entry for "(auto)?(class|method|property|function|func|mod|module)" in module "docutils\.parsers\.rst\.languages\.en"\.|Duplicate implicit target name: "kernels"\.)
 report=info
diff --git a/.scripts/ci/install.sh b/.scripts/ci/install.sh
index c6f3c53e0..4bdf01848 100755
--- a/.scripts/ci/install.sh
+++ b/.scripts/ci/install.sh
@@ -3,16 +3,16 @@
 set -ev
 
 python -m pip install --upgrade pip
-pip install pytest-cov codecov jax jaxlib
+pip install codecov
 
 if [[ "$OS" == "macos-latest" ]]; then
-  pip install -e".[test]"
+  pip install -e".[test,external]"
 elif [[ "$OS" == "ubuntu-latest" ]]; then
   if [[ "$USE_SLEPC" == "true" ]]; then
-      pip install -e".[krylov,test]"
+      pip install -e".[krylov,test,external]"
       python -c "import slepc; import petsc;"
   else
-      pip install -e".[test]"
+      pip install -e".[test,external]"
   fi
   pip install rpy2>=3.3.0
   Rscript --vanilla -e "library('mgcv')"
diff --git a/CONTRIBUTING.rst b/CONTRIBUTING.rst
new file mode 100644
index 000000000..d85b0ffa2
--- /dev/null
+++ b/CONTRIBUTING.rst
@@ -0,0 +1,253 @@
+Contributing guide
+==================
+
+Table of Contents
+=================
+- `Introduction`_
+
+  - `Setup`_
+  - `Codebase structure`_
+
+- `Adding new features`_
+
+  - `Adding external tools`_
+  - `Creating a new kernel`_
+
+- `Running tests`_
+- `Documentation`_
+
+  - `Building documentation`_
+  - `Writing documentation`_
+
+- `Tutorials/examples`_
+- `Making use of GitHub issues/discussions`_
+- `Maintainer notes`_
+
+  - `Making a new release`_
+
+- `Troubleshooting`_
+
+Introduction
+~~~~~~~~~~~~
+
+Setup
+-----
+Clone CellRank from source as::
+
+    git clone https://github.com/theislab/cellrank
+    cd cellrank
+
+Install development requirements in editable mode, as well as pre-commit as::
+
+    pip install -e'.[dev]'
+    pre-commit install
+
+Codebase structure
+------------------
+The CellRank is structured as follows:
+
+- `cellrank <cellrank>`_: the root of the package.
+
+  - `cellrank/datasets <cellrank/datasets>`__: contains datasets that are available for download.
+  - `cellrank/external <cellrank/external>`_: contains all external kernels/estimators/models.
+  - `cellrank/pl <cellrank/pl>`_: the plotting module, containing high level plotting functions.
+  - `cellrank/tl <cellrank/tl>`_: the tools module, containing high-level tool functions, linear solvers, etc.
+
+    - `cellrank/tl/kernels <cellrank/tl/kernels>`_: contains classes that create transition matrices based on
+      various input features, such as transcriptomic or spatial similarities, RNA velocity, pseudotime, etc.
+    - `cellrank/tl/estimators <cellrank/tl/estimators>`_: contains classes with perform computations on Markov chains
+      defined through the transition matrices obtained in the previous step through one or several kernels.
+      These computations usually involve high-level abstractions of the Markov chains, such as coarse-graining
+      into macrostates, finding initial and terminal states, estimating fate probabilities, etc.
+
+  - `cellrank/ul <cellrank/ul>`_: the utilities modules, containing mostly models for gene trend smoothing,
+    parallelization, etc.
+
+    - `cellrank/ul/models <cellrank/ul/models>`_: contains classes that compute smoothed gene expression trends.
+
+- `tests <tests>`_: unit tests, see `Running tests`_ for more information.
+- `docs <docs>`_: documentation, see `Documentation`_ for more.
+- `examples <examples>`__: the examples as seen in the
+  `documentation <https://cellrank.readthedocs.io/en/latest/auto_examples/index.html>`__, see `Tutorials/examples`_
+  for more information.
+
+Adding new features
+~~~~~~~~~~~~~~~~~~~
+We welcome every contribution to CellRank, whether it's a new kernel/estimator/model,
+or a new high-level function/modified functionality. Below, we mostly focus on kernels.
+
+Creating a new kernel
+---------------------
+Creating a new kernel is the primary way to interface with CellRank if you have created a method which computes a
+transition matrix among single-cells. This can be based on any input features that are measured in an experiment, such
+as transcriptional similarity, chromatin accessibility, RNA velocity, experimental timepoints, metabolic labeling, etc.
+
+If you have already implemented your method in a separate package, it can be wrapped in a CellRank kernel which will
+prompt the user to install you package upon initialization, hence boosting your package download numbers.
+Making your package available through a CellRank kernel will give it more visibility and allow users to benefit
+from all of CellRank's downstream functionalities available through estimators.
+
+If you'd like to contribute a new kernel to CellRank, either as a part of an external package or not,
+we have created a short, 5 minute `tutorial <https://cellrank.readthedocs.io/en/latest/creating_new_kernel.html>`_
+that shows you how to do it.
+
+Adding external tools
+---------------------
+If you already have/know of an external package that could be made available through CellRank, we'd love to hear!
+
+In short, external packages should be accessible through ``cellrank.external`` and a thin wrapper code should be placed
+under:
+
+- `cellrank/external/kernels <cellrank/external/kernels>`_: if adding an external kernel. As an example, see the
+  `OTKernel <cellrank/external/kernels/_statot_kernel.py>`_ from `statOT <https://github.com/zsteve/StationaryOT>`_.
+- `cellrank/external/estimators <cellrank/external/estimators>`_: if adding an external estimator.
+- `cellrank/external/models <cellrank/external/models>`_: if adding an external model.
+
+You can either open an `issue <https://github.com/theislab/cellrank/issues/new/choose>`__ with a suggestion or
+directly submit a `PR <https://github.com/theislab/cellrank/pulls>`_ containing the addition.
+Furthermore, external package dependencies should be added to `setup.py's <setup.py>`_ ``'external'``
+in ``extras_require`` and a corresponding entry should be made in the external API
+`documentation <docs/source/external_api.rst>`__.
+
+Running tests
+~~~~~~~~~~~~~
+We use ``tox`` alongside with ``pytest`` to automate the testing process. To run the tests, use the command below::
+
+    tox -e py{37,38,39}-{linux,macos}
+
+depending on the Python version(s) in your ``PATH`` and your operating system. For example, on Linux and Python 3.8,
+you would run::
+
+    tox -e py38-linux
+
+Note that during the first invocation, it can take several minutes to start the testing. This is due to the fact that
+PETSc/SLEPc needs to be built. To run only a subset of tests, run::
+
+    tox -e <environment> -- -k <name>
+
+where ``<name>`` can be a path to a test file/directory or a name of a test function/class, e.g.::
+
+    tox -e py38-linux -- -k tests/test_kernels.py  # run tests in the speciied file
+    tox -e py38-linux -- -k "TestExternal"  # run tests grouped in the `TestExternal` class
+
+Documentation
+~~~~~~~~~~~~~
+
+Building documentation
+----------------------
+In order to build the documentation, run one of the commands below, depending on whether you also want to build the
+examples::
+
+    tox -e docs  # builds the examples, takes longer (~10 mins)
+    tox -e shallow-docs  # does not build the examples
+
+If you need to clean the artifacts from previous documentation builds, run::
+
+    tox -e clean-docs
+
+Writing documentation
+---------------------
+We use ``numpy``-style docstrings for the documentation with the following additions and modifications:
+
+- no type hints in the docstring (optionally applies also for the return statement) should be used,
+  since all functions are required to have the type hints in their signatures.
+- when referring to some argument within the same docstring, enclose that reference in \`\`.
+- when referring to an argument of a class from within that class, use ``:paramref:`attribute```.
+- optional, but recommended: when referring to attributes of a foreign class, use ``:attr:`qualified_name```, such as
+  ``:attr:`anndata.AnnData.obs```.
+- use ``docrep`` for repeating documentation.
+
+Below is an example of how a docstring should look like::
+
+    from cellrank.ul._docs import d
+
+    @d.dedent  # using docrep to interpolate %(adata)s
+    def some_function(adata: AnnData, key: str) -> float:
+        """
+        This is a short one-line header.
+
+        Here you can add multi-paragraph explanation, if needed.
+
+        Parameters
+        ----------
+        %(adata)s
+        key
+            Some key in :attr:`anndata.AnnData.obs`.
+
+        Returns
+        --------
+        Some return description.
+        """
+
+
+Making use of GitHub issues/discussions
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Opening an `issue <https://github.com/theislab/cellrank/issues>`__ or
+starting a `discussion <https://github.com/theislab/cellrank/discussions>`_ is the primary way to get help.
+Issues are used mostly for feature requests or for fixing bugs, whereas discussions can be used to ask conceptual
+questions, algorithmic/biological questions or just to exchange ideas.
+
+Maintainer notes
+~~~~~~~~~~~~~~~~
+
+Making a new release
+--------------------
+New release is always created when a new tag is pushed to GitHub. When that happens, a new CI job starts the
+testing machinery. If all the tests pass, new release will be created on PyPI. Bioconda will automatically notice that
+a new release has been made and an automatic PR will be made to
+`bioconda-recipes <https://github.com/bioconda/bioconda-recipes/pulls>`_.
+Extra care has to be taken when updating runtime dependencies - this is not automatically picked up by Bioconda
+and a separate PR with the updated ``recipe.yaml`` will have to be made.
+
+To make creating new release as easy as possible, we use ``bump2version``, which can be installed as::
+
+    pip install bump2version
+
+Depending on what part of the version you want to update, you can run::
+
+    bump2version {major,minor,patch}
+
+By default, this will create a new tag, automatically update the ``__version__`` wherever necessary and commit the
+changes. Afterwards, you can just push the changes to upstream by running::
+
+    git push --atomic <branch> <tag>
+
+or set ``push.followtags=true`` in your git config and do a regular ``git push``.
+
+Tutorials/examples
+~~~~~~~~~~~~~~~~~~
+While our tutorials focus on an entire workflow or module of CellRank, i.e. using RNA velocity and similarity
+to compute terminal states, examples focus on a single function/method and show how it can be used in practice.
+
+The tutorials are hosted in a separate `repo <https://github.com/theislab/cellrank_notebooks>`_, whereas examples
+are hosted in this repo, under `examples <examples>`__. Both tutorials and examples use already preprocessed datasets
+from `cellrank/datasets <cellrank/datasets>`__, with precomputed attributes, such as velocities, pseudotime, etc.
+
+If you wish to contribute your own example (e.g. for an external kernel), you just need to write a ``.py`` file, similar
+to `this one <examples/other/compute_kernel_tricks.py>`_.
+The filenames should be prefixed with either ``compute_`` or ``plot_``, depending on what they do, i.e. whether they
+show a computational or a plotting functionality.
+
+Tutorials, on the other hand, are written as Jupyter notebooks. However, they are still tested on the CI to make sure
+they run properly with the newest version of CellRank. Since they require more effort to create than the examples,
+it's best to first start a new issue/discussion before adding them, see also `Making use of GitHub issues/discussions`_.
+
+Troubleshooting
+~~~~~~~~~~~~~~~
+- **I have problems with running some tox commands**
+
+  Try recreating the environment as::
+
+    tox -e <environment> --recreate
+
+  If this didn't work, you can purge the whole ``.tox`` directory as ``rm -rf .tox``.
+
+- **I can't commit because of pre-commit**
+
+  Sometimes, it can be hard to satisfy the linting step. You can temporarily bypass it by committing as::
+
+    git commit --no-verify
+
+- **I have an issue which this section does not solve**
+
+  Please see `Making use of GitHub issues/discussions`_ on how to create a new issue or how to start a discussion.
diff --git a/README.rst b/README.rst
index 42b528038..9e9b5d374 100644
--- a/README.rst
+++ b/README.rst
@@ -46,6 +46,8 @@ or via PyPI::
     pip install cellrank
     # or with extra libraries, useful for large datasets
     pip install 'cellrank[krylov]'
+    # or with external modules, see External API
+    pip install 'cellrank[external]'
 
 Why is it called "CellRank"?
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -71,15 +73,15 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
     :alt: Bioconda
 
 .. |CI| image:: https://img.shields.io/github/workflow/status/theislab/cellrank/CI/master
-    :target: https://github.com/theislab/cellrank
+    :target: https://github.com/theislab/cellrank/actions
     :alt: CI
 
 .. |Notebooks| image:: https://img.shields.io/github/workflow/status/theislab/cellrank_notebooks/CI/master?label=notebooks
-    :target: https://github.com/theislab/cellrank_notebooks
+    :target: https://github.com/theislab/cellrank_notebooks/actions
     :alt: CI-Notebooks
 
 .. |Docs|  image:: https://img.shields.io/readthedocs/cellrank
-    :target: https://cellrank.readthedocs.io/en/latest
+    :target: https://cellrank.readthedocs.io/en/stable
     :alt: Documentation
 
 .. |Downloads| image:: https://pepy.tech/badge/cellrank
@@ -91,25 +93,14 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
     :alt: Coverage
 
 .. _preprint: https://doi.org/10.1101/2020.10.19.345983
-
 .. _PageRank: https://en.wikipedia.org/wiki/PageRank#cite_note-1
-
 .. _La Manno et al. (2018): https://doi.org/10.1038/s41586-018-0414-6
-
 .. _Bergen et al. (2020): https://doi.org/10.1038/s41587-020-0591-3
-
 .. _pancreatic endocrinogenesis: https://doi.org/10.1242/dev.173849
-
 .. _scanpy: https://scanpy.readthedocs.io/en/latest/
-
 .. _scvelo: https://scvelo.readthedocs.io/
-
 .. _cellrank.org: https://cellrank.org
-
 .. _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
-
 .. _pyGPCCA: https://pygpcca.readthedocs.io/en/latest/
-
 .. _GPCCA18: https://doi.org/10.1021/acs.jctc.8b00079
diff --git a/cellrank/__init__.py b/cellrank/__init__.py
index 0cd263f0c..52debd01e 100644
--- a/cellrank/__init__.py
+++ b/cellrank/__init__.py
@@ -3,6 +3,7 @@
 import cellrank.ul
 import cellrank.logging
 import cellrank.datasets
+import cellrank.external
 from cellrank.settings import settings
 from cellrank.tl._read import read
 from cellrank.tl._constants import Lin
diff --git a/cellrank/datasets/__init__.py b/cellrank/datasets/__init__.py
index a74f2c284..f2d977e8a 100644
--- a/cellrank/datasets/__init__.py
+++ b/cellrank/datasets/__init__.py
@@ -1,6 +1,7 @@
 from cellrank.datasets._datasets import (
     lung,
     pancreas,
+    zebrafish,
     reprogramming,
     pancreas_preprocessed,
 )
diff --git a/cellrank/datasets/_datasets.py b/cellrank/datasets/_datasets.py
index 0c354875d..22fc1eaef 100644
--- a/cellrank/datasets/_datasets.py
+++ b/cellrank/datasets/_datasets.py
@@ -1,7 +1,7 @@
 """Datasets module."""
 
 import os
-from typing import Tuple, Union, TypeVar
+from typing import Any, Tuple, Union, TypeVar
 from pathlib import Path
 
 from scanpy import read
@@ -36,6 +36,10 @@ class ReprogrammingSubset(ModeEnum):  # noqa: D101
         "https://ndownloader.figshare.com/files/25503773",
         (104679, 22630),
     ),
+    "zebrafish": (
+        "https://ndownloader.figshare.com/files/27265280",
+        (2434, 23974),
+    ),
 }
 
 
@@ -79,7 +83,7 @@ def _load_dataset_from_url(
 @d.dedent
 def pancreas(
     path: Union[str, Path] = "datasets/endocrinogenesis_day15.5.h5ad",
-    **kwargs,
+    **kwargs: Any,
 ) -> AnnData:
     """
     Development of the murine pancreas at E15.5 from [Panc19]_.
@@ -108,7 +112,7 @@ def pancreas(
 @d.dedent
 def pancreas_preprocessed(
     path: Union[str, Path] = "datasets/endocrinogenesis_day15.5_preprocessed.h5ad",
-    **kwargs,
+    **kwargs: Any,
 ) -> AnnData:
     """
     Development of the murine pancreas at E15.5 from [Panc19]_, preprocessed according to the \
@@ -129,7 +133,7 @@ def pancreas_preprocessed(
 @d.dedent
 def lung(
     path: Union[str, Path] = "datasets/lung_regeneration.h5ad",
-    **kwargs,
+    **kwargs: Any,
 ) -> AnnData:
     """
     Regeneration of murine lung epithelial cells at 13 time points from [Lung20]_.
@@ -159,7 +163,7 @@ def lung(
 def reprogramming(
     subset: str = ReprogrammingSubset.FULL.s,
     path: Union[str, Path] = "datasets/reprogramming.h5ad",
-    **kwargs,
+    **kwargs: Any,
 ) -> AnnData:
     """
     Reprogramming of mouse embryonic fibroblasts to induced endoderm progenitors at 8 time points from [Morris18]_.
@@ -206,3 +210,26 @@ def reprogramming(
     raise NotImplementedError(
         f"Subsetting option `{subset.s!r}` is not yet implemented."
     )
+
+
+@d.dedent
+def zebrafish(
+    path: Union[str, Path] = "datasets/zebrafish.h5ad",
+    **kwargs: Any,
+) -> AnnData:
+    """
+    Zebrafish embryogenesis assayed using drop-seq, restricted to the axial mesoderm lineage from [Farrel18]_.
+
+    scRNA-seq time-series dataset comprising `2434` cells
+    which contains 12 time-points spanning 3.3 - 12 hours past fertilization.
+
+    Parameters
+    ----------
+    %(dataset.parameters)s
+
+    Returns
+    -------
+    %(adata)s
+    """
+
+    return _load_dataset_from_url(path, *_datasets["zebrafish"], **kwargs)
diff --git a/cellrank/external/__init__.py b/cellrank/external/__init__.py
new file mode 100644
index 000000000..e08ca9647
--- /dev/null
+++ b/cellrank/external/__init__.py
@@ -0,0 +1,1 @@
+from cellrank.external import models, kernels, estimators
diff --git a/cellrank/external/_error_mixin.py b/cellrank/external/_error_mixin.py
new file mode 100644
index 000000000..bfa049805
--- /dev/null
+++ b/cellrank/external/_error_mixin.py
@@ -0,0 +1,30 @@
+from typing import Any, Optional
+
+
+class ImportErrorMixin:
+    """
+    Mixin class which always raises :class:`ImportError`.
+
+    Subclasses can modify the message by overriding `__import_error_message__`.
+
+    Parameters
+    ----------
+    args
+        Ignored.
+    kwargs
+        Ignored.
+
+    Raises
+    ------
+    ImportError
+        Always.
+    """
+
+    __import_error_message__ = "Unable to import the class."
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        raise ImportError(self.__import_error_message__) from self.__error__
+
+    def __init_subclass__(cls, error: Optional[Exception] = None, **kwargs: Any):
+        super().__init_subclass__(**kwargs)
+        cls.__error__ = error
diff --git a/cellrank/external/estimators/__init__.py b/cellrank/external/estimators/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/cellrank/external/estimators/_import_error_estim.py b/cellrank/external/estimators/_import_error_estim.py
new file mode 100644
index 000000000..028fee767
--- /dev/null
+++ b/cellrank/external/estimators/_import_error_estim.py
@@ -0,0 +1,20 @@
+from typing import Any
+
+from cellrank.tl.estimators import BaseEstimator
+from cellrank.external._error_mixin import ImportErrorMixin
+
+
+class ErroredEstimator(ImportErrorMixin, BaseEstimator):
+    """
+    Utility estimator class which always throw :class:`ImportError` when instantiated.
+
+    Subclasses can modify the message by overriding `__import_error_message__`.
+    """
+
+    __import_error_message__ = "Unable to import the 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
+        raise NotImplementedError
diff --git a/cellrank/external/kernels/__init__.py b/cellrank/external/kernels/__init__.py
new file mode 100644
index 000000000..1052a81ec
--- /dev/null
+++ b/cellrank/external/kernels/__init__.py
@@ -0,0 +1,1 @@
+from cellrank.external.kernels._statot_kernel import OTKernel
diff --git a/cellrank/external/kernels/_import_error_kernel.py b/cellrank/external/kernels/_import_error_kernel.py
new file mode 100644
index 000000000..528b2cb13
--- /dev/null
+++ b/cellrank/external/kernels/_import_error_kernel.py
@@ -0,0 +1,23 @@
+from typing import Any
+
+from cellrank.tl.kernels import Kernel
+from cellrank.external._error_mixin import ImportErrorMixin
+from cellrank.tl.kernels._base_kernel import KernelExpression
+
+
+class ErroredKernel(ImportErrorMixin, Kernel):
+    """
+    Utility kernel class which always throw :class:`ImportError` when instantiated.
+
+    Subclasses can modify the message by overriding `__import_error_message__`.
+    """
+
+    __import_error_message__ = "Unable to import the kernel."
+
+    def compute_transition_matrix(  # noqa: D102
+        self, *args: Any, **kwargs: Any
+    ) -> KernelExpression:
+        raise NotImplementedError
+
+    def copy(self) -> KernelExpression:  # noqa: D102
+        raise NotImplementedError
diff --git a/cellrank/external/kernels/_statot_kernel.py b/cellrank/external/kernels/_statot_kernel.py
new file mode 100644
index 000000000..1c693c775
--- /dev/null
+++ b/cellrank/external/kernels/_statot_kernel.py
@@ -0,0 +1,138 @@
+from typing import Any, Union, Optional
+
+from typing_extensions import Literal
+
+from anndata import AnnData
+
+import numpy as np
+import pandas as pd
+
+from cellrank.ul._docs import d
+from cellrank.tl.estimators import GPCCA
+from cellrank.tl.kernels._precomputed_kernel import DummyKernel
+
+_error = None
+try:
+    from statot import OTKernel as OTKernel_
+except ImportError as e:
+    from cellrank.external.kernels._import_error_kernel import (
+        ErroredKernel as OTKernel_,
+    )
+
+    _error = e
+
+
+@d.dedent
+class OTKernel(OTKernel_, error=_error):
+    """
+    Stationary optimal transport kernel from [Zhang21]_.
+
+    This class requires the `statOT` package, which can be installed as `pip install statot POT`.
+
+    Parameters
+    ----------
+    %(adata)s
+    terminal_states
+        Categorical :class:`pandas.Series` where non-`NaN` values mark terminal states.
+        If `None`, terminal states are assumed to be already present in :paramref:`adata` ``['terminal_states']``.
+    g
+        Key in :attr:`anndata.AnnData.obs` containing relative growth rates for cells or the array itself.
+    cluster_key
+        If a key to cluster labels is given, `terminal_states` will be associated with these for naming and colors.
+    kwargs
+        Additional keyword arguments.
+    """
+
+    __import_error_message__ = "Unable to import the kernel. Please install `statOT` first as `pip install statot POT`."
+
+    def __init__(
+        self,
+        adata: AnnData,
+        g: Union[str, np.ndarray],
+        terminal_states: Optional[Union[str, pd.Series]] = None,
+        cluster_key: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        if terminal_states is not None:
+            dk = DummyKernel(adata, backward=False)
+            estim = GPCCA(dk, write_to_adata=True)
+            estim.set_terminal_states(terminal_states, cluster_key=cluster_key)
+
+        try:
+            super().__init__(adata, g=g, **kwargs)
+        except Exception as e:  # noqa: B902
+            raise RuntimeError("Unable to initialize the kernel.") from e
+
+    def compute_transition_matrix(
+        self,
+        eps: float,
+        dt: float,
+        basis: str = "X_pca",
+        cost_norm_method: Optional[str] = None,
+        method: Literal["ent", "quad", "unbal"] = "ent",
+        tol: float = 0.0,
+        thresh: float = 0.0,
+        maxiter: int = 5000,
+        C: Optional[np.ndarray] = None,
+        verbose: bool = False,
+        **kwargs: Any,
+    ) -> "OTKernel":
+        """
+        Compute transition matrix using stationary OT [Zhang21]_.
+
+        Parameters
+        ----------
+        eps
+            Regularization parameter.
+        dt
+            Choice of the time step over which to fit the model.
+        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
+            :func:`ot.utils.cost_normalization` for more information.
+        method
+            Choice of regularization. Valid options are:
+
+                - `'ent'`: entropy.
+                - `'quad'`: L2-norm.
+                - `'unbal'`: unbalanced transport (not yet implemented).
+
+        tol
+            Relative tolerance for OT solver convergence.
+        thresh
+            Threshold for output transition probabilities.
+        maxiter
+            Maximum number of iterations for OT solver.
+        C
+            Cost matrix for optimal transport problem.
+        verbose
+            Detailed output on convergence of OT solver.
+        kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        :class:`cellrank.external.kernels.OTKernel`
+            Makes :paramref:`transition_matrix` available.
+        """
+        if method not in ("ent", "quad", "unbal"):
+            raise ValueError(f"Invalid method `{method!r}`.")
+        if method == "unbal":
+            raise NotImplementedError("Method `'unbal'` is not yet implemented.")
+
+        super().compute_transition_matrix(
+            eps,
+            dt,
+            expr_key=basis,
+            cost_norm_method=cost_norm_method,
+            method=method,
+            tol=tol,
+            thresh=thresh,
+            maxiter=maxiter,
+            C=C,
+            verbose=verbose,
+            **kwargs,
+        )
+
+        return self
diff --git a/cellrank/external/models/__init__.py b/cellrank/external/models/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/cellrank/external/models/_import_error_model.py b/cellrank/external/models/_import_error_model.py
new file mode 100644
index 000000000..7868dfb03
--- /dev/null
+++ b/cellrank/external/models/_import_error_model.py
@@ -0,0 +1,41 @@
+from typing import Optional
+
+import numpy as np
+
+from cellrank.ul.models import BaseModel
+from cellrank.external._error_mixin import ImportErrorMixin
+
+
+class ErroredModel(ImportErrorMixin, BaseModel):
+    """
+    Utility model class which always throw :class:`ImportError` when instantiated.
+
+    Subclasses can modify the message by overriding `__import_error_message__`.
+    """
+
+    __import_error_message__ = "Unable to import the model."
+
+    def fit(  # noqa: D102
+        self,
+        x: Optional[np.ndarray] = None,
+        y: Optional[np.ndarray] = None,
+        w: Optional[np.ndarray] = None,
+        **kwargs,
+    ) -> "BaseModel":
+        raise NotImplementedError
+
+    def predict(  # noqa: D102
+        self,
+        x_test: Optional[np.ndarray] = None,
+        key_added: Optional[str] = "_x_test",
+        **kwargs,
+    ) -> np.ndarray:
+        raise NotImplementedError
+
+    def confidence_interval(  # noqa: D102
+        self, x_test: Optional[np.ndarray] = None, **kwargs
+    ) -> np.ndarray:
+        raise NotImplementedError
+
+    def copy(self) -> BaseModel:  # noqa: D102
+        raise NotImplementedError
diff --git a/cellrank/pl/_cluster_lineage.py b/cellrank/pl/_cluster_lineage.py
index 58d51f35f..c9b359408 100644
--- a/cellrank/pl/_cluster_lineage.py
+++ b/cellrank/pl/_cluster_lineage.py
@@ -119,6 +119,7 @@ def cluster_lineage(
 
             - ``key`` or ``lineage_{lineage}_trend`` - an :class:`anndata.AnnData` object of
               shape `(n_genes, n_points)` containing the clustered genes.
+
     """
 
     import scanpy as sc
diff --git a/cellrank/pl/_graph.py b/cellrank/pl/_graph.py
index 1f27dd122..7dbf66e4d 100644
--- a/cellrank/pl/_graph.py
+++ b/cellrank/pl/_graph.py
@@ -143,7 +143,7 @@ def graph(
         Title of the figure(s), one for each ``key``.
     %(plotting)s
     layout_kwargs
-        Additional kwargs for ``layout``.
+        Additional keyword arguments for ``layout``.
 
     Returns
     -------
diff --git a/cellrank/pl/_lineages.py b/cellrank/pl/_lineages.py
index 875f9355d..bc51855a3 100644
--- a/cellrank/pl/_lineages.py
+++ b/cellrank/pl/_lineages.py
@@ -73,7 +73,9 @@ def lineage_drivers(
     lineage: str,
     backward: bool = False,
     n_genes: int = 8,
+    ncols: Optional[int] = None,
     use_raw: bool = False,
+    title_fmt: str = "{gene} qval={qval:.4e}",
     **kwargs,
 ) -> None:
     """
@@ -114,4 +116,11 @@ def lineage_drivers(
     drivers.columns = [f"{lineage} corr", f"{lineage} qval"]
     mc._set(A.LIN_DRIVERS, drivers)
 
-    mc.plot_lineage_drivers(lineage, n_genes=n_genes, use_raw=use_raw, **kwargs)
+    mc.plot_lineage_drivers(
+        lineage,
+        n_genes=n_genes,
+        use_raw=use_raw,
+        ncols=ncols,
+        title_fmt=title_fmt,
+        **kwargs,
+    )
diff --git a/cellrank/tl/_constants.py b/cellrank/tl/_constants.py
index 54636b267..93b84d4ed 100644
--- a/cellrank/tl/_constants.py
+++ b/cellrank/tl/_constants.py
@@ -92,6 +92,13 @@ class AbsProbKey(PrettyEnum):
     BACKWARD = f"{DirPrefix.BACKWARD}_{TermStatesKey.BACKWARD}"
 
 
+class ThresholdScheme(PrettyEnum):
+    """Pseudotime thresholding scheme."""
+
+    SOFT = "soft"
+    HARD = "hard"
+
+
 def _transition(d: Union[str, Direction]) -> str:
     return f"T_{d}"
 
diff --git a/cellrank/tl/_lineage.py b/cellrank/tl/_lineage.py
index a32b8d4d9..4a7066581 100644
--- a/cellrank/tl/_lineage.py
+++ b/cellrank/tl/_lineage.py
@@ -733,10 +733,10 @@ def priming_degree(
         method
             The method used to compute the degree of lineage priming. Valid options are:
 
-                - `'kl_divergence'` - as in [Velten17]_, computes KL-divergence between the fate probabilities of a
-                   cell and the average fate probabilities. Computation of average fate probabilities can be restricted
-                   to a set of user-defined ``early_cells``.
-                - `'entropy'` - as in [Setty19]_, computes entropy over a cell's fate probabilities.
+                - `'kl_divergence'`: as in [Velten17]_, computes KL-divergence between the fate probabilities of a
+                  cell and the average fate probabilities. Computation of average fate probabilities can be restricted
+                  to a set of user-defined ``early_cells``.
+                - `'entropy'`: as in [Setty19]_, computes entropy over a cell's fate probabilities.
 
         early_cells
             Cell ids or a mask marking early cells. If `None`, use all cells. Only used when ``method='kl_divergence'``.
diff --git a/cellrank/tl/_utils.py b/cellrank/tl/_utils.py
index 5c42457f7..69219e1c7 100644
--- a/cellrank/tl/_utils.py
+++ b/cellrank/tl/_utils.py
@@ -18,6 +18,8 @@
 
 from statsmodels.stats.multitest import multipletests
 
+from anndata import AnnData
+
 import numpy as np
 import pandas as pd
 from pandas import Series
@@ -44,7 +46,6 @@
 from cellrank.tl._linear_solver import _solve_lin_system
 from cellrank.tl.kernels._utils import np_std, np_mean, _filter_kwargs
 
-AnnData = TypeVar("AnnData")
 ColorLike = TypeVar("ColorLike")
 GPCCA = TypeVar("GPCCA")
 CFLARE = TypeVar("CFLARE")
@@ -277,60 +278,6 @@ def _complex_warning(
     return X_
 
 
-def _bias_knn(
-    conn: csr_matrix, pseudotime: np.ndarray, n_neighbors: int, k: int = 3
-) -> csr_matrix:
-    """
-    Palantir Kernel utility function.
-
-    This function takes in symmetric connectivities and a pseudotime and removes edges that point "against" pseudotime,
-    in this way creating a directed graph. For each node, it always keeps the closest neighbors, making sure the graph
-    remains connected.
-
-    Parameters
-    ----------
-    conn
-        The nearest neighbor connectivities.
-    pseudotime
-        Pseudotemporal ordering of cells.
-    n_neighbors
-        Number of neighbors to keep
-    k
-        Number, alongside with ``n_neighbors`` which determined the threshold for candidate indices.
-
-    Returns
-    -------
-    :class:`scipy.sparse.csr_matrix`.
-        Biased connectivities according to the ``pseudotime``.
-    """
-
-    # set a threshold for the neighbors to keep
-    k_thresh = np.min([int(np.floor(n_neighbors / k)) - 1, 30])
-    conn_biased = conn.copy()
-
-    # loop over rows in the adjacency matrix
-    for i in range(conn.shape[0]):
-
-        # get indices, values and current pseudo t
-        row_data = conn[i, :].data
-        row_ixs = conn[i, :].indices
-        current_t = pseudotime[i]
-
-        # get the 'candidates' - ixs of nodes not in the k_thresh closest neighbors
-        p = np.flip(np.argsort(row_data))
-        sorted_ixs = row_ixs[p]
-        cand_ixs = sorted_ixs[k_thresh:]
-
-        # compare pseudotimes and set indices to zero
-        cand_t = pseudotime[cand_ixs]
-        rem_ixs = cand_ixs[cand_t < current_t]
-        conn_biased[i, rem_ixs] = 0
-
-    conn_biased.eliminate_zeros()
-
-    return conn_biased
-
-
 def _mat_mat_corr_sparse(
     X: csr_matrix,
     Y: np.ndarray,
@@ -441,7 +388,7 @@ def _correlation_test(
         Dataframe of shape ``(n_genes, n_lineages * 5)`` containing the following columns, 1 for each lineage:
 
             - ``{lineage} corr`` - correlation between the gene expression and absorption probabilities.
-            - ``{lineage} pval`` - calulated p-values for double-sided test.
+            - ``{lineage} pval`` - calculated p-values for double-sided test.
             - ``{lineage} qval`` - corrected p-values using Benjamini-Hochberg method at level `0.05`.
             - ``{lineage} ci low`` - lower bound of the ``confidence_level`` correlation confidence interval.
             - ``{lineage} ci high`` - upper bound of the ``confidence_level`` correlation confidence interval.
diff --git a/cellrank/tl/estimators/_base_estimator.py b/cellrank/tl/estimators/_base_estimator.py
index dee1fdb97..9678cc6b4 100644
--- a/cellrank/tl/estimators/_base_estimator.py
+++ b/cellrank/tl/estimators/_base_estimator.py
@@ -3,8 +3,7 @@
 from abc import ABC, abstractmethod
 from sys import version_info
 from copy import copy, deepcopy
-from math import ceil
-from typing import Any, Dict, Union, Mapping, TypeVar, Optional, Sequence
+from typing import Any, Dict, Tuple, Union, Mapping, TypeVar, Optional, Sequence
 from pathlib import Path
 from datetime import datetime
 
@@ -19,12 +18,14 @@
 from scipy.sparse import spmatrix
 from pandas.api.types import infer_dtype, is_categorical_dtype
 
+import matplotlib.pyplot as plt
 from matplotlib.colors import is_color_like
 
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import (
     TestMethod,
+    save_fig,
     _pairwise,
     _irreducible,
     _process_series,
@@ -132,6 +133,9 @@ def __init__(
         if read_from_adata:
             self._read_from_adata()
 
+    def __init_subclass__(cls, **kwargs: Any):
+        super().__init_subclass__()
+
     def _read_from_adata(self) -> None:
         self._set_or_debug(f"eig_{self._direction}", self.adata.uns, "_eig")
 
@@ -178,7 +182,7 @@ def _reconstruct_lineage(self, attr: PrettyEnum, obsm_key: str):
     @inject_docs(fs=P.TERM.s, fsp=P.TERM_PROBS.s)
     def set_terminal_states(
         self,
-        labels: Union[Series, Dict[str, Any]],
+        labels: Union[Series, Dict[str, Sequence[Any]]],
         cluster_key: Optional[str] = None,
         en_cutoff: Optional[float] = None,
         p_thresh: Optional[float] = None,
@@ -195,10 +199,10 @@ def set_terminal_states(
             belonging to a transient state or a :class:`dict`, where each key is the name of the recurrent class and
             values are list of cell names.
         cluster_key
-            If a key to cluster labels is given, :paramref:`{fs}` will ge associated with these for naming and colors.
+            If a key to cluster labels is given, :paramref:`{fs}` will be associated with these for naming and colors.
         %(en_cutoff_p_thresh)s
         add_to_existing
-            Whether to add thses categories to existing ones. Cells already belonging to recurrent classes will be
+            Whether to add these categories to existing ones. Cells already belonging to recurrent classes will be
             updated if there's an overlap.
             Throws an error if previous approximate recurrent classes have not been calculated.
 
@@ -783,7 +787,12 @@ def plot_lineage_drivers(
         self,
         lineage: str,
         n_genes: int = 8,
+        ncols: Optional[int] = None,
         use_raw: bool = False,
+        title_fmt: str = "{gene} qval={qval:.4e}",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
         **kwargs,
     ) -> None:
         """
@@ -794,9 +803,15 @@ def plot_lineage_drivers(
         lineage
             Lineage for which to plot the driver genes.
         n_genes
-            Number of genes to plot.
+            Top most correlated genes to plot.
+        ncols
+            Number of columns.
         use_raw
             Whether to look in :paramref:`adata` ``.raw.var`` or :paramref:`adata` ``.var``.
+        title_fmt
+            Title format. Possible keywords include `{gene}`, `{qval}`, `{corr}` for gene name,
+            q-value and correlation, respectively.
+        %(plotting)s
         kwargs
             Keyword arguments for :func:`scvelo.pl.scatter`.
 
@@ -805,6 +820,25 @@ def plot_lineage_drivers(
         %(just_plots)s
         """
 
+        def prepare_format(
+            gene: str,
+            *,
+            pval: Optional[float],
+            qval: Optional[float],
+            corr: Optional[float],
+        ) -> Dict[str, Any]:
+            kwargs = {}
+            if "{gene" in title_fmt:
+                kwargs["gene"] = gene
+            if "{pval" in title_fmt:
+                kwargs["pval"] = float(pval) if pval is not None else np.nan
+            if "{qval" in title_fmt:
+                kwargs["qval"] = float(qval) if qval is not None else np.nan
+            if "{corr" in title_fmt:
+                kwargs["corr"] = float(corr) if corr is not None else np.nan
+
+            return kwargs
+
         lin_drivers = self._get(P.LIN_DRIVERS)
 
         if lin_drivers is None:
@@ -821,34 +855,47 @@ def plot_lineage_drivers(
         if n_genes <= 0:
             raise ValueError(f"Expected `n_genes` to be positive, found `{n_genes}`.")
 
-        cmap = kwargs.pop("cmap", "viridis")
-        ncols = kwargs.pop("ncols", None)
-
-        geness = lin_drivers.sort_values(by=key, ascending=False).head(n_genes)
-        kwargs.setdefault(
-            "title",
-            [
-                f"{g} q-value={q}"
-                for g, q in dict(geness[f"{lineage} qval"].map("{:.4e}".format)).items()
-            ],
-        )
+        kwargs.pop("save", None)
+        genes = lin_drivers.sort_values(by=key, ascending=False).head(n_genes)
 
-        if ncols is None and len(geness) >= 8:
-            ncols = 4
+        ncols = 4 if ncols is None else ncols
+        nrows = int(np.ceil(len(genes) / ncols))
 
-        # TODO: scvelo can handle only < 20 plots, see https://github.com/theislab/scvelo/issues/252
-        geness = filter(len, np.array_split(geness.index, int(ceil(len(geness) / 8))))
+        fig, axes = plt.subplots(
+            ncols=ncols,
+            nrows=nrows,
+            dpi=dpi,
+            figsize=(ncols * 6, nrows * 4) if figsize is None else figsize,
+        )
+        axes = np.ravel([axes])
 
-        for genes in geness:
+        _i = 0
+        for _i, (gene, ax) in enumerate(zip(genes.index, axes)):
+            data = genes.loc[gene]
             scv.pl.scatter(
                 self.adata,
-                color=genes,
-                cmap=cmap,
+                color=gene,
                 ncols=ncols,
                 use_raw=use_raw,
+                ax=ax,
+                show=False,
+                title=title_fmt.format(
+                    **prepare_format(
+                        gene,
+                        pval=data.get(f"{lineage} pval", None),
+                        qval=data.get(f"{lineage} qval", None),
+                        corr=data.get(f"{lineage} corr", None),
+                    )
+                ),
                 **kwargs,
             )
 
+        for j in range(_i + 1, len(axes)):
+            axes[j].remove()
+
+        if save is not None:
+            save_fig(fig, save)
+
     def _detect_cc_stages(self, rc_labels: Series, p_thresh: float = 1e-15) -> None:
         """
         Detect cell-cycle driven start or endpoints.
@@ -980,13 +1027,34 @@ def _write_terminal_states(self, time=None) -> None:
             f"       `adata.obs[{self._term_key!r}]`\n"
             f"{extra_msg}"
             f"       `.{P.TERM_PROBS}`\n"
-            f"       `.{P.TERM}`",
+            f"       `.{P.TERM}`\n"
+            "    Finish",
             time=time,
         )
 
     @abstractmethod
-    def _fit_terminal_states(self, *args, **kwargs):
-        pass
+    def _fit_terminal_states(self, *args: Any, **kwargs: Any) -> None:
+        """
+        High level API helper method called inside :meth:`fit` that should compute the terminal states.
+
+        This method would usually call :meth:`compute_terminal_states` after all the functions that required beforehand.
+
+        Parameters
+        ----------
+        args
+            Positional arguments.
+        kwargs
+            Keyword arguments.
+
+        Returns
+        -------
+        None
+            Nothing, just sets the terminal states.
+
+        See also
+        --------
+        See :meth:`cellrank.tl.estimators.GPCCA._fit_terminal_states` for an example implementation.
+        """
 
     @d.dedent
     @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS, ap=P.ABS_PROBS, pd=P.PRIME_DEG)
diff --git a/cellrank/tl/estimators/_decomposition.py b/cellrank/tl/estimators/_decomposition.py
index 96cd45dc8..eb5275257 100644
--- a/cellrank/tl/estimators/_decomposition.py
+++ b/cellrank/tl/estimators/_decomposition.py
@@ -33,7 +33,7 @@ def _write_eig_to_adata(
         setattr(self, A.EIG.s, eig)
         self.adata.uns[f"eig_{self._direction}"] = eig
 
-        msg = f"Adding `.{P.EIG}`\n       `adata.uns['eig_{self._direction}']`"
+        msg = f"Adding `adata.uns['eig_{self._direction}']`\n       `.{P.EIG}`"
         if extra_msg is None:
             extra_msg = "\n    Finish"
 
@@ -264,8 +264,6 @@ def plot_spectrum(
         if save:
             save_fig(fig, save)
 
-        fig.show()
-
     def _plot_complex_spectrum(
         self,
         n: int,
diff --git a/cellrank/tl/estimators/_gpcca.py b/cellrank/tl/estimators/_gpcca.py
index 49819b893..4aee26487 100644
--- a/cellrank/tl/estimators/_gpcca.py
+++ b/cellrank/tl/estimators/_gpcca.py
@@ -744,8 +744,6 @@ def annotate_dist_ax(ax, data: np.ndarray, valfmt: str = "{x:.2f}"):
         if save:
             save_fig(fig, save)
 
-        fig.show()
-
     def _compute_one_macrostate(
         self,
         n_cells: int,
diff --git a/cellrank/tl/estimators/_property.py b/cellrank/tl/estimators/_property.py
index 373203192..a762e4838 100644
--- a/cellrank/tl/estimators/_property.py
+++ b/cellrank/tl/estimators/_property.py
@@ -793,17 +793,15 @@ class TerminalStates(Plottable):
     ]
 
     @abstractmethod
-    def set_terminal_states(self, *args, **kwargs) -> None:  # noqa
+    def set_terminal_states(self, *args: Any, **kwargs: Any) -> None:  # noqa
         pass
 
     @abstractmethod
-    def compute_terminal_states(self, *args, **kwargs) -> None:  # noqa
+    def compute_terminal_states(self, *args: Any, **kwargs: Any) -> None:  # noqa
         pass
 
-    # TODO: back compat
-
     @abstractmethod
-    def _write_terminal_states(self, *args, **kwargs) -> None:
+    def _write_terminal_states(self, *args: Any, **kwargs: Any) -> None:
         pass
 
 
diff --git a/cellrank/tl/kernels/__init__.py b/cellrank/tl/kernels/__init__.py
index dd8483c46..d05df49cb 100644
--- a/cellrank/tl/kernels/__init__.py
+++ b/cellrank/tl/kernels/__init__.py
@@ -1,11 +1,17 @@
-from cellrank.tl.kernels._base_kernel import Kernel, Constant
-from cellrank.tl.kernels._palantir_kernel import PalantirKernel
+from cellrank.tl.kernels._base_kernel import Kernel
 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
diff --git a/cellrank/tl/kernels/_base_kernel.py b/cellrank/tl/kernels/_base_kernel.py
index 9c16324ab..8f26b1f7d 100644
--- a/cellrank/tl/kernels/_base_kernel.py
+++ b/cellrank/tl/kernels/_base_kernel.py
@@ -1,5 +1,5 @@
 """Kernel module."""
-
+import warnings
 from abc import ABC, abstractmethod
 from copy import copy
 from typing import (
@@ -13,15 +13,27 @@
     Callable,
     Iterable,
     Optional,
+    Sequence,
 )
+from pathlib import Path
 from functools import reduce
 
+import scvelo as scv
+from scvelo.plotting.utils import default_size, plot_outline
+
 import numpy as np
 from scipy.sparse import spdiags, issparse, spmatrix, csr_matrix, isspmatrix_csr
+from pandas.api.types import infer_dtype
+from pandas.core.dtypes.common import is_categorical_dtype
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import LinearSegmentedColormap, to_hex
+from matplotlib.collections import LineCollection
 
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import (
+    save_fig,
     _connected,
     _normalize,
     _symmetric,
@@ -31,7 +43,8 @@
 )
 from cellrank.ul._utils import Pickleable, _write_graph_data
 from cellrank.tl._constants import Direction, _transition
-from cellrank.tl.kernels._utils import _filter_kwargs
+from cellrank.tl.kernels._utils import _get_basis, _filter_kwargs
+from cellrank.tl.kernels._random_walk import RandomWalk
 
 _ERROR_DIRECTION_MSG = "Can only combine kernels that have the same direction."
 _ERROR_EMPTY_CACHE_MSG = (
@@ -49,6 +62,7 @@
 _dtype = np.float64
 _cond_num_tolerance = 1e-15
 AnnData = TypeVar("AnnData")
+Indices_t = Optional[Union[Sequence[str], Dict[str, Union[str, Sequence[str]]]]]
 
 
 class KernelExpression(Pickleable, ABC):
@@ -69,6 +83,9 @@ def __init__(
         self._normalize = True
         self._parent = None
 
+    def __init_subclass__(cls, **kwargs: Any):
+        super().__init_subclass__()
+
     @property
     def condition_number(self):
         """Condition number of the transition matrix."""
@@ -193,7 +210,11 @@ def write_to_adata(self, key: Optional[str] = None) -> None:
         if key is None:
             key = _transition(self._direction)
 
-        self.adata.uns[f"{key}_params"] = str(self)
+        # retain the embedding info
+        self.adata.uns[f"{key}_params"] = {
+            **self.adata.uns.get(f"{key}_params", {}),
+            **{"params": self.params},
+        }
         _write_graph_data(self.adata, self.transition_matrix, key)
 
     @abstractmethod
@@ -215,6 +236,18 @@ def _maybe_compute_cond_num(self):
             else:
                 logg.info(f"Condition number is `{self._cond_num:.2e}`")
 
+    def _reuse_cache(
+        self, expected_params: Dict[str, Any], *, time: Optional[Any] = None
+    ) -> bool:
+        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
+
+        self._params = expected_params
+        return False
+
     @abstractmethod
     def _get_kernels(self) -> Iterable["Kernel"]:
         pass
@@ -224,6 +257,248 @@ def kernels(self) -> List["Kernel"]:
         """Get the kernels of the kernel expression, except for constants."""
         return list(set(self._get_kernels()))
 
+    def compute_projection(
+        self,
+        basis: str = "umap",
+        key_added: Optional[str] = None,
+        copy: bool = False,
+    ) -> Optional[np.ndarray]:
+        """
+        Compute a projection of the transition matrix in the embedding.
+
+        The projected matrix can be then visualized as::
+
+            scvelo.pl.velocity_embedding(adata, vkey='T_fwd', basis='umap')
+
+        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 :paramref:`adata` ``.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 :paramref:`adata` inplace.
+
+        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 scvelo.tools.velocity_embedding import quiver_autoscale
+
+        if self._transition_matrix is None:
+            raise RuntimeError(
+                "Compute transition matrix first as `.compute_transition_matrix()`."
+            )
+
+        start = logg.info(f"Projecting transition matrix onto `{basis}`")
+        emb = _get_basis(self.adata, basis)
+        T_emb = np.empty_like(emb)
+
+        with warnings.catch_warnings():
+            warnings.simplefilter("ignore")
+            for i, row in enumerate(self.transition_matrix):
+                dX = emb[row.indices] - emb[i, None]
+                if np.any(np.isnan(dX)):
+                    T_emb[i] = np.nan
+                else:
+                    dX /= np.linalg.norm(dX, axis=1)[:, None]
+                    dX = np.nan_to_num(dX)
+                    probs = row.data
+                    T_emb[i] = probs.dot(dX) - probs.mean() * dX.sum(0)
+
+        T_emb /= 3 * quiver_autoscale(np.nan_to_num(emb), T_emb)
+
+        if copy:
+            return T_emb
+
+        key = _transition(self._direction) if key_added is None else key_added
+        ukey = f"{key}_params"
+
+        embs = self.adata.uns.get(ukey, {}).get("embeddings", [])
+        if basis not in emb:
+            embs = list(embs) + [basis]
+            self.adata.uns[ukey] = self.adata.uns.get(ukey, {})
+            self.adata.uns[ukey]["embeddings"] = embs
+
+        key = key + "_" + basis
+        logg.info(
+            f"Adding `adata.obsm[{key!r}]`\n    Finish",
+            time=start,
+        )
+        self.adata.obsm[key] = T_emb
+
+    @d.dedent
+    def plot_random_walks(
+        self,
+        n_sims: int,
+        max_iter: Union[int, float] = 0.25,
+        seed: Optional[int] = None,
+        successive_hits: int = 0,
+        start_ixs: Indices_t = None,
+        stop_ixs: Indices_t = None,
+        basis: str = "umap",
+        cmap: Union[str, LinearSegmentedColormap] = "gnuplot",
+        linewidth: float = 1.0,
+        linealpha: float = 0.3,
+        ixs_legend_loc: Optional[str] = None,
+        n_jobs: Optional[int] = None,
+        backend: str = "loky",
+        show_progress_bar: bool = True,
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Plot random walks in an embedding.
+
+        This method simulates random walks on the Markov chain defined though the corresponding transition matrix. The
+        method is intended to give qualitative rather than quantitative insights into the transition matrix. Random
+        walks are simulated by iteratively choosing the next cell based on the current cell's transition probabilities.
+
+        Parameters
+        ----------
+        n_sims
+            Number of random walks to simulate.
+        %(rw_sim.parameters)s
+        start_ixs
+            Cells from which to sample the starting points. If `None`, use all cells.
+            %(rw_ixs)s
+            For example ``{'clusters': ['Ngn3 low EP', 'Ngn3 high EP']}`` means that starting points for random walks
+            will be samples uniformly from the these clusters.
+        stop_ixs
+            Cells which when hit, the random walk is terminated. If `None`, terminate after ``max_iters``.
+            %(rw_ixs)s
+            For example ``{'clusters': ['Alpha', 'Beta']}`` and ``succesive_hits=3`` means that the random walk will
+            stop prematurely after cells in the above specified clusters have been visited successively 3 times in a
+            row.
+        basis
+            Basis in :attr:`anndata.AnnData.obsm` to use as an embedding.
+        cmap
+            Colormap for the random walk lines.
+        linewidth
+            Width of the random walk lines.
+        linealpha
+            Alpha value of the random walk lines.
+        ixs_legend_loc
+            Legend location for the start/top indices.
+        %(parallel)s
+        %(plotting)s
+        kwargs
+            Keyword arguments for :func:`scvelo.pl.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        For each random walk, the first (last) cell is marked though a black (yellow) dot.
+        """
+
+        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)}`.")
+                cluster_key = next(iter(ixs.keys()))
+                if cluster_key not in self.adata.obs:
+                    raise KeyError(f"Unable to find `adata.obs[{cluster_key!r}]`.")
+                if not is_categorical_dtype(self.adata.obs[cluster_key]):
+                    raise TypeError(f"Expected `adata.obs[{cluster_key!r}]` to be categorical, "
+                                    f"found `{infer_dtype(self.adata.obs[cluster_key])}`.")
+                ixs = np.where(np.isin(self.adata.obs[cluster_key], ixs[cluster_key]))[0]
+                # 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)
+        sims = rw.simulate_many(
+            n_sims=n_sims,
+            max_iter=max_iter,
+            seed=seed,
+            n_jobs=n_jobs,
+            backend=backend,
+            successive_hits=successive_hits,
+            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)
+
     def __xor__(self, other: "KernelExpression") -> "KernelExpression":
         return self.__rxor__(other)
 
@@ -432,14 +707,14 @@ def adata(self, _adata: AnnData):
         self._adata = _adata
 
     def __repr__(self):
-        return f"{'~' if self.backward and self._parent is None else ''}<{self.__class__.__name__[:4]}>"
+        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__[:4]}[{params_fmt}]>"
+                f"<{self.__class__.__name__}[{params_fmt}]>"
             )
         return repr(self)
 
diff --git a/cellrank/tl/kernels/_connectivity_kernel.py b/cellrank/tl/kernels/_connectivity_kernel.py
index 93b1be6d4..7de71ca29 100644
--- a/cellrank/tl/kernels/_connectivity_kernel.py
+++ b/cellrank/tl/kernels/_connectivity_kernel.py
@@ -1,14 +1,11 @@
 """Connectivity kernel module."""
 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._base_kernel import (
-    _LOG_USING_CACHE,
-    _ERROR_EMPTY_CACHE_MSG,
-    AnnData,
-)
 
 
 @d.dedent
@@ -33,9 +30,7 @@ class ConnectivityKernel(Kernel):
     %(backward)s
     conn_key
         Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix, describing cell-cell similarity.
-    compute_cond_num
-        Whether to compute condition number of the transition matrix. Note that this might be costly,
-        since it does not use sparse implementation.
+    %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
     """
@@ -78,23 +73,15 @@ def compute_transition_matrix(
             Makes :paramref:`transition_matrix` available.
         """
 
-        start = logg.info(
-            f"Computing transition matrix based on `adata.obsp[{self._key!r}]`"
-        )
+        # fmt: off
+        start = logg.info(f"Computing transition matrix based on `adata.obsp[{self._key!r}]`")
 
-        params = {"dnorm": density_normalize, "key": self._key}
-        if params == self.params:
-            assert self.transition_matrix is not None, _ERROR_EMPTY_CACHE_MSG
-            logg.debug(_LOG_USING_CACHE)
-            logg.info("    Finish", time=start)
+        if self._reuse_cache({"dnorm": density_normalize, "key": self._key}, time=start):
             return self
 
-        self._params = params
-        self._compute_transition_matrix(
-            matrix=self._conn.copy(), density_normalize=density_normalize
-        )
-
+        self._compute_transition_matrix(matrix=self._conn.copy(), density_normalize=density_normalize)
         logg.info("    Finish", time=start)
+        # fmt: on
 
         return self
 
diff --git a/cellrank/tl/kernels/_cytotrace_kernel.py b/cellrank/tl/kernels/_cytotrace_kernel.py
new file mode 100644
index 000000000..24790696f
--- /dev/null
+++ b/cellrank/tl/kernels/_cytotrace_kernel.py
@@ -0,0 +1,239 @@
+from typing import Any
+
+from typing_extensions import Literal
+
+from anndata import AnnData
+
+import numpy as np
+from scipy.stats import gmean, hmean
+
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import _correlation_test_helper
+from cellrank.tl._constants import ModeEnum
+from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
+
+
+def _ct(key: str) -> str:
+    return f"ct_{key}"
+
+
+class CytoTRACEAggregation(ModeEnum):  # noqa
+    MEAN = "mean"
+    MEDIAN = "median"
+    GMEAN = "gmean"
+    HMEAN = "hmean"
+
+
+@d.dedent
+class CytoTRACEKernel(PseudotimeKernel):
+    """
+    Kernel which computes directed transition probabilities based on a KNN graph and the CytoTRACE score [Cyto20]_.
+
+    The KNN graph contains information about the (undirected) connectivities among cells, reflecting their similarity.
+    CytoTRACE can be used to estimate cellular plasticity and in turn, a pseudotemporal ordering of cells from more
+    plastic to less plastic states.
+    This kernel internally uses the :class:`cellrank.tl.kernels.PseudotimeKernel` to direct the KNN graph
+    on the basis of the CytoTRACE-derived pseudotime.
+
+    %(density_correction)s
+
+    Parameters
+    ----------
+    %(adata)s
+    %(backward)s
+    %(cytotrace.parameters)s
+
+    compute_cond_num
+        Whether to compute condition number of the transition matrix. Note that this might be costly,
+        since it does not use sparse implementation.
+
+    Examples
+    --------
+    Workflow::
+
+        import scvelo as scv
+        import cellrank as cr
+
+        adata = cr.datasets.pancreas()
+
+        sc.pp.filter_genes(adata, min_cells=10)
+        adata.raw = adata.copy()
+        sc.pp.normalize_total(adata)
+        sc.pp.log1p(adata)
+        sc.pp.highly_variable_genes(adata)
+
+        if 'spliced' not in adata.layers or 'unspliced' not in adata.layers:
+            # use the following trick to get scvelo's moments function working
+            adata.layers['spliced'] = adata.X
+            adata.layers['unspliced'] = adata.X
+
+        scv.pp.moments(adata, n_pcs=None, n_neighbors=None)
+    """
+
+    def __init__(
+        self,
+        adata: AnnData,
+        backward: bool = False,
+        layer: str = "Ms",
+        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        use_raw: bool = False,
+        compute_cond_num: bool = False,
+        check_connectivity: bool = False,
+    ):
+        super().__init__(
+            adata,
+            backward=backward,
+            time_key=_ct("pseudotime"),
+            compute_cond_num=compute_cond_num,
+            check_connectivity=check_connectivity,
+            layer=layer,
+            aggregation=aggregation,
+            use_raw=use_raw,
+        )
+        self._time_key = _ct("pseudotime")  # quirk or PT kernel
+
+    def _read_from_adata(
+        self,
+        time_key: str,
+        layer: str = "Ms",
+        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        use_raw: bool = True,
+        **kwargs: Any,
+    ):
+        self.compute_cytotrace(layer=layer, aggregation=aggregation, use_raw=use_raw)
+
+        super()._read_from_adata(time_key=time_key, **kwargs)
+
+    @d.get_sections(base="cytotrace", sections=["Parameters"])
+    def compute_cytotrace(
+        self,
+        layer: str = "Ms",
+        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        use_raw: bool = False,
+    ) -> None:
+        """
+        Re-implementation of the CytoTRACE algorithm [Cyto20]_ to estimate cellular plasticity.
+
+        Computes the number of genes expressed per cell and ranks genes according to their correlation with this
+        measure. Next, it selects to top-correlating genes and aggregates their (imputed) expression to obtain
+        the CytoTRACE score. A high score stands for high differentiation potential (naive, plastic cells) and
+        a low score stands for low differentiation potential (mature, differentiation cells).
+
+        Note that this will not exactly reproduce the results of the original CytoTRACE algorithm [Cyto20]_ because we
+        allow for any normalization and imputation techniques whereas CytoTRACE has build-in specific methods for that.
+
+        Parameters
+        ----------
+        layer
+            Key in :attr:`anndata.AnnData.layers` or `'X'` for :attr:`anndata.AnnData.X`
+            from where to get the expression.
+        aggregation
+            How to aggregate expression of the top-correlating genes. Valid options are:
+
+                - `'mean'`: arithmetic mean.
+                - `'median'`: median.
+                - `'gmean'`: geometric mean.
+                - `'hmean'`: harmonic mean.
+
+        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.
+
+        Returns
+        -------
+        Nothing, just modifies :attr:`anndata.AnnData.obs` with the following keys:
+
+            - `'ct_score'`: the normalized CytoTRACE score.
+            - `'ct_pseudotime'`: associated pseudotime, essentially `1 - CytoTRACE score`.
+            - `'ct_num_exp_genes'`: the number of genes expressed per cell, basis of the CytoTRACE score.
+
+        It also modifies :attr:`anndata.AnnData.var` with the following keys:
+
+            - `'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'`.
+        """
+        # check use_raw
+        aggregation = CytoTRACEAggregation(aggregation)
+        if use_raw and self.adata.raw is None:
+            logg.warning("`adata.raw` is `None`. Setting `use_raw=False`")
+            use_raw = False
+        if use_raw and self.adata.raw.n_vars != self.adata.n_vars:
+            logg.warning(
+                f"`adata.raw` has different number of genes ({self.adata.raw.n_vars}) "
+                f"than `adata` ({self.adata.n_vars}). Setting `use_raw=False`"
+            )
+            use_raw = False
+
+        adata_mraw = self.adata.raw if use_raw else self.adata
+        if layer != "X" and layer not in self.adata.layers:
+            raise KeyError(
+                f"Unable to find `{layer!r}` in `adata.layers`. "
+                f"Valid option are: `{sorted({'X'} | set(self.adata.layers.keys()))}`."
+            )
+
+        msg = f"Computing CytoTRACE score with `{self.adata.n_vars}` genes"
+        if self.adata.n_vars < 10000:
+            msg += ". Consider using more than `10000` genes"
+        start = logg.info(msg)
+
+        # compute number of expressed genes per cell
+        logg.debug(
+            f"Computing number of genes expressed per cell with `use_raw={use_raw}`"
+        )
+        num_exp_genes = np.array((adata_mraw.X > 0).sum(axis=1)).reshape(-1)
+        self.adata.obs[_ct("num_exp_genes")] = num_exp_genes
+
+        # fmt: off
+        # compute correlation with all genes
+        logg.debug("Correlating all genes with number of genes expressed per cell")
+        gene_corr, _, _, _ = _correlation_test_helper(adata_mraw.X.T, num_exp_genes[:, None])
+
+        # annotate the top 200 genes in terms of correlation
+        logg.debug("Finding the top `200` most correlated genes")
+        self.adata.var[_ct("gene_corr")] = gene_corr
+        top_200 = self.adata.var.sort_values(by=_ct("gene_corr"), ascending=False).index[:200]
+        self.adata.var[_ct("correlates")] = False
+        self.adata.var.loc[top_200, _ct("correlates")] = True
+
+        # compute mean/median over top 200 genes, aggregate over genes and shift to [0, 1] range
+        logg.debug(f"Aggregating imputed gene expression using aggregation `{aggregation}` in layer `{layer}`")
+        corr_mask = self.adata.var[_ct("correlates")]
+        imputed_exp = self.adata[:, corr_mask].X if layer == "X" else self.adata[:, corr_mask].layers[layer]
+
+        # aggregate across the top 200 genes
+        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
+
+        # scale to 0-1 range
+        cytotrace_score -= np.min(cytotrace_score)
+        cytotrace_score /= np.max(cytotrace_score)
+        self.adata.obs[_ct("score")] = cytotrace_score
+        self.adata.obs[_ct("pseudotime")] = 1 - cytotrace_score
+
+        self.adata.uns[_ct("params")] = {
+            "aggregation": aggregation.s,
+            "layer": layer,
+            "use_raw": use_raw,
+        }
+
+        logg.info(
+            f"Adding `adata.obs[{_ct('score')!r}]`\n"
+            f"       `adata.obs[{_ct('pseudotime')!r}]`\n"
+            f"       `adata.obs[{_ct('num_exp_genes')!r}]`\n"
+            f"       `adata.var[{_ct('gene_corr')!r}]`\n"
+            f"       `adata.var[{_ct('correlates')!r}]`\n"
+            f"       `adata.uns[{_ct('params')!r}]`\n"
+            f"    Finish",
+            time=start,
+        )
diff --git a/cellrank/tl/kernels/_palantir_kernel.py b/cellrank/tl/kernels/_palantir_kernel.py
deleted file mode 100644
index c19dd9ad0..000000000
--- a/cellrank/tl/kernels/_palantir_kernel.py
+++ /dev/null
@@ -1,176 +0,0 @@
-"""Palantir kernel module."""
-from copy import copy
-
-import numpy as np
-
-from cellrank import logging as logg
-from cellrank.ul._docs import d
-from cellrank.tl._utils import _bias_knn, _connected
-from cellrank.tl.kernels import Kernel
-from cellrank.tl._constants import Direction
-from cellrank.tl.kernels._base_kernel import (
-    _LOG_USING_CACHE,
-    _ERROR_EMPTY_CACHE_MSG,
-    AnnData,
-    _dtype,
-)
-
-
-@d.dedent
-class PalantirKernel(Kernel):
-    """
-    Kernel which computes transition probabilities in a similar way to *Palantir*, see [Setty19]_.
-
-    *Palantir* computes a KNN graph in gene expression space and a pseudotime, which it then uses to direct the edges of
-    the KNN graph, such that they are more likely to point into the direction of increasing pseudotime. To avoid
-    disconnecting the graph, it does not remove all edges that point into the direction of decreasing pseudotime
-    but keeps the ones that point to nodes inside a close radius. This radius is chosen according to the local density.
-
-    The implementation presented here won't exactly reproduce the original *Palantir* algorithm (see below)
-    but the results are qualitatively very similar.
-
-    %(density_correction)s
-
-    Parameters
-    ----------
-    %(adata)s
-    %(backward)s
-    time_key
-        Key in :paramref:`adata` ``.obs`` where the pseudotime is stored.
-    compute_cond_num
-        Whether to compute condition number of the transition matrix. Note that this might be costly,
-        since it does not use sparse implementation.
-    """
-
-    def __init__(
-        self,
-        adata: AnnData,
-        backward: bool = False,
-        time_key: str = "dpt_pseudotime",
-        compute_cond_num: bool = False,
-        check_connectivity: bool = False,
-    ):
-        super().__init__(
-            adata,
-            backward=backward,
-            time_key=time_key,
-            compute_cond_num=compute_cond_num,
-            check_connectivity=check_connectivity,
-        )
-        self._time_key = time_key
-
-    def _read_from_adata(self, **kwargs):
-        super()._read_from_adata(**kwargs)
-
-        time_key = kwargs.pop("time_key", "dpt_pseudotime")
-        if time_key not in self.adata.obs.keys():
-            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 np.any(np.isnan(self._pseudotime)):
-            raise ValueError("Encountered NaN values in pseudotime.")
-
-        logg.debug("Clipping the pseudotime to 0-1 range")
-        self._pseudotime = np.clip(self._pseudotime, 0, 1)
-
-    def compute_transition_matrix(
-        self,
-        k: int = 3,
-        density_normalize: bool = True,
-        check_irreducibility: bool = False,
-    ) -> "PalantirKernel":
-        """
-        Compute transition matrix based on KNN graph and pseudotemporal ordering.
-
-        This is a re-implementation of the Palantir algorithm by [Setty19]_.
-        Note that this won't exactly reproduce the original Palantir results, for three reasons:
-
-            - Palantir computes the KNN graph in a scaled space of diffusion components.
-            - Palantir uses its own pseudotime to bias the KNN graph which is not implemented here.
-            - Palantir uses a slightly different mechanism to ensure the graph remains connected when removing edges
-              that point into the "pseudotime past".
-
-        If you would like to reproduce the original results, please use the original Palantir algorithm.
-
-        Parameters
-        ----------
-        k
-            Number of neighbors to keep for each node, regardless of pseudotime.
-            This is done to ensure that the graph remains connected.
-        density_normalize
-            Whether or not to use the underlying KNN graph for density normalization.
-        check_irreducibility
-            Optional check for irreducibility of the final transition matrix.
-
-        Returns
-        -------
-        :class:`cellrank.tl.kernels.PalantirKernel`
-            Makes :paramref:`transition_matrix` available.
-        """
-
-        start = logg.info("Computing transition matrix based on Palantir-like kernel")
-
-        # get the connectivities and number of neighbors
-        n_neighbors = (
-            self.adata.uns.get("neighbors", {})
-            .get("params", {})
-            .get("n_neighbors", None)
-        )
-        if n_neighbors is None:
-            logg.warning(
-                "Could not find 'n_neighbors' in `adata.uns['neighbors']['params']`. Using an estimate"
-            )
-            n_neighbors = np.min(self._conn.sum(1))
-
-        params = dict(k=k, dnorm=density_normalize, n_neighs=n_neighbors)  # noqa
-        if params == self._params:
-            assert self.transition_matrix is not None, _ERROR_EMPTY_CACHE_MSG
-            logg.debug(_LOG_USING_CACHE)
-            logg.info("    Finish", time=start)
-            return self
-
-        self._params = params
-
-        # handle backward case and run biasing function
-        pseudotime = (
-            np.max(self.pseudotime) - self.pseudotime
-            if self._direction == Direction.BACKWARD
-            else self.pseudotime
-        )
-        biased_conn = _bias_knn(
-            conn=self._conn, pseudotime=pseudotime, n_neighbors=n_neighbors, k=k
-        ).astype(_dtype)
-        # make sure the biased graph is still connected
-        if not _connected(biased_conn):
-            logg.warning("Biased KNN graph is disconnected")
-
-        self._compute_transition_matrix(
-            matrix=biased_conn,
-            density_normalize=density_normalize,
-            check_irreducibility=check_irreducibility,
-        )
-
-        logg.info("    Finish", time=start)
-
-        return self
-
-    @property
-    def pseudotime(self) -> np.array:
-        """Pseudotemporal ordering of cells."""
-        return self._pseudotime
-
-    def copy(self) -> "PalantirKernel":
-        """Return a copy of self."""
-        pk = PalantirKernel(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) -> "PalantirKernel":
-        super().__invert__()
-        self._pseudotime = 1 - self.pseudotime
-        return self
diff --git a/cellrank/tl/kernels/_precomputed_kernel.py b/cellrank/tl/kernels/_precomputed_kernel.py
index 3f6b80548..c5f4f87bd 100644
--- a/cellrank/tl/kernels/_precomputed_kernel.py
+++ b/cellrank/tl/kernels/_precomputed_kernel.py
@@ -2,6 +2,8 @@
 from copy import copy
 from typing import Union, Optional
 
+from anndata import AnnData
+
 import numpy as np
 from scipy.sparse import eye as speye
 from scipy.sparse import spmatrix, csr_matrix
@@ -11,7 +13,7 @@
 from cellrank.ul._utils import _read_graph_data
 from cellrank.tl.kernels import Kernel
 from cellrank.tl._constants import Direction, _transition
-from cellrank.tl.kernels._base_kernel import _RTOL, AnnData, KernelExpression
+from cellrank.tl.kernels._base_kernel import _RTOL, KernelExpression
 
 
 @d.dedent
@@ -27,6 +29,7 @@ class PrecomputedKernel(Kernel):
         If `None`, try to determine the key based on ``backward``.
     %(adata)s
     %(backward)s
+    %(cond_num)s
     """
 
     def __init__(
@@ -128,7 +131,10 @@ def __invert__(self) -> "PrecomputedKernel":
         return self
 
     def __repr__(self):
-        return f"{'~' if self.backward and self._parent is None else ''}<Precomputed[origin={self._origin}]>"
+        return (
+            f"{'~' if self.backward and self._parent is None else ''}"
+            f"<{self.__class__.__name__}[origin={self._origin}]>"
+        )
 
     def __str__(self):
         return repr(self)
diff --git a/cellrank/tl/kernels/_pseudotime_kernel.py b/cellrank/tl/kernels/_pseudotime_kernel.py
new file mode 100644
index 000000000..65ecd8e12
--- /dev/null
+++ b/cellrank/tl/kernels/_pseudotime_kernel.py
@@ -0,0 +1,204 @@
+"""Pseudotime kernel module."""
+from copy import copy
+from typing import Any, Union, Callable, Optional
+
+from typing_extensions import Literal
+
+from anndata import AnnData
+
+import numpy as np
+
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import _connected
+from cellrank.tl.kernels import Kernel
+from cellrank.tl._constants import Direction, ThresholdScheme
+from cellrank.tl.kernels._base_kernel import _dtype
+from cellrank.tl.kernels._pseudotime_schemes import (
+    ThresholdSchemeABC,
+    HardThresholdScheme,
+    SoftThresholdScheme,
+    CustomThresholdScheme,
+)
+
+
+@d.dedent
+class PseudotimeKernel(Kernel):
+    """
+    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 (see
+    [Setty19]_, or to downweight them (see [VIA21]_).
+
+    %(density_correction)s
+
+    Parameters
+    ----------
+    %(adata)s
+    %(backward)s
+    time_key
+        Key in :paramref:`adata` ``.obs`` where the pseudotime is stored.
+    compute_cond_num
+        Whether to compute condition number of the transition matrix. Note that this might be costly,
+        since it does not use sparse implementation.
+    """
+
+    def __init__(
+        self,
+        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):
+        super()._read_from_adata(**kwargs)
+
+        if time_key not in self.adata.obs.keys():
+            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 np.any(np.isnan(self._pseudotime)):
+            raise ValueError("Encountered NaN values in pseudotime.")
+
+    @d.dedent
+    def compute_transition_matrix(
+        self,
+        threshold_scheme: Union[Literal["soft", "hard"], Callable] = "soft",
+        frac_to_keep: float = 0.3,
+        b: float = 20.0,
+        nu: float = 1.0,
+        percentile: Optional[int] = 95,
+        density_normalize: bool = True,
+        check_irreducibility: bool = False,
+        **kwargs: Any,
+    ) -> "PseudotimeKernel":
+        """
+        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 (see [Setty19]_)
+        or VIA (see [VIA21]_).
+
+        When using a `'hard'` thresholding scheme, this based on ideas by *Palantir* (see [Setty19]_) which removes some
+        edges that point against the direction of increasing pseudotime. To avoid disconnecting the graph, it does not
+        remove all edges that point against the direction of increasing pseudotime but keeps the ones that point to
+        cells inside a close radius. This radius is chosen according to the local cell density.
+
+        When using a `'soft'` thresholding scheme, this is based on ideas by *VIA* (see [VIA21]_) which downweights
+        edges that points against the direction 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.
+
+        Parameters
+        ----------
+        frac_to_keep
+            The `fract_to_keep` * n_neighbors 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'`.
+        %(soft_scheme_kernel)s
+        density_normalize
+            Whether or not to use the underlying KNN graph for density normalization.
+        check_irreducibility
+            Optional check for irreducibility of the final transition matrix.
+
+        Returns
+        -------
+        :class:`cellrank.tl.kernels.PseudotimeKernel`
+            Makes :paramref:`transition_matrix` available.
+        """
+        start = logg.info(f"Computing transition matrix based on `{self._time_key}`")
+
+        # get the connectivities and number of neighbors
+        n_neighbors = (
+            self.adata.uns.get("neighbors", {})
+            .get("params", {})
+            .get("n_neighbors", None)
+        )
+        if n_neighbors is None:
+            logg.warning(
+                "Could not find 'n_neighbors' in `adata.uns['neighbors']['params']`. Using an estimate"
+            )
+            n_neighbors = np.min(self._conn.sum(1))
+
+        if isinstance(threshold_scheme, str):
+            threshold_scheme = ThresholdScheme(threshold_scheme)
+            if threshold_scheme == ThresholdScheme.SOFT:
+                scheme = SoftThresholdScheme()
+                kwargs["b"], kwargs["nu"] = b, nu
+            elif threshold_scheme == ThresholdScheme.HARD:
+                scheme = HardThresholdScheme()
+                kwargs["frac_to_keep"], kwargs["n_neighs"] = frac_to_keep, n_neighbors
+            else:
+                raise NotImplementedError(
+                    f"Threshold scheme `{threshold_scheme}` is not yet implemented."
+                )
+        elif isinstance(threshold_scheme, ThresholdSchemeABC):
+            pass
+        elif callable(threshold_scheme):
+            scheme = CustomThresholdScheme(threshold_scheme)
+        else:
+            raise TypeError(
+                f"Expected `threshold_scheme` to be either a `str` or a `callable`, found `{type(threshold_scheme)}`."
+            )
+
+        # fmt: off
+        if self._reuse_cache({"dnorm": density_normalize, "scheme": str(threshold_scheme), **kwargs}, time=start):
+            return self
+
+        # handle backward case and run biasing function
+        pseudotime = (
+            np.max(self.pseudotime) - self.pseudotime
+            if self._direction == Direction.BACKWARD
+            else self.pseudotime
+        )
+
+        biased_conn = scheme.bias_knn(self._conn, pseudotime, **kwargs).astype(_dtype)
+        # fmt: on
+
+        # make sure the biased graph is still connected
+        if not _connected(biased_conn):
+            logg.warning("Biased KNN graph is disconnected")
+
+        self._compute_transition_matrix(
+            matrix=biased_conn,
+            density_normalize=density_normalize,
+            check_irreducibility=check_irreducibility,
+        )
+        logg.info("    Finish", time=start)
+
+        return self
+
+    @property
+    def pseudotime(self) -> 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
diff --git a/cellrank/tl/kernels/_pseudotime_schemes.py b/cellrank/tl/kernels/_pseudotime_schemes.py
new file mode 100644
index 000000000..17f6b34bb
--- /dev/null
+++ b/cellrank/tl/kernels/_pseudotime_schemes.py
@@ -0,0 +1,231 @@
+from abc import ABC, abstractmethod
+from typing import Any, Callable, Optional
+
+import numpy as np
+from scipy.sparse import csr_matrix
+
+from cellrank.ul._docs import d
+
+
+class ThresholdSchemeABC(ABC):
+    """Base class for all connectivity biasing schemes."""
+
+    @d.get_summary(base="pt_scheme")
+    @d.get_sections(base="pt_scheme", sections=["Parameters", "Returns"])
+    @abstractmethod
+    def __call__(
+        self,
+        cell_pseudotime: float,
+        neigh_pseudotime: np.ndarray,
+        neigh_conn: np.ndarray,
+        **kwargs: Any,
+    ) -> np.ndarray:
+        """
+        Calculate biased connections for a given cell.
+
+        Parameters
+        ----------
+        cell_pseudotime
+            Pseudotime of the current cell.
+        neigh_pseudotime
+            Array of shape ``(n_neighbors,)`` containing pseudotimes of neighbors.
+        neigh_conn
+            Array of shape ``(n_neighbors,)`` containing connectivities of the current cell and its neighbors.
+
+        Returns
+        -------
+        Array of shape ``(n_neighbors,)`` containing the biased connectivities.
+        """
+
+    def bias_knn(
+        self, conn: csr_matrix, pseudotime: np.ndarray, **kwargs: Any
+    ) -> csr_matrix:
+        """
+        Bias cell-cell connectivities of a KNN graph.
+
+        Parameters
+        ----------
+        conn
+            Sparse matrix of shape ``(n_cells, n_cells)`` containing the nearest neighbor connectivities.
+        pseudotime
+            Pseudotemporal ordering of cells.
+
+        Returns
+        -------
+        The biased connectivities.
+        """
+        conn_biased = conn.copy()
+
+        for i in range(conn.shape[0]):
+            row, start, end = conn[i], conn.indptr[i], conn.indptr[i + 1]
+
+            biased_row = self(
+                pseudotime[i], pseudotime[row.indices], row.data, **kwargs
+            )
+            if np.shape(biased_row) != row.data.shape:
+                raise ValueError(
+                    f"Expected row of shape `{row.data.shape}`, found `{np.shape(biased_row)}`."
+                )
+            conn_biased.data[start:end] = biased_row
+
+        conn_biased.eliminate_zeros()
+        return conn_biased
+
+    def __repr__(self):
+        return f"<{self.__class__.__name__}>"
+
+    def __str__(self):
+        return repr(self)
+
+
+class HardThresholdScheme(ThresholdSchemeABC):
+    """
+    Thresholding scheme inspired by Palantir [Setty19]_.
+
+    Note that this won't exactly reproduce the original Palantir results, for three reasons:
+
+        - Palantir computes the KNN graph in a scaled space of diffusion components.
+        - Palantir uses its own pseudotime to bias the KNN graph which is not implemented here.
+        - Palantir uses a slightly different mechanism to ensure the graph remains connected when removing edges
+          that point into the "pseudotime past".
+    """
+
+    @d.dedent
+    def __call__(
+        self,
+        cell_pseudotime: float,
+        neigh_pseudotime: np.ndarray,
+        neigh_conn: np.ndarray,
+        n_neighs: int,
+        frac_to_keep: float = 0.3,
+    ) -> np.ndarray:
+        """
+        Convert the undirected graph of cell-cell similarities into a directed one by removing "past" edges.
+
+        This uses a pseudotemporal measure to remove graph-edges that point into the pseudotime-past. For each cell,
+        it keeps the closest neighbors, even if they are in the pseudotime past, to make sure the graph remains
+        connected.
+
+        Parameters
+        ----------
+        %(pt_scheme.parameters)s
+        n_neighs
+            Number of neighbors to keep.
+        frac_to_keep
+            The `fract_to_keep` * n_neighbors closest neighbors (according to graph connectivities) are kept, no matter
+            whether they lie in the pseudotemporal past or future.
+
+        Returns
+        -------
+        %(pt_scheme.returns)s
+        """
+        k_thresh = max(0, min(30, int(np.floor(n_neighs * frac_to_keep)) - 1))
+
+        # below code does not work with argpartition
+        ixs = np.flip(np.argsort(neigh_conn))
+        close_ixs, far_ixs = ixs[:k_thresh], ixs[k_thresh:]
+
+        mask_keep = cell_pseudotime <= neigh_pseudotime[far_ixs]
+        far_ixs_keep = far_ixs[mask_keep]
+
+        biased_conn = np.zeros_like(neigh_conn)
+        biased_conn[close_ixs] = neigh_conn[close_ixs]
+        biased_conn[far_ixs_keep] = neigh_conn[far_ixs_keep]
+
+        return biased_conn
+
+
+class SoftThresholdScheme(ThresholdSchemeABC):
+    """
+    Thresholding scheme inspired by [VIA21]_.
+
+    The idea is to downweight edges that points against the direction 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.
+    """
+
+    @d.dedent
+    def __call__(
+        self,
+        cell_pseudotime: float,
+        neigh_pseudotime: np.ndarray,
+        neigh_conn: np.ndarray,
+        b: float = 20.0,
+        nu: float = 1.0,
+        perc: Optional[int] = 95,
+    ) -> np.ndarray:
+        """
+        Bias the connectivities by downweighting ones to past cells.
+
+        This function uses `generalized logistic regression
+        <https://en.wikipedia.org/wiki/Generalized_logistic_function>`_ to weight the past connectivities.
+
+        Parameters
+        ----------
+        %(pt_scheme.parameters)s
+        %(soft_scheme)s
+
+        Returns
+        -------
+        %(pt_scheme.returns)s
+        """
+        if perc is not None:
+            neigh_conn = np.clip(
+                neigh_conn,
+                np.percentile(neigh_conn, 100 - perc),
+                np.percentile(neigh_conn, perc),
+            )
+
+        past_ixs = np.where(neigh_pseudotime < cell_pseudotime)[0]
+        if not len(past_ixs):
+            return neigh_conn
+
+        weights = np.ones_like(neigh_conn)
+
+        dt = cell_pseudotime - neigh_pseudotime[past_ixs]
+        weights[past_ixs] = 2.0 / ((1.0 + np.exp(b * dt)) ** (1.0 / nu))
+
+        return neigh_conn * weights
+
+
+class CustomThresholdScheme(ThresholdSchemeABC):
+    """
+    Class that wraps a user supplied scheme.
+
+    Parameters
+    ----------
+    callback
+        Function which returns the biased connectivities.
+    """
+
+    def __init__(
+        self,
+        callback: Callable[
+            [float, np.ndarray, np.ndarray, np.ndarray, Any], np.ndarray
+        ],
+    ):
+        super().__init__()
+        self._callback = callback
+
+    @d.dedent
+    def __call__(
+        self,
+        cell_pseudotime: float,
+        neigh_pseudotime: np.ndarray,
+        neigh_conn: np.ndarray,
+        **kwargs: Any,
+    ) -> np.ndarray:
+        """
+        %(pt_scheme.summary)s
+
+        Parameters
+        ----------
+        %(pt_scheme.parameters)s
+        kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        %(pt_scheme.returns)s
+        """  # noqa: D400
+        return self._callback(cell_pseudotime, neigh_pseudotime, neigh_conn, **kwargs)
diff --git a/cellrank/tl/kernels/_random_walk.py b/cellrank/tl/kernels/_random_walk.py
new file mode 100644
index 000000000..fdcd287fb
--- /dev/null
+++ b/cellrank/tl/kernels/_random_walk.py
@@ -0,0 +1,191 @@
+from typing import Any, List, Union, Optional, Sequence
+from itertools import chain
+
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.ul._parallelize import parallelize
+
+
+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:
+        max_iter = (
+            int(np.ceil(max_iter * len(self._ixs)))
+            if isinstance(max_iter, float)
+            else max_iter
+        )
+        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)
+
+        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/_utils.py b/cellrank/tl/kernels/_utils.py
index 889aa6b46..adebfb751 100644
--- a/cellrank/tl/kernels/_utils.py
+++ b/cellrank/tl/kernels/_utils.py
@@ -2,6 +2,8 @@
 from typing import Tuple, Callable, Optional
 from inspect import signature
 
+from anndata import AnnData
+
 import numpy as np
 from numba import njit, prange
 from scipy.sparse import csr_matrix
@@ -233,3 +235,15 @@ def _calculate_starts(indptr: np.ndarray, ixs) -> np.ndarray:
 
     starts = np.cumsum(indptr[ixs + 1] - indptr[ixs])
     return np.hstack((np.array([0], dtype=starts.dtype), starts))
+
+
+def _get_basis(adata: AnnData, basis: str) -> np.ndarray:
+    try:
+        return adata.obsm[f"X_{basis}"]
+    except KeyError:
+        try:
+            return adata.obsm[basis]  # e.g. 'spatial'
+        except KeyError:
+            raise KeyError(
+                f"Unable to find a basis in `adata.obsm['X_{basis}']` or `adata.obsm[{basis!r}]`."
+            ) from None
diff --git a/cellrank/tl/kernels/_velocity_kernel.py b/cellrank/tl/kernels/_velocity_kernel.py
index e8f9c395f..e4a71a419 100644
--- a/cellrank/tl/kernels/_velocity_kernel.py
+++ b/cellrank/tl/kernels/_velocity_kernel.py
@@ -4,6 +4,7 @@
 from math import fsum
 from typing import Any, Union, Callable, Iterable, Optional
 
+from anndata import AnnData
 from scvelo.preprocessing.moments import get_moments
 
 import numpy as np
@@ -24,12 +25,7 @@
     _calculate_starts,
     _get_probs_for_zero_vec,
 )
-from cellrank.tl.kernels._base_kernel import (
-    _RTOL,
-    _LOG_USING_CACHE,
-    _ERROR_EMPTY_CACHE_MSG,
-    AnnData,
-)
+from cellrank.tl.kernels._base_kernel import _RTOL
 from cellrank.tl.kernels._velocity_schemes import Scheme, _get_scheme
 
 
@@ -67,9 +63,7 @@ class VelocityKernel(Kernel):
     gene_subset
         List of genes to be used to compute transition probabilities.
         By default, genes from :paramref:`adata` ``.var['velocity_genes']`` are used.
-    compute_cond_num
-        Whether to compute condition number of the transition matrix. Note that this might be costly,
-        since it does not use sparse implementation.
+    %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
     """
@@ -203,11 +197,6 @@ def compute_transition_matrix(
             raise TypeError(
                 f"Expected `scheme` to be a function, found `{type(scheme)!r}`."
             )
-        elif isinstance(scheme, type):
-            raise TypeError(
-                f"Expected `scheme` to be a function object, found function type `{type(scheme)!r}`. "
-                f"Try instantiating the type."
-            )
 
         if self.backward and mode != VelocityMode.DETERMINISTIC:
             logg.warning(
@@ -236,29 +225,19 @@ def compute_transition_matrix(
 
         if seed is None:
             seed = np.random.randint(0, 2 ** 16)
-        params = dict(  # noqa: C408
-            softmax_scale=softmax_scale, mode=mode.s, seed=seed, scheme=str(scheme)
-        )
+
+        # fmt: off
+        params = {"softmax_scale": softmax_scale, "mode": mode.s, "seed": seed, "scheme": str(scheme)}
         if self.backward:
             params["bwd_mode"] = backward_mode.s
-
-        # check whether we already computed such a transition matrix
-        if params == self._params:
-            assert self._transition_matrix is not None, _ERROR_EMPTY_CACHE_MSG
-            logg.debug(_LOG_USING_CACHE, time=start)
-            logg.info("    Finish", time=start)
+        if self._reuse_cache(params, time=start):
             return self
 
-        self._params = params
-
         # 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)
+        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`")
diff --git a/cellrank/ul/_docs.py b/cellrank/ul/_docs.py
index 3e20647c8..7326cdc02 100644
--- a/cellrank/ul/_docs.py
+++ b/cellrank/ul/_docs.py
@@ -59,7 +59,7 @@
     Match computed states against pre-computed clusters to annotate the states.
     For this, provide a key from :paramref:`adata` ``.obs`` where cluster labels have been computed.
 keys
-    Determines which %(initial_or_terminal) states to use by passing their names.
+    Determines which %(initial_or_terminal)s states to use by passing their names.
     Further, %(initial_or_terminal)s states can be combined. If e.g. the %(terminal)s states are
     ['Neuronal_1', 'Neuronal_1', 'Astrocytes', 'OPC'], then passing ``keys=['Neuronal_1, Neuronal_2', 'OPC']``
     means that the two neuronal %(terminal)s states are treated as one and the 'Astrocyte' state is excluded."""
@@ -124,7 +124,7 @@
 _en_cutoff_p_thresh = """\
 en_cutoff
     If ``cluster_key`` is given, this parameter determines when an approximate recurrent class will
-    be labelled as *'Unknown'*, based on the entropy of the distribution of cells over transcriptomic clusters.
+    be labeled as *'Unknown'*, based on the entropy of the distribution of cells over transcriptomic clusters.
 p_thresh
     If cell cycle scores were provided, a *Wilcoxon rank-sum test* is conducted to identify cell-cycle states.
     If the test returns a positive statistic and a p-value smaller than ``p_thresh``, a warning will be issued."""
@@ -138,12 +138,27 @@
 scheme
     Similarity scheme between cells as described in [Li2020]_. Can be one of the following:
 
-        - `{s.DOT_PRODUCT.s!r}` - :class:`cellrank.tl.kernels.DotProductScheme`.
-        - `{s.COSINE.s!r}` - :class:`cellrank.tl.kernels.CosineScheme`.
-        - `{s.CORRELATION.s!r}` - :class:`cellrank.tl.kernels.CorrelationScheme`.
-
-    Alternatively, any function can be passed as long as it follows the call signature of
-    :class:`cellrank.tl.kernels.SimilaritySchemeABC`."""
+        - `{s.DOT_PRODUCT.s!r}`: :class:`cellrank.tl.kernels.DotProductScheme`.
+        - `{s.COSINE.s!r}`: :class:`cellrank.tl.kernels.CosineScheme`.
+        - `{s.CORRELATION.s!r}`: :class:`cellrank.tl.kernels.CorrelationScheme`.
+
+    Alternatively, any function can be passed as long as it follows the signature of
+    :meth:`cellrank.tl.kernels.SimilaritySchemeABC.__call__`."""
+_cond_num = """\
+compute_cond_num
+    Whether to compute condition number of the transition matrix. Note that this might be costly,
+    since it does not use sparse implementation."""
+_soft_scheme_fmt = """\
+b
+    The growth rate of generalized logistic function.{}
+nu
+    Affects near which asymptote maximum growth occurs.{}
+perc
+    Percentile by which to clip the connectivities.{}"""
+_rw_ixs = """\
+Can be specified as either a :class:`dict` with a key corresponding to cluster key in
+    :attr:`anndata.AnnData.obs` and values to clusters or just a sequence of cell ids
+    in :attr:`anndata.AnnData.obs_names`."""
 
 
 def inject_docs(**kwargs):  # noqa
@@ -193,4 +208,10 @@ def decorator2(obj):
     plots_or_returns_models=_plots_or_returns_models,
     basis=_basis,
     velocity_scheme=_velocity_scheme,
+    cond_num=_cond_num,
+    soft_scheme=_soft_scheme_fmt.format("", "", ""),
+    soft_scheme_kernel=_soft_scheme_fmt.format(
+        *([" Only used when `threshold_scheme='soft'`."] * 3)
+    ),
+    rw_ixs=_rw_ixs,
 )
diff --git a/docs/Makefile b/docs/Makefile
index 48438322a..5e2c9eb66 100644
--- a/docs/Makefile
+++ b/docs/Makefile
@@ -20,8 +20,8 @@ help:
 	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
 
 clean:
-	rm -rf $(BUILDDIR)/*
-	rm -rf $(SOURCEDIR)/api/
-	rm -rf $(SOURCEDIR)/gen_modules/
-	rm -rf $(SOURCEDIR)/auto_examples/
-	rm -rf $(SOURCEDIR)/gen_modules/
+	@rm -rf $(BUILDDIR)/*
+	@rm -rf $(SOURCEDIR)/api/
+	@rm -rf $(SOURCEDIR)/gen_modules/
+	@rm -rf $(SOURCEDIR)/auto_examples/
+	@rm -rf $(SOURCEDIR)/gen_modules/
diff --git a/docs/requirements.txt b/docs/requirements.txt
index a2db5532d..b5f46437f 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,4 +1,4 @@
--r ../requirements.txt
+-r ../requirements.txt  # must include this because of examples
 bezier
 louvain>=0.6,!=0.6.2
 memory_profiler>=0.57.0
diff --git a/docs/source/_ext/edit_on_github.py b/docs/source/_ext/edit_on_github.py
deleted file mode 100644
index c610471e5..000000000
--- a/docs/source/_ext/edit_on_github.py
+++ /dev/null
@@ -1,45 +0,0 @@
-"""
-Loosely based on gist.github.com/MantasVaitkunas/7c16de233812adcb7028
-"""
-
-import os
-import warnings
-
-__licence__ = "BSD (3 clause)"
-
-
-def get_github_repo(app, path):
-    if path.endswith(".ipynb"):
-        return app.config.github_nb_repo, "/tutorials/"
-    return app.config.github_repo, "/docs/source/"
-
-
-def html_page_context(app, pagename, templatename, context, doctree):
-    # doctree is None - otherwise viewcode fails
-    if templatename != "page.html" or doctree is None:
-        return
-
-    if not app.config.github_repo:
-        warnings.warn("`github_repo` not specified")
-        return
-
-    if not app.config.github_nb_repo:
-        nb_repo = f"{app.config.github_repo}_notebooks"
-        warnings.warn("`github_nb_repo `not specified. Setting to `{nb_repo}`")
-        app.config.github_nb_repo = nb_repo
-
-    path = os.path.relpath(doctree.get("source"), app.builder.srcdir)
-    repo, conf_py_path = get_github_repo(app, path)
-
-    # For sphinx_rtd_theme.
-    context["display_github"] = True
-    context["github_user"] = "theislab"
-    context["github_version"] = "master"
-    context["github_repo"] = repo
-    context["conf_py_path"] = conf_py_path
-
-
-def setup(app):
-    app.add_config_value("github_nb_repo", "", True)
-    app.add_config_value("github_repo", "", True)
-    app.connect("html-page-context", html_page_context)
diff --git a/docs/source/_static/img/cellrank_software_package.png b/docs/source/_static/img/cellrank_software_package.png
new file mode 100644
index 000000000..4597eb123
Binary files /dev/null and b/docs/source/_static/img/cellrank_software_package.png differ
diff --git a/docs/source/api.rst b/docs/source/api.rst
index f07844bb2..0db257e80 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -10,12 +10,12 @@ CellRank offers two modes to interact with its core functionality:
 
 - high level mode, essentially calling :func:`cellrank.tl.terminal_states`, :func:`cellrank.tl.initial_states` and
   :func:`cellrank.tl.lineages`.
-  See our `high level tutorial  <https://cellrank.readthedocs.io/en/latest/pancreas_basic.html>`_
+  See our `CellRank basics tutorial <https://cellrank.readthedocs.io/en/stable/cellrank_basics.html>`_.
 - low level mode, interacting directly with the kernels defined in :class:`cellrank.tl.kernels.Kernel` and the
   estimators :class:`cellrank.tl.estimators.GPCCA` or :class:`cellrank.tl.estimators.CFLARE`.
   The division into kernels and estimators ensures that CellRank in broadly applicable, no matter how you have
   computed your transition matrix.
-  See our `low level tutorial <https://cellrank.readthedocs.io/en/latest/pancreas_advanced.html>`_.
+  See our `Kernels and estimators tutorial <https://cellrank.readthedocs.io/en/stable/kernels_and_estimators.html>`_.
 
 Additionally, there is a set of plotting functions which can be used downstream of either analysis mode.
 
@@ -52,7 +52,8 @@ Kernels are part of the low-level API and are used to estimate cell-to-cell tran
 
     tl.kernels.VelocityKernel
     tl.kernels.ConnectivityKernel
-    tl.kernels.PalantirKernel
+    tl.kernels.PseudotimeKernel
+    tl.kernels.CytoTRACEKernel
     tl.kernels.PrecomputedKernel
 
 Estimators
@@ -125,6 +126,7 @@ Datasets
     datasets.pancreas
     datasets.lung
     datasets.reprogramming
+    datasets.zebrafish
     datasets.pancreas_preprocessed
 
 
diff --git a/docs/source/classes.rst b/docs/source/classes.rst
index 897957528..9ffc9b3b9 100644
--- a/docs/source/classes.rst
+++ b/docs/source/classes.rst
@@ -30,19 +30,19 @@ Cosine similarity scheme
 ++++++++++++++++++++++++
 
 .. autoclass:: cellrank.tl.kernels.CosineScheme
-    :members:
+    :members: __call__, hessian
 
 Correlation scheme
 ++++++++++++++++++
 
 .. autoclass:: cellrank.tl.kernels.CorrelationScheme
-    :members:
+    :members: __call__, hessian
 
 Dot product scheme
 ++++++++++++++++++
 
 .. autoclass:: cellrank.tl.kernels.DotProductScheme
-    :members:
+    :members: __call__, hessian
 
 
 Connectivity Kernel
@@ -51,11 +51,31 @@ Connectivity Kernel
 .. autoclass:: cellrank.tl.kernels.ConnectivityKernel
     :members:
 
-Palantir Kernel
----------------
+Pseudotime Kernel
+-----------------
+
+.. autoclass:: cellrank.tl.kernels.PseudotimeKernel
+    :members:
+
+Hard threshold scheme
++++++++++++++++++++++
+
+.. autoclass:: cellrank.tl.kernels.HardThresholdScheme
+    :members:
+    :special-members: __call__
+
+Soft threshold scheme
++++++++++++++++++++++
 
-.. autoclass:: cellrank.tl.kernels.PalantirKernel
+.. autoclass:: cellrank.tl.kernels.SoftThresholdScheme
     :members:
+    :special-members: __call__
+
+CytoTRACE Kernel
+----------------
+
+.. autoclass:: cellrank.tl.kernels.CytoTRACEKernel
+    :members: compute_cytotrace, compute_transition_matrix
 
 Precomputed Kernel
 ------------------
@@ -105,12 +125,20 @@ Kernel
 
 Similarity scheme
 -----------------
+
 .. autoclass:: cellrank.tl.kernels.SimilaritySchemeABC
     :members:
     :special-members: __call__
     :inherited-members:
 
 
+Threshold scheme
+----------------
+.. autoclass:: cellrank.tl.kernels.ThresholdSchemeABC
+    :members:
+    :special-members: __call__
+    :inherited-members:
+
 BaseModel
 ---------
 
@@ -121,4 +149,4 @@ Lineage
 -------
 
 .. autoclass:: cellrank.tl.Lineage
-    :members: reduce, plot_pie, entropy, X, T, view, names, colors
+    :members: priming_degree, reduce, plot_pie, X, T, view, names, colors
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 8659a174a..0a5b02fe7 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -32,12 +32,17 @@
 needs_sphinx = "3.0"
 
 notebooks_url = "https://github.com/theislab/cellrank_notebooks/raw/master/tutorials/"
-for nb in ["pancreas_basic.ipynb", "pancreas_advanced.ipynb"]:
+for nb in [
+    "cellrank_basics.ipynb",
+    "kernels_and_estimators.ipynb",
+    "beyond_rna_velocity.ipynb",
+    "creating_new_kernel.ipynb",
+]:
     try:
         url = urljoin(notebooks_url, nb)
         urlretrieve(url, nb)
     except Exception as e:
-        logger.error(f"Unable to retrieve notebook: `{url}`. Reason: `{e}`.")
+        logger.error(f"Unable to retrieve notebook: `{url}`. Reason: `{e}`")
 
 
 # -- Project information -----------------------------------------------------
@@ -66,13 +71,12 @@
     "nbsphinx",
     "sphinx_copybutton",
     "sphinx_last_updated_by_git",
-    "edit_on_github",
 ]
 
 intersphinx_mapping = dict(
     anndata=("https://anndata.readthedocs.io/en/stable/", None),
     scanpy=("https://scanpy.readthedocs.io/en/stable/", None),
-    squidpy=("https://squidpy.readthedocs.io/", None),
+    squidpy=("https://squidpy.readthedocs.io/en/latest/", None),
     scvelo=("https://scvelo.readthedocs.io/", None),
     python=("https://docs.python.org/3", None),
     numpy=("https://docs.scipy.org/doc/numpy/", None),
@@ -87,6 +91,8 @@
     pygam=("https://pygam.readthedocs.io/en/latest/", None),
     jax=("https://jax.readthedocs.io/en/latest/", None),
     pygpcca=("https://pygpcca.readthedocs.io/en/latest/", None),
+    # external
+    ot=("https://pythonot.github.io/", None),
 )
 
 # Add any paths that contain templates here, relative to this directory.
@@ -248,9 +254,8 @@ def __repr__(self) -> str:
 html_theme = "sphinx_rtd_theme"
 html_static_path = ["_static"]
 html_theme_options = dict(navigation_depth=4, logo_only=True)
-github_repo = "cellrank"  # sets the html_context
-github_nb_repo = "cellrank_notebooks"
 html_show_sphinx = False
+html_show_sourcelink = False
 
 
 def setup(app):
diff --git a/docs/source/external_api.rst b/docs/source/external_api.rst
new file mode 100644
index 000000000..c19b9eef9
--- /dev/null
+++ b/docs/source/external_api.rst
@@ -0,0 +1,34 @@
+External API
+============
+
+This is a collection of external tools that are also accessible through CellRank. If you wish to contribute, please
+refer to our `contributing guide <https://github.com/theislab/cellrank/blob/master/CONTRIBUTING.rst>`_.
+
+Kernels
+~~~~~~~
+
+.. module:: cellrank.external
+.. currentmodule:: cellrank
+
+.. autosummary::
+    :toctree: api
+
+    external.kernels.OTKernel
+
+Estimators
+~~~~~~~~~~
+
+.. module:: cellrank.external
+.. currentmodule:: cellrank
+
+.. autosummary::
+    :toctree: api
+
+Models
+~~~~~~
+
+.. module:: cellrank.external
+.. currentmodule:: cellrank
+
+.. autosummary::
+    :toctree: api
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 996dd6d2e..a349a55f6 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -60,6 +60,7 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
 
     installation
     api
+    external_api
     classes
     release_notes
     references
@@ -76,8 +77,10 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
    :maxdepth: 2
    :hidden:
 
-   pancreas_basic
-   pancreas_advanced
+   cellrank_basics
+   kernels_and_estimators
+   beyond_rna_velocity
+   creating_new_kernel
 
 .. |PyPI| image:: https://img.shields.io/pypi/v/cellrank.svg
     :target: https://pypi.org/project/cellrank
@@ -88,15 +91,15 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
     :alt: Bioconda
 
 .. |CI| image:: https://img.shields.io/github/workflow/status/theislab/cellrank/CI/master
-    :target: https://github.com/theislab/cellrank
+    :target: https://github.com/theislab/cellrank/actions
     :alt: CI
 
 .. |Notebooks| image:: https://img.shields.io/github/workflow/status/theislab/cellrank_notebooks/CI/master?label=notebooks
-    :target: https://github.com/theislab/cellrank_notebooks
+    :target: https://github.com/theislab/cellrank_notebooks/actions
     :alt: CI-Notebooks
 
 .. |Docs|  image:: https://img.shields.io/readthedocs/cellrank
-    :target: https://cellrank.readthedocs.io/en/latest
+    :target: https://cellrank.readthedocs.io/en/stable
     :alt: Documentation
 
 .. |Downloads| image:: https://pepy.tech/badge/cellrank
@@ -108,15 +111,9 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
     :alt: Coverage
 
 .. _preprint: https://www.biorxiv.org/content/10.1101/2020.10.19.345983v1
-
 .. _PageRank: https://en.wikipedia.org/wiki/PageRank#cite_note-1
-
-.. _tutorial: https://cellrank.readthedocs.io/en/latest/pancreas_basic.html
-
-.. _scanpy: https://scanpy.readthedocs.io/en/latest/
-
+.. _tutorial: https://cellrank.readthedocs.io/en/stable/cellrank_basics.html
+.. _scanpy: https://scanpy.readthedocs.io/en/stable/
 .. _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
-
 .. _pyGPCCA: https://pygpcca.readthedocs.io/en/latest/
diff --git a/docs/source/installation.rst b/docs/source/installation.rst
index 26683e1d4..265283677 100644
--- a/docs/source/installation.rst
+++ b/docs/source/installation.rst
@@ -20,6 +20,8 @@ CellRank is also available on PyPI::
     pip install cellrank
     # or with extra libraries, useful for large datasets
     pip install 'cellrank[krylov]'
+    # or with external modules, see External API
+    pip install 'cellrank[external]'
 
 If an error occurs during ``pip install 'cellrank[krylov]'``, please consult the Dependencies_ section below.
 
@@ -27,13 +29,7 @@ Development Version
 ~~~~~~~~~~~~~~~~~~~
 To stay up-to-date with the newest version, run::
 
-    git clone https://github.com/theislab/cellrank
-    cd cellrank
-    pip install -e '.[dev]'
-
-``-e`` stands for ``--editable`` and makes sure that your environment is updated
-when you pull new changes from GitHub. The ``'[dev]'`` options installs requirements
-needed for development, because CellRank is bundled with an additional library.
+    pip install git+https://github.com/theislab/cellrank@dev
 
 Dependencies
 ~~~~~~~~~~~~
@@ -97,7 +93,6 @@ binder button at the top of each tutorial.
 
 If you run into issues, feel free to open a `GitHub issue`_ or send us an `email <mailto:info@cellrank.org>`_ .
 
-
 .. _`Miniconda`: https://conda.pydata.org/miniconda.html
 .. _`GitHub issue`: https://github.com/theislab/cellrank/issues/new
 .. _`binder`: https://mybinder.org/
diff --git a/docs/source/references.rst b/docs/source/references.rst
index 35c2f0c00..6b3313a13 100644
--- a/docs/source/references.rst
+++ b/docs/source/references.rst
@@ -2,21 +2,29 @@ References
 ----------
 
 .. [Bergen20] Bergen, V. *et al.* (2020),
-   *Generalizing RNA velocity to transient cell states through dynamical modeling*,
-   `Nature Biotechnology <https://doi.org/10.1038/s41587-020-0591-3>`__.
+    *Generalizing RNA velocity to transient cell states through dynamical modeling*,
+    `Nature Biotechnology <https://doi.org/10.1038/s41587-020-0591-3>`__.
 
 .. [Coifman05] Coifman *et al.* (2005),
     *Geometric diffusions as a tool for harmonic analysis and structure definition of data: Diffusion maps*,
     `PNAS <https://doi.org/10.1073/pnas.0500334102>`__.
 
+.. [Cyto20] Gulati *et al.* (2020),
+    *Single-cell transcriptional diversity is a hallmark of developmental potential*,
+    `Science <https://doi.org/10.1126/science.aax0249>`__.
+
+.. [Farrel18] Farrel, J. A. *et al.* (2018),
+    *Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis*,
+    `Science <https://doi.org/10.1126/science.aar3131>`__.
+
 .. [GPCCA18] Reuter, B. *et al.* (2018),
-   *Generalized Markov State Modeling Method for Nonequilibrium Biomolecular Dynamics: Exemplified on Amyloid β
-   Conformational Dynamics Driven by an Oscillating Electric Field.*,
-   `Journal of Chemical Theory and Computation <https://doi.org/10.1021/acs.jctc.8b00079>`__.
+    *Generalized Markov State Modeling Method for Nonequilibrium Biomolecular Dynamics: Exemplified on Amyloid β
+    Conformational Dynamics Driven by an Oscillating Electric Field.*,
+    `Journal of Chemical Theory and Computation <https://doi.org/10.1021/acs.jctc.8b00079>`__.
 
 .. [Haghverdi16] Haghverdi *et al.* (2016),
-   *Diffusion pseudotime robustly reconstructs branching cellular lineages*,
-   `Nature Methods <https://doi.org/10.1038/nmeth.3971>`__.
+    *Diffusion pseudotime robustly reconstructs branching cellular lineages*,
+    `Nature Methods <https://doi.org/10.1038/nmeth.3971>`__.
 
 .. [Li2020] Li, T. *et al.* (2020),
     *On the Mathematics of RNA Velocity I: Theoretical Analysis*,
@@ -62,6 +70,10 @@ References
     *Characterization of cell fate probabilities in single-cell data with Palantir*,
     `Nature Biotechnology <https://doi.org/10.1038/s41587-019-0068-4>`__.
 
+.. [VIA21] Shobana V. Stassen *et al.* (2021),
+    *VIA: Generalized and scalable trajectory inference in single-cell omics data*,
+    `bioRxiv <https://doi.org/10.1101/2021.02.10.430705>`__.
+
 .. [Velten17] Velten, L. *et al.* (2017),
     *Human haematopoietic stem cell lineage commitment is a continuous process.*,
     `Nat Cell Biol 19, 271–281 <https://doi.org/10.1038/ncb3493>`__.
@@ -82,3 +94,7 @@ References
     *PAGA: Graph abstraction reconciles clustering with trajectory inference through a topology preserving map of
     single cells*,
     `Genome Biology <https://doi.org/10.1186/s13059-019-1663-x>`__.
+
+.. [Zhang21] Zhang, S. *et al.* (2021),
+    *Optimal transport analysis reveals trajectories in steady-state systems*,
+    `bioRxiv <https://doi.org/10.1101/2021.03.02.433630>`__.
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 9c276eab3..264b5e2e7 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -6,15 +6,41 @@ Release Notes
 Version 1.0
 -----------
 
-1.1.0 :small:`2020-11-17`
+1.2.0 :small:`2021-02-02`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
 This release includes:
 
+.. rubric:: Additions
+
+- Completely **refactored the underlying code base of GPCCA** and set it up as it's own package called
+  `pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_ with documentation and an example. Going forwards, this will
+  ensure that one of the "engines" of CellRank is also easy to maintain to extend. Further, this will make CellRank's
+  installation more convenient by not needing to vendorize additional dependencies
+  `PR 472 <https://github.com/theislab/cellrank/pull/472>`_.
+- Add :func:`cellrank.pl.circular_projection` visualizing computed fate probabilities as done in [Velten17]_,
+  see :ref:`sphx_glr_auto_examples_plotting_plot_circular_embedding.py`.
+  `PR 459 <https://github.com/theislab/cellrank/pull/459>`_.
+- Allow legends not to be plotted by passing ``legend_loc="none"``, as done in `scVelo <https://scvelo.org>`_
+  `PR 470 <https://github.com/theislab/cellrank/pull/470>`_.
+
 .. rubric:: Bugfixes
 
-- Fix not vendorizing correct :mod:`msmtools` which sometimes caused densification of a sparse matrix.
-- Bump scanpy version requirement to 1.6 to fix plotting `PR 444 <https://github.com/theislab/cellrank/pull/444>`_.
+- Fix a bug when computing the Schur decomposition for reducible Markov chains
+  (*Schur vectors appear to not be D-orthogonal*). GPCCA requires the leading Schur vectors to be orthogonal w.r.t. a
+  symmetric, positive definite matrix :math:`D` `PR 453 <https://github.com/theislab/cellrank/pull/453>`_.
+- Fix not falling back to ``mode='monte_carlo'`` if no :mod:`jax` is found when using ``mode='stochastic'`` in
+  :meth:`cellrank.tl.kernels.VelocityKernel.compute_transition_matrix`
+  `PR 472 <https://github.com/theislab/cellrank/pull/472>`_.
+- Fix :mod:`pandas` ``v1.0.1`` indexing error in :func:`cellrank.tl.lineage_drivers`
+  `PR 475 <https://github.com/theislab/cellrank/pull/475>`_.
+- Fix not correctly propagating colors during aggregation in :class:`cellrank.tl.Lineage`
+  `PR 482 <https://github.com/theislab/cellrank/pull/482>`_.
+
+1.1.0 :small:`2020-11-17`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+This release includes:
 
 .. rubric:: Additions
 
@@ -30,6 +56,12 @@ This release includes:
   benchmarking results in our `CellRank preprint <https://doi.org/10.1101/2020.10.19.345983>`_. This is a reprogramming
   dataset from [Morris18]_.
 
+.. rubric:: Bugfixes
+
+- Fix not vendorizing correct :mod:`msmtools` which sometimes caused densification of a sparse matrix.
+- Bump scanpy version requirement to 1.6 to fix plotting `PR 444 <https://github.com/theislab/cellrank/pull/444>`_.
+
+
 1.0.0 :small:`2020-10-17`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 - Fix a bug when subsetting :class:`cellrank.tl.Lineage`
diff --git a/examples/other/compute_kernel_tricks.py b/examples/other/compute_kernel_tricks.py
index ff8b225d7..79ee0fac6 100644
--- a/examples/other/compute_kernel_tricks.py
+++ b/examples/other/compute_kernel_tricks.py
@@ -24,12 +24,12 @@
 #   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 [Manno18]_
 #   [Bergen20]_, but can also take into account uncertainty in RNA velocity.
-# - :class:`cellrank.tl.kernels.PalantirKernel` works similarly as in Palantir [Setty19]_ - it orients the edges of
+# - :class:`cellrank.tl.kernels.PseudotimeKernel` works similarly as in Palantir [Setty19]_ - it orients the edges of
 #   the KNN graph constructed in the expression space using the pseudotemporal ordering of cells, such as
 #   Diffusion Pseudotime (DPT) [Haghverdi16]_.
 ck = cr.tl.kernels.ConnectivityKernel(adata)
 vk = cr.tl.kernels.VelocityKernel(adata)
-pk = cr.tl.kernels.PalantirKernel(adata)
+pk = cr.tl.kernels.PseudotimeKernel(adata)
 
 # %%
 # Kernels can be easily combined with each other. All of the constants within parentheses
diff --git a/examples/plotting/plot_gene_trends.py b/examples/plotting/plot_gene_trends.py
index b059ae9af..bc5593420 100644
--- a/examples/plotting/plot_gene_trends.py
+++ b/examples/plotting/plot_gene_trends.py
@@ -74,7 +74,7 @@
 )
 
 # %%
-# We can also return the models, which can be useful to inspect the fitted models more granularly or when the
+# We can also return the models, which can be useful to inspect the fitted models more granular or when the
 # fitting has failed - such models will be returned as :class:`cellrank.ul.models.FailedModel`.
 #
 # Below we show what would happen if a model were to fail for some arbitrary gene and lineage combination.
diff --git a/requirements.txt b/requirements.txt
index b20fdfa3a..e6326f040 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,7 +10,7 @@ numba>=0.51.0,<0.53.0
 numpy>=1.2.0
 pandas>=1.2.0
 pygam>=0.8.0
-pygpcca>=1.0.0
+pygpcca>=1.0.2
 scanpy>=1.6.0
 scikit-learn>=0.21.3
 scipy>=1.2.0
diff --git a/setup.py b/setup.py
index aef00d33f..ff9800fef 100644
--- a/setup.py
+++ b/setup.py
@@ -19,7 +19,6 @@
         setup_requires=["setuptools_scm"],
         author=__author__,
         author_email=__email__,
-        email=__email__,
         maintainer=__maintainer__,
         maintainer_email=__email__,
         version=__version__,
@@ -39,18 +38,22 @@
             )
         ),
         extras_require=dict(
+            # `POT` is not requirement of `statot`
+            external=["statot>=0.0.14", "POT"],
             krylov=["pygpcca[slepc]"],
             test=[
                 "pytest>=6.1.1",
                 "pytest-mock>=3.5.1",
                 "pytest-xdist>=2.1.0",
+                "pytest-cov",
                 "Pillow",
                 "filelock",
-                "mock>=4.0.2",
                 "python-igraph",
                 "louvain==0.6.1",
                 "leidenalg==0.8.1",
-                "bezier",  # curved edges for `cellrank.pl.graph`
+                "bezier",
+                "jax",
+                "jaxlib",
             ],
             docs=[
                 r
@@ -60,7 +63,7 @@
                 )
                 if "requirements.txt" not in r
             ],
-            dev=["pre-commit>=2.9.0"],
+            dev=["pre-commit>=2.9.3", "tox>=3.23.0"],
         ),
         zip_safe=False,
         packages=find_packages(),
diff --git a/tox.ini b/tox.ini
index 7e55e9574..00d75b6e2 100644
--- a/tox.ini
+++ b/tox.ini
@@ -35,10 +35,9 @@ rst-directives =
 
 [doc8]
 max-line-length = 120
-ignore-path = .tox,docs/source/api,squidpy.egg-info,dist,build
+ignore-path = .tox,docs/source/api,cellrank.egg-info,dist,build
 # D000: D000 Unknown target name:
 ignore-path-errors = docs/source/index.rst;D000,docs/source/latest_additions.rst;D000,docs/source/release_notes.rst;D000
-
 quiet = 1
 
 [pytest]
@@ -46,3 +45,118 @@ python_files = test_*.py
 testpaths = tests/
 xfail_strict = true
 addopts = --ignore _ground_truth_figures --ignore _ground_truth_adatas -n auto -vv
+
+[coverage:run]
+branch = true
+parallel = true
+source = cellrank
+omit = */__init__.py
+
+[coverage:paths]
+source =
+    cellrank
+    */site-packages/cellrank
+
+[coverage:report]
+exclude_lines =
+    \#.*pragma:\s*no.?cover
+
+    ^if __name__ == .__main__.:$
+
+    ^\s*raise AssertionError\b
+    ^\s*raise NotImplementedError\b
+    ^\s*return NotImplemented\b
+show_missing = true
+precision = 2
+skip_empty = True
+sort = Miss
+
+[gh-actions]
+python =
+    3.7: py37
+    3.8: py38
+    3.9: py39
+
+[gh-actions:env]
+PLATFORM =
+    ubuntu-latest: linux
+    macos-latest: macos
+
+[tox]
+isolated_build = True
+envlist =
+    covclean
+    lint
+    py{37,38,39}-{linux,macos}
+    docs
+skip_missing_interpreters = true
+
+[testenv]
+platform =
+    linux: linux
+    macos: (osx|darwin)
+deps =
+    linux: rpy2>=3.3.0
+extras =
+    test
+    krylov
+    external
+passenv = TOXENV CI CODECOV_* GITHUB_ACTIONS PYTEST_FLAGS
+usedevelop = true
+commands =
+    pytest --cov --cov-append --cov-report=term-missing --cov-config={toxinidir}/tox.ini --ignore docs/ {posargs:-vv} {env:PYTEST_FLAGS:}
+
+[testenv:covclean]
+description = Clean coverage files.
+deps = coverage
+skip_install = True
+commands = coverage erase
+
+[testenv:coverage]
+description = Report the coverage difference.
+deps =
+    coverage
+    diff_cover
+skip_install = true
+depends = py{37,38,39}-{linux,macos}
+parallel_show_output = True
+commands =
+    coverage report --omit="tox/*"
+    coverage xml --omit="tox/*" -o {toxinidir}/coverage.xml
+    diff-cover --compare-branch origin/master {toxinidir}/coverage.xml
+
+[testenv:lint]
+description = Perform linting.
+basepython = python3.8
+deps = pre-commit>=2.9.3
+skip_install = true
+commands = pre-commit run --all-files --show-diff-on-failure {posargs:}
+
+[testenv:clean-docs]
+description = Clean the documentation artifacts.
+basepython = python3.8
+deps =
+skip_install = true
+changedir = {toxinidir}/docs
+whitelist_externals = make
+commands = make clean
+
+[testenv:docs]
+description = Build the documentation.
+basepython = python3.8
+skip_install = true
+deps = -r{toxinidir}/docs/requirements.txt
+whitelist_externals = sphinx-build
+commands =
+    sphinx-build --color -b html {toxinidir}/docs/source {toxinidir}/docs/build/html
+    python -c 'import pathlib; print(f"Documentation is available under:", pathlib.Path(f"{toxinidir}") / "docs" / "build" / "html" / "index.html")'
+
+[testenv:shallow-docs]
+description = Build the documentation without running the examples.
+basepython = python3.8
+skip_install = true
+deps = -r{toxinidir}/docs/requirements.txt
+whitelist_externals = sphinx-build
+commands =
+    sphinx-build --color -D plot_gallery=0 -b html {toxinidir}/docs/source {toxinidir}/docs/build/html
+    python -c 'import pathlib; print(f"Documentation is available under:", pathlib.Path(f"{toxinidir}") / "docs" / "build" / "html" / "index.html")'


Test patch:
diff --git a/docs/source/latest_additions.rst b/docs/source/latest_additions.rst
index c4eb5c4e2..bf155022e 100644
--- a/docs/source/latest_additions.rst
+++ b/docs/source/latest_additions.rst
@@ -1,32 +1,52 @@
 .. role:: small
 
-1.2.0 :small:`2021-02-02`
+1.3.0 :small:`2021-03-29`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
 
-This release includes:
+This release includes some major additions which make CellRank more applicable with and without RNA velocity
+information. In particular, it includes:
 
-.. rubric:: Bugfixes
+.. rubric:: Additions
 
-- Fix a bug when computing the Schur decomposition for reducible Markov chains
-  (*Schur vectors appear to not be D-orthogonal*). GPCCA requires the leading Schur vectors to be orthogonal w.r.t. a
-  symmetric, positive definite matrix :math:`D` `PR 453 <https://github.com/theislab/cellrank/pull/453>`_.
-- Fix not falling back to ``mode='monte_carlo'`` if no :mod:`jax` is found when using ``mode='stochastic'`` in
-  :meth:`cellrank.tl.kernels.VelocityKernel.compute_transition_matrix`
-  `PR 472 <https://github.com/theislab/cellrank/pull/472>`_.
-- Fix :mod:`pandas` ``v1.0.1`` indexing error in :func:`cellrank.tl.lineage_drivers`
-  `PR 475 <https://github.com/theislab/cellrank/pull/475>`_.
-- Fix not correctly propagating colors during aggregation in :class:`cellrank.tl.Lineage`
-  `PR 482 <https://github.com/theislab/cellrank/pull/482>`_.
+- Add new kernel :class:`cellrank.tl.kernels.CytoTRACEKernel` which computes cell-cell transition probabilities based
+  on the CytoTRACE score [Cyto20]_, a measure of differentiation potential,
+  `PR 527 <https://github.com/theislab/cellrank/pull/527>`_.
+- Add external API :mod:`cellrank.external` with a stationary optimal transport kernel
+  :class:`cellrank.external.kernels.OTKernel` contributed from [Zhang21]_, as well as a
+  `contributing guide <https://github.com/theislab/cellrank/blob/master/CONTRIBUTING.rst>`_,
+  `PR 522 <https://github.com/theislab/cellrank/pull/522>`_.
+- Rename ``cellrank.tl.kernels.PalantirKernel`` to :class:`cellrank.tl.kernels.PseudotimeKernel` and add
+  hard thresholding scheme inspired by [Setty19]_, a soft thresholding scheme inspired by [VIA21]_ and a custom scheme
+  when computing the transition matrix, see e.g. :class:`cellrank.tl.kernels.SoftThresholdScheme`
+  `PR 514 <https://github.com/theislab/cellrank/pull/514>`_.
+- Add more flexibility to :class:`cellrank.tl.kernels.ConnectivityKernel`, allowing it to use any cell-cell similarities
+  from :attr:`anndata.AnnData.obsp`, such as spatial similarities from :mod:`squidpy` [Palla21]_
+  `PR 501 <https://github.com/theislab/cellrank/pull/501>`_.
+- Revamp `Pancreas Advanced <https://cellrank.readthedocs.io/en/latest/pancreas_advanced.html>`_ tutorial
+  to showcase CellRank's modular structure of kernels and estimators.
+  `PR 32 <https://github.com/theislab/cellrank_notebooks/pull/32>`_.
+- Add 2 new tutorials:
 
-.. rubric:: Additions
+  - `Beyond RNA velocity <https://cellrank.readthedocs.io/en/latest/beyond_rna_velocity.html>`_: shows how to use
+    CellRank when no RNA velocity information is available.
+    `PR 32 <https://github.com/theislab/cellrank_notebooks/pull/32>`_
+  - `Creating a new kernel <https://cellrank.readthedocs.io/en/latest/creating_new_kernel.html>`_: explains how to
+    create your own custom kernel class that estimates cell-cell transition probabilities
+    `PR 31 <https://github.com/theislab/cellrank_notebooks/pull/31>`_.
+
+- Add projection of transition matrix onto an embedding :meth:`cellrank.tl.kernels.Kernel.compute_projection`
+- Add random walk simulation and visualization in an embedding :meth:`cellrank.tl.kernels.Kernel.plot_random_walks`
+  `PR 537 <https://github.com/theislab/cellrank/pull/537>`_.
+- Add :meth:`cellrank.tl.Lineage.priming_degree` `PR 502 <https://github.com/theislab/cellrank/pull/502>`_
+  which estimates a cell's plasticity/differentiation potential based on ideas by [Setty19]_ and [Velten17]_.
+- Add checks for transition matrix irreducibility `PR 516 <https://github.com/theislab/cellrank/pull/516>`_.
+- Add Zebrafish development dataset from [Farrel18]_ `PR 539 <https://github.com/theislab/cellrank/pull/539>`_.
+- Speed-up stationary distribution calculation in :mod:`pygpcca` `PR 22 <https://github.com/msmdev/pyGPCCA/pull/22>`_.
+
+.. rubric:: Bugfixes
 
-- Completely **refactored the underlying code base of GPCCA** and set it up as it's own package called
-  `pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_ with documentation and an example. Going forwards, this will
-  ensure that one of the "engines" of CellRank is also easy to maintain to extend. Further, this will make CellRank's
-  installation more convenient by not needing to vendorize additional dependencies
-  `PR 472 <https://github.com/theislab/cellrank/pull/472>`_.
-- Add :func:`cellrank.pl.circular_projection` visualizing computed fate probabilities as done in [Velten17]_,
-  see :ref:`sphx_glr_auto_examples_plotting_plot_circular_embedding.py`.
-  `PR 459 <https://github.com/theislab/cellrank/pull/459>`_.
-- Allow legends not to be plotted by passing ``legend_loc="none"``, as done in `scVelo <https://scvelo.org>`_
-  `PR 470 <https://github.com/theislab/cellrank/pull/470>`_.
+- Fix various bugs when plotting multiple gene trends `PR 487 <https://github.com/theislab/cellrank/pull/487>`_.
+- Fix gene trend smoothing not working for 1 lineage `PR 512 <https://github.com/theislab/cellrank/pull/512>`_.
+- Fix :mod:`pandas` error when computing macrostates `PR 513 <https://github.com/theislab/cellrank/pull/513>`_.
+- Remove malfunctioning *Edit on GitHub* from the documentation
+  `PR 538 <https://github.com/theislab/cellrank/pull/538>`_.
diff --git a/tests/_ground_truth_adatas/adata_200.h5ad b/tests/_ground_truth_adatas/adata_200.h5ad
index aed46a8c4..13892fe1f 100644
Binary files a/tests/_ground_truth_adatas/adata_200.h5ad and b/tests/_ground_truth_adatas/adata_200.h5ad differ
diff --git a/tests/_ground_truth_figures/drivers_backward.png b/tests/_ground_truth_figures/drivers_backward.png
new file mode 100644
index 000000000..b02731242
Binary files /dev/null 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
new file mode 100644
index 000000000..332fde86c
Binary files /dev/null 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
new file mode 100644
index 000000000..2a282bfae
Binary files /dev/null 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
new file mode 100644
index 000000000..8e167a359
Binary files /dev/null and b/tests/_ground_truth_figures/drivers_title_fmt.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 d28bb68b7..cf66c2be7 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/kernel_random_walk_basis.png b/tests/_ground_truth_figures/kernel_random_walk_basis.png
new file mode 100644
index 000000000..e30a6ea56
Binary files /dev/null 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
new file mode 100644
index 000000000..1bac74045
Binary files /dev/null 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
new file mode 100644
index 000000000..a4d7a072d
Binary files /dev/null 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
new file mode 100644
index 000000000..10b6812dd
Binary files /dev/null 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
new file mode 100644
index 000000000..7f511aa4f
Binary files /dev/null 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
new file mode 100644
index 000000000..e8d18630d
Binary files /dev/null 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
new file mode 100644
index 000000000..243249944
Binary files /dev/null and b/tests/_ground_truth_figures/kernel_random_walk_params.png differ
diff --git a/tests/_ground_truth_figures/scvelo_drivers_cmap.png b/tests/_ground_truth_figures/scvelo_drivers_cmap.png
deleted file mode 100644
index a70e96d1d..000000000
Binary files a/tests/_ground_truth_figures/scvelo_drivers_cmap.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_drivers_n_genes.png b/tests/_ground_truth_figures/scvelo_drivers_n_genes.png
deleted file mode 100644
index 0598dc8f4..000000000
Binary files a/tests/_ground_truth_figures/scvelo_drivers_n_genes.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png b/tests/_ground_truth_figures/scvelo_transition_matrix_projection.png
new file mode 100644
index 000000000..53f8c2366
Binary files /dev/null and b/tests/_ground_truth_figures/scvelo_transition_matrix_projection.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 37faad74a..a35c1c06a 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/_helpers.py b/tests/_helpers.py
index 2aa8210ed..993c47fa8 100644
--- a/tests/_helpers.py
+++ b/tests/_helpers.py
@@ -30,7 +30,7 @@ def _jax_not_installed() -> bool:
 
 
 def bias_knn(conn, pseudotime, n_neighbors, k=3):
-    k_thresh = np.min([int(np.floor(n_neighbors / k)) - 1, 30])
+    k_thresh = max(0, min(int(np.floor(n_neighbors / k)) - 1, 30))
     conn_biased = conn.copy()
 
     # check whether the original graph was connected
diff --git a/tests/test_external.py b/tests/test_external.py
new file mode 100644
index 000000000..9eb2d7c66
--- /dev/null
+++ b/tests/test_external.py
@@ -0,0 +1,51 @@
+import pytest
+
+from anndata import AnnData
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import spmatrix
+
+import cellrank.external as cre
+
+
+class TestOTKernel:
+    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]
+        terminal_states[ixs] = "GI"
+
+        ok = cre.kernels.OTKernel(
+            adata_large,
+            terminal_states=pd.Series(terminal_states).astype("category"),
+            g=np.ones((adata_large.n_obs,), dtype=np.float64),
+        )
+
+        with pytest.raises(
+            NotImplementedError, match="Method `'unbal'` is not yet implemented."
+        ):
+            ok.compute_transition_matrix(1, 0.001, method="unbal")
+
+    def test_no_terminal_states(self, adata_large: AnnData):
+        with pytest.raises(RuntimeError, match="Unable to initialize the kernel."):
+            cre.kernels.OTKernel(
+                adata_large,
+                g=np.ones((adata_large.n_obs,), dtype=np.float64),
+            )
+
+    def test_normal_run(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]
+        terminal_states[ixs] = "GI"
+
+        ok = cre.kernels.OTKernel(
+            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 isinstance(ok, cre.kernels.OTKernel)
+        assert isinstance(ok._transition_matrix, (np.ndarray, spmatrix))
+        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
+        assert isinstance(ok.params, dict)
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index 0406fdea0..cee8b6657 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1,4 +1,4 @@
-from typing import Tuple, Callable
+from typing import Tuple, Callable, Optional
 
 import pytest
 from _helpers import (
@@ -13,18 +13,26 @@
 from anndata import AnnData
 
 import numpy as np
+from pandas.core.dtypes.common import is_bool_dtype, is_integer_dtype
 
 import cellrank as cr
 from cellrank.tl._utils import _normalize
 from cellrank.ul._utils import _get_neighs, _get_neighs_params
 from cellrank.tl.kernels import (
-    Constant,
-    PalantirKernel,
     VelocityKernel,
+    CytoTRACEKernel,
+    PseudotimeKernel,
     PrecomputedKernel,
     ConnectivityKernel,
 )
-from cellrank.tl.kernels._base_kernel import KernelAdd, KernelMul, _is_bin_mult
+from cellrank.tl._constants import _transition
+from cellrank.tl.kernels._base_kernel import (
+    Constant,
+    KernelAdd,
+    KernelMul,
+    _is_bin_mult,
+)
+from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEAggregation, _ct
 
 _rtol = 1e-6
 
@@ -80,7 +88,7 @@ class TestInitializeKernel:
     def test_none_transition_matrix(self, adata: AnnData):
         vk = VelocityKernel(adata)
         ck = ConnectivityKernel(adata)
-        pk = PalantirKernel(adata, time_key="latent_time")
+        pk = PseudotimeKernel(adata, time_key="latent_time")
 
         assert vk._transition_matrix is None
         assert ck._transition_matrix is None
@@ -89,7 +97,7 @@ def test_none_transition_matrix(self, adata: AnnData):
     def test_not_none_transition_matrix_compute(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata).compute_transition_matrix()
-        pk = PalantirKernel(adata, time_key="latent_time").compute_transition_matrix()
+        pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix()
 
         assert vk.transition_matrix is not None
         assert ck.transition_matrix is not None
@@ -98,7 +106,7 @@ def test_not_none_transition_matrix_compute(self, adata: AnnData):
     def test_not_none_transition_matrix_accessor(self, adata: AnnData):
         vk = VelocityKernel(adata)
         ck = ConnectivityKernel(adata)
-        pk = PalantirKernel(adata, time_key="latent_time")
+        pk = PseudotimeKernel(adata, time_key="latent_time")
 
         assert vk.transition_matrix is not None
         assert ck.transition_matrix is not None
@@ -325,17 +333,20 @@ def test_adaptive_kernel_complex(self, adata: AnnData):
     def test_repr(self, adata: AnnData):
         rpr = repr(VelocityKernel(adata))
 
-        assert rpr == "<Velo>"
+        assert rpr == f"<{VelocityKernel.__name__}>"
 
     def test_repr_inv(self, adata: AnnData):
         rpr = repr(~VelocityKernel(adata))
 
-        assert rpr == "~<Velo>"
+        assert rpr == f"~<{VelocityKernel.__name__}>"
 
     def test_repr_inv_comb(self, adata: AnnData):
         rpr = repr(~(VelocityKernel(adata) + ConnectivityKernel(adata)))
 
-        assert rpr == "~((1 * <Velo>) + (1 * <Conn>))"
+        assert (
+            rpr
+            == f"~((1 * <{VelocityKernel.__name__}>) + (1 * <{ConnectivityKernel.__name__}>))"
+        )
 
     def test_str_repr_equiv_no_transition_matrix(self, adata: AnnData):
         vk = VelocityKernel(adata)
@@ -343,19 +354,24 @@ def test_str_repr_equiv_no_transition_matrix(self, adata: AnnData):
         rpr = repr(vk)
 
         assert string == rpr
-        assert string == "<Velo>"
+        assert string == f"<{VelocityKernel.__name__}>"
 
     def test_str(self, adata: AnnData):
         string = str(ConnectivityKernel(adata).compute_transition_matrix())
 
-        assert string == "<Conn[dnorm=True, key=connectivities]>"
+        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 == "~<Conn[dnorm=True, key=connectivities]>"
+        assert (
+            string
+            == f"~<{ConnectivityKernel.__name__}[dnorm=True, key=connectivities]>"
+        )
 
     def test_combination_correct_parameters(self, adata: AnnData):
         from cellrank.tl.kernels import CosineScheme
@@ -431,7 +447,7 @@ def test_precomputed_no_adata(self):
         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"] == "<Precomputed[origin='array']>"
+        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()
         )
@@ -536,23 +552,30 @@ def test_transition_forward_stoch_high_lvl(self, adata: AnnData):
 
         np.testing.assert_allclose(T_1.A, T_2.A, rtol=_rtol)
 
-    def test_palantir(self, adata: AnnData):
+    @pytest.mark.parametrize("dens_norm", [False, True])
+    def test_palantir(self, adata: AnnData, dens_norm: bool):
         conn = _get_neighs(adata, "connectivities")
         n_neighbors = _get_neighs_params(adata)["n_neighbors"]
         pseudotime = adata.obs["latent_time"]
 
         conn_biased = bias_knn(conn, pseudotime, n_neighbors)
-        T_1 = _normalize(conn_biased)
+        if dens_norm:
+            T_1 = density_normalization(conn_biased, conn)
+            T_1 = _normalize(T_1)
+        else:
+            T_1 = _normalize(conn_biased)
 
-        pk = PalantirKernel(adata, time_key="latent_time").compute_transition_matrix(
-            density_normalize=False
+        pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
+            density_normalize=dens_norm,
+            frac_to_keep=1 / 3.0,
+            threshold_scheme="hard",
         )
         T_2 = pk.transition_matrix
 
         np.testing.assert_allclose(T_1.A, T_2.A, rtol=_rtol)
 
     def test_palantir_inverse(self, adata: AnnData):
-        pk = PalantirKernel(adata, time_key="latent_time")
+        pk = PseudotimeKernel(adata, time_key="latent_time")
         pt = pk.pseudotime.copy()
 
         pk_inv = ~pk
@@ -561,22 +584,6 @@ def test_palantir_inverse(self, adata: AnnData):
         assert pk_inv.backward
         np.testing.assert_allclose(pt, 1 - pk_inv.pseudotime)
 
-    def test_palantir_dense_norm(self, adata: AnnData):
-        conn = _get_neighs(adata, "connectivities")
-        n_neighbors = _get_neighs_params(adata)["n_neighbors"]
-        pseudotime = adata.obs["latent_time"]
-
-        conn_biased = bias_knn(conn, pseudotime, n_neighbors)
-        T_1 = density_normalization(conn_biased, conn)
-        T_1 = _normalize(T_1)
-
-        pk = PalantirKernel(adata, time_key="latent_time").compute_transition_matrix(
-            density_normalize=True
-        )
-        T_2 = pk.transition_matrix
-
-        np.testing.assert_allclose(T_1.A, T_2.A, rtol=_rtol)
-
     def test_palantir_differ_dense_norm(self, adata: AnnData):
         conn = _get_neighs(adata, "connectivities")
         n_neighbors = _get_neighs_params(adata)["n_neighbors"]
@@ -586,7 +593,7 @@ def test_palantir_differ_dense_norm(self, adata: AnnData):
         T_1 = density_normalization(conn_biased, conn)
         T_1 = _normalize(T_1)
 
-        pk = PalantirKernel(adata, time_key="latent_time").compute_transition_matrix(
+        pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
             density_normalize=False
         )
         T_2 = pk.transition_matrix
@@ -917,7 +924,7 @@ def test_copy_cond_num(self, adata: AnnData):
         for KernelClass in [
             VelocityKernel,
             ConnectivityKernel,
-            PalantirKernel,
+            PseudotimeKernel,
             PrecomputedKernel,
         ]:
             if KernelClass is PrecomputedKernel:
@@ -955,7 +962,7 @@ def test_copy_connectivity_kernel(self, adata: AnnData):
         assert ck1.backward == ck2.backward
 
     def test_copy_palantir_kernel(self, adata: AnnData):
-        pk1 = PalantirKernel(adata).compute_transition_matrix()
+        pk1 = PseudotimeKernel(adata).compute_transition_matrix()
         pk2 = pk1.copy()
 
         np.testing.assert_array_equal(pk1.transition_matrix.A, pk2.transition_matrix.A)
@@ -1193,13 +1200,6 @@ def test_not_callable(self, adata: AnnData):
         ):
             vk.compute_transition_matrix(scheme=1311)
 
-    def test_not_callable_object(self, adata: AnnData):
-        vk = VelocityKernel(adata)
-        with pytest.raises(
-            TypeError, match="Expected `scheme` to be a function object"
-        ):
-            vk.compute_transition_matrix(scheme=CustomFunc)
-
     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."):
@@ -1276,3 +1276,194 @@ def test_custom_function_stochastic_no_hessian(self, adata: AnnData):
 
         assert vk.params["mode"] == "monte_carlo"
         assert vk.params["scheme"] == 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()
+
+    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")
+
+    def test_basis_prefix(self, adata: AnnData):
+        ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
+        ck.compute_projection(basis="X_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")
+        else:
+            ck.compute_projection(basis="umap")
+            ck.write_to_adata()
+
+        assert adata.uns[_transition(ck._direction) + "_params"] == {
+            "params": ck.params,
+            "embeddings": ["umap"],
+        }
+
+    @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)
+
+        key = key_added if key_added is not None else _transition(ck._direction)
+        ukey = f"{key}_params"
+        key = f"{key}_umap"
+
+        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 = _transition(ck._direction) + "_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
+
+        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, :]))
+
+
+class TestPseudotimeKernelScheme:
+    def test_invalid_scheme(self, adata: AnnData):
+        pk = PseudotimeKernel(adata)
+        with pytest.raises(ValueError, match="foo"):
+            pk.compute_transition_matrix(threshold_scheme="foo")
+
+    def test_invalid_custom_scheme(self, adata: AnnData):
+        pk = PseudotimeKernel(adata)
+        with pytest.raises(ValueError, match="Expected row of shape"):
+            pk.compute_transition_matrix(
+                threshold_scheme=lambda cpt, npt, ndist: np.ones(
+                    (len(ndist) - 1), dtype=np.float64
+                ),
+            )
+
+    def test_custom_scheme(self, adata: AnnData):
+        pk = PseudotimeKernel(adata)
+        pk.compute_transition_matrix(
+            threshold_scheme=lambda cpt, npt, ndist: np.ones(
+                (len(ndist)), dtype=np.float64
+            ),
+            density_normalize=False,
+        )
+
+        np.testing.assert_allclose(pk.transition_matrix.sum(1), 1.0)
+        for row in pk.transition_matrix:
+            np.testing.assert_allclose(row.data, 1 / len(row.data))
+
+    @pytest.mark.parametrize("scheme", ["hard", "soft"])
+    def test_scheme(self, adata: AnnData, scheme: str):
+        pk = PseudotimeKernel(adata)
+        pk.compute_transition_matrix(
+            threshold_scheme=scheme, frac_to_keep=0.3, b=10, nu=0.5
+        )
+
+        np.testing.assert_allclose(pk.transition_matrix.sum(1), 1.0)
+        assert pk.params["scheme"] == scheme
+        if scheme == "hard":
+            assert pk.params["frac_to_keep"] == 0.3
+            assert "b" not in pk.params
+            assert "nu" not in pk.params
+        elif scheme == "soft":
+            assert pk.params["b"] == 10
+            assert pk.params["nu"] == 0.5
+            assert "k" not in pk.params
+
+
+class TestCytoTRACEKernel:
+    @pytest.mark.parametrize("layer", ["X", "Ms", "foo"])
+    def test_layer(self, adata: AnnData, layer: str):
+        if layer == "foo":
+            with pytest.raises(KeyError, match=layer):
+                _ = CytoTRACEKernel(adata, layer=layer)
+        else:
+            _ = CytoTRACEKernel(adata, layer=layer)
+            assert adata.uns[_ct("params")]["layer"] == layer
+
+    @pytest.mark.parametrize("agg", list(CytoTRACEAggregation))
+    def test_aggregation(self, adata: AnnData, agg: CytoTRACEAggregation):
+        _ = CytoTRACEKernel(adata, aggregation=agg)
+        assert adata.uns[_ct("params")]["aggregation"] == agg.s
+
+    @pytest.mark.parametrize("use_raw", [False, True])
+    def test_raw(self, adata: AnnData, use_raw: bool):
+        _ = CytoTRACEKernel(adata, use_raw=use_raw)
+        assert adata.uns[_ct("params")]["use_raw"] == (
+            adata.raw.n_vars == adata.n_vars if use_raw else False
+        )
+
+    def test_correct_class(self, adata: AnnData):
+        k = CytoTRACEKernel(adata)
+        assert isinstance(k, PseudotimeKernel)
+        assert k._time_key == _ct("pseudotime")
+
+    def test_writes_params(self, adata: AnnData):
+        k = CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
+
+        assert adata.uns[_ct("params")] == {
+            "layer": "X",
+            "aggregation": "mean",
+            "use_raw": False,
+        }
+
+        assert np.all(adata.var[_ct("gene_corr")] <= 1.0)
+        assert np.all(-1 <= adata.var[_ct("gene_corr")])
+        assert is_bool_dtype(adata.var[_ct("correlates")])
+        assert adata.var[_ct("correlates")].sum() == min(200, adata.n_vars)
+
+        assert _ct("score") in adata.obs
+        assert _ct("pseudotime") in adata.obs
+        assert _ct("num_exp_genes") in adata.obs
+        assert is_integer_dtype(adata.obs[_ct("num_exp_genes")])
+        np.testing.assert_array_equal(k.pseudotime, adata.obs[_ct("pseudotime")].values)
+        np.testing.assert_array_equal(k.pseudotime.min(), 0.0)
+        np.testing.assert_array_equal(k.pseudotime.max(), 1.0)
+
+    def test_compute_transition_matrix(self, adata: AnnData):
+        k = CytoTRACEKernel(adata, 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")
+
+        pt = adata.obs[_ct("pseudotime")].values
+        np.testing.assert_array_equal(np.max(pt) - pt, k.pseudotime)
+
+    def test_regression_score(self, adata_large: AnnData):
+        _ = CytoTRACEKernel(adata_large, use_raw=False)
+
+        ct_score_expected = adata_large.obs["ct_score_ground_truth"].values
+        ct_score_actual = adata_large.obs[_ct("score")].values
+
+        np.testing.assert_array_equal(ct_score_actual, ct_score_expected)
+
+    def test_regression_gene_corr(self, adata_large: AnnData):
+        _ = CytoTRACEKernel(adata_large, use_raw=False)
+
+        gene_corr_expected = adata_large.var["gene_corr_ground_truth"].values
+        gene_corr_actual = adata_large.var[_ct("gene_corr")].values
+
+        np.testing.assert_array_equal(gene_corr_actual, gene_corr_expected)
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 8e6e51e7c..f6213c95b 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -2493,6 +2493,19 @@ def test_scvelo_gpcca_abs_probs_cont_not_same(self, mc: GPCCA, fpath: str):
             cluster_key="clusters", discrete=False, same_plot=False, dpi=DPI, save=fpath
         )
 
+    @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",
+            basis="umap",
+            arrow_length=6,
+            arrow_size=6,
+            dpi=DPI,
+            save=fpath,
+        )
+
 
 class TestLineages:
     @compare()
@@ -2652,17 +2665,28 @@ def test_pie_legend_kwargs(self, lineage: cr.tl.Lineage, fpath: str):
 
 class TestLineageDrivers:
     @compare()
-    def test_scvelo_drivers_n_genes(self, adata: AnnData, fpath: str):
+    def test_drivers_n_genes(self, adata: AnnData, fpath: str):
         cr.pl.lineage_drivers(adata, "0", n_genes=5, dpi=DPI, save=fpath)
 
     @compare(kind="bwd")
-    def test_scvelo_drivers_n_genes(self, adata: AnnData, fpath: str):
-        cr.pl.lineage_drivers(adata, "0", backward=True, n_genes=5, dpi=DPI, save=fpath)
+    def test_drivers_backward(self, adata: AnnData, fpath: str):
+        cr.pl.lineage_drivers(adata, "0", backward=True, ncols=2, dpi=DPI, save=fpath)
 
     @compare()
-    def test_scvelo_drivers_cmap(self, adata: AnnData, fpath: str):
+    def test_drivers_cmap(self, adata: AnnData, fpath: str):
         cr.pl.lineage_drivers(adata, "0", cmap="inferno", dpi=DPI, save=fpath)
 
+    @compare()
+    def test_drivers_title_fmt(self, adata: AnnData, fpath: str):
+        cr.pl.lineage_drivers(
+            adata,
+            "0",
+            cmap="inferno",
+            title_fmt="{gene} qval={qval} corr={corr}",
+            dpi=DPI,
+            save=fpath,
+        )
+
 
 class TestModel:
     @compare()
@@ -3090,3 +3114,58 @@ def test_proj_no_cbar(self, adata: AnnData, fpath: str):
         cr.pl.circular_projection(
             adata, keys=adata.var_names[0], colorbar=False, dpi=DPI, save=fpath
         )
+
+
+class TestPlotRandomWalk:
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_params(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=100,
+            max_iter=100,
+            seed=42,
+            start_ixs={"clusters": "OL"},
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_basis(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10, max_iter=100, seed=42, basis="pca", dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_cmap(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10, max_iter=100, seed=42, cmap="viridis", dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_line_width(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10, max_iter=100, seed=42, linewidth=2, dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_line_alpha(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10, max_iter=100, seed=42, linealpha=1, dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_kwargs(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10, max_iter=100, seed=42, color="none", dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_kernel_random_walk_ixs_legend_loc(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_random_walks(
+            n_sims=10,
+            max_iter=100,
+            seed=42,
+            ixs_legend_loc="top right out",
+            legend_loc="upper left",
+            dpi=DPI,
+            save=fpath,
+        )
diff --git a/tests/test_random_walk.py b/tests/test_random_walk.py
new file mode 100644
index 000000000..291a68193
--- /dev/null
+++ b/tests/test_random_walk.py
@@ -0,0 +1,90 @@
+import pytest
+
+import numpy as np
+
+from cellrank.tl.kernels._random_walk import RandomWalk
+
+
+class TestRandomWalk:
+    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)
+
+    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])
+
+    @pytest.mark.parametrize("kind", ["simulations", "iterations", "hits"])
+    def test_invalid_numbers(self, test_matrix_1: np.ndarray, kind: str):
+        rw = RandomWalk(test_matrix_1)
+        with pytest.raises(ValueError, match=kind):
+            if kind == "simulations":
+                rw.simulate_many(n_sims=0)
+            elif kind == "iterations":
+                rw.simulate_one(max_iter=1)
+            elif kind == "hits":
+                rw.simulate_one(successive_hits=-1)
+            else:
+                raise NotImplementedError(kind)
+
+    @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)
+
+        if n_sims == 1:
+            r1, r2 = rw1.simulate_one(seed=42), rw2.simulate_one(seed=42)
+        else:
+            r1, r2 = rw1.simulate_many(10, seed=42), rw2.simulate_many(10, seed=42)
+
+        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)
+
+        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)
+
+        assert res[0] == 1
+
+    def test_stop_ixs(self, test_matrix_1: np.ndarray):
+        res = RandomWalk(test_matrix_1, stop_ixs=[1]).simulate_one(
+            max_iter=1000, seed=42
+        )
+
+        assert len(res) <= 1000
+        assert res[-1] == 1
+        assert res[-2] != 1
+
+    def successive_hits(self, test_matrix_1: np.ndarray):
+        res = RandomWalk(test_matrix_1, stop_ixs=[1]).simulate_one(
+            max_iter=1000, seed=42, successive_hits=2
+        )
+
+        assert len(res) <= 1000
+        assert res[-1] == 1
+        assert res[-2] == 1
+        assert res[-3] == 1
+        assert res[-4] != 1
+
+    @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(
+            10, max_iter=100, seed=42, backend=backend, n_jobs=n_jobs
+        )
+
+        assert isinstance(res, list)
+        assert len(res) == 10
+        for i, r in enumerate(res):
+            assert isinstance(r, np.ndarray), i
+            assert np.issubdtype(r.dtype, np.integer), i
+            np.testing.assert_array_equal(r.shape, (101,))