Repository: theislab/cellrank (language: python)

Issue:
1.4 release
<!-- 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`?
- [x] Other?

<!-- Describe the feature requested below: -->
Meta issue to track things need to be done before the release:
- [x] close #647 
- [x] ~~update pre-commit & disable macOS testing & enable 3.9 - #645~~ after new scVelo release
- [x] close #639 
- [x] close #637 
- [x] close #636 
- [x] close #635 
- [x] ~~close #628~~ 
- [x] close #626
- [x] close #619
- [x] ~~close #603~~
- [x] close #597
- [x] ~~close #593~~ after new release



Discussion / hints:

Fix patch:
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
index 0318e436b..84f514273 100644
--- a/.github/PULL_REQUEST_TEMPLATE.md
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -1,19 +1,22 @@
-**IMPORTANT: Please search among the [Pull request](../pulls) before creating one.**
+**IMPORTANT: Please search among the [Pull requests](../pulls) before creating one.**
 
 ## Title
-<!--- Provide a general summary of your changes in the Title above -->
+<!-- Provide a general summary of your changes in the Title above -->
 
 ## Types of changes
-<!--- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
+<!-- What types of changes does your code introduce? Put an `x` in all the boxes that apply: -->
 - [ ] Bug fix (non-breaking change which fixes an issue)
 - [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)
 - [ ] New feature (non-breaking change which adds functionality)
 
 ## Description
-<!--- Clearly and concisely describe your changes. -->
+<!-- Clearly and concisely describe your changes. -->
 
 ## How has this been tested?
-<!--- Describe in detail how you've tested your changes. -->
+<!-- Describe in detail how you've tested your changes. -->
+
+## Have release notes been modified?
+<!-- Indicate whether release notes have been (or should) be modifies. -->
 
 ## Closes
-<!--- Type `closes #XXXX` in your comment to auto-close the issue that your PR fixes (if such). -->
+<!-- If applicable, type `closes #XXXX` in your comment to auto-close the issue that this PR fixes. -->
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index ecb7dd81e..a1977a8da 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -52,15 +52,21 @@ jobs:
         runs-on: ${{ matrix.os }}
         strategy:
             fail-fast: false
-            max-parallel: 3
+            max-parallel: 4
             matrix:
-                python: [3.7]
-                os: [ubuntu-latest, macos-latest]
+                python: [3.7]  # , 3.9]
+                os: [ubuntu-latest]
                 use_slepc: [false]
+                allow_failure: [false]
                 include:
                 -   os: ubuntu-latest
                     python: 3.8
                     use_slepc: true
+                    allow_failure: false
+                -   os: macos-latest
+                    python: 3.8
+                    use_slepc: false
+                    allow_failure: true
         env:
             OS: ${{ matrix.os }}
             PYTHON: ${{ matrix.python }}
@@ -81,8 +87,9 @@ jobs:
             uses: actions/cache@v2
             with:
                 path: ${{ steps.pip-cache.outputs.dir }}
-                key: pip-${{ runner.os }}-${{ env.pythonLocation }}-${{ hashFiles('**/requirements.txt') }}
+                key: pip-${{ runner.os }}-${{ env.pythonLocation }}-${{ hashFiles('**/requirements.txt') }}-${{ hashFiles('**/setup.py') }}
                 restore-keys: |
+                    pip-${{ runner.os }}-${{ env.pythonLocation }}-${{ hashFiles('**/requirements.txt') }}
                     pip-${{ runner.os }}-${{ env.pythonLocation }}-
 
         -   name: Before install
@@ -93,7 +100,8 @@ jobs:
                 ./.scripts/ci/install.sh
 
         -   name: Testing
-            timeout-minutes: 30
+            timeout-minutes: 60
+            continue-on-error: ${{ matrix.allow_failure }}
             run: |
                 python -m pytest --cov --cov-report=term-missing --cov-config=./.coveragerc --cov-report=xml -vv
 
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index 70cd08852..15cdf65c6 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,28 +7,28 @@ default_stages:
 minimum_pre_commit_version: 2.9.0
 repos:
 -   repo: https://github.com/psf/black
-    rev: 20.8b1
+    rev: 21.6b0
     hooks:
     -   id: black
         additional_dependencies: [toml]
 -   repo: https://github.com/timothycrosley/isort
-    rev: 5.7.0
+    rev: 5.9.1
     hooks:
     -   id: isort
         additional_dependencies: [toml]
         args: [--order-by-type]
 -   repo: https://github.com/asottile/yesqa
-    rev: v1.2.2
+    rev: v1.2.3
     hooks:
     -   id: yesqa
         additional_dependencies: [flake8-tidy-imports, flake8-docstrings, flake8-comprehensions, flake8-bugbear, flake8-blind-except]
 -   repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-    rev: v2.0.0
+    rev: v2.1.0
     hooks:
     -   id: pretty-format-yaml
         args: [--autofix, --indent, '4']
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v3.4.0
+    rev: v4.0.1
     hooks:
     -   id: detect-private-key
     -   id: check-merge-conflict
@@ -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.9.0
+    rev: 3.9.2
     hooks:
     -   id: flake8
         additional_dependencies: [flake8-tidy-imports, flake8-docstrings, flake8-comprehensions, flake8-bugbear, flake8-blind-except]
@@ -74,12 +74,12 @@ repos:
     -   id: blacken-docs
         additional_dependencies: [black==20.8b1]
 -   repo: https://github.com/asottile/pyupgrade
-    rev: v2.10.1
+    rev: v2.19.4
     hooks:
     -   id: pyupgrade
         args: [--py3-plus, --py37-plus]
 -   repo: https://github.com/pre-commit/pygrep-hooks
-    rev: v1.8.0
+    rev: v1.9.0
     hooks:
     -   id: python-no-eval
     -   id: python-use-type-annotations
diff --git a/.rstcheck.cfg b/.rstcheck.cfg
index 63b9ec02c..264401026 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|module)" 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|cite|bibliography)" 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 4bdf01848..8731e761a 100755
--- a/.scripts/ci/install.sh
+++ b/.scripts/ci/install.sh
@@ -3,18 +3,19 @@
 set -ev
 
 python -m pip install --upgrade pip
-pip install codecov
+pip install codecov adjustText
 
 if [[ "$OS" == "macos-latest" ]]; then
   pip install -e".[test,external]"
 elif [[ "$OS" == "ubuntu-latest" ]]; then
+  pip install numpy Cython  # 3.9 external (POT)
   if [[ "$USE_SLEPC" == "true" ]]; then
       pip install -e".[krylov,test,external]"
       python -c "import slepc; import petsc;"
   else
       pip install -e".[test,external]"
   fi
-  pip install rpy2>=3.3.0
+  pip install "rpy2>=3.3.0"
   Rscript --vanilla -e "library('mgcv')"
 else
   exit 42
diff --git a/cellrank/datasets/__init__.py b/cellrank/datasets/__init__.py
index f2d977e8a..4b6a80c00 100644
--- a/cellrank/datasets/__init__.py
+++ b/cellrank/datasets/__init__.py
@@ -2,6 +2,7 @@
     lung,
     pancreas,
     zebrafish,
-    reprogramming,
+    reprogramming_morris,
     pancreas_preprocessed,
+    reprogramming_schiebinger,
 )
diff --git a/cellrank/datasets/_datasets.py b/cellrank/datasets/_datasets.py
index 22fc1eaef..09618f1e9 100644
--- a/cellrank/datasets/_datasets.py
+++ b/cellrank/datasets/_datasets.py
@@ -5,7 +5,6 @@
 from pathlib import Path
 
 from scanpy import read
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._constants import ModeEnum
@@ -32,7 +31,7 @@ class ReprogrammingSubset(ModeEnum):  # noqa: D101
         "https://ndownloader.figshare.com/files/25038224",
         (24882, 24051),
     ),
-    "reprogramming": (
+    "reprogramming_morris": (
         "https://ndownloader.figshare.com/files/25503773",
         (104679, 22630),
     ),
@@ -40,6 +39,10 @@ class ReprogrammingSubset(ModeEnum):  # noqa: D101
         "https://ndownloader.figshare.com/files/27265280",
         (2434, 23974),
     ),
+    "reprogramming_schiebinger": (
+        "https://ndownloader.figshare.com/files/28618734",
+        (236285, 19089),
+    ),
 }
 
 
@@ -86,7 +89,7 @@ def pancreas(
     **kwargs: Any,
 ) -> AnnData:
     """
-    Development of the murine pancreas at E15.5 from [Panc19]_.
+    Development of the murine pancreas at E15.5 from :cite:`bastidas-ponce:19`.
 
     scRNA-seq dataset comprising 2531 cells recorded using 10x Chromium in a single time point. Data was filtered
     to remove heavily cycling populations and to focus on the late stages of endocrinogenesis.
@@ -105,7 +108,6 @@ def pancreas(
     -------
     %(adata)s
     """
-
     return _load_dataset_from_url(path, *_datasets["pancreas"], **kwargs)
 
 
@@ -115,7 +117,7 @@ def pancreas_preprocessed(
     **kwargs: Any,
 ) -> AnnData:
     """
-    Development of the murine pancreas at E15.5 from [Panc19]_, preprocessed according to the \
+    Development of the murine pancreas at E15.5 from :cite:`bastidas-ponce:19`, preprocessed according to the \
     `basic tutorial <https://cellrank.readthedocs.io/en/latest/pancreas_basic.html>`__.
 
     Parameters
@@ -126,7 +128,6 @@ def pancreas_preprocessed(
     -------
     %(adata)s
     """
-
     return _load_dataset_from_url(path, *_datasets["pancreas_preprocessed"], **kwargs)
 
 
@@ -136,12 +137,12 @@ def lung(
     **kwargs: Any,
 ) -> AnnData:
     """
-    Regeneration of murine lung epithelial cells at 13 time points from [Lung20]_.
+    Regeneration of murine lung epithelial cells at 13 time points from :cite:`strunz:20`.
 
-    scRNA-seq dataset comprising 24,051 cells recorded using Dropseq [Macosko15]_ at 13 time points spanning days
-    2-15 past lung bleomycin injury. Data was filtered to remove control cells as well as later time points which  are
-    more spaced out. We wanted to focus on the densely sampled days where RNA velocity [Manno18]_ [Bergen20]_ can be
-    used to predict the future cellular state.
+    scRNA-seq dataset comprising `24 051` cells recorded using Dropseq :cite:`macosko:15` at 13 time points spanning
+    days 2-15 past lung bleomycin injury. Data was filtered to remove control cells as well as later time points which
+    are more spaced out. We wanted to focus on the densely sampled days where RNA velocity :cite:`manno:18`
+    :cite:`bergen:20` can be used to predict the future cellular state.
 
     Contains raw spliced and un-spliced count data, low-dimensional embedding coordinates as well as original
     cluster annotations.
@@ -154,25 +155,25 @@ def lung(
     -------
     %(adata)s
     """
-
     return _load_dataset_from_url(path, *_datasets["lung"], **kwargs)
 
 
 @inject_docs(s=ReprogrammingSubset)
 @d.dedent
-def reprogramming(
+def reprogramming_morris(
     subset: str = ReprogrammingSubset.FULL.s,
-    path: Union[str, Path] = "datasets/reprogramming.h5ad",
+    path: Union[str, Path] = "datasets/reprogramming_morris.h5ad",
     **kwargs: Any,
 ) -> AnnData:
     """
-    Reprogramming of mouse embryonic fibroblasts to induced endoderm progenitors at 8 time points from [Morris18]_.
+    Reprogramming of mouse embryonic fibroblasts to induced endoderm progenitors at 8 time points from \
+    :cite:`morris:18`.
 
-    scRNA-seq dataset comprising `104,887` cell recorded using 10X Chromium and Dropseq [Macosko15]_ at 8 time points
-    spanning days 0-28 past reprogramming initiation.
+    scRNA-seq dataset comprising `104,887` cell recorded using 10X Chromium and Dropseq :cite:`macosko:15`
+    at 8 time points spanning days 0-28 past reprogramming initiation.
 
     Contains raw spliced and un-spliced count data, low-dimensional embedding coordinates as well as clonal information
-    from CellTagging [Morris18]_. Moreover, contains the following :attr:`anndata.AnnData.obs`: annotations:
+    from CellTagging :cite:`morris:18`. Moreover, it contains the following :attr:`anndata.AnnData.obs` annotations:
 
         - `'reprogramming_day'` - time-point information.
         - `'reprogramming'` - whether this clone is enriched for cells from successfully reprogrammed populations.
@@ -183,9 +184,9 @@ def reprogramming(
     subset
         Whether to return the full object or just a subset. Can be one of:
 
-            - `{s.FULL.s!r}` - return the complete dataset containing `104,887` cells.
-            - `{s.K85.s!r}` - return the subset as described in [Morris18]_ Fig. 1, containing `85,010` cells.
-            - `{s.K48.s!r}` - return the subset as described in [Morris18]_ Fig. 3, containing `48,515` cells.
+            - `{s.FULL.s!r}` - return the complete dataset containing `104 887` cells.
+            - `{s.K85.s!r}` - return the subset as described in :cite:`morris:18` Fig. 1, containing `85 010` cells.
+            - `{s.K48.s!r}` - return the subset as described in :cite:`morris:18` Fig. 3, containing `48 515` cells.
 
     %(dataset.parameters)s
 
@@ -195,10 +196,10 @@ def reprogramming(
 
     Notes
     -----
-    The dataset has approximately 1.5GiB and the subsetting is performed only locally after the full download.
+    The dataset has approximately 1.5GiB and the subsetting is performed locally after the full download.
     """
     subset = ReprogrammingSubset(subset)
-    adata = _load_dataset_from_url(path, *_datasets["reprogramming"], **kwargs)
+    adata = _load_dataset_from_url(path, *_datasets["reprogramming_morris"], **kwargs)
 
     if subset == ReprogrammingSubset.FULL:
         return adata
@@ -212,16 +213,51 @@ def reprogramming(
     )
 
 
+@d.dedent
+def reprogramming_schiebinger(
+    path: Union[str, Path] = "datasets/reprogramming_schiebinger.h5ad", **kwargs: Any
+) -> AnnData:
+    """
+    Reprogramming of mouse embryonic fibroblasts to induced pluripotent stem cells at 39 time points from \
+    :cite:`schiebinger:19`.
+
+    scRNA-seq dataset comprising `236 285` cell recorded using 10X Chromium
+    at 39 time points spanning days 0-18 past reprogramming initiation.
+
+    Contains total-counts normalized, log-transformed counts and low-dimensional embedding coordinates (force-directed).
+    Moreover, it contains the following :attr:`anndata.AnnData.obs` annotations:
+
+        - `'day'` - time-point information.
+        - `'serum'`/`'2i'` - whether this cell comes from the serum/2i condition.
+        - `'cell_sets'` - cluster labels.
+
+    Parameters
+    ----------
+    %(dataset.parameters)s
+
+    Returns
+    -------
+    %(adata)s
+
+    Notes
+    -----
+    The dataset has approximately 1.4GiB.
+    """
+    return _load_dataset_from_url(
+        path, *_datasets["reprogramming_schiebinger"], **kwargs
+    )
+
+
 @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]_.
+    Zebrafish embryogenesis assayed using drop-seq, restricted to the axial mesoderm lineage from :cite:`farrell:18`.
 
-    scRNA-seq time-series dataset comprising `2434` cells
-    which contains 12 time-points spanning 3.3 - 12 hours past fertilization.
+    scRNA-seq time-series dataset comprising `2434` cells which contains 12 time-points spanning 3.3-12 hours
+    past fertilization.
 
     Parameters
     ----------
@@ -231,5 +267,4 @@ def zebrafish(
     -------
     %(adata)s
     """
-
     return _load_dataset_from_url(path, *_datasets["zebrafish"], **kwargs)
diff --git a/cellrank/external/kernels/__init__.py b/cellrank/external/kernels/__init__.py
index 1052a81ec..a5425f0f5 100644
--- a/cellrank/external/kernels/__init__.py
+++ b/cellrank/external/kernels/__init__.py
@@ -1,1 +1,2 @@
-from cellrank.external.kernels._statot_kernel import OTKernel
+from cellrank.external.kernels._wot_kernel import WOTKernel
+from cellrank.external.kernels._statot_kernel import StationaryOTKernel
diff --git a/cellrank/external/kernels/_statot_kernel.py b/cellrank/external/kernels/_statot_kernel.py
index 1c693c775..bb0900d9e 100644
--- a/cellrank/external/kernels/_statot_kernel.py
+++ b/cellrank/external/kernels/_statot_kernel.py
@@ -3,14 +3,13 @@
 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
 
+import numpy as np
+import pandas as pd
+
 _error = None
 try:
     from statot import OTKernel as OTKernel_
@@ -23,9 +22,9 @@
 
 
 @d.dedent
-class OTKernel(OTKernel_, error=_error):
+class StationaryOTKernel(OTKernel_, error=_error):
     """
-    Stationary optimal transport kernel from [Zhang21]_.
+    Stationary optimal transport kernel from :cite:`zhang:21`.
 
     This class requires the `statOT` package, which can be installed as `pip install statot POT`.
 
@@ -34,7 +33,7 @@ class OTKernel(OTKernel_, error=_error):
     %(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']``.
+        If `None`, terminal states are assumed to be already present in :attr:`adata` ``['terminal_states']``.
     g
         Key in :attr:`anndata.AnnData.obs` containing relative growth rates for cells or the array itself.
     cluster_key
@@ -63,6 +62,14 @@ def __init__(
         except Exception as e:  # noqa: B902
             raise RuntimeError("Unable to initialize the kernel.") from e
 
+    def _read_from_adata(
+        self,
+        conn_key: Optional[str] = "connectivities",
+        read_conn: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super()._read_from_adata(conn_key=conn_key, read_conn=False, **kwargs)
+
     def compute_transition_matrix(
         self,
         eps: float,
@@ -76,9 +83,9 @@ def compute_transition_matrix(
         C: Optional[np.ndarray] = None,
         verbose: bool = False,
         **kwargs: Any,
-    ) -> "OTKernel":
+    ) -> "StationaryOTKernel":
         """
-        Compute transition matrix using stationary OT [Zhang21]_.
+        Compute transition matrix using stationary OT :cite:`zhang:21`.
 
         Parameters
         ----------
@@ -113,11 +120,13 @@ def compute_transition_matrix(
 
         Returns
         -------
-        :class:`cellrank.external.kernels.OTKernel`
-            Makes :paramref:`transition_matrix` available.
+        :class:`cellrank.external.kernels.StationaryOTKernel`
+            Makes :attr:`transition_matrix` available.
         """
         if method not in ("ent", "quad", "unbal"):
-            raise ValueError(f"Invalid method `{method!r}`.")
+            raise ValueError(
+                f"Invalid method `{method!r}`. Valid options are: `{['ent', 'quad', 'unbal']}`."
+            )
         if method == "unbal":
             raise NotImplementedError("Method `'unbal'` is not yet implemented.")
 
diff --git a/cellrank/external/kernels/_utils.py b/cellrank/external/kernels/_utils.py
new file mode 100644
index 000000000..545a387de
--- /dev/null
+++ b/cellrank/external/kernels/_utils.py
@@ -0,0 +1,118 @@
+from typing import Mapping, Sequence
+
+from typing_extensions import Literal
+
+
+class MarkerGenes:
+    """Helper class for :class:`cellrank.external.kernels.WOT` to get proliferation/apoptosis genes."""
+
+    # fmt: off
+    _proliferation_markers: Mapping[str, Sequence[str]] = {
+        # 97
+        "human": ("ANLN", "ANP32E", "ATAD2", "AURKA", "AURKB", "BIRC5", "BLM",
+                  "BRIP1", "BUB1", "CASP8AP2", "CBX5", "CCNB2", "CCNE2", "CDC20",
+                  "CDC25C", "CDC45", "CDC6", "CDCA2", "CDCA3", "CDCA7", "CDCA8",
+                  "CDK1", "CENPA", "CENPE", "CENPF", "CHAF1B", "CKAP2", "CKAP2L",
+                  "CKAP5", "CKS1B", "CKS2", "CLSPN", "CTCF", "DLGAP5", "DSCC1",
+                  "DTL", "E2F8", "ECT2", "EXO1", "FAM64A", "FEN1", "G2E3", "GAS2L3",
+                  "GINS2", "GMNN", "GTSE1", "HELLS", "HJURP", "HMGB2", "HMMR", "HN1",
+                  "KIF11", "KIF20B", "KIF23", "KIF2C", "LBR", "MCM2", "MCM4", "MCM5",
+                  "MCM6", "MKI67", "MLF1IP", "MSH2", "NASP", "NCAPD2", "NDC80",
+                  "NEK2", "NUF2", "NUSAP1", "PCNA", "POLA1", "POLD3", "PRIM1",
+                  "PSRC1", "RAD51", "RAD51AP1", "RANGAP1", "RFC2", "RPA2", "RRM1",
+                  "RRM2", "SLBP", "SMC4", "TACC3", "TIPIN", "TMPO", "TOP2A", "TPX2",
+                  "TTK", "TUBB4B", "TYMS", "UBE2C", "UBR7", "UHRF1", "UNG", "USP1",
+                  "WDR76"),
+        # 98
+        "mouse": ("Mcm4", "Smc4", "Gtse1", "Ttk", "Rangap1", "Ccnb2", "Cenpa",
+                  "Cenpe", "Cdca8", "Ckap2", "Rad51", "Pcna", "Ube2c", "Lbr",
+                  "Cenpf", "Birc5", "Dtl", "Dscc1", "Cbx5", "Usp1", "Hmmr", "Wdr76",
+                  "Ung", "Hn1", "Cks2", "Kif20b", "Cdk1", "Slbp", "Aurkb", "Kif11",
+                  "Cks1b", "Blm", "Msh2", "Gas2l3", "Tyms", "Hjurp", "Hells",
+                  "Prim1", "Uhrf1", "Ndc80", "Mcm6", "Rrm1", "Mlf1ip", "Top2a",
+                  "Hmgb2", "Ccne2", "G2e3", "Tmpo", "Nusap1", "Ncapd2", "Mcm2",
+                  "Kif2c", "Cdca2", "Nasp", "Gmnn", "Cdc6", "Pold3", "Ckap2l",
+                  "Fam64a", "Ubr7", "Fen1", "Bub1", "Brip1", "Atad2", "Psrc1",
+                  "Rrm2", "Tipin", "Casp8ap2", "Tubb4b", "Kif23", "Exo1", "Rfc2",
+                  "Pola1", "Mki67", "Tpx2", "Aurka", "Anln", "Chaf1b", "Hjurp",
+                  "Tacc3", "Mcm5", "Anp32e", "Dlgap5", "Ect2", "Nuf2", "Cdc45",
+                  "Ckap5", "Ctcf", "Clspn", "Cdca7", "Cdca3", "Rpa2", "Gins2",
+                  "E2f8", "Cdc25c", "Nek2", "Cdc20", "Rad51ap1"),
+    }
+    _apoptosis_markers: Mapping[str, Sequence[str]] = {
+        # 161
+        "human":
+            ("ADD1", "AIFM3", "ANKH", "ANXA1", "APP", "ATF3", "AVPR1A", "BAX",
+             "BCAP31", "BCL10", "BCL2L1", "BCL2L10", "BCL2L11", "BCL2L2", "BGN",
+             "BID", "BIK", "BIRC3", "BMF", "BMP2", "BNIP3L", "BRCA1", "BTG2",
+             "BTG3", "CASP1", "CASP2", "CASP3", "CASP4", "CASP6", "CASP7",
+             "CASP8", "CASP9", "CAV1", "CCNA1", "CCND1", "CCND2", "CD14", "CD2",
+             "CD38", "CD44", "CD69", "CDC25B", "CDK2", "CDKN1A", "CDKN1B",
+             "CFLAR", "CLU", "CREBBP", "CTH", "CTNNB1", "CYLD", "DAP", "DAP3",
+             "DCN", "DDIT3", "DFFA", "DIABLO", "DNAJA1", "DNAJC3", "DNM1L",
+             "DPYD", "EBP", "EGR3", "EMP1", "ENO2", "ERBB2", "ERBB3", "EREG",
+             "ETF1", "F2", "F2R", "FAS", "FASLG", "FDXR", "FEZ1", "GADD45A",
+             "GADD45B", "GCH1", "GNA15", "GPX1", "GPX3", "GPX4", "GSN", "GSR",
+             "GSTM1", "GUCY2D", "H1-0", "HGF", "HMGB2", "HMOX1", "HSPB1",
+             "IER3", "IFITM3", "IFNB1", "IFNGR1", "IGF2R", "IGFBP6", "IL18",
+             "IL1A", "IL1B", "IL6", "IRF1", "ISG20", "JUN", "KRT18", "LEF1",
+             "LGALS3", "LMNA", "LUM", "MADD", "MCL1", "MGMT", "MMP2", "NEDD9",
+             "NEFH", "PAK1", "PDCD4", "PDGFRB", "PEA15", "PLAT", "PLCB2",
+             "PLPPR4", "PMAIP1", "PPP2R5B", "PPP3R1", "PPT1", "PRF1", "PSEN1",
+             "PSEN2", "PTK2", "RARA", "RELA", "RETSAT", "RHOB", "RHOT2",
+             "RNASEL", "ROCK1", "SAT1", "SATB1", "SC5D", "SLC20A1", "SMAD7",
+             "SOD1", "SOD2", "SPTAN1", "SQSTM1", "TAP1", "TGFB2", "TGFBR3",
+             "TIMP1", "TIMP2", "TIMP3", "TNF", "TNFRSF12A", "TNFSF10", "TOP2A",
+             "TSPO", "TXNIP", "VDAC2", "WEE1", "XIAP"),
+        # 193
+        "mouse": ("Ercc5", "Serpinb5", "Inhbb", "Steap3", "Btg2", "Phlda3", "Tnni1",
+                  "Rgs16", "Ier5", "Slc19a2", "Adck3", "Ephx1", "Ptpn14", "Atf3",
+                  "Notch1", "Rxra", "Ralgds", "Ak1", "Stom", "Ddb2", "Cd82", "Il1a",
+                  "Pcna", "Bmp2", "Trib3", "Procr", "Blcap", "Ada", "Fgf13", "Irak1",
+                  "Tspyl2", "Sat1", "Zmat3", "Hspa4l", "Slc7a11", "Tm4sf1", "Rap2b",
+                  "Fbxw7", "S100a4", "S100a10", "Txnip", "Nhlh2", "Dnttip2", "Clca2",
+                  "Wwp1", "Klf4", "Ikbkap", "Cdkn2a", "Cdkn2b", "Jun", "Slc35d1",
+                  "Plk3", "Rnf19b", "Sfn", "Fuca1", "Epha2", "Wrap73", "Mxd4",
+                  "Rchy1", "Iscu", "Triap1", "Prkab1", "Trafd1", "Pom121", "Pdgfa",
+                  "Gadd45a", "Vamp8", "Retsat", "Tprkb", "Tgfa", "Mxd1", "Sec61a1",
+                  "Xpc", "Ccnd2", "H2afj", "Ldhb", "Lrmp", "Tm7sf3", "Tgfb1",
+                  "Sertad3", "Cebpa", "Klk8", "Bax", "Ppp1r15a", "Rpl18", "Aen",
+                  "Rrp8", "Ccp110", "Nupr1", "Ptpre", "Hras", "Eps8l2", "Ctsd",
+                  "Cd81", "Perp", "Rps12", "Tpd52l1", "Sesn1", "Foxo3", "Ddit4",
+                  "Zfp365", "Prmt2", "Mknk2", "Dram1", "Apaf1", "Btg1", "Mdm2",
+                  "Ddit3", "Gls2", "Dgka", "Cdkn2aip", "Hmox1", "Rrad", "Cdh13",
+                  "Osgin1", "Cgrrf1", "Abhd4", "Kif13b", "Rb1", "Nudt15", "Tsc22d1",
+                  "Casp1", "St14", "Ei24", "Vwa5a", "Zbtb16", "Rps27l", "Mapkapk3",
+                  "Ip6k2", "Tcn2", "Lif", "Upp1", "Ccng1", "Cyfip2", "Gnb2l1",
+                  "Hint1", "Gm2a", "Hist3h2a", "Alox8", "Trp53", "Tax1bp3", "Traf4",
+                  "Cdk5r1", "Ppm1d", "Rad51c", "Tob1", "Krt17", "Hexim1", "Fdxr",
+                  "Itgb4", "Sphk1", "Rhbdf2", "Baiap2", "Dcxr", "Hist1h1c", "Ninj1",
+                  "Nol8", "F2r", "Ankra2", "Plk2", "Sdc1", "Gpx2", "Zfp36l1", "Fos",
+                  "Ccnk", "Jag2", "Ndrg1", "Pmm1", "Plxnb2", "Vdr", "Csrnp2",
+                  "Acvr1b", "Sp1", "Abat", "Socs1", "Abcc5", "Trp63", "Fam162a",
+                  "App", "Rab40c", "Bak1", "Def6", "Cdkn1a", "Tap1", "Ier3", "Polh",
+                  "Ccnd3", "Hbegf", "Hdac3", "Rad9a", "Ctsf", "Slc3a2", "Fas"),
+    }
+    # fmt: on
+
+    @classmethod
+    def proliferation_markers(
+        cls, organism: Literal["human", "mouse"]
+    ) -> Sequence[str]:
+        """Get proliferation markers for ``organism``."""
+        try:
+            return cls._proliferation_markers[organism]
+        except KeyError:
+            raise NotImplementedError(
+                f"Proliferation markers for `{organism}` are not yet implemented."
+            ) from None
+
+    @classmethod
+    def apoptosis_markers(cls, organism: Literal["human", "mouse"]) -> Sequence[str]:
+        """Get apoptosis markers for ``organism``."""
+        try:
+            return cls._apoptosis_markers[organism]
+        except KeyError:
+            raise NotImplementedError(
+                f"Apoptosis markers for `{organism}` are not yet implemented."
+            ) from None
diff --git a/cellrank/external/kernels/_wot_kernel.py b/cellrank/external/kernels/_wot_kernel.py
new file mode 100644
index 000000000..b9348f005
--- /dev/null
+++ b/cellrank/external/kernels/_wot_kernel.py
@@ -0,0 +1,503 @@
+from types import MappingProxyType
+from typing import Any, Dict, Tuple, Union, Mapping, Optional
+
+from tqdm.auto import tqdm
+from typing_extensions import Literal
+
+import scanpy as sc
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank.ul._docs import d, inject_docs
+from cellrank.tl._utils import _maybe_subset_hvgs
+from cellrank.external.kernels._utils import MarkerGenes
+from cellrank.tl.kernels._exp_time_kernel import LastTimePoint
+
+import numpy as np
+import pandas as pd
+
+_error = None
+try:
+    import wot
+
+    from cellrank.tl.kernels import TransportMapKernel as Kernel
+except ImportError as e:
+    from cellrank.external.kernels._import_error_kernel import ErroredKernel as Kernel
+
+    _error = e
+    wot = None
+
+
+class nstr(str):  # used for params+cache (precomputed cost matrices)
+    """String class that is not equal to any other string."""
+
+    def __eq__(self, other: str) -> bool:
+        return False
+
+
+@d.dedent
+class WOTKernel(Kernel, error=_error):
+    """
+    Waddington optimal transport kernel from :cite:`schiebinger:19`.
+
+    This class requires the `wot` package, which can be installed as
+    `pip install git+https://github.com/broadinstitute/wot`.
+
+    Parameters
+    ----------
+    %(adata)s
+    %(backward)s
+    time_key
+        Key in :attr:`adata` ``.obs`` where experimental time is stored.
+        The experimental time can be of either of a numeric or an ordered categorical type.
+    %(cond_num)s
+
+    Examples
+    --------
+    Workflow::
+
+        # import packages, load data
+        import scanpy as sc
+        import cellrank as cr
+        adata = cr.datasets.lung()
+
+        # filter, normalize and annotate highly variable genes
+        sc.pp.filter_genes(adata, min_cells=10)
+        sc.pp.normalize_total(adata)
+        sc.pp.log1p(adata)
+        sc.pp.highly_variable_genes(adata)
+
+        # estimate proliferation and apoptosis from gene sets (see. e.g. WOT tutorial for example lists)
+        proliferation_genes = ...
+        apoptosis_genes = ...
+        sc.tl.score_genes(adata, gene_list=proliferation_genes, score_name='proliferation')
+        sc.tl.score_genes(adata, gene_list=apoptosis_genes, score_name='apoptosis')
+
+        # initialize kernel, estimate initial growth rate based on scores from above
+        from cellrank.external.kernels import WOTKernel
+        ot = WOTKernel(adata, time_key='day')
+        ot.compute_initial_growth_rates(proliferation_key='proliferation',
+                                        apoptosis_key='apoptosis',
+                                        key_added='initial_growth_rates')
+
+        # compute transport maps, aggregate into one single transition matrix
+        ot.compute_transition_matrix(growth_rate_key='initial_growth_rates', growth_iters=3)
+    """
+
+    __import_error_message__ = (
+        "Unable to import the kernel. Please install `wot` first as "
+        "`pip install git+https://github.com/broadinstitute/wot`."
+    )
+
+    def __init__(
+        self,
+        adata: AnnData,
+        backward: bool = False,
+        time_key: str = "exp_time",
+        compute_cond_num: bool = False,
+    ):
+        super().__init__(
+            adata,
+            backward=backward,
+            time_key=time_key,
+            compute_cond_num=compute_cond_num,
+        )
+
+        # WOT's requirements
+        cats = self.experimental_time.cat.categories
+        self._exp_time = self.experimental_time.cat.rename_categories(
+            dict(zip(cats, map(float, cats)))
+        )
+        self.adata.obs[self._time_key] = self.experimental_time
+        self._growth_rates = None
+
+    def _read_from_adata(
+        self,
+        conn_key: Optional[str] = "connectivities",
+        read_conn: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        super()._read_from_adata(conn_key=conn_key, read_conn=False, **kwargs)
+
+    def compute_initial_growth_rates(
+        self,
+        proliferation_key: Optional[str] = None,
+        apoptosis_key: Optional[str] = None,
+        organism: Optional[Literal["human", "mouse"]] = None,
+        beta_min: float = 0.3,
+        beta_max: float = 1.7,
+        delta_min: float = 0.3,
+        delta_max: float = 1.7,
+        key_added: Optional[str] = None,
+        **kwargs: Any,
+    ) -> Optional[pd.Series]:
+        r"""
+        Estimate initial growth rates using a birth-death process as described in :cite:`schiebinger:19`.
+
+        The doubling time is defined as :math:`\frac{\ln 2}{\beta - \delta}` (similarly defined for half-time).
+        The logistic function is used to transform the birth/death rate scores and to smoothly interpolate between
+        specified minimum and maximum birth/death rates.
+
+        Parameters
+        ----------
+        proliferation_key
+            Key in :attr:`adata` ``.obs`` where the birth rate score is saved.
+        apoptosis_key
+            Key in :attr:`adata` ``.obs`` where the death rate score is saved.
+        organism
+            Organism for which to calculate the birth/death scores, if they cannot be found in :attr:`adata`.
+            In this case, :func:`scanpy.tl.score_genes` is used to calculate the scores based on an organism-dependent
+            set of marker genes for proliferation and apoptosis.
+        beta_min
+            Minimum birth rate.
+        beta_max
+            Maximum birth rate.
+        delta_min
+            Minimum death rate.
+        delta_max
+            Maximum death rate.
+        key_added
+            Key in :attr:`adata` ``.obs`` where to add the estimated growth rates. If `None`, just return them.
+        kwargs
+            Keyword arguments for :func:`scanpy.tl.score_genes`. Only used when ``proliferation_key``
+            or ``apoptosis_key`` cannot be found in :attr:`adata` ``.obs``.
+
+        Returns
+        -------
+        :class:`pandas.Series`
+            The estimated initial growth rates if ``key_added=None``, otherwise `None`.
+
+        Notes
+        -----
+        If you don't have access to proliferation/apoptosis gene sets, you can use the ones defined in :mod:`cellrank`
+        for a specific organism. Alternatively, you can also use WOT without an estimate of initial growth rates. In
+        that case, make sure to use several iterations in
+        :meth:`cellrank.external.kernels.WOTKernel.compute_transition_matris` by increasing the `growth_iters`
+        parameter. A value around 3 works well in most cases.
+
+        The markers used here were taken from the following sources:
+
+            - human, proliferation - :cite:`tirosh:16:science`.
+            - human, apoptosis - `Hallmark Apoptosis, MSigDB <https://www.gsea-msigdb.org/gsea/msigdb/cards/HALLMARK_APOPTOSIS>`_.
+            - mouse, proliferation - :cite:`tirosh:16:nature`.
+            - mouse, apoptosis - `Hallmark P53 Pathway, MSigDB <https://www.gsea-msigdb.org/gsea/msigdb/cards/HALLMARK_P53_PATHWAY>`_.
+
+        For more information about WOT, see the official `tutorial <https://broadinstitute.github.io/wot/tutorial/>`_.
+        """  # noqa: E501
+
+        def get_scores(
+            key: str,
+            *,
+            kind: Literal["proliferation", "apoptosis"],
+            organism: Optional[Literal["human", "mouse"]],
+        ) -> np.ndarray:
+            try:
+                return np.asarray(self.adata.obs[key])
+            except KeyError:
+                if organism is None:
+                    raise KeyError(
+                        f"Unable to find `{kind}` scores in `adata.obs[{kind!r}]`. Consider specifying `organism=...`."
+                    ) from None
+
+                logg.info(f"Computing `{kind}` scores")
+                score_name = f"{kind}_score" if key is None else key
+
+                sc.tl.score_genes(
+                    self.adata,
+                    gene_list=getattr(MarkerGenes, (f"{kind}_markers"))(organism),
+                    score_name=score_name,
+                    **kwargs,
+                )
+
+                return get_scores(score_name, kind=kind, organism=None)
+
+        def logistic(x: np.ndarray, L: float, k: float, x0: float = 0) -> np.ndarray:
+            return L / (1 + np.exp(-k * (x - x0)))
+
+        def gen_logistic(
+            p: np.ndarray, beta_max: float, beta_min: float, center: float, width: float
+        ) -> np.ndarray:
+            return beta_min + logistic(p, L=beta_max - beta_min, k=4 / width, x0=center)
+
+        def beta(
+            p: np.ndarray,
+            beta_max: float = 1.7,
+            beta_min: float = 0.3,
+            center: float = 0.25,
+        ) -> np.ndarray:
+            return gen_logistic(p, beta_max, beta_min, center, width=0.5)
+
+        def delta(
+            a: np.ndarray,
+            delta_max: float = 1.7,
+            delta_min: float = 0.3,
+            center: float = 0.1,
+        ) -> np.ndarray:
+            return gen_logistic(a, delta_max, delta_min, center, width=0.2)
+
+        birth = beta(
+            get_scores(proliferation_key, kind="proliferation", organism=organism),
+            beta_min=beta_min,
+            beta_max=beta_max,
+        )
+        death = delta(
+            get_scores(apoptosis_key, kind="apoptosis", organism=organism),
+            delta_min=delta_min,
+            delta_max=delta_max,
+        )
+        gr = np.exp(birth - death)
+
+        if key_added is None:
+            return pd.Series(gr, index=self.adata.obs_names)
+        self.adata.obs[key_added] = gr
+
+    @inject_docs(ltp=LastTimePoint)
+    def compute_transition_matrix(
+        self,
+        cost_matrices: Optional[
+            Union[str, Mapping[Tuple[float, float], np.ndarray]]
+        ] = None,
+        lambda1: float = 1,
+        lambda2: float = 50,
+        epsilon: float = 0.05,
+        growth_iters: int = 1,
+        solver: Literal["fixed_iters", "duality_gap"] = "duality_gap",
+        growth_rate_key: Optional[str] = None,
+        use_highly_variable: Optional[Union[str, bool]] = True,
+        last_time_point: Literal["uniform", "diagonal", "connectivities"] = "uniform",
+        conn_kwargs: Mapping[str, Any] = MappingProxyType({}),
+        **kwargs: Any,
+    ) -> "WOTKernel":
+        """
+        Compute transition matrix using Waddington OT :cite:`schiebinger:19`.
+
+        Computes transport maps linking together pairs of time points for time-series single cell data using unbalanced
+        optimal transport, taking into account cell birth and death rates. From the sequence of transition maps linking
+        pairs of sequential time points, we construct one large transition matrix which contains the normalized
+        transport maps as blocks on the 1st upper diagonal.
+
+        Parameters
+        ----------
+        cost_matrices
+            Cost matrices for each consecutive pair of time points.
+            If a :class:`str`, it specifies a key in :attr:`adata` ``.layers`` or :attr:`adata` ``.obsm``
+            containing cell features that are used to compute cost matrices. If `None`, use `WOT`'s default, i.e.
+            compute distances in PCA space derived from :attr:`adata` ``.X`` for each time point pair separately.
+        lambda1
+            Regularization parameter for the marginal constraint on :math:`p`, the transport map row sums.
+            Smaller value is useful when precise information about the growth rate is not present.
+        lambda2
+            Regularization parameter for the marginal constraint on :math:`q`, the transport map column sums.
+        epsilon
+            Entropy regularization parameter. Larger value gives more entropic descendant distributions.
+        growth_iters
+            Number of iterations for growth rate estimates. If growth rates are not known, consider using more
+            iterations.
+        solver
+            Which solver to use.
+        growth_rate_key
+            Key in :attr:`adata` ``.obs`` where initial cell growth rates are stored.
+            See :meth:`cellrank.external.kernels.WOTKernel.compute_initial_growth_rates` on how to estimate them.
+        use_highly_variable
+            Key in :attr:`adata` ``.var`` where highly variable genes are stored.
+            If `True`, use `'highly_variable'`. If `None`, use all genes.
+        last_time_point
+            How to define transitions within the last time point. Valid options are:
+
+                - `{ltp.UNIFORM.s!r}` - row-normalized matrix of 1s for transitions within the last time point.
+                - `{ltp.DIAGONAL.s!r}` - diagonal matrix with 1s on the diagonal.
+                - `{ltp.CONNECTIVITIES.s!r}` - use transitions from :class:`cellrank.tl.kernels.ConnectivityKernel`
+                  derived from the last time point subset of :attr:`adata`.
+        conn_kwargs
+            Keyword arguments for :func:`scanpy.pp.neighbors`, when using ``last_time_point={ltp.CONNECTIVITIES.s!r}``.
+            Can contain `'density_normalize'` for
+            :meth:`cellrank.tl.kernels.ConnectivityKernel.compute_transition_matrix`.
+        kwargs
+            Additional keyword arguments for OT configuration.
+
+        Returns
+        -------
+        :class:`cellrank.external.kernels.WOTKernel`
+            Makes :attr:`transition_matrix`, :attr:`transport_maps` and :attr:`growth_rates` available.
+            It also modifies :attr:`anndata.AnnData.obs` with the following key:
+
+                - `'estimated_growth_rates'` - the estimated final growth rates.
+
+        Notes
+        -----
+        For more information about WOT, see the official `tutorial <https://broadinstitute.github.io/wot/tutorial/>`_.
+        """
+        # disallow these params (because they e.g. subset the data)
+        _ = kwargs.pop("cell_day_filter", None)
+        _ = kwargs.pop("covariate_field", None)
+        _ = kwargs.pop("ncounts", None)
+        _ = kwargs.pop("ncells", None)
+        kwargs["lambda1"] = lambda1
+        kwargs["lambda2"] = lambda2
+        kwargs["epsilon"] = epsilon
+        kwargs["growth_iters"] = max(growth_iters, 1)
+        last_time_point = LastTimePoint(last_time_point)
+
+        start = logg.info(
+            "Computing transition matrix using Waddington optimal transport"
+        )
+
+        adata = _maybe_subset_hvgs(self.adata, use_highly_variable=use_highly_variable)
+        cost_matrices, cmat_param = self._generate_cost_matrices(adata, cost_matrices)
+        if self._reuse_cache(
+            {
+                "cost_matrices": cmat_param,
+                "solver": solver,
+                "growth_rate_key": growth_rate_key,
+                "use_highly_variable": use_highly_variable,
+                "last_time_point": last_time_point.s,
+                **kwargs,
+            },
+            time=start,
+        ):
+            return self
+
+        tmap = self._compute_pairwise_tmaps(
+            adata,
+            cost_matrices=cost_matrices,
+            solver=solver,
+            growth_rate_field=growth_rate_key,
+            **kwargs,
+        )
+        tmap = self._restich_tmaps(tmap, last_time_point, conn_kwargs=conn_kwargs)
+        self._growth_rates = tmap.obs
+
+        self._compute_transition_matrix(
+            matrix=tmap.X,
+            density_normalize=False,
+            check_irreducibility=False,
+        )
+        self.adata.obs["estimated_growth_rates"] = self.growth_rates[f"g{growth_iters}"]
+
+        logg.info("    Finish", time=start)
+
+        return self
+
+    def _compute_pairwise_tmaps(
+        self,
+        adata: AnnData,
+        cost_matrices: Optional[
+            Union[str, Mapping[Tuple[float, float], np.ndarray]]
+        ] = None,
+        solver: Literal["fixed_iters", "duality_gap"] = "duality_gap",
+        growth_rate_field: Optional[str] = None,
+        **kwargs,
+    ) -> Dict[Tuple[float, float], AnnData]:
+        self._ot_model = wot.ot.OTModel(
+            adata,
+            day_field=self._time_key,
+            covariate_field=None,
+            growth_rate_field=growth_rate_field,
+            solver=solver,
+            **kwargs,
+        )
+
+        self._tmaps: Dict[Tuple[float, float], AnnData] = {}
+        start = logg.info(
+            f"Computing transport maps for `{len(cost_matrices)}` time pairs"
+        )
+        for tpair, cost_matrix in tqdm(cost_matrices.items(), unit="time pair"):
+            tmap: AnnData = self._ot_model.compute_transport_map(
+                *tpair, cost_matrix=cost_matrix
+            )
+            nans = int(np.sum(np.isnan(tmap.X)))
+            if nans:
+                raise ValueError(
+                    f"Encountered `{nans}` NaN values for time pair `{tpair}`."
+                )
+            self._tmaps[tpair] = tmap
+
+        logg.info("    Finish", time=start)
+
+        return self._tmaps
+
+    def _generate_cost_matrices(
+        self,
+        adata: AnnData,
+        cost_matrices: Optional[
+            Union[str, Mapping[Tuple[float, float], np.ndarray]]
+        ] = None,
+    ) -> Tuple[Mapping[Tuple[float, float], Optional[np.ndarray]], str]:
+        timepoints = self.experimental_time.cat.categories
+        timepoints = list(zip(timepoints[:-1], timepoints[1:]))
+
+        if cost_matrices is None:
+            logg.info("Using default cost matrices")
+            return {tpair: None for tpair in timepoints}, "default"
+
+        if isinstance(cost_matrices, dict):
+            logg.info("Using precomputed cost matrices")
+
+            cmats = {}
+            for tpair in timepoints:
+                if tpair not in cost_matrices:
+                    logg.warning(
+                        f"Unable to find cost matrix for pair `{tpair}`. Using default"
+                    )
+                cmats[tpair] = cmat = cost_matrices.get(tpair, None)
+                if cmat is not None:
+                    n_start = len(np.where(self.experimental_time == tpair[0])[0])
+                    n_end = len(np.where(self.experimental_time == tpair[1])[0])
+                    try:
+                        if cmat.shape != (n_start, n_end):
+                            raise ValueError(
+                                f"Expected cost matrix for time pair `{tpair}` to be "
+                                f"of shape `{(n_start, n_end)}`, found `{cmat.shape}`."
+                            )
+                    except AttributeError:
+                        logg.warning(
+                            f"Unable to verify whether supplied cost matrix for time pair `{tpair}` "
+                            f"has the correct shape `{(n_start, n_end)}`"
+                        )
+
+            # prevent equality comparison when comparing with cache
+            return cmats, nstr("precomputed")
+
+        if isinstance(cost_matrices, str):
+            logg.info(f"Computing cost matrices using `{cost_matrices!r}` key")
+            if cost_matrices == "X":
+                cost_matrices = None
+
+            try:
+                features = adata._get_X(layer=cost_matrices)
+                modifier = "layer"
+            except KeyError:
+                try:
+                    features = adata.obsm[cost_matrices]
+                    modifier = "obsm"
+                except KeyError:
+                    raise KeyError(
+                        f"Unable to find key `{cost_matrices!r}` in `adata.layers` or `adata.obsm`."
+                    ) from None
+
+            cmats = {}
+            for tpair in tqdm(timepoints, unit="cost matrix"):
+                start_ixs = np.where(self.experimental_time == tpair[0])[0]
+                end_ixs = np.where(self.experimental_time == tpair[1])[0]
+
+                # being sparse is handled in WOT's function below
+                cmats[tpair] = wot.ot.OTModel.compute_default_cost_matrix(
+                    features[start_ixs], features[end_ixs]
+                )
+
+            return cmats, f"{modifier}:{cost_matrices}"
+
+        raise NotImplementedError(
+            f"Specifying cost matrices as "
+            f"`{type(cost_matrices).__name__}` is not yet implemented."
+        )
+
+    @property
+    def growth_rates(self) -> Optional[pd.DataFrame]:
+        """Estimated cell growth rates for each growth rate iteration."""
+        return self._growth_rates
+
+    def __invert__(self) -> "WOTKernel":
+        super().__invert__()
+        # because WOT reads from `adata`
+        self.adata.obs[self._time_key] = self.experimental_time
+        return self
diff --git a/cellrank/external/models/_import_error_model.py b/cellrank/external/models/_import_error_model.py
index 7868dfb03..f4029af2f 100644
--- a/cellrank/external/models/_import_error_model.py
+++ b/cellrank/external/models/_import_error_model.py
@@ -1,10 +1,10 @@
 from typing import Optional
 
-import numpy as np
-
 from cellrank.ul.models import BaseModel
 from cellrank.external._error_mixin import ImportErrorMixin
 
+import numpy as np
+
 
 class ErroredModel(ImportErrorMixin, BaseModel):
     """
diff --git a/cellrank/pl/__init__.py b/cellrank/pl/__init__.py
index ed0547ef2..09ebf6ea5 100644
--- a/cellrank/pl/__init__.py
+++ b/cellrank/pl/__init__.py
@@ -1,6 +1,7 @@
 from cellrank.pl._graph import graph
 from cellrank.pl._heatmap import heatmap
 from cellrank.pl._lineages import lineages, lineage_drivers
+from cellrank.pl._log_odds import log_odds
 from cellrank.pl._gene_trend import gene_trends
 from cellrank.pl._cluster_fates import cluster_fates
 from cellrank.pl._cluster_lineage import cluster_lineage
diff --git a/cellrank/pl/_circular_projection.py b/cellrank/pl/_circular_projection.py
index 0e9d092b1..7c49c775f 100644
--- a/cellrank/pl/_circular_projection.py
+++ b/cellrank/pl/_circular_projection.py
@@ -6,15 +6,6 @@
 
 import scvelo as scv
 from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from sklearn.metrics import pairwise_distances
-
-import matplotlib.pyplot as plt
-from matplotlib.colors import LogNorm, LinearSegmentedColormap
-from matplotlib.collections import LineCollection
-
 from cellrank import logging as logg
 from cellrank.tl import Lineage
 from cellrank.ul._docs import d
@@ -24,6 +15,14 @@
 from cellrank.tl._lineage import PrimingDegree
 from cellrank.tl._constants import ModeEnum, AbsProbKey
 
+import numpy as np
+import pandas as pd
+from sklearn.metrics import pairwise_distances
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import LogNorm, LinearSegmentedColormap
+from matplotlib.collections import LineCollection
+
 
 class LineageOrder(ModeEnum):  # noqa: D101
     DEFAULT = "default"
@@ -97,10 +96,10 @@ def circular_projection(
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[int] = None,
     save: Optional[Union[str, Path]] = None,
-    **kwargs,
+    **kwargs: Any,
 ):
     r"""
-    Plot absorption probabilities on a circular embedding as done in [Velten17]_.
+    Plot absorption probabilities on a circular embedding as in :cite:`velten:17`.
 
     Parameters
     ----------
@@ -108,9 +107,9 @@ def circular_projection(
     keys
         Keys in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names`. Additional keys are:
 
-            - `'kl_divergence'` - as in [Velten17]_, computes KL-divergence between the fate probabilities of a cell
-              and the average fate probabilities. See ``early_cells`` for more information.
-            - `'entropy'` - as in [Setty19]_, computes entropy over a cells fate probabilities.
+            - `'kl_divergence'` - as in :cite:`velten:17`, computes KL-divergence between the fate probabilities
+              of a cell and the average fate probabilities. See ``early_cells`` for more information.
+            - `'entropy'` - as in :cite:`setty:19`, computes entropy over a cells fate probabilities.
 
     %(backward)s
     lineages
@@ -123,7 +122,7 @@ def circular_projection(
         Can be one of the following:
 
             - `None` - it will determined automatically, based on the number of lineages.
-            - `'optimal'` - order the lineages optimally by solving the Travelling salesman problem (TSP).
+            - `'optimal'` - order lineages optimally by solving the Travelling salesman problem (TSP).
               Recommended for <= `20` lineages.
             - `'default'` - use the order as specified in ``lineages``.
 
@@ -131,7 +130,7 @@ def circular_projection(
         Metric to use when constructing pairwise distance matrix when ``lineage_order = 'optimal'``. For available
         options, see :func:`sklearn.metrics.pairwise_distances`.
     normalize_by_mean
-        If `True`, normalize each lineage by its mean probability, as done in [Velten17]_.
+        If `True`, normalize each lineage by its mean probability, as done in :cite:`velten:17`.
     ncols
         Number of columns when plotting multiple ``keys``.
     space
@@ -164,8 +163,8 @@ def circular_projection(
     %(just_plots)s
         Also updates ``adata`` with the following fields:
 
-            - :attr:`anndata.AnnData.obsm` ``['{key_added}']``: the circular projection.
-            - :attr:`anndata.AnnData.obs` ``['to_{initial,terminal}_states_{method}']``: the priming degree,
+            - :attr:`anndata.AnnData.obsm` ``['{key_added}']`` - the circular projection.
+            - :attr:`anndata.AnnData.obs` ``['to_{initial,terminal}_states_{method}']`` - the priming degree,
               if a method is present in ``keys``.
     """
     if labeldistance is not None and labeldistance < 0:
@@ -209,8 +208,8 @@ def circular_projection(
 
     probs: Lineage = adata.obsm[lineage_key][lineages]
     n_lin = probs.shape[1]
-    if n_lin <= 2:
-        raise ValueError(f"Expected at least `3` lineages, found `{n_lin}`")
+    if n_lin < 3:
+        raise ValueError(f"Expected at least `3` lineages, found `{n_lin}`.")
 
     X = probs.X.copy()
     if normalize_by_mean:
@@ -220,7 +219,9 @@ def circular_projection(
         X = np.nan_to_num(X, nan=1.0 / n_lin, copy=False)
 
     if lineage_order is None:
-        lineage_order = LineageOrder.OPTIMAL if n_lin <= 15 else LineageOrder.DEFAULT
+        lineage_order = (
+            LineageOrder.OPTIMAL if 3 < n_lin <= 20 else LineageOrder.DEFAULT
+        )
         logg.debug(f"Set ordering to `{lineage_order}`")
     lineage_order = LineageOrder(lineage_order)
 
diff --git a/cellrank/pl/_cluster_fates.py b/cellrank/pl/_cluster_fates.py
index 1f0d3aeaf..649cca5a3 100644
--- a/cellrank/pl/_cluster_fates.py
+++ b/cellrank/pl/_cluster_fates.py
@@ -6,6 +6,13 @@
 from pathlib import Path
 from collections import OrderedDict as odict
 
+from cellrank import logging as logg
+from cellrank.ul._docs import d, inject_docs
+from cellrank.pl._utils import _position_legend
+from cellrank.tl._utils import RandomKeys, save_fig, _unique_order_preserving
+from cellrank.ul._utils import valuedispatch
+from cellrank.tl._constants import ModeEnum, DirPrefix, AbsProbKey, TerminalStatesPlot
+
 import numpy as np
 import pandas as pd
 from scipy.sparse import csr_matrix
@@ -15,13 +22,6 @@
 import matplotlib.pyplot as plt
 from matplotlib import cm
 
-from cellrank import logging as logg
-from cellrank.ul._docs import d, inject_docs
-from cellrank.pl._utils import _position_legend
-from cellrank.tl._utils import RandomKeys, save_fig, _unique_order_preserving
-from cellrank.ul._utils import valuedispatch
-from cellrank.tl._constants import ModeEnum, DirPrefix, AbsProbKey, TerminalStatesPlot
-
 AnnData = TypeVar("AnnData")
 
 
@@ -59,7 +59,8 @@ def cluster_fates(
     Plot aggregate lineage probabilities at a cluster level.
 
     This can be used to investigate how likely a certain cluster is to go to the %(terminal)s states,or in turn to have
-    descended from the %(initial)s states. For mode `{m.PAGA.s!r}` and `{m.PAGA_PIE.s!r}`, we use *PAGA*, see [Wolf19]_.
+    descended from the %(initial)s states.
+    For mode `{m.PAGA.s!r}` and `{m.PAGA_PIE.s!r}`, we use *PAGA*, see :cite:`wolf:19`.
 
     Parameters
     ----------
diff --git a/cellrank/pl/_cluster_lineage.py b/cellrank/pl/_cluster_lineage.py
index c9b359408..1b39fc42b 100644
--- a/cellrank/pl/_cluster_lineage.py
+++ b/cellrank/pl/_cluster_lineage.py
@@ -1,15 +1,11 @@
 """Cluster lineages module."""
 
 from types import MappingProxyType
-from typing import Dict, Tuple, Union, TypeVar, Optional, Sequence
+from typing import Dict, Tuple, Union, Optional, Sequence
 from pathlib import Path
-from collections import Iterable
-
-import numpy as np
-from sklearn.preprocessing import StandardScaler
-
-import matplotlib.pyplot as plt
 
+import scanpy as sc
+from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.pl._utils import (
@@ -20,14 +16,19 @@
     _time_range_type,
     _create_callbacks,
     _input_model_type,
+    _get_sorted_colors,
     _return_model_type,
 )
 from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul._utils import _get_n_cores, _check_collection
 from cellrank.tl._constants import _DEFAULT_BACKEND, AbsProbKey
 
-AnnData = TypeVar("AnnData")
-Queue = TypeVar("Queue")
+import numpy as np
+from sklearn.preprocessing import StandardScaler
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import ListedColormap, is_color_like
+from matplotlib.gridspec import GridSpec, GridSpecFromSubplotSpec
 
 
 @d.dedent
@@ -41,6 +42,9 @@ def cluster_lineage(
     clusters: Optional[Sequence[str]] = None,
     n_points: int = 200,
     time_key: str = "latent_time",
+    covariate_key: Optional[Union[str, Sequence[str]]] = None,
+    ratio: float = 0.05,
+    cmap: Optional[str] = "viridis",
     norm: bool = True,
     recompute: bool = False,
     callback: _callback_type = None,
@@ -64,7 +68,7 @@ def cluster_lineage(
     """
     Cluster gene expression trends within a lineage and plot the clusters.
 
-    This function is based on Palantir, see [Setty19]_. It can be used to discover modules of genes that drive
+    This function is based on Palantir, see :cite:`setty:19`. It can be used to discover modules of genes that drive
     development along a given lineage. Consider running this function on a subset of genes which are potential
     lineage drivers, identified e.g. by running :func:`cellrank.tl.lineage_drivers`.
 
@@ -76,7 +80,7 @@ def cluster_lineage(
     lineage
         Name of the lineage for which to cluster the genes.
     %(backward)s
-    %(time_ranges)s
+    %(time_range)s
     clusters
         Cluster identifiers to plot. If `None`, all clusters will be considered. Useful when
         plotting previously computed clusters.
@@ -84,6 +88,12 @@ def cluster_lineage(
         Number of points used for prediction.
     time_key
         Key in ``adata.obs`` where the pseudotime is stored.
+    covariate_key
+        Key(s) in ``adata.obs`` containing observations to be plotted at the bottom of each plot.
+    ratio
+        Height ratio of each covariate in ``covariate_key``.
+    cmap
+        Colormap to use for continuous covariates in ``covariate_key``.
     norm
         Whether to z-normalize each trend to have zero mean, unit variance.
     recompute
@@ -122,9 +132,6 @@ def cluster_lineage(
 
     """
 
-    import scanpy as sc
-    from anndata import AnnData as _AnnData
-
     lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
     if lineage_key not in adata.obsm:
         raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")
@@ -165,8 +172,10 @@ def cluster_lineage(
             logg.debug("Normalizing trends")
             _ = StandardScaler(copy=False).fit_transform(trends)
 
-        trends = _AnnData(trends.T)
+        mod = next(mod for tmp in all_models.values() for mod in tmp.values())
+        trends = AnnData(trends.T)
         trends.obs_names = genes
+        trends.var["x_test"] = x_test = mod.x_test
 
         # sanity check
         if trends.n_obs != len(genes):
@@ -210,6 +219,7 @@ def cluster_lineage(
         all_models = None
         logg.info(f"Loading data from `adata.uns[{key!r}]`")
         trends = adata.uns[key]
+        x_test = trends.var["x_test"]
 
     if "clusters" not in trends.obs:
         raise KeyError("Unable to find the clustering in `trends.obs['clusters']`.")
@@ -223,24 +233,18 @@ def cluster_lineage(
                 f"Valid options are `{list(trends.obs['clusters'].cat.categories)}`."
             )
 
-    nrows = int(np.ceil(len(clusters) / ncols))
-    fig, axes = plt.subplots(
-        nrows,
-        ncols,
-        figsize=(ncols * 10, nrows * 10) if figsize is None else figsize,
-        sharey=sharey,
-        dpi=dpi,
-    )
-
-    if not isinstance(axes, Iterable):
-        axes = [axes]
-    axes = np.ravel(axes)
+    def plot_cluster(row, col, cluster, sharey_ax: Optional[str] = None):
+        gss = GridSpecFromSubplotSpec(
+            row_delta,
+            1,
+            subplot_spec=gs[row : row + row_delta, col],
+            hspace=0,
+            height_ratios=[1] + [ratio] * (row_delta - 1),
+        )
+        ax = fig.add_subplot(gss[0, 0], sharey=sharey_ax)
 
-    j = 0
-    for j, (ax, c) in enumerate(zip(axes, clusters)):  # noqa
         data = trends[trends.obs["clusters"] == c].X
-        mean, sd = np.mean(data, axis=0), np.var(data, axis=0)
-        sd = np.sqrt(sd)
+        mean, sd = np.mean(data, axis=0), np.std(data, axis=0)
 
         for i in range(data.shape[0]):
             ax.plot(data[i], color="gray", lw=0.5)
@@ -252,14 +256,57 @@ def cluster_lineage(
             range(len(mean)), mean - sd, mean + sd, color="black", alpha=0.1
         )
 
-        ax.set_title(f"Cluster {c}")
+        ax.set_title(f"cluster {cluster}")
         ax.set_xticks([])
-
-        if not sharey:
+        if sharey:
             ax.set_yticks([])
+        ax.margins(0)
+
+        if covariate_colors is not None:
+            for i, colors in enumerate(covariate_colors):
+                ax_clusters = fig.add_subplot(gss[i + 1, 0])
+                if is_color_like(colors[0]):  # e.g. categorical
+                    cm = ListedColormap(colors, N=len(colors))
+                    ax_clusters.imshow(np.arange(cm.N)[None, :], cmap=cm, aspect="auto")
+                else:
+                    cm = plt.get_cmap(cmap)
+                    ax_clusters.imshow(colors[None, :], cmap=cm, aspect="auto")
+                ax_clusters.set_xticks([])
+                ax_clusters.set_yticks([])
+
+            ax_clusters.set_xticks(np.linspace(0, len(colors), 5))
+            ax_clusters.set_xticklabels(
+                [f"{v:.3f}" for v in np.linspace(tmin, tmax, 5)]
+            )
+
+        return ax if sharey else None
 
-    for j in range(j + 1, len(axes)):
-        axes[j].remove()
+    nrows = int(np.ceil(len(clusters) / ncols))
+    fig = plt.figure(
+        dpi=dpi,
+        figsize=(ncols * 10, nrows * 10) if figsize is None else figsize,
+        tight_layout=True,
+    )
+
+    if covariate_key is None:
+        covariate_colors, row_delta = None, 1
+    else:
+        tmin, tmax = np.min(x_test), np.max(x_test)
+        covariate_colors = _get_sorted_colors(
+            adata,
+            covariate_key,
+            time_key,
+            tmin=tmin,
+            tmax=tmax,
+        )
+        row_delta = len(covariate_colors) + 1
+    gs = GridSpec(nrows=nrows * row_delta, ncols=ncols, figure=fig)
+
+    row, sharey_ax = 0, None
+    for i, c in enumerate(clusters):
+        if i == ncols:
+            row += row_delta
+        sharey_ax = plot_cluster(row, i % ncols, c, sharey_ax=sharey_ax)
 
     if save is not None:
         save_fig(fig, save)
diff --git a/cellrank/pl/_gene_trend.py b/cellrank/pl/_gene_trend.py
index 43db78236..7fe8d846b 100644
--- a/cellrank/pl/_gene_trend.py
+++ b/cellrank/pl/_gene_trend.py
@@ -3,13 +3,6 @@
 from typing import List, Tuple, Union, Mapping, TypeVar, Optional, Sequence
 from pathlib import Path
 
-import numpy as np
-import pandas as pd
-
-import matplotlib
-import matplotlib.pyplot as plt
-from matplotlib import cm
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.pl._utils import (
@@ -27,6 +20,13 @@
 from cellrank.ul._utils import _get_n_cores, _check_collection
 from cellrank.tl._constants import _DEFAULT_BACKEND, AbsProbKey
 
+import numpy as np
+import pandas as pd
+
+import matplotlib
+import matplotlib.pyplot as plt
+from matplotlib import cm
+
 AnnData = TypeVar("AnnData")
 
 
@@ -91,13 +91,15 @@ def gene_trends(
         Key in ``adata.layers`` or `'X'` for ``adata.X`` where the data is stored.
     time_key
         Key in ``adata.obs`` where the pseudotime is stored.
-    %(time_ranges)s
+    %(time_range)s
+
+        This can also be specified on per-lineage basis.
     transpose
         If ``same_plot=True``, group the trends by ``lineages`` instead of ``genes``. This enforces ``hide_cells=True``.
         If ``same_plot=False``, show ``lineages`` in rows and ``genes`` in columns.
     %(model_callback)s
     conf_int
-        Whether to compute and show confidence interval. If the :paramref:`model` is :class:`cellrank.ul.models.GAMR`,
+        Whether to compute and show confidence interval. If the ``model`` is :class:`cellrank.ul.models.GAMR`,
         it can also specify the confidence level, the default is `0.95`.
     same_plot
         Whether to plot all lineages for each gene in the same plot.
diff --git a/cellrank/pl/_graph.py b/cellrank/pl/_graph.py
index 7dbf66e4d..0bb3cc14a 100644
--- a/cellrank/pl/_graph.py
+++ b/cellrank/pl/_graph.py
@@ -5,6 +5,12 @@
 from typing import Dict, Tuple, Union, TypeVar, Callable, Optional, Sequence
 from pathlib import Path
 
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import save_fig
+from cellrank.ul._utils import _read_graph_data
+from cellrank.tl._constants import _colors
+
 import numpy as np
 import pandas as pd
 from scipy.sparse import issparse, spmatrix
@@ -18,12 +24,6 @@
 from matplotlib.collections import LineCollection
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-from cellrank import logging as logg
-from cellrank.ul._docs import d
-from cellrank.tl._utils import save_fig
-from cellrank.ul._utils import _read_graph_data
-from cellrank.tl._constants import _colors
-
 AnnData = TypeVar("AnnData")
 
 
diff --git a/cellrank/pl/_heatmap.py b/cellrank/pl/_heatmap.py
index 09895c472..e27e8c8d6 100644
--- a/cellrank/pl/_heatmap.py
+++ b/cellrank/pl/_heatmap.py
@@ -6,18 +6,7 @@
 from pathlib import Path
 from collections import Iterable, defaultdict
 
-import numpy as np
-import pandas as pd
-from pandas.api.types import is_categorical_dtype
-from scipy.ndimage.filters import convolve
-
-import matplotlib as mpl
-import matplotlib.colors as mcolors
-import matplotlib.pyplot as plt
-from matplotlib import cm
-from matplotlib.ticker import FormatStrFormatter
-from mpl_toolkits.axes_grid1.inset_locator import inset_axes
-
+from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.pl._utils import (
@@ -29,15 +18,25 @@
     _create_callbacks,
     _input_model_type,
     _return_model_type,
+    _get_categorical_colors,
 )
 from cellrank.tl._utils import save_fig, _min_max_scale, _unique_order_preserving
 from cellrank.ul._utils import _get_n_cores, valuedispatch, _check_collection
-from cellrank.tl._colors import _create_categorical_colors
 from cellrank.tl._constants import _DEFAULT_BACKEND, ModeEnum, AbsProbKey
 
+import numpy as np
+import pandas as pd
+from scipy.ndimage.filters import convolve
+
+import matplotlib as mpl
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from matplotlib.ticker import FormatStrFormatter
+from mpl_toolkits.axes_grid1.inset_locator import inset_axes
+
 _N_XTICKS = 10
 
-AnnData = TypeVar("AnnData")
 Cmap = TypeVar("Cmap")
 Norm = TypeVar("Norm")
 Ax = TypeVar("Ax")
@@ -104,10 +103,12 @@ def heatmap(
             - `{m.GENES.s!r}` - group by ``lineages`` for each gene in ``genes``.
     time_key
         Key in ``adata.obs`` where the pseudotime is stored.
-    %(time_ranges)s
+    %(time_range)s
+
+        This can also be specified on per-lineage basis.
     %(model_callback)s
     cluster_key
-        Key(s) in ``adata.obs`` containing categorical observations to be plotted on top of the heatmap.
+        Key(s) in ``adata.obs`` containing categorical observations to be plotted at the top of the heatmap.
         Only available when ``mode={m.LINEAGES.s!r}``.
     show_absorption_probabilities
         Whether to also plot absorption probabilities alongside the smoothed expression.
@@ -118,7 +119,7 @@ def heatmap(
         Whether to keep the gene order for later lineages after the first was sorted.
         Only available when ``cluster_genes=False`` and ``mode={m.LINEAGES.s!r}``.
     scale
-        Whether to normalize the gene expression `0-1` range.
+        Whether to normalize the gene expression `[0, 1]` range.
     n_convolve
         Size of the convolution window when smoothing absorption probabilities.
     show_all_genes
@@ -199,28 +200,8 @@ def create_col_colors(lname: str, rng: np.ndarray) -> Tuple[np.ndarray, Cmap, No
         )
 
     def create_col_categorical_color(cluster_key: str, rng: np.ndarray) -> np.ndarray:
-        if not is_categorical_dtype(adata.obs[cluster_key]):
-            raise TypeError(
-                f"Expected `adata.obs[{cluster_key!r}]` to be categorical, "
-                f"found `{adata.obs[cluster_key].dtype.name!r}`."
-            )
-
-        color_key = f"{cluster_key}_colors"
-        if color_key not in adata.uns:
-            logg.warning(
-                f"Color key `{color_key!r}` not found in `adata.uns`. Creating new colors"
-            )
-            colors = _create_categorical_colors(
-                len(adata.obs[cluster_key].cat.categories)
-            )
-            adata.uns[color_key] = colors
-        else:
-            colors = adata.uns[color_key]
-
-        time_series = adata.obs[time_key]
-        ixs = find_indices(time_series, rng)
-
-        mapper = dict(zip(adata.obs[cluster_key].cat.categories, colors))
+        colors, mapper = _get_categorical_colors(adata, cluster_key)
+        ixs = find_indices(adata.obs[time_key], rng)
 
         return np.array(
             [mcolors.to_hex(mapper[v]) for v in adata[ixs, :].obs[cluster_key].values]
diff --git a/cellrank/pl/_lineages.py b/cellrank/pl/_lineages.py
index bc51855a3..e8815376c 100644
--- a/cellrank/pl/_lineages.py
+++ b/cellrank/pl/_lineages.py
@@ -1,8 +1,6 @@
 """Module for plotting lineage-related stuff."""
 from typing import Union, Optional, Sequence
 
-import pandas as pd
-
 import cellrank.logging as logg
 from cellrank.ul._docs import d
 from cellrank.pl._utils import AnnData
@@ -11,6 +9,8 @@
 from cellrank.tl.estimators._constants import A, P
 from cellrank.tl.kernels._precomputed_kernel import DummyKernel
 
+import pandas as pd
+
 
 @d.dedent
 def lineages(
diff --git a/cellrank/pl/_log_odds.py b/cellrank/pl/_log_odds.py
new file mode 100644
index 000000000..1abc9f7fd
--- /dev/null
+++ b/cellrank/pl/_log_odds.py
@@ -0,0 +1,318 @@
+from copy import copy
+from typing import Any, Tuple, Union, Iterable, Optional, Sequence
+from pathlib import Path
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.pl._utils import _position_legend, _get_categorical_colors
+from cellrank.tl._utils import save_fig, _unique_order_preserving
+from cellrank.tl._constants import AbsProbKey
+
+import numpy as np
+import pandas as pd
+
+import seaborn as sns
+import matplotlib.pyplot as plt
+from matplotlib.cm import ScalarMappable
+from matplotlib.axes import Axes
+from matplotlib.colors import Normalize, to_hex
+
+
+@d.dedent
+def log_odds(
+    adata: AnnData,
+    lineage_1: str,
+    lineage_2: Optional[str] = None,
+    time_key: str = "exp_time",
+    backward: bool = False,
+    keys: Optional[Union[str, Sequence[str]]] = None,
+    threshold: Optional[Union[float, Sequence]] = None,
+    threshold_color: str = "red",
+    layer: Optional[str] = None,
+    use_raw: bool = False,
+    size: float = 2.0,
+    cmap: str = "viridis",
+    alpha: Optional[float] = 0.8,
+    ncols: Optional[int] = None,
+    fontsize: Optional[Union[float, str]] = None,
+    xticks_step_size: Optional[int] = 1,
+    legend_loc: Optional[str] = "best",
+    jitter: Union[bool, float] = True,
+    seed: Optional[int] = None,
+    figsize: Optional[Tuple[float, float]] = None,
+    dpi: Optional[int] = None,
+    save: Optional[Union[str, Path]] = None,
+    show: bool = True,
+    **kwargs: Any,
+) -> Optional[Union[Axes, Sequence[Axes]]]:
+    """
+    Plot log-odds ratio between lineages.
+
+    Log-odds are plotted as a function of the experimental time.
+
+    Parameters
+    ----------
+    %(adata)s
+    lineage_1
+        The first lineage for which to compute the log-odds.
+    lineage_2
+        The second lineage for which to compute the log-odds. If `None`, use the rest of the lineages.
+    time_key
+        Key in :attr:`anndata.AnnData.obs` containing the experimental time.
+    %(backward)s
+    keys
+        Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names`.
+    threshold
+        Visualize whether total expression per cell is greater than ``threshold``.
+        If a :class:`typing.Sequence`, it should be the same length as ``keys``.
+    threshold_color
+        Color to use when plotting thresholded expression values.
+    layer
+        Which layer to use to get expression values. If `None` or `'X'`, use :attr:`anndata.AnnData.X`.
+    use_raw
+        Whether to access :attr:`anndata.AnnData.raw`. If `True`, ``layer`` is ignored.
+    size
+        Size of the dots.
+    cmap
+        Colormap to use for continuous variables in ``keys``.
+    alpha
+        Alpha values for the dots.
+    ncols
+        Number of columns.
+    fontsize
+        Size of the font for the title, x- and y-label.
+    xticks_step_size
+        Show only every n-th ticks on x-axis. If `None`, don't show any ticks.
+    legend_loc
+        Position of the legend. If `None`, do not show the legend.
+    jitter
+        Amount of jitter to apply along x-axis.
+    seed
+        Seed for ``jitter`` to ensure reproducibility.
+    %(plotting)s
+    show
+        If `False`, return :class:`matplotlib.pyplot.Axes` or a sequence of them.
+    kwargs
+        Keyword arguments for :func:`seaborn.stripplot`.
+
+    Returns
+    -------
+    :class:`matplotlib.pyplot.Axes`
+        The axis object(s) if ``show=False``.
+    %(just_plots)s
+    """
+    from cellrank.tl.kernels._utils import _ensure_numeric_ordered
+
+    def decorate(
+        ax: Axes, *, title: Optional[str] = None, show_ylabel: bool = True
+    ) -> None:
+        ax.set_xlabel(time_key, fontsize=fontsize)
+        ax.set_title(title, fontdict={"fontsize": fontsize})
+        ax.set_ylabel(ylabel if show_ylabel else "", fontsize=fontsize)
+
+        if xticks_step_size is None:
+            ax.set_xticks([])
+        else:
+            step = max(1, xticks_step_size)
+            ax.set_xticks(np.arange(0, n_cats, step))
+            ax.set_xticklabels(df[time_key].cat.categories[::step])
+
+    def cont_palette(values: np.ndarray) -> Tuple[np.ndarray, ScalarMappable]:
+        cm = copy(plt.get_cmap(cmap))
+        cm.set_bad("grey")
+        sm = ScalarMappable(
+            cmap=cm, norm=Normalize(vmin=np.nanmin(values), vmax=np.nanmax(values))
+        )
+        return np.array([to_hex(v) for v in (sm.to_rgba(values))]), sm
+
+    def get_data(
+        key: str,
+        thresh: Optional[float] = None,
+    ) -> Tuple[
+        Optional[str], Optional[np.ndarray], Optional[np.ndarray], ScalarMappable
+    ]:
+        try:
+            _, palette = _get_categorical_colors(adata, key)
+            df[key] = adata.obs[key].values[mask]
+            df[key] = df[key].cat.remove_unused_categories()
+            try:
+                # seaborn doesn't like numeric categories
+                df[key] = df[key].astype(float)
+                palette = {float(k): v for k, v in palette.items()}
+            except ValueError:
+                pass
+            # otherwise seaborn plots all
+            palette = {k: palette[k] for k in df[key].unique()}
+            hue, thresh_mask, sm = key, None, None
+        except TypeError:
+            palette, hue, thresh_mask, sm = (
+                cont_palette(adata.obs[key].values[mask])[0],
+                None,
+                None,
+                None,
+            )
+        except KeyError:
+            try:
+                # fmt: off
+                if thresh is None:
+                    values = adata.raw.obs_vector(key) if use_raw else adata.obs_vector(key, layer=layer)
+                    palette, sm = cont_palette(values)
+                    hue, thresh_mask = None, None
+                else:
+                    if use_raw:
+                        values = np.asarray(adata.raw[:, key].X[mask].sum(1)).squeeze()
+                    elif layer not in (None, "X"):
+                        values = np.asarray(adata[:, key].layers[layer][mask].sum(1)).squeeze()
+                    else:
+                        values = np.asarray(adata[:, key].X[mask].sum(1)).squeeze()
+                    thresh_mask = values > thresh
+                    hue, palette, sm = None, None, None
+                # fmt: on
+            except KeyError as e:
+                raise e from None
+
+        return hue, palette, thresh_mask, sm
+
+    np.random.seed(seed)
+    _ = kwargs.pop("orient", None)
+    if use_raw and adata.raw is None:
+        logg.warning("No raw attribute set. Setting `use_raw=False`")
+        use_raw = False
+
+    # define log-odds
+    ln_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
+    if ln_key not in adata.obsm:
+        raise KeyError(f"Lineages key `{ln_key!r}` not found in `adata.obsm`.")
+    time = _ensure_numeric_ordered(adata, time_key)
+    order = time.cat.categories[:: -1 if backward else 1]
+
+    fate1 = adata.obsm[ln_key][lineage_1].X.squeeze(-1)
+    if lineage_2 is None:
+        fate2 = 1 - fate1
+        ylabel = rf"$\log{{\frac{{{lineage_1}}}{{rest}}}}$"
+    else:
+        fate2 = adata.obsm[ln_key][lineage_2].X.squeeze(-1)
+        ylabel = rf"$\log{{\frac{{{lineage_1}}}{{{lineage_2}}}}}$"
+
+    # fmt: off
+    df = pd.DataFrame(
+        {
+            "log_odds": np.log(np.divide(fate1, fate2, where=fate2 != 0, out=np.zeros_like(fate1) + 1e-12)),
+            time_key: time,
+        }
+    )
+    mask = (fate1 != 0) & (fate2 != 0)
+    df = df[mask]
+    n_cats = len(df[time_key].cat.categories)
+    # fmt: on
+
+    if keys is None:
+        if figsize is None:
+            figsize = np.array([n_cats, n_cats * 4 / 6]) / 2
+
+        fig, ax = plt.subplots(figsize=figsize, dpi=dpi, tight_layout=True)
+        ax = sns.stripplot(
+            time_key,
+            "log_odds",
+            data=df,
+            order=order,
+            jitter=jitter,
+            color="k",
+            size=size,
+            ax=ax,
+            **kwargs,
+        )
+        decorate(ax)
+        if save is not None:
+            save_fig(fig, save)
+        return None if show else ax
+
+    if isinstance(keys, str):
+        keys = (keys,)
+    if not len(keys):
+        raise ValueError("No keys have been selected.")
+    keys = _unique_order_preserving(keys)
+
+    if not isinstance(threshold, Iterable):
+        threshold = (threshold,) * len(keys)
+    if len(threshold) != len(keys):
+        raise ValueError(
+            f"Expected `threshold` to be of length `{len(keys)}`, found `{len(threshold)}`."
+        )
+
+    ncols = max(len(keys) if ncols is None else ncols, 1)
+    nrows = int(np.ceil(len(keys) / ncols))
+    if figsize is None:
+        figsize = np.array([n_cats * ncols, n_cats * nrows * 4 / 6]) / 2
+
+    fig, axes = plt.subplots(
+        nrows=nrows,
+        ncols=ncols,
+        figsize=figsize,
+        dpi=dpi,
+        constrained_layout=True,
+        sharey="all",
+    )
+    axes = np.ravel([axes])
+
+    i = 0
+    for i, (key, ax, thresh) in enumerate(zip(keys, axes, threshold)):
+        hue, palette, thresh_mask, sm = get_data(key, thresh)
+        show_ylabel = i % ncols == 0
+
+        ax = sns.stripplot(
+            time_key,
+            "log_odds",
+            data=df if thresh_mask is None else df[~thresh_mask],
+            hue=hue,
+            order=order,
+            jitter=jitter,
+            color="black",
+            palette=palette,
+            size=size,
+            alpha=alpha if alpha is not None else None if thresh_mask is None else 0.8,
+            ax=ax,
+            **kwargs,
+        )
+        if thresh_mask is not None:
+            sns.stripplot(
+                time_key,
+                "log_odds",
+                data=df if thresh_mask is None else df[thresh_mask],
+                hue=hue,
+                order=order,
+                jitter=jitter,
+                color=threshold_color,
+                palette=palette,
+                size=size * 2,
+                alpha=0.9,
+                ax=ax,
+                **kwargs,
+            )
+            key = rf"${key} > {thresh}$"
+        if sm is not None:
+            cax = ax.inset_axes([1.02, 0, 0.025, 1], transform=ax.transAxes)
+            fig.colorbar(sm, ax=ax, cax=cax)
+        else:
+            if legend_loc in (None, "none"):
+                legend = ax.get_legend()
+                if legend is not None:
+                    legend.remove()
+            else:
+                handles, labels = ax.get_legend_handles_labels()
+                if len(handles):
+                    _position_legend(
+                        ax, legend_loc=legend_loc, handles=handles, labels=labels
+                    )
+
+        decorate(ax, title=key, show_ylabel=show_ylabel)
+
+    for ax in axes[i + 1 :]:
+        ax.remove()
+    axes = axes[: i + 1]
+
+    if save is not None:
+        save_fig(fig, save)
+
+    return None if show else axes[0] if len(axes) == 1 else axes
diff --git a/cellrank/pl/_utils.py b/cellrank/pl/_utils.py
index 38b078c44..1d8318e3b 100644
--- a/cellrank/pl/_utils.py
+++ b/cellrank/pl/_utils.py
@@ -4,6 +4,7 @@
 from typing import (
     Any,
     Dict,
+    List,
     Tuple,
     Union,
     Mapping,
@@ -16,16 +17,6 @@
 from itertools import combinations
 from collections import namedtuple, defaultdict
 
-import numpy as np
-import pandas as pd
-from pandas.api.types import is_categorical_dtype
-
-import matplotlib as mpl
-import matplotlib.colors as mcolors
-import matplotlib.pyplot as plt
-from matplotlib import cm
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.tl._utils import save_fig, _unique_order_preserving
@@ -35,6 +26,16 @@
 from cellrank.ul._parallelize import parallelize
 from cellrank.ul.models._base_model import ColorType
 
+import numpy as np
+import pandas as pd
+from pandas.api.types import infer_dtype, is_numeric_dtype, is_categorical_dtype
+
+import matplotlib as mpl
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+from matplotlib import cm
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
 AnnData = TypeVar("AnnData")
 Queue = TypeVar("Queue")
 Graph = TypeVar("Graph")
@@ -679,12 +680,12 @@ def _trends_helper(
 
     if show_prob and same_plot:
         minns, maxxs = zip(
-            *[
+            *(
                 models[gene][n]._return_min_max(
                     show_conf_int=kwargs.get("conf_int", False),
                 )
                 for n in lineage_names
-            ]
+            )
         )
         minn, maxx = min(minns), max(maxxs)
         kwargs["loc"] = legend_loc
@@ -1216,3 +1217,70 @@ def _held_karp(dists: np.ndarray) -> Tuple[float, np.ndarray]:
     path.append(0)
 
     return opt, np.array(path)[::-1]
+
+
+def _get_categorical_colors(
+    adata: AnnData, cluster_key: str
+) -> Tuple[np.ndarray, Mapping[str, str]]:
+    if cluster_key not in adata.obs:
+        raise KeyError(f"Unable to find data in `adata.obs[{cluster_key!r}].`")
+    if not is_categorical_dtype(adata.obs[cluster_key]):
+        raise TypeError(
+            f"Expected `adata.obs[{cluster_key!r}]` to be categorical, "
+            f"found `{infer_dtype(adata.obs[cluster_key])}`."
+        )
+
+    color_key = f"{cluster_key}_colors"
+    try:
+        colors = adata.uns[color_key]
+    except KeyError:
+        adata.uns[color_key] = colors = _create_categorical_colors(
+            len(adata.obs[cluster_key].cat.categories)
+        )
+    mapper = dict(zip(adata.obs[cluster_key].cat.categories, colors))
+    mapper[np.nan] = "grey"
+
+    return colors, mapper
+
+
+def _get_sorted_colors(
+    adata: AnnData,
+    cluster_key: Union[str, Sequence[str]],
+    time_key: Optional[str] = None,
+    tmin: float = -np.inf,
+    tmax: float = np.inf,
+) -> List[np.ndarray]:
+    if time_key is not None:
+        if time_key not in adata.obs:
+            raise KeyError(f"Unable to find time in `adata.obs[{time_key!r}]`.")
+        adata = adata[(adata.obs[time_key] >= tmin) & (adata.obs[time_key] <= tmax)]
+        if not adata.n_obs:
+            raise ValueError(
+                f"Specified time range `{[tmin, tmax]}` does not contain any data."
+            )
+        order = np.argsort(adata.obs[time_key].values)
+    else:
+        order = np.arange(adata.n_obs)
+
+    if isinstance(cluster_key, str):
+        cluster_key = (cluster_key,)
+    cluster_key = _unique_order_preserving(cluster_key)
+
+    res = []
+    for ck in cluster_key:
+        try:
+            colors, mapper = _get_categorical_colors(adata, ck)
+            res.append(
+                np.array(
+                    [mcolors.to_hex(mapper[v]) for v in adata.obs[ck].values[order]]
+                )
+            )
+        except TypeError:
+            if not is_numeric_dtype(adata.obs[ck]):
+                raise TypeError(
+                    f"Expected `adata.obs[{cluster_key!r}]` to be numeric, "
+                    f"found `{infer_dtype(adata.obs[cluster_key])}`."
+                )
+            res.append(np.asarray(adata.obs[ck])[order])
+
+    return res
diff --git a/cellrank/settings/_settings.py b/cellrank/settings/_settings.py
index cd5e5f102..2341a400e 100644
--- a/cellrank/settings/_settings.py
+++ b/cellrank/settings/_settings.py
@@ -2,7 +2,6 @@
 from copy import copy
 
 from scanpy import settings
-
 from cellrank.logging._logging import _RootLogger, _LogFormatter
 
 
diff --git a/cellrank/tl/_colors.py b/cellrank/tl/_colors.py
index ce0d24140..af2dd00e6 100644
--- a/cellrank/tl/_colors.py
+++ b/cellrank/tl/_colors.py
@@ -2,6 +2,8 @@
 
 from typing import Any, List, Tuple, Union, Iterable, Optional, Sequence
 
+from cellrank import logging as logg
+
 import numpy as np
 from pandas import Series, DataFrame, to_numeric
 from scipy.stats import entropy
@@ -11,8 +13,6 @@
 from matplotlib import cm
 from matplotlib import colors as mcolors
 
-from cellrank import logging as logg
-
 
 def _create_colors(
     base_color: Union[str, Tuple[float, float, float]],
@@ -105,7 +105,7 @@ def _create_categorical_colors(n_categories: Optional[int] = None):
         n_categories = max_cats
     if n_categories > max_cats:
         raise ValueError(
-            f"Maximum number of colors ({max_cats}) exceeded: `{n_categories}`."
+            f"Number of categories `{n_categories}` exceeded the maximum number of colors `{max_cats}`."
         )
 
     colors = []
@@ -140,7 +140,7 @@ def _get_black_or_white(value: float, cmap) -> str:
     if not (0.0 <= value <= 1.0):
         raise ValueError(f"Value must be in range `[0, 1]`, found `{value}`.")
 
-    r, g, b, *_ = [int(c * 255) for c in cmap(value)]
+    r, g, b, *_ = (int(c * 255) for c in cmap(value))
     return _contrasting_color(r, g, b)
 
 
diff --git a/cellrank/tl/_init_term_states.py b/cellrank/tl/_init_term_states.py
index 3007ec87a..1db43d44e 100644
--- a/cellrank/tl/_init_term_states.py
+++ b/cellrank/tl/_init_term_states.py
@@ -20,15 +20,15 @@
 
 
 _docstring = """\
-Find {direction} states of a dynamic process of single cells based on RNA velocity [Manno18]_.
+Find {direction} states of a dynamic process of single cells based on RNA velocity :cite:`manno:18`.
 
 The function models dynamic cellular processes as a Markov chain, where the transition matrix is computed based
 on the velocity vector of each individual cell. Based on this Markov chain, we provide two estimators
 to compute {direction} states, both of which are based on spectral methods.
 
 For the estimator :class:`cellrank.tl.estimators.GPCCA`, cells are fuzzily clustered into macrostates,
-using Generalized Perron Cluster Cluster Analysis [GPCCA18]_. In short, this coarse-grains the Markov chain into a set
-of macrostates representing the slow time-scale dynamics, i.e. transitions between these macrostates are rare.
+using Generalized Perron Cluster Cluster Analysis :cite:`reuter:18`. In short, this coarse-grains the Markov chain into
+a set of macrostates representing the slow time-scale dynamics, i.e. transitions between these macrostates are rare.
 The most stable ones of these will represent {direction}, while the others represent intermediate macrostates.
 
 For the estimator :class:`cellrank.tl.estimators.CFLARE`, cells are filtered into transient/recurrent cells using the
@@ -49,6 +49,8 @@
 key
     Key in ``adata.obsp`` where the transition matrix is saved.
     If not found, compute a new one using :func:`cellrank.tl.transition_matrix`.
+force_recompute
+    Whether to always recompute the transition matrix even if one exists.
 show_plots
     Whether to show plots of the spectrum and eigenvectors in the embedding.
 %(n_jobs)s
@@ -80,16 +82,18 @@ def _initial_terminal(
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
+    force_recompute: bool = False,
     show_plots: bool = False,
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
     **kwargs,
 ) -> Optional[Union[AnnData, BaseEstimator]]:
-
     _check_estimator_type(estimator)
 
     try:
+        if force_recompute:
+            raise KeyError("Forcing transition matrix recomputation.")
         kernel = PrecomputedKernel(key, adata=adata, backward=backward)
         write_to_adata = False  # no need to write
         logg.info("Using precomputed transition matrix")
@@ -164,12 +168,13 @@ def initial_states(  # noqa: D103
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
+    force_recompute: bool = False,
     show_plots: bool = False,
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
     **kwargs,
-) -> Optional[AnnData]:
+) -> Optional[Union[AnnData, BaseEstimator]]:
 
     return _initial_terminal(
         adata,
@@ -180,6 +185,7 @@ def initial_states(  # noqa: D103
         n_states=n_states,
         cluster_key=cluster_key,
         key=key,
+        force_recompute=force_recompute,
         show_plots=show_plots,
         copy=copy,
         return_estimator=return_estimator,
@@ -202,12 +208,13 @@ def terminal_states(  # noqa: D103
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
+    force_recompute: bool = False,
     show_plots: bool = False,
     copy: bool = False,
     return_estimator: bool = False,
     fit_kwargs: Mapping = MappingProxyType({}),
     **kwargs,
-) -> Optional[AnnData]:
+) -> Optional[Union[AnnData, BaseEstimator]]:
 
     return _initial_terminal(
         adata,
@@ -217,6 +224,7 @@ def terminal_states(  # noqa: D103
         n_states=n_states,
         cluster_key=cluster_key,
         key=key,
+        force_recompute=force_recompute,
         show_plots=show_plots,
         copy=copy,
         return_estimator=return_estimator,
diff --git a/cellrank/tl/_lineage.py b/cellrank/tl/_lineage.py
index 4a7066581..3e019ff91 100644
--- a/cellrank/tl/_lineage.py
+++ b/cellrank/tl/_lineage.py
@@ -9,13 +9,6 @@
 
 from typing_extensions import Literal
 
-import numpy as np
-import pandas as pd
-from scipy.stats import entropy
-
-import matplotlib.colors as c
-import matplotlib.pyplot as plt
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.tl._utils import save_fig, _convert_lineage_name, _unique_order_preserving
@@ -26,6 +19,13 @@
 )
 from cellrank.tl._constants import Lin, ModeEnum
 
+import numpy as np
+import pandas as pd
+from scipy.stats import entropy
+
+import matplotlib.colors as c
+import matplotlib.pyplot as plt
+
 ColorLike = TypeVar("ColorLike")
 _ERROR_NOT_ITERABLE = "Expected `{}` to be iterable, found type `{}`."
 _ERROR_WRONG_SIZE = "Expected `{}` to be of size `{{}}`, found `{{}}`."
@@ -733,11 +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
+                - `'kl_divergence'`: as in :cite:`velten:17`, 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.
-
+                - `'entropy'`: as in :cite:`setty:19`, 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/_lineages.py b/cellrank/tl/_lineages.py
index cdc40339d..eb22d73be 100644
--- a/cellrank/tl/_lineages.py
+++ b/cellrank/tl/_lineages.py
@@ -2,8 +2,6 @@
 
 from typing import Union, TypeVar, Optional, Sequence
 
-import pandas as pd
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.tl._utils import TestMethod
@@ -13,6 +11,8 @@
 from cellrank.tl.estimators._constants import P
 from cellrank.tl.kernels._precomputed_kernel import DummyKernel
 
+import pandas as pd
+
 AnnData = TypeVar("AnnData")
 
 
diff --git a/cellrank/tl/_linear_solver.py b/cellrank/tl/_linear_solver.py
index 3af6d105d..9e0b5ed47 100644
--- a/cellrank/tl/_linear_solver.py
+++ b/cellrank/tl/_linear_solver.py
@@ -2,6 +2,11 @@
 from typing import List, Tuple, Union, TypeVar, Optional
 from functools import singledispatch
 
+from cellrank import logging as logg
+from cellrank.ul._utils import _get_n_cores
+from cellrank.tl._constants import _DEFAULT_BACKEND
+from cellrank.ul._parallelize import parallelize
+
 import numpy as np
 from scipy.linalg import solve
 from scipy.sparse import eye as speye
@@ -15,11 +20,6 @@
 )
 from scipy.sparse.linalg import gmres, lgmres, gcrotmk, bicgstab
 
-from cellrank import logging as logg
-from cellrank.ul._utils import _get_n_cores
-from cellrank.tl._constants import _DEFAULT_BACKEND
-from cellrank.ul._parallelize import parallelize
-
 _DEFAULT_SOLVER = "gmres"
 _PETSC_ERROR_MSG_SHOWN = False
 _PETSC_ERROR_MSG = (
diff --git a/cellrank/tl/_read.py b/cellrank/tl/_read.py
index 045172a81..60df5805a 100644
--- a/cellrank/tl/_read.py
+++ b/cellrank/tl/_read.py
@@ -4,9 +4,6 @@
 from pathlib import Path
 
 from scvelo import read as scv_read
-
-from matplotlib.colors import is_color_like
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.tl._colors import _create_categorical_colors
@@ -19,6 +16,8 @@
     _lin_names,
 )
 
+from matplotlib.colors import is_color_like
+
 AnnData = TypeVar("AnnData")
 
 
diff --git a/cellrank/tl/_transition_matrix.py b/cellrank/tl/_transition_matrix.py
index 472b7d0a9..ee18dd844 100644
--- a/cellrank/tl/_transition_matrix.py
+++ b/cellrank/tl/_transition_matrix.py
@@ -25,7 +25,7 @@ def transition_matrix(
     backward_mode: str = BackwardMode.TRANSPOSE.s,
     scheme: str = Scheme.CORRELATION.s,
     softmax_scale: Optional[float] = None,
-    weight_connectivities: Optional[float] = 0.2,
+    weight_connectivities: float = 0.2,
     density_normalize: bool = True,
     key: Optional[str] = None,
     **kwargs,
@@ -46,8 +46,7 @@ def transition_matrix(
     xkey
         Key in ``adata.layers`` where expected gene expression counts are stored.
     conn_key
-        Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix, describing cell-cell similarity. Only
-        used when ``weight_connectivities > 0``.
+        Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix, describing cell-cell similarity.
     gene_subset
         List of genes to be used to compute transition probabilities.
         By default, genes from ``adata.var['velocity_genes']`` are used.
@@ -59,7 +58,7 @@ def transition_matrix(
         Weight given to similarities as opposed to velocities. Must be in `[0, 1]`.
     density_normalize
         Whether to use density correction when computing the transition probabilities based on similarities.
-        Density correction is done as by [Haghverdi16]_.
+        Density correction is done as by :cite:`haghverdi:16`.
     %(write_to_adata.parameters)s
     kwargs
         Keyword arguments for :meth:`cellrank.tl.kernels.VelocityKernel.compute_transition_matrix`.
@@ -74,7 +73,12 @@ def transition_matrix(
 
     def compute_velocity_kernel() -> VelocityKernel:
         return VelocityKernel(
-            adata, backward=backward, vkey=vkey, xkey=xkey, gene_subset=gene_subset
+            adata,
+            backward=backward,
+            vkey=vkey,
+            xkey=xkey,
+            gene_subset=gene_subset,
+            conn_key=conn_key,
         ).compute_transition_matrix(
             softmax_scale=softmax_scale,
             mode=mode,
@@ -83,32 +87,27 @@ def compute_velocity_kernel() -> VelocityKernel:
             **kwargs,
         )
 
-    if weight_connectivities is not None:
-        if 0 < weight_connectivities < 1:
-            vk = compute_velocity_kernel()
-            logg.info(
-                f"Using a connectivity kernel with weight `{weight_connectivities}`"
-            )
-            ck = ConnectivityKernel(
-                adata, backward=backward, conn_key=conn_key
-            ).compute_transition_matrix(density_normalize=density_normalize)
-            final = (
-                (1 - weight_connectivities) * vk + weight_connectivities * ck
-            ).compute_transition_matrix()
-        elif weight_connectivities == 0:
-            final = compute_velocity_kernel()
-        elif weight_connectivities == 1:
-            final = ConnectivityKernel(
-                adata,
-                backward=backward,
-                conn_key=conn_key,
-            ).compute_transition_matrix(density_normalize=density_normalize)
-        else:
-            raise ValueError(
-                f"Parameter `weight_connectivities` must be in range `[0, 1]`, found `{weight_connectivities}`."
-            )
+    if 0 < weight_connectivities < 1:
+        vk = compute_velocity_kernel()
+        logg.info(f"Using a connectivity kernel with weight `{weight_connectivities}`")
+        ck = ConnectivityKernel(
+            adata, backward=backward, conn_key=conn_key
+        ).compute_transition_matrix(density_normalize=density_normalize)
+        final = (
+            (1 - weight_connectivities) * vk + weight_connectivities * ck
+        ).compute_transition_matrix()
+    elif weight_connectivities == 0:
+        final = compute_velocity_kernel()
+    elif weight_connectivities == 1:
+        final = ConnectivityKernel(
+            adata,
+            backward=backward,
+            conn_key=conn_key,
+        ).compute_transition_matrix(density_normalize=density_normalize)
     else:
-        final = vk
+        raise ValueError(
+            f"Parameter `weight_connectivities` must be in range `[0, 1]`, found `{weight_connectivities}`."
+        )
 
     final.write_to_adata(key=key)
 
diff --git a/cellrank/tl/_utils.py b/cellrank/tl/_utils.py
index 69219e1c7..0d835d6e6 100644
--- a/cellrank/tl/_utils.py
+++ b/cellrank/tl/_utils.py
@@ -19,20 +19,6 @@
 from statsmodels.stats.multitest import multipletests
 
 from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from pandas import Series
-from scipy.stats import norm
-from numpy.linalg import norm as d_norm
-from scipy.sparse import eye as speye
-from scipy.sparse import diags, issparse, spmatrix, csr_matrix, isspmatrix_csr
-from sklearn.cluster import KMeans
-from pandas.api.types import infer_dtype, is_categorical_dtype
-from scipy.sparse.linalg import norm as sparse_norm
-
-import matplotlib.colors as mcolors
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.ul._utils import _get_neighs, _has_neighs, _get_neighs_params
@@ -46,6 +32,19 @@
 from cellrank.tl._linear_solver import _solve_lin_system
 from cellrank.tl.kernels._utils import np_std, np_mean, _filter_kwargs
 
+import numpy as np
+import pandas as pd
+from pandas import Series
+from scipy.stats import norm
+from numpy.linalg import norm as d_norm
+from scipy.sparse import eye as speye
+from scipy.sparse import diags, issparse, spmatrix, csr_matrix, isspmatrix_csr
+from sklearn.cluster import KMeans
+from pandas.api.types import infer_dtype, is_bool_dtype, is_categorical_dtype
+from scipy.sparse.linalg import norm as sparse_norm
+
+import matplotlib.colors as mcolors
+
 ColorLike = TypeVar("ColorLike")
 GPCCA = TypeVar("GPCCA")
 CFLARE = TypeVar("CFLARE")
@@ -359,7 +358,7 @@ def _correlation_test(
     **kwargs,
 ) -> pd.DataFrame:
     """
-    Perform a stastical test.
+    Perform a statistical test.
 
     Return NaN for genes which don't vary across cells.
 
@@ -384,14 +383,13 @@ def _correlation_test(
 
     Returns
     -------
-    :class:`pandas.DataFrame`
-        Dataframe of shape ``(n_genes, n_lineages * 5)`` containing the following columns, 1 for each lineage:
+    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`` - 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.
+        - ``{lineage} corr`` - correlation between the gene expression and absorption probabilities.
+        - ``{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.
     """
 
     corr, pvals, ci_low, ci_high = _correlation_test_helper(
@@ -403,6 +401,9 @@ def _correlation_test(
         confidence_level=confidence_level,
         **kwargs,
     )
+    invalid = np.sum((corr < -1) | (corr > 1))
+    if invalid:
+        raise ValueError(f"Found `{invalid}` correlations that are not in `[0, 1]`.")
 
     res = pd.DataFrame(corr, index=gene_names, columns=[f"{c} corr" for c in Y.names])
     for idx, c in enumerate(Y.names):
@@ -661,7 +662,7 @@ def _partition(
 
     If *G* describes the state space of a Markov chain, then communication classes are often
     characterized as either recurrent or transient. Intuitively, once the process enters a recurrent class, it will
-    never leave it again. See [Tolver16]_ for more formal definition.
+    never leave it again.
 
     Parameters
     ----------
@@ -758,7 +759,9 @@ def _symmetric(
     return d_norm((matrix - matrix.T), ord=ord) < eps
 
 
-def _normalize(X: Union[np.ndarray, spmatrix]) -> Union[np.ndarray, spmatrix]:
+def _normalize(
+    X: Union[np.ndarray, spmatrix], eps: float = 1e-12
+) -> Union[np.ndarray, spmatrix]:
     """
     Row-normalizes an array to sum to 1.
 
@@ -766,19 +769,21 @@ def _normalize(X: Union[np.ndarray, spmatrix]) -> Union[np.ndarray, spmatrix]:
     ----------
     X
         Array to be normalized.
+    eps
+        To avoid division by zero.
 
     Returns
     -------
-    :class:`numpy.ndarray` or :class:`scipy.sparse.spmatrx`
+    :class:`numpy.ndarray` or :class:`scipy.sparse.spmatrix`
         The normalized array.
     """
 
     with warnings.catch_warnings():
         warnings.simplefilter("ignore")
         if issparse(X):
-            return X.multiply(csr_matrix(1.0 / np.abs(X).sum(1)))
+            return X.multiply(csr_matrix(1.0 / (np.abs(X).sum(1) + eps)))
         X = np.array(X)
-        return X / X.sum(1)[:, None]
+        return X / (X.sum(1)[:, None] + eps)
 
 
 def _get_connectivities(
@@ -904,15 +909,15 @@ def save_fig(
 
 
 def _convert_to_categorical_series(
-    rc_classes: Dict[Union[int, str], Sequence[Union[int, str]]], cell_names: List[str]
+    term_states: Dict[Union[int, str], Sequence[Union[int, str]]], cell_names: List[str]
 ) -> Series:
     """
-    Convert a mapping of recurrent classes to cells to a :class:`pandas.Series`.
+    Convert a mapping of terminal states to cells to a :class:`pandas.Series`.
 
     Parameters
     ----------
-    rc_classes
-        Recurrent classes in the following format: `{'rc_0': ['cell_0', 'cell_1', ...], ...}`.
+    term_states
+        Terminal states in the following format: `{'state_0': ['cell_0', 'cell_1', ...], ...}`.
     cell_names
         List of valid cell names, usually taken from ``adata.obs_names``.
 
@@ -924,28 +929,32 @@ def _convert_to_categorical_series(
 
     cnames = set(cell_names)
     mapper, expected_size = {}, 0
-    for rc, cells in rc_classes.items():
+    for ts, cells in term_states.items():
         if not len(cells):
-            logg.warning(f"No cells selected for category `{rc!r}`.")
+            logg.warning(f"No cells selected for category `{ts!r}`. Skipping")
             continue
         cells = [c if isinstance(c, str) else cell_names[c] for c in cells]
         rest = set(cells) - cnames
         if rest:
             raise ValueError(f"Invalid cell names `{list(rest)}`.")
-        mapper[str(rc)] = cells
+        mapper[str(ts)] = cells
         expected_size += 1
 
     if len(mapper) != expected_size:
         raise ValueError(
-            "All recurrent class labels are being converted to strings, ensure "
+            "All terminal states are being converted to strings, ensure "
             "that there are no conflicting keys, such as `0` and `'0'`."
         )
 
-    rc_labels = Series([np.nan] * len(cell_names), index=cell_names)
-    for rc, cells in mapper.items():
-        rc_labels[cells] = rc
+    term_states = Series([np.nan] * len(cell_names), index=cell_names)
+    for ts, cells in mapper.items():
+        term_states[cells] = ts
 
-    return rc_labels.astype("category")
+    term_states = term_states.astype("category")
+    if not len(term_states.cat.categories):
+        raise ValueError("No categories have been selected.")
+
+    return term_states
 
 
 def _merge_categorical_series(
@@ -1621,3 +1630,25 @@ def _create_initial_terminal_annotations(
     # write to AnnData
     adata.obs[key_added] = cats_merged
     adata.uns[f"{key_added}_colors"] = colors_merged
+
+
+def _maybe_subset_hvgs(
+    adata: AnnData, use_highly_variable: Optional[Union[bool, str]]
+) -> AnnData:
+    if use_highly_variable in (None, False):
+        return adata
+    key = "highly_variable" if use_highly_variable is True else use_highly_variable
+
+    if key not in adata.var.keys():
+        logg.warning(f"Unable to find HVGs in `adata.var[{key!r}]`. Using all genes")
+        return adata
+
+    if not is_bool_dtype(adata.var[key]):
+        logg.warning(
+            f"Expected `adata.var[{key!r}]` to be of bool dtype, "
+            f"found `{infer_dtype(adata.var[key])}`. Using all genes"
+        )
+        return adata
+
+    logg.info(f"Using `{np.sum(adata.var[key])}` HVGs from `adata.var[{key!r}]`")
+    return adata[:, adata.var[key]]
diff --git a/cellrank/tl/estimators/_base_estimator.py b/cellrank/tl/estimators/_base_estimator.py
index 4c9e0c5a6..376797d47 100644
--- a/cellrank/tl/estimators/_base_estimator.py
+++ b/cellrank/tl/estimators/_base_estimator.py
@@ -3,24 +3,24 @@
 from abc import ABC, abstractmethod
 from sys import version_info
 from copy import copy, deepcopy
-from typing import Any, Dict, Tuple, Union, Mapping, TypeVar, Optional, Sequence
+from typing import (
+    Any,
+    Dict,
+    Tuple,
+    Union,
+    Mapping,
+    TypeVar,
+    Iterable,
+    Optional,
+    Sequence,
+)
 from pathlib import Path
 from datetime import datetime
 
 from typing_extensions import Literal
 
+import scanpy as sc
 import scvelo as scv
-
-import numpy as np
-import pandas as pd
-from pandas import Series
-from scipy.stats import ranksums
-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 (
@@ -58,6 +58,19 @@
 from cellrank.tl.kernels._base_kernel import KernelExpression
 from cellrank.tl.estimators._constants import A, P
 
+import numpy as np
+import pandas as pd
+from pandas import Series
+from scipy.stats import ranksums
+from scipy.sparse import spmatrix
+from pandas.api.types import infer_dtype, is_categorical_dtype
+
+import matplotlib.pyplot as plt
+from matplotlib import rc_context, patheffects
+from matplotlib.axes import Axes
+from matplotlib.colors import is_color_like
+from matplotlib.patches import ArrowStyle
+
 AnnData = TypeVar("AnnData")
 _COMP_TERM_STATES_MSG = (
     "Compute terminal states first as `.compute_terminal_states()` or"
@@ -77,18 +90,18 @@ class BaseEstimator(LineageEstimatorMixin, Partitioner, ABC):
         Either a :class:`cellrank.tl.kernels.Kernel` object, an :class:`anndata.AnnData` object which
         stores the transition matrix in ``.obsp`` attribute or :mod:`numpy` or :mod:`scipy` array.
     inplace
-        Whether to modify :paramref:`adata` object inplace or make a copy.
+        Whether to modify :attr:`adata` object inplace or make a copy.
     read_from_adata
-        Whether to read available attributes in :paramref:`adata`, if present.
+        Whether to read available attributes in :attr:`adata`, if present.
     obsp_key
         Key in ``obj.obsp`` when ``obj`` is an :class:`anndata.AnnData` object.
     g2m_key
-        Key in :paramref:`adata` ``.obs``. Can be used to detect cell-cycle driven start- or endpoints.
+        Key in :attr:`adata` ``.obs``. Can be used to detect cell-cycle driven start- or endpoints.
     s_key
-        Key in :paramref:`adata` ``.obs``. Can be used to detect cell-cycle driven start- or endpoints.
+        Key in :attr:`adata` ``.obs``. Can be used to detect cell-cycle driven start- or endpoints.
     write_to_adata
-        Whether to write the transition matrix to :paramref:`adata` ``.obsp`` and the parameters to
-        :paramref:`adata` ``.uns``.
+        Whether to write the transition matrix to :attr:`adata` ``.obsp`` and the parameters to
+        :attr:`adata` ``.uns``.
     %(write_to_adata.parameters)s
         Only used when ``write_to_adata=True``.
     """
@@ -137,6 +150,23 @@ def __init_subclass__(cls, **kwargs: Any):
         super().__init_subclass__()
 
     def _read_from_adata(self) -> None:
+        def _check_term_states():
+            term_states = self._get(A.TERM)
+            colors = self._get(A.TERM_COLORS)
+
+            if term_states is None or colors is None:
+                return
+            if not is_categorical_dtype(term_states):
+                raise TypeError(
+                    f"Expected `adata.obs[{self._term_key!r}]` to be `categorical`, "
+                    f"found `{infer_dtype(term_states)}`."
+                )
+            if len(term_states.cat.categories) != len(colors):
+                raise ValueError(
+                    f"Expected `adata.uns[{_colors(self._term_key)}]` to have "
+                    f"`{len(term_states.cat.categories)}` colors, found `{len(colors)}`."
+                )
+
         self._set_or_debug(f"eig_{self._direction}", self.adata.uns, "_eig")
 
         self._set_or_debug(self._g2m_key, self.adata.obs, "_G2M_score")
@@ -145,18 +175,23 @@ def _read_from_adata(self) -> None:
         self._set_or_debug(self._term_key, self.adata.obs, A.TERM.s)
         self._set_or_debug(_probs(self._term_key), self.adata.obs, A.TERM_PROBS)
         self._set_or_debug(_colors(self._term_key), self.adata.uns, A.TERM_COLORS)
+        _check_term_states()
 
         self._reconstruct_lineage(A.ABS_PROBS, self._abs_prob_key)
         self._set_or_debug(_pd(self._abs_prob_key), self.adata.obs, A.PRIME_DEG)
 
     def _reconstruct_lineage(self, attr: PrettyEnum, obsm_key: str):
         self._set_or_debug(obsm_key, self.adata.obsm, attr)
-        names = self._set_or_debug(_lin_names(self._term_key), self.adata.uns)
-        colors = self._set_or_debug(_colors(self._term_key), self.adata.uns)
-
         probs = self._get(attr)
 
         if probs is not None:
+            names = self._set_or_debug(_lin_names(self._term_key), self.adata.uns)
+            colors = self._set_or_debug(_colors(self._term_key), self.adata.uns)
+            if len(names) != len(colors):
+                raise ValueError(
+                    f"Expected names and colors to be of same length, found `{names}` and `{colors}`."
+                )
+
             if len(names) != probs.shape[1]:
                 if isinstance(probs, Lineage):
                     names = probs.names
@@ -167,18 +202,41 @@ def _reconstruct_lineage(self, attr: PrettyEnum, obsm_key: str):
                     )
                     names = [f"Lineage {i}" for i in range(probs.shape[1])]
             if len(colors) != probs.shape[1] or not all(
-                map(lambda c: isinstance(c, str) and is_color_like(c), colors)
+                is_color_like(c) for c in colors
             ):
                 if isinstance(probs, Lineage):
                     colors = probs.colors
                 else:
                     logg.warning(
-                        f"Expected lineage colors to be of length `{probs.shape[1]}`, found `{len(names)}`. "
+                        f"Expected lineage colors to be of length `{probs.shape[1]}`, found `{len(colors)}`. "
                         f"Creating new colors"
                     )
                     colors = _create_categorical_colors(probs.shape[1])
+
+            colors = _convert_to_hex_colors(colors)
             self._set(attr, Lineage(probs, names=names, colors=colors))
 
+            # ensure that the cont. lineage names match with the categorical values
+            term_states = self._get(A.TERM)
+            if term_states is not None:
+                direction = "initial" if self.kernel.backward else "terminal"
+                if len(names) != len(term_states.cat.categories):
+                    self._set(attr, None)
+                    logg.warning(
+                        f"Expected to find `{len(names)}` {direction} states "
+                        f"(from adata.uns[{_lin_names(self._term_key)!r}]), "
+                        f"found `{len(term_states.cat.categories)}` (from adata.obsm[{obsm_key!r}]). Skipping`.{attr}`"
+                    )
+                    return
+                if tuple(term_states.cat.categories) != tuple(names):
+                    self._set(attr, None)
+                    logg.warning(
+                        f"Expected to find `{names}` {direction} states "
+                        f"(from adata.uns[{_lin_names(self._term_key)!r}]), "
+                        f"found `{term_states.cat.categories}` (from adata.obs[{obsm_key!r}). Skipping `{attr}`"
+                    )
+                    return
+
             self.adata.obsm[obsm_key] = self._get(attr)
             self.adata.uns[_lin_names(self._term_key)] = names
             self.adata.uns[_colors(self._term_key)] = colors
@@ -195,29 +253,38 @@ def set_terminal_states(
         **kwargs,
     ) -> None:
         """
-        Set the approximate recurrent classes, if they are known a priori.
+        Manually define terminal states.
 
         Parameters
         ----------
         labels
-            Either a categorical :class:`pandas.Series` with index as cell names, where `NaN` marks marks a cell
-            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.
+            Defines the terminal states. Valid options are:
+
+                - categorical :class:`pandas.Series` where each category corresponds to one terminal state.
+                  `NaN` entries denote cells that do not belong to any terminal state, i.e. these are either initial or
+                  transient cells.
+                - :class:`dict` where keys are terminal states and values are lists of cell barcodes corresponding to
+                  annotations in :attr:`adata` ``.obs_names``.
+                  If only 1 key is provided, values should correspond to terminal state clusters if a categorical
+                  :class:`pandas.Series` can be found in :attr:`adata` ``.obs``.
+
         cluster_key
-            If a key to cluster labels is given, :paramref:`{fs}` will be associated with these for naming and colors.
+            Key from :attr:`adata.obs` where categorical cluster labels are stored. These are used to associate names
+            and colors with each terminal state. Each terminal state will be given the name and color corresponding to
+            the cluster it mostly overlaps with.
         %(en_cutoff_p_thresh)s
         add_to_existing
-            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.
+            Whether the new terminal states should be added to pre-existing ones. Cells already assigned to a terminal
+            state will be re-assigned to the new terminal state if there's a conflict between old and new annotations.
+            This throws an error if no previous annotations corresponding to terminal states have been found.
 
         Returns
         -------
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
         """
 
         self._set_categorical_labels(
@@ -238,19 +305,19 @@ def rename_terminal_states(
         self, new_names: Mapping[str, str], update_adata: bool = True
     ) -> None:
         """
-        Rename the names of :paramref:`{ts}`.
+        Rename the names of :attr:`{ts}`.
 
         Parameters
         ----------
         new_names
             Mapping where keys are the old names and the values are the new names. New names must be unique.
         update_adata
-            Whether to update underlying :paramref:`adata` object as well or not.
+            Whether to update underlying :attr:`adata` object as well or not.
 
         Returns
         -------
         None
-            Nothing, just updates the names of :paramref:`{ts}`.
+            Nothing, just updates the names of :attr:`{ts}`.
         """
 
         term_states = self._get(P.TERM)
@@ -327,7 +394,7 @@ def compute_absorption_probabilities(
         Compute absorption probabilities of a Markov chain.
 
         For each cell, this computes the probability of it reaching any of the approximate recurrent classes defined
-        by :paramref:`{fs}`.
+        by :attr:`{fs}`.
 
         Parameters
         ----------
@@ -376,8 +443,8 @@ def compute_absorption_probabilities(
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`{abs_prob}` - probabilities of being absorbed into the terminal states.
-                - :paramref:`{lat}` - mean times until absorption to subset absorbing states and optionally
+                - :attr:`{abs_prob}` - probabilities of being absorbed into the terminal states.
+                - :attr:`{lat}` - mean times until absorption to subset absorbing states and optionally
                   their variances saved as ``'{{lineage}} mean'`` and ``'{{lineage}} var'``, respectively,
                   for each subset of absorbing states specified in ``time_to_absorption``.
         """
@@ -449,6 +516,15 @@ def compute_absorption_probabilities(
         q = t[trans_indices, :][:, trans_indices]
         s = t[trans_indices, :][:, rec_indices]
 
+        # take individual solutions and piece them together to get absorption probabilities towards the classes
+        macro_ix_helper = np.cumsum(
+            [0] + [len(indices) for indices in lookup_dict.values()]
+        )
+        s = np.concatenate(
+            [s[:, np.arange(a, b)].sum(axis=1) for a, b in _pairwise(macro_ix_helper)],
+            axis=1,
+        )
+
         # check for irreducibility
         if check_irreducibility:
             if self.is_irreducible is None:
@@ -462,7 +538,7 @@ def compute_absorption_probabilities(
         logg.debug(f"Found `{n_cells}` cells and `{s.shape[1]}` absorbing states")
 
         # solve the linear system of equations
-        mat_x = _solve_lin_system(
+        _abs_classes = _solve_lin_system(
             q,
             s,
             solver=solver,
@@ -495,18 +571,6 @@ def compute_absorption_probabilities(
         else:
             abs_time_means = None
 
-        # take individual solutions and piece them together to get absorption probabilities towards the classes
-        macro_ix_helper = np.cumsum(
-            [0] + [len(indices) for indices in lookup_dict.values()]
-        )
-        _abs_classes = np.concatenate(
-            [
-                mat_x[:, np.arange(a, b)].sum(1)[:, None]
-                for a, b in _pairwise(macro_ix_helper)
-            ],
-            axis=1,
-        )
-
         # for recurrent states, set their self-absorption probability to one
         abs_classes = np.zeros((len(self), n_macrostates))
         rec_classes_full = {
@@ -517,6 +581,21 @@ def compute_absorption_probabilities(
             abs_classes[trans_indices, col] = _abs_classes[:, col]
             abs_classes[cl_indices, col] = 1
 
+        mask = abs_classes >= 0
+        if not np.all(mask):
+            raise ValueError(
+                f"`{np.sum(~mask)}` value(s) are negative. Try decreasing the tolerance as `tol=...`, "
+                f"specifying a preconditioner as `preconditioner=...` or "
+                f"use a direct solver as `solver='direct'` if the matrix is small."
+            )
+        mask = np.isclose(abs_classes.sum(1), 1.0, rtol=1e-3)
+        if not np.all(mask):
+            raise ValueError(
+                f"`{np.sum(~mask)}` value(s) do not sum to 1 (rtol=1e-3). Try decreasing the tolerance as `tol=...`, "
+                f"specifying a preconditioner as `preconditioner=...` or "
+                f"use a direct solver as `solver='direct'` if the matrix is small."
+            )
+
         self._set(
             A.ABS_PROBS,
             Lineage(
@@ -615,22 +694,21 @@ def compute_lineage_drivers(
         Parameters
         ----------
         lineages
-            Either a set of lineage names from :paramref:`absorption_probabilities` `.names` or `None`,
+            Either a set of lineage names from :attr:`absorption_probabilities` `.names` or `None`,
             in which case all lineages are considered.
         method
-            Mode to use when calculating p-values and confidence intervals. Can be one of:
+            Mode to use when calculating p-values and confidence intervals. Valid options are:
 
-                - {tm.FISCHER.s!r} - use Fischer transformation [Fischer21]_.
+                - {tm.FISCHER.s!r} - use Fischer transformation :cite:`fischer:21`.
                 - {tm.PERM_TEST.s!r} - use permutation test.
-
         cluster_key
-            Key from :paramref:`adata` ``.obs`` to obtain cluster annotations. These are considered for ``clusters``.
+            Key from :attr:`adata` ``.obs`` to obtain cluster annotations. These are considered for ``clusters``.
         clusters
             Restrict the correlations to these clusters.
         layer
-            Key from :paramref:`adata` ``.layers``.
+            Key from :attr:`adata` ``.layers``.
         use_raw
-            Whether or not to use :paramref:`adata` ``.raw`` to correlate gene expression.
+            Whether or not to use :attr:`adata` ``.raw`` to correlate gene expression.
             If using a layer other than ``.X``, this must be set to `False`.
         confidence_level
             Confidence level for the confidence interval calculation. Must be in `[0, 1]`.
@@ -646,23 +724,16 @@ def compute_lineage_drivers(
         Returns
         -------
         %(correlation_test.returns)s
-            Only if ``return_drivers=True``.
-
-        None
-            Updates :paramref:`adata` ``.var`` or :paramref:`adata` ``.raw.var``, depending ``use_raw`` with:
+        Only if ``return_drivers=True``.
 
-                - ``'{{direction}} {{lineage}} corr'`` - the potential lineage drivers.
-                - ``'{{direction}} {{lineage}} qval'`` - the corrected p-values.
+        Otherwise, updates :attr:`adata` ``.var`` or :attr:`adata` ``.raw.var``, depending ``use_raw`` with:
 
-            Updates the following fields:
+            - ``'{{direction}} {{lineage}} corr'`` - the potential lineage drivers.
+            - ``'{{direction}} {{lineage}} qval'`` - the corrected p-values.
 
-                - :paramref:`{lin_drivers}` - same as the returned values.
+        Also updates the following fields:
 
-        References
-        ----------
-        .. [Fischer21] Fisher, R. A. (1921),
-            *On the “probable error” of a coefficient of correlation deduced from a small sample.*,
-            `Metron 1 3–32 <http://hdl.handle.net/2440/15169>`__.
+            - :attr:`{lin_drivers}` - same as the returned values.
         """
 
         # check that lineage probs have been computed
@@ -737,9 +808,10 @@ def compute_lineage_drivers(
             f"layer `{layer}` with `use_raw={use_raw}`"
         )
 
+        lin_probs = lin_probs[lineages]
         drivers = _correlation_test(
             data,
-            lin_probs[lineages],
+            lin_probs,
             gene_names=var_names,
             method=method,
             n_perms=n_perms,
@@ -749,8 +821,8 @@ def compute_lineage_drivers(
         )
         self._set(A.LIN_DRIVERS, drivers)
 
-        corrs, qvals = [f"{lin} corr" for lin in lineages], [
-            f"{lin} qval" for lin in lineages
+        corrs, qvals = [f"{lin} corr" for lin in lin_probs.names], [
+            f"{lin} qval" for lin in lin_probs.names
         ]
         if use_raw:
             self.adata.raw.var[[f"{prefix} {col}" for col in corrs]] = drivers[corrs]
@@ -760,7 +832,9 @@ def compute_lineage_drivers(
             self.adata.var[[f"{prefix} {col}" for col in qvals]] = drivers[qvals]
 
         field = "raw.var" if use_raw else "var"
-        keys_added = [f"`adata.{field}['{prefix} {lin} corr']`" for lin in lineages]
+        keys_added = [
+            f"`adata.{field}['{prefix} {lin} corr']`" for lin in lin_probs.names
+        ]
 
         logg.info(
             f"Adding `.{P.LIN_DRIVERS}`\n       "
@@ -798,7 +872,7 @@ def plot_lineage_drivers(
         ncols
             Number of columns.
         use_raw
-            Whether to look in :paramref:`adata` ``.raw.var`` or :paramref:`adata` ``.var``.
+            Whether to look in :attr:`adata` ``.raw.var`` or :attr:`adata` ``.var``.
         title_fmt
             Title format. Possible keywords include `{gene}`, `{qval}`, `{corr}` for gene name,
             q-value and correlation, respectively.
@@ -831,7 +905,6 @@ def prepare_format(
             return kwargs
 
         lin_drivers = self._get(P.LIN_DRIVERS)
-
         if lin_drivers is None:
             raise RuntimeError(
                 f"Compute `.{P.LIN_DRIVERS}` first as `.compute_lineage_drivers()`."
@@ -887,6 +960,192 @@ def prepare_format(
         if save is not None:
             save_fig(fig, save)
 
+    @d.dedent
+    def plot_lineage_drivers_correlation(
+        self,
+        lineage_x: str,
+        lineage_y: str,
+        color: Optional[str] = None,
+        gene_sets: Optional[Dict[str, Iterable]] = None,
+        gene_sets_colors: Optional[Iterable] = None,
+        use_raw: bool = False,
+        cmap: str = "RdYlBu_r",
+        fontsize: int = 12,
+        adjust_text: bool = False,
+        legend_loc: Optional[str] = "best",
+        figsize: Optional[Tuple[float, float]] = (4, 4),
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        show: bool = True,
+        **kwargs: Any,
+    ) -> Optional[Axes]:
+        """
+        Show scatter plot of gene-correlations between two lineages.
+
+        Optionally, you can pass a :class:`dict` of gene names that will be annotated in the plot.
+
+        Parameters
+        ----------
+        lineage_x
+            Name of the lineage on the x-axis.
+        lineage_y
+            Name of the lineage on the y-axis.
+        color
+            Key in :attr:`adata` ``.var``.
+        gene_sets
+            Gene sets annotations of the form `{'gene_set_name': ['gene_1', 'gene_2'], ...}`.
+        gene_sets_colors
+            List of colors where each entry corresponds to a gene set from ``genes_sets``.
+            If `None` and keys in ``gene_sets`` correspond to lineage names, use the lineage colors.
+            Otherwise, use default colors.
+        use_raw
+            Whether to access :attr:`adata` ``.raw.var`` or :attr:`adata` ``.var``.
+        cmap
+            Colormap to use.
+        fontsize
+            Size of the text when plotting ``gene_sets``.
+        adjust_text
+            Whether to automatically adjust text in order to reduce overlap.
+        legend_loc
+            Position of the legend. If `None`, don't show the legend.
+            Only used when ``gene_sets!=None``.
+        %(plotting)s
+        show
+            If `False`, return :class:`matplotlib.pyplot.Axes`.
+        kwargs
+            Keyword arguments for :func:`scanpy.pl.scatter`.
+
+        Returns
+        -------
+        :class:`matplotlib.pyplot.Axes`
+            The axis object if ``show=False``.
+        %(just_plots)s
+
+        Notes
+        -----
+        This plot is based on the following
+        `notebook <https://github.com/theislab/gastrulation_analysis/blob/main/6_cellrank.ipynb>`_ by Maren Büttner.
+        """
+        from cellrank.pl._utils import _position_legend
+
+        if use_raw and self.adata.raw is None:
+            logg.warning("No raw attribute set. Setting `use_raw=False`")
+            use_raw = False
+        adata = self.adata.raw if use_raw else self.adata
+
+        # silent assumption: `.compute_lineage_drivers()` always writes to AnnData
+        key1, key2 = f"to {lineage_x} corr", f"to {lineage_y} corr"
+        if key1 not in adata.var or key2 not in adata.var:
+            haystack = "adata.raw.var" if use_raw else "adata.var"
+            raise RuntimeError(
+                f"Unable to find correlations in `{haystack}[{key1!r}]` or `{haystack}[{key2!r}]`."
+                f"Compute `.{P.LIN_DRIVERS}` first as "
+                f"`.compute_lineage_drivers([{lineage_x!r}, {lineage_y!r}], use_raw={use_raw})`."
+            )
+        # produce the actual scatter plot
+        ctx = {"figure.figsize": figsize, "figure.dpi": dpi}
+        for key in list(ctx.keys()):
+            if ctx[key] is None:
+                del ctx[key]
+        with rc_context(ctx):
+            ax = sc.pl.scatter(
+                adata.to_adata() if hasattr(adata, "to_adata") else adata,
+                x=key1,
+                y=key2,
+                color_map=cmap,
+                use_raw=False,
+                show=False,
+                color=color,
+                **kwargs,
+            )
+        fig = ax.figure
+
+        # add some lines to highlight the origin
+        xmin, xmax = np.nanmin(adata.var[key1]), np.nanmax(adata.var[key1])
+        ymin, ymax = np.nanmin(adata.var[key2]), np.nanmax(adata.var[key2])
+        ax.hlines(0, xmin=xmin, xmax=xmax, color="grey", alpha=0.5, zorder=-1)
+        ax.vlines(0, ymin=ymin, ymax=ymax, color="grey", alpha=0.5, zorder=-1)
+
+        # annotate the passed set of genes
+        if gene_sets is not None:
+            if gene_sets_colors is None:
+                try:
+                    # fmt: off
+                    sets = list(gene_sets.keys())
+                    gene_sets_colors = self.adata.obsm[self._abs_prob_key][sets].colors
+                    # fmt: on
+                except KeyError:
+                    logg.warning(
+                        "Unable to determine gene sets colors from lineages. Using default colors"
+                    )
+                    gene_sets_colors = _create_categorical_colors(len(gene_sets))
+            if len(gene_sets_colors) != len(gene_sets):
+                raise ValueError(
+                    f"Expected `gene_sets_colors` to be of length `{len(gene_sets)}`, "
+                    f"found `{len(gene_sets_colors)}`."
+                )
+
+            path_effect = [
+                patheffects.Stroke(linewidth=2, foreground="w", alpha=0.8),
+                patheffects.Normal(),
+            ]
+            annots = []
+            for (key, values), color in zip(gene_sets.items(), gene_sets_colors):
+                arrowprops = (
+                    {
+                        "arrowstyle": ArrowStyle.CurveFilledB(
+                            head_width=0.1, head_length=0.2
+                        ),
+                        "fc": color,
+                        "ec": color,
+                    }
+                    if adjust_text
+                    else None
+                )
+                values = set(values) & set(adata.var_names)
+                for value in values:
+                    x = adata.var.loc[value, key1]
+                    y = adata.var.loc[value, key2]
+
+                    annot = ax.annotate(
+                        value,
+                        xy=(x, y),
+                        va="top",
+                        ha="left",
+                        path_effects=path_effect,
+                        arrowprops=arrowprops,
+                        size=fontsize,
+                        c=color,
+                    )
+                    annots.append(annot)
+                if values:
+                    ax.scatter([], [], color=color, label=key)
+
+            if adjust_text:
+                try:
+                    import adjustText
+
+                    start = logg.info("Adjusting text position")
+                    adjustText.adjust_text(
+                        annots,
+                        x=adata.var[key1].values,
+                        y=adata.var[key2].values,
+                        ax=ax,
+                    )
+                    logg.info("    Finish", time=start)
+                except ImportError:
+                    logg.error(
+                        "Please install `adjustText` first as `pip install adjustText`"
+                    )
+            if len(annots) and legend_loc not in (None, "none"):
+                _position_legend(ax, legend_loc=legend_loc)
+
+        if save is not None:
+            save_fig(fig, path=save)
+
+        if not show:
+            return ax
+
     def _detect_cc_stages(self, rc_labels: Series, p_thresh: float = 1e-15) -> None:
         """
         Detect cell-cycle driven start or endpoints.
@@ -935,6 +1194,22 @@ def _set_categorical_labels(
         add_to_existing: bool = False,
     ) -> None:
         if isinstance(categories, dict):
+            if len(categories) == 1 and is_categorical_dtype(
+                self.adata.obs.get(next(iter(categories.keys())), None)
+            ):
+                key = next(iter(categories.keys()))
+                if isinstance(categories[key], str) or not isinstance(
+                    categories[key], Iterable
+                ):
+                    vals = (categories[key],)
+                else:
+                    vals = categories[key]
+
+                clusters = self.adata.obs[key]
+                categories = {
+                    cat: self.adata[clusters == cat].obs_names for cat in vals
+                }
+
             categories = _convert_to_categorical_series(
                 categories, list(self.adata.obs_names)
             )
@@ -1004,19 +1279,9 @@ def _write_terminal_states(self, time=None) -> None:
             self._get(P.TERM).cat.categories
         )
 
-        extra_msg = ""
-        if getattr(self, A.TERM_ABS_PROBS.s, None) is not None and hasattr(
-            self, "_term_abs_prob_key"
-        ):
-            # checking for None because terminal states can be set using `set_terminal_states`
-            # without the probabilities in GPCCA
-            self.adata.obsm[self._term_abs_prob_key] = self._get(A.TERM_ABS_PROBS)
-            extra_msg = f"       `adata.obsm[{self._term_abs_prob_key!r}]`\n"
-
         logg.info(
             f"Adding `adata.obs[{_probs(self._term_key)!r}]`\n"
             f"       `adata.obs[{self._term_key!r}]`\n"
-            f"{extra_msg}"
             f"       `.{P.TERM_PROBS}`\n"
             f"       `.{P.TERM}`\n"
             "    Finish",
@@ -1072,10 +1337,10 @@ def fit(
         None
             Nothing, just makes available the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
-                - :paramref:`{ap}`
-                - :paramref:`{pd}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
+                - :attr:`{ap}`
+                - :attr:`{pd}`
         """
         self._fit_terminal_states(**kwargs)
         if compute_absorption_probabilities:
@@ -1118,7 +1383,7 @@ def _set_or_debug(
             logg.debug(f"Unable to set attribute `.{attr}`, skipping")
 
     def copy(self) -> "BaseEstimator":
-        """Return a copy of self, including the underlying :paramref:`adata` object."""
+        """Return a copy of self, including the underlying :attr:`adata` object."""
         k = deepcopy(self.kernel)  # ensure we copy the adata object
         res = type(self)(k, read_from_adata=False)
         for k, v in self.__dict__.items():
diff --git a/cellrank/tl/estimators/_cflare.py b/cellrank/tl/estimators/_cflare.py
index 7f0fccd77..19833334a 100644
--- a/cellrank/tl/estimators/_cflare.py
+++ b/cellrank/tl/estimators/_cflare.py
@@ -1,10 +1,6 @@
 """Clustering and Filtering of Left and Right Eigenvectors based on Markov chains."""
 from typing import List, Tuple, Union, Optional, Sequence
 
-import numpy as np
-from pandas import Series
-from scipy.stats import zscore
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import (
@@ -17,6 +13,10 @@
 from cellrank.tl.estimators._decomposition import Eigen
 from cellrank.tl.estimators._base_estimator import BaseEstimator
 
+import numpy as np
+from pandas import Series
+from scipy.stats import zscore
+
 
 @d.dedent
 class CFLARE(BaseEstimator, Eigen):
@@ -67,7 +67,7 @@ def compute_terminal_states(
         method
             Method to be used for clustering. Must be one of `'louvain'`, `'leiden'` or `'kmeans'`.
         cluster_key
-            If a key to cluster labels is given, :paramref:`{fs}` will get associated with these for naming and colors.
+            If a key to cluster labels is given, :attr:`{fs}` will get associated with these for naming and colors.
         n_clusters_kmeans
             If `None`, this is set to ``use + 1``.
         n_neighbors
@@ -94,8 +94,8 @@ def compute_terminal_states(
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
         """
 
         def _compute_macrostates_prob() -> Series:
@@ -301,10 +301,10 @@ def fit(
         None
             Nothing, just makes available the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
-                - :paramref:`{ap}`
-                - :paramref:`{pd}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
+                - :attr:`{ap}`
+                - :attr:`{pd}`
         """
         super().fit(
             n_lineages=n_lineages,
diff --git a/cellrank/tl/estimators/_decomposition.py b/cellrank/tl/estimators/_decomposition.py
index 3a90a2ede..b46c039d8 100644
--- a/cellrank/tl/estimators/_decomposition.py
+++ b/cellrank/tl/estimators/_decomposition.py
@@ -6,6 +6,13 @@
 from pygpcca import GPCCA as _GPCCA
 from pygpcca._sorted_schur import _check_conj_split
 
+from cellrank import logging as logg
+from cellrank.ul._docs import d, inject_docs
+from cellrank.tl._utils import save_fig, _eigengap
+from cellrank.tl.estimators._utils import Metadata, _delegate
+from cellrank.tl.estimators._property import Property, KernelHolder, VectorPlottable
+from cellrank.tl.estimators._constants import A, F, P
+
 import numpy as np
 from scipy.sparse import issparse
 from scipy.sparse.linalg import eigs
@@ -14,13 +21,6 @@
 from matplotlib.ticker import MultipleLocator, FormatStrFormatter
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-from cellrank import logging as logg
-from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import save_fig, _eigengap
-from cellrank.tl.estimators._utils import Metadata, _delegate
-from cellrank.tl.estimators._property import Property, KernelHolder, VectorPlottable
-from cellrank.tl.estimators._constants import A, F, P
-
 EPS = np.finfo(np.float64).eps
 
 
@@ -80,7 +80,7 @@ def compute_eigendecomposition(
         None
             Nothing, but updates the following field:
 
-                - :paramref:`{prop}`
+                - :attr:`{prop}`
         """
 
         def get_top_k_evals():
@@ -431,9 +431,9 @@ def compute_schur(
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`{schur_vectors}`
-                - :paramref:`{schur_matrix}`
-                - :paramref:`{eigendec}`
+                - :attr:`{schur_vectors}`
+                - :attr:`{schur_matrix}`
+                - :attr:`{eigendec}`
         """
         if n_components < 2:
             raise ValueError(
diff --git a/cellrank/tl/estimators/_gpcca.py b/cellrank/tl/estimators/_gpcca.py
index db299db93..332b8cbeb 100644
--- a/cellrank/tl/estimators/_gpcca.py
+++ b/cellrank/tl/estimators/_gpcca.py
@@ -1,16 +1,9 @@
-"""Generalized Perron Cluster Cluster Analysis [GPCCA18]_."""
+"""Generalized Perron Cluster Cluster Analysis :cite:`reuter:18`."""
 
 from types import MappingProxyType
 from typing import Any, Dict, List, Tuple, Union, Mapping, Iterable, Optional, Sequence
 from pathlib import Path
 
-import numpy as np
-import pandas as pd
-
-import matplotlib as mpl
-import matplotlib.colors as mcolors
-import matplotlib.pyplot as plt
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import (
@@ -19,7 +12,7 @@
     _fuzzy_to_discrete,
     _series_from_one_hot_matrix,
 )
-from cellrank.tl._colors import _get_black_or_white
+from cellrank.tl._colors import _get_black_or_white, _create_categorical_colors
 from cellrank.tl._lineage import Lineage
 from cellrank.tl._constants import TermStatesKey, _probs, _colors, _lin_names
 from cellrank.tl.estimators._utils import Metadata, _print_insufficient_number_of_cells
@@ -28,11 +21,21 @@
 from cellrank.tl.estimators._decomposition import Eigen, Schur
 from cellrank.tl.estimators._base_estimator import BaseEstimator
 
+import numpy as np
+import pandas as pd
+from pandas.api.types import infer_dtype
+from pandas.core.dtypes.common import is_categorical_dtype
+
+import matplotlib as mpl
+import matplotlib.colors as mcolors
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+
 
 @d.dedent
 class GPCCA(BaseEstimator, Macrostates, Schur, Eigen):
     """
-    Generalized Perron Cluster Cluster Analysis [GPCCA18]_ as implemented in `pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_.
+    Generalized Perron Cluster Cluster Analysis :cite:`reuter:18` as implemented in `pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_.
 
     Coarse-grains a discrete Markov chain into a set of macrostates and computes coarse-grained transition probabilities
     among the macrostates. Each macrostate corresponds to an area of the state space, i.e. to a subset of cells. The
@@ -41,7 +44,7 @@ class GPCCA(BaseEstimator, Macrostates, Schur, Eigen):
     minimal overlap between macrostates in a certain inner-product sense. Once the macrostates have been computed,
     we project the large transition matrix onto a coarse-grained transition matrix among the macrostates via
     a Galerkin projection. This projection is based on invariant subspaces of the original transition matrix which
-    are obtained using the real Schur decomposition [GPCCA18]_.
+    are obtained using the real Schur decomposition :cite:`reuter:18`.
 
     Parameters
     ----------
@@ -109,11 +112,11 @@ def compute_macrostates(
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`{msp}`
-                - :paramref:`{ms}`
-                - :paramref:`{schur}`
-                - :paramref:`{coarse_T}`
-                - :paramref:`{coarse_stat}`
+                - :attr:`{msp}`
+                - :attr:`{ms}`
+                - :attr:`{schur}`
+                - :attr:`{coarse_T}`
+                - :attr:`{coarse_stat}`
         """
 
         was_from_eigengap = False
@@ -168,7 +171,7 @@ def compute_macrostates(
         if self._gpcca._p_X.shape[1] < n_states:
             logg.warning(
                 f"Requested more macrostates `{n_states}` than available "
-                f"Schur vectors `{self._gpcca.X.shape[1]}`. Recomputing the decomposition"
+                f"Schur vectors `{self._gpcca._p_X.shape[1]}`. Recomputing the decomposition"
             )
 
         start = logg.info(f"Computing `{n_states}` macrostates")
@@ -261,8 +264,8 @@ def set_terminal_states_from_macrostates(
         None
             Nothing, just updates the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
         """
 
         if not isinstance(n_cells, int):
@@ -383,8 +386,8 @@ def compute_terminal_states(
         None
             Nothing, just updates the following fields:
 
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
         """
 
         if len(self._get(P.MACRO).cat.categories) == 1:
@@ -437,21 +440,21 @@ def compute_gdpt(
         self, n_components: int = 10, key_added: str = "gdpt_pseudotime", **kwargs
     ):
         """
-        Compute generalized Diffusion pseudotime from [Haghverdi16]_ using the real Schur decomposition.
+        Compute generalized Diffusion pseudotime from :cite:`haghverdi:16` using the real Schur decomposition.
 
         Parameters
         ----------
         n_components
             Number of real Schur vectors to consider.
         key_added
-            Key in :paramref:`adata` ``.obs`` where to save the pseudotime.
+            Key in :attr:`adata` ``.obs`` where to save the pseudotime.
         kwargs
             Keyword arguments for :meth:`cellrank.tl.GPCCA.compute_schur` if Schur decomposition is not found.
 
         Returns
         -------
         None
-            Nothing, just updates :paramref:`adata` ``.obs[key_added]`` with the computed pseudotime.
+            Nothing, just updates :attr:`adata` ``.obs[key_added]`` with the computed pseudotime.
         """
 
         def _get_dpt_row(e_vals: np.ndarray, e_vecs: np.ndarray, i: int):
@@ -738,12 +741,126 @@ def annotate_dist_ax(ax, data: np.ndarray, valfmt: str = "{x:.2f}"):
 
         if annotate:
             annotate_heatmap(im)
-            annotate_dist_ax(stat_ax, coarse_stat_d.values)
-            annotate_dist_ax(init_ax, coarse_init_d)
+            if show_stationary_dist:
+                annotate_dist_ax(stat_ax, coarse_stat_d.values)
+            if show_initial_dist:
+                annotate_dist_ax(init_ax, coarse_init_d)
 
         if save:
             save_fig(fig, save)
 
+    @d.dedent
+    def plot_macrostate_composition(
+        self,
+        key: str,
+        width: float = 0.8,
+        title: Optional[str] = None,
+        labelrot: float = 45,
+        legend_loc: Optional[str] = "upper right out",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        show: bool = True,
+    ) -> Optional[Axes]:
+        """
+        Plot stacked histogram of macrostates over categorical annotations.
+
+        Parameters
+        ----------
+        %(adata)s
+        key
+            Key from :attr:`adata` ``.obs`` containing categorical annotations.
+        width
+            Bar width in `[0, 1]`.
+        title
+            Title of the figure. If `None`, create one automatically.
+        labelrot
+            Rotation of labels on x-axis.
+        legend_loc
+            Position of the legend. If `None`, don't show legend.
+        %(plotting)s
+        show
+            If `False`, return :class:`matplotlib.pyplot.Axes`.
+
+        Returns
+        -------
+        :class:`matplotlib.pyplot.Axes`
+            The axis object if ``show=False``.
+        %(just_plots)s
+        """
+        from cellrank.pl._utils import _position_legend
+
+        macrostates = self._get(P.MACRO)
+        if macrostates is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+        if key not in self.adata.obs:
+            raise KeyError(f"Key `{key}` not found in `adata.obs`.")
+        if not is_categorical_dtype(self.adata.obs[key]):
+            raise TypeError(
+                f"Expected `adata.obs[{key!r}]` to be `categorical`, "
+                f"found `{infer_dtype(self.adata.obs[key])}`."
+            )
+
+        mask = ~macrostates.isnull()
+        df = (
+            pd.DataFrame({"macrostates": macrostates, key: self.adata.obs[key]})[mask]
+            .groupby([key, "macrostates"])
+            .size()
+        )
+        try:
+            cats_colors = self.adata.uns[f"{key}_colors"]
+        except KeyError:
+            cats_colors = _create_categorical_colors(
+                len(self.adata.obs[key].cat.categories)
+            )
+        cat_color_mapper = dict(zip(self.adata.obs[key].cat.categories, cats_colors))
+        x_indices = np.arange(len(macrostates.cat.categories))
+        bottom = np.zeros_like(x_indices, dtype=np.float32)
+
+        width = min(1, max(0, width))
+        fig, ax = plt.subplots(figsize=figsize, dpi=dpi, tight_layout=True)
+        for cat, color in cat_color_mapper.items():
+            frequencies = df.loc[cat]
+            # do not add to legend if category is missing
+            if np.sum(frequencies) > 0:
+                ax.bar(
+                    x_indices,
+                    frequencies,
+                    width,
+                    label=cat,
+                    color=color,
+                    bottom=bottom,
+                    ec="black",
+                    lw=0.5,
+                )
+                bottom += np.array(frequencies)
+
+        ax.set_xticks(x_indices)
+        ax.set_xticklabels(
+            # assuming at least 1 category
+            frequencies.index,
+            rotation=labelrot,
+            ha="center" if labelrot in (0, 90) else "right",
+        )
+        y_max = bottom.max()
+        ax.set_ylim([0, y_max + 0.05 * y_max])
+        ax.set_yticks(np.linspace(0, y_max, 5))
+        ax.margins(0.05)
+
+        ax.set_xlabel("macrostate")
+        ax.set_ylabel("frequency")
+        if title is None:
+            title = f"distribution over {key}"
+        ax.set_title(title)
+        if legend_loc not in (None, "none"):
+            _position_legend(ax, legend_loc=legend_loc)
+
+        if save is not None:
+            save_fig(fig, save)
+
+        if not show:
+            return ax
+
     def _compute_one_macrostate(
         self,
         n_cells: int,
@@ -870,13 +987,13 @@ def _set_macrostates(
             Fuzzy clustering.
         %(n_cells)s
         cluster_key
-            Key from :paramref:`adata` ``.obs`` to get reference cluster annotations.
+            Key from :attr:`adata` ``.obs`` to get reference cluster annotations.
         en_cutoff
             Threshold to decide when we we want to warn the user about an uncertain name mapping. This happens when
             one fuzzy state overlaps with several reference clusters, and the most likely cells are distributed almost
             evenly across the reference clusters.
         p_thresh
-            Only used to detect cell cycle stages. These have to be present in :paramref:`adata` ``.obs`` as
+            Only used to detect cell cycle stages. These have to be present in :attr:`adata` ``.obs`` as
             `'G2M_score'` and `'S_score'`.
         check_row_sums
             Check whether rows in `memberships` sum to `1`.
@@ -1071,12 +1188,12 @@ def fit(
         None
             Nothing, just makes available the following fields:
 
-                - :paramref:`{msp}`
-                - :paramref:`{ms}`
-                - :paramref:`{fsp}`
-                - :paramref:`{fs}`
-                - :paramref:`{ap}`
-                - :paramref:`{pd}`
+                - :attr:`{msp}`
+                - :attr:`{ms}`
+                - :attr:`{fsp}`
+                - :attr:`{fs}`
+                - :attr:`{ap}`
+                - :attr:`{pd}`
         """
 
         super().fit(
@@ -1157,7 +1274,7 @@ def _set_initial_states_from_macrostates(
         Returns
         -------
         None
-            Nothing, just writes to :paramref:`adata`:
+            Nothing, just writes to :attr:`adata`:
 
                 - ``.obs[{key!r}]`` - probability of being an initial state.
                 - ``.obs[{probs_key!r}]`` - top ``n_cells`` from each initial state.
@@ -1208,3 +1325,25 @@ def _write_initial_states(
             f"Adding `adata.obs[{_probs(key)!r}]`\n       `adata.obs[{key!r}]`\n",
             time=time,
         )
+
+    def _write_terminal_states(self, time=None) -> None:
+        super()._write_terminal_states(time=time)
+
+        term_abs_probs = self._get(A.TERM_ABS_PROBS)
+        if term_abs_probs is None:
+            # possibly remove previous value if it's inconsistent
+            term_abs_probs = self.adata.obsm.get(self._term_abs_prob_key, None)
+
+        if term_abs_probs is not None:
+            new = list(self._get(P.TERM).cat.categories)
+            old = list(term_abs_probs.names)
+            if term_abs_probs.shape[1] == len(new) and new == old:
+                self.adata.obsm[self._term_abs_prob_key] = term_abs_probs
+            else:
+                logg.warning(
+                    f"Removing previously computed `adata.obsm[{self._term_abs_prob_key!r}]` because the "
+                    f"names mismatch `{new}` (new), `{old}` (old)."
+                )
+
+                self._set(A.TERM_ABS_PROBS, None)
+                self.adata.obsm.pop(self._term_abs_prob_key, None)
diff --git a/cellrank/tl/estimators/_property.py b/cellrank/tl/estimators/_property.py
index d66fbc7f0..cfcf58f98 100644
--- a/cellrank/tl/estimators/_property.py
+++ b/cellrank/tl/estimators/_property.py
@@ -5,17 +5,8 @@
 from inspect import isabstract
 
 import scvelo as scv
-from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from scipy.sparse import issparse, spmatrix
-from pandas.api.types import is_categorical_dtype
-
-import matplotlib as mpl
-from matplotlib import cm
-
 import cellrank.logging as logg
+from anndata import AnnData
 from cellrank.ul._docs import d, _initial, _terminal
 from cellrank.tl._utils import RandomKeys, _make_cat, _partition, _complex_warning
 from cellrank.ul._utils import Pickleable
@@ -34,6 +25,14 @@
 from cellrank.tl.kernels._base_kernel import KernelExpression
 from cellrank.tl.estimators._constants import META_KEY, A, F, P
 
+import numpy as np
+import pandas as pd
+from scipy.sparse import issparse, spmatrix
+from pandas.api.types import is_categorical_dtype
+
+import matplotlib as mpl
+from matplotlib import cm
+
 
 def is_abstract(classname: str) -> bool:  # TODO: determine the necessity of this
     """
@@ -143,9 +142,9 @@ def __new__(cls, clsname, superclasses, attributedict):
         elif len(compute_md) == 0:
             raise ValueError("No metadata found.")
         else:
-            compute_md, *metadata = [
+            compute_md, *metadata = (
                 Metadata(attr=md) if isinstance(md, str) else md for md in compute_md
-            ]
+            )
 
         prop_name = PropertyMeta.update_attributes(compute_md, attributedict)
         plot_name = str(compute_md.plot_fmt).format(prop_name)
@@ -298,7 +297,7 @@ def _plot_vectors(
         abs_value
             Whether to take the absolute value before plotting.
         cluster_key
-            Key in :paramref:`adata` ``.obs`` for plotting categorical observations.
+            Key in :attr:`adata` ``.obs`` for plotting categorical observations.
         %(basis)s
         kwargs
             Keyword arguments for :func:`scvelo.pl.scatter`.
@@ -409,7 +408,7 @@ def _plot_discrete(
         lineages
             Plot only these lineages. If `None`, plot all lineages.
         cluster_key
-            Key from :paramref:`adata` ``.obs`` for plotting categorical observations.
+            Key from :attr:`adata` ``.obs`` for plotting categorical observations.
         same_plot
             Whether to plot the lineages on the same plot or separately.
         title
@@ -543,10 +542,10 @@ def _plot_continuous(
         lineages
             Plot only these lineages. If `None`, plot all lineages.
         cluster_key
-            Key from :paramref:`adata` ``.obs`` for plotting categorical observations.
+            Key from :attr:`adata` ``.obs`` for plotting categorical observations.
         %(time_mode)s
         time_key
-            Key from :paramref:`adata` ``.obs`` to use as a pseudotime ordering of the cells.
+            Key from :attr:`adata` ``.obs`` to use as a pseudotime ordering of the cells.
         title
             Either `None`, in which case titles are ``'{to,from} {terminal,initial} {state}'``,
             or an array of titles, one per lineage.
@@ -873,9 +872,9 @@ def compute_partition(self) -> None:
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`recurrent_classes`
-                - :paramref:`transient_classes`
-                - :paramref:`is_irreducible`
+                - :attr:`recurrent_classes`
+                - :attr:`transient_classes`
+                - :attr:`is_irreducible`
         """
 
         start = logg.info("Computing communication classes")
diff --git a/cellrank/tl/kernels/__init__.py b/cellrank/tl/kernels/__init__.py
index d05df49cb..817773ef8 100644
--- a/cellrank/tl/kernels/__init__.py
+++ b/cellrank/tl/kernels/__init__.py
@@ -1,4 +1,8 @@
 from cellrank.tl.kernels._base_kernel import Kernel
+from cellrank.tl.kernels._exp_time_kernel import (
+    TransportMapKernel,
+    ExperimentalTimeKernel,
+)
 from cellrank.tl.kernels._velocity_kernel import VelocityKernel
 from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEKernel
 from cellrank.tl.kernels._velocity_schemes import (
diff --git a/cellrank/tl/kernels/_base_kernel.py b/cellrank/tl/kernels/_base_kernel.py
index 8f26b1f7d..50463657d 100644
--- a/cellrank/tl/kernels/_base_kernel.py
+++ b/cellrank/tl/kernels/_base_kernel.py
@@ -9,7 +9,6 @@
     Type,
     Tuple,
     Union,
-    TypeVar,
     Callable,
     Iterable,
     Optional,
@@ -19,17 +18,7 @@
 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 anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import (
@@ -38,14 +27,24 @@
     _normalize,
     _symmetric,
     _get_neighs,
-    _has_neighs,
     _irreducible,
 )
 from cellrank.ul._utils import Pickleable, _write_graph_data
+from scvelo.plotting.utils import default_size, plot_outline
 from cellrank.tl._constants import Direction, _transition
 from cellrank.tl.kernels._utils import _get_basis, _filter_kwargs
+from cellrank.tl.kernels._tmat_flow import FlowPlotter
 from cellrank.tl.kernels._random_walk import RandomWalk
 
+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
+
 _ERROR_DIRECTION_MSG = "Can only combine kernels that have the same direction."
 _ERROR_EMPTY_CACHE_MSG = (
     "Fatal error: tried to used cached values, but the cache was empty."
@@ -61,7 +60,6 @@
 _n_dec = 2
 _dtype = np.float64
 _cond_num_tolerance = 1e-15
-AnnData = TypeVar("AnnData")
 Indices_t = Optional[Union[Sequence[str], Dict[str, Union[str, Sequence[str]]]]]
 
 
@@ -87,7 +85,7 @@ def __init_subclass__(cls, **kwargs: Any):
         super().__init_subclass__()
 
     @property
-    def condition_number(self):
+    def condition_number(self) -> Optional[int]:
         """Condition number of the transition matrix."""
         return self._cond_num
 
@@ -96,7 +94,7 @@ def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
         """
         Return row-normalized transition matrix.
 
-        If not present, it is computed, if all the underlying kernels have been initialized.
+        If not present, it is computed iff all underlying kernels have been initialized.
         """
 
         if self._parent is None and self._transition_matrix is None:
@@ -123,7 +121,7 @@ def adata(self) -> AnnData:
 
     @adata.setter
     @abstractmethod
-    def adata(self, value):
+    def adata(self, value: AnnData) -> None:
         pass
 
     @property
@@ -153,7 +151,7 @@ def transition_matrix(self, value: Union[np.ndarray, spmatrix]) -> None:
         Returns
         -------
         None
-            Nothing, just updates the :paramref:`transition_matrix` and optionally normalizes it.
+            Nothing, just updates the :attr:`transition_matrix` and optionally normalizes it.
         """
         should_norm = ~np.isclose(value.sum(1), 1.0, rtol=_RTOL).all()
 
@@ -166,13 +164,15 @@ def transition_matrix(self, value: Union[np.ndarray, spmatrix]) -> None:
             )
 
     @abstractmethod
-    def compute_transition_matrix(self, *args, **kwargs) -> "KernelExpression":
+    def compute_transition_matrix(
+        self, *args: Any, **kwargs: Any
+    ) -> "KernelExpression":
         """
         Compute a transition matrix.
 
         Parameters
         ----------
-        *args
+        args
             Positional arguments.
         kwargs
             Keyword arguments.
@@ -188,13 +188,13 @@ def compute_transition_matrix(self, *args, **kwargs) -> "KernelExpression":
     @d.dedent
     def write_to_adata(self, key: Optional[str] = None) -> None:
         """
-        Write the transition matrix and parameters used for computation to the underlying :paramref:`adata` object.
+        Write the transition matrix and parameters used for computation to the underlying :attr:`adata` object.
 
         Parameters
         ----------
         key
-            Key used when writing transition matrix to :paramref:`adata`.
-            If `None`, the ``key`` is set to `'T_bwd'` if :paramref:`backward` is `True`, else `'T_fwd'`.
+            Key used when writing transition matrix to :attr:`adata`.
+            If `None`, the ``key`` is set to `'T_bwd'` if :attr:`backward` is `True`, else `'T_fwd'`.
 
         Returns
         -------
@@ -219,9 +219,10 @@ def write_to_adata(self, key: Optional[str] = None) -> None:
 
     @abstractmethod
     def copy(self) -> "KernelExpression":
-        """Return a copy of itself. Note that the underlying :paramref:`adata` object is not copied."""
+        """Return a copy of itself. Note that the underlying :attr:`adata` object is not copied."""
 
-    def _maybe_compute_cond_num(self):
+    def _maybe_compute_cond_num(self) -> None:
+        """Optionally compute condition number."""
         if self._compute_cond_num and self._cond_num is None:
             logg.debug(f"Computing condition number of `{repr(self)}`")
             self._cond_num = np.linalg.cond(
@@ -273,12 +274,12 @@ def compute_projection(
         Parameters
         ----------
         basis
-            Basis in :attr:`anndata.AnnData.obsm` for which to compute the projection.
+            Basis in :attr:`adata` ``.obsm`` for which to compute the projection.
         key_added
-            If not `None` and ``copy=False``, save the result to :paramref:`adata` ``.obsm['{key_added}']``.
+            If not `None` and ``copy=False``, save the result to :attr:`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.
+            Whether to return the projection or modify :attr:`adata` inplace.
 
         Returns
         -------
@@ -393,7 +394,7 @@ def plot_random_walks(
         Returns
         -------
         %(just_plots)s
-        For each random walk, the first (last) cell is marked though a black (yellow) dot.
+        For each random walk, the first/last cell is marked by the start/end colors of ``cmap``.
         """
 
         def create_ixs(ixs: Indices_t, *, kind: str) -> Optional[np.ndarray]:
@@ -499,6 +500,86 @@ def create_ixs(ixs: Indices_t, *, kind: str) -> Optional[np.ndarray]:
         if save is not None:
             save_fig(fig, save)
 
+    @d.get_full_description(base="plot_single_flow")
+    @d.get_sections(base="plot_single_flow", sections=["Parameters", "Returns"])
+    @d.dedent
+    def plot_single_flow(
+        self,
+        cluster: str,
+        cluster_key: str,
+        time_key: str,
+        clusters: Optional[Sequence[Any]] = None,
+        time_points: Optional[Sequence[Union[int, float]]] = None,
+        min_flow: float = 0,
+        remove_empty_clusters: bool = True,
+        ascending: Optional[bool] = False,
+        legend_loc: Optional[str] = "upper right out",
+        alpha: Optional[float] = 0.8,
+        xticks_step_size: Optional[int] = 1,
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        show: bool = True,
+    ) -> Optional[plt.Axes]:
+        """
+        Visualize outgoing flow from a cluster of cells :cite:`mittnenzweig:21`.
+
+        Parameters
+        ----------
+        cluster
+            Cluster for which to visualize outgoing compute_flow.
+        cluster_key
+            Key in :attr:`adata` ``.obs`` where clustering is stored.
+        time_key
+            Key in :attr:`adata` ``.obs`` where experimental time is stored.
+        clusters
+            Visualize flow only for these clusters. If `None`, use all clusters.
+        time_points
+            Visualize flow only for these time points. If `None`, use all time points.
+        %(flow.parameters)s
+        %(plotting)s
+        show
+            If `False`, return :class:`matplotlib.pyplot.Axes`.
+
+        Returns
+        -------
+        :class:`matplotlib.pyplot.Axes`
+            The axis object if ``show=False``.
+        %(just_plots)s
+
+        Notes
+        -----
+        This function is a Python reimplementation of the following
+        `original R function <https://github.com/tanaylab/embflow/blob/main/scripts/generate_paper_figures/plot_vein.r>`_
+        with some minor stylistic differences.
+        This function will not recreate the results from :cite:`mittnenzweig:21`, because there the Metacell model
+        :cite:`baran:19` was used to compute the flow, whereas here the transition matrix is used.
+        """  # noqa: E501
+        if self._transition_matrix is None:
+            raise RuntimeError(
+                "Compute transition matrix first as `.compute_transition_matrix()`."
+            )
+
+        fp = FlowPlotter(self.adata, self.transition_matrix, cluster_key, time_key)
+        fp = fp.prepare(cluster, clusters, time_points)
+
+        ax = fp.plot(
+            min_flow=min_flow,
+            remove_empty_clusters=remove_empty_clusters,
+            ascending=ascending,
+            alpha=alpha,
+            xticks_step_size=xticks_step_size,
+            legend_loc=legend_loc,
+            figsize=figsize,
+            dpi=dpi,
+        )
+
+        if save is not None:
+            save_fig(ax.figure, save)
+
+        if not show:
+            return ax
+
     def __xor__(self, other: "KernelExpression") -> "KernelExpression":
         return self.__rxor__(other)
 
@@ -691,9 +772,7 @@ def adata(self) -> AnnData:
         return self._adata
 
     @adata.setter
-    def adata(self, _adata: AnnData):
-        from anndata import AnnData
-
+    def adata(self, _adata: AnnData) -> None:
         if not isinstance(_adata, AnnData):
             raise TypeError(
                 f"Expected argument of type `anndata.AnnData`, found `{type(_adata).__name__!r}`."
@@ -784,7 +863,7 @@ def adata(self) -> AnnData:
         return self._kexprs[0].adata
 
     @adata.setter
-    def adata(self, _adata: AnnData):
+    def adata(self, _adata: AnnData) -> None:
         self._kexprs[0].adata = _adata
 
     def __invert__(self) -> "NaryKernelExpression":
@@ -828,12 +907,11 @@ class Kernel(UnaryKernelExpression, ABC):
     ----------
     %(adata)s
     %(backward)s
-    compute_cond_num
-        Whether to compute the condition number of the transition matrix. For large matrices, this can be very slow.
+    %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
     kwargs
-        Keyword arguments which can specify key to be read from :paramref:`adata` object.
+        Keyword arguments which can specify key to be read from :attr:`adata` object.
     """
 
     def __init__(
@@ -842,20 +920,44 @@ def __init__(
         backward: bool = False,
         compute_cond_num: bool = False,
         check_connectivity: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ):
         super().__init__(
             adata, backward, op_name=None, compute_cond_num=compute_cond_num, **kwargs
         )
         self._read_from_adata(check_connectivity=check_connectivity, **kwargs)
 
-    def _read_from_adata(self, key: str = "connectivities", **kwargs):
-        """Import the base-KNN graph and optionally check for symmetry and connectivity."""
+    # TODO: move to a mixin class
+    def _read_from_adata(
+        self,
+        conn_key: Optional[str] = "connectivities",
+        read_conn: bool = True,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Import the base-KNN graph and optionally check for symmetry and connectivity.
+
+        Parameters
+        ----------
+        conn_key
+            Key in :attr:`anndata.AnnData.uns` where connectivities are stored.
+        read_conn
+            Whether to read connectivities or set them to `None`. Useful when not exposing density normalization or
+            when KNN connectivities are not used to compute the transition matrix.
+        kwargs
+            Additional keyword arguments.
 
-        if not _has_neighs(self.adata):
-            raise KeyError("Compute KNN graph first as `scanpy.pp.neighbors()`.")
+        Returns
+        -------
+        Nothing, just sets :attr:`_conn`.
+        """
+        if not read_conn:
+            self._conn = None
+            return
 
-        self._conn = _get_neighs(self.adata, key).astype(_dtype)
+        self._conn = _get_neighs(
+            self.adata, mode="connectivities", key=conn_key
+        ).astype(_dtype)
 
         check_connectivity = kwargs.pop("check_connectivity", False)
         if check_connectivity:
@@ -879,23 +981,29 @@ def _density_normalize(
 
         Parameters
         ----------
-        other:
+        other
             Matrix to normalize.
 
         Returns
         -------
         :class:`np.ndarray` or :class:`scipy.sparse.spmatrix`
             Density normalized transition matrix.
+
+        Raises
+        ------
+        ValueError
+            If KNN connectivities are not set.
         """
+        if self._conn is None:
+            raise ValueError(
+                "Unable to density normalize the transition matrix "
+                "because KNN connectivities are not set."
+            )
 
         logg.debug("Density-normalizing the transition matrix")
 
         q = np.asarray(self._conn.sum(axis=0))
-
-        if not issparse(other):
-            Q = np.diag(1.0 / q)
-        else:
-            Q = spdiags(1.0 / q, 0, other.shape[0], other.shape[0])
+        Q = spdiags(1.0 / q, 0, other.shape[0], other.shape[0])
 
         return Q @ other @ Q
 
@@ -904,13 +1012,23 @@ def _get_kernels(self) -> Iterable["Kernel"]:
 
     def _compute_transition_matrix(
         self,
-        matrix: spmatrix,
+        matrix: Union[np.ndarray, spmatrix],
         density_normalize: bool = True,
         check_irreducibility: bool = False,
     ):
-        # density correction based on node degrees in the KNN graph
-        matrix = csr_matrix(matrix) if not isspmatrix_csr(matrix) else matrix
+        if matrix.shape[0] != matrix.shape[1]:
+            raise ValueError(f"Expected a square matrix, found `{matrix.shape}`.")
+        if matrix.shape[0] != self.adata.n_obs:
+            raise ValueError(
+                f"Expected matrix to be of shape `{(self.adata.n_obs, self.adata.n_obs)}`, "
+                f"found `{matrix.shape}`."
+            )
+
+        matrix = matrix.astype(_dtype)
+        if issparse(matrix) and not isspmatrix_csr(matrix):
+            matrix = csr_matrix(matrix)
 
+        # density correction based on node degrees in the KNN graph
         if density_normalize:
             matrix = self._density_normalize(matrix)
 
@@ -919,7 +1037,7 @@ def _compute_transition_matrix(
         if len(problematic_indices):
             logg.warning(
                 f"Detected `{len(problematic_indices)}` absorbing states in the transition matrix. "
-                f"This matrix won't be irreducible."
+                f"This matrix won't be irreducible"
             )
             matrix[problematic_indices, problematic_indices] = 1.0
 
@@ -940,7 +1058,7 @@ class Constant(Kernel):
     ----------
     %(adata)s
     value
-        Constant value by which to multiply Must be a positive number.
+        Constant value by which to multiply. Must be a positive number.
     %(backward)s
     """
 
@@ -956,11 +1074,11 @@ def __init__(
             raise ValueError(f"Expected the constant to be positive, found `{value}`.")
         self._recalculate(value)
 
-    def _recalculate(self, value):
+    def _recalculate(self, value) -> None:
         self._transition_matrix = value
         self._params = {"value": value}
 
-    def _read_from_adata(self, **kwargs):
+    def _read_from_adata(self, **kwargs: Any) -> None:
         pass
 
     def compute_transition_matrix(self, *args, **kwargs) -> "Constant":
@@ -1169,7 +1287,7 @@ def _get_expr_and_constant(k: KernelMul) -> Tuple[KernelExpression, Union[int, f
         )
     if len(k) != 2:
         raise ValueError(
-            f"Expected expression to be binary, found, `{len(k)}` subexpressions."
+            f"Expected expression to be binary, found `{len(k)}` subexpressions."
         )
     e1, e2 = k[0], k[1]
 
diff --git a/cellrank/tl/kernels/_connectivity_kernel.py b/cellrank/tl/kernels/_connectivity_kernel.py
index 7de71ca29..75d80eb7d 100644
--- a/cellrank/tl/kernels/_connectivity_kernel.py
+++ b/cellrank/tl/kernels/_connectivity_kernel.py
@@ -2,7 +2,6 @@
 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
@@ -15,8 +14,8 @@ class ConnectivityKernel(Kernel):
 
     As a measure of similarity, we currently support:
 
-        - transcriptomic similarities, computed using e.g. :func:`scanpy.pp.neighbors`, see [Wolf18]_.
-        - spatial similarities, computed using e.g. :func:`squidpy.gr.spatial_neighbors`, see [Palla21]_.
+        - transcriptomic similarities, computed using e.g. :func:`scanpy.pp.neighbors`, see :cite:`wolf:18`.
+        - spatial similarities, computed using e.g. :func:`squidpy.gr.spatial_neighbors`, see :cite:`palla:21`.
 
     The resulting transition matrix is symmetric and thus cannot be used to learn about the direction of the biological
     process. To include this direction, consider combining with a velocity-derived transition matrix via
@@ -29,7 +28,7 @@ class ConnectivityKernel(Kernel):
     %(adata)s
     %(backward)s
     conn_key
-        Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix, describing cell-cell similarity.
+        Key in :attr:`anndata.AnnData.obsp` to obtain the connectivity matrix describing cell-cell similarity.
     %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
@@ -48,7 +47,7 @@ def __init__(
             backward=backward,
             compute_cond_num=compute_cond_num,
             check_connectivity=check_connectivity,
-            key=conn_key,
+            conn_key=conn_key,
         )
         self._key = conn_key
 
@@ -70,7 +69,7 @@ def compute_transition_matrix(
         Returns
         -------
         :class:`cellrank.tl.kernels.ConnectivityKernel`
-            Makes :paramref:`transition_matrix` available.
+            Makes :attr:`transition_matrix` available.
         """
 
         # fmt: off
diff --git a/cellrank/tl/kernels/_cytotrace_kernel.py b/cellrank/tl/kernels/_cytotrace_kernel.py
index 24790696f..ca208dd00 100644
--- a/cellrank/tl/kernels/_cytotrace_kernel.py
+++ b/cellrank/tl/kernels/_cytotrace_kernel.py
@@ -3,16 +3,15 @@
 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
 
+import numpy as np
+from scipy.stats import gmean, hmean
+
 
 def _ct(key: str) -> str:
     return f"ct_{key}"
@@ -28,7 +27,8 @@ class CytoTRACEAggregation(ModeEnum):  # noqa
 @d.dedent
 class CytoTRACEKernel(PseudotimeKernel):
     """
-    Kernel which computes directed transition probabilities based on a KNN graph and the CytoTRACE score [Cyto20]_.
+    Kernel which computes directed transition probabilities based on a KNN graph and the CytoTRACE score \
+    :cite:`gulati:20`.
 
     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
@@ -44,31 +44,40 @@ class CytoTRACEKernel(PseudotimeKernel):
     %(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.
+    %(cond_num)s
+    check_connectivity
+        Check whether the underlying KNN graph is connected.
+    kwargs
+        Keyword arguments for :class:`cellrank.tl.kernels.PseudotimeKernel`.
 
-    Examples
-    --------
+    Example
+    -------
     Workflow::
 
+        # import packages and load data
         import scvelo as scv
         import cellrank as cr
-
         adata = cr.datasets.pancreas()
 
+        # standard pre-processing
         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)
 
+        # CytoTRACE by default uses imputed data - a simple way to compute KNN-imputed data is to use scVelo's moments
+        # function. However, note that this function expects `spliced` counts because it's designed for RNA velocity,
+        # so we're using a simple hack here:
         if 'spliced' not in adata.layers or 'unspliced' not in adata.layers:
-            # 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)
+        # compute KNN-imputation using scVelo's moments function
+        scv.pp.moments(adata)
+
+        # import and initialize the CytoTRACE kernel, compute transition matrix - done!
+        from cellrank.tl.kernels import CytoTRACEKernel
+        ctk = CytoTRACEKernel(adata).compute_transition_matrix()
     """
 
     def __init__(
@@ -80,6 +89,7 @@ def __init__(
         use_raw: bool = False,
         compute_cond_num: bool = False,
         check_connectivity: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             adata,
@@ -90,6 +100,7 @@ def __init__(
             layer=layer,
             aggregation=aggregation,
             use_raw=use_raw,
+            **kwargs,
         )
         self._time_key = _ct("pseudotime")  # quirk or PT kernel
 
@@ -100,7 +111,7 @@ def _read_from_adata(
         aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
         use_raw: bool = True,
         **kwargs: Any,
-    ):
+    ) -> None:
         self.compute_cytotrace(layer=layer, aggregation=aggregation, use_raw=use_raw)
 
         super()._read_from_adata(time_key=time_key, **kwargs)
@@ -113,16 +124,13 @@ def compute_cytotrace(
         use_raw: bool = False,
     ) -> None:
         """
-        Re-implementation of the CytoTRACE algorithm [Cyto20]_ to estimate cellular plasticity.
+        Re-implementation of the CytoTRACE algorithm :cite:`gulati:20` 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
@@ -153,6 +161,11 @@ def compute_cytotrace(
             - `'ct_gene_corr'`: the correlation as specified above.
             - `'ct_correlates'`: indication of the genes used to compute the CytoTRACE score, i.e. the ones that
               correlated best with `'num_exp_genes'`.
+
+        Notes
+        -----
+        This will not exactly reproduce the results of the original CytoTRACE algorithm :cite:`gulati:20` because we
+        allow for any normalization and imputation techniques whereas CytoTRACE has built-in specific methods for that.
         """
         # check use_raw
         aggregation = CytoTRACEAggregation(aggregation)
diff --git a/cellrank/tl/kernels/_exp_time_kernel.py b/cellrank/tl/kernels/_exp_time_kernel.py
new file mode 100644
index 000000000..88debb15c
--- /dev/null
+++ b/cellrank/tl/kernels/_exp_time_kernel.py
@@ -0,0 +1,191 @@
+"""Experimental time kernel module."""
+from abc import ABC
+from copy import copy
+from types import MappingProxyType
+from typing import Any, Dict, Tuple, Mapping, Optional
+
+import scanpy as sc
+from cellrank.ul._docs import d
+from cellrank.tl._utils import _normalize
+from cellrank.tl.kernels import Kernel
+from cellrank.tl._constants import ModeEnum
+from cellrank.tl.kernels._utils import _ensure_numeric_ordered
+from cellrank.tl.kernels._base_kernel import AnnData
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import bmat, spdiags
+
+
+class LastTimePoint(ModeEnum):  # noqa
+    UNIFORM = "uniform"
+    DIAGONAL = "diagonal"
+    CONNECTIVITIES = "connectivities"
+
+
+@d.dedent
+class ExperimentalTimeKernel(Kernel, ABC):
+    """
+    Kernel base class which computes directed transition probabilities based on experimental time.
+
+    %(density_correction)s
+
+    Parameters
+    ----------
+    %(adata)s
+    %(backward)s
+    time_key
+        Key in :attr:`adata` ``.obs`` where experimental time is stored.
+        The experimental time can be of either of a numeric or an ordered categorical type.
+    %(cond_num)s
+    """
+
+    def __init__(
+        self,
+        adata: AnnData,
+        backward: bool = False,
+        time_key: str = "exp_time",
+        compute_cond_num: bool = False,
+    ):
+        super().__init__(
+            adata,
+            backward=backward,
+            time_key=time_key,
+            compute_cond_num=compute_cond_num,
+            check_connectivity=False,
+        )
+        self._time_key = time_key
+
+    def _read_from_adata(self, **kwargs: Any) -> None:
+        super()._read_from_adata(**kwargs)
+
+        time_key = kwargs.pop("time_key", "exp_time")
+        self._exp_time = _ensure_numeric_ordered(self.adata, time_key)
+        self.adata.obs[time_key] = self.experimental_time.values
+
+    @d.dedent
+    def plot_single_flow(
+        self,
+        cluster: str,
+        cluster_key: str,
+        time_key: Optional[str] = None,
+        *args: Any,
+        **kwargs: Any,
+    ) -> None:
+        """
+        %(plot_single_flow.full_desc)s
+
+        Parameters
+        ----------
+        %(plot_single_flow.parameters)s
+
+        Returns
+        -------
+        %(plot_single_flow.returns)s
+        """  # noqa: D400
+        if time_key is None:
+            time_key = self._time_key
+        return super().plot_single_flow(cluster, cluster_key, time_key, *args, **kwargs)
+
+    @property
+    def experimental_time(self) -> pd.Series:
+        """Experimental time."""
+        return self._exp_time
+
+    def copy(self) -> "ExperimentalTimeKernel":
+        """Return a copy of self."""
+        pk = type(self)(self.adata, backward=self.backward, time_key=self._time_key)
+        pk._exp_time = copy(self.experimental_time)
+        pk._params = copy(self._params)
+        pk._cond_num = self.condition_number
+        pk._transition_matrix = copy(self._transition_matrix)
+
+        return pk
+
+    def __invert__(self) -> "ExperimentalTimeKernel":
+        super().__invert__()
+        if len(self.experimental_time.cat.categories) > 1:
+            minn, maxx = self.experimental_time.min(), self.experimental_time.max()
+            self._exp_time = pd.Series(
+                pd.Categorical(
+                    maxx - np.array(self.experimental_time) + minn, ordered=True
+                ),
+                index=self.experimental_time.index,
+            )
+        return self
+
+
+class TransportMapKernel(ExperimentalTimeKernel, ABC):
+    """Kernel base class which computes transition matrix based on transport maps for consecutive time pairs."""
+
+    def __init__(self, *args: Any, **kwargs: Any):
+        super().__init__(*args, **kwargs)
+        self._tmaps = None
+
+    def _restich_tmaps(
+        self,
+        tmaps: Mapping[Tuple[float, float], AnnData],
+        last_time_point: LastTimePoint = LastTimePoint.DIAGONAL,
+        conn_kwargs: Mapping[str, Any] = MappingProxyType({}),
+        normalize: bool = True,
+    ) -> AnnData:
+        from cellrank.tl.kernels import ConnectivityKernel
+
+        conn_kwargs = dict(conn_kwargs)
+        conn_kwargs["copy"] = False
+        _ = conn_kwargs.pop("key_added", None)
+        density_normalize = conn_kwargs.pop("density_normalize", True)
+
+        blocks = [[None] * (len(tmaps) + 1) for _ in range(len(tmaps) + 1)]
+        nrows, ncols = 0, 0
+        obs_names, obs = [], []
+
+        for i, tmap in enumerate(tmaps.values()):
+            blocks[i][i + 1] = _normalize(tmap.X) if normalize else tmap.X
+            nrows += tmap.n_obs
+            ncols += tmap.n_vars
+            obs_names.extend(tmap.obs_names)
+            obs.append(tmap.obs)
+        obs_names.extend(tmap.var_names)
+
+        n = self.adata.n_obs - nrows
+        if last_time_point == LastTimePoint.DIAGONAL:
+            blocks[-1][-1] = spdiags([1] * n, 0, n, n)
+        elif last_time_point == LastTimePoint.UNIFORM:
+            blocks[-1][-1] = np.ones((n, n)) / float(n)
+        elif last_time_point == LastTimePoint.CONNECTIVITIES:
+            adata_subset = self.adata[tmap.var_names].copy()
+            sc.pp.neighbors(adata_subset, **conn_kwargs)
+            blocks[-1][-1] = (
+                ConnectivityKernel(adata_subset)
+                .compute_transition_matrix(density_normalize)
+                .transition_matrix
+            )
+        else:
+            raise NotImplementedError(
+                f"Last time point mode `{last_time_point}` is not yet implemented."
+            )
+
+        # prevent the last block from disappearing
+        n = blocks[0][1].shape[0]
+        blocks[0][0] = spdiags([], 0, n, n)
+
+        tmp = AnnData(bmat(blocks, format="csr"))
+        tmp.obs_names = obs_names
+        tmp.var_names = obs_names
+        tmp = tmp[self.adata.obs_names, :][:, self.adata.obs_names]
+
+        tmp.obs = pd.merge(
+            tmp.obs,
+            pd.concat(obs),
+            left_index=True,
+            right_index=True,
+            how="left",
+        )
+
+        return tmp
+
+    @property
+    def transport_maps(self) -> Optional[Dict[Tuple[float, float], AnnData]]:
+        """Transport maps for consecutive time pairs."""
+        return self._tmaps
diff --git a/cellrank/tl/kernels/_precomputed_kernel.py b/cellrank/tl/kernels/_precomputed_kernel.py
index c5f4f87bd..2a7a1369e 100644
--- a/cellrank/tl/kernels/_precomputed_kernel.py
+++ b/cellrank/tl/kernels/_precomputed_kernel.py
@@ -1,13 +1,8 @@
 """Precomputed kernel module."""
 from copy import copy
-from typing import Union, Optional
+from typing import Any, Union, Optional
 
 from anndata import AnnData
-
-import numpy as np
-from scipy.sparse import eye as speye
-from scipy.sparse import spmatrix, csr_matrix
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 from cellrank.ul._utils import _read_graph_data
@@ -15,6 +10,10 @@
 from cellrank.tl._constants import Direction, _transition
 from cellrank.tl.kernels._base_kernel import _RTOL, KernelExpression
 
+import numpy as np
+from scipy.sparse import eye as speye
+from scipy.sparse import spmatrix, csr_matrix
+
 
 @d.dedent
 class PrecomputedKernel(Kernel):
@@ -24,12 +23,15 @@ class PrecomputedKernel(Kernel):
     Parameters
     ----------
     transition_matrix
-        Row-normalized transition matrix or a key in :paramref:`adata` ``.obsp``
-        or a :class:`cellrank.tl.kernels.KernelExpression` with the computed transition matrix.
+        Row-normalized transition matrix or a key in :attr:`adata` ``.obsp``
+        or a :class:`cellrank.tl.kernels.KernelExpression` with a precomputed transition matrix.
         If `None`, try to determine the key based on ``backward``.
     %(adata)s
+        If `None`, a temporary placeholder object is created.
     %(backward)s
     %(cond_num)s
+    kwargs
+        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
     """
 
     def __init__(
@@ -40,6 +42,7 @@ def __init__(
         adata: Optional[AnnData] = None,
         backward: bool = False,
         compute_cond_num: bool = False,
+        **kwargs: Any,
     ):
         from anndata import AnnData as _AnnData
 
@@ -97,14 +100,16 @@ def __init__(
                 csr_matrix((transition_matrix.shape[0], 1), dtype=np.float32)
             )
 
-        super().__init__(adata, backward=backward, compute_cond_num=compute_cond_num)
+        super().__init__(
+            adata, backward=backward, compute_cond_num=compute_cond_num, **kwargs
+        )
 
         self._params = params
         self._transition_matrix = csr_matrix(transition_matrix)
         self._maybe_compute_cond_num()
 
-    def _read_from_adata(self, **kwargs):
-        pass
+    def _read_from_adata(self, **kwargs: Any) -> None:
+        self._conn = None
 
     @d.dedent
     def copy(self) -> "PrecomputedKernel":
diff --git a/cellrank/tl/kernels/_pseudotime_kernel.py b/cellrank/tl/kernels/_pseudotime_kernel.py
index 65ecd8e12..a07902839 100644
--- a/cellrank/tl/kernels/_pseudotime_kernel.py
+++ b/cellrank/tl/kernels/_pseudotime_kernel.py
@@ -5,14 +5,11 @@
 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._constants import ThresholdScheme
 from cellrank.tl.kernels._base_kernel import _dtype
 from cellrank.tl.kernels._pseudotime_schemes import (
     ThresholdSchemeABC,
@@ -21,6 +18,8 @@
     CustomThresholdScheme,
 )
 
+import numpy as np
+
 
 @d.dedent
 class PseudotimeKernel(Kernel):
@@ -28,20 +27,18 @@ 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
+    Pseudotime can be used to either remove edges that point against the direction of increasing pseudotime
+    :cite:`setty:19`, or to downweight them :cite:`stassen:21`.
 
     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.
+        Key in :attr:`adata` ``.obs`` where the pseudotime is stored.
+    %(cond_num)s
+    kwargs
+        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
     """
 
     def __init__(
@@ -63,13 +60,15 @@ def __init__(
         )
         self._time_key = time_key
 
-    def _read_from_adata(self, time_key: str, **kwargs):
+    def _read_from_adata(self, time_key: str, **kwargs: Any) -> None:
         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 self.backward:
+            self._pseudotime = np.max(self.pseudotime) - self.pseudotime
 
         if np.any(np.isnan(self._pseudotime)):
             raise ValueError("Encountered NaN values in pseudotime.")
@@ -77,27 +76,28 @@ def _read_from_adata(self, time_key: str, **kwargs):
     @d.dedent
     def compute_transition_matrix(
         self,
-        threshold_scheme: Union[Literal["soft", "hard"], Callable] = "soft",
+        threshold_scheme: Union[Literal["soft", "hard"], Callable] = "hard",
         frac_to_keep: float = 0.3,
-        b: float = 20.0,
-        nu: float = 1.0,
-        percentile: Optional[int] = 95,
-        density_normalize: bool = True,
+        b: float = 10.0,
+        nu: float = 0.5,
         check_irreducibility: bool = False,
+        n_jobs: Optional[int] = None,
+        backend: str = "loky",
+        show_progress_bar: bool = True,
         **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]_).
+        Depending on the choice of the `thresholding_scheme`, this is based on ideas by either *Palantir*
+        :cite:`setty:19` or *VIA* :cite:`stassen:21`.
 
-        When using a `'hard'` thresholding scheme, this based on ideas by *Palantir* (see [Setty19]_) which removes some
+        When using a `'hard'` thresholding scheme, this based on ideas by *Palantir* :cite:`setty:19` 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
+        When using a `'soft'` thresholding scheme, this is based on ideas by *VIA* :cite:`stassen:21` 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.
@@ -109,15 +109,16 @@ def compute_transition_matrix(
             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.
+        %(parallel)s
+        kwargs
+            Keyword arguments for ``threshold_scheme``.
 
         Returns
         -------
         :class:`cellrank.tl.kernels.PseudotimeKernel`
-            Makes :paramref:`transition_matrix` available.
+            Makes :attr:`transition_matrix` available.
         """
         start = logg.info(f"Computing transition matrix based on `{self._time_key}`")
 
@@ -137,7 +138,8 @@ def compute_transition_matrix(
             threshold_scheme = ThresholdScheme(threshold_scheme)
             if threshold_scheme == ThresholdScheme.SOFT:
                 scheme = SoftThresholdScheme()
-                kwargs["b"], kwargs["nu"] = b, nu
+                kwargs["b"] = b
+                kwargs["nu"] = nu
             elif threshold_scheme == ThresholdScheme.HARD:
                 scheme = HardThresholdScheme()
                 kwargs["frac_to_keep"], kwargs["n_neighs"] = frac_to_keep, n_neighbors
@@ -146,7 +148,7 @@ def compute_transition_matrix(
                     f"Threshold scheme `{threshold_scheme}` is not yet implemented."
                 )
         elif isinstance(threshold_scheme, ThresholdSchemeABC):
-            pass
+            scheme = threshold_scheme
         elif callable(threshold_scheme):
             scheme = CustomThresholdScheme(threshold_scheme)
         else:
@@ -155,26 +157,26 @@ def compute_transition_matrix(
             )
 
         # fmt: off
-        if self._reuse_cache({"dnorm": density_normalize, "scheme": str(threshold_scheme), **kwargs}, time=start):
+        if self._reuse_cache({"dnorm": False, "scheme": str(threshold_scheme), **kwargs}, time=start):
             return self
+        # fmt: on
 
-        # 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,
+            self.pseudotime,
+            n_jobs=n_jobs,
+            backend=backend,
+            show_progress_bar=show_progress_bar,
+            **kwargs,
         )
 
-        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,
+            density_normalize=False,
             check_irreducibility=check_irreducibility,
         )
         logg.info("    Finish", time=start)
diff --git a/cellrank/tl/kernels/_pseudotime_schemes.py b/cellrank/tl/kernels/_pseudotime_schemes.py
index 17f6b34bb..9503f5bd5 100644
--- a/cellrank/tl/kernels/_pseudotime_schemes.py
+++ b/cellrank/tl/kernels/_pseudotime_schemes.py
@@ -1,11 +1,12 @@
 from abc import ABC, abstractmethod
-from typing import Any, Callable, Optional
+from typing import Any, Tuple, Callable, Optional
+
+from cellrank.ul._docs import d
+from cellrank.ul._parallelize import parallelize
 
 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."""
@@ -37,8 +38,50 @@ def __call__(
         Array of shape ``(n_neighbors,)`` containing the biased connectivities.
         """
 
+    def _bias_knn_helper(
+        self,
+        ixs: np.ndarray,
+        conn: csr_matrix,
+        pseudotime: np.ndarray,
+        queue=None,
+        **kwargs: Any,
+    ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
+
+        indices, indptr, data = [], [], []
+
+        for i in ixs:
+            row = conn[i]
+            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)}`."
+                )
+
+            data.extend(biased_row)
+            indices.extend(row.indices)
+            indptr.append(conn.indptr[i])
+
+            if queue is not None:
+                queue.put(1)
+
+        if i == conn.shape[0] - 1:
+            indptr.append(conn.indptr[-1])
+        if queue is not None:
+            queue.put(None)
+
+        return np.array(data), np.array(indices), np.array(indptr)
+
+    @d.dedent
     def bias_knn(
-        self, conn: csr_matrix, pseudotime: np.ndarray, **kwargs: Any
+        self,
+        conn: csr_matrix,
+        pseudotime: np.ndarray,
+        n_jobs: Optional[int] = None,
+        backend: str = "loky",
+        show_progress_bar: bool = True,
+        **kwargs: Any,
     ) -> csr_matrix:
         """
         Bias cell-cell connectivities of a KNN graph.
@@ -49,27 +92,29 @@ def bias_knn(
             Sparse matrix of shape ``(n_cells, n_cells)`` containing the nearest neighbor connectivities.
         pseudotime
             Pseudotemporal ordering of cells.
+        %(parallel)s
 
         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
+        res = parallelize(
+            self._bias_knn_helper,
+            np.arange(conn.shape[0]),
+            as_array=False,
+            unit="cell",
+            n_jobs=n_jobs,
+            backend=backend,
+            show_progress_bar=show_progress_bar,
+        )(conn, pseudotime, **kwargs)
+        data, indices, indptr = zip(*res)
+
+        conn = csr_matrix(
+            (np.concatenate(data), np.concatenate(indices), np.concatenate(indptr))
+        )
+        conn.eliminate_zeros()
+
+        return conn
 
     def __repr__(self):
         return f"<{self.__class__.__name__}>"
@@ -80,13 +125,13 @@ def __str__(self):
 
 class HardThresholdScheme(ThresholdSchemeABC):
     """
-    Thresholding scheme inspired by Palantir [Setty19]_.
+    Thresholding scheme inspired by *Palantir* :cite:`setty:19`.
 
-    Note that this won't exactly reproduce the original Palantir results, for three reasons:
+    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
+        - *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".
     """
 
@@ -119,9 +164,7 @@ def __call__(
         -------
         %(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
+        k_thresh = max(0, min(30, int(np.floor(n_neighs * frac_to_keep))))
         ixs = np.flip(np.argsort(neigh_conn))
         close_ixs, far_ixs = ixs[:k_thresh], ixs[k_thresh:]
 
@@ -137,7 +180,7 @@ def __call__(
 
 class SoftThresholdScheme(ThresholdSchemeABC):
     """
-    Thresholding scheme inspired by [VIA21]_.
+    Thresholding scheme inspired by :cite:`stassen:21`.
 
     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
@@ -150,9 +193,8 @@ def __call__(
         cell_pseudotime: float,
         neigh_pseudotime: np.ndarray,
         neigh_conn: np.ndarray,
-        b: float = 20.0,
-        nu: float = 1.0,
-        perc: Optional[int] = 95,
+        b: float = 10.0,
+        nu: float = 0.5,
     ) -> np.ndarray:
         """
         Bias the connectivities by downweighting ones to past cells.
@@ -169,13 +211,6 @@ def __call__(
         -------
         %(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
diff --git a/cellrank/tl/kernels/_random_walk.py b/cellrank/tl/kernels/_random_walk.py
index fdcd287fb..baa8838bd 100644
--- a/cellrank/tl/kernels/_random_walk.py
+++ b/cellrank/tl/kernels/_random_walk.py
@@ -1,13 +1,13 @@
 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
 
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+
 
 class RandomWalk:
     """
@@ -136,7 +136,8 @@ def _simulate_many(
             if queue is not None:
                 queue.put(1)
 
-        queue.put(None)
+        if queue is not None:
+            queue.put(None)
 
         return res
 
diff --git a/cellrank/tl/kernels/_tmat_flow.py b/cellrank/tl/kernels/_tmat_flow.py
new file mode 100644
index 000000000..eaf6d2715
--- /dev/null
+++ b/cellrank/tl/kernels/_tmat_flow.py
@@ -0,0 +1,575 @@
+from typing import Any, List, Tuple, Union, Mapping, Optional, Sequence
+from functools import lru_cache
+from dataclasses import dataclass
+
+from statsmodels.nonparametric.smoothers_lowess import lowess
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import _unique_order_preserving
+from cellrank.tl._colors import _create_categorical_colors
+from cellrank.tl.kernels._utils import _ensure_numeric_ordered
+
+import numpy as np
+import pandas as pd
+from scipy.stats import logistic
+from scipy.sparse import issparse, spmatrix
+from pandas.api.types import infer_dtype
+from scipy.interpolate import interp1d
+from pandas.core.dtypes.common import is_categorical_dtype
+
+import matplotlib.pyplot as plt
+from matplotlib.colors import to_rgb
+from matplotlib.collections import PolyCollection
+
+Numeric_t = Union[float, int]
+
+
+@dataclass(frozen=True)
+class Point:  # noqa: D101
+    x: float
+    xt: float
+
+
+@d.dedent
+class FlowPlotter:
+    """
+    Class that plots outgoing flow for a specific cluster :cite:`mittnenzweig:21`.
+
+    It should be able to recreate (to a high degree) figures such as Fig. 4a in the above mentioned paper.
+
+    Parameters
+    ----------
+    %(adata)s
+    tmat
+        Matrix of shape ``(adata.n_obs, adata.n_obs)``.
+    cluster_key
+        Key in :attr:`adata` ``.obs`` where clustering is stored.
+    time_key
+        Key in :attr:`adata` ``.obs`` where experimental time is stored.
+    """
+
+    TIME_KEY = "time"
+
+    def __init__(
+        self,
+        adata: AnnData,
+        tmat: Union[np.ndarray, spmatrix],
+        cluster_key: str,
+        time_key: str,
+    ):
+        self._adata = adata
+        self._tmat = tmat
+        self._ckey = cluster_key
+        self._tkey = time_key
+
+        self._cluster: Optional[str] = None
+        self._clusters: Optional[Sequence[Any]] = None
+
+        self._flow: Optional[pd.DataFrame] = None
+        self._cmat: Optional[pd.DataFrame] = None
+
+        if self._ckey not in self._adata.obs:
+            raise KeyError(f"Unable to find clusters in `adata.obs[{self._ckey!r}]`.")
+        if not is_categorical_dtype(self._adata.obs[self._ckey]):
+            raise TypeError(
+                f"Expected `adata.obs[{self._ckey!r}]` to be categorical, "
+                f"found `{infer_dtype(self._adata.obs[self._ckey])}`."
+            )
+        self._adata.obs[self._tkey] = _ensure_numeric_ordered(self._adata, self._tkey)
+
+    def prepare(
+        self,
+        cluster: str,
+        clusters: Optional[Sequence[Any]] = None,
+        time_points: Optional[Sequence[Numeric_t]] = None,
+    ) -> "FlowPlotter":
+        """
+        Prepare itself for plotting by computing flow and contingency matrix.
+
+        Parameters
+        ----------
+        cluster
+            Source cluster for flow calculation.
+        clusters
+            Target clusters for flow calculation. If `None`, use all clusters.
+        time_points
+            Restrict flow calculation only to these time points. If `None`, use all time points.
+
+        Returns
+        -------
+        Modifies and return self.
+        """
+        if clusters is None:
+            self._clusters = self.clusters.cat.categories
+        else:
+            clusters = _unique_order_preserving([cluster] + list(clusters))
+            mask = self.clusters.isin(clusters).values
+
+            self._adata = self._adata[mask]
+            if not self._adata.n_obs:
+                raise ValueError("No valid clusters have been selected.")
+            self._tmat = self._tmat[mask, :][:, mask]
+            self._clusters = [c for c in clusters if c in self.clusters.cat.categories]
+
+        if cluster not in self._clusters:
+            raise ValueError(f"Invalid source cluster `{cluster!r}`.")
+
+        if len(self._clusters) < 2:
+            raise ValueError(
+                f"Expected at least `2` clusters, found `{len(clusters)}`."
+            )
+
+        if time_points is not None:
+            time_points = _unique_order_preserving(time_points)
+            if len(time_points) < 2:
+                raise ValueError(
+                    f"Expected at least `2` time points, found `{len(time_points)}`."
+                )
+
+            mask = self.time.isin(time_points)
+
+            self._adata = self._adata[mask]
+            if not self._adata.n_obs:
+                raise ValueError("No valid time points have been selected.")
+            self._tmat = self._tmat[mask, :][:, mask]
+
+        time_points = list(
+            zip(self.time.cat.categories[:-1], self.time.cat.categories[1:])
+        )
+
+        logg.info(
+            f"Computing flow from `{cluster}` into `{len(self._clusters) - 1}` cluster(s) "
+            f"in `{len(time_points)}` time points"
+        )
+        self._cluster = cluster
+        self._cmat = self.compute_contingency_matrix()
+        self._flow = self.compute_flow(time_points, cluster)
+
+        return self
+
+    def compute_flow(
+        self,
+        time_points: Sequence[Tuple[Numeric_t, Numeric_t]],
+        cluster: Optional[str] = None,
+    ) -> pd.DataFrame:
+        """
+        Compute outgoing flow.
+
+        Parameters
+        ----------
+        time_points
+            Time point pair for which to calculate the flow.
+        cluster
+            Cluster for which to calculate the outgoing flow. If `None`, calculate the flow for all clusters.
+
+        Returns
+        -------
+        Dataframe of shape ``(n_time_points, n_clusters)`` if ``cluster!=None`` or
+        a dataframe of shape ``(n_time_points * n_clusters, n_clusters)`` otherwise.
+        The dataframe's index is a multi-index and the 1st level corresponds to time, the 2nd level to source clusters.
+        """
+
+        def default_helper(t1: Numeric_t, t2: Numeric_t) -> pd.DataFrame:
+            subset, row_cls, col_cls = self._get_time_subset(t1, t2)
+
+            df = pd.DataFrame(subset.A if issparse(subset) else subset)
+            df = df.groupby(row_cls).sum().T.groupby(col_cls).sum().T
+
+            res = pd.DataFrame(np.zeros((n, n)), index=categories, columns=categories)
+            res.loc[df.index, df.columns] = df
+            res.fillna(0, inplace=True)
+
+            return res
+
+        def cluster_helper(t1: Numeric_t, t2: Numeric_t) -> pd.DataFrame:
+            subset, row_cls, col_cls = self._get_time_subset(t1, t2, cluster=cluster)
+
+            df = pd.DataFrame(subset.A if issparse(subset) else subset).sum(0)
+            df = df.groupby(col_cls).sum()
+            df = pd.DataFrame([df], index=[cluster], columns=df.index)
+
+            res = pd.DataFrame(np.zeros((1, n)), index=[cluster], columns=categories)
+            res.loc[df.index, df.columns] = df
+            res.fillna(0, inplace=True)
+
+            return res
+
+        categories = self.clusters.cat.categories
+        n = len(categories)
+        callback = cluster_helper if cluster is not None else default_helper
+        flows, times = [], []
+
+        for t1, t2 in time_points:
+            flow = callback(t1, t2)
+            times.extend([t1] * len(flow))
+            flows.append(flow)
+
+        flow = pd.concat(flows)
+        flow.set_index([times, flow.index], inplace=True)
+        flow /= flow.sum(1)[:, None]
+        flow.fillna(0, inplace=True)
+
+        return flow
+
+    def compute_contingency_matrix(self) -> pd.DataFrame:
+        """Row-normalized contingency matrix of shape ``(n_clusters, n_time_points)``."""
+        cmat = pd.crosstab(self.clusters, self.time)
+        return (cmat / cmat.sum(0)[None, :]).fillna(0)
+
+    @d.get_sections(base="flow", sections=["Parameters"])
+    def plot(
+        self,
+        min_flow: float = 0,
+        remove_empty_clusters: bool = True,
+        ascending: Optional[bool] = False,
+        alpha: float = 0.8,
+        xticks_step_size: Optional[int] = 1,
+        legend_loc: Optional[str] = "upper right out",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+    ) -> plt.Axes:
+        """
+        Plot outgoing flow.
+
+        Parameters
+        ----------
+        min_flow
+            Only show flow edges with flow greater than this value. Flow values are always in `[0, 1]`.
+        remove_empty_clusters
+            Whether to remove clusters with no incoming flow edges.
+        ascending
+            Whether to sort the cluster by ascending or descending incoming flow.
+            If `None`, use the order as in defined by ``clusters``.
+        alpha
+            Alpha value for cell proportions.
+        xticks_step_size
+            Show only every n-th ticks on x-axis. If `None`, don't show any ticks.
+        legend_loc
+            Position of the legend. If `None`, do not show the legend.
+
+        Returns
+        -------
+        The axes object.
+        """
+        if self._flow is None or self._cmat is None:
+            raise RuntimeError(
+                "Compute flow and contingency matrix first as `.prepare()`."
+            )
+
+        flow, cmat = self._flow, self._cmat
+        try:
+            if remove_empty_clusters:
+                self._remove_min_clusters(min_flow)
+            logg.info(
+                f"Plotting flow from `{self._cluster}` into `{len(self._flow.columns) - 1}` cluster(s) "
+                f"in `{len(self._cmat.columns) - 1}` time points"
+            )
+            return self._plot(
+                self._rename_times(),
+                ascending=ascending,
+                min_flow=min_flow,
+                alpha=alpha,
+                xticks_step_size=xticks_step_size,
+                legend_loc=legend_loc,
+                figsize=figsize,
+                dpi=dpi,
+            )
+        finally:
+            self._flow = flow
+            self._cmat = cmat
+
+    def _get_time_subset(
+        self, t1: Numeric_t, t2: Numeric_t, cluster: Optional[str] = None
+    ) -> Tuple[Union[np.ndarray, spmatrix], pd.Series, pd.Series]:
+        if cluster is None:
+            row_ixs = np.where(self.time == t1)[0]
+        else:
+            row_ixs = np.where((self.time == t1) & (self.clusters == cluster))[0]
+
+        col_ixs = np.where(self.time == t2)[0]
+        row_cls = self.clusters.values[row_ixs]
+        col_cls = self.clusters.values[col_ixs]
+
+        return self._tmat[row_ixs, :][:, col_ixs], row_cls, col_cls
+
+    def _remove_min_clusters(self, min_flow: float) -> None:
+        logg.debug("Removing clusters with no incoming flow edges")
+        columns = (self._flow.loc[(slice(None), self._cluster), :] > min_flow).any()
+        columns = columns[columns].index
+        if not len(columns):
+            raise ValueError(
+                "After removing clusters with no incoming flow edges, none remain."
+            )
+        self._flow = self._flow[columns]
+
+    def _rename_times(self) -> Sequence[Numeric_t]:
+        # make sure we have enough horizontal space to draw the flow (i.e. time points are at least 1 unit apart)
+        old_times = self._cmat.columns
+        tmp = np.array(old_times)
+        tmp = (tmp - tmp.min()) / (tmp.max() - tmp.min())
+        tmp /= np.min(tmp[1:] - tmp[:-1])
+        time_mapper = dict(zip(old_times, tmp))
+        self._flow.index = pd.MultiIndex.from_tuples(
+            [(time_mapper[t], c) for t, c in self._flow.index]
+        )
+        self._cmat.columns = tmp
+        return old_times
+
+    def _order_clusters(
+        self, cluster: str, ascending: Optional[bool] = False
+    ) -> Tuple[List[Any], List[Any]]:
+        if ascending is not None:
+            tmp = [[], []]
+            total_flow = (
+                self._flow.loc[(slice(None), cluster), :]
+                .sum()
+                .sort_values(ascending=ascending)
+            )
+            for i, c in enumerate(c for c in total_flow.index if c != cluster):
+                tmp[i % 2].append(c)
+            return tmp[0][::-1], tmp[1]
+
+        clusters = [c for c in self._clusters if c != cluster]
+        return clusters[: len(clusters) // 2], clusters[len(clusters) // 2 :]
+
+    def _calculate_y_offsets(
+        self, clusters: Sequence[Any], delta: float = 0.2
+    ) -> Mapping[Any, float]:
+        offset = [0]
+        for i in range(1, len(clusters)):
+            offset.append(
+                offset[-1]
+                + delta
+                + np.max(self._cmat.loc[clusters[i]] + self._cmat.loc[clusters[i - 1]])
+            )
+        return dict(zip(clusters, offset))
+
+    def _plot_smoothed_proportion(
+        self,
+        ax: plt.Axes,
+        clusters: Sequence[Any],
+        y_offset: Mapping[Any, float],
+        alpha: float = 0.8,
+    ) -> Tuple[Mapping[Any, np.ndarray], Mapping[Any, PolyCollection]]:
+        start_t, end_t = self._cmat.columns.min(), self._cmat.columns.max()
+        x = np.array(self._cmat.columns)  # fitting
+        # extrapolation
+        e = np.linspace(start_t, end_t, int(1 + (end_t - start_t) * 100))
+
+        smoothed_proportion, handles = {}, {}
+        for clust in clusters:
+            y = self._cmat.loc[clust]
+            f = interp1d(x, y)
+            fe = f(e)
+            lo = lowess(fe, e, frac=0.3, is_sorted=True, return_sorted=False)
+            smoothed_proportion[clust] = lo
+
+            handles[clust] = ax.fill_between(
+                e,
+                y_offset[clust] + lo,
+                y_offset[clust] - lo,
+                color=self.cmap[clust],
+                label=clust,
+                alpha=alpha,
+                edgecolor=None,
+            )
+
+        return smoothed_proportion, handles
+
+    def _draw_flow_edge(
+        self,
+        ax,
+        x1: Point,
+        x2: Point,
+        y1: Point,
+        y2: Point,
+        start_color: Tuple[float, float, float],
+        end_color: Tuple[float, float, float],
+        flow: float,
+        alpha: float = 0.8,
+    ) -> None:
+        # transcribed from: https://github.com/tanaylab/embflow/blob/main/scripts/generate_paper_figures/plot_vein.r
+        dx = x2.xt - x1.x
+        dy = y2.xt - y1.x
+        dxt = x2.x - x1.x
+        dyt = y2.x - y1.xt
+
+        start_color = np.asarray(to_rgb(start_color))
+        end_color = np.asarray(to_rgb(end_color))
+        delta = 0.05
+
+        beta0 = _lcdf(0)
+        beta_f = _lcdf(1) - _lcdf(0)
+
+        rs = np.arange(0, 1, delta)
+        beta = (_lcdf(rs) - beta0) / beta_f
+        beta5 = (_lcdf(rs + delta) - beta0) / beta_f
+
+        sx1 = x1.x + rs * dx
+        sy1 = y1.x + beta * dy
+        sx2 = x1.x + (rs + delta) * dx
+        sy2 = y1.x + beta5 * dy
+
+        sx1t = x1.x + flow + rs * dxt
+        sy1t = y1.xt + beta * dyt
+        sx2t = x1.x + flow + (rs + delta) * dxt
+        sy2t = y1.xt + beta5 * dyt
+
+        xs = np.c_[sx1, sx2, sx2t, sx1t]
+        ys = np.c_[sy1, sy2, sy2t, sy1t]
+
+        start_alpha, end_alpha = 0.2, alpha
+        if start_alpha > end_alpha:
+            start_alpha, end_alpha = end_alpha, start_alpha
+        col = np.c_[
+            (start_color * (1 - rs[:, None])) + (end_color * rs[:, None]),
+            np.linspace(start_alpha, end_alpha, len(rs)),
+        ]
+
+        for x, y, c in zip(xs, ys, col):
+            ax.fill(x, y, c=c, edgecolor=None)
+
+    def _plot(
+        self,
+        old_times: Sequence[Numeric_t],
+        ascending: Optional[bool],
+        min_flow: float = 0,
+        alpha: float = 0.8,
+        xticks_step_size: Optional[int] = 1,
+        legend_loc: Optional[str] = "upper right out",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+    ) -> plt.Axes:
+        from cellrank.pl._utils import _position_legend
+
+        def r(num: float) -> int:
+            return max(0, int(round(num, 2) * 100) - 1)
+
+        def draw_edges(
+            curr_t: Numeric_t,
+            next_t: Numeric_t,
+            clusters: Sequence[Any],
+            *,
+            bottom: bool,
+        ):
+            smooth_cluster = float(smoothed_proportions[self._cluster][r(curr_t)])
+            flow = self._flow.loc[curr_t]
+            for clust in clusters:
+                fl = flow.loc[self._cluster, clust]
+                if fl > min_flow:
+                    fl = np.clip(fl, 0, 0.95)
+                    smooth_cluster_fl = smoothed_proportions[self._cluster][
+                        r(curr_t + fl)
+                    ]
+
+                    if bottom:
+                        self._draw_flow_edge(
+                            ax,
+                            x1=Point(curr_t, 0),
+                            x2=Point(next_t - fl, next_t - fl - 0.05),
+                            y1=Point(
+                                cluster_offset - smooth_cluster,
+                                cluster_offset - smooth_cluster_fl,
+                            ),
+                            y2=Point(
+                                y_offset[clust]
+                                + smoothed_proportions[clust][r(next_t)],
+                                y_offset[clust]
+                                + smoothed_proportions[clust][r(next_t - fl - 0.05)],
+                            ),
+                            flow=fl,
+                            start_color=self.cmap[self._cluster],
+                            end_color=self.cmap[clust],
+                            alpha=alpha,
+                        )
+                    else:
+                        self._draw_flow_edge(
+                            ax,
+                            x1=Point(curr_t + fl, 0),
+                            x2=Point(next_t - 0.05, next_t),
+                            y1=Point(
+                                cluster_offset + smooth_cluster_fl,
+                                cluster_offset + smooth_cluster,
+                            ),
+                            y2=Point(
+                                y_offset[clust]
+                                - smoothed_proportions[clust][r(next_t - fl - 0.05)],
+                                y_offset[clust]
+                                - smoothed_proportions[clust][r(next_t)],
+                            ),
+                            flow=-fl,
+                            start_color=self.cmap[self._cluster],
+                            end_color=self.cmap[clust],
+                            alpha=alpha,
+                        )
+
+        if xticks_step_size is not None:
+            xticks_step_size = max(1, xticks_step_size)
+        times = self._cmat.columns
+        fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
+
+        clusters_bottom, clusters_top = self._order_clusters(self._cluster, ascending)
+        all_clusters = clusters_bottom + [self._cluster] + clusters_top
+
+        y_offset = self._calculate_y_offsets(all_clusters)
+        cluster_offset = y_offset[self._cluster]
+
+        smoothed_proportions, handles = self._plot_smoothed_proportion(
+            ax, all_clusters, y_offset, alpha=alpha
+        )
+
+        for curr_t, next_t in zip(times[:-1], times[1:]):
+            draw_edges(curr_t, next_t, clusters_bottom, bottom=True)
+            draw_edges(curr_t, next_t, clusters_top, bottom=False)
+
+        ax.margins(0.025)
+        ax.set_title(self._cluster)
+        ax.set_xlabel(self._tkey)
+        ax.set_ylabel(self._ckey)
+        if xticks_step_size is None:
+            ax.set_xticks([])
+        else:
+            ax.set_xticks(times[::xticks_step_size])
+            ax.set_xticklabels(old_times[::xticks_step_size])
+        ax.set_yticks([])
+
+        if legend_loc not in (None, "none"):
+            _position_legend(
+                ax,
+                legend_loc=legend_loc,
+                handles=[handles[c] for c in all_clusters[::-1]],
+            )
+
+        return ax
+
+    @property
+    def clusters(self) -> pd.Series:
+        """Clusters."""
+        return self._adata.obs[self._ckey]
+
+    @property
+    def time(self) -> pd.Series:
+        """Time points."""
+        return self._adata.obs[self._tkey]
+
+    @property
+    @lru_cache(1)
+    def cmap(self) -> Mapping[str, Any]:
+        """Colormap for :attr:`clusters`."""
+        return dict(
+            zip(
+                self.clusters.cat.categories,
+                self._adata.uns.get(
+                    f"{self._ckey}_colors",
+                    _create_categorical_colors(len(self.clusters.cat.categories)),
+                ),
+            )
+        )
+
+
+def _lcdf(
+    x: Union[int, float, np.ndarray], loc: float = 0.5, scale: float = 0.2
+) -> float:
+    return logistic.cdf(x, loc=loc, scale=scale)
diff --git a/cellrank/tl/kernels/_utils.py b/cellrank/tl/kernels/_utils.py
index adebfb751..35af477ea 100644
--- a/cellrank/tl/kernels/_utils.py
+++ b/cellrank/tl/kernels/_utils.py
@@ -3,10 +3,14 @@
 from inspect import signature
 
 from anndata import AnnData
+from cellrank import logging as logg
 
 import numpy as np
+import pandas as pd
 from numba import njit, prange
 from scipy.sparse import csr_matrix
+from pandas.api.types import infer_dtype
+from pandas.core.dtypes.common import is_numeric_dtype, is_categorical_dtype
 
 jit_kwargs = {"nogil": True, "cache": True, "fastmath": True}
 
@@ -247,3 +251,48 @@ def _get_basis(adata: AnnData, basis: str) -> np.ndarray:
             raise KeyError(
                 f"Unable to find a basis in `adata.obsm['X_{basis}']` or `adata.obsm[{basis!r}]`."
             ) from None
+
+
+def _ensure_numeric_ordered(adata: AnnData, key: str) -> pd.Series:
+    if key not in adata.obs.keys():
+        raise KeyError(f"Unable to find data in `adata.obs[{key!r}]`.")
+
+    exp_time = adata.obs[key].copy()
+    if not is_numeric_dtype(np.asarray(exp_time)):
+        try:
+            exp_time = np.asarray(exp_time).astype(float)
+        except ValueError as e:
+            raise TypeError(
+                f"Unable to convert `adata.obs[{key!r}]` of type `{infer_dtype(adata.obs[key])}` to `float`."
+            ) from e
+
+    if not is_categorical_dtype(exp_time):
+        logg.debug(f"Converting `adata.obs[{key!r}]` to `categorical`")
+        exp_time = np.asarray(exp_time)
+        categories = sorted(set(exp_time[~np.isnan(exp_time)]))
+        if len(categories) > 100:
+            raise ValueError(
+                f"Unable to convert `adata.obs[{key!r}]` to `categorical` since it "
+                f"would create `{len(categories)}` categories."
+            )
+        exp_time = pd.Series(
+            pd.Categorical(
+                exp_time,
+                categories=categories,
+                ordered=True,
+            )
+        )
+
+    if not exp_time.cat.ordered:
+        logg.warning("Categories are not ordered. Using ascending order")
+        exp_time.cat = exp_time.cat.as_ordered()
+
+    exp_time = pd.Series(pd.Categorical(exp_time, ordered=True), index=adata.obs_names)
+    if exp_time.isnull().any():
+        raise ValueError("Series contains NaN value(s).")
+
+    n_cats = len(exp_time.cat.categories)
+    if n_cats < 2:
+        raise ValueError(f"Expected to find at least `2` categories, found `{n_cats}`.")
+
+    return exp_time
diff --git a/cellrank/tl/kernels/_velocity_kernel.py b/cellrank/tl/kernels/_velocity_kernel.py
index e4a71a419..eea484195 100644
--- a/cellrank/tl/kernels/_velocity_kernel.py
+++ b/cellrank/tl/kernels/_velocity_kernel.py
@@ -5,11 +5,6 @@
 from typing import Any, Union, Callable, Iterable, Optional
 
 from anndata import AnnData
-from scvelo.preprocessing.moments import get_moments
-
-import numpy as np
-from scipy.sparse import issparse, csr_matrix
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.ul._utils import valuedispatch
@@ -25,9 +20,13 @@
     _calculate_starts,
     _get_probs_for_zero_vec,
 )
+from scvelo.preprocessing.moments import get_moments
 from cellrank.tl.kernels._base_kernel import _RTOL
 from cellrank.tl.kernels._velocity_schemes import Scheme, _get_scheme
 
+import numpy as np
+from scipy.sparse import issparse, csr_matrix
+
 
 class VelocityMode(ModeEnum):  # noqa
     DETERMINISTIC = "deterministic"
@@ -46,10 +45,10 @@ class VelocityKernel(Kernel):
     """
     Kernel which computes a transition matrix based on RNA velocity.
 
-    This borrows ideas from both [Manno18]_ and [Bergen20]_. In short, for each cell *i*, we compute transition
-    probabilities :math:`p_{i, j}` to each cell *j* in the neighborhood of *i*. The transition probabilities are
-    computed as a multinomial logistic regression where the weights :math:`w_j` (for all *j*) are given by the vector
-    that connects cell *i* with cell *j* in gene expression space, and the features :math:`x_i` are given
+    This borrows ideas from both :cite:`manno:18` and :cite:`bergen:20`. In short, for each cell *i*, we compute
+    transition probabilities :math:`p_{i, j}` to each cell *j* in the neighborhood of *i*. The transition probabilities
+    are computed as a multinomial logistic regression where the weights :math:`w_j` (for all *j*) are given
+    by the vector that connects cell *i* with cell *j* in gene expression space, and the features :math:`x_i` are given
     by the velocity vector :math:`v_i` of cell *i*.
 
     Parameters
@@ -57,15 +56,17 @@ class VelocityKernel(Kernel):
     %(adata)s
     %(backward)s
     vkey
-        Key in :paramref:`adata` ``.uns`` where the velocities are stored.
+        Key in :attr:`adata` ``.uns`` where the velocities are stored.
     xkey
-        Key in :paramref:`adata` ``.layers`` where expected gene expression counts are stored.
+        Key in :attr:`adata` ``.layers`` where expected gene expression counts are stored.
     gene_subset
         List of genes to be used to compute transition probabilities.
-        By default, genes from :paramref:`adata` ``.var['velocity_genes']`` are used.
+        By default, genes from :attr:`adata` ``.var['velocity_genes']`` are used.
     %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
+    kwargs
+        Keyword arguments for :class:`cellrank.tl.kernels.Kernel`.
     """
 
     def __init__(
@@ -77,6 +78,7 @@ def __init__(
         gene_subset: Optional[Iterable] = None,
         compute_cond_num: bool = False,
         check_connectivity: bool = False,
+        **kwargs: Any,
     ):
         super().__init__(
             adata,
@@ -86,13 +88,14 @@ def __init__(
             gene_subset=gene_subset,
             compute_cond_num=compute_cond_num,
             check_connectivity=check_connectivity,
+            **kwargs,
         )
         self._vkey = vkey  # for copy
         self._xkey = xkey
         self._gene_subset = gene_subset
         self._logits = None
 
-    def _read_from_adata(self, **kwargs):
+    def _read_from_adata(self, **kwargs: Any) -> None:
         super()._read_from_adata(**kwargs)
 
         # check whether velocities have been computed
@@ -184,8 +187,8 @@ def compute_transition_matrix(
         :class:`cellrank.tl.kernels.VelocityKernel`
             Makes available the following fields:
 
-                - :paramref:`transition_matrix`.
-                - :paramref:`logits`.
+                - :attr:`transition_matrix`.
+                - :attr:`logits`.
         """
         mode = VelocityMode(mode)
         backward_mode = BackwardMode(backward_mode)
diff --git a/cellrank/tl/kernels/_velocity_schemes.py b/cellrank/tl/kernels/_velocity_schemes.py
index d50874449..ad313dbe0 100644
--- a/cellrank/tl/kernels/_velocity_schemes.py
+++ b/cellrank/tl/kernels/_velocity_schemes.py
@@ -2,14 +2,14 @@
 from typing import Any, Tuple
 from functools import partial
 
-import numpy as np
-from numba import njit
-
 from cellrank.ul._docs import d
 from cellrank.ul._utils import valuedispatch
 from cellrank.tl._constants import ModeEnum
 from cellrank.tl.kernels._utils import norm, np_mean, jit_kwargs
 
+import numpy as np
+from numba import njit
+
 
 class Scheme(ModeEnum):  # noqa: D101
     DOT_PRODUCT = "dot_product"
@@ -212,7 +212,7 @@ def __str__(self):
 
 class SimilaritySchemeHessian(SimilaritySchemeABC, Hessian):
     """
-    Base class for all similarity schemes as defined in [Li2020]_.
+    Base class for all similarity schemes as defined in :cite:`li:20`.
 
     Parameters
     ----------
@@ -274,7 +274,7 @@ def __getattribute__(self, name: str) -> Any:
 
 class DotProductScheme(SimilaritySchemeHessian):
     r"""
-    Dot product scheme as defined in eq. (4.9) of [Li2020]_.
+    Dot product scheme as defined in eq. (4.9) :cite:`li:20`.
 
         :math:`v(s_i, s_j) = g(\delta_{i, j}^T v_i)`
 
@@ -288,7 +288,7 @@ def __init__(self):
 
 class CosineScheme(SimilaritySchemeHessian):
     r"""
-    Cosine similarity scheme as defined in eq. (4.7) of [Li2020]_.
+    Cosine similarity scheme as defined in eq. (4.7) :cite:`li:20`.
 
         :math:`v(s_i, s_j) = g(cos(\delta_{i, j}, v_i))`
 
@@ -302,7 +302,7 @@ def __init__(self):
 
 class CorrelationScheme(SimilaritySchemeHessian):
     r"""
-    Pearson correlation scheme as defined in eq. (4.8) of [Li2020]_.
+    Pearson correlation scheme as defined in eq. (4.8) :cite:`li:20`.
 
         :math:`v(s_i, s_j) = g(corr(\delta_{i, j}, v_i))`
 
diff --git a/cellrank/ul/_docs.py b/cellrank/ul/_docs.py
index 7326cdc02..5927257dd 100644
--- a/cellrank/ul/_docs.py
+++ b/cellrank/ul/_docs.py
@@ -31,7 +31,7 @@
     Model based on :class:`cellrank.ul.models.BaseModel` to fit.
 
     If a :class:`dict`, gene and lineage specific models can be specified. Use ``'*'`` to indicate
-    all genes or lineages, for example ``{'Map2': {'*': ...}, 'Dcx': {'Alpha': ..., '*': ...}}``."""
+    all genes or lineages, for example ``{'gene_1': {'*': ...}, 'gene_2': {'lineage_1': ..., '*': ...}}``."""
 _just_plots = """\
 None
     Nothing, just plots the figure. Optionally saves it based on ``save``."""
@@ -57,15 +57,15 @@
     Number of lineages. If `None`, it will be determined automatically.
 cluster_key
     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.
+    For this, provide a key from :attr:`adata` ``.obs`` where cluster labels have been computed.
 keys
     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."""
 _density_correction = (
-    "Optionally, we apply a density correction as described in [Coifman05]_, "
-    "where we use the implementation of [Haghverdi16]_."
+    "Optionally, we apply a density correction as described in :cite:`coifman:05`, "
+    "where we use the implementation of :cite:`haghverdi:16`."
 )
 _time_range = """\
 time_range
@@ -74,7 +74,7 @@
         - If a :class:`tuple`, it specifies the minimum and maximum pseudotime. Both values can be `None`,
           in which case the minimum is the earliest pseudotime and the maximum is automatically determined.
         - If a :class:`float`, it specifies the maximum pseudotime."""
-_time_ranges = f"{_time_range}\n    This can also be specified on per-lineage basis."
+
 _velocity_mode = """\
 mode
     How to compute transition probabilities. Valid options are:
@@ -85,7 +85,7 @@
         - `{m.SAMPLING.s!r}` - sample 1 transition matrix from the velocity distribution."""
 _velocity_backward_mode = """\
 backward_mode
-    Only matters if initialized as :paramref:`backward` `=True`.  Valid options are:
+    Only matters if initialized as :attr:`backward` ``=True``.  Valid options are:
 
         - `{b.TRANSPOSE.s!r}` - compute transitions from neighboring cells `j` to cell `i`.
         - `{b.NEGATE.s!r}` - negate the velocity vector."""
@@ -102,7 +102,7 @@
 callback
     Function which takes a :class:`cellrank.ul.models.BaseModel` and some keyword arguments
     for :meth:`cellrank.ul.models.BaseModel.prepare` and returns the prepared model.
-    Can be specified in gene- and lineage-specific manner, similarly to :paramref:`model`."""
+    Can be specified in gene- and lineage-specific manner, similarly to :attr:`model`."""
 _genes = """\
 genes
     Genes in ``adata.var_names`` or in ``adata.raw.var_names``, if ``use_raw=True``."""
@@ -117,7 +117,7 @@
         - `'embedding'` - plot the embedding while coloring in the absorption probabilities.
         - `'time'` - plot the pseudotime on x-axis and the absorption probabilities on y-axis."""
 _write_to_adata = """\
-Updates the :paramref:`adata` with the following fields:
+Updates the :attr:`adata` with the following fields:
 
         - ``.obsp['{{key}}']`` - the transition matrix.
         - ``.uns['{{key}}_params']`` - parameters used for calculation."""
@@ -136,7 +136,7 @@
     Basis to use when ``mode='embedding'``. If `None`, use `'umap'`."""
 _velocity_scheme = """\
 scheme
-    Similarity scheme between cells as described in [Li2020]_. Can be one of the following:
+    Similarity scheme between cells as described in :cite:`li:20`. 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`.
@@ -152,9 +152,7 @@
 b
     The growth rate of generalized logistic function.{}
 nu
-    Affects near which asymptote maximum growth occurs.{}
-perc
-    Percentile by which to clip the connectivities.{}"""
+    Affects near which asymptote maximum growth occurs.{}"""
 _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
@@ -194,7 +192,6 @@ def decorator2(obj):
     copy=_copy,
     density_correction=_density_correction,
     time_range=_time_range,
-    time_ranges=_time_ranges,
     velocity_mode=_velocity_mode,
     velocity_backward_mode=_velocity_backward_mode,
     velocity_backward_mode_high_lvl=_velocity_backward_mode_high_lvl,
diff --git a/cellrank/ul/_parallelize.py b/cellrank/ul/_parallelize.py
index f3ee94740..58382a538 100644
--- a/cellrank/ul/_parallelize.py
+++ b/cellrank/ul/_parallelize.py
@@ -6,11 +6,11 @@
 
 import joblib as jl
 
+from cellrank.ul._utils import _get_n_cores
+
 import numpy as np
 from scipy.sparse import issparse, spmatrix
 
-from cellrank.ul._utils import _get_n_cores
-
 
 def parallelize(
     callback: Callable[[Any], Any],
@@ -20,7 +20,7 @@ def parallelize(
     unit: str = "",
     as_array: bool = True,
     use_ixs: bool = False,
-    backend: str = "multiprocessing",
+    backend: str = "loky",
     extractor: Optional[Callable[[Any], Any]] = None,
     show_progress_bar: bool = True,
 ) -> Any:
@@ -52,11 +52,18 @@ def parallelize(
 
     Returns
     -------
-        The result depending on ``callable``, ``extractor`` and ``as_array``.
+    The result depending on ``callable``, ``extractor`` and ``as_array``.
     """
 
     if show_progress_bar:
-        from tqdm.auto import tqdm
+        try:
+            import ipywidgets  # noqa: F401
+            from tqdm.auto import tqdm
+        except ImportError:
+            try:
+                from tqdm.std import tqdm
+            except ImportError:
+                tqdm = None
     else:
         tqdm = None
 
@@ -68,7 +75,7 @@ def update(pbar, queue, n_total):
             except EOFError as e:
                 if not n_finished != n_total:
                     raise RuntimeError(
-                        f"Finished only `{n_finished} out of `{n_total}` tasks.`"
+                        f"Finished only `{n_finished}` out of `{n_total}` tasks.`"
                     ) from e
                 break
             assert res in (None, (1, None), 1)  # (None, 1) means only 1 job
@@ -119,11 +126,11 @@ def wrapper(*args, **kwargs):
         n_split = n_jobs
 
     if issparse(collection):
+        n_split = max(1, min(n_split, collection.shape[0]))
         if n_split == collection.shape[0]:
             collections = [collection[[ix], :] for ix in range(collection.shape[0])]
         else:
             step = collection.shape[0] // n_split
-
             ixs = [
                 np.arange(i * step, min((i + 1) * step, collection.shape[0]))
                 for i in range(n_split)
@@ -136,6 +143,8 @@ def wrapper(*args, **kwargs):
     else:
         collections = list(filter(len, np.array_split(collection, n_split)))
 
+    n_split = len(collections)
+    n_jobs = min(n_jobs, n_split)
     pass_queue = not hasattr(callback, "py_func")  # we'd be inside a numba function
 
     return wrapper
diff --git a/cellrank/ul/_utils.py b/cellrank/ul/_utils.py
index fe6139495..cfc72fa43 100644
--- a/cellrank/ul/_utils.py
+++ b/cellrank/ul/_utils.py
@@ -6,12 +6,12 @@
 from functools import wraps, update_wrapper
 from multiprocessing import cpu_count
 
-import numpy as np
-from scipy.sparse import issparse, spmatrix
-
 from cellrank import logging as logg
 from cellrank.ul._docs import d
 
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+
 AnnData = TypeVar("AnnData")
 
 
@@ -173,34 +173,56 @@ def _minmax(
     return float(np.nanmin(data)), float(np.nanmax(data))
 
 
-def _has_neighs(adata: AnnData) -> bool:
-    return "neighbors" in adata.uns.keys()
+def _modify_neigh_key(key: Optional[str]) -> str:
+    if key in (None, "connectivities", "distances"):
+        return "neighbors"
 
+    if key.endswith("_connectivities"):
+        key = key[:-15]
+    elif key.endswith("_distances"):
+        key = key[:-10]
 
-def _get_neighs(adata: AnnData, mode: str = "distances") -> Union[np.ndarray, spmatrix]:
-    try:
-        return _read_graph_data(adata, mode)
-    except KeyError:
-        return _read_graph_data(adata, "neighbors")[mode]
+    return key
+
+
+def _has_neighs(adata: AnnData, key: Optional[str] = None) -> bool:
+    return _modify_neigh_key(key) in adata.uns.keys()
 
 
-def _get_neighs_params(adata: AnnData) -> Dict[str, Any]:
-    return adata.uns["neighbors"]["params"]
+def _get_neighs(
+    adata: AnnData, mode: str = "distances", key: Optional[str] = None
+) -> Union[np.ndarray, spmatrix]:
+    if key is None:
+        res = _read_graph_data(adata, mode)  # legacy behavior
+    else:
+        try:
+            res = _read_graph_data(adata, key)
+            assert isinstance(res, (np.ndarray, spmatrix))
+        except (KeyError, AssertionError):
+            res = _read_graph_data(adata, f"{_modify_neigh_key(key)}_{mode}")
+
+    if not isinstance(res, (np.ndarray, spmatrix)):
+        raise TypeError(
+            f"Expected to find `numpy.ndarray` or `scipy.sparse.spmatrix`, found `{type(res)}`."
+        )
+
+    return res
+
+
+def _get_neighs_params(adata: AnnData, key: str = "neighbors") -> Dict[str, Any]:
+    return adata.uns.get(key, {}).get("params")
 
 
 def _read_graph_data(adata: AnnData, key: str) -> Union[np.ndarray, spmatrix]:
     """
     Read graph data from :mod:`anndata`.
 
-    :module`AnnData` >=0.7 stores `(n_obs x n_obs)` matrices in `.obsp` rather than `.uns`.
-    This is for backward compatibility.
-
     Parameters
     ----------
     adata
         Annotated data object.
     key
-        Key from either ``adata.uns`` or ``adata.obsp``.
+        Key in ``adata.obsp``.
 
     Returns
     -------
@@ -208,15 +230,11 @@ def _read_graph_data(adata: AnnData, key: str) -> Union[np.ndarray, spmatrix]:
         The graph data.
     """
 
-    if hasattr(adata, "obsp") and adata.obsp is not None and key in adata.obsp.keys():
-        logg.debug(f"Reading key `{key!r}` from `adata.obsp`")
+    logg.debug(f"Reading key `{key!r}` from `adata.obsp`")
+    if key in adata.obsp.keys():
         return adata.obsp[key]
 
-    if key in adata.uns.keys():
-        logg.debug(f"Reading key `{key!r}` from `adata.uns`")
-        return adata.uns[key]
-
-    raise KeyError(f"Unable to find key `{key!r}` in `adata.obsp` or `adata.uns`.")
+    raise KeyError(f"Unable to find key `{key!r}` in `adata.obsp`.")
 
 
 def _write_graph_data(
diff --git a/cellrank/ul/models/_base_model.py b/cellrank/ul/models/_base_model.py
index 3c0246246..649165866 100644
--- a/cellrank/ul/models/_base_model.py
+++ b/cellrank/ul/models/_base_model.py
@@ -19,6 +19,12 @@
 
 import wrapt
 
+from cellrank import logging as logg
+from cellrank.tl import Lineage
+from cellrank.ul._docs import d
+from cellrank.tl._utils import save_fig
+from cellrank.ul._utils import Pickleable, _minmax, valuedispatch, _densify_squeeze
+from cellrank.tl._constants import ModeEnum, AbsProbKey
 from scanpy.plotting._utils import add_colors_for_categorical_sample_annotation
 
 import numpy as np
@@ -34,13 +40,6 @@
 from matplotlib.colors import to_rgb, is_color_like
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-from cellrank import logging as logg
-from cellrank.tl import Lineage
-from cellrank.ul._docs import d
-from cellrank.tl._utils import save_fig
-from cellrank.ul._utils import Pickleable, _minmax, valuedispatch, _densify_squeeze
-from cellrank.tl._constants import ModeEnum, AbsProbKey
-
 AnnData = TypeVar("AnnData")
 _dup_spaces = re.compile(r" +")  # used on repr for underlying model's repr
 ArrayLike = Union[np.ndarray, spmatrix, List, Tuple]
@@ -289,19 +288,19 @@ def x_test(self) -> np.ndarray:
     @property
     @d.get_summary(base="base_model_y_test")
     def y_test(self) -> np.ndarray:
-        """Prediction values of shape `(n_samples,)` for :paramref:`x_test`."""
+        """Prediction values of shape `(n_samples,)` for :attr:`x_test`."""
         return self._y_test
 
     @property
     @d.get_summary(base="base_model_x_hat")
     def x_hat(self) -> np.ndarray:
-        """Filtered independent variables used when calculating default confidence interval, usually same as :paramref:`x`."""  # noqa
+        """Filtered independent variables used when calculating default confidence interval, usually same as :attr:`x`."""  # noqa
         return self._x_hat
 
     @property
     @d.get_summary(base="base_model_y_hat")
     def y_hat(self) -> np.ndarray:
-        """Filtered dependent variables used when calculating default confidence interval, usually same as :paramref:`y`."""  # noqa
+        """Filtered dependent variables used when calculating default confidence interval, usually same as :attr:`y`."""  # noqa
         return self._y_hat
 
     @property
@@ -333,18 +332,18 @@ def prepare(
         Parameters
         ----------
         gene
-            Gene in :paramref:`adata` ``.var_names`` or in :paramref:`adata` ``.raw.var_names``.
+            Gene in :attr:`adata` ``.var_names`` or in :attr:`adata` ``.raw.var_names``.
         lineage
-            Name of a lineage in :paramref:`adata` ``.obsm[lineage_key]``. If `None`, all weights will be set to `1`.
+            Name of a lineage in :attr:`adata` ``.obsm['{lineage_key}']``. If `None`, all weights will be set to `1`.
         %(backward)s
         %(time_range)s
         data_key
-            Key in :paramref:`adata` ``.layers`` or `'X'` for :paramref:`adata` ``.X``.
+            Key in :attr:`adata` ``.layers`` or `'X'` for :attr:`adata` ``.X``.
             If ``use_raw=True``, it's always set to `'X'`.
         time_key
-            Key in :paramref:`adata` ``.obs`` where the pseudotime is stored.
+            Key in :attr:`adata` ``.obs`` where the pseudotime is stored.
         use_raw
-            Whether to access :paramref:`adata` ``.raw`` or not.
+            Whether to access :attr:`adata` ``.raw`` or not.
         threshold
             Consider only cells with weights > ``threshold`` when estimating the test endpoint.
             If `None`, use the median of the weights.
@@ -361,17 +360,17 @@ def prepare(
         None
             Nothing, but updates the following fields:
 
-                - :paramref:`x` - %(base_model_x.summary)s
-                - :paramref:`y` - %(base_model_y.summary)s
-                - :paramref:`w` - %(base_model_w.summary)s
+                - :attr:`x` - %(base_model_x.summary)s
+                - :attr:`y` - %(base_model_y.summary)s
+                - :attr:`w` - %(base_model_w.summary)s
 
-                - :paramref:`x_all` - %(base_model_x_all.summary)s
-                - :paramref:`y_all` - %(base_model_y_all.summary)s
-                - :paramref:`w_all` - %(base_model_w_all.summary)s
+                - :attr:`x_all` - %(base_model_x_all.summary)s
+                - :attr:`y_all` - %(base_model_y_all.summary)s
+                - :attr:`w_all` - %(base_model_w_all.summary)s
 
-                - :paramref:`x_test` - %(base_model_x_test.summary)s
+                - :attr:`x_test` - %(base_model_x_test.summary)s
 
-                - :paramref:`prepared` - %(base_model_prepared.summary)s
+                - :attr:`prepared` - %(base_model_prepared.summary)s
         """
 
         self._use_raw = use_raw
@@ -580,13 +579,13 @@ def fit(
         Parameters
         ----------
         x
-            Independent variables, array of shape `(n_samples, 1)`. If `None`, use :paramref:`x`.
+            Independent variables, array of shape `(n_samples, 1)`. If `None`, use :attr:`x`.
         y
-            Dependent variables, array of shape `(n_samples, 1)`. If `None`, use :paramref:`y`.
+            Dependent variables, array of shape `(n_samples, 1)`. If `None`, use :attr:`y`.
         w
-            Optional weights of :paramref:`x`, array of shape `(n_samples,)`. If `None`, use :paramref:`w`.
+            Optional weights of :attr:`x`, array of shape `(n_samples,)`. If `None`, use :attr:`w`.
         kwargs
-            Keyword arguments for underlying :paramref:`model`'s fitting function.
+            Keyword arguments for underlying :attr:`model`'s fitting function.
 
         Returns
         -------
@@ -629,18 +628,18 @@ def predict(
         Parameters
         ----------
         x_test
-            Array of shape `(n_samples,)` used for prediction. If `None`, use :paramref:`x_test`.
+            Array of shape `(n_samples,)` used for prediction. If `None`, use :attr:`x_test`.
         key_added
-            Attribute name where to save the :paramref:`x_test` for later use. If `None`, don't save it.
+            Attribute name where to save the :attr:`x_test` for later use. If `None`, don't save it.
         kwargs
-            Keyword arguments for underlying :paramref:`model`'s prediction method.
+            Keyword arguments for underlying :attr:`model`'s prediction method.
 
         Returns
         -------
         :class:`numpy.ndarray`
             Updates and returns the following:
 
-                - :paramref:`y_test` - %(base_model_y_test.summary)s
+                - :attr:`y_test` - %(base_model_y_test.summary)s
         """
 
     @abstractmethod
@@ -654,15 +653,15 @@ def confidence_interval(
         """
         Calculate the confidence interval.
 
-        Use :meth:`default_confidence_interval` function if underlying :paramref:`model` has not method
+        Use :meth:`default_confidence_interval` function if underlying :attr:`model` has not method
         for confidence interval calculation.
 
         Parameters
         ----------
         x_test
-            Array of shape `(n_samples,)` used for confidence interval calculation. If `None`, use :paramref:`x_test`.
+            Array of shape `(n_samples,)` used for confidence interval calculation. If `None`, use :attr:`x_test`.
         kwargs
-            Keyword arguments for underlying :paramref:`model`'s confidence method
+            Keyword arguments for underlying :attr:`model`'s confidence method
             or for :meth:`default_confidence_interval`.
 
         Returns
@@ -670,7 +669,7 @@ def confidence_interval(
         :class:`numpy.ndarray`
             Updates the following fields:
 
-                - :paramref:`conf_int` - %(base_model_conf_int.summary)s
+                - :attr:`conf_int` - %(base_model_conf_int.summary)s
         """
 
     @d.dedent
@@ -680,9 +679,9 @@ def default_confidence_interval(
         **kwargs,
     ) -> np.ndarray:
         """
-        Calculate the confidence interval, if the underlying :paramref:`model` has no method for it.
+        Calculate the confidence interval, if the underlying :attr:`model` has no method for it.
 
-        This formula is taken from [DeSalvo70]_, eq. 5.
+        This formula is taken from :cite:`desalvo:70`, eq. 5.
 
         Parameters
         ----------
@@ -691,14 +690,8 @@ def default_confidence_interval(
         Returns
         -------
         %(base_model_ci.returns)s
-                - :paramref:`x_hat` - %(base_model_x_hat.summary)s
-                - :paramref:`y_hat` - %(base_model_y_hat.summary)s
-
-        References
-        ----------
-        .. [DeSalvo70] DeSalvo, J. S. (1970),
-            *Standard Error of Forecast in Multiple Regression: Proof of a Useful Result.*,
-            `RAND Corporation <https://www.rand.org/pubs/papers/P4365.html>`__.
+                - :attr:`x_hat` - %(base_model_x_hat.summary)s
+                - :attr:`y_hat` - %(base_model_y_hat.summary)s
         """
 
         use_ixs = self.w > 0
@@ -799,7 +792,7 @@ def plot(
             Note that this will require 1 additional model fit.
         lineage_probability_conf_int
             Whether to compute and show smoothed lineage probability confidence interval.
-            If :paramref:`self` is :class:`cellrank.ul.models.GAMR`, it can also specify the confidence level,
+            If :attr:`self` is :class:`cellrank.ul.models.GAMR`, it can also specify the confidence level,
             the default is `0.95`. Only used when ``show_lineage_probability=True``.
         lineage_probability_color
             Color to use when plotting the smoothed ``lineage_probability``.
@@ -1023,7 +1016,7 @@ def _reshape_and_retype(self, arr: np.ndarray) -> np.ndarray:
         Returns
         -------
         :class:`np.ndarray`
-            Array of shape `(n, 1)` with dtype set to :paramref:`_dtype`.
+            Array of shape `(n, 1)` with dtype set to :attr:`_dtype`.
         """
 
         if arr.ndim not in (1, 2):
@@ -1269,14 +1262,14 @@ def _return_min_max(self, show_conf_int: bool):
 
 class FailedModel(BaseModel):
     """
-    Model representing a failure of the original :paramref:`model`.
+    Model representing a failure of the original :attr:`model`.
 
     Parameters
     ----------
     model
         The original model which has failed.
     exc
-        The exception that caused the :paramref:`model` to fail or a :class:`str` containing the message.
+        The exception that caused the :attr:`model` to fail or a :class:`str` containing the message.
         In the latter case, :meth:`cellrank.ul.models.FailedModel.reraise` a :class:`RuntimeError` with that message.
         If `None`, :class`UnknownModelError` will eventually be raised.
     """
@@ -1404,7 +1397,7 @@ class FittedModel(BaseModel):
         If `None`, always sets `hide_cells=True`` in :meth:`cellrank.ul.models.BaseModel.plot`.
     w_all
         %(base_model_w_all.summary)s
-        If `None` and :paramref:`x_all` and :paramref:`y_all` are present, it will be set an array of `1`.
+        If `None` and :attr:`x_all` and :attr:`y_all` are present, it will be set an array of `1`.
     """
 
     def __init__(
diff --git a/cellrank/ul/models/_gamr_model.py b/cellrank/ul/models/_gamr_model.py
index ea1510e80..9cf7f6660 100644
--- a/cellrank/ul/models/_gamr_model.py
+++ b/cellrank/ul/models/_gamr_model.py
@@ -2,16 +2,16 @@
 from copy import copy, deepcopy
 from typing import Any, Tuple, Union, Optional
 
-import numpy as np
-import pandas as pd
-from scipy.stats import norm
-
 from cellrank.ul._docs import d, inject_docs
 from cellrank.ul.models import BaseModel
 from cellrank.tl._constants import ModeEnum
 from cellrank.ul.models._utils import _OFFSET_KEY, _get_offset, _get_knotlocs
 from cellrank.ul.models._base_model import AnnData, FailedModel
 
+import numpy as np
+import pandas as pd
+from scipy.stats import norm
+
 _r_lib = None
 _r_lib_name = None
 
@@ -35,7 +35,7 @@ class GAMR(BaseModel):
         Number of knots.
     distribution
         Distribution family in `rpy2.robjects.r`, such as `'gaussian'` or `'nb'` for negative binomial.
-        If `'nb'`, raw count data in :paramref:`adata` ``.raw`` is always used.
+        If `'nb'`, raw count data in :attr:`adata` ``.raw`` is always used.
     basis
         Basis for the smoothing term.
         See `here <https://www.rdocumentation.org/packages/mgcv/versions/1.8-33/topics/s>`__ for valid options.
@@ -46,7 +46,7 @@ class GAMR(BaseModel):
             - `{kloc.DENSITY.s!r}` - position the knots based on the density of the pseudotime.
     offset
         Offset term for the GAM. Only available when ``distribution='nb'``. If `'default'`, it is calculated
-        according to [Robinson10]_. The values are saved in :paramref:`adata` ``.obs[{key!r}]``.
+        according to :cite:`robinson:10`. The values are saved in :attr:`adata` ``.obs[{key!r}]``.
         If `None`, no offset is used.
     smoothing_penalty
         Penalty for the smoothing term. The larger the value, the smoother the fitted curve.
@@ -177,11 +177,11 @@ def fit(
         Returns
         -------
         :class:`cellrank.ul.models.GAMR`
-            Fits the model and returns self. Updates the following fields by filtering out `0` weights :paramref:`w`:
+            Fits the model and returns self. Updates the following fields by filtering out `0` weights :attr:`w`:
 
-                - :paramref:`x` - %(base_model_x.summary)s
-                - :paramref:`y` - %(base_model_y.summary)s
-                - :paramref:`w` - %(base_model_w.summary)s
+                - :attr:`x` - %(base_model_x.summary)s
+                - :attr:`y` - %(base_model_y.summary)s
+                - :attr:`w` - %(base_model_w.summary)s
         """  # noqa
 
         from rpy2.robjects import Formula, r, pandas2ri
diff --git a/cellrank/ul/models/_pygam_model.py b/cellrank/ul/models/_pygam_model.py
index 4a58f6b94..98e7c5c75 100644
--- a/cellrank/ul/models/_pygam_model.py
+++ b/cellrank/ul/models/_pygam_model.py
@@ -5,6 +5,13 @@
 from typing import Union, Mapping, Optional
 from collections import defaultdict
 
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.ul.models import BaseModel
+from cellrank.tl._constants import ModeEnum
+from cellrank.tl.kernels._utils import _filter_kwargs
+from cellrank.ul.models._base_model import AnnData
+
 import numpy as np
 from pygam import GAM as pGAM
 from pygam import (
@@ -17,13 +24,6 @@
     s,
 )
 
-from cellrank import logging as logg
-from cellrank.ul._docs import d
-from cellrank.ul.models import BaseModel
-from cellrank.tl._constants import ModeEnum
-from cellrank.tl.kernels._utils import _filter_kwargs
-from cellrank.ul.models._base_model import AnnData
-
 _r_lib = None
 _r_lib_name = None
 
diff --git a/cellrank/ul/models/_sklearn_model.py b/cellrank/ul/models/_sklearn_model.py
index be452af14..156e87532 100644
--- a/cellrank/ul/models/_sklearn_model.py
+++ b/cellrank/ul/models/_sklearn_model.py
@@ -2,13 +2,13 @@
 from typing import Iterable, Optional
 from inspect import signature
 
-import numpy as np
-from sklearn.base import BaseEstimator
-
 from cellrank.ul._docs import d
 from cellrank.ul.models import BaseModel
 from cellrank.ul.models._base_model import AnnData
 
+import numpy as np
+from sklearn.base import BaseEstimator
+
 
 @d.dedent
 class SKLearnModel(BaseModel):
@@ -21,10 +21,10 @@ class SKLearnModel(BaseModel):
     model
         Instance of the underlying :mod:`sklearn` estimator, such as :class:`sklearn.svm.SVR`.
     weight_name
-        Name of the weight argument for :paramref:`model` ``.fit``. If `None`, to determine it automatically.
+        Name of the weight argument for :attr:`model` ``.fit``. If `None`, to determine it automatically.
         If and empty string, no weights will be used.
     ignore_raise
-        Do not raise an exception if weight argument is not found in the fitting function of :paramref:`model`.
+        Do not raise an exception if weight argument is not found in the fitting function of :attr:`model`.
         This is useful in case when weight is passed in ``**kwargs`` and cannot be determined from signature.
     """
 
@@ -161,7 +161,7 @@ def _find_func(
         default: Optional[str] = None,
     ) -> Optional[str]:
         """
-        Find a function in :paramref:`model` from given names.
+        Find a function in :attr:`model` from given names.
 
         If `None` is found, use :paramref;`default` or raise a :class:`RuntimeError`.
 
@@ -193,12 +193,12 @@ def _find_arg_name(
         self, func_name: Optional[str], param_names: Iterable[str]
     ) -> Optional[str]:
         """
-        Find an argument in :paramref:`model`'s ``func_name``.
+        Find an argument in :attr:`model`'s ``func_name``.
 
         Parameters
         ----------
         func_name
-            Function name of :paramref:`model`.
+            Function name of :attr:`model`.
         param_names
             Parameter names to search. The first one found is returned.
 
diff --git a/cellrank/ul/models/_utils.py b/cellrank/ul/models/_utils.py
index 184c49f1e..d9ffc63a4 100644
--- a/cellrank/ul/models/_utils.py
+++ b/cellrank/ul/models/_utils.py
@@ -1,16 +1,16 @@
 from typing import Union, TypeVar, Optional, Sequence
 
-import numpy as np
-from numba import njit, prange
-from scipy.sparse import issparse, spmatrix
-from sklearn.utils.sparsefuncs import csc_median_axis_0
-
 import cellrank.logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.ul._utils import valuedispatch
 from cellrank.tl._constants import ModeEnum
 from cellrank.ul._parallelize import parallelize
 
+import numpy as np
+from numba import njit, prange
+from scipy.sparse import issparse, spmatrix
+from sklearn.utils.sparsefuncs import csc_median_axis_0
+
 AnnData = TypeVar("AnnData")
 _OFFSET_KEY = "cellrank_offset"
 
@@ -139,7 +139,7 @@ def _calculate_norm_factors(
     method
         One of following:
 
-            - method={m.TMM.s!r} - weighted trimmed mean of M-values from [Robinson10]_.
+            - method={m.TMM.s!r} - weighted trimmed mean of M-values from :cite:`robinson:10`.
             - method={m.RLE.s!r} - relative log expression from [Anders10]_.
             - method={m.UPPER_QUANT.s!r} - upper-quartile normalization method from [Bullard10]_.
             - method={m.NONE.s!r} - all the factors are set to 1.
diff --git a/docs/requirements.txt b/docs/requirements.txt
index b5f46437f..3c118ac3c 100644
--- a/docs/requirements.txt
+++ b/docs/requirements.txt
@@ -1,15 +1,17 @@
 -r ../requirements.txt  # must include this because of examples
 bezier
+ipython
+ipywidgets
 louvain>=0.6,!=0.6.2
-memory_profiler>=0.57.0
-nbsphinx>=0.7
+memory_profiler>=0.58.0
+nbsphinx>=0.8
 pypandoc
 python-igraph
-sphinx>=3.2.1
+sphinx>=4
 sphinx-autodoc-annotation
 sphinx-autodoc-typehints>=1.10.3
 sphinx-gallery
-sphinx-paramlinks
 sphinx_copybutton
 sphinx_last_updated_by_git
 sphinx_rtd_theme
+sphinxcontrib-bibtex>=2.3.0
diff --git a/docs/source/_templates/autosummary/class.rst b/docs/source/_templates/autosummary/class.rst
index 748806cff..828c2a2dd 100644
--- a/docs/source/_templates/autosummary/class.rst
+++ b/docs/source/_templates/autosummary/class.rst
@@ -34,6 +34,6 @@
 
    .. _sphx_glr_backref_{{fullname}}:
 
-   .. minigallery:: {{fullname}}
-       :add-heading: Examples
-       :heading-level: -
+.. minigallery:: {{fullname}}
+   :add-heading: Examples
+   :heading-level: -
diff --git a/docs/source/api.rst b/docs/source/api.rst
index 0db257e80..0bfecf0ee 100644
--- a/docs/source/api.rst
+++ b/docs/source/api.rst
@@ -44,9 +44,6 @@ Kernels
 -------
 Kernels are part of the low-level API and are used to estimate cell-to-cell transitions.
 
-.. module:: cellrank.tl
-.. currentmodule:: cellrank
-
 .. autosummary::
     :toctree: api
 
@@ -60,9 +57,6 @@ Estimators
 ----------
 Estimators predict cell fates using the transitions derived from :ref:`Kernels`.
 
-.. module:: cellrank.tl
-.. currentmodule:: cellrank
-
 .. autosummary::
     :toctree: api
 
@@ -84,6 +78,7 @@ Plotting
     pl.lineage_drivers
     pl.circular_projection
     pl.gene_trends
+    pl.log_odds
     pl.heatmap
     pl.cluster_lineage
     pl.cluster_fates
@@ -125,7 +120,8 @@ Datasets
 
     datasets.pancreas
     datasets.lung
-    datasets.reprogramming
+    datasets.reprogramming_morris
+    datasets.reprogramming_schiebinger
     datasets.zebrafish
     datasets.pancreas_preprocessed
 
diff --git a/docs/source/classes.rst b/docs/source/classes.rst
index 9ffc9b3b9..9df52db06 100644
--- a/docs/source/classes.rst
+++ b/docs/source/classes.rst
@@ -7,13 +7,14 @@ Estimators
 GPCCA
 -----
 .. autoclass:: cellrank.tl.estimators.GPCCA
+    :noindex:
     :members:
     :inherited-members:
 
 CFLARE
 ------
-
 .. autoclass:: cellrank.tl.estimators.CFLARE
+    :noindex:
     :members:
     :inherited-members:
 
@@ -22,65 +23,59 @@ Kernels
 
 Velocity Kernel
 ---------------
-
 .. autoclass:: cellrank.tl.kernels.VelocityKernel
+    :noindex:
     :members:
 
-Cosine similarity scheme
+Cosine Similarity Scheme
 ++++++++++++++++++++++++
-
 .. autoclass:: cellrank.tl.kernels.CosineScheme
     :members: __call__, hessian
 
-Correlation scheme
+Correlation Scheme
 ++++++++++++++++++
-
 .. autoclass:: cellrank.tl.kernels.CorrelationScheme
     :members: __call__, hessian
 
-Dot product scheme
+Dot Product Scheme
 ++++++++++++++++++
-
 .. autoclass:: cellrank.tl.kernels.DotProductScheme
     :members: __call__, hessian
 
-
 Connectivity Kernel
 -------------------
-
 .. autoclass:: cellrank.tl.kernels.ConnectivityKernel
+    :noindex:
     :members:
 
 Pseudotime Kernel
 -----------------
-
 .. autoclass:: cellrank.tl.kernels.PseudotimeKernel
+    :noindex:
     :members:
 
-Hard threshold scheme
+Hard Threshold Scheme
 +++++++++++++++++++++
-
 .. autoclass:: cellrank.tl.kernels.HardThresholdScheme
     :members:
     :special-members: __call__
 
-Soft threshold scheme
+Soft Threshold Scheme
 +++++++++++++++++++++
-
 .. autoclass:: cellrank.tl.kernels.SoftThresholdScheme
     :members:
     :special-members: __call__
 
 CytoTRACE Kernel
 ----------------
-
 .. autoclass:: cellrank.tl.kernels.CytoTRACEKernel
+    :noindex:
     :members: compute_cytotrace, compute_transition_matrix
 
 Precomputed Kernel
 ------------------
-
 .. autoclass:: cellrank.tl.kernels.PrecomputedKernel
+    :noindex:
     :members:
 
 Models
@@ -88,22 +83,22 @@ Models
 
 GAM
 ---
-
 .. autoclass:: cellrank.ul.models.GAM
+    :noindex:
     :members:
     :inherited-members:
 
 SKLearnModel
 ------------
-
 .. autoclass:: cellrank.ul.models.SKLearnModel
+    :noindex:
     :members:
     :inherited-members:
 
 GAMR
 ----
-
 .. autoclass:: cellrank.ul.models.GAMR
+    :noindex:
     :members:
     :inherited-members:
 
@@ -112,27 +107,35 @@ Base Classes
 
 BaseEstimator
 -------------
-
 .. autoclass:: cellrank.tl.estimators.BaseEstimator
     :members:
 
 Kernel
 ------
-
 .. autoclass:: cellrank.tl.kernels.Kernel
     :members:
     :inherited-members:
 
-Similarity scheme
------------------
+ExperimentalTime Kernel
+-----------------------
+.. autoclass:: cellrank.tl.kernels.ExperimentalTimeKernel
+    :members:
+    :inherited-members:
+
+TransportMap Kernel
+-------------------
+.. autoclass:: cellrank.tl.kernels.TransportMapKernel
+    :members:
+    :inherited-members:
 
+Similarity Scheme
+-----------------
 .. autoclass:: cellrank.tl.kernels.SimilaritySchemeABC
     :members:
     :special-members: __call__
     :inherited-members:
 
-
-Threshold scheme
+Threshold Scheme
 ----------------
 .. autoclass:: cellrank.tl.kernels.ThresholdSchemeABC
     :members:
@@ -141,12 +144,10 @@ Threshold scheme
 
 BaseModel
 ---------
-
 .. autoclass:: cellrank.ul.models.BaseModel
     :members:
 
 Lineage
 -------
-
 .. autoclass:: cellrank.tl.Lineage
     :members: priming_degree, reduce, plot_pie, X, T, view, names, colors
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 0a5b02fe7..9c8ff4158 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -3,7 +3,7 @@
 import logging
 from pathlib import Path
 from datetime import datetime
-from collections import ChainMap, defaultdict
+from collections import ChainMap
 from urllib.parse import urljoin
 from urllib.request import urlretrieve
 
@@ -36,6 +36,7 @@
     "cellrank_basics.ipynb",
     "kernels_and_estimators.ipynb",
     "beyond_rna_velocity.ipynb",
+    "real_time.ipynb",
     "creating_new_kernel.ipynb",
 ]:
     try:
@@ -65,12 +66,14 @@
     "sphinx.ext.viewcode",
     "sphinx_autodoc_typehints",
     "sphinx.ext.intersphinx",
-    "sphinx_paramlinks",
     "sphinx.ext.autosummary",
     "sphinx_gallery.gen_gallery",
     "nbsphinx",
     "sphinx_copybutton",
     "sphinx_last_updated_by_git",
+    "sphinxcontrib.bibtex",
+    # https://github.com/spatialaudio/nbsphinx/issues/24
+    "IPython.sphinxext.ipython_console_highlighting",
 ]
 
 intersphinx_mapping = dict(
@@ -91,7 +94,6 @@
     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),
 )
 
@@ -105,11 +107,16 @@
 # directories to ignore when looking for source files.
 # This pattern also affects html_static_path and html_extra_path.
 exclude_patterns = [
-    "**.md5",
-    "**.py",
+    "auto_*/**.ipynb",
+    "auto_*/**.md5",
+    "auto_*/**.py",
     "**.ipynb_checkpoints",
-]  # ignore anything that isn't .rst
-suppress_warnings = ["ref.citation"]
+]
+
+# bibliography
+bibtex_bibfiles = ["references.bib"]
+bibtex_reference_style = "author_year"
+bibtex_default_style = "alpha"
 
 # -- Notebooks
 nbsphinx_execute_arguments = [
@@ -158,15 +165,16 @@ class ExplicitSubsectionOrder(_SortKey):
 
     _order = ChainMap(
         {
-            example_dir / "estimators" / "compute_eigendecomposition.py": 0,
-            example_dir / "estimators" / "compute_schur_vectors.py": 10,
-            example_dir / "estimators" / "compute_schur_matrix.py": 20,
-            example_dir / "estimators" / "compute_macrostates.py": 30,
-            example_dir / "estimators" / "compute_coarse_T.py": 40,
-            example_dir / "estimators" / "compute_terminal_states_gpcca.py": 50,
-            example_dir / "estimators" / "compute_abs_probs.py": 60,
-            example_dir / "estimators" / "compute_lineage_drivers.py": 70,
-            example_dir / "estimators" / "compute_fit.py": 80,
+            example_dir / "kernels" / "plot_projection.py": 0,
+            example_dir / "kernels" / "plot_random_walks.py": 10,
+            example_dir / "kernels" / "compute_kernel_tricks.py": 30,
+        },
+        {
+            example_dir / "estimators" / "compute_macrostates.py": 0,
+            example_dir / "estimators" / "compute_coarse_T.py": 10,
+            example_dir / "estimators" / "compute_terminal_states_gpcca.py": 20,
+            example_dir / "estimators" / "compute_abs_probs.py": 30,
+            example_dir / "estimators" / "compute_lineage_drivers.py": 40,
         },
         {
             example_dir / "plotting" / "plot_initial_states.py": 0,
@@ -174,20 +182,17 @@ class ExplicitSubsectionOrder(_SortKey):
             example_dir / "plotting" / "plot_lineages.py": 20,
             example_dir / "plotting" / "plot_lineage_drivers.py": 30,
             example_dir / "plotting" / "plot_directed_paga.py": 40,
-            example_dir / "plotting" / "plot_gene_trends.py": 50,
-            example_dir / "plotting" / "plot_heatmap.py": 60,
-            example_dir / "plotting" / "plot_cluster_lineage.py": 70,
-            example_dir / "plotting" / "plot_cluster_fates.py": 80,
-            example_dir / "plotting" / "plot_graph.py": 90,
+            example_dir / "plotting" / "plot_circular_embedding.py": 50,
+            example_dir / "plotting" / "plot_gene_trends.py": 60,
+            example_dir / "plotting" / "plot_heatmap.py": 70,
+            example_dir / "plotting" / "plot_cluster_lineage.py": 80,
+            example_dir / "plotting" / "plot_cluster_fates.py": 90,
+            example_dir / "plotting" / "plot_graph.py": 100,
+        },
+        {
+            example_dir / "other" / "plot_model.py": 0,
+            example_dir / "other" / "compute_lineage_tricks.py": 10,
         },
-        defaultdict(
-            lambda: 1000,
-            {
-                example_dir / "other" / "plot_model.py": 0,
-                example_dir / "other" / "compute_lineage_tricks.py": 10,
-                example_dir / "other" / "compute_kernel_tricks.py": 20,
-            },
-        ),
     )
 
     def __call__(self, filename: str) -> int:
@@ -213,7 +218,8 @@ def __repr__(self) -> str:
     "within_subsection_order": ExplicitSubsectionOrder,
     "subsection_order": ExplicitOrder(
         [
-            rel_example_dir / "estimators",  # really must be relative
+            rel_example_dir / "kernels",
+            rel_example_dir / "estimators",
             rel_example_dir / "plotting",
             rel_example_dir / "other",
         ]
@@ -237,7 +243,6 @@ def __repr__(self) -> str:
 
 autosummary_generate = True
 autodoc_member_order = "bysource"
-# autodoc_default_flags = ['members']
 typehints_fully_qualified = False
 napoleon_google_docstring = False
 napoleon_numpy_docstring = True
@@ -259,4 +264,4 @@ def __repr__(self) -> str:
 
 
 def setup(app):
-    app.add_stylesheet("css/custom.css")
+    app.add_css_file("css/custom.css")
diff --git a/docs/source/external_api.rst b/docs/source/external_api.rst
index c19b9eef9..8ce851a8f 100644
--- a/docs/source/external_api.rst
+++ b/docs/source/external_api.rst
@@ -6,29 +6,21 @@ refer to our `contributing guide <https://github.com/theislab/cellrank/blob/mast
 
 Kernels
 ~~~~~~~
-
 .. module:: cellrank.external
 .. currentmodule:: cellrank
 
 .. autosummary::
     :toctree: api
 
-    external.kernels.OTKernel
+    external.kernels.StationaryOTKernel
+    external.kernels.WOTKernel
 
 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 70e6a2ef3..239db203c 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -9,16 +9,16 @@ CellRank for directed single-cell fate mapping
 
 **CellRank** is a toolkit to uncover cellular dynamics based on Markov state modeling of single-cell data. It contains
 two main modules: `kernels`_ compute cell-cell transition probabilities and `estimators`_ generate hypothesis based on
-these. Our kernels work with a variety of input data including `RNA velocity`_ (see [Manno18]_ and [Bergen20]_),
+these. Our kernels work with a variety of input data including `RNA velocity`_ :cite:`manno:18` and :cite:`bergen:20`,
 `cellular similarity`_ (both transcriptomic and spatial) and `pseudotime`_, among others. Our `VelocityKernel`_
 takes into account **uncertainty in the velocities** and allows you to aggregate the short-range fate relations
 given by RNA velocity into longer trends along the phenotypic manifold. Our main estimator is
-*Generalized Perron Cluster Cluster Analysis* (G-PCCA) [GPCCA18]_ which coarse-grains the Markov chain
+*Generalized Perron Cluster Cluster Analysis* (G-PCCA) :cite:`reuter:18` which coarse-grains the Markov chain
 into a set of macrostates which represent initial, terminal and intermediate states. For each transient cell,
 we compute its fate probability towards any terminal state. We show an example of such a fate map in the figure above,
-which has been computed using the data of [Panc19]_. CellRank combines `kernels`_ and `estimators`_ with a powerful
-`plotting API`_, enabling you to visualize e.g. smooth `gene expression trends`_ along lineages or fate-informed
-`circular embeddings`_, to name just a few.
+which has been computed using the data of :cite:`bastidas-ponce:19`. CellRank combines `kernels`_ and `estimators`_
+with a powerful `plotting API`_, enabling you to visualize e.g. smooth `gene expression trends`_ along lineages
+or fate-informed `circular embeddings`_, to name just a few.
 
 CellRank scales to large cell numbers, is fully compatible with `scanpy`_ and `scvelo`_ and is easy to use.
 
@@ -29,9 +29,11 @@ API. Most biological systems require a bit more control, so be sure to check out
 which allows to unlock the full power of CellRank. If you want to see individual functions in action, visit our
 `gallery`_.
 
+To use CellRank without RNA velocity information, check out the `beyond RNA velocity tutorial`_ as well as the
+`time-series tutorial`_.
+
 Latest additions
 ^^^^^^^^^^^^^^^^
-
 .. include:: latest_additions.rst
 
 Manuscript
@@ -96,6 +98,7 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
    cellrank_basics
    kernels_and_estimators
    beyond_rna_velocity
+   real_time
    creating_new_kernel
 
 .. |PyPI| image:: https://img.shields.io/pypi/v/cellrank.svg
@@ -144,6 +147,8 @@ CellRank was developed in collaboration between the `Theislab`_ and the `Peerlab
 .. _basic tutorial: https://cellrank.readthedocs.io/en/stable/cellrank_basics.html
 .. _kernel tutorial: https://cellrank.readthedocs.io/en/stable/creating_new_kernel.html
 .. _kernels and estimators tutorial: https://cellrank.readthedocs.io/en/stable/kernels_and_estimators.html
+.. _beyond RNA velocity tutorial: https://cellrank.readthedocs.io/en/stable/beyond_rna_velocity.html
+.. _time-series tutorial: https://cellrank.readthedocs.io/en/stable/real_time.html
 
 .. _scanpy: https://scanpy.readthedocs.io/en/stable/
 .. _scvelo: https://scvelo.readthedocs.io/
diff --git a/docs/source/references.bib b/docs/source/references.bib
new file mode 100644
index 000000000..5ceac8505
--- /dev/null
+++ b/docs/source/references.bib
@@ -0,0 +1,438 @@
+@article{bergen:20,
+    author = {Bergen, V. and Lange, M. and Peidli, S. and Wolf, F.A. and Theis, F.J.},
+    title = {{Generalizing RNA velocity to transient cell states through dynamical modeling.}},
+    journal = {Nat. Biotechnol.},
+    location = {75 Varick St, 9th Flr, New York, Ny 10013-1917 Usa},
+    publisher = {Nature Publishing Group},
+    volume = {38},
+    pages = {1408?1414},
+    year = {2020},
+    issn = {0733-222X},
+    eissn = {1546-1696},
+    doi = {10.1038/s41587-020-0591-3},
+}
+
+@article{coifman:05,
+	author = {Coifman, R. R. and Lafon, S. and Lee, A. B. and Maggioni, M. and Nadler, B. and Warner, F. and Zucker, S. W.},
+	title = {Geometric diffusions as a tool for harmonic analysis and structure definition of data: Diffusion maps},
+	volume = {102},
+	number = {21},
+	pages = {7426--7431},
+	year = {2005},
+	doi = {10.1073/pnas.0500334102},
+	publisher = {National Academy of Sciences},
+	issn = {0027-8424},
+	journal = {Proceedings of the National Academy of Sciences},
+}
+
+@article{gulati:20,
+	author = {Gulati, Gunsagar S. and Sikandar, Shaheen S. and Wesche, Daniel J. and Manjunath, Anoop and Bharadwaj, Anjan and Berger, Mark J. and Ilagan, Francisco and Kuo, Angera~H. and Hsieh, Robert W. and Cai, Shang and Zabala, Maider and Scheeren, Ferenc A. and Lobo, Neethan A. and Qian, Dalong and Yu, Feiqiao B. and Dirbas, Frederick M. and Clarke, Michael F. and Newman, Aaron M.},
+	title = {Single-cell transcriptional diversity is a hallmark of developmental potential},
+	volume = {367},
+	number = {6476},
+	pages = {405--411},
+	year = {2020},
+	doi = {10.1126/science.aax0249},
+	publisher = {American Association for the Advancement of Science},
+	issn = {0036-8075},
+	journal = {Science},
+}
+
+@article{farrell:18,
+	author = {Farrell, Jeffrey A. and Wang, Yiqun and Riesenfeld, Samantha J. and Shekhar, Karthik and Regev, Aviv and Schier, Alexander F.},
+	title = {Single-cell reconstruction of developmental trajectories during zebrafish embryogenesis},
+	volume = {360},
+	number = {6392},
+	elocation-id = {eaar3131},
+	year = {2018},
+	doi = {10.1126/science.aar3131},
+	publisher = {American Association for the Advancement of Science},
+	issn = {0036-8075},
+	journal = {Science},
+}
+
+@article{reuter:18,
+    author = {Reuter, Bernhard and Weber, Marcus and Fackeldey, Konstantin and Röblitz, Susanna and Garcia, Martin E.},
+    title = {Generalized Markov State Modeling Method for Nonequilibrium Biomolecular Dynamics: Exemplified on Amyloid β Conformational Dynamics Driven by an Oscillating Electric Field},
+    journal = {Journal of Chemical Theory and Computation},
+    volume = {14},
+    number = {7},
+    pages = {3579-3594},
+    year = {2018},
+    doi = {10.1021/acs.jctc.8b00079},
+    note  = {PMID: 29812922},
+}
+
+@article{haghverdi:16,
+    author = {Haghverdi, L. and B{\"u}ttner, M. and Wolf, F.A. and Buettner, F. and Theis, F.J.},
+    title = {{Diffusion pseudotime robustly reconstructs lineage branching.}},
+    journal = {Nat. Methods},
+    location = {London},
+    publisher = {Nature Publishing Group},
+    volume = {13},
+    number = {10},
+    pages = {845-848},
+    year = {2016},
+    issn = {1548-7091},
+    eissn = {1548-7105},
+    doi = {10.1038/nmeth.3971},
+}
+
+@article{li:20,
+    author = {Li , Tiejun and Shi , Jifan and Wu , Yichong and Zhou , Peijie},
+    title = {On the Mathematics of RNA Velocity I: Theoretical Analysis},
+    journal = {CSIAM Transactions on Applied Mathematics},
+    year = {2021},
+    volume = {2},
+    number = {1},
+    pages = {1--55},
+    issn = {2708-0579},
+    doi = {10.4208/csiam-am.SO-2020-0001},
+}
+
+@article{strunz:20,
+    author = {Strunz, M. and Simon, L. and Ansari, M. and Kathiriya, J.J. and Angelidis, I. and Mayr, C. and Tsidiridis, G. and Lange, M. and Mattner, L. and Yee, M. and Ogar, P. and Sengupta, A. and Kukhtevich, I. and Schneider, R. and Zhao, Z. and Voss, C. and St{\"o}ger, T. and Neumann, J.H.L. and Hilgendorff, A. and Behr, J. and O'Reilly, M. and Lehmann, M. and Burgstaller, G. and K{\"o}nigshoff, M. and Chapman, H.A. and Theis, F.J. and Schiller, H. B.},
+    title = {{Alveolar regeneration through a Krt8+transitional stem cell state that persists in human lung fibrosis.}},
+    journal = {Nat. Commun.},
+    location = {Macmillan Building, 4 Crinan St, London N1 9xw, England},
+    publisher = {Nature Publishing Group},
+    volume = {11},
+    number = {1},
+    year = {2020},
+    issn = {2041-1723},
+    eissn = {2041-1723},
+    doi = {10.1038/s41467-020-17358-3},
+}
+
+@article{macosko:15,
+    author = {Evan Z. Macosko and Anindita Basu and Rahul Satija and James Nemesh and Karthik Shekhar and Melissa Goldman and Itay Tirosh and Allison R. Bialas and Nolan Kamitaki and Emily M. Martersteck and John J. Trombetta and David A. Weitz and Joshua R. Sanes and Alex K. Shalek and Aviv Regev and Steven A. McCarroll},
+    title = {Highly Parallel Genome-wide Expression Profiling of Individual Cells Using Nanoliter Droplets},
+    journal = {Cell},
+    volume = {161},
+    number = {5},
+    pages = {1202-1214},
+    year = {2015},
+    issn = {0092-8674},
+    doi = {10.1016/j.cell.2015.05.002},
+}
+
+@article{vandijk:18,
+    author = {van Dijk, David and Sharma, Roshan and Nainys, Juozas and Yim, Kristina and Kathail, Pooja and Carr, Ambrose J. and Burdziak, Cassandra and Moon, Kevin R. and Chaffer, Christine L. and Pattabiraman, Diwakar and Bierie, Brian and Mazutis, Linas and Wolf, Guy and Krishnaswamy, Smita and Pe'er, Dana},
+    title = {Recovering Gene Interactions from Single-Cell Data Using Data Diffusion},
+    journal = {Cell},
+    year = {2018},
+    month = {Jul},
+    day = {26},
+    publisher = {Elsevier},
+    volume = {174},
+    number = {3},
+    pages = {716-729.e27},
+    issn = {0092-8674},
+    doi = {10.1016/j.cell.2018.05.061},
+}
+
+@article{manno:18,
+    added-at = {2018-08-11T16:49:03.000+0200},
+    author = {La Manno, Gioele and Soldatov, Ruslan and Zeisel, Amit and Braun, Emelie and Hochgerner, Hannah and Petukhov, Viktor and Lidschreiber, Katja and Kastriti, Maria E. and Lönnerberg, Peter and Furlan, Alessandro and Fan, Jean and Borm, Lars E. and Liu, Zehua and van Bruggen, David and Guo, Jimin and He, Xiaoling and Barker, Roger and Sundström, Erik and Castelo-Branco, Gonçalo and Cramer, Patrick and Adameyko, Igor and Linnarsson, Sten and Kharchenko, Peter V.},
+    doi = {10.1038/s41586-018-0414-6},
+    interhash = {9139d1fd43303666673c91dcab8b69a6},
+    intrahash = {44c48e1742cc878db81d591358193289},
+    issn = {14764687},
+    journal = {Nature},
+    keywords = {RNA_abundance cellular_machinery transcription},
+    refid = {La Manno2018},
+    timestamp = {2018-08-11T16:49:03.000+0200},
+    title = {{RNA} velocity of single cells},
+    year = {2018},
+}
+
+@article{morris:18,
+    author = "Biddy, Brent A. and Kong, Wenjun and Kamimoto, Kenji and Guo, Chuner and Waye, Sarah E. and Sun, Tao and Morris, Samantha A.",
+    title = "Single-cell mapping of lineage and identity in direct reprogramming",
+    journal = "Nature",
+    year = "2018",
+    month = "Dec",
+    day = "01",
+    volume = "564",
+    number = "7735",
+    pages = "219--224",
+    issn = "1476-4687",
+    doi = "10.1038/s41586-018-0744-4",
+}
+
+@article{palla:21,
+	author = {Palla, Giovanni and Spitzer, Hannah and Klein, Michal and Fischer, David and Schaar, Anna Christina and Kuemmerle, Louis Benedikt and Rybakov, Sergei and Ibarra, Ignacio L. and Holmberg, Olle and Virshup, Isaac and Lotfollahi, Mohammad and Richter, Sabrina and Theis, Fabian J.},
+	title = {Squidpy: a scalable framework for spatial single cell analysis},
+	elocation-id = {2021.02.19.431994},
+	year = {2021},
+	doi = {10.1101/2021.02.19.431994},
+	publisher = {Cold Spring Harbor Laboratory},
+	journal = {bioRxiv},
+}
+
+@article{bastidas-ponce:19,
+    author = {Bastidas-Ponce, A. and Tritschler, S. and Dony, L. and Scheibner, K. and Tarquis-Medina, M. and Salinno, C. and Schirge, S. and Burtscher, I. and B{\"o}ttcher, A. and Theis, F.J. and Lickert, H. and Bakhti, M.},
+    title = {{Comprehensive single cell mRNA profiling reveals a detailed roadmap for pancreatic endocrinogenesis.}},
+    journal = {Development},
+    location = {Bidder Building, Station Rd, Histon, Cambridge Cb24 9lf, England},
+    publisher = {Company Biologists Ltd},
+    volume = {146},
+    number = {12},
+    year = {2019},
+    issn = {0950-1991},
+    eissn = {1477-9129},
+    doi = {10.1242/dev.173849},
+}
+
+@article{reuter:19,
+    author = {Reuter,Bernhard  and Fackeldey,Konstantin  and Weber,Marcus },
+    title = {Generalized Markov modeling of nonreversible molecular kinetics},
+    journal = {The Journal of Chemical Physics},
+    volume = {150},
+    number = {17},
+    pages = {174103},
+    year = {2019},
+    doi = {10.1063/1.5064530},
+}
+
+@article{robinson:10,
+    author = "Robinson, Mark D.
+    and Oshlack, Alicia",
+    title = "A scaling normalization method for differential expression analysis of RNA-seq data",
+    journal = "Genome Biology",
+    year = "2010",
+    month = "Mar",
+    day = "02",
+    volume = "11",
+    number = "3",
+    pages = "R25",
+    issn = "1474-760X",
+    doi = "10.1186/gb-2010-11-3-r25",
+}
+
+@article{schiebinger:19,
+    author = {Schiebinger, Geoffrey and Shu, Jian and Tabaka, Marcin and Cleary, Brian and Subramanian, Vidya and Solomon, Aryeh and Gould, Joshua and Liu, Siyan and Lin, Stacie and Berube, Peter and Lee, Lia and Chen, Jenny and Brumbaugh, Justin and Rigollet, Philippe and Hochedlinger, Konrad and Jaenisch, Rudolf and Regev, Aviv and Lander, Eric S.},
+    title = {Optimal-Transport Analysis of Single-Cell Gene Expression Identifies Developmental Trajectories in Reprogramming},
+    journal = {Cell},
+    year = {2019},
+    month = {Feb},
+    day = {07},
+    publisher = {Elsevier},
+    volume = {176},
+    number = {4},
+    pages = {928-943.e22},
+    issn = {0092-8674},
+    doi = {10.1016/j.cell.2019.01.006},
+}
+
+@article{setty:19,
+    author = "Setty, Manu and Kiseliovas, Vaidotas and Levine, Jacob and Gayoso, Adam and Mazutis, Linas and Pe'er, Dana",
+    title = "Characterization of cell fate probabilities in single-cell data with Palantir",
+    journal = "Nature Biotechnology",
+    year = "2019",
+    month = "Apr",
+    day = "01",
+    volume = "37",
+    number = "4",
+    pages = "451--460",
+    issn = "1546-1696",
+    doi = "10.1038/s41587-019-0068-4",
+}
+
+@article{stassen:21,
+	author = {Stassen, Shobana V. and Yip, Gwinky G. K. and Wong, Kenneth K. Y. and Ho, Joshua W. K. and Tsia, Kevin K.},
+	title = {VIA: Generalized and scalable trajectory inference in single-cell omics data},
+	elocation-id = {2021.02.10.430705},
+	year = {2021},
+	doi = {10.1101/2021.02.10.430705},
+	publisher = {Cold Spring Harbor Laboratory},
+	journal = {bioRxiv},
+}
+
+@article{velten:17,
+    author = "Velten, Lars and Haas, Simon F. and Raffel, Simon and Blaszkiewicz, Sandra and Islam, Saiful and Hennig, Bianca P. and Hirche, Christoph and Lutz, Christoph and Buss, Eike C. and Nowak, Daniel and Boch, Tobias and Hofmann, Wolf-Karsten and Ho, Anthony D. and Huber, Wolfgang and Trumpp, Andreas and Essers, Marieke A. G. and Steinmetz, Lars M.",
+    title = "Human haematopoietic stem cell lineage commitment is a continuous process",
+    journal = "Nature Cell Biology",
+    year = "2017",
+    month = "Apr",
+    day = "01",
+    volume = "19",
+    number = "4",
+    pages = "271--281",
+    issn = "1476-4679",
+    doi = "10.1038/ncb3493",
+}
+
+@article{brandts:02,
+    title={Matlab code for sorting real Schur forms},
+    author={Brandts, Jan H},
+    journal={Numerical linear algebra with applications},
+    volume={9},
+    number={3},
+    pages={249--261},
+    year={2002},
+    doi={10.1002/nla.274},
+    publisher={Wiley Online Library}
+}
+
+@article{weinreb:18,
+	author = {Weinreb, Caleb and Wolock, Samuel and Tusi, Betsabeh K. and Socolovsky, Merav and Klein, Allon M.},
+	title = {Fundamental limits on dynamic inference from single-cell snapshots},
+	volume = {115},
+	number = {10},
+	pages = {E2467--E2476},
+	year = {2018},
+	doi = {10.1073/pnas.1714723115},
+	publisher = {National Academy of Sciences},
+	issn = {0027-8424},
+	journal = {Proceedings of the National Academy of Sciences},
+}
+
+@article{wolf:18,
+    author = "Wolf, F. Alexander and Angerer, Philipp and Theis, Fabian J.",
+    title = "SCANPY: large-scale single-cell gene expression data analysis",
+    journal = "Genome Biology",
+    year = "2018",
+    month = "Feb",
+    day = "06",
+    volume = "19",
+    number = "1",
+    pages = "15",
+    issn = "1474-760X",
+    doi = "10.1186/s13059-017-1382-0",
+}
+
+@article{wolf:19,
+    author = "Wolf, F. Alexander and Hamey, Fiona K. and Plass, Mireya and Solana, Jordi and Dahlin, Joakim S. and G{\"o}ttgens, Berthold and Rajewsky, Nikolaus and Simon, Lukas and Theis, Fabian J.",
+    title = "PAGA: graph abstraction reconciles clustering with trajectory inference through a topology preserving map of single cells",
+    journal = "Genome Biology",
+    year = "2019",
+    month = "Mar",
+    day = "19",
+    volume = "20",
+    number = "1",
+    pages = "59",
+    issn = "1474-760X",
+    doi = "10.1186/s13059-019-1663-x",
+}
+
+@article{zhang:21,
+	author = {Zhang, Stephen and Afanassiev, Anton and Greenstreet, Laura and Matsumoto, Tetsuya and Schiebinger, Geoffrey},
+	title = {Optimal transport analysis reveals trajectories in steady-state systems},
+	elocation-id = {2021.03.02.433630},
+	year = {2021},
+	doi = {10.1101/2021.03.02.433630},
+	publisher = {Cold Spring Harbor Laboratory},
+	journal = {bioRxiv},
+}
+
+@online{desalvo:70,
+    author = {DeSalvo, J. S.},
+    title = {Standard Error of Forecast in Multiple Regression: Proof of a Useful Result.},
+    year = {1970},
+    url = {https://www.rand.org/pubs/papers/P4365.html},
+}
+
+@online{fischer:21,
+    author = {Fisher, R. A.},
+    title = {On the “probable error” of a coefficient of correlation deduced from a small sample.},
+    journal = {Metron},
+    volume = {1},
+    year = {1921},
+    url = {http://hdl.handle.net/2440/15169},
+}
+
+@article{mittnenzweig:21,
+    title = {A single-embryo, single-cell time-resolved model for mouse gastrulation},
+    journal = {Cell},
+    volume = {184},
+    number = {11},
+    pages = {2825-2842.e22},
+    year = {2021},
+    issn = {0092-8674},
+    doi = {10.1016/j.cell.2021.04.004},
+    author = {Markus Mittnenzweig and Yoav Mayshar and Saifeng Cheng and Raz Ben-Yair and Ron Hadas and Yoach Rais and
+              Elad Chomsky and Netta Reines and Anna Uzonyi and Lior Lumerman and Aviezer Lifshitz and Zohar Mukamel and
+              Ayelet-Hashahar Orenbuch and Amos Tanay and Yonatan Stelzer},
+    keywords = {scRNA-seq, developmental biology, mouse gastrulation, network flow model, trajectory inference,
+                cell fate decisions, tetraploid complementation assay, chimera assay},
+}
+
+@article{baran:19,
+    author="Baran, Yael and Bercovich, Akhiad and Sebe-Pedros, Arnau and Lubling, Yaniv and Giladi, Amir and
+            Chomsky, Elad and Meir, Zohar and Hoichman, Michael and Lifshitz, Aviezer and Tanay, Amos",
+    title = "MetaCell: analysis of single-cell RNA-seq data using K-nn graph partitions",
+    journal = "Genome Biology",
+    year = "2019",
+    month = "Oct",
+    day = "11",
+    volume = "20",
+    number = "1",
+    pages = "206",
+    issn = "1474-760X",
+    doi = "10.1186/s13059-019-1812-2",
+}
+
+@article{tirosh:16:nature,
+    author="Tirosh, Itay and Venteicher, Andrew S. and Hebert, Christine and Escalante, Leah E. and Patel, Anoop P. and
+            Yizhak, Keren and Fisher, Jonathan M. and Rodman, Christopher and Mount, Christopher and Filbin, Mariella G.
+            and Neftel, Cyril and Desai, Niyati and Nyman, Jackson and Izar, Benjamin and Luo, Christina C. and Francis,
+            Joshua M. and Patel, Aanand A. and Onozato, Maristela L. and Riggi, Nicolo and Livak, Kenneth J. and
+            Gennert, Dave and Satija, Rahul and Nahed, Brian V. and Curry, William T. and Martuza, Robert L. and
+            Mylvaganam, Ravindra and Iafrate, A. John and Frosch, Matthew P. and Golub, Todd R. and Rivera, Miguel N.
+            and Getz, Gad and Rozenblatt-Rosen, Orit and Cahill, Daniel P. and Monje, Michelle and Bernstein, Bradley E.
+            and Louis, David N. and Regev, Aviv and Suv{\`a}, Mario L.",
+    title="Single-cell RNA-seq supports a developmental hierarchy in human oligodendroglioma",
+    journal="Nature",
+    year="2016",
+    month="Nov",
+    day="01",
+    volume="539",
+    number="7628",
+    pages="309--313",
+    issn="1476-4687",
+    doi="10.1038/nature20123",
+}
+
+@article{tirosh:16:science,
+	author = {Tirosh, Itay and Izar, Benjamin and Prakadan, Sanjay M. and Wadsworth, Marc H. and Treacy,
+	          Daniel and Trombetta, John J. and Rotem, Asaf and Rodman, Christopher and Lian, Christine and Murphy,
+	          George and Fallahi-Sichani, Mohammad and Dutton-Regester, Ken and Lin, Jia-Ren and Cohen, Ofir and Shah,
+	          Parin and Lu, Diana and Genshaft, Alex S. and Hughes, Travis K. and Ziegler, Carly G. K. and Kazer,
+	          Samuel W. and Gaillard, Aleth and Kolb, Kellie E. and Villani, Alexandra-Chlo{\'e} and Johannessen,
+	          Cory M. and Andreev, Aleksandr Y. and Van Allen, Eliezer M. and Bertagnolli, Monica and Sorger, Peter K.
+	          and Sullivan, Ryan J. and Flaherty, Keith T. and Frederick, Dennie T. and Jan{\'e}-Valbuena, Judit and
+	          Yoon, Charles H. and Rozenblatt-Rosen, Orit and Shalek, Alex K. and Regev, Aviv and Garraway, Levi A.},
+	title = {Dissecting the multicellular ecosystem of metastatic melanoma by single-cell RNA-seq},
+	volume = {352},
+	number = {6282},
+	pages = {189--196},
+	year = {2016},
+	doi = {10.1126/science.aad0501},
+	publisher = {American Association for the Advancement of Science},
+	issn = {0036-8075},
+	journal = {Science}
+}
+
+@article{lange:20,
+    title={CellRank for directed single-cell fate mapping},
+    author={Lange, Marius and Bergen, Volker and Klein, Michal and Setty, Manu and Reuter, Bernhard and Bakhti, Mostafa
+            and Lickert, Heiko and Ansari, Meshal and Schniering, Janine and Schiller, Herbert B and others},
+    journal={BioRxiv},
+    year={2020},
+    doi={10.1101/2020.10.19.345983},
+    publisher={Cold Spring Harbor Laboratory}
+}
+
+@article{nowotschin:19,
+    title={The emergent landscape of the mouse gut endoderm at single-cell resolution},
+    author={Nowotschin, Sonja and Setty, Manu and Kuo, Ying-Yi and Liu, Vincent and Garg, Vidur and Sharma, Roshan and
+            Simon, Claire S and Saiz, Nestor and Gardner, Rui and Boutet, St{\'e}phane C and others},
+    journal={Nature},
+    volume={569},
+    number={7756},
+    pages={361--367},
+    year={2019},
+    doi={10.1038/s41586-019-1127-1},
+    publisher={Nature Publishing Group}
+}
diff --git a/docs/source/references.rst b/docs/source/references.rst
index 6b3313a13..52d4c0fb7 100644
--- a/docs/source/references.rst
+++ b/docs/source/references.rst
@@ -1,100 +1,5 @@
 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>`__.
-
-.. [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>`__.
-
-.. [Haghverdi16] Haghverdi *et al.* (2016),
-    *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*,
-    `bioRxiv <https://doi.org/10.1101/2020.09.19.304584>`__.
-
-.. [Lung20] Strunz, M. *et al.* (2020),
-    *Alveolar regeneration through a Krt8+ transitional stem cell state that persists in human lung fibrosis*,
-    `Nature Communications <https://doi.org/10.1038/s41467-020-17358-3>`__.
-
-.. [Macosko15] Macosko, E. Z. *et al.*, (2015),
-    *Highly Parallel Genome-wide Expression Profiling of Individual Cells Using Nanoliter Droplets*,
-    `Cell <https://doi.org/10.1016/j.cell.2015.05.002>`__.
-
-.. [MAGIC18] Van Dijk, D. *et al.* (2018),
-    *Recovering Gene Interactions from Single-Cell Data Using Data Diffusion*,
-    `Cell <https://doi.org/10.1016/j.cell.2018.05.061>`__.
-
-.. [Manno18] La Manno *et al.* (2018),
-    *RNA velocity of single cells*,
-    `Nature <https://doi.org/10.1038/s41586-018-0414-6>`__.
-
-.. [Morris18] Biddy, B.A., Kong, W., Kamimoto, K. *et al.* (2018),
-    *Single-cell mapping of lineage and identity in direct reprogramming.*,
-    `Nature 564, 219–224 <https://doi.org/10.1038/s41586-018-0744-4>`__.
-
-.. [Palla21] Palla *et al.* (2021),
-    *Squidpy: a scalable framework for spatial single cell analysis*,
-    `bioRxiv <https://doi.org/10.1101/2021.02.19.431994>`__.
-
-.. [Panc19] Bastidas-Ponce *et al.* (2019),
-    *Comprehensive single cell mRNA profiling reveals a detailed roadmap for pancreatic endocrinogenesis*,
-    `Development <https://doi.org/10.1242/dev.173849>`__.
-
-.. [Reuter19] Reuter, B. *et al.* (2019),
-    *Generalized Markov modeling of nonreversible molecular kinetics*,
-    `The Journal of Chemical Physics <https://doi.org/10.1063/1.5064530>`__.
-
-.. [Robinson10] Robinson, M. D. *et al.* (2010),
-    *A scaling normalization method for differential expression analysis of RNA-seq data*,
-    `Genome Biology <https://doi.org/10.1186/gb-2010-11-3-r25>`__.
-
-.. [Setty19] Setty *et al.* (2019),
-    *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>`__.
-
-.. [Tolver16] Tolver (2016),
-    *An introduction to Markov chains*,
-    `Recuperado el 15 <http://old.math.ku.dk/noter/filer/stoknoter.pdf>`__.
-
-.. [Weinreb18] Weinreb *et al.* (2018),
-    *Fundamental limits on dynamic inference from single-cell snapshots*,
-    `PNAS <https://doi.org/10.1073/pnas.1714723115>`__.
-
-.. [Wolf18] Wolf *et al.* (2018),
-    *Scanpy: large-scale single-cell gene expression data analysis*,
-    `Genome Biology <https://doi.org/10.1186/s13059-017-1382-0>`__.
-
-.. [Wolf19] Wolf *et al.* (2019),
-    *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>`__.
+.. bibliography::
+    :cited:
diff --git a/docs/source/release_notes.rst b/docs/source/release_notes.rst
index 264b5e2e7..02ef3cc3b 100644
--- a/docs/source/release_notes.rst
+++ b/docs/source/release_notes.rst
@@ -6,9 +6,67 @@ Release Notes
 Version 1.0
 -----------
 
-1.2.0 :small:`2021-02-02`
+1.3.1 :small:`2021-04-09`
+~~~~~~~~~~~~~~~~~~~~~~~~~
+.. rubric:: Bugfixes
+
+- Fix estimator's lineages color/names mismatch when reading from :class:`anndata.AnnData`
+  `PR 556 <https://github.com/theislab/cellrank/pull/556>`_.
+- Remove heuristics used to determine which solver to use `PR 558 <https://github.com/theislab/cellrank/pull/558>`_.
+
+1.3.0 :small:`2021-03-29`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
+This release includes some major additions which make CellRank more applicable with and without RNA velocity
+information. In particular, it includes:
+
+.. rubric:: Additions
+
+- Add new kernel :class:`cellrank.tl.kernels.CytoTRACEKernel` which computes cell-cell transition probabilities based
+  on the CytoTRACE score :cite:`gulati:20`, 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 :cite:`zhang:21`, 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 :cite:`setty:19`, a soft thresholding scheme inspired by :cite:`stassen:21` 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` :cite:`palla:21`
+  `PR 501 <https://github.com/theislab/cellrank/pull/501>`_.
+- Revamp `Pancreas Advanced <https://cellrank.readthedocs.io/en/stable/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:
+
+  - `Beyond RNA velocity <https://cellrank.readthedocs.io/en/stable/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/stable/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 :cite:`setty:19`
+  and :cite:`velten:17`.
+- Add checks for transition matrix irreducibility `PR 516 <https://github.com/theislab/cellrank/pull/516>`_.
+- Add Zebrafish development dataset from :cite:`farrell:18` `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
+
+- 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>`_.
 
+1.2.0 :small:`2021-02-02`
+~~~~~~~~~~~~~~~~~~~~~~~~~
 This release includes:
 
 .. rubric:: Additions
@@ -18,7 +76,7 @@ This release includes:
   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]_,
+- Add :func:`cellrank.pl.circular_projection` visualizing computed fate probabilities as done in :cite:`velten:17`,
   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>`_
@@ -39,7 +97,6 @@ This release includes:
 
 1.1.0 :small:`2020-11-17`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
-
 This release includes:
 
 .. rubric:: Additions
@@ -49,12 +106,12 @@ This release includes:
   or an exact, permutation based test. Corrects for multiple-testing.
 - :meth:`cellrank.tl.kernels.VelocityKernel.compute_transition_matrix` now allows different metrics to be used to
   compare velocity vectors with expression-differences across neighboring cells. We add cosine-correlation and
-  dot-product schemes and we allow the user to input their own scheme. It has been shown recently by [Li2020]_
+  dot-product schemes and we allow the user to input their own scheme. It has been shown recently by :cite:`li:20`
   that the choice of metric can lead to slightly different results. Users can now also supply their own scheme as long
   as it follows the signature of :class:`cellrank.tl.kernels.SimilaritySchemeABC`.
 - :func:`cellrank.datasets.reprogramming` has been added to allow for easy reproducibility of the time & memory
   benchmarking results in our `CellRank preprint <https://doi.org/10.1101/2020.10.19.345983>`_. This is a reprogramming
-  dataset from [Morris18]_.
+  dataset from :cite:`morris:18`.
 
 .. rubric:: Bugfixes
 
@@ -63,7 +120,7 @@ This release includes:
 
 
 1.0.0 :small:`2020-10-17`
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~
 - Fix a bug when subsetting :class:`cellrank.tl.Lineage`
 - Add renaming terminal states :meth:`cellrank.tl.estimators.BaseEstimator.rename_terminal_states`
 - Enable negative binomial distribution for :class:`cellrank.ul.models.GAMR`
@@ -71,12 +128,10 @@ This release includes:
 - Add :class:`cellrank.ul.models.FailedModel` inspired by the maybe monad
 - Allow returning models when doing bulk fitting
 - Add ``transpose`` parameter for :func:`cellrank.pl.gene_trends`
-- Various other minor bugfixes
+- Various minor bugfixes
 
 1.0.0-rc.11 :small:`2020-09-25`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-- Rename ``metastable states`` to ``macrostates``
 - Rename ``final states`` to ``terminal states``
 - Fix pickling if :class:`cellrank.tl.estimators.BaseEstimator`
 - Fix various color bugs
@@ -86,25 +141,21 @@ This release includes:
 
 1.0.0-rc.0 :small:`2020-07-15`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-- Fix pickling of :class:`cellrank.tl.Lineage` improvements
+- Fix pickling of :class:`cellrank.tl.Lineage`
 - Add additional options to :func:`cellrank.pl.heatmap`
 - Updated documentation
 
 1.0.0-b.8 :small:`2020-07-12`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
 - Add installation options for PETSc and SLEPc
-- Add iterative solver for absorption proabilities
+- Add iterative solver for absorption probabilities
 - Add minor :class:`cellrank.tl.Lineage` improvements
 - Fix docstring issues
 
 1.0.0-b.2 :small:`2020-07-02`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-- Fix installation by including future-fstrings.
+- Fix installation by including future-fstrings
 
 1.0.0-b.1 :small:`2020-07-02`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-- Initial beta pre-release.
+- Initial beta pre-release
diff --git a/examples/README.rst b/examples/README.rst
index d12df612c..16e6d0b7a 100644
--- a/examples/README.rst
+++ b/examples/README.rst
@@ -1,4 +1,3 @@
 Examples
 ========
-
 This section contains various code snippets that demonstrate :mod:`cellrank`'s usage.
diff --git a/examples/estimators/README.rst b/examples/estimators/README.rst
index d317c6486..07916d6ec 100644
--- a/examples/estimators/README.rst
+++ b/examples/estimators/README.rst
@@ -1,4 +1,3 @@
 Estimators
 ----------
-
-Below is a gallery of examples for :class:`cellrank.tl.estimators.GPCCA` and :class:`cellrank.tl.estimators.CFLARE`.
+Below are examples of :class:`cellrank.tl.estimators.GPCCA` and :class:`cellrank.tl.estimators.CFLARE`.
diff --git a/examples/estimators/compute_abs_probs.py b/examples/estimators/compute_abs_probs.py
index cec39721d..bb090b6ca 100644
--- a/examples/estimators/compute_abs_probs.py
+++ b/examples/estimators/compute_abs_probs.py
@@ -11,7 +11,6 @@
 """
 
 import scvelo as scv
-
 import cellrank as cr
 
 adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
diff --git a/examples/estimators/compute_coarse_T.py b/examples/estimators/compute_coarse_T.py
index 0baef8f64..7416fb644 100644
--- a/examples/estimators/compute_coarse_T.py
+++ b/examples/estimators/compute_coarse_T.py
@@ -4,9 +4,9 @@
 
 This example shows how to compute and plot the coarse-grained transition matrix among the set of macrostates.
 
-We computed the macrostates using Generalized Perron Cluster Cluster Analysis [GPCCA18]_ [Reuter19]_, see the example
-:ref:`sphx_glr_auto_examples_estimators_compute_macrostates.py`. The coarse-grained transition matrix shows
-transitional behavior among the set of macrostates and may be used to classify these states as initial,
+We computed the macrostates using Generalized Perron Cluster Cluster Analysis :cite:`reuter:18` :cite:`reuter:19`,
+see the example :ref:`sphx_glr_auto_examples_estimators_compute_macrostates.py`. The coarse-grained transition matrix
+shows transitional behavior among the set of macrostates and may be used to classify these states as initial,
 intermediate or terminal.
 """
 
@@ -23,8 +23,7 @@
 g = cr.tl.estimators.GPCCA(k)
 
 # %%
-# Next, we compute the Schur vectors. See :ref:`sphx_glr_auto_examples_estimators_compute_schur_vectors.py` for
-# more information.
+# Next, we compute the Schur vectors.
 g.compute_schur(n_components=6)
 
 # %%
diff --git a/examples/estimators/compute_eigendecomposition.py b/examples/estimators/compute_eigendecomposition.py
deleted file mode 100644
index 792b4abf1..000000000
--- a/examples/estimators/compute_eigendecomposition.py
+++ /dev/null
@@ -1,50 +0,0 @@
-"""
-Compute eigendecomposition
---------------------------
-
-This example shows how to compute and plot the spectrum of a transition matrix.
-
-The spectrum of a rectangular matrix :math:`T` is given by its `eigenvalues
-<https://en.wikipedia.org/wiki/Eigendecomposition_of_a_matrix>`_. For a general, real valued matrix :math:`T`,
-eigenvalues will either be real or appear in pairs of complex conjugates. CellRank can plot the spectrum in the complex
-plane or restrict the visualization to the real part of the eigenvalues, which may be helpful to spot the
-`eigengap <https://en.wikipedia.org/wiki/Eigengap>`_.
-"""
-
-import cellrank as cr
-
-adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
-adata
-
-# %%
-# First, we prepare the kernel using the high-level pipeline and the :class:`cellrank.tl.estimators.GPCCA` estimator.
-k = cr.tl.transition_matrix(
-    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
-)
-g = cr.tl.estimators.GPCCA(k)
-
-# %%
-# In this example, we visualize the eigenvectors in an embedding, therefore we specify ``only_evals=False``,
-# but for the method :meth:`cellrank.tl.estimators.GPCCA.plot_spectrum`, they are not necessary, because
-# there we plot only the eigenvalues.
-g.compute_eigendecomposition(k=20, only_evals=False)
-g.eigendecomposition
-
-# %%
-# Let's start first by plotting the real spectrum. The `eigengap` is shown by the vertical line. Below
-# we plot only the first 10 real eigenvalues, sorted by their real part.
-g.plot_spectrum(10, real_only=True)
-
-# %%
-# Next, we plot the eigenvalues in the complex plane.
-g.plot_spectrum(real_only=False)
-
-# %%
-# Finally, we can plot the left and right eigenvectors in an embedding. We can also specify which or how many
-# eigenvectors to plot using the parameter ``use``. If not specified, vectors up to the `eigengap` are shown.
-g.plot_eigendecomposition(left=False)
-g.plot_eigendecomposition(left=True, use=2)
-
-# %%
-# Left and right eigenvectors are used in :class:`cellrank.tl.estimators.CFLARE` estimator to compute the initial or
-# the states of the process.
diff --git a/examples/estimators/compute_fit.py b/examples/estimators/compute_fit.py
deleted file mode 100644
index 810955cb5..000000000
--- a/examples/estimators/compute_fit.py
+++ /dev/null
@@ -1,39 +0,0 @@
-"""
-Fit estimator
--------------
-
-This example shows how to fit an estimator in order to compute the lineages.
-
-CellRank is composed of :mod:`cellrank.tl.kernels` and :mod:`cellrank.tl.estimators`. Kernels compute transition
-matrices on the basis of directed data, given by i.e. RNA velocity [Manno18]_ [Bergen20]_ while estimators perform
-inference making use of kernels.
-
-Here, we show how to fit an estimator, given a kernel. An estimator may be used e.g. to identify initial and terminal
-states or to compute lineage probabilities towards the terminal states, for each individual cell.
-The :meth:`cellrank.tl.estimators.BaseEstimator.fit` method is a combination of individual methods to compute terminal
-states and absorption probabilities.
-
-To have maximum control over the inference performed, use these individual methods directly.
-"""
-
-import cellrank as cr
-
-adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
-adata
-
-# %%
-# First, we prepare the kernel using high-level pipeline and the :class:`cellrank.tl.estimators.GPCCA` estimator.
-k = cr.tl.transition_matrix(
-    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
-)
-g = cr.tl.estimators.GPCCA(k)
-
-# %%
-# Afterwards, we simply call :meth:`cellrank.tl.estimators.BaseEstimator.fit`. It offers a quick and easy way
-# to compute the terminal states and optionally the absorption probabilities by following similar steps as defined in
-# :ref:`sphx_glr_auto_examples_estimators_compute_terminal_states_gpcca.py`.
-g.fit(n_lineages=3, cluster_key="clusters", compute_absorption_probabilities=True)
-
-# %%
-# In order to verify that the absorption probabilities have been computed, we plot them below.
-g.plot_absorption_probabilities()
diff --git a/examples/estimators/compute_macrostates.py b/examples/estimators/compute_macrostates.py
index 45daf6b9a..70d9234b5 100644
--- a/examples/estimators/compute_macrostates.py
+++ b/examples/estimators/compute_macrostates.py
@@ -4,12 +4,10 @@
 
 This example shows how to compute and plot the macrostates.
 
-For the computation of macrostates, we adapted the popular Generalized Perron Cluster Cluster Analysis [GPCCA18]_
-[Reuter19]_ method to the single cell context. We provide a scalable implementation which can decompose datasets of
-100k+ cells into their dominant dynamical macrostates in just a few minutes. GPCCA relies on the real Schur
-decomposition to handle non-symmetric transition matrices as they arise from RNA velocity information, see
-:ref:`sphx_glr_auto_examples_estimators_compute_schur_vectors.py` and
-:ref:`sphx_glr_auto_examples_estimators_compute_schur_matrix.py`.
+For the computation of macrostates, we adapted the popular Generalized Perron Cluster Cluster Analysis :cite:`reuter:18`
+:cite:`reuter:19` method to the single cell context. We provide a scalable implementation which can decompose datasets
+of 100k+ cells into their dominant dynamical macrostates in just a few minutes. GPCCA relies on the real Schur
+decomposition to handle non-symmetric transition matrices as they arise from RNA velocity information.
 """
 
 import cellrank as cr
@@ -35,7 +33,7 @@
 # on the membership degree. By default, 30 cells are selected.
 #
 # Lastly, the parameter ``n_states`` can also be estimated by using either the `eigengap` or the `minChi` criterion from
-# [GPCCA18]_.
+# :cite:`reuter:18`.
 g.compute_macrostates(n_states=3, cluster_key="clusters")
 
 # %%
diff --git a/examples/estimators/compute_schur_matrix.py b/examples/estimators/compute_schur_matrix.py
deleted file mode 100644
index 30c5e7da8..000000000
--- a/examples/estimators/compute_schur_matrix.py
+++ /dev/null
@@ -1,37 +0,0 @@
-"""
-Compute Schur matrix
---------------------
-
-This examples show how to compute and plot the `Schur matrix <https://en.wikipedia.org/wiki/Schur_decomposition>`_.
-
-The Schur matrix is equivalent to the diagonal matrix of eigenvalues for diagonalizable matrices. For the real
-Schur decomposition, it is a real, quasi-upper triangular matrix with 1x1 and 2x2 blocks on the diagonal.
-For a rectangular, real matrix :math:`T`, it holds: :math:`T = Q R Q^T`, where :math:`R` is the real Schur matrix and
-:math:`Q` is an orthogonal matrix containing the real Schur vectors as columns. 1x1 blocks on the diagonal of :math:`R`
-correspond to real eigenvalues of :math:`T` whereas 2x2 blocks on the diagonal of :math:`R` correspond to pairs of
-complex conjugate eigenvalues of :math:`T`.
-"""
-
-import cellrank as cr
-
-adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
-adata
-
-# %%
-# First, we prepare the kernel using the high-level pipeline and the :class:`cellrank.tl.estimators.GPCCA` estimator.
-k = cr.tl.transition_matrix(
-    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
-)
-g = cr.tl.estimators.GPCCA(k)
-
-# %%
-# To compute the Schur vectors, run the code below. Parameter ``n_components`` controls how many vectors to compute.
-g.compute_schur(n_components=6)
-
-# %%
-# Finally, we are ready to plot the Schur matrix. The real Schur matrix is quasi-upper triangular,
-# which means there may be 1x1 or 2x2 blocks on the diagonal.
-g.plot_schur_matrix()
-
-# %%
-# For more information about the Schur vectors, see :ref:`sphx_glr_auto_examples_estimators_compute_schur_vectors.py`.
diff --git a/examples/estimators/compute_schur_vectors.py b/examples/estimators/compute_schur_vectors.py
deleted file mode 100644
index 1ded786ca..000000000
--- a/examples/estimators/compute_schur_vectors.py
+++ /dev/null
@@ -1,43 +0,0 @@
-"""
-Compute Schur vectors
----------------------
-
-This example shows how to compute and plot the `Schur vectors <https://en.wikipedia.org/wiki/Schur_decomposition>`_.
-
-In CellRank, we work with the real Schur decomposition because the transition matrices we are concerned with are non
-symmetrical (think of the velocity flow). In general, for these matrices the eigendecomposition does not exist.
-You may still compute the top eigenvectors, however, these will in general be complex valued and hard to interpret.
-The Schur decomposition is the natural generalization of the eigendecomposition for non-symmetric matrices and it’s
-real variant ensures that the top Schur vectors are only real and may be interpreted in the context of metastabilities.
-"""
-
-import cellrank as cr
-
-adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
-adata
-
-# %%
-# First, we prepare the kernel using the high-level pipeline and the :class:`cellrank.tl.estimators.GPCCA` estimator.
-k = cr.tl.transition_matrix(
-    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
-)
-g = cr.tl.estimators.GPCCA(k)
-
-# %%
-# To compute the Schur vectors, simply run the code below. The parameter ``n_components`` controls how many vectors
-# to compute. We can also specify prior initial distribution over all cells as ``initial_distribution``, by default
-# it is uniform.
-g.compute_schur(n_components=4)
-
-# %%
-# Lastly, we plot the vectors in the UMAP embedding (default).
-g.plot_schur()
-
-# %%
-# Note that above, only 3 vectors are shown, because the 1st Schur vector is a constant vector of ones.
-# This can be verified visually by plotting the first Schur vector.
-g.plot_schur(use=[0])
-
-# %%
-# Method :meth:`cellrank.tl.estimators.GPCCA.compute_schur` also computes the Schur matrix,
-# see :ref:`sphx_glr_auto_examples_estimators_compute_schur_matrix.py`.
diff --git a/examples/estimators/compute_terminal_states_gpcca.py b/examples/estimators/compute_terminal_states_gpcca.py
index 81aabdf1e..ab8ce6d67 100644
--- a/examples/estimators/compute_terminal_states_gpcca.py
+++ b/examples/estimators/compute_terminal_states_gpcca.py
@@ -4,8 +4,8 @@
 
 This example shows how to compute and plot the terminal states using the :class:`cellrank.tl.estimators.GPCCA`.
 
-This estimator makes use of Generalized Perron Cluster Cluster Analysis [GPCCA18]_ [Reuter19]_ as implemented in
-`pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_.
+This estimator makes use of Generalized Perron Cluster Cluster Analysis :cite:`reuter:18` :cite:`reuter:19`
+as implemented in `pyGPCCA <https://pygpcca.readthedocs.io/en/latest/>`_.
 """
 
 import cellrank as cr
@@ -62,7 +62,7 @@
 # Compute terminal states
 # ^^^^^^^^^^^^^^^^^^^^^^^
 # Lastly, :meth:`cellrank.tl.estimators.GPCCA.compute_terminal_states` which also makes use of the coarse-grained
-# transition matrix :paramref:`cellrank.tl.estimators.GPCCA.coarse_T` of the macrostates or the `eigengap`
+# transition matrix :attr:`cellrank.tl.estimators.GPCCA.coarse_T` of the macrostates or the `eigengap`
 # statistic.
 #
 # In the example below, we use ``method='eigenap'`` which selects the number of states based on the `eigengap`. The
diff --git a/examples/kernels/README.rst b/examples/kernels/README.rst
new file mode 100644
index 000000000..7b25f0a82
--- /dev/null
+++ b/examples/kernels/README.rst
@@ -0,0 +1,3 @@
+Kernels
+-------
+Below are examples of :class:`cellrank.tl.kernels.Kernel`.
diff --git a/examples/other/compute_kernel_tricks.py b/examples/kernels/compute_kernel_tricks.py
similarity index 82%
rename from examples/other/compute_kernel_tricks.py
rename to examples/kernels/compute_kernel_tricks.py
index 79ee0fac6..18d7d4c98 100644
--- a/examples/other/compute_kernel_tricks.py
+++ b/examples/kernels/compute_kernel_tricks.py
@@ -5,8 +5,8 @@
 This example shows some niche, but useful functionalities of :class:`cellrank.tl.kernels.Kernel`.
 
 CellRank is split into :mod:`cellrank.tl.kernels` and :mod:`cellrank.tl.estimators`. Kernels compute transition matrices
-based on some inputs, like RNA velocity [Bergen20]_, [Manno18]_, while estimators perform inference based on a given
-kernel, e.g. they compute initial and terminal cells and fate probabilities.
+based on some inputs, like RNA velocity :cite:`manno:18` :cite:`bergen:20`, while estimators perform inference based on
+a given kernel, e.g. they compute initial and terminal cells and fate probabilities.
 
 Here, will will dive a bit deeper into the how these kernel objects work.
 """
@@ -20,13 +20,13 @@
 # First, we create some kernels which will be used to compute the cell-to-cell transition matrix:
 #
 # - :class:`cellrank.tl.kernels.ConnectivityKernel` computes the transition matrix using the KNN graph from
-#   :func:`scanpy.pp.neighbors` [Wolf18]_. Note that this kernel is by itself directionless and
+#   :func:`scanpy.pp.neighbors` :cite:`wolf:18`. Note that this kernel is by itself directionless and
 #   should be used in conjunction with e.g. :class:`cellrank.tl.kernels.VelocityKernel`.
-# - :class:`cellrank.tl.kernels.VelocityKernel` uses RNA velocity for the transition matrix computation [Manno18]_
-#   [Bergen20]_, but can also take into account uncertainty in RNA velocity.
-# - :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]_.
+# - :class:`cellrank.tl.kernels.VelocityKernel` uses RNA velocity for the transition matrix computation :cite:`manno:18`
+#   :cite:`bergen:20`, but can also take into account uncertainty in RNA velocity.
+# - :class:`cellrank.tl.kernels.PseudotimeKernel` works similarly as in Palantir :cite:`setty:19` - it orients the edges
+#   of the KNN graph constructed in the expression space using the pseudotemporal ordering of cells, such as
+#   Diffusion Pseudotime (DPT) :cite:`haghverdi:16`.
 ck = cr.tl.kernels.ConnectivityKernel(adata)
 vk = cr.tl.kernels.VelocityKernel(adata)
 pk = cr.tl.kernels.PseudotimeKernel(adata)
@@ -54,7 +54,7 @@
 # since the handles point to the original objects.
 #
 # Note that in the 2nd :func:`print` statement, we access the private attribute - that's because accessing
-# :paramref:`cellrank.tl.kernels.Kernel.transition_matrix` computes the transition matrix with default values.
+# :attr:`cellrank.tl.kernels.Kernel.transition_matrix` computes the transition matrix with default values.
 # This happens only with basic kernels and not the kernel expressions and only if they are not part of a
 # larger expression.
 ck.compute_transition_matrix()
diff --git a/examples/kernels/plot_projection.py b/examples/kernels/plot_projection.py
new file mode 100644
index 000000000..69caf3c7c
--- /dev/null
+++ b/examples/kernels/plot_projection.py
@@ -0,0 +1,34 @@
+"""
+Plot projection
+---------------
+
+This example shows how to project a transition matrix onto a low-dimensional embedding.
+"""
+
+import scvelo as scv
+import cellrank as cr
+
+adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
+adata
+
+# %%
+# First, we create the kernel using the high-level function :func:`cellrank.tl.transition_matrix`, but any
+# instance of :class:`cellrank.tl.kernels.Kernel` would do.
+k = cr.tl.transition_matrix(
+    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
+)
+k
+
+# %%
+# Next, we compute the projection of the transition matrix onto an embedding, in this case UMAP.
+# The projected transition matrix can be found in ``adata.obsm['T_fwd_umap']``.
+k.compute_projection(basis="umap")
+
+# %%
+# The projected transition matrix can now be visualized using :mod:`scvelo`.
+# Below we use :func:`scvelo.pl.velocity_embedding_stream` to visualize it as streamlines.
+scv.pl.velocity_embedding_stream(adata, vkey="T_fwd", basis="umap")
+
+# %%
+# Another way of visualizing the transition matrix is to simulate random walks on Markov chain and
+# plot them in an embedding, see :ref:`sphx_glr_auto_examples_kernels_plot_random_walks.py` for more information.
diff --git a/examples/kernels/plot_random_walks.py b/examples/kernels/plot_random_walks.py
new file mode 100644
index 000000000..73c969d8b
--- /dev/null
+++ b/examples/kernels/plot_random_walks.py
@@ -0,0 +1,51 @@
+"""
+Plot random walks
+-----------------
+
+This example shows how to plot random walks on a Markov chain in an embedding.
+"""
+
+import cellrank as cr
+
+adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
+adata
+
+# %%
+# First, we create the kernel using the high-level function :func:`cellrank.tl.transition_matrix`, but any
+# instance of :class:`cellrank.tl.kernels.Kernel` would do.
+k = cr.tl.transition_matrix(
+    adata, weight_connectivities=0.2, softmax_scale=4, show_progress_bar=False
+)
+k
+
+# %%
+# Next, we simulate and plot random walks. Below we simulate 100 random walks of length 100 where starting points
+# are uniformly sampled from the `Ngn3 low EP` cluster.
+k.plot_random_walks(
+    100,
+    start_ixs={"clusters": "Ngn3 low EP"},
+    max_iter=100,
+    show_progress_bar=False,
+    ixs_legend_loc="best",
+    seed=42,
+)
+
+
+# %%
+# Apart from ``start_ixs``, we can also specify ``stop_ixs`` and ``successive_hits`` to control when random walks
+# stop. Random walk is stopped if the maximum number of iterations is reached or when states in ``stop_ixs``
+# is visited successively ``successive_hits`` times.
+k.plot_random_walks(
+    100,
+    start_ixs={"clusters": "Ngn3 low EP"},
+    stop_ixs={"clusters": ["Alpha", "Beta"]},
+    max_iter=100,
+    successive_hits=5,
+    show_progress_bar=False,
+    cmap="viridis",
+    seed=42,
+)
+
+# %%
+# Another way of visualizing the transition matrix is to project it onto an embedding,
+# see :ref:`sphx_glr_auto_examples_kernels_plot_projection.py` for more information.
diff --git a/examples/other/README.rst b/examples/other/README.rst
index 9773eba5a..dc334a9d5 100644
--- a/examples/other/README.rst
+++ b/examples/other/README.rst
@@ -1,5 +1,4 @@
-Others
-------
-
-Below are examples not covered in other examples such as :class:`cellrank.tl.Lineage`
-or :class:`cellrank.ul.models.BaseModel`.
+Other
+-----
+Below are examples not covered in other sections, such as
+:class:`cellrank.tl.Lineage` or :class:`cellrank.ul.models.BaseModel`.
diff --git a/examples/other/compute_lineage_tricks.py b/examples/other/compute_lineage_tricks.py
index eb1c3f6e4..e38ed1dbc 100644
--- a/examples/other/compute_lineage_tricks.py
+++ b/examples/other/compute_lineage_tricks.py
@@ -10,13 +10,13 @@
 
 We use it primarily to store either the fate probabilities or the macrostates memberships, see
 :ref:`sphx_glr_auto_examples_estimators_compute_abs_probs.py` or
-:ref:`sphx_glr_auto_examples_estimators_compute_macrostates.py`, to learn how to compute them.
+:ref:`sphx_glr_auto_examples_estimators_compute_macrostates.py` to learn how to compute them.
 """
 
-import numpy as np
-
 import cellrank as cr
 
+import numpy as np
+
 np.random.seed(42)
 
 # %%
@@ -33,7 +33,7 @@
 
 # %%
 # In some cases, this behavior is not desirable or can have unintended consequences. To access the underlying
-# :class:`numpy` array, use the :paramref:`cellrank.tl.Lineage.X` attribute.
+# :class:`numpy` array, use the :attr:`cellrank.tl.Lineage.X` attribute.
 lin.X
 
 # %%
@@ -51,12 +51,12 @@
 
 # %%
 # Most of the :mod:`numpy` methods are supported by the :class:`cellrank.tl.Lineage`. One can also calculate the
-# entropy, which in [Setty19]_ is defined as the differentiation potential of cells.
+# entropy, which in :cite:`setty:19` is defined as the differentiation potential of cells.
 lin.entropy(axis=1)
 
 # %%
 # When subsetting the lineage and not selecting all of them, they will no longer sum to 1 and cannot be
-# interpreted as a probability distribution. We offer a method :paramref:`cellrank.tl.Lineage.reduce` which
+# interpreted as a probability distribution. We offer a method :attr:`cellrank.tl.Lineage.reduce` which
 # can be used to solve this issue. Below we show only one out of many normalization techniques.
 lin.reduce("foo, quux", "baz", normalize_weights="softmax")
 
diff --git a/examples/other/plot_model.py b/examples/other/plot_model.py
index 866a792ad..ffdd64832 100644
--- a/examples/other/plot_model.py
+++ b/examples/other/plot_model.py
@@ -10,10 +10,10 @@
 relatively straightforward to derive a confidence interval around the main trend.
 """
 
-from sklearn.svm import SVR
-
 import cellrank as cr
 
+from sklearn.svm import SVR
+
 adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
 adata
 
@@ -46,11 +46,11 @@
 )
 # %%
 # CellRank allows any pseudotime from ``adata.obs`` to be passed via the ``time_key``, e.g. Diffusion Pseudotime (DPT)
-# [Haghverdi16]_, :mod:`scvelo`’s latent time [Bergen20]_, Palantir’s pseudotime [Setty19]_, etc.
+# :cite:`haghverdi:16`, :mod:`scvelo`’s latent time :cite:`bergen:20`, Palantir’s pseudotime :cite:`setty:19`, etc.
 #
 # Further, CellRank accepts imputed gene expression values stored in ``adata.layers`` via the ``data_key``, i.e. you
-# can pass MAGIC [MAGIC18]_ imputed data, :func:`scvelo.pp.moments` [Bergen20]_ (used below) or any other form of
-# imputation.
+# can pass MAGIC :cite:`vandijk:18` imputed data, :func:`scvelo.pp.moments` :cite:`bergen:20` (used below) or any other
+# form of imputation.
 
 # %%
 # Once the model has been prepared, it is ready for fitting and prediction.
diff --git a/examples/plotting/README.rst b/examples/plotting/README.rst
index c2e10bbae..67673b94a 100644
--- a/examples/plotting/README.rst
+++ b/examples/plotting/README.rst
@@ -1,4 +1,3 @@
 Plotting
 --------
-
-Below are examples for :mod:`cellrank.pl` plotting module.
+Below are examples of :mod:`cellrank.pl` plotting module.
diff --git a/examples/plotting/plot_circular_embedding.py b/examples/plotting/plot_circular_embedding.py
index b63fe07cf..f01e8cfec 100644
--- a/examples/plotting/plot_circular_embedding.py
+++ b/examples/plotting/plot_circular_embedding.py
@@ -2,7 +2,7 @@
 Plot circular embedding
 -----------------------
 
-This example shows how to plot absorption probabilities using circular a posteriori projection as used in [Velten17]_.
+This example shows how to plot absorption probabilities using circular a posteriori projection :cite:`velten:17`.
 """
 
 import cellrank as cr
diff --git a/examples/plotting/plot_cluster_fates.py b/examples/plotting/plot_cluster_fates.py
index 2d1e410b4..1349796aa 100644
--- a/examples/plotting/plot_cluster_fates.py
+++ b/examples/plotting/plot_cluster_fates.py
@@ -7,14 +7,13 @@
 """
 
 import scanpy as sc
-
 import cellrank as cr
 
 adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
 adata
 
 # %%
-# First, we compute the terminal states and the lineages, as well as :mod:`scanpy`'s PAGA [Wolf19]_.
+# First, we compute the terminal states and the lineages, as well as :mod:`scanpy`'s PAGA :cite:`wolf:19`.
 cr.tl.terminal_states(
     adata,
     cluster_key="clusters",
diff --git a/examples/plotting/plot_directed_paga.py b/examples/plotting/plot_directed_paga.py
index bacc594c8..78e631d0d 100644
--- a/examples/plotting/plot_directed_paga.py
+++ b/examples/plotting/plot_directed_paga.py
@@ -2,16 +2,15 @@
 Plot directed PAGA
 ------------------
 
-This example shows how to compute and plot a directed version of the popular PAGA algorithm [Wolf19]_.
+This example shows how to compute and plot a directed version of the popular PAGA algorithm :cite:`wolf:19`.
 
 In classical PAGA plot, nodes correspond to clusters and edge thickness denotes transcriptomic similarity.
-We introduce a new directed version of PAGA where directed edges reflect local velocity flow. We add the possibility
-to include prior information in the form of a pseudotemporal ordering or initial/terminal state annotation to restrict
-the possible edge set. We further replace nodes by pie charts that show average CellRank fate probabilities.
+We introduce a new directed version of PAGA where directed edges reflect local velocity compute_flow. We add the
+possibility to include prior information in the form of a pseudotemporal ordering or initial/terminal state annotation
+to restrict the possible edge set. We further replace nodes by pie charts that show average CellRank fate probabilities.
 """
 
 import scvelo as scv
-
 import cellrank as cr
 
 adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
diff --git a/examples/plotting/plot_graph.py b/examples/plotting/plot_graph.py
index 5941e0011..90cb6bdcc 100644
--- a/examples/plotting/plot_graph.py
+++ b/examples/plotting/plot_graph.py
@@ -5,10 +5,10 @@
 This functions show how to plot graph structures, such as the transition matrix.
 """
 
-import numpy as np
-
 import cellrank as cr
 
+import numpy as np
+
 adata = cr.datasets.pancreas_preprocessed("../example.h5ad")
 adata
 
diff --git a/examples/plotting/plot_lineages.py b/examples/plotting/plot_lineages.py
index 341d084e6..25b4a63e8 100644
--- a/examples/plotting/plot_lineages.py
+++ b/examples/plotting/plot_lineages.py
@@ -41,10 +41,10 @@
 
 # %%
 # Lastly, we can also plot the absorption probabilities separately, one plot for each lineage.
-# By default this also shows the differentiation potential, defined in [Setty19]_ as the entropy over
+# By default this also shows the differentiation potential, defined in :cite:`setty:19` as the entropy over
 # the absorption probabilities.
 cr.pl.lineages(adata, same_plot=False)
 
 # %%
-# To see how to compute and plot the lineage driver genes,
+# To see how to compute and plot the lineage driver genes, please
 # refer to :ref:`sphx_glr_auto_examples_plotting_plot_lineage_drivers.py`.
diff --git a/pyproject.toml b/pyproject.toml
index 26678cc95..bf2638bbf 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -2,13 +2,9 @@
 requires = ['setuptools', 'setuptools_scm', 'wheel']
 build-backend = 'setuptools.build_meta'
 
-[dev-dependencies]
-black = { version = "^19.10b0", python = "^3.6" }
-pytest = { version = "^5.3", python = "^3.6" }
-
 [tool.black]
 line-length = 88
-target-version = ['py36', 'py37', 'py38']
+target-version = ['py37', 'py38', 'py39']
 include = '\.pyi?$'
 exclude = '''
 
@@ -36,7 +32,7 @@ line_length = 88
 multi_line_output = 3
 include_trailing_comma = true
 use_parentheses = true
-known_bio="scanpy,scvelo,anndata"
+known_bio="scanpy,cellrank,scvelo,anndata"
 known_plot="mpl_toolkits,matplotlib,seaborn"
 known_num="numpy,numba,pandas,scipy,sklearn,statmodels,networkx,pygam"
 sections = "FUTURE,STDLIB,THIRDPARTY,BIO,NUM,PLOT,FIRSTPARTY,LOCALFOLDER"
diff --git a/requirements.txt b/requirements.txt
index e6326f040..4ecd6b487 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,8 +1,4 @@
-anndata>=0.7.2
 docrep>=0.3.0
-future_fstrings
-h5py>=2.0.0
-ipywidgets>=7.5.1
 joblib>=0.13.1
 matplotlib>=3.3.0
 networkx>=2.2
@@ -11,12 +7,11 @@ numpy>=1.2.0
 pandas>=1.2.0
 pygam>=0.8.0
 pygpcca>=1.0.2
-scanpy>=1.6.0
-scikit-learn>=0.21.3
+scanpy>=1.7.2
+scikit-learn>=0.24.0
 scipy>=1.2.0
-scvelo>=0.2.2
+scvelo>=0.2.3
 seaborn>=0.10.0
 setuptools>=41.0.1
 typing_extensions
-tzlocal>=1.5.1
 wrapt>=1.12.1
diff --git a/setup.py b/setup.py
index ede961fc5..2a9224762 100644
--- a/setup.py
+++ b/setup.py
@@ -26,7 +26,7 @@
         long_description=Path("README.rst").read_text("utf-8"),
         url="https://github.com/theislab/cellrank",
         project_urls={
-            "Documentation": "https://cellrank.readthedocs.io/en/latest",
+            "Documentation": "https://cellrank.readthedocs.io/en/stable",
             "Source Code": "https://github.com/theislab/cellrank",
         },
         download_url="https://github.com/theislab/cellrank",
@@ -39,7 +39,7 @@
         ),
         extras_require=dict(
             # `POT` is not requirement of `statot`
-            external=["statot>=0.0.14", "POT"],
+            external=["statot>=0.0.14", "POT", "wot"],
             krylov=["pygpcca[slepc]"],
             test=[
                 "pytest>=6.1.1",
@@ -49,8 +49,8 @@
                 "Pillow",
                 "filelock",
                 "python-igraph",
-                "louvain==0.6.1",
-                "leidenalg==0.8.1",
+                "louvain>=0.6,!=0.6.2",
+                "leidenalg",
                 "bezier",
                 "jax",
                 "jaxlib",
@@ -90,7 +90,6 @@
             "Operating System :: POSIX :: Linux",
             "Typing :: Typed",
             "Programming Language :: Python :: 3",
-            "Programming Language :: Python :: 3.6",
             "Programming Language :: Python :: 3.7",
             "Programming Language :: Python :: 3.8",
             "Programming Language :: Python :: 3.9",
diff --git a/tox.ini b/tox.ini
index 00d75b6e2..9d830a65e 100644
--- a/tox.ini
+++ b/tox.ini
@@ -96,6 +96,7 @@ platform =
     linux: linux
     macos: (osx|darwin)
 deps =
+    adjustText
     linux: rpy2>=3.3.0
 extras =
     test


Test patch:
diff --git a/docs/source/latest_additions.rst b/docs/source/latest_additions.rst
index bf155022e..55bb36692 100644
--- a/docs/source/latest_additions.rst
+++ b/docs/source/latest_additions.rst
@@ -1,52 +1,57 @@
 .. role:: small
 
-1.3.0 :small:`2021-03-29`
+1.4.0 :small:`2021-06-30`
 ~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This release includes some major additions which make CellRank more applicable with and without RNA velocity
-information. In particular, it includes:
+This release includes:
 
 .. rubric:: Additions
 
-- 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:
-
-  - `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>`_.
+- Add new external kernel based on Waddington optimal transport :cite:`schiebinger:19`
+  :class:`cellrank.external.kernels.WOTKernel` `PR 595 <https://github.com/theislab/cellrank/pull/595>`_.
+- Add new tutorial `Time series datasets <https://cellrank.readthedocs.io/en/stable/real_time.html>`_ that
+  shows how to use the external Waddington OT kernel `PR 35 <https://github.com/theislab/cellrank_notebooks/pull/35>`_.
+- Add reprogramming dataset :func:`cellrank.datasets.reprogramming_schiebinger` from :cite:`schiebinger:19`
+  `PR 631 <https://github.com/theislab/cellrank/pull/631>`_.
+- Add :func:`cellrank.pl.log_odds` to plot log-odds ratio between 2 lineages sorted by experimental time.
+  `PR 642 <https://github.com/theislab/cellrank/pull/642>`_.
+- Add :meth:`cellrank.tl.estimators.GPCCA.plot_macrostate_composition` to visualize macrostate composition over a
+  categorical variable `PR 641 <https://github.com/theislab/cellrank/pull/641>`_.
+- Add :meth:`cellrank.tl.estimators.BaseEstimator.plot_lineage_drivers_correlation` to plot lineage driver correlation
+  with 2 lineages in a scatter plot `PR 640 <https://github.com/theislab/cellrank/pull/640>`_.
+- Add :meth:`cellrank.tl.kernels.Kernel.plot_single_flow` based on :cite:`mittnenzweig:21` to visualize outgoing
+  transition matrix flow in clusters across experimental time `PR 615 <https://github.com/theislab/cellrank/pull/615>`_.
+- Dramatically speed-up the computation of fate probabilities (>5x) `PR 638 <https://github.com/theislab/cellrank/pull/638>`_.
+- Remove 4 technical examples and add 2 new examples `PR 602 <https://github.com/theislab/cellrank/pull/602>`_:
+
+  - :ref:`sphx_glr_auto_examples_kernels_plot_projection.py` - transition matrix project onto an embedding.
+  - :ref:`sphx_glr_auto_examples_kernels_plot_random_walks.py` - simulation of random walks on a Markov chain.
+
+- Add option to visualize cell-level covariates in :func:`cellrank.pl.cluster_lineage`
+  `PR 634 <https://github.com/theislab/cellrank/pull/634>`_.
+- Add option to force-recompute transition matrix in :func:`cellrank.tl.initial_states` and
+  :func:`cellrank.tl.terminal_states` `PR 577 <https://github.com/theislab/cellrank/pull/577>`_.
+- Change :class:`cellrank.tl.kernels.PseudotimeKernel` defaults and prune available parameters
+  of soft thresholding scheme `PR 583 <https://github.com/theislab/cellrank/pull/583>`_.
+- Parallelize transition matrix computation in :class:`cellrank.tl.kernels.PseudotimeKernel`
+  `PR 587 <https://github.com/theislab/cellrank/pull/587>`_.
+- Prune *requirements.txt* `PR 571 <https://github.com/theislab/cellrank/pull/571>`_.
+- Add small improvements to documentation `PR 584 <https://github.com/theislab/cellrank/pull/584>`_
+  `PR 601 <https://github.com/theislab/cellrank/pull/601>`_ `PR 605 <https://github.com/theislab/cellrank/pull/605>`_
+  `PR 648 <https://github.com/theislab/cellrank/issues/648>`_ `PR 35 <https://github.com/theislab/cellrank_notebooks/pull/35>`_.
 
 .. rubric:: Bugfixes
 
-- 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>`_.
+- Fix estimator's inconsistent state when reading from :class:`anndata.AnnData`
+  `PR 563 <https://github.com/theislab/cellrank/pull/563>`_.
+- Fix not checking whether probabilities sum to 1 in
+  :meth:`cellrank.tl.estimators.BaseEstimator.compute_absorption_probabilities`
+  `PR 566 <https://github.com/theislab/cellrank/pull/566>`_.
+- Fix always forcing sparse transition matrix in :class:`cellrank.tl.kernels.Kernel`
+  `PR 586 <https://github.com/theislab/cellrank/pull/586>`_.
+- Fix passing custom connectivity key in :class:`cellrank.tl.kernels.Kernel`
+  `PR 590 <https://github.com/theislab/cellrank/pull/590>`_.
+- Fix kernels in :mod:`cellrank.external` always requiring connectivities
+  `PR 600 <https://github.com/theislab/cellrank/pull/600>`_.
+- Fix parallelization of sparse matrix with too many jobs `PR 633 <https://github.com/theislab/cellrank/pull/633>`_.
+- Fix plotting coarse-grained transition matrix when no stationary distribution is found
+  `Issue 594 <https://github.com/theislab/cellrank/issues/594>`_.
diff --git a/tests/_ground_truth_figures/cluster_lineage.png b/tests/_ground_truth_figures/cluster_lineage.png
index a0306e4fb..1a5bbf532 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage.png and b/tests/_ground_truth_figures/cluster_lineage.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png b/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png
index ca89d763e..73e93af20 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png and b/tests/_ground_truth_figures/cluster_lineage_2_failed_genes.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_bwd.png b/tests/_ground_truth_figures/cluster_lineage_bwd.png
index 047fa86f8..874886a8e 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_bwd.png and b/tests/_ground_truth_figures/cluster_lineage_bwd.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates.png b/tests/_ground_truth_figures/cluster_lineage_covariates.png
new file mode 100644
index 000000000..3d5045bb3
Binary files /dev/null and b/tests/_ground_truth_figures/cluster_lineage_covariates.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png b/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png
new file mode 100644
index 000000000..4b528f0c6
Binary files /dev/null and b/tests/_ground_truth_figures/cluster_lineage_covariates_cmap.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png b/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png
new file mode 100644
index 000000000..0cbb34f42
Binary files /dev/null and b/tests/_ground_truth_figures/cluster_lineage_covariates_ratio.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_data_key.png b/tests/_ground_truth_figures/cluster_lineage_data_key.png
index eb7ccaf1c..320de1d33 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_data_key.png and b/tests/_ground_truth_figures/cluster_lineage_data_key.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_leiden.png b/tests/_ground_truth_figures/cluster_lineage_leiden.png
index a0306e4fb..1a5bbf532 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_leiden.png and b/tests/_ground_truth_figures/cluster_lineage_leiden.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_no_norm.png b/tests/_ground_truth_figures/cluster_lineage_no_norm.png
index 799a21b62..21323f9e3 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_no_norm.png and b/tests/_ground_truth_figures/cluster_lineage_no_norm.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_random_state.png b/tests/_ground_truth_figures/cluster_lineage_random_state.png
index a0306e4fb..1a5bbf532 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_random_state.png and b/tests/_ground_truth_figures/cluster_lineage_random_state.png differ
diff --git a/tests/_ground_truth_figures/cluster_lineage_raw.png b/tests/_ground_truth_figures/cluster_lineage_raw.png
index d5fff3e33..211574a2e 100644
Binary files a/tests/_ground_truth_figures/cluster_lineage_raw.png and b/tests/_ground_truth_figures/cluster_lineage_raw.png differ
diff --git a/tests/_ground_truth_figures/driver_corr.png b/tests/_ground_truth_figures/driver_corr.png
new file mode 100644
index 000000000..c0a796dc3
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_adjust_text.png b/tests/_ground_truth_figures/driver_corr_adjust_text.png
new file mode 100644
index 000000000..6e841088f
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_adjust_text.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_cmap.png b/tests/_ground_truth_figures/driver_corr_cmap.png
new file mode 100644
index 000000000..94e11098a
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_cmap.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_color.png b/tests/_ground_truth_figures/driver_corr_color.png
new file mode 100644
index 000000000..aeb3753df
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_color.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_fontsize.png b/tests/_ground_truth_figures/driver_corr_fontsize.png
new file mode 100644
index 000000000..8226ecfe5
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_fontsize.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets.png b/tests/_ground_truth_figures/driver_corr_gene_sets.png
new file mode 100644
index 000000000..444576d3a
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_gene_sets.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png
new file mode 100644
index 000000000..23461b073
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_legend_loc.png b/tests/_ground_truth_figures/driver_corr_legend_loc.png
new file mode 100644
index 000000000..555b5744c
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_lenged_loc.png b/tests/_ground_truth_figures/driver_corr_lenged_loc.png
new file mode 100644
index 000000000..555b5744c
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_lenged_loc.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_return_ax.png b/tests/_ground_truth_figures/driver_corr_return_ax.png
new file mode 100644
index 000000000..0cbdddfa9
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_return_ax.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_use_raw.png b/tests/_ground_truth_figures/driver_corr_use_raw.png
new file mode 100644
index 000000000..d09bb54d8
Binary files /dev/null and b/tests/_ground_truth_figures/driver_corr_use_raw.png differ
diff --git a/tests/_ground_truth_figures/flow_alpha.png b/tests/_ground_truth_figures/flow_alpha.png
new file mode 100644
index 000000000..0e6b91862
Binary files /dev/null and b/tests/_ground_truth_figures/flow_alpha.png differ
diff --git a/tests/_ground_truth_figures/flow_cluster_ascending.png b/tests/_ground_truth_figures/flow_cluster_ascending.png
new file mode 100644
index 000000000..8fc9cddd1
Binary files /dev/null and b/tests/_ground_truth_figures/flow_cluster_ascending.png differ
diff --git a/tests/_ground_truth_figures/flow_cluster_descending.png b/tests/_ground_truth_figures/flow_cluster_descending.png
new file mode 100644
index 000000000..6dfafdca8
Binary files /dev/null and b/tests/_ground_truth_figures/flow_cluster_descending.png differ
diff --git a/tests/_ground_truth_figures/flow_clusters_subset.png b/tests/_ground_truth_figures/flow_clusters_subset.png
new file mode 100644
index 000000000..38f20949d
Binary files /dev/null and b/tests/_ground_truth_figures/flow_clusters_subset.png differ
diff --git a/tests/_ground_truth_figures/flow_explicit_cluster_order.png b/tests/_ground_truth_figures/flow_explicit_cluster_order.png
new file mode 100644
index 000000000..463213764
Binary files /dev/null and b/tests/_ground_truth_figures/flow_explicit_cluster_order.png differ
diff --git a/tests/_ground_truth_figures/flow_legend_loc.png b/tests/_ground_truth_figures/flow_legend_loc.png
new file mode 100644
index 000000000..64cd1b056
Binary files /dev/null and b/tests/_ground_truth_figures/flow_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png b/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png
new file mode 100644
index 000000000..10317cd85
Binary files /dev/null and b/tests/_ground_truth_figures/flow_min_flow_keep_empty_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png b/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png
new file mode 100644
index 000000000..2cd45b681
Binary files /dev/null and b/tests/_ground_truth_figures/flow_min_flow_remove_empty_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_no_xticks.png b/tests/_ground_truth_figures/flow_no_xticks.png
new file mode 100644
index 000000000..59321da57
Binary files /dev/null and b/tests/_ground_truth_figures/flow_no_xticks.png differ
diff --git a/tests/_ground_truth_figures/flow_return_ax.png b/tests/_ground_truth_figures/flow_return_ax.png
new file mode 100644
index 000000000..6dfafdca8
Binary files /dev/null and b/tests/_ground_truth_figures/flow_return_ax.png differ
diff --git a/tests/_ground_truth_figures/flow_source_clusters.png b/tests/_ground_truth_figures/flow_source_clusters.png
new file mode 100644
index 000000000..ee24238d0
Binary files /dev/null and b/tests/_ground_truth_figures/flow_source_clusters.png differ
diff --git a/tests/_ground_truth_figures/flow_time_categories_too_close.png b/tests/_ground_truth_figures/flow_time_categories_too_close.png
new file mode 100644
index 000000000..8248fc56c
Binary files /dev/null and b/tests/_ground_truth_figures/flow_time_categories_too_close.png differ
diff --git a/tests/_ground_truth_figures/log_odds.png b/tests/_ground_truth_figures/log_odds.png
new file mode 100644
index 000000000..9d27f3074
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds.png differ
diff --git a/tests/_ground_truth_figures/log_odds_alpha.png b/tests/_ground_truth_figures/log_odds_alpha.png
new file mode 100644
index 000000000..e8b5ae918
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_alpha.png differ
diff --git a/tests/_ground_truth_figures/log_odds_bwd.png b/tests/_ground_truth_figures/log_odds_bwd.png
new file mode 100644
index 000000000..927079a54
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_bwd.png differ
diff --git a/tests/_ground_truth_figures/log_odds_categorical_keys.png b/tests/_ground_truth_figures/log_odds_categorical_keys.png
new file mode 100644
index 000000000..81fa89806
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_categorical_keys.png differ
diff --git a/tests/_ground_truth_figures/log_odds_cmap.png b/tests/_ground_truth_figures/log_odds_cmap.png
new file mode 100644
index 000000000..65d0df501
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_cmap.png differ
diff --git a/tests/_ground_truth_figures/log_odds_continuous_keys.png b/tests/_ground_truth_figures/log_odds_continuous_keys.png
new file mode 100644
index 000000000..efabdec5b
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_continuous_keys.png differ
diff --git a/tests/_ground_truth_figures/log_odds_fontsize.png b/tests/_ground_truth_figures/log_odds_fontsize.png
new file mode 100644
index 000000000..f6a489b9d
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_fontsize.png differ
diff --git a/tests/_ground_truth_figures/log_odds_jitter.png b/tests/_ground_truth_figures/log_odds_jitter.png
new file mode 100644
index 000000000..9d27f3074
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_jitter.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs.png b/tests/_ground_truth_figures/log_odds_kwargs.png
new file mode 100644
index 000000000..bd8205fe0
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_kwargs.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png b/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png
new file mode 100644
index 000000000..b9b102a3d
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_kwargs_return_ax.png differ
diff --git a/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png b/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png
new file mode 100644
index 000000000..a607ee332
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_kwargs_return_axes.png differ
diff --git a/tests/_ground_truth_figures/log_odds_layer.png b/tests/_ground_truth_figures/log_odds_layer.png
new file mode 100644
index 000000000..95060f0db
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_layer.png differ
diff --git a/tests/_ground_truth_figures/log_odds_legend_loc.png b/tests/_ground_truth_figures/log_odds_legend_loc.png
new file mode 100644
index 000000000..1164e085a
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/log_odds_multiple_threshold.png b/tests/_ground_truth_figures/log_odds_multiple_threshold.png
new file mode 100644
index 000000000..50a8a887e
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_multiple_threshold.png differ
diff --git a/tests/_ground_truth_figures/log_odds_ncols.png b/tests/_ground_truth_figures/log_odds_ncols.png
new file mode 100644
index 000000000..6ab163eaf
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_ncols.png differ
diff --git a/tests/_ground_truth_figures/log_odds_rest.png b/tests/_ground_truth_figures/log_odds_rest.png
new file mode 100644
index 000000000..823083da6
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_rest.png differ
diff --git a/tests/_ground_truth_figures/log_odds_size.png b/tests/_ground_truth_figures/log_odds_size.png
new file mode 100644
index 000000000..e76d815dc
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_size.png differ
diff --git a/tests/_ground_truth_figures/log_odds_threshold.png b/tests/_ground_truth_figures/log_odds_threshold.png
new file mode 100644
index 000000000..1152040af
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_threshold.png differ
diff --git a/tests/_ground_truth_figures/log_odds_threshold_color.png b/tests/_ground_truth_figures/log_odds_threshold_color.png
new file mode 100644
index 000000000..493af94b2
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_threshold_color.png differ
diff --git a/tests/_ground_truth_figures/log_odds_use_raw.png b/tests/_ground_truth_figures/log_odds_use_raw.png
new file mode 100644
index 000000000..d52a402ff
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_use_raw.png differ
diff --git a/tests/_ground_truth_figures/log_odds_xticks_steps_size.png b/tests/_ground_truth_figures/log_odds_xticks_steps_size.png
new file mode 100644
index 000000000..f10db7b4e
Binary files /dev/null and b/tests/_ground_truth_figures/log_odds_xticks_steps_size.png differ
diff --git a/tests/_ground_truth_figures/msc_default.png b/tests/_ground_truth_figures/msc_default.png
new file mode 100644
index 000000000..0981fbc9f
Binary files /dev/null and b/tests/_ground_truth_figures/msc_default.png differ
diff --git a/tests/_ground_truth_figures/msc_labelrot.png b/tests/_ground_truth_figures/msc_labelrot.png
new file mode 100644
index 000000000..d4417854d
Binary files /dev/null and b/tests/_ground_truth_figures/msc_labelrot.png differ
diff --git a/tests/_ground_truth_figures/msc_legend_loc.png b/tests/_ground_truth_figures/msc_legend_loc.png
new file mode 100644
index 000000000..90771dab8
Binary files /dev/null and b/tests/_ground_truth_figures/msc_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/msc_title.png b/tests/_ground_truth_figures/msc_title.png
new file mode 100644
index 000000000..ae7e99213
Binary files /dev/null and b/tests/_ground_truth_figures/msc_title.png differ
diff --git a/tests/_ground_truth_figures/msc_width.png b/tests/_ground_truth_figures/msc_width.png
new file mode 100644
index 000000000..1051fb42d
Binary files /dev/null and b/tests/_ground_truth_figures/msc_width.png differ
diff --git a/tests/_helpers.py b/tests/_helpers.py
index 993c47fa8..ebfde57a5 100644
--- a/tests/_helpers.py
+++ b/tests/_helpers.py
@@ -8,16 +8,15 @@
 
 import scanpy as sc
 import scvelo as scv
+import cellrank as cr
 from anndata import AnnData
+from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
 
 import numpy as np
 import pandas as pd
 from sklearn.svm import SVR
 from scipy.sparse import spdiags, issparse, csr_matrix
 
-import cellrank as cr
-from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
-
 
 def _jax_not_installed() -> bool:
     try:
diff --git a/tests/conftest.py b/tests/conftest.py
index 607e1f9af..6b2a38552 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -13,17 +13,16 @@
 from _helpers import create_model
 
 import scanpy as sc
+import cellrank as cr
 from anndata import AnnData
+from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
+from cellrank.tl._constants import AbsProbKey
+from cellrank.tl.estimators import GPCCA, CFLARE
 
 import numpy as np
 
 import matplotlib
 
-import cellrank as cr
-from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
-from cellrank.tl._constants import AbsProbKey
-from cellrank.tl.estimators import GPCCA, CFLARE
-
 matplotlib.use("Agg")
 np.random.seed(42)
 
@@ -214,6 +213,13 @@ def lineage():
     return x / x.sum(1)
 
 
+@pytest.fixture
+def kernel(adata_large: AnnData):
+    vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+    ck = ConnectivityKernel(adata_large).compute_transition_matrix()
+    return (0.8 * vk + 0.2 * ck).compute_transition_matrix()
+
+
 # removes overly verbose logging errors for rpy2
 # see: https://github.com/pytest-dev/pytest/issues/5502#issuecomment-647157873
 def pytest_sessionfinish(session, exitstatus):
diff --git a/tests/test_cflare.py b/tests/test_cflare.py
index fd69d870b..6c829e359 100644
--- a/tests/test_cflare.py
+++ b/tests/test_cflare.py
@@ -5,14 +5,9 @@
 import pytest
 from _helpers import assert_estimators_equal
 
-from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from pandas.api.types import is_categorical_dtype
-
 import cellrank as cr
 import cellrank.tl.kernels._precomputed_kernel
+from anndata import AnnData
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
 from cellrank.tl._constants import (
     Direction,
@@ -26,6 +21,10 @@
 )
 from cellrank.tl.estimators._constants import A, P
 
+import numpy as np
+import pandas as pd
+from pandas.api.types import is_categorical_dtype
+
 EPS = np.finfo(np.float64).eps
 
 
@@ -590,6 +589,55 @@ def test_manual_approx_rc_set(self, adata_large):
         assert (adata.obs[f"{TermStatesKey.FORWARD}"][zero_mask] == "foo").all()
         assert pd.isna(adata.obs[f"{TermStatesKey.FORWARD}"][~zero_mask]).all()
 
+    def test_abs_probs_do_not_sum_to_1(self, adata_large: AnnData, mocker):
+        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
+        terminal_kernel = 0.8 * vk + 0.2 * ck
+
+        mc = cr.tl.estimators.CFLARE(terminal_kernel)
+        mc.compute_eigendecomposition(k=5)
+        mc.set_terminal_states(
+            {"x": adata_large.obs_names[:3], "y": adata_large.obs_names[3:6]}
+        )
+
+        n_term = np.sum(~pd.isnull(mc.terminal_states))
+        abs_prob = np.zeros((adata_large.n_obs - n_term, n_term))
+        abs_prob[:, 0] = 1.0
+        abs_prob[0, 0] = 1.01
+        mocker.patch(
+            "cellrank.tl.estimators._base_estimator._solve_lin_system",
+            return_value=abs_prob,
+        )
+
+        with pytest.raises(
+            ValueError, match=r"`1` value\(s\) do not sum to 1 \(rtol=1e-3\)."
+        ):
+            mc.compute_absorption_probabilities()
+
+    def test_abs_probs_negative(self, adata_large: AnnData, mocker):
+        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
+        terminal_kernel = 0.8 * vk + 0.2 * ck
+
+        mc = cr.tl.estimators.CFLARE(terminal_kernel)
+        mc.compute_eigendecomposition(k=5)
+        mc.set_terminal_states(
+            {"x": adata_large.obs_names[:3], "y": adata_large.obs_names[3:6]}
+        )
+
+        n_term = np.sum(~pd.isnull(mc.terminal_states))
+        abs_prob = np.zeros((adata_large.n_obs - n_term, n_term))
+        abs_prob[:, 0] = 1.0
+        abs_prob[0, 0] = -0.5
+        abs_prob[0, 1] = -1.5
+        mocker.patch(
+            "cellrank.tl.estimators._base_estimator._solve_lin_system",
+            return_value=abs_prob,
+        )
+
+        with pytest.raises(ValueError, match=r"`2` value\(s\) are negative."):
+            mc.compute_absorption_probabilities()
+
 
 class TestCFLAREIO:
     def test_copy(self, adata_cflare_fwd: Tuple[AnnData, cr.tl.estimators.CFLARE]):
diff --git a/tests/test_colors.py b/tests/test_colors.py
index 681708429..68448f48a 100644
--- a/tests/test_colors.py
+++ b/tests/test_colors.py
@@ -1,13 +1,13 @@
 import pytest
 
+from cellrank.tl._colors import _map_names_and_colors, _create_categorical_colors
+
 import numpy as np
 import pandas as pd
 from pandas.api.types import is_categorical_dtype
 
 from matplotlib.colors import is_color_like
 
-from cellrank.tl._colors import _map_names_and_colors, _create_categorical_colors
-
 
 class TestColors:
     def test_create_categorical_colors_too_many_colors(self):
diff --git a/tests/test_datasets.py b/tests/test_datasets.py
index da1e10c67..30452e32e 100644
--- a/tests/test_datasets.py
+++ b/tests/test_datasets.py
@@ -2,9 +2,8 @@
 
 import pytest
 
-from anndata import AnnData
-
 import cellrank as cr
+from anndata import AnnData
 
 
 @pytest.mark.skipif(
diff --git a/tests/test_external.py b/tests/test_external.py
index 9eb2d7c66..5878a86ed 100644
--- a/tests/test_external.py
+++ b/tests/test_external.py
@@ -1,21 +1,43 @@
+from typing import Optional
+
 import pytest
 
+import scanpy as sc
+import cellrank.external as cre
 from anndata import AnnData
+from cellrank.tl.kernels import ConnectivityKernel
+from cellrank.external.kernels._utils import MarkerGenes
+from cellrank.external.kernels._wot_kernel import LastTimePoint
 
 import numpy as np
 import pandas as pd
 from scipy.sparse import spmatrix
-
-import cellrank.external as cre
+from pandas.core.dtypes.common import is_categorical_dtype
 
 
 class TestOTKernel:
+    def test_no_connectivities(self, adata_large: AnnData):
+        del adata_large.obsp["connectivities"]
+        terminal_states = np.full((adata_large.n_obs,), fill_value=np.nan, dtype=object)
+        ixs = np.where(adata_large.obs["clusters"] == "Granule immature")[0]
+        terminal_states[ixs] = "GI"
+
+        ok = cre.kernels.StationaryOTKernel(
+            adata_large,
+            terminal_states=pd.Series(terminal_states).astype("category"),
+            g=np.ones((adata_large.n_obs,), dtype=np.float64),
+        )
+        ok = ok.compute_transition_matrix(1, 0.001)
+
+        assert ok._conn is None
+        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
+
     def test_method_not_implemented(self, adata_large: AnnData):
         terminal_states = np.full((adata_large.n_obs,), fill_value=np.nan, dtype=object)
         ixs = np.where(adata_large.obs["clusters"] == "Granule immature")[0]
         terminal_states[ixs] = "GI"
 
-        ok = cre.kernels.OTKernel(
+        ok = cre.kernels.StationaryOTKernel(
             adata_large,
             terminal_states=pd.Series(terminal_states).astype("category"),
             g=np.ones((adata_large.n_obs,), dtype=np.float64),
@@ -28,7 +50,7 @@ def test_method_not_implemented(self, adata_large: AnnData):
 
     def test_no_terminal_states(self, adata_large: AnnData):
         with pytest.raises(RuntimeError, match="Unable to initialize the kernel."):
-            cre.kernels.OTKernel(
+            cre.kernels.StationaryOTKernel(
                 adata_large,
                 g=np.ones((adata_large.n_obs,), dtype=np.float64),
             )
@@ -38,14 +60,186 @@ def test_normal_run(self, adata_large: AnnData):
         ixs = np.where(adata_large.obs["clusters"] == "Granule immature")[0]
         terminal_states[ixs] = "GI"
 
-        ok = cre.kernels.OTKernel(
+        ok = cre.kernels.StationaryOTKernel(
             adata_large,
             terminal_states=pd.Series(terminal_states).astype("category"),
             g=np.ones((adata_large.n_obs,), dtype=np.float64),
         )
         ok = ok.compute_transition_matrix(1, 0.001)
 
-        assert isinstance(ok, cre.kernels.OTKernel)
+        assert isinstance(ok, cre.kernels.StationaryOTKernel)
+        assert isinstance(ok._transition_matrix, (np.ndarray, spmatrix))
+        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
+        assert isinstance(ok.params, dict)
+
+
+@pytest.mark.skip("wot on PyPI doesn't support passing cost matrices")
+class TestWOTKernel:
+    def test_no_connectivities(self, adata_large: AnnData):
+        del adata_large.obsp["connectivities"]
+        ok = cre.kernels.WOTKernel(
+            adata_large, time_key="age(days)"
+        ).compute_transition_matrix()
+
+        assert ok._conn is None
+        np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
+
+    def test_inversion_updates_adata(self, adata_large: AnnData):
+        key = "age(days)"
+        ok = cre.kernels.WOTKernel(adata_large, time_key=key)
+        assert is_categorical_dtype(adata_large.obs[key])
+        assert adata_large.obs[key].cat.ordered
+        np.testing.assert_array_equal(ok.experimental_time, adata_large.obs[key])
+        orig_time = ok.experimental_time
+
+        ok = ~ok
+
+        inverted_time = ok.experimental_time
+        assert is_categorical_dtype(adata_large.obs[key])
+        assert adata_large.obs[key].cat.ordered
+        np.testing.assert_array_equal(ok.experimental_time, adata_large.obs[key])
+        np.testing.assert_array_equal(
+            orig_time.cat.categories, inverted_time.cat.categories
+        )
+        np.testing.assert_array_equal(orig_time.index, inverted_time.index)
+
+        with pytest.raises(AssertionError):
+            np.testing.assert_array_equal(orig_time, inverted_time)
+
+    @pytest.mark.parametrize("cmat", [None, "Ms", "X_pca", "good_shape", "bad_shape"])
+    def test_cost_matrices(self, adata_large: AnnData, cmat: str):
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        if isinstance(cmat, str) and "shape" in cmat:
+            cost_matrices = {
+                (12.0, 35.0): np.random.normal(size=(73 + ("bad" in cmat), 127))
+            }
+        else:
+            cost_matrices = cmat
+
+        if cmat == "bad_shape":
+            with pytest.raises(ValueError, match=r"Expected cost matrix for time pair"):
+                ok.compute_transition_matrix(cost_matrices=cost_matrices)
+        else:
+            ok = ok.compute_transition_matrix(cost_matrices=cost_matrices)
+            param = ok.params["cost_matrices"]
+            if cmat == "Ms":
+                assert param == "layer:Ms"
+            elif cmat == "X_pca":
+                assert param == "obsm:X_pca"
+            elif cmat == "good_shape":
+                # careful, param is `nstr`, which does not equal anything
+                assert str(param) == "precomputed"
+            else:
+                assert param == "default"
+
+    @pytest.mark.parametrize("n_iters", [3, 5])
+    def test_growth_rates(self, adata_large: AnnData, n_iters: int):
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        ok = ok.compute_transition_matrix(growth_iters=n_iters)
+
+        assert isinstance(ok.growth_rates, pd.DataFrame)
+        np.testing.assert_array_equal(adata_large.obs_names, ok.growth_rates.index)
+        np.testing.assert_array_equal(
+            ok.growth_rates.columns, [f"g{i}" for i in range(n_iters + 1)]
+        )
+        np.testing.assert_array_equal(
+            adata_large.obs["estimated_growth_rates"], ok.growth_rates[f"g{n_iters}"]
+        )
+        assert ok.params["growth_iters"] == n_iters
+
+    @pytest.mark.parametrize("key_added", [None, "gr"])
+    def test_birth_death_process(self, adata_large: AnnData, key_added: Optional[str]):
+        np.random.seed(42)
+        adata_large.obs["foo"] = np.random.normal(size=(adata_large.n_obs,))
+        adata_large.obs["bar"] = np.random.normal(size=(adata_large.n_obs,))
+
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        gr = ok.compute_initial_growth_rates("foo", "bar", key_added=key_added)
+
+        if key_added is None:
+            assert isinstance(gr, pd.Series)
+            np.testing.assert_array_equal(gr.index, adata_large.obs_names)
+        else:
+            assert gr is None
+            assert "gr" in adata_large.obs
+
+    @pytest.mark.parametrize("ltp", list(LastTimePoint))
+    def test_last_time_point(self, adata_large: AnnData, ltp: LastTimePoint):
+        key = "age(days)"
+        ok = cre.kernels.WOTKernel(adata_large, time_key=key).compute_transition_matrix(
+            last_time_point=ltp.s, conn_kwargs={"n_neighbors": 11}
+        )
+        ixs = np.where(adata_large.obs[key] == 35.0)[0]
+
+        T = ok.transition_matrix[ixs, :][:, ixs].A
+        if ltp == LastTimePoint.UNIFORM:
+            np.testing.assert_allclose(T, np.ones_like(T) / float(len(ixs)))
+        elif ltp == LastTimePoint.DIAGONAL:
+            np.testing.assert_allclose(T, np.eye(len(ixs)))
+        elif ltp == LastTimePoint.CONNECTIVITIES:
+            adata_subset = adata_large[adata_large.obs[key] == 35.0]
+            sc.pp.neighbors(adata_subset, n_neighbors=11)
+            T_actual = (
+                ConnectivityKernel(adata_subset)
+                .compute_transition_matrix()
+                .transition_matrix.A
+            )
+            np.testing.assert_allclose(T, T_actual)
+
+    @pytest.mark.parametrize("organism", ["human", "mouse"])
+    def test_compute_scores_default(self, adata_large: AnnData, organism: str):
+        pk, ak = "p_score", "a_score"
+        if organism == "human":
+            adata_large.var_names = adata_large.var_names.str.upper()
+
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        assert pk not in ok.adata.obs
+        assert ak not in ok.adata.obs
+
+        ok.compute_initial_growth_rates(
+            organism=organism, proliferation_key=pk, apoptosis_key=ak, use_raw=False
+        )
+
+        assert pk in ok.adata.obs
+        assert ak in ok.adata.obs
+
+    def test_normal_run(self, adata_large: AnnData):
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        ok = ok.compute_transition_matrix()
+
+        assert isinstance(ok, cre.kernels.WOTKernel)
         assert isinstance(ok._transition_matrix, (np.ndarray, spmatrix))
         np.testing.assert_allclose(ok.transition_matrix.sum(1), 1.0)
         assert isinstance(ok.params, dict)
+        assert isinstance(ok.growth_rates, pd.DataFrame)
+        assert isinstance(ok.transport_maps, dict)
+        np.testing.assert_array_equal(adata_large.obs_names, ok.growth_rates.index)
+        np.testing.assert_array_equal(ok.growth_rates.columns, ["g0", "g1"])
+        assert isinstance(ok.transport_maps[12.0, 35.0], AnnData)
+
+    def test_copy(self, adata_large: AnnData):
+        ok = cre.kernels.WOTKernel(adata_large, time_key="age(days)")
+        ok = ok.compute_transition_matrix()
+
+        ok2 = ok.copy()
+
+        assert isinstance(ok2, cre.kernels.WOTKernel)
+        assert ok is not ok2
+        np.testing.assert_array_equal(ok.transition_matrix.A, ok2.transition_matrix.A)
+
+
+class TestGetMarkers:
+    @pytest.mark.parametrize("kind", ["proliferation", "apoptosis"])
+    @pytest.mark.parametrize("organism", ["human", "mouse", "foo"])
+    def test_get_markers(self, organism: str, kind: str):
+        if organism == "foo":
+            with pytest.raises(NotImplementedError, match=r""):
+                getattr(MarkerGenes, f"{kind}_markers")(organism)
+        else:
+            markers = getattr(MarkerGenes, f"{kind}_markers")(organism)
+            assert isinstance(markers, tuple)
+            assert np.all([isinstance(marker, str) for marker in markers])
+            if organism == "human":
+                np.testing.assert_array_equal(
+                    markers, [marker.upper() for marker in markers]
+                )
diff --git a/tests/test_gpcca.py b/tests/test_gpcca.py
index 1dfcd32c6..e70dea3a0 100644
--- a/tests/test_gpcca.py
+++ b/tests/test_gpcca.py
@@ -1,18 +1,13 @@
 import os
 from copy import deepcopy
-from typing import Tuple
+from typing import List, Tuple, Union
 from tempfile import TemporaryDirectory
 
 import pytest
 from _helpers import assert_array_nan_equal, assert_estimators_equal
 
-from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from pandas.testing import assert_series_equal
-
 import cellrank as cr
+from anndata import AnnData
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
 from cellrank.tl._constants import (
     Direction,
@@ -26,6 +21,10 @@
 )
 from cellrank.tl.estimators._constants import A, P
 
+import numpy as np
+import pandas as pd
+from pandas.testing import assert_series_equal
+
 
 def _check_eigdecomposition(mc: cr.tl.estimators.GPCCA) -> None:
     assert isinstance(mc._get(P.EIG), dict)
@@ -580,6 +579,27 @@ def test_set_terminal_states_from_macrostates_invalid_name(
         with pytest.raises(KeyError):
             mc.set_terminal_states_from_macrostates(names=["foobar"])
 
+    @pytest.mark.parametrize("values", ["Astrocytes", ["Astrocytes", "OPC"]])
+    def test_set_terminal_states_clusters(
+        self, adata_large: AnnData, values: Union[str, List[str]]
+    ):
+        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
+        terminal_kernel = 0.8 * vk + 0.2 * ck
+
+        to_remove = list(
+            set(adata_large.obs["clusters"].cat.categories)
+            - ({values} if isinstance(values, str) else set(values))
+        )
+        expected = adata_large.obs["clusters"].cat.remove_categories(to_remove)
+
+        mc = cr.tl.estimators.GPCCA(terminal_kernel)
+
+        mc.set_terminal_states({"clusters": values})
+        pd.testing.assert_series_equal(
+            expected, mc.terminal_states, check_category_order=False, check_names=False
+        )
+
     def test_compute_terminal_states_invalid_method(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -900,6 +920,25 @@ def test_compute_priming_clusters(self, adata_large: AnnData):
         )
         assert_series_equal(mc._get(A.PRIME_DEG), deg1)
 
+    def test_recompute_terminal_states_different_n_states(self, adata_large: AnnData):
+        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
+        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
+        terminal_kernel = 0.8 * vk + 0.2 * ck
+
+        mc = cr.tl.estimators.GPCCA(terminal_kernel)
+        mc.compute_schur(n_components=10, method="krylov")
+        mc.compute_macrostates(n_states=2)
+        mc.set_terminal_states_from_macrostates()
+
+        assert len(mc.terminal_states.cat.categories) == 2
+        assert adata_large.obsm[mc._term_abs_prob_key].shape == (adata_large.n_obs, 2)
+
+        mc.set_terminal_states({"foo": adata_large.obs_names[:20]})
+
+        assert len(mc.terminal_states.cat.categories) == 1
+        assert mc._get(A.TERM_ABS_PROBS) is None
+        assert mc._term_abs_prob_key not in adata_large.obsm
+
 
 class TestGPCCAIO:
     def test_copy(self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]):
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index cee8b6657..f41f138dc 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1,3 +1,4 @@
+from copy import copy
 from typing import Tuple, Callable, Optional
 
 import pytest
@@ -9,13 +10,10 @@
     random_transition_matrix,
 )
 
+import scanpy as sc
+import cellrank as cr
 from scanpy import Neighbors
 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 (
@@ -27,13 +25,20 @@
 )
 from cellrank.tl._constants import _transition
 from cellrank.tl.kernels._base_kernel import (
+    Kernel,
     Constant,
     KernelAdd,
     KernelMul,
+    _dtype,
     _is_bin_mult,
 )
 from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEAggregation, _ct
 
+import numpy as np
+from scipy.sparse import eye as speye
+from scipy.sparse import isspmatrix_csr
+from pandas.core.dtypes.common import is_bool_dtype, is_integer_dtype
+
 _rtol = 1e-6
 
 
@@ -57,6 +62,22 @@ def hessian(
         return np.zeros((D.shape[0], v.shape[0], v.shape[0]))
 
 
+class CustomKernel(Kernel):
+    def compute_transition_matrix(
+        self, sparse: bool = False, dnorm: bool = False
+    ) -> "KernelExpression":
+        if sparse:
+            tmat = speye(self.adata.n_obs, dtype=np.float32)
+        else:
+            tmat = np.eye(self.adata.n_obs, dtype=np.float32)
+
+        self._compute_transition_matrix(tmat, density_normalize=dnorm)
+        return self
+
+    def copy(self) -> "KernelExpression":
+        return copy(self)
+
+
 class InvalidFuncProbs(cr.tl.kernels.SimilaritySchemeABC):
     def __call__(
         self, v: np.ndarray, D: np.ndarray, _softmax_scale: float = 1.0
@@ -421,6 +442,45 @@ def test_precomputed_from_kernel_no_transition(self, adata: AnnData):
         with pytest.raises(ValueError):
             PrecomputedKernel(vk)
 
+    @pytest.mark.parametrize(
+        "clazz",
+        [
+            ConnectivityKernel,
+            VelocityKernel,
+            PseudotimeKernel,
+            CytoTRACEKernel,
+            PrecomputedKernel,
+        ],
+    )
+    @pytest.mark.parametrize("key_added", [None, "foo"])
+    def test_kernel_reads_correct_connectivities(
+        self, adata: AnnData, key_added: Optional[str], clazz: type
+    ):
+        del adata.uns["neighbors"]
+        del adata.obsp["connectivities"]
+        del adata.obsp["distances"]
+
+        sc.pp.neighbors(adata, key_added=key_added)
+        kwargs = {"adata": adata, "conn_key": key_added}
+
+        if clazz == PseudotimeKernel:
+            kwargs["time_key"] = "latent_time"
+        elif clazz == PrecomputedKernel:
+            adata.obsp["foo"] = np.eye(adata.n_obs)
+            kwargs["transition_matrix"] = "foo"
+        conn = (
+            adata.obsp["connectivities"]
+            if key_added is None
+            else adata.obsp[f"{key_added}_connectivities"]
+        )
+
+        k = clazz(**kwargs)
+
+        if isinstance(k, PrecomputedKernel):
+            assert k._conn is None
+        else:
+            np.testing.assert_array_equal(k._conn.A, conn.A)
+
     def test_precomputed_from_kernel(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(
             mode="stochastic",
@@ -497,6 +557,18 @@ def test_precomputed_sum(self, adata: AnnData):
             expected.toarray(), actual.transition_matrix.toarray()
         )
 
+    @pytest.mark.parametrize("dnorm", [False, True])
+    @pytest.mark.parametrize("sparse", [False, True])
+    def test_custom_preserves_type(self, adata: AnnData, sparse: bool, dnorm: bool):
+        c = CustomKernel(adata).compute_transition_matrix(sparse=sparse, dnorm=dnorm)
+
+        if sparse:
+            assert isspmatrix_csr(c.transition_matrix)
+        else:
+            assert isinstance(c.transition_matrix, np.ndarray)
+
+        assert c.transition_matrix.dtype == _dtype
+
     def test_write_adata(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
         vk.write_to_adata()
@@ -552,29 +624,39 @@ def test_transition_forward_stoch_high_lvl(self, adata: AnnData):
 
         np.testing.assert_allclose(T_1.A, T_2.A, rtol=_rtol)
 
-    @pytest.mark.parametrize("dens_norm", [False, True])
-    def test_palantir(self, adata: AnnData, dens_norm: bool):
+    # only to 15 because in kernel, if a row sums to 0, abs. states are created
+    # this happens because k_thresh = frac_to_keep = 0
+    @pytest.mark.parametrize("k", range(1, 15))
+    def test_pseudotime_frac_to_keep(self, adata: AnnData, k: int):
         conn = _get_neighs(adata, "connectivities")
         n_neighbors = _get_neighs_params(adata)["n_neighbors"]
         pseudotime = adata.obs["latent_time"]
+        k_thresh = max(0, min(int(np.floor(n_neighbors / k)) - 1, 30))
 
-        conn_biased = bias_knn(conn, pseudotime, n_neighbors)
-        if dens_norm:
-            T_1 = density_normalization(conn_biased, conn)
-            T_1 = _normalize(T_1)
-        else:
-            T_1 = _normalize(conn_biased)
+        conn_biased = bias_knn(conn.copy(), pseudotime, n_neighbors, k=k)
+        T_1 = _normalize(conn_biased)
 
         pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
-            density_normalize=dens_norm,
-            frac_to_keep=1 / 3.0,
+            frac_to_keep=k_thresh / float(n_neighbors),
             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):
+    def test_pseudotime_parallelize(self, adata: AnnData):
+        pk1 = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
+            n_jobs=None
+        )
+        pk2 = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
+            n_jobs=2
+        )
+
+        np.testing.assert_allclose(
+            pk1.transition_matrix.A, pk2.transition_matrix.A, rtol=_rtol
+        )
+
+    def test_pseudotime_inverse(self, adata: AnnData):
         pk = PseudotimeKernel(adata, time_key="latent_time")
         pt = pk.pseudotime.copy()
 
@@ -584,7 +666,7 @@ def test_palantir_inverse(self, adata: AnnData):
         assert pk_inv.backward
         np.testing.assert_allclose(pt, 1 - pk_inv.pseudotime)
 
-    def test_palantir_differ_dense_norm(self, adata: AnnData):
+    def test_pseudotime_differ_dense_norm(self, adata: AnnData):
         conn = _get_neighs(adata, "connectivities")
         n_neighbors = _get_neighs_params(adata)["n_neighbors"]
         pseudotime = adata.obs["latent_time"]
@@ -593,9 +675,7 @@ def test_palantir_differ_dense_norm(self, adata: AnnData):
         T_1 = density_normalization(conn_biased, conn)
         T_1 = _normalize(T_1)
 
-        pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix(
-            density_normalize=False
-        )
+        pk = PseudotimeKernel(adata, time_key="latent_time").compute_transition_matrix()
         T_2 = pk.transition_matrix
 
         assert not np.allclose(T_1.A, T_2.A, rtol=_rtol)
@@ -751,7 +831,7 @@ def test_connectivities_key_kernel(self, adata: AnnData):
         T_cr = ck.transition_matrix
 
         assert key == ck.params["key"]
-        np.testing.assert_array_equal(T_cr.A, adata.obsp[key])
+        np.testing.assert_array_equal(T_cr, adata.obsp[key])
 
         del adata.obsp[key]
 
@@ -1365,7 +1445,6 @@ def test_custom_scheme(self, adata: AnnData):
             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)
@@ -1467,3 +1546,62 @@ def test_regression_gene_corr(self, adata_large: AnnData):
         gene_corr_actual = adata_large.var[_ct("gene_corr")].values
 
         np.testing.assert_array_equal(gene_corr_actual, gene_corr_expected)
+
+
+class TestSingleFlow:
+    def test_no_transition_matrix(self, kernel: Kernel):
+        kernel._transition_matrix = None
+        with pytest.raises(RuntimeError, match=r"Compute transition matrix first as"):
+            kernel.plot_single_flow("Astrocytes", "clusters", "age(days)")
+
+    def test_invalid_cluster_key(self, kernel: Kernel):
+        with pytest.raises(KeyError, match=r"Unable to find clusters in"):
+            kernel.plot_single_flow("Astrocytes", "foo", "age(days)")
+
+    def test_invalid_source_cluster(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"Invalid source cluster"):
+            kernel.plot_single_flow("foo", "clusters", "age(days)")
+
+    def test_too_few_invalid_clusters(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"Expected at least `2` clusters"):
+            kernel.plot_single_flow(
+                "Astrocytes", "clusters", "age(days)", clusters=["foo", "bar", "baz"]
+            )
+
+    def test_all_invalid_clusters(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"No valid clusters have been selected."):
+            kernel.plot_single_flow(
+                "quux", "clusters", "age(days)", clusters=["foo", "bar", "baz"]
+            )
+
+    def test_invalid_time_key(self, kernel: Kernel):
+        with pytest.raises(
+            KeyError, match=r"Unable to find data in `adata.obs\['foo'\]`."
+        ):
+            kernel.plot_single_flow("Astrocytes", "clusters", "foo")
+
+    def test_too_few_valid_timepoints(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"Expected at least `2` time points"):
+            kernel.plot_single_flow(
+                "Astrocytes", "clusters", "age(days)", time_points=["35"]
+            )
+
+    def test_all_invalid_times(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"No valid time points"):
+            kernel.plot_single_flow(
+                "Astrocytes", "clusters", "age(days)", time_points=[0, 1, 2]
+            )
+
+    def test_time_key_cannot_be_coerced_to_numeric(self, kernel: Kernel):
+        with pytest.raises(TypeError, match=r"Unable to convert .* to `float`."):
+            kernel.plot_single_flow("Astrocytes", "clusters", "clusters")
+
+    def test_remove_empty_clusters_none_remain(self, kernel: Kernel):
+        with pytest.raises(ValueError, match=r"After removing clusters with no"):
+            kernel.plot_single_flow(
+                "Astrocytes",
+                "clusters",
+                "age(days)",
+                min_flow=np.inf,
+                remove_empty_clusters=True,
+            )
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index a14d7516c..fee830674 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -6,16 +6,16 @@
 
 import pytest
 
-import numpy as np
-from pandas import DataFrame
-
-import matplotlib.colors as colors
-
 from cellrank.tl import Lineage
 from cellrank.tl._colors import _compute_mean_color, _create_categorical_colors
 from cellrank.tl._lineage import _HT_CELLS, LineageView, PrimingDegree
 from cellrank.tl._constants import Lin
 
+import numpy as np
+from pandas import DataFrame
+
+import matplotlib.colors as colors
+
 
 class SimpleHTMLValidator(HTMLParser):
     _expected_tags = {"table", "div", "tr", "th", "td", "p"}
diff --git a/tests/test_linear_solver.py b/tests/test_linear_solver.py
index 23e8e36b4..a6fd97e7e 100644
--- a/tests/test_linear_solver.py
+++ b/tests/test_linear_solver.py
@@ -1,15 +1,15 @@
 import pytest
 
-import numpy as np
-from scipy.sparse import eye as speye
-from scipy.sparse import random, csr_matrix
-
 from cellrank.tl._linear_solver import (
     _petsc_mat_solve,
     _solve_lin_system,
     _create_petsc_matrix,
 )
 
+import numpy as np
+from scipy.sparse import eye as speye
+from scipy.sparse import random, csr_matrix
+
 
 def _petsc_not_installed() -> bool:
     try:
diff --git a/tests/test_logging.py b/tests/test_logging.py
index bd1f01c56..2feed34e3 100644
--- a/tests/test_logging.py
+++ b/tests/test_logging.py
@@ -7,7 +7,6 @@
 import pytest
 
 from scanpy import Verbosity
-
 from cellrank import logging as logg
 from cellrank import settings as settings
 
diff --git a/tests/test_model.py b/tests/test_model.py
index 23d5ce3ec..c85a2013b 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -7,12 +7,6 @@
 from _helpers import gamr_skip, create_model, assert_models_equal
 
 from anndata import AnnData
-
-import numpy as np
-from pygam import ExpectileGAM
-from scipy.stats import rankdata
-from sklearn.svm import SVR
-
 from cellrank.tl import Lineage
 from cellrank.ul.models import GAM, GAMR, FittedModel, SKLearnModel
 from cellrank.tl._constants import AbsProbKey
@@ -27,6 +21,11 @@
 from cellrank.ul.models._base_model import FailedModel, UnknownModelError
 from cellrank.ul.models._pygam_model import GamDistribution, GamLinkFunction, _gams
 
+import numpy as np
+from pygam import ExpectileGAM
+from scipy.stats import rankdata
+from sklearn.svm import SVR
+
 
 class TestModel:
     def test_wrong_type(self):
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index b0f0439a0..c4bf926d5 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -1,11 +1,7 @@
 import pytest
 
-from anndata import AnnData
-
-import numpy as np
-from pandas.api.types import is_categorical_dtype
-
 import cellrank as cr
+from anndata import AnnData
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
 from cellrank.tl._constants import (
     Direction,
@@ -19,6 +15,9 @@
 )
 from cellrank.tl.estimators._constants import P
 
+import numpy as np
+from pandas.api.types import is_categorical_dtype
+
 
 def _assert_has_all_keys(adata: AnnData, direction: Direction):
     assert _transition(direction) in adata.obsp.keys()
@@ -195,9 +194,9 @@ def test_bwd_pipeline_gpcca(self, adata: AnnData):
 
         _assert_has_all_keys(adata, Direction.BACKWARD)
 
-    def test_multiple_read_write(self, adata: AnnData):
+    def test_multiple_read_write_diff(self, adata: AnnData):
         for n_states in range(2, 5):
-            cr.tl.terminal_states(
+            e = cr.tl.terminal_states(
                 adata,
                 estimator=cr.tl.estimators.GPCCA,
                 cluster_key="clusters",
@@ -205,7 +204,10 @@ def test_multiple_read_write(self, adata: AnnData):
                 show_plots=True,
                 n_states=n_states,
                 fit_kwargs={"n_cells": 5},
+                return_estimator=True,
             )
+            assert e.absorption_probabilities is None
+
             cr.tl.lineages(adata, backward=False)
             cr.pl.lineages(adata)
 
@@ -216,6 +218,26 @@ def test_multiple_read_write(self, adata: AnnData):
             for name in probs.names:
                 assert "Lineage" not in name, probs.names
 
+    def test_multiple_read_write_same(self, adata: AnnData):
+        e = cr.tl.terminal_states(
+            adata,
+            estimator=cr.tl.estimators.GPCCA,
+            cluster_key="clusters",
+            method="brandts",
+            show_plots=True,
+            n_states=2,
+            fit_kwargs={"n_cells": 5},
+            return_estimator=True,
+        )
+        cr.tl.lineages(adata, backward=False)
+        cr.pl.lineages(adata)
+
+        e.set_terminal_states({"foo": adata.obs_names[:20]})
+
+        e = cr.tl.estimators.GPCCA(adata, read_from_adata=True, obsp_key="T_fwd")
+        assert e.macrostates is None
+        assert e.absorption_probabilities is None
+
 
 class TestLowLevelPipeline:
     def test_fwd_pipelne_cflare(self, adata: AnnData):
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index f6213c95b..e50da6fa5 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -12,7 +12,12 @@
 from packaging import version
 
 import scvelo as scv
+import cellrank as cr
 from anndata import AnnData
+from cellrank.tl import Lineage
+from cellrank.ul.models import GAMR
+from cellrank.tl._constants import AbsProbKey, _pd
+from cellrank.tl.estimators import GPCCA, CFLARE
 
 import numpy as np
 import pandas as pd
@@ -20,15 +25,10 @@
 from pandas.api.types import is_categorical_dtype
 
 import matplotlib.cm as cm
+import matplotlib.pyplot as plt
 from matplotlib.testing import setup
 from matplotlib.testing.compare import compare_images
 
-import cellrank as cr
-from cellrank.tl import Lineage
-from cellrank.ul.models import GAMR
-from cellrank.tl._constants import AbsProbKey, _pd
-from cellrank.tl.estimators import GPCCA, CFLARE
-
 setup()
 
 HERE: str = Path(__file__).parent
@@ -106,6 +106,8 @@ def _assert_equal(fpath: str) -> None:
             fpath += ".png"
         if dirname is not None:
             for file in os.listdir(FIGS / dirname):
+                if "-diff" in file:
+                    continue
                 _compare_images(GT_FIGS / dirname / file, FIGS / dirname / file)
         else:
             _compare_images(GT_FIGS / fpath, FIGS / fpath)
@@ -723,6 +725,53 @@ def test_cluster_lineage_key(self, adata_cflare: AnnData):
         assert adata_cflare.uns["foobar"].shape == (10, 200)
         assert is_categorical_dtype(adata_cflare.uns["foobar"].obs["clusters"])
 
+    @compare()
+    def test_cluster_lineage_covariates(self, adata: AnnData, fpath: str):
+        model = create_model(adata)
+        cr.pl.cluster_lineage(
+            adata,
+            model,
+            GENES[:10],
+            "1",
+            covariate_key=["clusters", "latent_time"],
+            random_state=0,
+            time_key="latent_time",
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare()
+    def test_cluster_lineage_covariates_cmap(self, adata: AnnData, fpath: str):
+        model = create_model(adata)
+        cr.pl.cluster_lineage(
+            adata,
+            model,
+            GENES[:10],
+            "1",
+            covariate_key="latent_time",
+            cmap="inferno",
+            random_state=0,
+            time_key="latent_time",
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare()
+    def test_cluster_lineage_covariates_ratio(self, adata: AnnData, fpath: str):
+        model = create_model(adata)
+        cr.pl.cluster_lineage(
+            adata,
+            model,
+            GENES[:10],
+            "1",
+            covariate_key="latent_time",
+            ratio=0.25,
+            random_state=0,
+            time_key="latent_time",
+            dpi=DPI,
+            save=fpath,
+        )
+
 
 class TestHeatmap:
     @compare(dirname="heatmap_lineages")
@@ -3169,3 +3218,562 @@ def test_kernel_random_walk_ixs_legend_loc(self, mc: GPCCA, fpath: str):
             dpi=DPI,
             save=fpath,
         )
+
+
+class TestPlotSingleFlow:
+    @compare(kind="gpcca")
+    def test_flow_source_clusters(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Neuroblast",
+            "clusters",
+            "age(days)",
+            clusters=["OPC", "Endothelial", "OL"],
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_clusters_subset(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            clusters=["OPC", "Endothelial", "OL"],
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_min_flow_remove_empty_clusters(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            min_flow=0.2,
+            remove_empty_clusters=True,
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_min_flow_keep_empty_clusters(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            min_flow=0.2,
+            remove_empty_clusters=False,
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_cluster_ascending(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes", "clusters", "age(days)", ascending=True, dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_cluster_descending(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes", "clusters", "age(days)", ascending=False, dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_explicit_cluster_order(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            ascending=None,
+            clusters=["OPC", "OL"],
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_legend_loc(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            legend_loc="upper left out",
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_alpha(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes", "clusters", "age(days)", alpha=0.3, dpi=DPI, save=fpath
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_no_xticks(self, mc: GPCCA, fpath: str):
+        mc.kernel.plot_single_flow(
+            "Astrocytes",
+            "clusters",
+            "age(days)",
+            xticks_step_size=None,
+            dpi=DPI,
+            save=fpath,
+        )
+
+    @compare(kind="gpcca")
+    def test_flow_time_categories_too_close(self, mc: GPCCA, fpath: str):
+        mc.adata.obs["day"] = (
+            mc.adata.obs["age(days)"]
+            .cat.rename_categories(
+                {
+                    "12": 0.1,
+                    "35": 0.291,
+                }
+            )
+            .values
+        )
+        mc.kernel.plot_single_flow("Astrocytes", "clusters", "day", dpi=DPI, save=fpath)
+
+    @compare(kind="gpcca")
+    def test_flow_return_ax(self, mc: GPCCA, fpath: str):
+        ax = mc.kernel.plot_single_flow(
+            "Astrocytes", "clusters", "age(days)", show=False, dpi=DPI, save=fpath
+        )
+        assert isinstance(ax, plt.Axes)
+
+
+class TestPlotDriverCorrelation:
+    @compare(kind="gpcca")
+    def test_driver_corr(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "1", "2", dpi=DPI, save=fpath, title="bar", size=100
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_color(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0", "1", dpi=DPI, save=fpath, color="to 2 corr"
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_gene_sets(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0", "1", dpi=DPI, save=fpath, gene_sets={"0": mc.adata.var_names[:3]}
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_gene_sets_colors(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0",
+            "1",
+            dpi=DPI,
+            save=fpath,
+            gene_sets={"0": mc.adata.var_names[:3], "1": [mc.adata.var_names[4]]},
+            gene_sets_colors=["red", "black"],
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_legend_loc(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0",
+            "1",
+            dpi=DPI,
+            save=fpath,
+            gene_sets={"0": mc.adata.var_names[:3], "1": [mc.adata.var_names[4]]},
+            legend_loc="lower center out",
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_use_raw(self, mc: GPCCA, fpath: str):
+        mc.compute_lineage_drivers(cluster_key="clusters", use_raw=True)
+        mc.plot_lineage_drivers_correlation(
+            "0", "1", dpi=DPI, save=fpath, use_raw=True, color="to 1 qval"
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_cmap(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0", "1", dpi=DPI, save=fpath, color="to 1 qval", cmap="inferno"
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_fontsize(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0",
+            "1",
+            dpi=DPI,
+            save=fpath,
+            gene_sets={"foo": mc.adata.var_names[4:6]},
+            fontsize=20,
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_adjust_text(self, mc: GPCCA, fpath: str):
+        mc.plot_lineage_drivers_correlation(
+            "0",
+            "1",
+            dpi=DPI,
+            save=fpath,
+            gene_sets={"bar": mc.adata.var_names[:3]},
+            adjust_text=True,
+        )
+
+    @compare(kind="gpcca")
+    def test_driver_corr_return_ax(self, mc: GPCCA, fpath: str):
+        ax = mc.plot_lineage_drivers_correlation(
+            "2", "0", dpi=DPI, save=fpath, show=False
+        )
+        assert isinstance(ax, plt.Axes)
+
+
+class TestLogOdds:
+    @compare()
+    def test_log_odds(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare(kind="bwd")
+    def test_log_odds_bwd(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            backward=True,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_rest(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "2",
+            None,
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_continuous_keys(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[:3],
+            figsize=(4, 3),
+            size=4,
+        )
+
+    @compare()
+    def test_log_odds_categorical_keys(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=["clusters", "clusters_enlarged"],
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_threshold(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[:3],
+            threshold=0.5,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_multiple_threshold(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[:3],
+            threshold=[0.7, 0.2, 0.3],
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_threshold_color(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[:3],
+            threshold=0.5,
+            threshold_color="blue",
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_layer(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[3:6],
+            layer="Ms",
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_use_raw(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.raw.var_names[3:6],
+            use_raw=True,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_size(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys="clusters",
+            size=20,
+            figsize=(4, 3),
+        )
+
+    @compare()
+    def test_log_odds_cmap(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=adata.var_names[:2],
+            size=10,
+            cmap="inferno",
+            figsize=(4, 3),
+        )
+
+    @compare()
+    def test_log_odds_alpha(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys="clusters",
+            alpha=0.5,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_ncols(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=["clusters", adata.var_names[-1]],
+            ncols=1,
+            figsize=(3, 4),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_fontsize(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys="clusters",
+            fontsize=25,
+            figsize=(3, 4),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_xticks_steps_size(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys="clusters",
+            xticks_step_size=None,
+            figsize=(3, 4),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_legend_loc(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            keys=["clusters", adata.var_names[-1]],
+            legend_loc="upper right out",
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+        )
+
+    @compare()
+    def test_log_odds_jitter(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "0",
+            "1",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            figsize=(4, 3),
+            size=10,
+            seed=42,
+            jitter=1,
+        )
+
+    @compare()
+    def test_log_odds_kwargs_return_ax(self, adata: AnnData, fpath: str):
+        ax = cr.pl.log_odds(
+            adata,
+            "1",
+            "2",
+            "age(days)",
+            keys="clusters",
+            dpi=DPI,
+            save=fpath,
+            show=False,
+            edgecolor="red",
+            figsize=(4, 3),
+            size=4,
+        )
+        assert isinstance(ax, plt.Axes)
+
+    @compare()
+    def test_log_odds_kwargs_return_axes(self, adata: AnnData, fpath: str):
+        axes = cr.pl.log_odds(
+            adata,
+            "1",
+            "2",
+            "age(days)",
+            keys=adata.var_names[:3],
+            dpi=DPI,
+            save=fpath,
+            ncols=2,
+            show=False,
+            figsize=(4, 3),
+            size=4,
+        )
+        assert isinstance(axes, np.ndarray)
+        assert axes.shape == (3,)
+        assert np.all([isinstance(ax, plt.Axes) for ax in axes])
+
+    @compare()
+    def test_log_odds_kwargs(self, adata: AnnData, fpath: str):
+        cr.pl.log_odds(
+            adata,
+            "1",
+            "2",
+            "age(days)",
+            dpi=DPI,
+            save=fpath,
+            linewidth=5,
+            edgecolor="red",
+            figsize=(4, 3),
+            size=4,
+        )
+
+
+class TestMacrostateComposition:
+    @compare(kind="gpcca")
+    def test_msc_default(self, mc: GPCCA, fpath: str):
+        mc.plot_macrostate_composition("clusters", dpi=DPI, save=fpath)
+
+    @compare(kind="gpcca")
+    def test_msc_width(self, mc: GPCCA, fpath: str):
+        mc.plot_macrostate_composition("clusters", dpi=DPI, save=fpath, width=0.2)
+
+    @compare(kind="gpcca")
+    def test_msc_title(self, mc: GPCCA, fpath: str):
+        mc.plot_macrostate_composition("clusters", dpi=DPI, save=fpath, title="foobar")
+
+    @compare(kind="gpcca")
+    def test_msc_labelrot(self, mc: GPCCA, fpath: str):
+        mc.plot_macrostate_composition("clusters", dpi=DPI, save=fpath, labelrot=0)
+
+    @compare(kind="gpcca")
+    def test_msc_legend_loc(self, mc: GPCCA, fpath: str):
+        mc.plot_macrostate_composition(
+            "clusters_enlarged", dpi=DPI, save=fpath, legend_loc="upper left out"
+        )
diff --git a/tests/test_random_walk.py b/tests/test_random_walk.py
index 291a68193..5d48d2cbd 100644
--- a/tests/test_random_walk.py
+++ b/tests/test_random_walk.py
@@ -1,9 +1,9 @@
 import pytest
 
-import numpy as np
-
 from cellrank.tl.kernels._random_walk import RandomWalk
 
+import numpy as np
+
 
 class TestRandomWalk:
     def test_matrix_not_row_stochastic(self, test_matrix_4: np.ndarray):
diff --git a/tests/test_read.py b/tests/test_read.py
index 69971ee30..d8e2ed865 100644
--- a/tests/test_read.py
+++ b/tests/test_read.py
@@ -4,16 +4,15 @@
 import pytest
 
 import scanpy as sc
-from anndata import AnnData
-
-import numpy as np
-
 import cellrank as cr
+from anndata import AnnData
 from cellrank.tl._colors import _create_categorical_colors
 from cellrank.tl._lineage import Lineage
 from cellrank.tl._constants import AbsProbKey, _colors, _lin_names
 from cellrank.tl.estimators._cflare import CFLARE
 
+import numpy as np
+
 
 def test_fwd():
     def wrapper(func: Callable) -> Callable:
diff --git a/tests/test_tools.py b/tests/test_tools.py
index f3f842a44..872ba95b9 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -1,14 +1,15 @@
+from typing import Optional, Sequence
+
 import pytest
 
+import cellrank as cr
 from anndata import AnnData
+from cellrank.tl._constants import AbsProbKey, TermStatesKey, _probs
+from cellrank.tl.kernels._base_kernel import KernelAdd
 
 import numpy as np
 import pandas as pd
 
-import cellrank as cr
-from cellrank.tl._constants import AbsProbKey, TermStatesKey, _probs
-from cellrank.tl.kernels._base_kernel import KernelAdd
-
 
 class TestLineages:
     def test_no_root_cells(self, adata: AnnData):
@@ -55,26 +56,22 @@ def test_normal_run_raw(self, adata_cflare):
             assert np.all(adata_cflare.raw.var[f"to {name} qval"] >= 0)
             assert np.all(adata_cflare.raw.var[f"to {name} qval"] <= 1.0)
 
-    def test_return_drivers(self, adata_cflare):
+    @pytest.mark.parametrize("lineages", [None, ["0"], ["0, 1"]])
+    def test_return_drivers(self, adata_cflare, lineages: Optional[Sequence[str]]):
+        cr.tl.terminal_states(adata_cflare, n_states=3)
         cr.tl.lineages(adata_cflare)
-        res = cr.tl.lineage_drivers(adata_cflare, return_drivers=True, use_raw=False)
+        res = cr.tl.lineage_drivers(
+            adata_cflare, return_drivers=True, use_raw=False, lineages=lineages
+        )
+        exp_lineages = adata_cflare.obsm["to_terminal_states"]
+        if lineages is not None:
+            exp_lineages = exp_lineages[lineages]
+        exp_lineages = exp_lineages.names
 
         assert isinstance(res, pd.DataFrame)
         assert res.shape[0] == adata_cflare.n_vars
-        assert {
-            "0 corr",
-            "0 pval",
-            "0 qval",
-            "0 ci low",
-            "0 ci high",
-            "1 corr",
-            "1 pval",
-            "1 qval",
-            "1 ci low",
-            "1 ci high",
-        } == set(res.columns)
 
-        for name in ["0", "1"]:
+        for name in exp_lineages:
             assert np.all(adata_cflare.var[f"to {name} corr"] >= -1.0)
             assert np.all(adata_cflare.var[f"to {name} corr"] <= 1.0)
             # res is sorted
@@ -226,6 +223,25 @@ def test_invalid_invalid_mode(self, adata: AnnData):
         with pytest.raises(ValueError):
             cr.tl.initial_states(adata, mode="foobar")
 
+    @pytest.mark.parametrize("force_recompute", [False, True])
+    def test_force_recompute(self, adata: AnnData, force_recompute: bool):
+        cr.tl.terminal_states(adata, n_states=5, fit_kwargs=dict(n_cells=5))
+        tmat1 = adata.obsp["T_fwd"]
+
+        cr.tl.terminal_states(
+            adata,
+            n_states=5,
+            fit_kwargs=dict(n_cells=5),
+            force_recompute=force_recompute,
+        )
+        tmat2 = adata.obsp["T_fwd"]
+        if force_recompute:
+            assert tmat1 is not tmat2
+            np.testing.assert_allclose(tmat1.A, tmat2.A)
+            np.testing.assert_allclose(tmat1.sum(1), 1.0)
+        else:
+            assert tmat1 is tmat2
+
 
 class TestTransitionMatrix:
     def test_invalid_velocity_key(self, adata: AnnData):
@@ -260,7 +276,7 @@ def test_connectivities_conn_key(self, adata: AnnData):
 
         assert isinstance(ck, cr.tl.kernels.ConnectivityKernel)
 
-        np.testing.assert_array_equal(ck.transition_matrix.A, adata.obsp[key])
+        np.testing.assert_array_equal(ck.transition_matrix, adata.obsp[key])
 
     def test_only_velocity(self, adata: AnnData):
         vk = cr.tl.transition_matrix(adata, weight_connectivities=0)
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 47f1b3b8b..5ca3dd441 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -1,15 +1,10 @@
+from typing import Any
+
 import pytest
 from _helpers import create_model, assert_array_nan_equal, jax_not_installed_skip
 
 import scanpy as sc
 from anndata import AnnData
-
-import numpy as np
-import pandas as pd
-from numba import njit
-from scipy.sparse import diags, random, csr_matrix
-from pandas.api.types import is_categorical_dtype
-
 from cellrank.tl import Lineage
 from cellrank.pl._utils import (
     _create_models,
@@ -30,6 +25,7 @@
 )
 from cellrank.ul.models import GAM, BaseModel
 from cellrank.tl._colors import _compute_mean_color
+from cellrank.ul._parallelize import parallelize
 from cellrank.tl.kernels._utils import (
     norm,
     np_max,
@@ -46,6 +42,13 @@
     _predict_transition_probabilities_numpy,
 )
 
+import numpy as np
+import pandas as pd
+from numba import njit
+from scipy.sparse import rand as srand
+from scipy.sparse import diags, random, csr_matrix
+from pandas.api.types import is_categorical_dtype
+
 
 class TestToolsUtils:
     def test_merge_not_categorical(self):
@@ -1071,3 +1074,23 @@ def test_random_normal_1_sample(self):
         x = _random_normal(np.array([0]), np.array([1]), 1)
 
         assert x.shape == (1, 1)
+
+
+class TestParallelize:
+    @pytest.mark.parametrize("n_jobs", [1, 3, 4])
+    def test_more_jobs_than_work(self, n_jobs: int):
+        def callback(data, **_: Any):
+            assert isinstance(data, csr_matrix)
+            assert data.shape[1] == 100
+
+            return [42] * data.shape[0]
+
+        res = parallelize(
+            callback,
+            collection=srand(3, 100, format="csr"),
+            n_jobs=n_jobs,
+            show_progress_bar=False,
+            extractor=np.concatenate,
+        )()
+
+        np.testing.assert_array_equal(res, 42)