Repository: theislab/cellrank (language: python)

Issue:
Refactor estimators
<!-- 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: -->

TODOs:
- [ ] rephrase it as a stateful machine (easier to handle IO from/to adata)?
- [ ] remove `PropertyMeta` (was used during development due to frequent naming changes)
- [ ] related to ^, get rid of https://github.com/theislab/cellrank/blob/master/cellrank/tl/estimators/_constants.py
- [ ] change default tolerance for solvers in abs. prob. calculation/add preconditioner?
- [ ] allow option to serialize self to a file without `anndata`


Discussion / hints:
For the preconditioners, we should probably do a quick benchmark, i.e. take the three datasets we have on Figshare and compare `gmres` with and without preconditioner to the direct solver. Our assumption here is that the direct solver gives us ground truth. Probably, we can't do this for the morris data because it's way too large for the direct solver. 
Hi @michalk8, what's a realistic timeline for this?
> Hi @michalk8, what's a realistic timeline for this?

Next week.

Fix patch:
diff --git a/cellrank/__init__.py b/cellrank/__init__.py
index 9507e110d..8ff6cfadb 100644
--- a/cellrank/__init__.py
+++ b/cellrank/__init__.py
@@ -1,3 +1,4 @@
+# TODO: in 2.0, remove .tl
 import cellrank.pl
 import cellrank.tl
 import cellrank.ul
@@ -9,7 +10,7 @@
 import cellrank.estimators
 from cellrank.settings import settings
 from cellrank.tl._read import read
-from cellrank.tl._constants import Lin
+from cellrank.tl._lineage import Lin
 
 __author__ = ", ".join(["Marius Lange", "Michal Klein"])
 __maintainer__ = ", ".join(["Marius Lange", "Michal Klein"])
diff --git a/cellrank/_key.py b/cellrank/_key.py
new file mode 100644
index 000000000..82270abb3
--- /dev/null
+++ b/cellrank/_key.py
@@ -0,0 +1,103 @@
+from typing import Any, Callable, Optional
+
+
+class cprop:
+    """Class property."""
+
+    def __init__(self, f: Callable[..., str]):
+        self.f = f
+
+    def __get__(self, obj: Any, owner: Any) -> str:
+        return self.f(owner)
+
+
+class Key:
+    """Class which manages keys in :class:`anndata.AnnData`."""
+
+    @classmethod
+    def backward(cls, bwd: bool) -> str:
+        return "bwd" if bwd else "fwd"
+
+    @classmethod
+    def where(cls, bwd: bool) -> str:
+        return "from" if bwd else "to"
+
+    @classmethod
+    def initial(cls, bwd: bool) -> str:
+        return "initial" if bwd else "terminal"
+
+    @classmethod
+    def cytotrace(cls, key: str) -> str:
+        return f"ct_{key}"
+
+    class obs:
+        @classmethod
+        def probs(cls, key: str) -> str:
+            return f"{key}_probs"
+
+        @classmethod
+        def macrostates(cls, bwd: bool) -> str:
+            return f"macrostates_{Key.backward(bwd)}"
+
+        @classmethod
+        def term_states(cls, bwd: bool) -> str:
+            return f"{Key.initial(bwd)}_states"
+
+        @classmethod
+        def priming_degree(cls, bwd: bool) -> str:
+            return f"priming_degree_{Key.backward(bwd)}"
+
+    class obsm:
+        @classmethod
+        def memberships(cls, key: str) -> str:
+            return f"{key}_memberships"
+
+        @classmethod
+        def schur_vectors(cls, bwd: bool) -> str:
+            return f"schur_vectors_{Key.backward(bwd)}"
+
+        @classmethod
+        def macrostates(cls, bwd: bool) -> str:
+            return f"macrostates_{Key.backward(bwd)}"
+
+        @classmethod
+        def abs_probs(cls, bwd: bool) -> str:
+            return ("from" if bwd else "to") + "_" + Key.obs.term_states(bwd)
+
+        @classmethod
+        def abs_times(cls, bwd: bool) -> str:
+            return f"absorption_times_{Key.backward(bwd)}"
+
+    class varm:
+        @classmethod
+        def lineage_drivers(cls, bwd: bool):
+            return ("initial" if bwd else "terminal") + "_lineage_drivers"
+
+    class uns:
+        @classmethod
+        def kernel(cls, bwd: bool, key: Optional[str] = None) -> str:
+            return key if key is not None else f"T_{Key.backward(bwd)}"
+
+        @classmethod
+        def estimator(cls, bwd: bool, key: Optional[str] = None) -> str:
+            return key if key is not None else f"{Key.backward(bwd)}_estimator"
+
+        @classmethod
+        def names(cls, key: str) -> str:
+            return f"{key}_names"
+
+        @classmethod
+        def colors(cls, key: str) -> str:
+            return f"{key}_colors"
+
+        @classmethod
+        def eigen(cls, bwd: bool) -> str:
+            return f"eigendecomposition_{Key.backward(bwd)}"
+
+        @classmethod
+        def schur_matrix(cls, bwd: bool) -> str:
+            return f"schur_matrix_{Key.backward(bwd)}"
+
+        @classmethod
+        def coarse(cls, bwd: bool) -> str:
+            return f"coarse_{Key.backward(bwd)}"
diff --git a/cellrank/datasets/_datasets.py b/cellrank/datasets/_datasets.py
index bf9a09c94..685a61f55 100644
--- a/cellrank/datasets/_datasets.py
+++ b/cellrank/datasets/_datasets.py
@@ -1,20 +1,19 @@
-"""Datasets module."""
-
-from typing import Any, Tuple, Union, TypeVar
+from typing import Any, Tuple, Union
+from typing_extensions import Literal
 
 import os
+from enum import auto
 from pathlib import Path
 
 from scanpy import read
+from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._constants import ModeEnum
-
-AnnData = TypeVar("AnnData")
 
 
 class ReprogrammingSubset(ModeEnum):  # noqa: D101
-    FULL = "full"
+    FULL = auto()
     K48 = "48k"
     K85 = "85k"
 
@@ -162,7 +161,7 @@ def lung(
 @inject_docs(s=ReprogrammingSubset)
 @d.dedent
 def reprogramming_morris(
-    subset: str = ReprogrammingSubset.FULL.s,
+    subset: Literal["full", "48k", "85k"] = ReprogrammingSubset.FULL,
     path: Union[str, Path] = "datasets/reprogramming_morris.h5ad",
     **kwargs: Any,
 ) -> AnnData:
@@ -170,7 +169,7 @@ def reprogramming_morris(
     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 :cite:`macosko:15`
+    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
@@ -185,9 +184,9 @@ def reprogramming_morris(
     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 :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.
+            - `{s.FULL!r}` - return the complete dataset containing `104 887` cells.
+            - `{s.K85!r}` - return the subset as described in :cite:`morris:18` Fig. 1, containing `85 010` cells.
+            - `{s.K48!r}` - return the subset as described in :cite:`morris:18` Fig. 3, containing `48 515` cells.
 
     %(dataset.parameters)s
 
@@ -209,9 +208,7 @@ def reprogramming_morris(
     if subset == ReprogrammingSubset.K85:
         return adata[~adata.obs["timecourse"].isnull()].copy()
 
-    raise NotImplementedError(
-        f"Subsetting option `{subset.s!r}` is not yet implemented."
-    )
+    raise NotImplementedError(f"Subsetting option `{subset!r}` is not yet implemented.")
 
 
 @d.dedent
diff --git a/cellrank/estimators.py b/cellrank/estimators.py
index 5dfadbf0a..304d88830 100644
--- a/cellrank/estimators.py
+++ b/cellrank/estimators.py
@@ -1,3 +1,2 @@
-from cellrank.tl.estimators._gpcca import GPCCA
-from cellrank.tl.estimators._cflare import CFLARE
 from cellrank.tl.estimators._base_estimator import BaseEstimator
+from cellrank.tl.estimators.terminal_states import GPCCA, CFLARE
diff --git a/cellrank/external/kernels/_statot_kernel.py b/cellrank/external/kernels/_statot_kernel.py
index 8f2280ae8..75cad1ddd 100644
--- a/cellrank/external/kernels/_statot_kernel.py
+++ b/cellrank/external/kernels/_statot_kernel.py
@@ -2,9 +2,8 @@
 from typing_extensions import Literal
 
 from anndata import AnnData
+from cellrank._key import Key
 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
@@ -31,12 +30,11 @@ class StationaryOTKernel(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 :attr:`adata` ``['terminal_states']``.
+        Key in :attr:`anndata.AnnData.obs` or a categorical :class:`pandas.Series` where non-`NaN` values
+        mark terminal states.
+        If `None`, terminal states are assumed to be present in :attr:`anndata.AnnData.obs` ``['terminal_states']``.
     g
         Key in :attr:`anndata.AnnData.obs` containing relative growth rates for cells or the array itself.
-    cluster_key
-        If a key to cluster labels is given, `terminal_states` will be associated with these for naming and colors.
     kwargs
         Additional keyword arguments.
     """
@@ -48,13 +46,12 @@ def __init__(
         adata: AnnData,
         g: Union[str, np.ndarray],
         terminal_states: Optional[Union[str, pd.Series]] = None,
-        cluster_key: Optional[str] = None,
         **kwargs: Any,
     ):
-        if terminal_states is not None:
-            dk = DummyKernel(adata, backward=False)
-            estim = GPCCA(dk, write_to_adata=True)
-            estim.set_terminal_states(terminal_states, cluster_key=cluster_key)
+        if isinstance(terminal_states, str):
+            adata.obs[Key.obs.term_states(bwd=False)] = adata.obs[terminal_states]
+        elif isinstance(terminal_states, pd.Series):
+            adata.obs[Key.obs.term_states(bwd=False)] = terminal_states
 
         try:
             super().__init__(adata, g=g, **kwargs)
@@ -100,9 +97,9 @@ def compute_transition_matrix(
         method
             Choice of regularization. Valid options are:
 
-                - `'ent'`: entropy.
-                - `'quad'`: L2-norm.
-                - `'unbal'`: unbalanced transport (not yet implemented).
+                - `'ent'` - entropy.
+                - `'quad'` - L2-norm.
+                - `'unbal'` - unbalanced transport (not yet implemented).
 
         tol
             Relative tolerance for OT solver convergence.
diff --git a/cellrank/external/kernels/_wot_kernel.py b/cellrank/external/kernels/_wot_kernel.py
index 0b5922153..3c72e3cfb 100644
--- a/cellrank/external/kernels/_wot_kernel.py
+++ b/cellrank/external/kernels/_wot_kernel.py
@@ -142,9 +142,9 @@ def compute_initial_growth_rates(
         Parameters
         ----------
         proliferation_key
-            Key in :attr:`adata` ``.obs`` where the birth rate score is saved.
+            Key in :attr:`anndata.AnnData.obs` where the birth rate score is saved.
         apoptosis_key
-            Key in :attr:`adata` ``.obs`` where the death rate score is saved.
+            Key in :attr:`anndata.AnnData.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
@@ -161,19 +161,19 @@ def compute_initial_growth_rates(
             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``.
+            or ``apoptosis_key`` cannot be found in :attr:`adata.AnnData.obs`.
 
         Returns
         -------
         :class:`pandas.Series`
-            The estimated initial growth rates if ``key_added=None``, otherwise `None`.
+            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`
+        :meth:`cellrank.external.kernels.WOTKernel.compute_transition_matrix` 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:
@@ -205,7 +205,7 @@ def get_scores(
 
                 sc.tl.score_genes(
                     self.adata,
-                    gene_list=getattr(MarkerGenes, (f"{kind}_markers"))(organism),
+                    gene_list=getattr(MarkerGenes, f"{kind}_markers")(organism),
                     score_name=score_name,
                     **kwargs,
                 )
@@ -265,7 +265,9 @@ def compute_transition_matrix(
         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",
+        last_time_point: Literal[
+            "uniform", "diagonal", "connectivities"
+        ] = LastTimePoint.UNIFORM,
         threshold: Optional[Union[float, Literal["auto"]]] = "auto",
         conn_kwargs: Mapping[str, Any] = MappingProxyType({}),
         **kwargs: Any,
@@ -282,9 +284,9 @@ def compute_transition_matrix(
         ----------
         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``
+            If a :class:`str`, it specifies a key in :attr:`anndata.AnnData.layers` or :attr:`anndata.AnnData.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.
+            compute distances in PCA space derived from :attr:`anndata.AnnData.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.
@@ -298,17 +300,17 @@ def compute_transition_matrix(
         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.
+            Key in :attr:`anndata.AnnData.obs` where initial cell growth rates are stored.
+            See :meth:`compute_initial_growth_rates` on how to estimate them.
         use_highly_variable
-            Key in :attr:`adata` ``.var`` where highly variable genes are stored.
+            Key in :attr:`anndata.AnnData.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`
+                - `{ltp.UNIFORM!r}` - row-normalized matrix of 1s for transitions within the last time point.
+                - `{ltp.DIAGONAL!r}` - diagonal matrix with 1s on the diagonal.
+                - `{ltp.CONNECTIVITIES!r}` - use transitions from :class:`cellrank.tl.kernels.ConnectivityKernel`
                   derived from the last time point subset of :attr:`adata`.
         threshold
             How to remove small non-zero values from the transition matrix. Valid options are:
@@ -319,17 +321,17 @@ def compute_transition_matrix(
                 - `None` - do not threshold.
 
         conn_kwargs
-            Keyword arguments for :func:`scanpy.pp.neighbors`, when using ``last_time_point={ltp.CONNECTIVITIES.s!r}``.
+            Keyword arguments for :func:`scanpy.pp.neighbors`, when using ``last_time_point = {ltp.CONNECTIVITIES!r}``.
             Can contain `'density_normalize'` for
             :meth:`cellrank.tl.kernels.ConnectivityKernel.compute_transition_matrix`.
         kwargs
-            Additional keyword arguments for OT configuration.
+            Additional keyword arguments for optimal transport 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:
+            Also modifies :attr:`anndata.AnnData.obs` with the following key:
 
                 - `'estimated_growth_rates'` - the estimated final growth rates.
 
@@ -361,7 +363,7 @@ def compute_transition_matrix(
                 "solver": solver,
                 "growth_rate_key": growth_rate_key,
                 "use_highly_variable": use_highly_variable,
-                "last_time_point": last_time_point.s,
+                "last_time_point": last_time_point,
                 "threshold": threshold,
                 **kwargs,
             },
diff --git a/cellrank/pl/_circular_projection.py b/cellrank/pl/_circular_projection.py
index c57721333..f7195cbfe 100644
--- a/cellrank/pl/_circular_projection.py
+++ b/cellrank/pl/_circular_projection.py
@@ -1,6 +1,7 @@
 from typing import Any, Tuple, Union, Mapping, Callable, Optional, Sequence
 from typing_extensions import Literal
 
+from enum import auto
 from types import MappingProxyType
 from pathlib import Path
 
@@ -8,12 +9,14 @@
 from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.tl import Lineage
+from cellrank._key import Key
+from scanpy._utils import deprecated_arg_names
+from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d
 from cellrank.pl._utils import _held_karp
 from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul._utils import _check_collection
 from cellrank.tl._lineage import PrimingDegree
-from cellrank.tl._constants import ModeEnum, AbsProbKey
 
 import numpy as np
 import pandas as pd
@@ -25,13 +28,13 @@
 
 
 class LineageOrder(ModeEnum):  # noqa: D101
-    DEFAULT = "default"
-    OPTIMAL = "optimal"
+    DEFAULT = auto()
+    OPTIMAL = auto()
 
 
 class LabelRot(ModeEnum):  # noqa: D101
-    DEFAULT = "default"
-    BEST = "best"
+    DEFAULT = auto()
+    BEST = auto()
 
 
 Metric_T = Union[str, Callable, np.ndarray, pd.DataFrame]
@@ -44,22 +47,11 @@ def _get_distances(data: Union[np.ndarray, Lineage], metric: Metric_T) -> np.nda
 
     if isinstance(metric, str) or callable(metric):
         metric = pairwise_distances(data.T, metric=metric)
-    elif isinstance(metric, np.ndarray):
+    elif isinstance(metric, (pd.DataFrame, np.ndarray)):
         shape = (data.shape[1], data.shape[1])
         if metric.shape != shape:
             raise ValueError(
-                f"Expected a `numpy.ndarray` of shape `{shape}`, found `{metric.shape}`."
-            )
-    elif isinstance(metric, pd.DataFrame):
-        if np.any(metric.index != data.names):
-            raise ValueError(
-                f"Expected `pandas.DataFrame` to have the following index `{list(data.names)}`, "
-                f"found `{list(metric.columns)}`."
-            )
-        if np.any(metric.columns != data.names):
-            raise ValueError(
-                f"Expected `pandas.DataFrame` to have the following columns `{list(data.names)}`, "
-                f"found `{list(metric.columns)}`."
+                f"Expected an `numpy.array` or `pandas.DataFrame` of shape `{shape}`, found `{metric.shape}`."
             )
     else:
         raise TypeError(
@@ -76,6 +68,7 @@ def _get_optimal_order(data: Lineage, metric: Metric_T) -> Tuple[float, np.ndarr
 
 
 @d.dedent
+@deprecated_arg_names({"labeldistance": "label_distance", "labelrot": "label_rot"})
 def circular_projection(
     adata: AnnData,
     keys: Union[str, Sequence[str]],
@@ -89,8 +82,8 @@ def circular_projection(
     space: float = 0.25,
     use_raw: bool = False,
     text_kwargs: Mapping[str, Any] = MappingProxyType({}),
-    labeldistance: float = 1.25,
-    labelrot: Union[Literal["default", "best"], float] = "best",
+    label_distance: float = 1.25,
+    label_rot: Union[Literal["default", "best"], float] = "best",
     show_edges: bool = True,
     key_added: Optional[str] = None,
     figsize: Optional[Tuple[float, float]] = None,
@@ -124,7 +117,7 @@ def circular_projection(
             - `None` - it will determined automatically, based on the number of lineages.
             - `'optimal'` - order lineages optimally by solving the Travelling salesman problem (TSP).
               Recommended for <= `20` lineages.
-            - `'default'` - use the order as specified in ``lineages``.
+            - `'default'` - use the order as specified by ``lineages``.
 
     metric
         Metric to use when constructing pairwise distance matrix when ``lineage_order = 'optimal'``. For available
@@ -139,9 +132,9 @@ def circular_projection(
         Whether to access :attr:`anndata.AnnData.raw` when there are ``keys`` in :attr:`anndata.AnnData.var_names`.
     text_kwargs
         Keyword arguments for :func:`matplotlib.pyplot.text`.
-    labeldistance
+    label_distance
         Distance at which the lineage labels will be drawn.
-    labelrot
+    label_rot
         How to rotate the labels. Valid options are:
 
             - `'best'` - rotate labels so that they are easily readable.
@@ -167,13 +160,14 @@ def 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:
-        raise ValueError(f"Expected `delta` to be positive, found `{labeldistance}`.")
+    if label_distance is not None and label_distance < 0:
+        raise ValueError(
+            f"Expected `label_distance` to be positive, found `{label_distance}`."
+        )
 
-    if labelrot is None:
-        labelrot = LabelRot.DEFAULT
-    if isinstance(labelrot, str):
-        labelrot = LabelRot(labelrot)
+    if label_rot is None:
+        label_rot = LabelRot.DEFAULT
+    label_rot = LabelRot(label_rot)
 
     suffix = "bwd" if backward else "fwd"
     if key_added is None:
@@ -188,25 +182,25 @@ def circular_projection(
     ) + _check_collection(
         adata, keys, "var_names", key_name="Gene", raise_exc=False, use_raw=use_raw
     )
-    haystack = {s.s for s in PrimingDegree}
+    haystack = set(PrimingDegree)
     keys = keys_ + [k for k in keys if k in haystack]
     keys = _unique_order_preserving(keys)
 
     if not len(keys):
         raise ValueError("No valid keys have been selected.")
 
-    lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
+    lineage_key = Key.obsm.abs_probs(backward)
     if lineage_key not in adata.obsm:
         raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")
 
-    probs = adata.obsm[lineage_key]
+    probs: Lineage = adata.obsm[lineage_key]
 
     if isinstance(lineages, str):
         lineages = (lineages,)
     elif lineages is None:
         lineages = probs.names
 
-    probs: Lineage = adata.obsm[lineage_key][lineages]
+    probs = adata.obsm[lineage_key][lineages]
     n_lin = probs.shape[1]
     if n_lin < 3:
         raise ValueError(f"Expected at least `3` lineages, found `{n_lin}`.")
@@ -292,7 +286,7 @@ def circular_projection(
 
         patches, texts = ax.pie(
             np.ones_like(angle_vec),
-            labeldistance=labeldistance,
+            labeldistance=label_distance,
             rotatelabels=True,
             labels=probs.names[::-1],
             startangle=-360 / len(angle_vec) / 2,
@@ -305,14 +299,14 @@ def circular_projection(
 
         # clockwise
         for color, text in zip(probs.colors[::-1], texts):
-            if isinstance(labelrot, (int, float)):
-                text.set_rotation(labelrot)
-            elif labelrot == LabelRot.BEST:
+            if isinstance(label_rot, (int, float)):
+                text.set_rotation(label_rot)
+            elif label_rot == LabelRot.BEST:
                 rot = text.get_rotation()
                 text.set_rotation(rot + 90 + (1 - rot // 180) * 180)
-            elif labelrot != LabelRot.DEFAULT:
+            elif label_rot != LabelRot.DEFAULT:
                 raise NotImplementedError(
-                    f"Label rotation `{labelrot}` is not yet implemented."
+                    f"Label rotation `{label_rot}` is not yet implemented."
                 )
             text.set_color(color)
 
diff --git a/cellrank/pl/_cluster_fates.py b/cellrank/pl/_cluster_fates.py
index 7d5078bae..9561acbfa 100644
--- a/cellrank/pl/_cluster_fates.py
+++ b/cellrank/pl/_cluster_fates.py
@@ -1,20 +1,22 @@
-"""Cluster fatess and similarity module."""
-
-from typing import Any, Tuple, Union, Mapping, TypeVar, Optional, Sequence
+from typing import Any, Tuple, Union, Mapping, Optional, Sequence
+from typing_extensions import Literal
 
+from enum import auto
 from math import ceil
 from types import MappingProxyType
 from pathlib import Path
 from collections import OrderedDict as odict
 
+from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
 from scanpy.plotting import violin
 from scvelo.plotting import paga
+from cellrank.tl._enum import ModeEnum
 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
@@ -26,23 +28,23 @@
 from seaborn import heatmap, clustermap
 from matplotlib import cm
 
-AnnData = TypeVar("AnnData")
-
 
-class ClusterFatesMode(ModeEnum):  # noqa
-    BAR = "bar"
-    PAGA = "paga"
-    PAGA_PIE = "paga_pie"
-    VIOLIN = "violin"
-    HEATMAP = "heatmap"
-    CLUSTERMAP = "clustermap"
+class ClusterFatesMode(ModeEnum):  # noqa: D101
+    BAR = auto()
+    PAGA = auto()
+    PAGA_PIE = auto()
+    VIOLIN = auto()
+    HEATMAP = auto()
+    CLUSTERMAP = auto()
 
 
 @d.dedent
 @inject_docs(m=ClusterFatesMode)
 def cluster_fates(
     adata: AnnData,
-    mode: str = ClusterFatesMode.PAGA_PIE.s,
+    mode: Literal[
+        "bar", "paga", "paga_pie", "violin", "heatmap", "clustermap"
+    ] = ClusterFatesMode.PAGA_PIE,
     backward: bool = False,
     lineages: Optional[Union[str, Sequence[str]]] = None,
     cluster_key: Optional[str] = "clusters",
@@ -64,7 +66,7 @@ def cluster_fates(
 
     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 :cite:`wolf:19`.
+    For mode `{m.PAGA!r}` and `{m.PAGA_PIE!r}`, we use *PAGA*, see :cite:`wolf:19`.
 
     Parameters
     ----------
@@ -72,12 +74,12 @@ def cluster_fates(
     mode
         Type of plot to show. Valid options are:
 
-            - `{m.BAR.s!r}` - barplot, one panel per cluster. The whiskers correspond to the standard error of the mean.
-            - `{m.PAGA.s!r}` - scanpy's PAGA, one per %(initial_or_terminal)s state, colored in by fate.
-            - `{m.PAGA_PIE.s!r}` - scanpy's PAGA with pie charts indicating aggregated fates.
-            - `{m.VIOLIN.s!r}` - violin plots, one per %(initial_or_terminal)s state.
-            - `{m.HEATMAP.s!r}` - a heatmap, showing average fates per cluster.
-            - `{m.CLUSTERMAP.s!r}` - same as a heatmap, but with a dendrogram.
+            - `{m.BAR!r}` - barplot, one panel per cluster. The whiskers correspond to the standard error of the mean.
+            - `{m.PAGA!r}` - scanpy's PAGA, one per %(initial_or_terminal)s state, colored in by fate.
+            - `{m.PAGA_PIE!r}` - scanpy's PAGA with pie charts indicating aggregated fates.
+            - `{m.VIOLIN!r}` - violin plots, one per %(initial_or_terminal)s state.
+            - `{m.HEATMAP!r}` - a heatmap, showing average fates per cluster.
+            - `{m.CLUSTERMAP!r}` - same as a heatmap, but with a dendrogram.
     %(backward)s
     lineages
         Lineages for which to visualize absorption probabilities. If `None`, use all lineages.
@@ -86,27 +88,27 @@ def cluster_fates(
     clusters
         Clusters to visualize. If `None`, all clusters will be plotted.
     basis
-        Basis for scatterplot to use when ``mode={m.PAGA_PIE.s!r}``. If `None`, don't show the scatterplot.
+        Basis for scatterplot to use when ``mode = {m.PAGA_PIE!r}``. If `None`, don't show the scatterplot.
     cbar
-        Whether to show colorbar when ``mode={m.PAGA_PIE.s!r}``.
+        Whether to show colorbar when ``mode = {m.PAGA_PIE!r}``.
     ncols
-        Number of columns when ``mode={m.BAR.s!r}`` or ``mode={m.PAGA.s!r}``.
+        Number of columns when ``mode = {m.BAR!r}`` or ``mode = {m.PAGA!r}``.
     sharey
-        Whether to share y-axis when ``mode={m.BAR.s!r}``.
+        Whether to share y-axis when ``mode = {m.BAR!r}``.
     fmt
-        Format when using ``mode={m.HEATMAP.s!r}`` or ``mode={m.CLUSTERMAP.s!r}``.
+        Format when using ``mode = {m.HEATMAP!r}`` or ``mode = {m.CLUSTERMAP!r}``.
     xrot
         Rotation of the labels on the x-axis.
     figsize
         Size of the figure.
     legend_kwargs
         Keyword arguments for :func:`matplotlib.axes.Axes.legend`, such as `'loc'` for legend position.
-        For ``mode={m.PAGA_PIE.s!r}`` and ``basis='...'``, this controls the placement of the
+        For ``mode = {m.PAGA_PIE!r}`` and ``basis = '...'``, this controls the placement of the
         absorption probabilities legend.
     %(plotting)s
     kwargs
         Keyword arguments for :func:`scvelo.pl.paga`, :func:`scanpy.pl.violin` or :func:`matplotlib.pyplot.bar`,
-        depending on the value of ``mode``.
+        depending on the ``mode``.
 
     Returns
     -------
@@ -129,7 +131,7 @@ def _():
         gs = plt.GridSpec(n_rows, cols, figure=fig, wspace=0.5, hspace=0.5)
 
         ax = None
-        colors = list(adata.obsm[lk][:, lin_names].colors)
+        colors = list(adata.obsm[lineage_key][:, lin_names].colors)
 
         for g, k in zip(gs, d.keys()):
             current_ax = fig.add_subplot(g, sharey=ax)
@@ -148,7 +150,7 @@ def _():
             current_ax.set_xticks(np.arange(len(lin_names)))
             current_ax.set_xticklabels(lin_names, rotation=xrot)
             if not is_all:
-                current_ax.set_xlabel(points)
+                current_ax.set_xlabel(term_states)
             current_ax.set_ylabel("absorption probability")
             current_ax.set_title(k)
 
@@ -185,7 +187,7 @@ def _():
             colors = [v[0][i] for v in d.values()]
             kwargs["ax"] = ax
             kwargs["colors"] = tuple(colors)
-            kwargs["title"] = f"{dir_prefix} {lineage_name}"
+            kwargs["title"] = f"{direction} {lineage_name}"
 
             vmin = np.min(colors + [vmin])
             vmax = np.max(colors + [vmax])
@@ -210,7 +212,7 @@ def _():
 
     @plot.register(ClusterFatesMode.PAGA_PIE)
     def _():
-        colors = list(adata.obsm[lk][:, lin_names].colors)
+        colors = list(adata.obsm[lineage_key][:, lin_names].colors)
         colors = {i: odict(zip(colors, mean)) for i, (mean, _) in enumerate(d.values())}
 
         fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
@@ -263,7 +265,7 @@ def _():
             handles = []
             for lineage_name, color in zip(lin_names, colors[0].keys()):
                 handles += [ax.scatter([], [], label=lineage_name, c=color)]
-            if len(colors[0].keys()) != len(adata.obsm[lk].names):
+            if len(colors[0].keys()) != len(adata.obsm[lineage_key].names):
                 handles += [ax.scatter([], [], label="Rest", c="grey")]
 
             second_legend = _position_legend(
@@ -271,7 +273,7 @@ def _():
                 legend_loc=legend_kwargs["loc"],
                 handles=handles,
                 **{k: v for k, v in legend_kwargs.items() if k != "loc"},
-                title=points,
+                title=term_states,
             )
             fig.add_artist(second_legend)
 
@@ -307,8 +309,8 @@ def _():
         with RandomKeys(adata, len(lin_names), where="obs") as keys:
             _i = 0
             for _i, (name, key, ax) in enumerate(zip(lin_names, keys, axes)):
-                adata.obs[key] = adata.obsm[lk][name].X
-                ax.set_title(f"{dir_prefix} {name}")
+                adata.obs[key] = adata.obsm[lineage_key][name].X
+                ax.set_title(f"{direction} {name}")
                 violin(
                     adata, ylabel="absorption probability", keys=key, ax=ax, **kwargs
                 )
@@ -318,41 +320,39 @@ def _():
         return fig
 
     def plot_violin_no_cluster_key():
-        from anndata import AnnData as _AnnData
-
         kwargs.pop("ax", None)
         kwargs.pop("keys", None)  # don't care
         kwargs.pop("save", None)
 
         kwargs["show"] = False
-        kwargs["groupby"] = points
+        kwargs["groupby"] = term_states
         kwargs["xlabel"] = None
         kwargs["rotation"] = xrot
 
-        data = np.ravel(adata.obsm[lk].X.T)[..., np.newaxis]
-        tmp = _AnnData(csr_matrix(data.shape, dtype=np.float32))
+        data = np.ravel(adata.obsm[lineage_key].X.T)[..., np.newaxis]
+        tmp = AnnData(csr_matrix(data.shape, dtype=np.float32))
         tmp.obs["absorption probability"] = data
-        tmp.obs[points] = (
+        tmp.obs[term_states] = (
             pd.Series(
                 np.concatenate(
                     [
-                        [f"{dir_prefix.lower()} {n}"] * adata.n_obs
-                        for n in adata.obsm[lk].names
+                        [f"{direction.lower()} {n}"] * adata.n_obs
+                        for n in adata.obsm[lineage_key].names
                     ]
                 )
             )
             .astype("category")
             .values
         )
-        tmp.obs[points] = tmp.obs[points].cat.reorder_categories(
-            [f"{dir_prefix.lower()} {n}" for n in adata.obsm[lk].names]
+        tmp.obs[term_states] = tmp.obs[term_states].cat.reorder_categories(
+            [f"{direction.lower()} {n}" for n in adata.obsm[lineage_key].names]
         )
-        tmp.uns[f"{points}_colors"] = adata.obsm[lk].colors
+        tmp.uns[f"{term_states}_colors"] = adata.obsm[lineage_key].colors
 
         fig, ax = plt.subplots(
             figsize=figsize if figsize is not None else (8, 6), dpi=dpi
         )
-        ax.set_title(points.capitalize())
+        ax.set_title(term_states.capitalize())
 
         violin(tmp, keys=["absorption probability"], ax=ax, **kwargs)
 
@@ -383,7 +383,7 @@ def _():
                 vmin=vmin,
                 vmax=vmax,
                 fmt=fmt,
-                row_colors=adata.obsm[lk][lin_names].colors,
+                row_colors=adata.obsm[lineage_key][lin_names].colors,
                 dendrogram_ratio=(
                     0.15 * data.shape[0] / max_size,
                     0.15 * data.shape[1] / max_size,
@@ -431,14 +431,9 @@ def _():
             f"`{ClusterFatesMode.BAR!r}` and `{ClusterFatesMode.VIOLIN!r}`, found `mode={mode!r}`."
         )
 
-    if backward:
-        lk = AbsProbKey.BACKWARD.s
-        points = TerminalStatesPlot.BACKWARD.s
-        dir_prefix = DirPrefix.BACKWARD.s
-    else:
-        lk = AbsProbKey.FORWARD.s
-        points = TerminalStatesPlot.FORWARD.s
-        dir_prefix = DirPrefix.FORWARD.s
+    lineage_key = Key.obsm.abs_probs(backward)
+    term_states = Key.obs.term_states(backward)
+    direction = Key.where(backward)
 
     if cluster_key is not None:
         is_all = False
@@ -461,10 +456,10 @@ def _():
             clusters = list(adata.obs[cluster_key].cat.categories)
     else:
         is_all = True
-        clusters = [points]
+        clusters = [term_states]
 
-    if lk not in adata.obsm:
-        raise KeyError(f"Lineage key `{lk!r}` not found in `adata.obsm`.")
+    if lineage_key not in adata.obsm:
+        raise KeyError(f"Lineage key `{lineage_key!r}` not found in `adata.obsm`.")
 
     if lineages is not None:
         if isinstance(lineages, str):
@@ -472,8 +467,8 @@ def _():
         lin_names = _unique_order_preserving(lineages)
     else:
         # must be list for `sc.pl.violin`, else cats str
-        lin_names = list(adata.obsm[lk].names)
-    _ = adata.obsm[lk][lin_names]
+        lin_names = list(adata.obsm[lineage_key].names)
+    _ = adata.obsm[lineage_key][lin_names]
 
     if mode == mode.VIOLIN and not is_all:
         adata = adata[np.isin(adata.obs[cluster_key], clusters)].copy()
@@ -486,7 +481,7 @@ def _():
             else (adata.obs[cluster_key] == name).values
         )
         mask = np.array(mask, dtype=bool)
-        data = adata.obsm[lk][mask, lin_names].X
+        data = adata.obsm[lineage_key][mask, lin_names].X
         mean = np.nanmean(data, axis=0)
         std = np.nanstd(data, axis=0) / np.sqrt(data.shape[0])
         d[name] = [mean, std]
diff --git a/cellrank/pl/_cluster_lineage.py b/cellrank/pl/_cluster_lineage.py
index 5728d0851..cf892a9c7 100644
--- a/cellrank/pl/_cluster_lineage.py
+++ b/cellrank/pl/_cluster_lineage.py
@@ -1,6 +1,4 @@
-"""Cluster lineages module."""
-
-from typing import Dict, Tuple, Union, Optional, Sequence
+from typing import Any, Dict, Tuple, Union, Optional, Sequence
 
 from types import MappingProxyType
 from pathlib import Path
@@ -8,6 +6,8 @@
 import scanpy as sc
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.tl._enum import _DEFAULT_BACKEND, Backend_t
 from cellrank.ul._docs import d
 from cellrank.pl._utils import (
     _fit_bulk,
@@ -22,7 +22,6 @@
 )
 from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul._utils import _genesymbols, _get_n_cores, _check_collection
-from cellrank.tl._constants import _DEFAULT_BACKEND, AbsProbKey
 
 import numpy as np
 from sklearn.preprocessing import StandardScaler
@@ -56,7 +55,7 @@ def cluster_lineage(
     random_state: Optional[int] = None,
     show_progress_bar: bool = True,
     n_jobs: Optional[int] = 1,
-    backend: str = _DEFAULT_BACKEND,
+    backend: Backend_t = _DEFAULT_BACKEND,
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[int] = None,
     save: Optional[Union[str, Path]] = None,
@@ -64,7 +63,7 @@ def cluster_lineage(
     neighbors_kwargs: Dict = MappingProxyType({"use_rep": "X"}),
     clustering_kwargs: Dict = MappingProxyType({}),
     return_models: bool = False,
-    **kwargs,
+    **kwargs: Any,
 ) -> Optional[_return_model_type]:
     """
     Cluster gene expression trends within a lineage and plot the clusters.
@@ -88,9 +87,9 @@ def cluster_lineage(
     n_points
         Number of points used for prediction.
     time_key
-        Key in ``adata.obs`` where the pseudotime is stored.
+        Key in :attr:`anndata.AnnData.obs` where the pseudotime is stored.
     covariate_key
-        Key(s) in ``adata.obs`` containing observations to be plotted at the bottom of each plot.
+        Key(s) in :attr:`anndata.AnnData.obs` containing observations to be plotted at the bottom of each plot.
     %(gene_symbols)s
     ratio
         Height ratio of each covariate in ``covariate_key``.
@@ -106,7 +105,8 @@ def cluster_lineage(
     sharey
         Whether to share y-axis across multiple plots.
     key
-        Key in ``adata.uns`` where to save the results. If `None`, it will be saved as ``lineage_{lineage}_trend`` .
+        Key in :attr:`anndata.AnnData.uns` where to save the results.
+        If `None`, it will be saved as ``'lineage_{lineage}_trend'`` .
     random_state
         Random seed for reproducibility.
     %(parallel)s
@@ -131,13 +131,15 @@ def cluster_lineage(
               shape `(n_genes, n_points)` containing the clustered genes.
     """
 
-    def plot_cluster(row, col, cluster, sharey_ax: Optional[str] = None):
+    def plot_cluster(
+        row: int, col: int, cluster: str, sharey_ax: Optional[str] = None
+    ) -> Optional[plt.Axes]:
         gss = GridSpecFromSubplotSpec(
             row_delta,
             1,
             subplot_spec=gs[row : row + row_delta, col],
             hspace=0,
-            height_ratios=[1] + [ratio] * (row_delta - 1),
+            height_ratios=[1.0] + [ratio] * (row_delta - 1),
         )
         ax = fig.add_subplot(gss[0, 0], sharey=sharey_ax)
 
@@ -180,7 +182,7 @@ def plot_cluster(row, col, cluster, sharey_ax: Optional[str] = None):
         return ax if sharey else None
 
     use_raw = kwargs.get("use_raw", False)
-    lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
+    lineage_key = Key.obsm.abs_probs(backward)
     if lineage_key not in adata.obsm:
         raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")
 
diff --git a/cellrank/pl/_gene_trend.py b/cellrank/pl/_gene_trend.py
index 361cf8d14..28f020bbb 100644
--- a/cellrank/pl/_gene_trend.py
+++ b/cellrank/pl/_gene_trend.py
@@ -1,4 +1,3 @@
-"""Gene trend module."""
 from typing import Any, List, Tuple, Union, Mapping, Optional, Sequence
 
 from types import MappingProxyType
@@ -6,6 +5,8 @@
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.tl._enum import _DEFAULT_BACKEND, Backend_t
 from cellrank.ul._docs import d
 from cellrank.pl._utils import (
     _fit_bulk,
@@ -20,7 +21,6 @@
 )
 from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul._utils import _genesymbols, _get_n_cores, _check_collection
-from cellrank.tl._constants import _DEFAULT_BACKEND, AbsProbKey
 
 import numpy as np
 import pandas as pd
@@ -65,7 +65,7 @@ def gene_trends(
     suptitle: Optional[str] = None,
     return_models: bool = False,
     n_jobs: Optional[int] = 1,
-    backend: str = _DEFAULT_BACKEND,
+    backend: Backend_t = _DEFAULT_BACKEND,
     show_progress_bar: bool = True,
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[int] = None,
@@ -89,16 +89,17 @@ def gene_trends(
         Names of the lineages to plot. If `None`, plot all lineages.
     %(backward)s
     data_key
-        Key in ``adata.layers`` or `'X'` for ``adata.X`` where the data is stored.
+        Key in :attr:`anndata.AnnData.layers` or `'X'` for :attr:`anndata.AnnData.X` where the data is stored.
     time_key
-        Key in ``adata.obs`` where the pseudotime is stored.
+        Key in :attr:`anndata.AnnData.obs` where the pseudotime is stored.
     %(time_range)s
 
         This can also be specified on per-lineage basis.
     %(gene_symbols)s
     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.
+        If ``same_plot = True``, group the trends by ``lineages`` instead of ``genes``.
+        This forces ``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 ``model`` is :class:`cellrank.ul.models.GAMR`,
@@ -112,10 +113,10 @@ def gene_trends(
         This can improve visualization. Can be specified individually for each lineage.
     lineage_cmap
         Categorical colormap to use when coloring in the lineages. If `None` and ``same_plot``,
-        use the corresponding colors in ``adata.uns``, otherwise use `'black'`.
+        use the corresponding colors in :attr:`anndata.AnnData.uns`, otherwise use `'black'`.
     abs_prob_cmap
         Continuous colormap to use when visualizing the absorption probabilities for each lineage.
-        Only used when ``same_plot=False``.
+        Only used when ``same_plot = False``.
     cell_color
         Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var_names` used for coloring the cells.
     cell_alpha
@@ -127,7 +128,8 @@ def gene_trends(
     lw
         Line width of the smoothed values.
     cbar
-        Whether to show colorbar. Always shown when percentiles for lineages differ. Only used when ``same_plot=False``.
+        Whether to show colorbar. Always shown when percentiles for lineages differ.
+        Only used when ``same_plot = False``.
     margins
         Margins around the plot.
     sharex
@@ -137,11 +139,11 @@ def gene_trends(
     gene_as_title
         Whether to show gene names as titles instead on y-axis.
     legend_loc
-        Location of the legend displaying lineages. Only used when `same_plot=True`.
+        Location of the legend displaying lineages. Only used when `same_plot = True`.
     obs_legend_loc
         Location of the legend when ``cell_color`` corresponds to a categorical variable.
     ncols
-        Number of columns of the plot when plotting multiple genes. Only used when ``same_plot=True``.
+        Number of columns of the plot when plotting multiple genes. Only used when ``same_plot = True``.
     suptitle
         Suptitle of the figure.
     %(return_models)s
@@ -167,12 +169,12 @@ def gene_trends(
         use_raw=kwargs.get("use_raw", False),
     )
 
-    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`.")
+    lineage_key = Key.obsm.abs_probs(backward)
+    if lineage_key not in adata.obsm:
+        raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")
 
     if lineages is None:
-        lineages = adata.obsm[ln_key].names
+        lineages = adata.obsm[lineage_key].names
     elif isinstance(lineages, str):
         lineages = [lineages]
     elif all(ln is None for ln in lineages):  # no lineage, all the weights are 1
@@ -211,7 +213,7 @@ def gene_trends(
     )
 
     lineages = sorted(lineages)
-    tmp = adata.obsm[ln_key][lineages].colors
+    tmp = adata.obsm[lineage_key][lineages].colors
     if lineage_cmap is None and not transpose:
         lineage_cmap = tmp
 
@@ -279,7 +281,7 @@ def gene_trends(
                 and plot_kwargs.get("lineage_probability", False)
                 and transpose
             ):
-                lpc = adata.obsm[ln_key][gene].colors[0]
+                lpc = adata.obsm[lineage_key][gene].colors[0]
             else:
                 lpc = None
 
diff --git a/cellrank/pl/_graph.py b/cellrank/pl/_graph.py
index a3477206d..31a3f2b65 100644
--- a/cellrank/pl/_graph.py
+++ b/cellrank/pl/_graph.py
@@ -1,6 +1,5 @@
-"""Graph plotting module."""
-
-from typing import Dict, Tuple, Union, Callable, Optional, Sequence
+from typing import Any, Dict, Tuple, Union, Callable, Optional, Sequence
+from typing_extensions import Literal
 
 from copy import deepcopy
 from types import MappingProxyType
@@ -8,10 +7,10 @@
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
 from cellrank.ul._docs import d
-from cellrank.tl._utils import save_fig
+from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul._utils import _read_graph_data
-from cellrank.tl._constants import _colors
 
 import numpy as np
 import pandas as pd
@@ -32,8 +31,10 @@ def graph(
     data: Union[AnnData, np.ndarray, spmatrix],
     graph_key: Optional[str] = None,
     ixs: Optional[Union[range, np.array]] = None,
-    layout: Union[str, Dict, Callable] = "umap",
-    keys: Sequence[str] = ("incoming",),
+    layout: Union[str, Dict[str, Any], Callable[..., np.ndarray]] = "umap",
+    keys: Sequence[
+        Union[str, Literal["incoming", "outgoing", "self_loops"]]
+    ] = "incoming",
     keylocs: Union[str, Sequence[str]] = "uns",
     node_size: float = 400,
     labels: Optional[Union[Sequence[str], Sequence[Sequence[str]]]] = None,
@@ -59,7 +60,7 @@ def graph(
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[int] = None,
     save: Optional[Union[str, Path]] = None,
-    layout_kwargs: Dict = MappingProxyType({}),
+    layout_kwargs: Dict[str, Any] = MappingProxyType({}),
 ) -> None:
     """
     Plot a graph, visualizing incoming and outgoing edges or self-transitions.
@@ -74,22 +75,23 @@ def graph(
     data
         The graph data to be plotted.
     graph_key
-        Key in ``adata.obsp`` or ``adata.uns`` where the graph is stored. Only used
-        when ``data`` is :class:`~anndata.Anndata` object.
+        Key in :attr:`anndata.AnnData.obsp` where the graph is stored.
+        Only used when ``data`` is :class:`anndata.Anndata` object.
     ixs
         Subset of indices of the graph to visualize.
     layout
         Layout to use for graph drawing.
 
-        - If :class:`str`, search for embedding in ``adata.obsm['X_{layout}']``.
+        - If :class:`str`, search for embedding in :attr:`anndata.AnnData.obsm` ``['X_{layout}']``.
           Use ``layout_kwargs={'components': [0, 1]}`` to select components.
-        - If :class:`dict`, keys should be values in interval ``[0, len(ixs))``
+        - If :class:`dict`, keys should be values in interval `[0, len(ixs))`
           and values `(x, y)` pairs corresponding to node positions.
     keys
-        Keys in ``adata.obs``, ``adata.obsm`` or ``adata.obsp`` to color the nodes.
+        Keys in :attr:`anndata.AnnData.obs`, :attr:`anndata.AnnData.obsm` or :attr:`anndata.AnnData.obsm`
+        used to color the nodes.
 
-        - If `'incoming'`, `'outgoing'` or `'self_loops'`, visualize reduction (see ``edge_reductions``)
-          for each node based on incoming or outgoing edges, respectively.
+        If `'incoming'`, `'outgoing'` or `'self_loops'`, visualize reduction (see ``edge_reductions``)
+        for each node based on incoming or outgoing edges, respectively.
     keylocs
         Locations of ``keys``. Can be any attribute of ``data`` if it's :class:`anndata.AnnData` object.
     node_size
@@ -261,8 +263,9 @@ def remove_low_weight_edges():
         )
         logg.debug(f"Setting self loop radius fraction to `{self_loop_radius_frac}`")
 
-    if not isinstance(keys, (tuple, list)):
+    if isinstance(keys, str):
         keys = [keys]
+    keys = _unique_order_preserving(keys)
 
     if not isinstance(keylocs, (tuple, list)):
         keylocs = [keylocs] * len(keys)
@@ -477,7 +480,7 @@ def remove_low_weight_edges():
             if keyloc in ("obs", "obsm"):
                 values = values[ixs]
             categories = values.cat.categories
-            color_key = _colors(key)
+            color_key = Key.uns.colors(key)
             if color_key in data.uns:
                 mapper = dict(zip(categories, data.uns[color_key]))
             else:
diff --git a/cellrank/pl/_heatmap.py b/cellrank/pl/_heatmap.py
index 3617eb153..31d0cc0fd 100644
--- a/cellrank/pl/_heatmap.py
+++ b/cellrank/pl/_heatmap.py
@@ -1,14 +1,16 @@
-"""Heatmap module."""
-
-from typing import Any, Dict, List, Tuple, Union, TypeVar, Optional, Sequence
+from typing import Any, Dict, List, Tuple, Union, Optional, Sequence
+from typing_extensions import Literal
 
 import os
+from enum import auto
 from math import fabs
 from pathlib import Path
 from collections import Iterable, defaultdict
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum, Backend_t
 from cellrank.ul._docs import d, inject_docs
 from cellrank.pl._utils import (
     _fit_bulk,
@@ -28,7 +30,6 @@
     valuedispatch,
     _check_collection,
 )
-from cellrank.tl._constants import _DEFAULT_BACKEND, ModeEnum, AbsProbKey
 
 import numpy as np
 import pandas as pd
@@ -44,15 +45,10 @@
 
 _N_XTICKS = 10
 
-Cmap = TypeVar("Cmap")
-Norm = TypeVar("Norm")
-Ax = TypeVar("Ax")
-Fig = TypeVar("Fig")
-
 
-class HeatmapMode(ModeEnum):  # noqa
-    GENES = "genes"
-    LINEAGES = "lineages"
+class HeatmapMode(ModeEnum):  # noqa: D101
+    GENES = auto()
+    LINEAGES = auto()
 
 
 @d.dedent
@@ -64,7 +60,7 @@ def heatmap(
     genes: Sequence[str],
     lineages: Optional[Union[str, Sequence[str]]] = None,
     backward: bool = False,
-    mode: str = HeatmapMode.LINEAGES.s,
+    mode: Literal["genes", "lineages"] = HeatmapMode.LINEAGES,
     time_key: str = "latent_time",
     time_range: Optional[Union[_time_range_type, List[_time_range_type]]] = None,
     callback: _callback_type = None,
@@ -84,7 +80,7 @@ def heatmap(
     return_genes: bool = False,
     return_models: bool = False,
     n_jobs: Optional[int] = 1,
-    backend: str = _DEFAULT_BACKEND,
+    backend: Backend_t = _DEFAULT_BACKEND,
     show_progress_bar: bool = True,
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[int] = None,
@@ -107,26 +103,26 @@ def heatmap(
     mode
         Valid options are:
 
-            - `{m.LINEAGES.s!r}` - group by ``genes`` for each lineage in ``lineages``.
-            - `{m.GENES.s!r}` - group by ``lineages`` for each gene in ``genes``.
+            - `{m.LINEAGES!r}` - group by ``genes`` for each lineage in ``lineages``.
+            - `{m.GENES!r}` - group by ``lineages`` for each gene in ``genes``.
     time_key
-        Key in ``adata.obs`` where the pseudotime is stored.
+        Key in attr:`anndata.AnnData.obs` where the pseudotime is stored.
     %(time_range)s
 
         This can also be specified on per-lineage basis.
     %(gene_symbols)s
     %(model_callback)s
     cluster_key
-        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}``.
+        Key(s) in :attr:`anndata.AnnData.obs` containing categorical observations to be plotted on the top of heatmap.
+        Only available when ``mode = {m.LINEAGES!r}``.
     show_absorption_probabilities
         Whether to also plot absorption probabilities alongside the smoothed expression.
-        Only available when ``mode={m.LINEAGES.s!r}``.
+        Only available when ``mode = {m.LINEAGES!r}``.
     cluster_genes
-        Whether to cluster genes using :func:`seaborn.clustermap` when ``mode='lineages'``.
+        Whether to cluster genes using :func:`seaborn.clustermap` when ``mode = 'lineages'``.
     keep_gene_order
         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}``.
+        Only available when ``cluster_genes = False`` and ``mode = {m.LINEAGES!r}``.
     scale
         Whether to normalize the gene expression `[0, 1]` range.
     n_convolve
@@ -136,7 +132,7 @@ def heatmap(
     cbar
         Whether to show the colorbar.
     lineage_height
-        Height of a bar when ``mode={m.GENES.s!r}``.
+        Height of a bar when ``mode = {m.GENES!r}``.
     fontsize
         Size of the title's font.
     xlabel
@@ -144,9 +140,9 @@ def heatmap(
     cmap
         Colormap to use when visualizing the smoothed expression.
     dendrogram
-        Whether to show dendrogram when ``cluster_genes=True``.
+        Whether to show dendrogram when ``cluster_genes = True``.
     return_genes
-        Whether to return the sorted or clustered genes. Only available when ``mode={m.LINEAGES.s!r}``.
+        Whether to return the sorted or clustered genes. Only available when ``mode = {m.LINEAGES!r}``.
     %(return_models)s
     %(parallel)s
     %(plotting)s
@@ -157,7 +153,7 @@ def heatmap(
     -------
     %(plots_or_returns_models)s
     :class:`pandas.DataFrame`
-        If ``return_genes=True`` and ``mode={m.LINEAGES.s!r}``, returns :class:`pandas.DataFrame`
+        If ``return_genes = True`` and ``mode = {m.LINEAGES!r}``, returns :class:`pandas.DataFrame`
         containing the clustered or sorted genes.
     """
 
@@ -168,8 +164,8 @@ def find_nearest(array: np.ndarray, value: float) -> int:
                 ix == len(array)
                 or fabs(value - array[ix - 1]) < fabs(value - array[ix])
             ):
-                return ix - 1
-            return ix
+                return int(ix - 1)
+            return int(ix)
 
         series = series[np.argsort(series.values)]
 
@@ -187,7 +183,9 @@ def subset_lineage(lname: str, rng: np.ndarray) -> np.ndarray:
 
         return lin
 
-    def create_col_colors(lname: str, rng: np.ndarray) -> Tuple[np.ndarray, Cmap, Norm]:
+    def create_col_colors(
+        lname: str, rng: np.ndarray
+    ) -> Tuple[np.ndarray, mcolors.Colormap, mcolors.Normalize]:
         color = adata.obsm[lineage_key][lname].colors[0]
         lin = subset_lineage(lname, rng)
 
@@ -217,10 +215,10 @@ def create_col_categorical_color(cluster_key: str, rng: np.ndarray) -> np.ndarra
     def create_cbar(
         ax,
         x_delta: float,
-        cmap: Cmap,
-        norm: Norm,
+        cmap: mcolors.Colormap,
+        norm: mcolors.Normalize,
         label: Optional[str] = None,
-    ) -> Ax:
+    ) -> plt.Axes:
         cax = inset_axes(
             ax,
             width="1%",
@@ -241,11 +239,11 @@ def create_cbar(
         return cax
 
     @valuedispatch
-    def _plot_heatmap(mode: HeatmapMode) -> Fig:
+    def _plot_heatmap(mode: HeatmapMode) -> plt.Figure:
         raise NotImplementedError(mode.value)
 
     @_plot_heatmap.register(HeatmapMode.GENES)
-    def _() -> Tuple[Fig, None]:
+    def _() -> Tuple[plt.Figure, None]:
         def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None:
             colors = colors if cmap is None else cmap(colors)
 
@@ -363,7 +361,7 @@ def color_fill_rec(ax, xs, y1, y2, colors=None, cmap=cmap, **kwargs) -> None:
         return fig, None
 
     @_plot_heatmap.register(HeatmapMode.LINEAGES)
-    def _() -> Tuple[List[Fig], pd.DataFrame]:
+    def _() -> Tuple[List[plt.Figure], pd.DataFrame]:
         data_t = defaultdict(dict)  # transpose
         for gene, lns in data.items():
             for ln, y in lns.items():
@@ -511,7 +509,7 @@ def _() -> Tuple[List[Fig], pd.DataFrame]:
 
     mode = HeatmapMode(mode)
 
-    lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
+    lineage_key = Key.obsm.abs_probs(backward)
     if lineage_key not in adata.obsm:
         raise KeyError(f"Lineages key `{lineage_key!r}` not found in `adata.obsm`.")
 
@@ -554,7 +552,7 @@ def _() -> Tuple[List[Fig], pd.DataFrame]:
 
     xlabel = time_key if xlabel is None else xlabel
 
-    logg.debug(f"Plotting `{mode.s!r}` heatmap")
+    logg.debug(f"Plotting `{mode!r}` heatmap")
     fig, genes = _plot_heatmap(mode)
 
     if save is not None and fig is not None:
@@ -572,11 +570,11 @@ def _() -> Tuple[List[Fig], pd.DataFrame]:
         return all_models
 
 
-def _get_ax_bbox(fig: Fig, ax: Ax):
+def _get_ax_bbox(fig: plt.Figure, ax: plt.Axes):
     return ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
 
 
-def _set_ax_height_to_cm(fig: Fig, ax: Ax, height: float) -> None:
+def _set_ax_height_to_cm(fig: plt.Figure, ax: plt.Axes, height: float) -> None:
     from mpl_toolkits.axes_grid1 import Size, Divider
 
     height /= 2.54  # cm to inches
diff --git a/cellrank/pl/_init_term_states.py b/cellrank/pl/_init_term_states.py
index c0f3d2dcf..c29e0929d 100644
--- a/cellrank/pl/_init_term_states.py
+++ b/cellrank/pl/_init_term_states.py
@@ -1,14 +1,11 @@
-"""Module used for finding initial and terminal states."""
-from typing import Union, TypeVar, Optional, Sequence
+from typing import Union, Optional, Sequence
+from typing_extensions import Literal
 
+from anndata import AnnData
+from cellrank._key import Key
+from scanpy._utils import deprecated_arg_names
 from cellrank.ul._docs import d, _initial, _terminal, inject_docs
-from cellrank.tl._constants import TermStatesKey, TerminalStatesPlot
-from cellrank.tl.estimators import GPCCA
-from cellrank.tl.estimators._constants import P
-from cellrank.tl.kernels._precomputed_kernel import DummyKernel
-
-AnnData = TypeVar("AnnData")
-
+from cellrank.tl.estimators import GPCCA, CFLARE
 
 _find_docs = """\
 Plot {direction} states uncovered by :class:`cellrank.tl.{fn_name}`.
@@ -21,14 +18,14 @@
     Only available when {direction} were estimated by :class:`cellrank.tl.estimators.GPCCA`.
 states
     Subset of {direction} states to plot. If `None`, plot all {direction} states.
-cluster_key
+color
     If given, plot cluster annotations left of the {direction} states.
 %(time_mode)s
 time_key
-    Key in ``adata.obs`` where the pseudotime is stored.
+    Key in :attr:`anndata.AnnData.obs` where the pseudotime is stored.
 %(basis)s
 kwargs
-    Keyword arguments for :meth:`cellrank.tl.estimators.BaseEstimator.plot_terminal_states`.
+    Keyword arguments for :meth:`cellrank.tl.estimators.GPCCA.plot_terminal_states`.
 
 Returns
 -------
@@ -36,27 +33,29 @@
 """
 
 
+@deprecated_arg_names({"cluster_key": "color"})
 def _initial_terminal(
     adata: AnnData,
     backward: bool = False,
     discrete: bool = False,
     states: Optional[Union[str, Sequence[str]]] = None,
-    cluster_key: Optional[str] = None,
-    mode: str = "embedding",
+    color: Optional[str] = None,
+    mode: Literal["embedding", "time"] = "embedding",
     time_key: str = "latent_time",
     **kwargs,
 ) -> None:
+    mc = GPCCA.from_adata(adata, obsp_key=Key.uns.kernel(backward))
+    # in case there are no term. states memberships
+    if mc.terminal_states_memberships is None:
+        mc = CFLARE.from_adata(adata, obsp_key=Key.uns.kernel(backward))
 
-    pk = DummyKernel(adata=adata, backward=backward)
-    mc = GPCCA(pk, read_from_adata=True, write_to_adata=False)
-
-    if mc._get(P.TERM) is None:
+    if mc.terminal_states is None:
         raise RuntimeError(
             f"Compute {_initial if backward else _terminal} states first as "
-            f"`cellrank.tl.compute_{TermStatesKey.BACKWARD if backward else TermStatesKey.FORWARD}()`."
+            f"`cellrank.tl.{'initial' if backward else 'terminal'}_states()`."
         )
 
-    n_states = len(mc._get(P.TERM).cat.categories)
+    n_states = len(mc.terminal_states.cat.categories)
     if n_states == 1 or (
         states is not None and (isinstance(states, str) or len(states) == 1)
     ):
@@ -65,11 +64,7 @@ def _initial_terminal(
     if kwargs.get("title", None) is None:
         if discrete:
             if kwargs.get("same_plot", True):
-                kwargs["title"] = (
-                    TerminalStatesPlot.BACKWARD.s
-                    if backward
-                    else TerminalStatesPlot.FORWARD.s
-                )
+                kwargs["title"] = "initial states" if backward else "terminal states"
         elif (
             mode == "embedding"
             and kwargs.get("title", None) is None
@@ -81,17 +76,13 @@ def _initial_terminal(
                 )
             )
         ):
-            kwargs["title"] = (
-                TerminalStatesPlot.BACKWARD.s
-                if backward
-                else TerminalStatesPlot.FORWARD.s
-            )
+            kwargs["title"] = "initial states" if backward else "terminal states"
 
-    _ = kwargs.pop("lineages", None)
+    _ = kwargs.pop("states", None)
 
     mc.plot_terminal_states(
-        lineages=states,
-        cluster_key=cluster_key,
+        states=states,
+        color=color,
         mode=mode,
         time_key=time_key,
         discrete=discrete,
@@ -103,16 +94,16 @@ def _initial_terminal(
 @inject_docs(
     __doc__=_find_docs.format(
         direction=_initial,
-        fn_name=TermStatesKey.BACKWARD.s,
-        title=TerminalStatesPlot.BACKWARD.s,
+        fn_name="initial",
+        title="initial states",
     )
 )
 def initial_states(  # noqa: D103
     adata: AnnData,
     discrete: bool = False,
     states: Optional[Union[str, Sequence[str]]] = None,
-    cluster_key: Optional[str] = None,
-    mode: str = "embedding",
+    color: Optional[str] = None,
+    mode: Literal["embedding", "time"] = "embedding",
     time_key: str = "latent_time",
     **kwargs,
 ) -> Optional[AnnData]:
@@ -122,7 +113,7 @@ def initial_states(  # noqa: D103
         backward=True,
         discrete=discrete,
         states=states,
-        cluster_key=cluster_key,
+        color=color,
         mode=mode,
         time_key=time_key,
         **kwargs,
@@ -133,16 +124,16 @@ def initial_states(  # noqa: D103
 @inject_docs(
     __doc__=_find_docs.format(
         direction=_terminal,
-        fn_name=TermStatesKey.FORWARD.s,
-        title=TerminalStatesPlot.FORWARD.s,
+        fn_name="terminal",
+        title="terminal states",
     )
 )
 def terminal_states(  # noqa: D103
     adata: AnnData,
     discrete: bool = False,
     states: Optional[Union[str, Sequence[str]]] = None,
-    cluster_key: Optional[str] = None,
-    mode: str = "embedding",
+    color: Optional[str] = None,
+    mode: Literal["embedding", "time"] = "embedding",
     time_key: str = "latent_time",
     **kwargs,
 ) -> Optional[AnnData]:
@@ -152,7 +143,7 @@ def terminal_states(  # noqa: D103
         backward=False,
         discrete=discrete,
         states=states,
-        cluster_key=cluster_key,
+        color=color,
         mode=mode,
         time_key=time_key,
         **kwargs,
diff --git a/cellrank/pl/_lineages.py b/cellrank/pl/_lineages.py
index f6795882b..c4d7c432c 100644
--- a/cellrank/pl/_lineages.py
+++ b/cellrank/pl/_lineages.py
@@ -1,15 +1,11 @@
-"""Module for plotting lineage-related stuff."""
 from typing import Any, Union, Optional, Sequence
+from typing_extensions import Literal
 
 import cellrank.logging as logg
+from cellrank._key import Key
 from cellrank.ul._docs import d
 from cellrank.pl._utils import AnnData
-from cellrank.tl._constants import DirPrefix
-from cellrank.tl.estimators import GPCCA
-from cellrank.tl.estimators._constants import A, P
-from cellrank.tl.kernels._precomputed_kernel import DummyKernel
-
-import pandas as pd
+from cellrank.tl.estimators.terminal_states._cflare import CFLARE
 
 
 @d.dedent
@@ -17,8 +13,8 @@ def lineages(
     adata: AnnData,
     lineages: Optional[Union[str, Sequence[str]]] = None,
     backward: bool = False,
-    cluster_key: Optional[str] = None,
-    mode: str = "embedding",
+    color: Optional[str] = None,
+    mode: Literal["embedding", "time"] = "embedding",
     time_key: str = "latent_time",
     **kwargs: Any,
 ) -> None:
@@ -36,31 +32,32 @@ def lineages(
     lineages
         Plot only these lineages. If `None`, plot all lineages.
     %(backward)s
-    cluster_key
+    color
         If given, plot cluster annotations left of the lineage probabilities.
     %(time_mode)s
     time_key
         Key in ``adata.obs`` where the pseudotime is stored.
     %(basis)s
     kwargs
-        Keyword arguments for :meth:`cellrank.tl.estimators.BaseEstimator.plot_absorption_probabilities`.
+        Keyword arguments for :meth:`cellrank.tl.estimators.GPCCA.plot_absorption_probabilities`.
 
     Returns
     -------
     %(just_plots)s
     """
 
-    pk = DummyKernel(adata, backward=backward)
-    mc = GPCCA(pk, read_from_adata=True, write_to_adata=False)
-    if mc._get(P.ABS_PROBS) is None:
+    # use CFLARE (no macrostates required)
+    mc = CFLARE.from_adata(adata, obsp_key=Key.uns.kernel(backward))
+    if mc.absorption_probabilities is None:
         raise RuntimeError(
             f"Compute absorption probabilities first as `cellrank.tl.lineages(..., backward={backward})`."
         )
 
     # plot using the MC object
+    color = kwargs.pop("cluster_key", color)
     mc.plot_absorption_probabilities(
-        lineages=lineages,
-        cluster_key=cluster_key,
+        states=lineages,
+        color=color,
         mode=mode,
         time_key=time_key,
         **kwargs,
@@ -92,30 +89,24 @@ def lineage_drivers(
     %(just_plots)s
     """
 
-    pk = DummyKernel(adata, backward=backward)
-    mc = GPCCA(pk, read_from_adata=True, write_to_adata=False)
+    mc = CFLARE.from_adata(adata, obsp_key=Key.uns.kernel(backward))
 
     if use_raw and adata.raw is None:
-        logg.warning("No raw attribute set. Using `adata.var` instead")
+        logg.warning("No raw attribute set. Using `use_raw=False`")
         use_raw = False
 
-    direction = DirPrefix.BACKWARD if backward else DirPrefix.FORWARD
-    needle = f"{direction} {lineage} corr"
-
-    haystack = adata.raw.var if use_raw else adata.var
+    key = Key.varm.lineage_drivers(backward)
+    haystack = (adata.raw if use_raw else adata).varm
 
-    if needle not in haystack:
+    if key not in haystack:
         raise RuntimeError(
-            f"Unable to find lineage drivers in "
-            f"`{'adata.raw.var' if use_raw else 'adata.var'}[{needle!r}]`. "
+            "Unable to find lineage drivers in "
+            f"`{'adata.raw.varm' if use_raw else 'adata.varm'}[{key!r}]`. "
             f"Compute lineage drivers first as `cellrank.tl.lineage_drivers(lineages={lineage!r}, "
             f"use_raw={use_raw}, backward={backward}).`"
         )
 
-    drivers = pd.DataFrame(haystack[[needle, f"{direction} {lineage} qval"]])
-    drivers.columns = [f"{lineage} corr", f"{lineage} qval"]
-    mc._set(A.LIN_DRIVERS, drivers)
-
+    mc._lineage_drivers = haystack[key]
     mc.plot_lineage_drivers(
         lineage,
         n_genes=n_genes,
diff --git a/cellrank/pl/_log_odds.py b/cellrank/pl/_log_odds.py
index fe85f22dd..6c614e2ff 100644
--- a/cellrank/pl/_log_odds.py
+++ b/cellrank/pl/_log_odds.py
@@ -5,10 +5,10 @@
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
 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
@@ -100,7 +100,7 @@ def log_odds(
     Returns
     -------
     :class:`matplotlib.pyplot.Axes`
-        The axis object(s) if ``show=False``.
+        The axis object(s) if ``show = False``.
     %(just_plots)s
     """
     from cellrank.tl.kernels._utils import _ensure_numeric_ordered
@@ -182,18 +182,18 @@ def get_data(
         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`.")
+    lineage_key = Key.obsm.abs_probs(backward)
+    if lineage_key not in adata.obsm:
+        raise KeyError(f"Lineages key `{lineage_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)
+    fate1 = adata.obsm[lineage_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)
+        fate2 = adata.obsm[lineage_key][lineage_2].X.squeeze(-1)
         ylabel = rf"$\log{{\frac{{{lineage_1}}}{{{lineage_2}}}}}$"
 
     # fmt: off
diff --git a/cellrank/pl/_utils.py b/cellrank/pl/_utils.py
index 8f60bb130..a74fbf660 100644
--- a/cellrank/pl/_utils.py
+++ b/cellrank/pl/_utils.py
@@ -1,5 +1,3 @@
-"""Utility functions for CellRank plotting."""
-
 from typing import (
     Any,
     Dict,
@@ -18,12 +16,13 @@
 from itertools import combinations
 from collections import namedtuple, defaultdict
 
+from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import _DEFAULT_BACKEND
 from cellrank.ul._docs import d
 from cellrank.tl._utils import save_fig, _unique_order_preserving
 from cellrank.ul.models import GAMR, BaseModel, FailedModel, SKLearnModel
 from cellrank.tl._colors import _create_categorical_colors
-from cellrank.tl._constants import _DEFAULT_BACKEND
 from cellrank.ul._parallelize import parallelize
 from cellrank.ul.models._base_model import ColorType
 
@@ -37,7 +36,6 @@
 from matplotlib import cm
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-AnnData = TypeVar("AnnData")
 Queue = TypeVar("Queue")
 Graph = TypeVar("Graph")
 
@@ -574,7 +572,7 @@ def _trends_helper(
     legend_loc: Optional[str] = "best",
     fig: mpl.figure.Figure = None,
     axes: Union[mpl.axes.Axes, Sequence[mpl.axes.Axes]] = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> None:
     """
     Plot an expression gene for some lineages.
@@ -1096,25 +1094,21 @@ def composition(
     figsize: Optional[Tuple[float, float]] = None,
     dpi: Optional[float] = None,
     save: Optional[Union[str, Path]] = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> None:
     """
     Plot a pie chart for categorical annotation.
 
-    .. image:: https://raw.githubusercontent.com/theislab/cellrank/master/resources/images/composition.png
-       :width: 400px
-       :align: center
-
     Parameters
     ----------
     %(adata)s
     key
-        Key in ``adata.obs`` containing categorical observation.
+        Key in :attr:`anndata.AnnData.obs` containing categorical observation.
     fontsize
         Font size for the pie chart labels.
     %(plotting)s
-    **kwargs
-        Keyworded arguments for :func:`matplotlib.pyplot.pie`.
+    kwargs
+        Keyword arguments for :func:`matplotlib.pyplot.pie`.
 
     Returns
     -------
@@ -1122,21 +1116,22 @@ def composition(
     """
 
     if key not in adata.obs:
-        raise KeyError(f"Key `{key!r}` not found in `adata.obs`.")
+        raise KeyError(f"Data not found in `adata.obs[{key!r}]`.")
     if not is_categorical_dtype(adata.obs[key]):
-        raise TypeError(f"Observation `adata.obs[{key!r}]` is not categorical.")
+        raise TypeError(
+            f"Expected `adata.obs[{key!r}]` is not `categorical`, "
+            f"found `{infer_dtype(adata.obs[key])}`."
+        )
 
-    cats = adata.obs[key].cat.categories
     colors = adata.uns.get(f"{key}_colors", None)
-    x = [np.sum(adata.obs[key] == cl) for cl in cats]
-    cats_frac = x / np.sum(x)
+    x = adata.obs[key].value_counts()
 
     # plot these fractions in a pie plot
     fig, ax = plt.subplots(figsize=figsize, dpi=dpi)
 
     ax.pie(
-        x=cats_frac,
-        labels=cats,
+        x=x,
+        labels=x.index,
         colors=colors,
         textprops={"fontsize": fontsize},
         **kwargs,
diff --git a/cellrank/tl/_constants.py b/cellrank/tl/_constants.py
deleted file mode 100644
index 7b6709186..000000000
--- a/cellrank/tl/_constants.py
+++ /dev/null
@@ -1,168 +0,0 @@
-"""Module containing constants."""
-
-from typing import Any, Union, Callable
-
-from abc import ABC, ABCMeta
-from enum import Enum, EnumMeta
-from functools import wraps
-
-from cellrank.ul._docs import _initial, _terminal
-
-_DEFAULT_BACKEND = "loky"
-
-
-class PrettyEnum(Enum):
-    """Enum with a pretty __str__ and __repr__."""
-
-    def __repr__(self):
-        return str(self)
-
-    def __str__(self):
-        return str(self.value)
-
-    @property
-    def v(self) -> Any:
-        """Return the value."""
-        return self.value
-
-    @property
-    def s(self) -> str:
-        """Return the value as string."""
-        return str(self)
-
-    @property
-    def fmt(self) -> Callable:
-        """Return the formatting function given by the string value."""
-        return self.s.format
-
-
-class Lin(PrettyEnum):
-    """Lineage aggregation type."""
-
-    REST = "rest"
-    OTHERS = "others"
-
-
-class Direction(PrettyEnum):
-    """Direction of the process."""
-
-    FORWARD = "fwd"
-    BACKWARD = "bwd"
-
-
-class DirectionPlot(PrettyEnum):
-    """Pretty direction names for plotting."""
-
-    FORWARD = "forward"
-    BACKWARD = "backward"
-
-
-class DirPrefix(PrettyEnum):
-    """Direction prefix."""
-
-    FORWARD = "to"
-    BACKWARD = "from"
-
-
-class TermStatesKey(PrettyEnum):
-    """State key in `adata.obs`."""
-
-    FORWARD = f"{_terminal}_states"
-    BACKWARD = f"{_initial}_states"
-
-
-class TerminalStatesPlot(PrettyEnum):
-    """Pretty state names for pl."""
-
-    FORWARD = TermStatesKey.FORWARD.s.replace("_", " ")
-    BACKWARD = TermStatesKey.BACKWARD.s.replace("_", " ")
-
-
-class Macro(PrettyEnum):
-    """Key for scaled, subsampled macrostates in `adata.obsm`."""
-
-    FORWARD = f"macrostates_{Direction.FORWARD}"
-    BACKWARD = f"macrostates_{Direction.BACKWARD}"
-
-
-# TermStatesKey and AbsProbKey must have the same suffix `_..._states` because of model.prepare
-class AbsProbKey(PrettyEnum):
-    """Lineage key in `adata.obsm`."""
-
-    FORWARD = f"{DirPrefix.FORWARD}_{TermStatesKey.FORWARD}"
-    BACKWARD = f"{DirPrefix.BACKWARD}_{TermStatesKey.BACKWARD}"
-
-
-class ThresholdScheme(PrettyEnum):
-    """Pseudotime thresholding scheme."""
-
-    SOFT = "soft"
-    HARD = "hard"
-
-
-def _transition(d: Union[str, Direction]) -> str:
-    return f"T_{d}"
-
-
-def _lin_names(k: Union[str, AbsProbKey]) -> str:
-    return f"{k}_names"
-
-
-def _colors(k: Union[str, AbsProbKey, TermStatesKey]) -> str:
-    return f"{k}_colors"
-
-
-def _probs(k: Union[str, TermStatesKey]) -> str:
-    return f"{k}_probs"
-
-
-def _pd(k: Union[str, AbsProbKey]) -> str:
-    return f"{k}_priming_degree"
-
-
-def _pretty_raise_enum(cls, fun):
-    @wraps(fun)
-    def wrapper(*args, **kwargs):
-        try:
-            return fun(*args, **kwargs)
-        except ValueError as e:
-            _cls, value, *_ = args
-            e.args = (cls._format(value),)
-            raise e
-
-    if not issubclass(cls, ErrorFormatterABC):
-        raise TypeError(f"Class `{cls}` must be subtype of `ErrorFormatterABC`.")
-    elif not len(cls.__members__):
-        # empty enum, for class hierarchy
-        return fun
-
-    return wrapper
-
-
-class ABCEnumMeta(EnumMeta, ABCMeta):  # noqa
-    def __call__(cls, *args, **kw):  # noqa
-        if getattr(cls, "__error_format__", None) is None:
-            raise TypeError(
-                f"Can't instantiate class `{cls.__name__}` "
-                f"without `__error_format__` class attribute."
-            )
-        return super().__call__(*args, **kw)
-
-    def __new__(cls, clsname, superclasses, attributedict):  # noqa
-        res = super().__new__(cls, clsname, superclasses, attributedict)
-        res.__new__ = _pretty_raise_enum(res, res.__new__)
-        return res
-
-
-class ErrorFormatterABC(ABC):  # noqa
-    __error_format__ = "Invalid option `{!r}` for `{}`. Valid options are: `{}`."
-
-    @classmethod
-    def _format(cls, value):
-        return cls.__error_format__.format(
-            value, cls.__name__, [m.value for m in cls.__members__.values()]
-        )
-
-
-class ModeEnum(ErrorFormatterABC, PrettyEnum, metaclass=ABCEnumMeta):  # noqa
-    pass
diff --git a/cellrank/tl/_enum.py b/cellrank/tl/_enum.py
new file mode 100644
index 000000000..f4c02ce03
--- /dev/null
+++ b/cellrank/tl/_enum.py
@@ -0,0 +1,75 @@
+from typing import Any, Dict, Type, Tuple, Callable
+from typing_extensions import Literal
+
+from abc import ABC, ABCMeta
+from enum import Enum, EnumMeta
+from functools import wraps
+
+_DEFAULT_BACKEND = "loky"
+Backend_t = Literal["loky", "multiprocessing", "threading"]
+
+
+class PrettyEnum(Enum):
+    """Enum with a pretty :meth:`__str__` and :meth:`__repr__`."""
+
+    @property
+    def v(self) -> Any:
+        """Alias for :attr`value`."""
+        return self.value
+
+    def __repr__(self) -> str:
+        return f"{self.value!r}"
+
+    def __str__(self) -> str:
+        return f"{self.value!s}"
+
+
+def _pretty_raise_enum(cls: Type["ErrorFormatterABC"], func: Callable) -> Callable:
+    @wraps(func)
+    def wrapper(*args: Any, **kwargs: Any) -> "ErrorFormatterABC":
+        try:
+            return func(*args, **kwargs)
+        except ValueError as e:
+            _cls, value, *_ = args
+            e.args = (cls._format(value),)
+            raise e
+
+    if not issubclass(cls, ErrorFormatterABC):
+        raise TypeError(f"Class `{cls}` must be subtype of `ErrorFormatterABC`.")
+    elif not len(cls.__members__):
+        # empty enum, for class hierarchy
+        return func
+
+    return wrapper
+
+
+class ABCEnumMeta(EnumMeta, ABCMeta):  # noqa: D101
+    def __call__(cls, *args, **kwargs):  # noqa
+        if getattr(cls, "__error_format__", None) is None:
+            raise TypeError(
+                f"Can't instantiate class `{cls.__name__}` "
+                f"without `__error_format__` class attribute."
+            )
+        return super().__call__(*args, **kwargs)
+
+    def __new__(  # noqa: D102
+        cls, clsname: str, superclasses: Tuple[type], attributedict: Dict[str, Any]
+    ):
+        res = super().__new__(cls, clsname, superclasses, attributedict)
+        res.__new__ = _pretty_raise_enum(res, res.__new__)
+        return res
+
+
+class ErrorFormatterABC(ABC):  # noqa: D101
+    __error_format__ = "Invalid option `{!r}` for `{}`. Valid options are: `{}`."
+
+    @classmethod
+    def _format(cls, value) -> str:
+        return cls.__error_format__.format(
+            value, cls.__name__, [m.value for m in cls.__members__.values()]
+        )
+
+
+class ModeEnum(str, ErrorFormatterABC, PrettyEnum, metaclass=ABCEnumMeta):  # noqa: D101
+    def _generate_next_value_(self, start, count, last_values):
+        return str(self).lower()
diff --git a/cellrank/tl/_init_term_states.py b/cellrank/tl/_init_term_states.py
index 16acbfe64..3123a18f7 100644
--- a/cellrank/tl/_init_term_states.py
+++ b/cellrank/tl/_init_term_states.py
@@ -1,8 +1,8 @@
-"""Module used for finding initial and terminal states."""
-from typing import Union, Mapping, TypeVar, Optional
+from typing import Any, Union, Mapping, Optional, Sequence
 
 from types import MappingProxyType
 
+from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d, _initial, _terminal, inject_docs
 from cellrank.tl._utils import (
@@ -11,16 +11,11 @@
     _info_if_obs_keys_categorical_present,
 )
 from cellrank.tl.kernels import PrecomputedKernel
-from cellrank.tl._constants import TermStatesKey
-from cellrank.tl.estimators import GPCCA, CFLARE
+from cellrank.tl.estimators import GPCCA, CFLARE, TermStatesEstimator
 from cellrank.tl._transition_matrix import transition_matrix
-from cellrank.tl.estimators._constants import P
 from cellrank.tl.kernels._velocity_kernel import BackwardMode, VelocityMode
 from cellrank.tl.estimators._base_estimator import BaseEstimator
 
-AnnData = TypeVar("AnnData")
-
-
 _docstring = """\
 Find {direction} states of a dynamic process of single cells based on RNA velocity :cite:`manno:18`.
 
@@ -75,12 +70,69 @@
 """
 
 
+def _fit(
+    estim: TermStatesEstimator,
+    n_lineages: Optional[int] = None,
+    keys: Optional[Sequence[str]] = None,
+    cluster_key: Optional[str] = None,
+    compute_absorption_probabilities: bool = True,
+    **kwargs: Any,
+) -> TermStatesEstimator:
+    if isinstance(estim, CFLARE):
+        estim.compute_eigendecomposition(k=20 if n_lineages is None else n_lineages + 1)
+        if n_lineages is None:
+            n_lineages = estim.eigendecomposition["eigengap"] + 1
+
+        estim.compute_terminal_states(
+            use=n_lineages,
+            cluster_key=cluster_key,
+            n_clusters_kmeans=n_lineages,
+            method=kwargs.pop("method", "kmeans"),
+            **kwargs,
+        )
+    elif isinstance(estim, GPCCA):
+        if n_lineages is None or n_lineages == 1:
+            estim.compute_eigendecomposition()
+            if n_lineages is None:
+                n_lineages = estim.eigendecomposition["eigengap"] + 1
+
+        if n_lineages > 1:
+            estim.compute_schur(n_lineages, method=kwargs.pop("method", "krylov"))
+
+        try:
+            estim.compute_macrostates(
+                n_states=n_lineages, cluster_key=cluster_key, **kwargs
+            )
+        except ValueError:
+            logg.warning(
+                f"Computing `{n_lineages}` macrostates cuts through a block of complex conjugates. "
+                f"Increasing `n_lineages` to {n_lineages + 1}"
+            )
+            estim.compute_macrostates(
+                n_states=n_lineages + 1, cluster_key=cluster_key, **kwargs
+            )
+
+        fs_kwargs = {"n_cells": kwargs["n_cells"]} if "n_cells" in kwargs else {}
+
+        if n_lineages is None:
+            estim.compute_terminal_states(method="eigengap", **fs_kwargs)
+        else:
+            estim.set_terminal_states_from_macrostates(**fs_kwargs)
+    else:
+        raise NotImplementedError(type(estim))
+
+    if compute_absorption_probabilities:
+        estim.compute_absorption_probabilities(keys=keys)
+
+    return estim
+
+
 def _initial_terminal(
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
     backward: bool = False,
-    mode: str = VelocityMode.DETERMINISTIC.s,
-    backward_mode: str = BackwardMode.TRANSPOSE.s,
+    mode: str = VelocityMode.DETERMINISTIC,
+    backward_mode: str = BackwardMode.TRANSPOSE,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
@@ -93,11 +145,12 @@ def _initial_terminal(
 ) -> Optional[Union[AnnData, BaseEstimator]]:
     _check_estimator_type(estimator)
 
+    if copy:
+        adata = adata.copy()
     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")
     except KeyError:
         # compute kernel object
@@ -108,16 +161,8 @@ def _initial_terminal(
             backward_mode=backward_mode,
             **kwargs,
         )
-        write_to_adata = True
 
-    # create estimator object
-    mc = estimator(
-        kernel,
-        read_from_adata=False,
-        inplace=not copy,
-        key=key,
-        write_to_adata=write_to_adata,
-    )
+    mc = estimator(kernel)
 
     if cluster_key is None:
         _info_if_obs_keys_categorical_present(
@@ -126,7 +171,8 @@ def _initial_terminal(
             msg_fmt="Found categorical observation in `adata.obs[{!r}]`. Consider specifying it as `cluster_key`.",
         )
 
-    mc.fit(
+    _fit(
+        mc,
         n_lineages=n_states,
         cluster_key=cluster_key,
         compute_absorption_probabilities=False,
@@ -136,12 +182,9 @@ def _initial_terminal(
     if show_plots:
         mc.plot_spectrum(real_only=True)
         if isinstance(mc, CFLARE):
-            mc.plot_eigendecomposition(abs_value=True, perc=[0, 98], use=n_states)
             mc.plot_terminal_states(discrete=True, same_plot=False)
         elif isinstance(mc, GPCCA):
-            n_states = len(mc._get(P.MACRO).cat.categories)
-            if n_states > 1:
-                mc.plot_schur()
+            n_states = len(mc.macrostates.cat.categories)
             mc.plot_terminal_states(discrete=True, same_plot=False)
             if n_states > 1:
                 mc.plot_coarse_T()
@@ -159,15 +202,15 @@ def _initial_terminal(
 @inject_docs(
     __doc__=_docstring.format(
         direction=_initial,
-        key_added=TermStatesKey.BACKWARD.s,
+        key_added="initial_states",
         bwd_mode="\n%(velocity_backward_mode_high_lvl)s",
     )
 )
 def initial_states(  # noqa: D103
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
-    mode: str = VelocityMode.DETERMINISTIC.s,
-    backward_mode: str = BackwardMode.TRANSPOSE.s,
+    mode: str = VelocityMode.DETERMINISTIC,
+    backward_mode: str = BackwardMode.TRANSPOSE,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
@@ -202,13 +245,13 @@ def initial_states(  # noqa: D103
 @d.dedent
 @inject_docs(
     __doc__=_docstring.format(
-        direction=_terminal, key_added=TermStatesKey.FORWARD.s, bwd_mode=""
+        direction=_terminal, key_added="terminal_states", bwd_mode=""
     )
 )
 def terminal_states(  # noqa: D103
     adata: AnnData,
     estimator: type(BaseEstimator) = GPCCA,
-    mode: str = VelocityMode.DETERMINISTIC.s,
+    mode: str = VelocityMode.DETERMINISTIC,
     n_states: Optional[int] = None,
     cluster_key: Optional[str] = None,
     key: Optional[str] = None,
diff --git a/cellrank/tl/_lineage.py b/cellrank/tl/_lineage.py
index 9ba5de07e..eea50e4c9 100644
--- a/cellrank/tl/_lineage.py
+++ b/cellrank/tl/_lineage.py
@@ -1,27 +1,30 @@
-"""Lineage class module."""
-
 from typing import List, Tuple, Union, Mapping, TypeVar, Callable, Iterable, Optional
 from typing_extensions import Literal
 
+from copy import copy
+from enum import auto
 from types import FunctionType, MappingProxyType
 from inspect import signature
 from pathlib import Path
 from functools import wraps
 from itertools import combinations
 
+from anndata import AnnData
 from cellrank import logging as logg
-from cellrank.ul._docs import d
+from cellrank._key import Key
+from cellrank.tl._enum import ModeEnum
+from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import save_fig, _convert_lineage_name, _unique_order_preserving
 from cellrank.tl._colors import (
     _get_bg_fg_colors,
     _compute_mean_color,
     _create_categorical_colors,
 )
-from cellrank.tl._constants import Lin, ModeEnum
 
 import numpy as np
 import pandas as pd
 from scipy.stats import entropy
+from pandas.api.types import infer_dtype, is_categorical_dtype
 
 import matplotlib.colors as c
 import matplotlib.pyplot as plt
@@ -36,11 +39,39 @@
 _ORDER = "C"
 
 
-class PrimingDegree(ModeEnum):
-    """Priming degree method."""
+class PrimingDegree(ModeEnum):  # noqa: D101
+    KL_DIVERGENCE = auto()
+    ENTROPY = auto()
+
+
+class Lin(ModeEnum):  # noqa: D101
+    REST = auto()
+    OTHERS = auto()
+
+
+class DistanceMeasure(ModeEnum):  # noqa: D101
+    COSINE_SIM = auto()
+    WASSERSTEIN_DIST = auto()
+    KL_DIV = auto()
+    JS_DIV = auto()
+    MUTUAL_INFO = auto()
+    EQUAL = auto()
+
+
+class NormWeights(ModeEnum):  # noqa: D101
+    SCALE = auto()
+    SOFTMAX = auto()
 
-    KL_DIVERGENCE = "kl_divergence"
-    ENTROPY = "entropy"
+
+class Reduction(ModeEnum):  # noqa: D101
+    DIST = auto()
+    SCALE = auto()
+
+
+class LinKind(ModeEnum):  # noqa: D101
+    MACROSTATES = auto()
+    TERM_STATES = auto()
+    ABS_PROBS = auto()
 
 
 def _at_least_2d(array: np.ndarray, dim: int):
@@ -276,7 +307,7 @@ def update_entries(key):
         lin_kind = [_ for _ in mixtures if isinstance(_, Lin)]
         if len(lin_kind) > 1:
             raise ValueError(
-                f"`Lin` enum is allowed only once in the expression, found `{lin_kind}`."
+                f"`Lin` enum is allowed only once in the expression, found `{len(lin_kind)}`."
             )
 
         keys = [
@@ -319,7 +350,9 @@ def update_entries(key):
                 if keys:
                     names[-1] = str(lin_kind)
             else:
-                raise ValueError(f"Invalid `Lin` enum `{lin_kind}`.")
+                raise NotImplementedError(
+                    f"Mixing `{lin_kind}` is not yet implemented."
+                )
 
         res = np.stack(res, axis=-1)
         return Lineage(res, names=names, colors=colors)
@@ -725,20 +758,20 @@ def priming_degree(
         """
         Compute the degree of lineage priming.
 
-        This method computes how naive vs. committed each individual cell is.
-        It returns a score where 0 stands for naive and 1 stands for committed.
+        It returns a score in `[0, 1]` where `0` stands for naive and `1` stands for committed.
 
         Parameters
         ----------
         method
             The method used to compute the degree of lineage priming. Valid options are:
 
-                - `'kl_divergence'`: as in :cite:`velten:17`, computes KL-divergence between the fate probabilities of
+                - `'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 :cite:`setty:19`, 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'``.
+            Cell IDs or a mask marking early cells. If `None`, use all cells.
+            Only used when ``method = 'kl_divergence'``.
 
         Returns
         -------
@@ -862,14 +895,15 @@ def plot_pie(
             save_fig(fig, save)
 
     @d.dedent
+    @inject_docs(m=Reduction, dm=DistanceMeasure, nw=NormWeights)
     def reduce(
         self,
         *keys: str,
-        mode: str = "dist",
+        mode: Literal["dist", "scale"] = Reduction.DIST,
         dist_measure: Literal[
-            "cosine_sim", "wasserstein_dist", "kl_div", "js_div", "mutual_inf", "equal"
-        ] = "mutual_info",
-        normalize_weights: Literal["scale", "softmax"] = "softmax",
+            "cosine_sim", "wasserstein_dist", "kl_div", "js_div", "mutual_info", "equal"
+        ] = DistanceMeasure.MUTUAL_INFO,
+        normalize_weights: Literal["scale", "softmax"] = NormWeights.SOFTMAX,
         softmax_scale: float = 1,
         return_weights: bool = False,
     ) -> Union["Lineage", Tuple["Lineage", Optional[pd.DataFrame]]]:
@@ -882,34 +916,44 @@ def reduce(
             List of keys that define the states, to which this object will be reduced by projecting the values
             of the other states.
         mode
-            Whether to use a distance measure to compute weights - `'dist'`, or just rescale - `'scale'`.
+            Reduction mode to use. Valid options are:
+
+                - `{m.DIST!r}` - use a distance measure ``dist_measure`` to compute weights.
+                - `{m.SCALE!r}` - just rescale the values.
         dist_measure
             Used to quantify similarity between query and reference states. Valid options are:
 
-                - `'cosine_sim'` - cosine similarity.
-                - `'wasserstein_dist'` - Wasserstein distance.
-                - `'kl_div'` - Kullback–Leibler divergence.
-                - `'js_div'` - Jensen–Shannon divergence.
-                - `'mutual_inf'` - mutual information.
-                - `'equal'` - equally redistribute the mass among the rest.
+                - `{dm.COSINE_SIM!r}` - cosine similarity.
+                - `{dm.WASSERSTEIN_DIST!r}` - Wasserstein distance.
+                - `{dm.KL_DIV!r}` - Kullback–Leibler divergence.
+                - `{dm.JS_DIV!r}` - Jensen–Shannon divergence.
+                - `{dm.MUTUAL_INFO!r}` - mutual information.
+                - `{dm.EQUAL!r}` - equally redistribute the mass among the rest.
+
+            Only use when ``mode = {m.DIST!r}``.
         normalize_weights
             How to row-normalize the weights. Valid options are:
 
-                - `'scale'` - divide by the sum.
-                - `'softmax'`- use a softmax.
+                - `{nw.SCALE!r}` - divide by the sum.
+                - `{nw.SOFTMAX!r}`- use a softmax.
+
+            Only use when ``mode = {m.DIST!r}``.
         softmax_scale
             Scaling factor in the softmax, used for normalizing the weights to sum to `1`.
         return_weights
-            If `True`, a :class:`pandas.DataFrame` of the weights used for the projection is returned.
+            If `True`, a :class:`pandas.DataFrame` of the weights used for the projection is also returned.
+            If ``mode = {m.SCALE!r}``, the weights will be `None`.
 
         Returns
         -------
-        :class:`cellrank.tl.Lineage`, :class:`pandas.DataFrame`
-            Lineage object, reduced to the %(initial_or_terminal)s states. If a reduction is not possible,
-            returns just a copy of self.The weights used for the projection of shape
-            ``(n_query, n_reference)``, if ``return_weights=True``.
+        Lineage object, reduced to the %(initial_or_terminal)s states.
+        The weights used for the projection of shape ``(n_query, n_reference)``, if ``return_weights = True``.
         """
 
+        mode = Reduction(mode)
+        dist_measure = DistanceMeasure(dist_measure)
+        normalize_weights = NormWeights(normalize_weights)
+
         if self._is_transposed:
             raise RuntimeError("This method works only on non-transposed lineages.")
 
@@ -929,9 +973,9 @@ def reduce(
             )
 
         # check input parameters
-        if return_weights and mode == "scale":
+        if return_weights and mode == Reduction.SCALE:
             logg.warning(
-                "If `mode=='scale'`, no weights are computed. Returning `None`"
+                f"If `mode={mode!r}`, no weights are computed. Returning `None`"
             )
 
         reference = self[:, keys]
@@ -946,22 +990,26 @@ def reduce(
 
         query = self[:, rest]
 
-        if mode == "dist":
+        if mode == Reduction.SCALE:
+            reference = _row_normalize(reference)
+        elif mode == Reduction.DIST:
             # compute a set of weights of shape (n_query x n_reference)
-            if dist_measure == "cosine_sim":
+            if dist_measure == DistanceMeasure.COSINE_SIM:
                 weights = _cosine_sim(reference.X, query.X)
-            elif dist_measure == "wasserstein_dist":
+            elif dist_measure == DistanceMeasure.WASSERSTEIN_DIST:
                 weights = _wasserstein_dist(reference.X, query.X)
-            elif dist_measure == "kl_div":
+            elif dist_measure == DistanceMeasure.KL_DIV:
                 weights = _kl_div(reference.X, query.X)
-            elif dist_measure == "js_div":
+            elif dist_measure == DistanceMeasure.JS_DIV:
                 weights = _js_div(reference.X, query.X)
-            elif dist_measure == "mutual_info":
+            elif dist_measure == DistanceMeasure.MUTUAL_INFO:
                 weights = _mutual_info(reference.X, query.X)
-            elif dist_measure == "equal":
+            elif dist_measure == DistanceMeasure.EQUAL:
                 weights = _row_normalize(np.ones((query.shape[1], reference.shape[1])))
             else:
-                raise ValueError(f"Invalid distance measure `{dist_measure!r}`.")
+                raise NotImplementedError(
+                    f"Distance measure `{dist_measure}` is not yet implemented."
+                )
 
             # make some checks on the weights
             if weights.shape != (query.shape[1], reference.shape[1]):
@@ -980,13 +1028,13 @@ def reduce(
                 logg.warning("Weights matrix contains exact zeros.")
 
             # normalize the weights to row-sum to one
-            if normalize_weights == "scale":
+            if normalize_weights == NormWeights.SCALE:
                 weights_n = _row_normalize(weights)
-            elif normalize_weights == "softmax":
+            elif normalize_weights == NormWeights.SOFTMAX:
                 weights_n = _softmax(_row_normalize(weights), softmax_scale)
             else:
-                raise ValueError(
-                    f"Normalization method `{normalize_weights!r}` not found. Valid options are: `'scale', 'softmax'`."
+                raise NotImplementedError(
+                    f"Normalization method `{normalize_weights}` is yet implemented."
                 )
 
             # check that the weights row-sum to one now
@@ -996,12 +1044,9 @@ def reduce(
             # use the weights to re-distribute probability mass form query to reference
             for i, w in enumerate(weights_n):
                 reference += np.dot(query[:, i].X, w[None, :])
-
-        elif mode == "scale":
-            reference = _row_normalize(reference)
         else:
-            raise ValueError(
-                f"Invalid mode `{mode!r}`. Valid options are: `'dist', 'scale'`."
+            raise NotImplementedError(
+                f"Reduction mode `{mode}` is not yet implemented."
             )
 
         # check that the lineages row-sum to one now
@@ -1010,7 +1055,7 @@ def reduce(
 
         # potentially create a weights-df and return everything
         if return_weights:
-            if mode == "dist":
+            if mode == Reduction.DIST:
                 return (
                     reference,
                     pd.DataFrame(
@@ -1021,6 +1066,91 @@ def reduce(
 
         return reference
 
+    @classmethod
+    @d.dedent
+    @inject_docs(lk=LinKind)
+    def from_adata(
+        cls,
+        adata: AnnData,
+        backward: bool = False,
+        kind: Literal["macrostates", "term_states", "abs_probs"] = LinKind.ABS_PROBS,
+    ) -> "Lineage":
+        """
+        Reconstruct :class:`cellrank.tl.Lineage` from :class:`anndata.AnnData`.
+
+        Parameters
+        ----------
+        %(adata)s
+        %(backward)s
+        kind
+            Which kind of object to reconstruct. Valid options are:
+
+                - `{lk.MACROSTATES!r}`- macrostates memberships from :class:`cellrank.tl.estimators.GPCCA`.
+                - `{lk.TERM_STATES!r}`- terminal states memberships from :class:`cellrank.tl.estimators.GPCCA`.
+                - `{lk.ABS_PROBS!r}`- the absorption probabilities.
+
+        Returns
+        -------
+        The reconstructed lineage object.
+        """
+        kind = LinKind(kind)
+        if kind == LinKind.MACROSTATES:
+            nkey = Key.obs.macrostates(backward)
+            key = Key.obsm.memberships(nkey)
+        elif kind == LinKind.TERM_STATES:
+            nkey = Key.obs.term_states(backward)
+            key = Key.obsm.memberships(nkey)
+        elif kind == LinKind.ABS_PROBS:
+            nkey = Key.obs.term_states(backward)
+            key = Key.obsm.abs_probs(backward)
+        else:
+            raise NotImplementedError(f"Lineage kind `{kind}` is not yet implemented.")
+
+        ckey = Key.uns.colors(nkey)
+
+        data: Optional[Union[np.ndarray, Lineage]] = copy(adata.obsm.get(key, None))
+        if data is None:
+            raise KeyError(f"Unable to find lineage data in `adata.obsm[{key!r}]`.")
+        if isinstance(data, Lineage):
+            return data
+        if data.ndim != 2:
+            raise ValueError(f"Expected 2 dimensional data, found `{data.ndim}`.")
+
+        states = adata.obs.get(nkey, None)
+        if states is None:
+            logg.warning(
+                f"Unable to find states in `adata.obs[{nkey!r}]`. Using default names"
+            )
+        elif not is_categorical_dtype(states):
+            logg.warning(
+                f"Expected `adata.obs[{key!r}]` to be `categorical`, "
+                f"found `{infer_dtype(adata.obs[nkey])}`. Using default names"
+            )
+        else:
+            states = list(states.cat.categories)
+            if len(states) != data.shape[1]:
+                logg.warning(
+                    f"Expected to find `{data.shape[1]}` names, found `{len(states)}`. "
+                    f"Using default names"
+                )
+        if states is None or len(states) != data.shape[1]:
+            states = [str(i) for i in range(data.shape[1])]
+
+        colors = adata.uns.get(ckey, None)
+        if colors is None:
+            logg.warning(
+                f"Unable to find colors in `adata.uns[{ckey!r}]`. "
+                f"Using default colors"
+            )
+        elif len(colors) != data.shape[1]:
+            logg.warning(
+                f"Expected to find `{data.shape[1]}` colors, found `{len(colors)}`. "
+                f"Using default colors"
+            )
+            colors = None
+
+        return Lineage(data, names=states, colors=colors)
+
 
 class LineageView(Lineage):
     """View of :class:`cellrank.tl.Lineage`."""
diff --git a/cellrank/tl/_lineages.py b/cellrank/tl/_lineages.py
index 41bb918b5..a369c7418 100644
--- a/cellrank/tl/_lineages.py
+++ b/cellrank/tl/_lineages.py
@@ -1,20 +1,14 @@
-"""Lineages module."""
+from typing import Any, Union, Optional, Sequence
+from typing_extensions import Literal
 
-from typing import Union, TypeVar, Optional, Sequence
-
-from cellrank import logging as logg
+from anndata import AnnData
+from cellrank._key import Key
 from cellrank.ul._docs import d
 from cellrank.tl._utils import TestMethod, _deprecate
-from cellrank.tl.kernels import PrecomputedKernel
-from cellrank.tl._constants import AbsProbKey, TermStatesKey, TerminalStatesPlot
-from cellrank.tl.estimators import GPCCA
-from cellrank.tl.estimators._constants import P
-from cellrank.tl.kernels._precomputed_kernel import DummyKernel
+from cellrank.tl.estimators import CFLARE
 
 import pandas as pd
 
-AnnData = TypeVar("AnnData")
-
 
 @_deprecate(version="2.0")
 @d.dedent
@@ -23,7 +17,7 @@ def lineages(
     backward: bool = False,
     copy: bool = False,
     return_estimator: bool = False,
-    **kwargs,
+    **kwargs: Any,
 ) -> Optional[AnnData]:
     """
     Compute probabilistic lineage assignment using RNA velocity.
@@ -47,7 +41,7 @@ def lineages(
     copy
         Whether to update the existing ``adata`` object or to return a copy.
     return_estimator
-        Whether to return the estimator. Only available when ``copy=False``.
+        Whether to return the estimator. Only available when ``copy = False``.
     kwargs
         Keyword arguments for :meth:`cellrank.tl.estimators.BaseEstimator.compute_absorption_probabilities`.
 
@@ -57,37 +51,26 @@ def lineages(
         Depending on ``copy`` and ``return_estimator``, either updates the existing ``adata`` object,
         returns its copy or returns the estimator.
     """
-
-    if backward:
-        lin_key = AbsProbKey.BACKWARD
-        fs_key = TermStatesKey.BACKWARD
-        fs_key_pretty = TerminalStatesPlot.BACKWARD
-    else:
-        lin_key = AbsProbKey.FORWARD
-        fs_key = TermStatesKey.FORWARD
-        fs_key_pretty = TerminalStatesPlot.FORWARD
+    if copy:
+        adata = adata.copy()
 
     try:
-        pk = PrecomputedKernel(adata=adata, backward=backward)
+        mc = CFLARE.from_adata(adata, obsp_key=Key.uns.kernel(backward))
+        assert mc.adata is adata
     except KeyError as e:
         raise RuntimeError(
             f"Compute transition matrix first as `cellrank.tl.transition_matrix(..., backward={backward})`."
         ) from e
 
-    start = logg.info(f"Computing lineage probabilities towards {fs_key_pretty.s}")
-    mc = GPCCA(
-        pk, read_from_adata=True, inplace=not copy
-    )  # GPCCA is more general than CFLARE, in terms of what is saves
-    if mc._get(P.TERM) is None:
+    if mc.terminal_states is None:
+        fs_key = Key.obs.term_states(backward)
         raise RuntimeError(
-            f"Compute the states first as `cellrank.tl.{fs_key.s}(..., backward={backward})`."
+            f"Compute the states first as `cellrank.tl.{fs_key}(..., backward={backward})`."
         )
 
     # compute the absorption probabilities
     mc.compute_absorption_probabilities(**kwargs)
 
-    logg.info(f"Adding lineages to `adata.obsm[{lin_key.s!r}]`\n    Finish", time=start)
-
     return mc.adata if copy else mc if return_estimator else None
 
 
@@ -96,18 +79,17 @@ def lineages(
 def lineage_drivers(
     adata: AnnData,
     backward: bool = False,
-    lineages: Optional[Union[Sequence, str]] = None,
-    method: str = TestMethod.FISCHER.s,
+    lineages: Optional[Union[str, Sequence[str]]] = None,
+    method: Literal["fischer", "perm_test"] = TestMethod.FISCHER,
     cluster_key: Optional[str] = None,
-    clusters: Optional[Union[Sequence, str]] = None,
+    clusters: Optional[Union[str, Sequence[str]]] = None,
     layer: str = "X",
     use_raw: bool = False,
     confidence_level: float = 0.95,
     n_perms: int = 1000,
     seed: Optional[int] = None,
-    return_drivers: bool = True,
-    **kwargs,
-) -> Optional[pd.DataFrame]:
+    **kwargs: Any,
+) -> pd.DataFrame:
     """
     %(lineage_drivers.full_desc)s
 
@@ -120,16 +102,10 @@ def lineage_drivers(
     Returns
     -------
     %(lineage_drivers.returns)s
-
-    References
-    ----------
-    %(lineage_drivers.references)s
     """  # noqa: D400
 
-    # create dummy kernel and estimator
-    pk = DummyKernel(adata, backward=backward)
-    g = GPCCA(pk, read_from_adata=True, write_to_adata=False)
-    if g._get(P.ABS_PROBS) is None:
+    g = CFLARE.from_adata(adata, obsp_key=Key.uns.kernel(backward))
+    if g.absorption_probabilities is None:
         raise RuntimeError(
             f"Compute absorption probabilities first as `cellrank.tl.lineages(..., backward={backward})`."
         )
@@ -145,6 +121,5 @@ def lineage_drivers(
         confidence_level=confidence_level,
         n_perms=n_perms,
         seed=seed,
-        return_drivers=return_drivers,
         **kwargs,
     )
diff --git a/cellrank/tl/_linear_solver.py b/cellrank/tl/_linear_solver.py
index 9905c6490..069cc0f90 100644
--- a/cellrank/tl/_linear_solver.py
+++ b/cellrank/tl/_linear_solver.py
@@ -4,8 +4,8 @@
 from functools import singledispatch
 
 from cellrank import logging as logg
+from cellrank.tl._enum import _DEFAULT_BACKEND
 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
@@ -169,7 +169,7 @@ def _(
 
     ksp.solve(b, x)
 
-    return x.getArray().copy().squeeze(), int(ksp.converged)
+    return np.atleast_1d(x.getArray().copy().squeeze()), int(ksp.converged)
 
 
 @_solve_many_sparse_problems_petsc.register(csc_matrix)
@@ -192,7 +192,7 @@ def _(
 
         ksp.solve(b, x)
 
-        xs.append(x.getArray().copy().squeeze())
+        xs.append(np.atleast_1d(x.getArray().copy().squeeze()))
         converged += ksp.converged
 
         if queue is not None:
@@ -249,7 +249,7 @@ def _solve_many_sparse_problems(
         x, info = solver(mat_a, b.toarray().flatten(), tol=tol, x0=None, **kwargs)
 
         # append solution and info
-        x_list.append(x)
+        x_list.append(np.atleast_1d(x))
         n_converged += info == 0
 
         if queue is not None:
@@ -261,7 +261,7 @@ def _solve_many_sparse_problems(
     return np.stack(x_list, axis=1), n_converged
 
 
-def _petsc_mat_solve(
+def _petsc_direct_solve(
     mat_a: Union[np.ndarray, spmatrix],
     mat_b: Optional[Union[spmatrix, np.ndarray]] = None,
     tol: float = 1e-5,
@@ -298,7 +298,7 @@ def _petsc_mat_solve(
                     f"did not converge"
                 )
 
-            return res
+            return res[:, None]
 
         B = _create_petsc_matrix(mat_b, as_dense=True)
 
@@ -430,7 +430,7 @@ def extractor(
     if solver == "direct":
         if use_petsc:
             logg.debug("Solving the linear system directly using `PETSc`")
-            return _petsc_mat_solve(
+            return _petsc_direct_solve(
                 mat_a, mat_b, solver=solver, preconditioner=preconditioner, tol=tol
             )
 
@@ -461,7 +461,6 @@ def extractor(
             f"`tol={tol}`"
         )
 
-        # can't pass PETSc matrix - not pickleable
         mat_x, n_converged = parallelize(
             _solve_many_sparse_problems_petsc,
             mat_b,
diff --git a/cellrank/tl/_mixins/__init__.py b/cellrank/tl/_mixins/__init__.py
new file mode 100644
index 000000000..e82aa4f59
--- /dev/null
+++ b/cellrank/tl/_mixins/__init__.py
@@ -0,0 +1,3 @@
+from cellrank.tl._mixins._io import IOMixin
+from cellrank.tl._mixins._kernel import KernelMixin
+from cellrank.tl._mixins._anndata import AnnDataMixin
diff --git a/cellrank/tl/_mixins/_anndata.py b/cellrank/tl/_mixins/_anndata.py
new file mode 100644
index 000000000..a727bfa9e
--- /dev/null
+++ b/cellrank/tl/_mixins/_anndata.py
@@ -0,0 +1,70 @@
+from typing import Any, Optional
+
+from abc import ABC, abstractmethod
+
+from anndata import AnnData
+from cellrank.ul._docs import d
+
+
+class AnnDataMixin(ABC):
+    """Mixin that allows for serialization from/to :class:`anndata.AnnData`."""
+
+    @abstractmethod
+    @d.dedent
+    def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
+        """
+        Populate attributes of self from :class:`anndata.AnnData`.
+
+        Parameters
+        ----------
+        %(adata)s
+        kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        `True` if the deserialization was successful, otherwise `False`.
+        """
+        return True
+
+    @property
+    @abstractmethod
+    def adata(self) -> AnnData:
+        """Annotated data object."""
+
+    @adata.setter
+    @abstractmethod
+    def adata(self, adata: Optional[AnnData]) -> None:
+        ...
+
+    @abstractmethod
+    def __len__(self) -> int:
+        ...
+
+    @abstractmethod
+    @d.get_full_description(base="to_adata")
+    def to_adata(self, **kwargs: Any) -> AnnData:
+        """Serialize self to :class:`anndata.Anndata`."""
+
+    @classmethod
+    @d.get_full_description(base="from_adata")
+    @d.get_sections(base="from_adata", sections=["Returns"])
+    @d.dedent
+    def from_adata(cls, adata: AnnData, **kwargs: Any) -> "AnnDataMixin":
+        """
+        Deserialize self from :class:`anndata.AnnData`.
+
+        Parameters
+        ----------
+        %(adata)s
+        kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        The deserialized object.
+        """
+        obj = cls(adata, **kwargs)
+        obj._read_from_adata(adata)
+
+        return obj
diff --git a/cellrank/tl/_mixins/_io.py b/cellrank/tl/_mixins/_io.py
new file mode 100644
index 000000000..22aabe430
--- /dev/null
+++ b/cellrank/tl/_mixins/_io.py
@@ -0,0 +1,136 @@
+from typing import Any, Tuple, Union, Optional
+from typing_extensions import Protocol
+
+import pickle
+from pathlib import Path
+from contextlib import contextmanager
+
+from anndata import AnnData
+from cellrank import logging as logg
+
+
+class IOMixinProtocol(Protocol):  # noqa: D101
+    @property
+    def shape(self) -> Tuple[int, ...]:  # noqa: D102
+        ...
+
+    @property
+    def adata(self) -> AnnData:  # noqa: D102
+        ...
+
+    @adata.setter
+    def adata(self, adata: AnnData) -> None:
+        ...
+
+
+class IOMixin:
+    """Mixin that allows for serialization from/to files using :mod:`pickle`."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+    @property
+    @contextmanager
+    def _remove_adata(self) -> None:
+        """Temporarily remove :attr:`adata`, if present."""
+        adata = getattr(self, "adata", None)
+
+        try:
+            if adata is not None:
+                self.adata = None
+            yield
+        finally:
+            if adata is not None:
+                self.adata = adata
+
+    def write(
+        self,
+        fname: Union[str, Path],
+        write_adata: bool = True,
+        ext: Optional[str] = "pickle",
+    ) -> None:
+        """
+        Serialize self to a file.
+
+        Parameters
+        ----------
+        fname
+            Filename where to save the object.
+        write_adata
+            Whether to save :attr:`adata` object or not, if present.
+        ext
+            Filename extension to use. If `None`, don't append any extension.
+
+        Returns
+        -------
+        Nothing, just writes itself to a file using :mod:`pickle`.
+        """
+
+        fname = str(fname)
+        if ext is not None:
+            if not ext.startswith("."):
+                ext = "." + ext
+            if not fname.endswith(ext):
+                fname += ext
+
+        logg.info(f"Writing `{self}` to `{fname}`")
+
+        if write_adata:
+            with open(fname, "wb") as fout:
+                pickle.dump(self, fout)
+            return
+
+        with self._remove_adata:
+            with open(fname, "wb") as fout:
+                pickle.dump(self, fout)
+
+    @staticmethod
+    def read(
+        fname: Union[str, Path], adata: Optional[AnnData] = None, copy: bool = False
+    ) -> "IOMixin":
+        """
+        Deserialize self from a file.
+
+        Parameters
+        ----------
+        fname
+            Filename from which to read the object.
+        adata
+            :class:`anndata.AnnData` object to assign to the saved object.
+            Only used when the saved object has :attr:`adata` and it was saved without it.
+        copy
+            Whether to copy ``adata`` before assigning it or not. If ``adata`` is a view, it is always copied.
+
+        Returns
+        -------
+        The deserialized object.
+        """
+
+        with open(fname, "rb") as fin:
+            obj: IOMixinProtocol = pickle.load(fin)
+
+        if hasattr(obj, "adata"):
+            if isinstance(obj.adata, AnnData):
+                if adata is not None:
+                    logg.warning(
+                        "Ignoring supplied `adata` object because it is already present"
+                    )
+                return obj
+
+            if not isinstance(adata, AnnData):
+                raise TypeError(
+                    "This object was saved without its `adata` object. "
+                    "Please supply one as `adata=...`."
+                )
+
+            if obj.shape[0] != len(adata):
+                raise ValueError(
+                    f"Expected `adata` to be of length `{len(adata)}`, found `{obj.shape[0]}`."
+                )
+            if copy or adata.is_view:
+                adata = adata.copy()
+
+            obj.adata = adata
+            return obj
+
+        return obj
diff --git a/cellrank/tl/_mixins/_kernel.py b/cellrank/tl/_mixins/_kernel.py
new file mode 100644
index 000000000..784506074
--- /dev/null
+++ b/cellrank/tl/_mixins/_kernel.py
@@ -0,0 +1,51 @@
+from typing import Any, Tuple, Union, TypeVar, Optional
+
+from abc import ABC
+
+from anndata import AnnData
+
+import numpy as np
+from scipy.sparse import spmatrix
+
+KernelExpression = TypeVar("KernelExpression")
+
+
+class KernelMixin(ABC):
+    """Mixin that exposes various properties of :class:`cellrank.kernels.KernelExpression`."""
+
+    def __init__(self, kernel: "KernelExpression", **kwargs: Any):
+        super().__init__(**kwargs)
+        self._kernel = kernel
+        self._n_obs = self.kernel.adata.n_obs
+
+    @property
+    def kernel(self) -> "KernelExpression":
+        """Underlying kernel expression."""
+        return self._kernel
+
+    @property
+    def adata(self) -> AnnData:
+        """Annotated data object."""
+        return self.kernel.adata
+
+    @property
+    def shape(self) -> Tuple[int, int]:
+        """Shape of the kernel."""
+        return self.kernel.shape
+
+    @adata.setter
+    def adata(self, adata: Optional[AnnData]) -> None:
+        self.kernel.adata = adata
+
+    def __len__(self) -> int:
+        return self._n_obs
+
+    @property
+    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
+        """Transition matrix of :attr:`kernel`."""
+        return self.kernel.transition_matrix
+
+    @property
+    def backward(self) -> bool:
+        """Direction of :attr:`kernel`."""
+        return self.kernel.backward
diff --git a/cellrank/tl/_read.py b/cellrank/tl/_read.py
index d42e333cd..ef7bc865e 100644
--- a/cellrank/tl/_read.py
+++ b/cellrank/tl/_read.py
@@ -1,32 +1,25 @@
-"""IO module."""
-
-from typing import Union, Callable, Optional
+from typing import Any, Union, Callable, Optional
 
 from pathlib import Path
 
 from scvelo import read as scv_read
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
 from cellrank.ul._docs import d
+from cellrank.tl._utils import _deprecate
 from cellrank.tl._colors import _create_categorical_colors
 from cellrank.tl._lineage import Lineage
-from cellrank.tl._constants import (
-    Direction,
-    AbsProbKey,
-    DirectionPlot,
-    _colors,
-    _lin_names,
-)
 
 from matplotlib.colors import is_color_like
 
 
 @d.dedent
+@_deprecate(version="2.0")
 def read(
     path: Union[Path, str],
-    key: Optional[str] = None,
     read_callback: Callable = scv_read,
-    **kwargs,
+    **kwargs: Any,
 ) -> AnnData:
     """
     Read file and return :class:`anndata.AnnData` object.
@@ -35,9 +28,6 @@ def read(
     ----------
     path
         Path to the annotated data object.
-    key
-        Key in ``adata.obsm`` where the :class:`cellrank.tl.Lineage` is stored.
-        If `None`, it is determined automatically.
     read_callback
         Function that actually reads the :class:`anndata.AnnData` object, such as
         :func:`scvelo.read` (default) or :func:`scanpy.read`.
@@ -49,40 +39,31 @@ def read(
     %(adata)s
     """
 
-    def maybe_create_lineage(
-        direction: Union[str, Direction], pretty_name: Optional[str] = None
-    ):
-        if isinstance(direction, Direction):
-            lin_key = str(
-                AbsProbKey.FORWARD
-                if direction == Direction.FORWARD
-                else AbsProbKey.BACKWARD
-            )
-        else:
-            lin_key = direction
-
+    def maybe_create_lineage(backward: bool, pretty_name: Optional[str] = None) -> None:
+        lin_key = Key.obsm.abs_probs(backward)
         pretty_name = "" if pretty_name is None else (pretty_name + " ")
-        names_key, colors_key = _lin_names(lin_key), _colors(lin_key)
+        names_key = Key.obs.term_states(backward)
+        colors_key = Key.uns.colors(names_key)
 
         if lin_key in adata.obsm.keys():
             n_cells, n_lineages = adata.obsm[lin_key].shape
             logg.info(f"Creating {pretty_name}`Lineage` from `adata.obsm[{lin_key!r}]`")
 
-            if names_key not in adata.uns.keys():
+            if names_key not in adata.obs:
                 logg.warning(
                     f"    Lineage names not found in `adata.uns[{names_key!r}]`, creating new names"
                 )
                 names = [f"Lineage {i}" for i in range(n_lineages)]
-            elif len(adata.uns[names_key]) != n_lineages:
+            elif len(adata.obs[names_key].cat.categories) != n_lineages:
                 logg.warning(
                     f"    Lineage names are don't have the required length ({n_lineages}), creating new names"
                 )
                 names = [f"Lineage {i}" for i in range(n_lineages)]
             else:
                 logg.info("    Successfully loaded names")
-                names = adata.uns[names_key]
+                names = list(adata.obs[names_key].cat.categories)
 
-            if colors_key not in adata.uns.keys():
+            if colors_key not in adata.uns:
                 logg.warning(
                     f"    Lineage colors not found in `adata.uns[{colors_key!r}]`, creating new colors"
                 )
@@ -111,10 +92,7 @@ def maybe_create_lineage(
 
     adata = read_callback(path, **kwargs)
 
-    if key is None:
-        maybe_create_lineage(Direction.FORWARD, pretty_name=DirectionPlot.FORWARD.s)
-        maybe_create_lineage(Direction.BACKWARD, pretty_name=DirectionPlot.BACKWARD.s)
-    else:
-        maybe_create_lineage(key)
+    maybe_create_lineage(False, pretty_name="forward")
+    maybe_create_lineage(True, pretty_name="backward")
 
     return adata
diff --git a/cellrank/tl/_transition_matrix.py b/cellrank/tl/_transition_matrix.py
index 4d5cf1b61..29f474242 100644
--- a/cellrank/tl/_transition_matrix.py
+++ b/cellrank/tl/_transition_matrix.py
@@ -1,7 +1,7 @@
-"""Transition matrix module."""
-
-from typing import TypeVar, Iterable, Optional
+from typing import Union, Callable, Iterable, Optional
+from typing_extensions import Literal
 
+from anndata import AnnData
 from cellrank import logging as logg
 from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import _deprecate
@@ -10,8 +10,6 @@
 from cellrank.tl.kernels._velocity_kernel import BackwardMode, VelocityMode
 from cellrank.tl.kernels._velocity_schemes import Scheme
 
-AnnData = TypeVar("AnnData")
-
 
 @_deprecate(version="2.0")
 @inject_docs(m=VelocityMode, b=BackwardMode, s=Scheme)  # don't swap the order
@@ -23,9 +21,13 @@ def transition_matrix(
     xkey: str = "Ms",
     conn_key: str = "connectivities",
     gene_subset: Optional[Iterable] = None,
-    mode: str = VelocityMode.DETERMINISTIC.s,
-    backward_mode: str = BackwardMode.TRANSPOSE.s,
-    scheme: str = Scheme.CORRELATION.s,
+    mode: Literal[
+        "deterministic", "stochastic", "sampling", "monte_carlo"
+    ] = VelocityMode.DETERMINISTIC,
+    backward_mode: Literal["transpose", "negate"] = BackwardMode.TRANSPOSE,
+    scheme: Union[
+        Literal["dot_product", "cosine", "correlation"], Callable
+    ] = Scheme.CORRELATION,
     softmax_scale: Optional[float] = None,
     weight_connectivities: float = 0.2,
     density_normalize: bool = True,
diff --git a/cellrank/tl/_utils.py b/cellrank/tl/_utils.py
index c24a1ca77..ecd1cba11 100644
--- a/cellrank/tl/_utils.py
+++ b/cellrank/tl/_utils.py
@@ -1,4 +1,3 @@
-"""Utility functions for CellRank tools."""
 from typing import (
     Any,
     Dict,
@@ -23,6 +22,7 @@
 import scanpy as sc
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import ModeEnum
 from cellrank.ul._docs import d
 from cellrank.ul._utils import _get_neighs, _has_neighs, _get_neighs_params
 from cellrank.tl._colors import (
@@ -30,7 +30,6 @@
     _convert_to_hex_colors,
     _insert_categorical_colors,
 )
-from cellrank.tl._constants import ModeEnum
 from cellrank.ul._parallelize import parallelize
 from cellrank.tl._linear_solver import _solve_lin_system
 from cellrank.tl.kernels._utils import np_std, np_mean, _filter_kwargs
@@ -83,7 +82,7 @@ def __init__(self, adata: AnnData, n: Optional[int] = None, where: str = "obs"):
 
     def _generate_random_keys(self):
         def generator():
-            return f"CELLRANK_RANDOM_COL_{np.random.randint(2 ** 16)}"
+            return f"RNG_COL_{np.random.randint(2 ** 16)}"
 
         where = getattr(self._adata, self._where)
         names, seen = [], set(where.keys())
@@ -103,7 +102,9 @@ def __enter__(self):
     def __exit__(self, exc_type, exc_val, exc_tb):
         for key in self._keys:
             try:
-                self._adata.obs.drop(key, axis="columns", inplace=True)
+                getattr(self._adata, self._where).drop(
+                    key, axis="columns", inplace=True
+                )
             except KeyError:
                 pass
             if self._where == "obs":
@@ -362,6 +363,7 @@ def _perm_test(
 
 
 @d.get_sections(base="correlation_test", sections=["Returns"])
+@d.dedent
 def _correlation_test(
     X: Union[np.ndarray, spmatrix],
     Y: "Lineage",  # noqa: F821
@@ -370,7 +372,7 @@ def _correlation_test(
     confidence_level: float = 0.95,
     n_perms: Optional[int] = None,
     seed: Optional[int] = None,
-    **kwargs,
+    **kwargs: Any,
 ) -> pd.DataFrame:
     """
     Perform a statistical test.
@@ -390,21 +392,20 @@ def _correlation_test(
     confidence_level
         Confidence level for the confidence interval calculation. Must be in `[0, 1]`.
     n_perms
-        Number of permutations if ``method='perm_test'``.
+        Number of permutations if ``method = 'perm_test'``.
     seed
-        Random seed if ``method='perm_test'``.
-    kwargs
-        Keyword arguments for :func:`cellrank.ul._parallelize.parallelize`.
+        Random seed if ``method = 'perm_test'``.
+    %(parallel)s
 
     Returns
     -------
-    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, one 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(
@@ -420,23 +421,17 @@ def _correlation_test(
     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])
+    res = pd.DataFrame(corr, index=gene_names, columns=[f"{c}_corr" for c in Y.names])
     for idx, c in enumerate(Y.names):
-        res[f"{c} pval"] = pvals[:, idx]
-        res[f"{c} qval"] = multipletests(pvals[:, idx], alpha=0.05, method="fdr_bh")[1]
-        res[f"{c} ci low"] = ci_low[:, idx]
-        res[f"{c} ci high"] = ci_high[:, idx]
-
-    res = res[
-        [
-            f"{c} {stat}"
-            for c in Y.names
-            for stat in ("corr", "pval", "qval", "ci low", "ci high")
-        ]
-    ]
-    res.sort_values(by=[f"{c} corr" for c in Y.names], ascending=False, inplace=True)
+        res[f"{c}_pval"] = pvals[:, idx]
+        res[f"{c}_qval"] = multipletests(pvals[:, idx], alpha=0.05, method="fdr_bh")[1]
+        res[f"{c}_ci_low"] = ci_low[:, idx]
+        res[f"{c}_ci_high"] = ci_high[:, idx]
 
-    return res
+    # fmt: off
+    res = res[[f"{c}_{stat}" for c in Y.names for stat in ("corr", "pval", "qval", "ci_low", "ci_high")]]
+    return res.sort_values(by=[f"{c}_corr" for c in Y.names], ascending=False)
+    # fmt: on
 
 
 def _correlation_test_helper(
@@ -553,7 +548,7 @@ def _filter_cells(distances: spmatrix, rc_labels: Series, n_matches_min: int) ->
 
     if not is_categorical_dtype(rc_labels):
         raise TypeError(
-            f"Argument `categories` must be a categorical variable, found `{infer_dtype(rc_labels)}`."
+            f"Expected `categories` be `categorical`, found `{infer_dtype(rc_labels)}`."
         )
 
     # retrieve knn graph
@@ -574,10 +569,10 @@ def _filter_cells(distances: spmatrix, rc_labels: Series, n_matches_min: int) ->
 
     freqs_new = np.array([np.sum(rc_labels == cl) for cl in cls])
 
-    if any(freqs_new / freqs_orig < 0.5):
+    if np.any((freqs_new / freqs_orig) < 0.5):
         logg.warning(
             "Consider lowering  'n_matches_min' or "
-            "increasing 'n_neighbors_filtering'. This filters out too many cells."
+            "increasing 'n_neighbors_filtering'. This filters out too many cells"
         )
 
     return rc_labels
@@ -586,8 +581,8 @@ def _filter_cells(distances: spmatrix, rc_labels: Series, n_matches_min: int) ->
 def _cluster_X(
     X: Union[np.ndarray, spmatrix],
     n_clusters: int,
-    method: Literal["leiden", "kmeans"] = "kmeans",
-    n_neighbors: int = 15,
+    method: Literal["leiden", "kmeans"] = "leiden",
+    n_neighbors: int = 20,
     resolution: float = 1.0,
 ) -> List[Any]:
     """
@@ -596,7 +591,7 @@ def _cluster_X(
     Parameters
     ----------
     X
-        Data matrix of shape `n_samples x n_features`.
+        Matrix of shape ``n_samples x n_features``.
     n_clusters
         Number of clusters to use.
     method
@@ -611,10 +606,11 @@ def _cluster_X(
     :class:`list`
         List of cluster labels of length `n_samples`.
     """
-
-    # make sure data is at least 2D
+    if X.shape[0] == 1:
+        # sc.tl.leiden issue
+        return [0]
     if X.ndim == 1:
-        X = X.reshape((1, -1))
+        X = X[:, None]
 
     if method == "kmeans":
         kmeans = KMeans(n_clusters=n_clusters).fit(X)
@@ -626,7 +622,7 @@ def _cluster_X(
         labels = adata_dummy.obs[method]
     else:
         raise NotImplementedError(
-            f"Invalid method `{method!r}`. Valid options are: `'kmeans'` or `'leiden'`."
+            f"Invalid method `{method}`. Valid options are `kmeans` or `leiden`."
         )
 
     return list(labels)
@@ -1440,10 +1436,10 @@ def _calculate_lineage_absorption_time_means(
     Q: csr_matrix,
     R: csr_matrix,
     trans_indices: np.ndarray,
-    n: int,
     ixs: Dict[str, np.ndarray],
     lineages: Dict[Sequence[str], str],
-    **kwargs,
+    index: pd.Index,
+    **kwargs: Any,
 ) -> pd.DataFrame:
     """
     Calculate the mean time until absorption and optionally its variance for specific lineages or their combinations.
@@ -1473,8 +1469,9 @@ def _calculate_lineage_absorption_time_means(
 
         Uses more efficient implementation if compute the time for all lineages.
     """
+    n = len(index)
+    res = pd.DataFrame(index=index)
 
-    res = pd.DataFrame()
     if len(lineages) == 1 and set(next(iter(lineages.keys()))) == set(ixs.keys()):
         # use faster implementation in this case
         name = ", ".join(ixs.keys())
@@ -1539,8 +1536,7 @@ def _calculate_lineage_absorption_time_means(
             ).squeeze()
             assert np.all(v >= 0), f"Encountered negative variance: `{v[v < 0]}`."
 
-            var = np.empty(n, dtype=np.float64)
-            var[:] = np.inf
+            var = np.full(n, fill_value=np.nan, dtype=np.float64)
             var[ix] = 0
             var[trans_indices] = v
 
diff --git a/cellrank/tl/estimators/__init__.py b/cellrank/tl/estimators/__init__.py
index 5dfadbf0a..ccf65e503 100644
--- a/cellrank/tl/estimators/__init__.py
+++ b/cellrank/tl/estimators/__init__.py
@@ -1,3 +1,2 @@
-from cellrank.tl.estimators._gpcca import GPCCA
-from cellrank.tl.estimators._cflare import CFLARE
 from cellrank.tl.estimators._base_estimator import BaseEstimator
+from cellrank.tl.estimators.terminal_states import GPCCA, CFLARE, TermStatesEstimator
diff --git a/cellrank/tl/estimators/_base_estimator.py b/cellrank/tl/estimators/_base_estimator.py
index 0c71e6235..1ca958dbd 100644
--- a/cellrank/tl/estimators/_base_estimator.py
+++ b/cellrank/tl/estimators/_base_estimator.py
@@ -1,1300 +1,468 @@
-"""Abstract base class for all kernel-holding estimators."""
-from typing import (
-    Any,
-    Dict,
-    Tuple,
-    Union,
-    Mapping,
-    TypeVar,
-    Iterable,
-    Optional,
-    Sequence,
-)
+from typing import Any, Dict, List, Tuple, Union, Mapping, Callable, Optional, Sequence
 from typing_extensions import Literal
 
 from abc import ABC, abstractmethod
-from copy import copy, deepcopy
-from pathlib import Path
-from datetime import datetime
-
-import scanpy as sc
-import scvelo as scv
-from cellrank import logging as logg
-from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import (
-    TestMethod,
-    save_fig,
-    _pairwise,
-    _irreducible,
-    _process_series,
-    _correlation_test,
-    _get_cat_and_null_indices,
-    _merge_categorical_series,
-    _convert_to_categorical_series,
-    _calculate_lineage_absorption_time_means,
-)
-from cellrank.tl._colors import (
-    _map_names_and_colors,
-    _convert_to_hex_colors,
-    _create_categorical_colors,
-)
+from copy import copy as copy_
+from copy import deepcopy
+from inspect import Parameter, signature, getmembers, currentframe
+from contextlib import contextmanager
+
+from anndata import AnnData
+from cellrank._key import Key
+from cellrank.ul._docs import d
+from cellrank.tl._mixins import IOMixin, KernelMixin, AnnDataMixin
+from cellrank.tl.kernels import PrecomputedKernel
 from cellrank.tl._lineage import Lineage
-from cellrank.tl._constants import (
-    Macro,
-    DirPrefix,
-    AbsProbKey,
-    PrettyEnum,
-    TermStatesKey,
-    _pd,
-    _probs,
-    _colors,
-    _lin_names,
-)
-from cellrank.tl._linear_solver import _solve_lin_system
-from cellrank.tl.estimators._property import Partitioner, LineageEstimatorMixin
-from cellrank.tl.kernels._base_kernel import KernelExpression
-from cellrank.tl.estimators._constants import A, P
+from cellrank.tl.kernels._base_kernel import Kernel, KernelExpression
 
 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
+from scipy.sparse import spmatrix, csr_matrix
 
-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"
-    "set them manually as `.set_terminal_states()`."
+Attr_t = (
+    Literal["X", "raw", "layers", "obs", "var", "obsm", "varm", "obsp", "varp", "uns"],
 )
 
 
 @d.get_sections(base="base_estimator", sections=["Parameters"])
-@d.dedent
-class BaseEstimator(LineageEstimatorMixin, Partitioner, ABC):
+class BaseEstimator(IOMixin, KernelMixin, AnnDataMixin, ABC):
     """
     Base class for all estimators.
 
     Parameters
     ----------
     obj
-        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 :attr:`adata` object inplace or make a copy.
-    read_from_adata
-        Whether to read available attributes in :attr:`adata`, if present.
+        Can be one of the following:
+
+            - :class:`cellrank.tl.kernels.Kernel` - kernel object.
+            - :class:`anndata.AnnData` - annotated data object containing transition matrix in
+              :attr:`anndata.AnnData.obsp`.
+            - :class:`numpy.ndarray` - row-normalized sparse transition matrix.
+            - :class:`scipy.sparse.spmatrix` - row-normalized sparse transition matrix.
     obsp_key
-        Key in ``obj.obsp`` when ``obj`` is an :class:`anndata.AnnData` object.
-    g2m_key
-        Key in :attr:`adata` ``.obs``. Can be used to detect cell-cycle driven start- or endpoints.
-    s_key
-        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 :attr:`adata` ``.obsp`` and the parameters to
-        :attr:`adata` ``.uns``.
-    %(write_to_adata.parameters)s
-        Only used when ``write_to_adata=True``.
+        Key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
+        Only used when ``obj`` is an :class:`anndata.AnnData` object.
     """
 
     def __init__(
         self,
-        obj: Union[KernelExpression, AnnData, spmatrix, np.ndarray],
-        inplace: bool = True,
-        read_from_adata: bool = False,
+        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
         obsp_key: Optional[str] = None,
-        g2m_key: Optional[str] = "G2M_score",
-        s_key: Optional[str] = "S_score",
-        write_to_adata: bool = True,
-        key: Optional[str] = None,
     ):
-        from anndata import AnnData
-
-        super().__init__(obj, obsp_key=obsp_key, key=key, write_to_adata=write_to_adata)
-
-        if isinstance(obj, (KernelExpression, AnnData)) and not inplace:
-            self.kernel._adata = self.adata.copy()
-
-        if self.kernel.backward:
-            self._term_key = TermStatesKey.BACKWARD.s
-            self._abs_prob_key = AbsProbKey.BACKWARD.s
-            self._term_abs_prob_key = Macro.BACKWARD.s
-        else:
-            self._term_key = TermStatesKey.FORWARD.s
-            self._abs_prob_key = AbsProbKey.FORWARD.s
-            self._term_abs_prob_key = Macro.FORWARD.s
-
-        self._key_added = key
-        self._g2m_key = g2m_key
-        self._s_key = s_key
-
-        self._G2M_score = None
-        self._S_score = None
-
-        self._absorption_time_mean = None
-        self._absorption_time_var = None
-
-        if read_from_adata:
-            self._read_from_adata()
-
-    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):
+        if isinstance(obj, Kernel):
+            if obj._transition_matrix is None:
                 raise ValueError(
-                    f"Expected `adata.uns[{_colors(self._term_key)}]` to have "
-                    f"`{len(term_states.cat.categories)}` colors, found `{len(colors)}`."
+                    "Compute transition matrix first as `.compute_transition_matrix()`."
                 )
-
-        self._set_or_debug(f"eig_{self._direction}", self.adata.uns, "_eig")
-
-        self._set_or_debug(self._g2m_key, self.adata.obs, "_G2M_score")
-        self._set_or_debug(self._s_key, self.adata.obs, "_S_score")
-
-        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)
-        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):
+            kernel = obj
+        elif isinstance(obj, KernelExpression):
+            # this will fail if not all kernels have transition matrix computed
+            kernel = obj.compute_transition_matrix()
+        elif isinstance(obj, (np.ndarray, spmatrix)):
+            kernel = PrecomputedKernel(obj)
+        elif isinstance(obj, AnnData):
+            if obsp_key is None:
                 raise ValueError(
-                    f"Expected names and colors to be of same length, found `{names}` and `{colors}`."
+                    "Specify `obsp_key=...` when supplying an `AnnData` object."
                 )
-
-            if len(names) != probs.shape[1]:
-                if isinstance(probs, Lineage):
-                    names = probs.names
-                else:
-                    logg.warning(
-                        f"Expected lineage names to be of length `{probs.shape[1]}`, found `{len(names)}`. "
-                        f"Creating new names"
-                    )
-                    names = [f"Lineage {i}" for i in range(probs.shape[1])]
-            if len(colors) != probs.shape[1] or not all(
-                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(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
-
-    @d.dedent
-    @inject_docs(fs=P.TERM.s, fsp=P.TERM_PROBS.s)
-    def set_terminal_states(
-        self,
-        labels: Union[Series, Dict[str, Sequence[Any]]],
-        cluster_key: Optional[str] = None,
-        en_cutoff: Optional[float] = None,
-        p_thresh: Optional[float] = None,
-        add_to_existing: bool = False,
-        **kwargs,
-    ) -> None:
-        """
-        Manually define terminal states.
-
-        Parameters
-        ----------
-        labels
-            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
-            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 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:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-        """
-
-        self._set_categorical_labels(
-            attr_key=A.TERM.v,
-            color_key=A.TERM_COLORS.v,
-            pretty_attr_key=P.TERM.v,
-            add_to_existing_error_msg=_COMP_TERM_STATES_MSG,
-            categories=labels,
-            cluster_key=cluster_key,
-            en_cutoff=en_cutoff,
-            p_thresh=p_thresh,
-            add_to_existing=add_to_existing,
-        )
-        self._write_terminal_states(time=kwargs.get("time", None))
-
-    @inject_docs(ts=P.TERM.s)
-    def rename_terminal_states(
-        self, new_names: Mapping[str, str], update_adata: bool = True
-    ) -> None:
-        """
-        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 :attr:`adata` object as well or not.
-
-        Returns
-        -------
-        None
-            Nothing, just updates the names of :attr:`{ts}`.
-        """
-
-        term_states = self._get(P.TERM)
-
-        if term_states is None:
-            raise RuntimeError(_COMP_TERM_STATES_MSG)
-
-        if not isinstance(new_names, Mapping):
-            raise TypeError(f"Expected a `Mapping` type, found `{type(new_names)!r}`.")
-        if not len(new_names):
-            return
-
-        new_names = {k: str(v) for k, v in new_names.items()}
-
-        mask = np.isin(list(new_names.keys()), term_states.cat.categories)
-        if not np.all(mask):
-            raise ValueError(
-                f"Invalid old terminal states names: `{np.array(list(new_names.keys()))[~mask]}`."
-            )
-
-        names_after_renaming = [new_names.get(n, n) for n in term_states.cat.categories]
-        if len(set(names_after_renaming)) != len(term_states.cat.categories):
-            raise ValueError(
-                f"After renaming, the names will not be unique: `{names_after_renaming}`."
+            elif obsp_key not in obj.obsp:
+                raise KeyError(
+                    f"Unable to find transition matrix in `adata.obsp[{obsp_key!r}]`."
+                )
+            kernel = PrecomputedKernel(obsp_key, adata=obj)
+        else:
+            raise TypeError(
+                f"Expected an object of type `KernelExpression`, `numpy.ndarray`, `scipy.sparse.spmatrix` "
+                f"or `anndata.AnnData`, got `{type(obj).__name__}`."
             )
 
-        self._set(A.TERM, term_states.cat.rename_categories(new_names))
+        super().__init__(kernel=kernel)
 
-        memberships = (
-            self._get(A.TERM_ABS_PROBS) if hasattr(self, A.TERM_ABS_PROBS.s) else None
+        self._params: Dict[str, Any] = {}
+        self._shadow_adata = AnnData(
+            X=csr_matrix(self.adata.shape, dtype=self.adata.X.dtype),
+            obs=self.adata.obs[[]].copy(),
+            var=self.adata.var[[]].copy(),
+            raw=None if self.adata.raw is None else self.adata.raw.to_adata(),
         )
-        if memberships is not None:  # GPCCA
-            memberships.names = [new_names.get(n, n) for n in memberships.names]
-            self._set(A.TERM_ABS_PROBS, memberships)
 
-        # we can be just computing it and it's not yet saved in adata
-        if (
-            update_adata
-            and self._term_key in self.adata.obs
-            and _lin_names(self._term_key) in self.adata.uns
-        ):
-            # fmt: off
-            self.adata.obs[self._term_key] = self.adata.obs[self._term_key].cat.rename_categories(new_names)
-            self.adata.uns[_lin_names(self._term_key)] = np.array(self.adata.obs[self._term_key].cat.categories)
-            # fmt: on
+    def __init_subclass__(cls, **kwargs: Any):
+        super().__init_subclass__()
 
-    @inject_docs(
-        abs_prob=P.ABS_PROBS,
-        fs=P.TERM.s,
-        lat=P.LIN_ABS_TIMES,
-    )
-    def compute_absorption_probabilities(
+    def _set(
         self,
-        keys: Optional[Sequence[str]] = None,
-        check_irreducibility: bool = False,
-        solver: str = "gmres",
-        use_petsc: bool = True,
-        time_to_absorption: Optional[
-            Union[
-                str,
-                Sequence[Union[str, Sequence[str]]],
-                Dict[Union[str, Sequence[str]], str],
-            ]
+        attr: Optional[str] = None,
+        obj: Optional[Union[pd.DataFrame, Mapping[str, Any]]] = None,
+        key: Optional[str] = None,
+        value: Optional[
+            Union[np.ndarray, pd.Series, pd.DataFrame, Lineage, AnnData, Dict[str, Any]]
         ] = None,
-        n_jobs: Optional[int] = None,
-        backend: str = "loky",
-        show_progress_bar: bool = True,
-        tol: float = 1e-6,
-        preconditioner: Optional[str] = None,
+        copy: bool = True,
+        shadow_only: bool = False,
     ) -> None:
         """
-        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 :attr:`{fs}`.
+        Set an attribute and optionally update ``obj[{key}]``.
 
         Parameters
         ----------
-        keys
-            Keys defining the recurrent classes.
-        check_irreducibility:
-            Check whether the transition matrix is irreducible.
-        solver
-            Solver to use for the linear problem. Options are `'direct', 'gmres', 'lgmres', 'bicgstab' or 'gcrotmk'`
-            when ``use_petsc=False`` or one of :class:`petsc4py.PETSc.KPS.Type` otherwise.
-
-            Information on the :mod:`scipy` iterative solvers can be found in :func:`scipy.sparse.linalg` or for
-            :mod:`petsc4py` solver `here <https://petsc.org/release/overview/linear_solve_table/>`__.
-        use_petsc
-            Whether to use solvers from :mod:`petsc4py` or :mod:`scipy`. Recommended for large problems.
-            If no installation is found, defaults to :func:`scipy.sparse.linalg.gmres`.
-        time_to_absorption
-            Whether to compute mean time to absorption and its variance to specific absorbing states.
-
-            If a :class:`dict`, can be specified as ``{{'Alpha': 'var', ...}}`` to also compute variance.
-            In case when states are a :class:`tuple`, time to absorption will be computed to the subset of these states,
-            such as ``[('Alpha', 'Beta'), ...]`` or ``{{('Alpha', 'Beta'): 'mean', ...}}``.
-            Can be specified as ``'all'`` to compute it to any absorbing state in ``keys``, which is more efficient
-            than listing all absorbing states.
-
-            It might be beneficial to disable the progress bar as ``show_progress_bar=False``, because many linear
-            systems are being solved.
-        n_jobs
-            Number of parallel jobs to use when using an iterative solver. When ``use_petsc=True`` or for
-            quickly-solvable problems, we recommend higher number (>=8) of jobs in order to fully saturate the cores.
-        backend
-            Which backend to use for multiprocessing. See :class:`joblib.Parallel` for valid options.
-        show_progress_bar
-            Whether to show progress bar when the solver isn't a direct one.
-        tol
-            Convergence tolerance for the iterative solver. The default is fine for most cases, only consider
-            decreasing this for severely ill-conditioned matrices.
-        preconditioner
-            Preconditioner to use, only available when ``use_petsc=True``. For available values, see
-            `here <https://petsc.org/release/docs/manual/ksp/?highlight=pctype#preconditioners>`__.
-            We recommended the `'ilu'` preconditioner for badly conditioned problems.
+        attr
+            Attribute to set. Only updated when we're not in the shadow. If `None`, don't update anything.
+            See :attr:`_in_shadow` and ``obj`` for more information.
+        obj
+            Object which to update with ``value`` alongside the ``attr`.
+            Usually, an attribute of :attr:`adata` is passed here.
+        key
+            Key in ``obj`` to update with ``value``. Only used when ``obj != None``.
+        value
+            Value to set. If `None` and ``key != None``, it removes the values under ``obj[key]``, if present.
+        copy
+            Whether to copy the ``value`` before setting it in ``obj``.
+        shadow_only
+            Whether or not to update the ``obj`` if we are not in the shadow.
 
         Returns
         -------
-        None
-            Nothing, but updates the following fields:
+        Nothing, just optionally updates ``attr`` and/or ``obj[{key}]``.
 
-                - :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``.
+        Raises
+        ------
+        AttributeError
+            If ``attr`` doesn't exist.
         """
+        if not self._in_shadow:
+            if attr is not None:
+                if not hasattr(self, attr):
+                    raise AttributeError(attr)
+                setattr(self, attr, value)
+            if shadow_only:
+                return
 
-        if self._get(P.TERM) is None:
-            raise RuntimeError(_COMP_TERM_STATES_MSG)
-        if keys is not None:
-            keys = sorted(set(keys))
-
-        start = logg.info("Computing absorption probabilities")
-
-        # get the transition matrix
-        t = self.transition_matrix
-        if not self.issparse:
-            logg.warning(
-                "Attempting to solve a potentially large linear system with dense transition matrix"
-            )
-
-        # process the current annotations according to `keys`
-        terminal_states_, colors_ = _process_series(
-            series=self._get(P.TERM), keys=keys, colors=self._get(A.TERM_COLORS)
-        )
-        # warn in case only one state is left
-        keys = list(terminal_states_.cat.categories)
-        if len(keys) == 1:
-            logg.warning(
-                "There is only 1 recurrent class, all cells will have probability 1 of going there"
-            )
-
-        lin_abs_times = {}
-        if time_to_absorption is not None:
-            if isinstance(time_to_absorption, (str, tuple)):
-                time_to_absorption = [time_to_absorption]
-            if not isinstance(time_to_absorption, dict):
-                time_to_absorption = {ln: "mean" for ln in time_to_absorption}
-            for ln, moment in time_to_absorption.items():
-                if moment not in ("mean", "var"):
-                    raise ValueError(
-                        f"Moment must be either `'mean'` or `'var'`, found `{moment!r}` for `{ln!r}`."
-                    )
-
-                seen = set()
-                if isinstance(ln, str):
-                    ln = tuple(keys) if ln == "all" else (ln,)
-                sorted_ln = tuple(sorted(ln))  # preserve the user order
-
-                if sorted_ln not in seen:
-                    seen.add(sorted_ln)
-                    for lin in ln:
-                        if lin not in keys:
-                            raise ValueError(
-                                f"Invalid absorbing state `{lin!r}` in `{ln}`. "
-                                f"Valid options are `{list(terminal_states_.cat.categories)}`."
-                            )
-                    lin_abs_times[tuple(ln)] = moment
-
-            # define the dimensions of this problem
-        n_cells = t.shape[0]
-        n_macrostates = len(terminal_states_.cat.categories)
-
-        # get indices corresponding to recurrent and transient states
-        rec_indices, trans_indices, lookup_dict = _get_cat_and_null_indices(
-            terminal_states_
-        )
-        if not len(trans_indices):
-            raise RuntimeError("Cannot proceed - Markov chain is irreducible.")
-
-        # create Q (restriction transient-transient), S (restriction transient-recurrent)
-        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,
-        )
+        if obj is None:
+            return
 
-        # check for irreducibility
-        if check_irreducibility:
-            if self.is_irreducible is None:
-                self._is_irreducible = _irreducible(self.transition_matrix)
+        if key is not None:
+            if value is None:
+                try:
+                    del obj[key]
+                except KeyError:
+                    pass
             else:
-                if not self.is_irreducible:
-                    logg.warning("Transition matrix is not irreducible")
-                else:
-                    logg.debug("Transition matrix is irreducible")
-
-        logg.debug(f"Found `{n_cells}` cells and `{s.shape[1]}` absorbing states")
-
-        # solve the linear system of equations
-        _abs_classes = _solve_lin_system(
-            q,
-            s,
-            solver=solver,
-            use_petsc=use_petsc,
-            n_jobs=n_jobs,
-            backend=backend,
-            tol=tol,
-            use_eye=True,
-            show_progress_bar=show_progress_bar,
-            preconditioner=preconditioner,
-        )
-
-        if time_to_absorption is not None:
-            abs_time_means = _calculate_lineage_absorption_time_means(
-                q,
-                t[trans_indices, :][:, rec_indices],
-                trans_indices,
-                n=t.shape[0],
-                ixs=lookup_dict,
-                lineages=lin_abs_times,
-                solver=solver,
-                use_petsc=use_petsc,
-                n_jobs=n_jobs,
-                backend=backend,
-                tol=tol,
-                show_progress_bar=show_progress_bar,
-                preconditioner=preconditioner,
-            )
-            abs_time_means.index = self.adata.obs_names
-        else:
-            abs_time_means = None
-
-        # for recurrent states, set their self-absorption probability to one
-        abs_classes = np.zeros((len(self), n_macrostates))
-        rec_classes_full = {
-            cl: np.where(terminal_states_ == cl)[0]
-            for cl in terminal_states_.cat.categories
-        }
-        for col, cl_indices in enumerate(rec_classes_full.values()):
-            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(
-                abs_classes,
-                names=terminal_states_.cat.categories,
-                colors=colors_,
-            ),
-        )
-
-        extra_msg = ""
-        if abs_time_means is not None:
-            self._set(A.LIN_ABS_TIMES, abs_time_means)
-            extra_msg = f"       `.{P.LIN_ABS_TIMES}`\n"
+                obj[key] = copy_(value) if copy else value
 
-        self._write_absorption_probabilities(time=start, extra_msg=extra_msg)
-
-    @d.dedent
-    def compute_lineage_priming(
+    def _get(
         self,
-        method: Literal["kl_divergence", "entropy"] = "kl_divergence",
-        early_cells: Optional[Union[Mapping[str, Sequence[str]], Sequence[str]]] = None,
-    ) -> pd.Series:
+        attr: str,
+        obj: Union[pd.DataFrame, Mapping[str, Any]],
+        key: str,
+        where: Optional[Literal["obs", "obsm", "var", "varm", "uns"]] = None,
+        dtype: Optional[Union[type, Tuple[type, ...]]] = None,
+        copy: bool = True,
+        allow_missing: bool = False,
+    ) -> None:
         """
-        %(lin_pd.full_desc)s
+        Get data from an object and set an attribute.
 
         Parameters
         ----------
-        %(lin_pd.parameters)s
-            Cell ids or a mask marking early cells. If `None`, use all cells. Only used when ``method='kl_divergence'``.
-            If a :class:`dict`, the key species a cluster key in :attr:`anndata.AnnData.obs` and the values
-            specify cluster labels containing early cells.
+        attr
+            Attribute to set.
+        obj
+            Object from which to extract the data.
+        key
+            Key in ``obj`` where the data is stored.
+        where
+            Attribute of :attr:`_shadow_adata` where to save the extracted data. If `None`, don't update it.
+        dtype
+            Valid type(s) of the extracted data.
+        copy
+            Copy the data before setting the ``attr``.
+        allow_missing
+            Whether or not to allow ``key`` to be missing in ``obj``.
 
         Returns
         -------
-        %(lin_pd.returns)s
-        """  # noqa: D400
-        abs_probs: Optional[Lineage] = self._get(P.ABS_PROBS)
-        if abs_probs is None:
-            raise RuntimeError(
-                "Compute absorption probabilities first as `.compute_absorption_probabilities()`."
-            )
-        if isinstance(early_cells, dict):
-            if len(early_cells) != 1:
-                raise ValueError(
-                    f"Expected a dictionary with only 1 key, found `{len(early_cells)}`."
-                )
-            key = next(iter(early_cells.keys()))
-            if key not in self.adata.obs:
-                raise KeyError(f"Unable to find clustering in `adata.obs[{key!r}]`.")
-            early_cells = self.adata.obs[key].isin(early_cells[key])
-        elif early_cells is not None:
-            early_cells = np.asarray(early_cells)
-            if not np.issubdtype(early_cells.dtype, np.bool_):
-                early_cells = np.isin(self.adata.obs_names, early_cells)
-
-        values = pd.Series(
-            abs_probs.priming_degree(method, early_cells), index=self.adata.obs_names
-        )
-
-        self._set(A.PRIME_DEG, values)
-        self.adata.obs[_pd(self._abs_prob_key)] = values
-
-        logg.info(
-            f"Adding `adata.obs[{_pd(self._abs_prob_key)!r}]`\n"
-            f"       `.{P.PRIME_DEG}`"
-        )
+        Nothing, just updates ``attr`` with the extracted values and optionally :attr:`_shadow_adata`.
+
+        Raises
+        ------
+        AttributeError
+            If ``attr`` doesn't exist.
+        TypeError
+            If ``dtype != None`` and the extracted values are not instances of ``dtype``.
+        KeyError
+            If ``allow_missing = False`` and ``key`` was not found in ``obj``.
+        """
+        if not hasattr(self, attr):
+            raise AttributeError(attr)
 
-        return values
+        try:
+            data = obj[key]
+            if dtype is not None and not isinstance(data, dtype):
+                raise TypeError(
+                    f"Expected `.{attr}` to be of type `{dtype}`, found `{type(data).__name__}`."
+                )
+            if copy:
+                data = copy_(data)
+            setattr(self, attr, data)
+            if where is not None:
+                getattr(self._shadow_adata, where)[key] = data
+        except KeyError:
+            if not allow_missing:
+                raise
+
+    @property
+    @contextmanager
+    def _shadow(self) -> None:
+        """
+        Temporarily set :attr:`adata` to :attr:`_shadow_adata`.
 
-    @d.get_sections(
-        base="lineage_drivers", sections=["Parameters", "Returns", "References"]
-    )
-    @d.get_full_description(base="lineage_drivers")
-    @d.dedent
-    @inject_docs(lin_drivers=P.LIN_DRIVERS, tm=TestMethod)
-    def compute_lineage_drivers(
+        Used to construct the serialization object in :meth:`to_adata`.
+        """
+        if self._in_shadow:
+            yield
+        else:
+            adata = self.adata
+            try:
+                self.adata = self._shadow_adata
+                yield
+            finally:
+                self.adata = adata
+
+    @property
+    def _in_shadow(self) -> bool:
+        """Return `True` if :attr:`adata` is :attr:`_shadow_adata`."""
+        return self.adata is self._shadow_adata
+
+    def _create_params(
         self,
-        lineages: Optional[Union[str, Sequence]] = None,
-        method: str = TestMethod.FISCHER.s,
-        cluster_key: Optional[str] = None,
-        clusters: Optional[Union[str, Sequence]] = None,
-        layer: str = "X",
-        use_raw: bool = False,
-        confidence_level: float = 0.95,
-        n_perms: int = 1000,
-        seed: Optional[int] = None,
-        return_drivers: bool = True,
-        **kwargs,
-    ) -> Optional[pd.DataFrame]:
+        locs: Optional[Mapping[str, Any]] = None,
+        func: Optional[Callable] = None,
+        remove: Sequence[str] = (),
+    ) -> Dict[str, Any]:
         """
-        Compute driver genes per lineage.
-
-        Correlates gene expression with lineage probabilities, for a given lineage and set of clusters.
-        Often, it makes sense to restrict this to a set of clusters which are relevant for the specified lineages.
+        Create parameters of interest from a function call.
 
         Parameters
         ----------
-        lineages
-            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. Valid options are:
-
-                - {tm.FISCHER.s!r} - use Fischer transformation :cite:`fischer:21`.
-                - {tm.PERM_TEST.s!r} - use permutation test.
-        cluster_key
-            Key from :attr:`adata` ``.obs`` to obtain cluster annotations. These are considered for ``clusters``.
-        clusters
-            Restrict the correlations to these clusters.
-        layer
-            Key from :attr:`adata` ``.layers``.
-        use_raw
-            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]`.
-        n_perms
-            Number of permutations to use when ``method={tm.PERM_TEST.s!r}``.
-        seed
-            Random seed when ``method={tm.PERM_TEST.s!r}``.
-        return_drivers
-            Whether to return the drivers. This also contains the lower and upper ``confidence_level`` confidence
-            interval bounds.
-        %(parallel)s
+        locs
+            Environment from which to get the parameters. If `None`, get the caller's environment.
+        func
+            Function of interest. If `None`, use the caller.
+        remove
+            Keys in ``locs`` which should not be included in the result.
 
         Returns
         -------
-        %(correlation_test.returns)s
-        Only if ``return_drivers=True``.
-
-        Otherwise, updates :attr:`adata` ``.var`` or :attr:`adata` ``.raw.var``, depending ``use_raw`` with:
+        The parameters as a :class:`dict`.
 
-            - ``'{{direction}} {{lineage}} corr'`` - the potential lineage drivers.
-            - ``'{{direction}} {{lineage}} qval'`` - the corrected p-values.
-
-        Also updates the following fields:
-
-            - :attr:`{lin_drivers}` - same as the returned values.
+        Notes
+        -----
+        *args/**kwargs are always ignored and the values in ``locs`` are not copied.
         """
-
-        # check that lineage probs have been computed
-        method = TestMethod(method)
-        abs_probs = self._get(P.ABS_PROBS)
-        prefix = DirPrefix.BACKWARD if self.kernel.backward else DirPrefix.FORWARD
-
-        if abs_probs is None:
-            raise RuntimeError(
-                "Compute absorption probabilities first as `.compute_absorption_probabilities()`."
-            )
-
-        if abs_probs.shape[1] == 1:
-            logg.warning(
-                "There is only 1 lineage present. Using stationary distribution instead"
-            )
-            stat_dist = self._get(P.EIG).get("stationary_dist", None)
-            if stat_dist is None:
-                raise RuntimeError(
-                    f"No stationary distribution found in `.{P.EIG.s}['stationary_dist']`."
+        frame = currentframe()
+        try:
+            if locs is None:
+                locs = frame.f_back.f_locals
+            if func is None or True:
+                name = frame.f_back.f_code.co_name
+                func = dict(getmembers(self)).get(name, None)
+            if not callable(func):
+                raise TypeError(
+                    f"Expected `func` to be `callable`, found `{type(func).__name__}`."
                 )
-            abs_probs = Lineage(
-                stat_dist,
-                names=abs_probs.names,
-                colors=abs_probs.colors,
-            )
-
-        # check all lin_keys exist in self.lin_names
-        if isinstance(lineages, str):
-            lineages = [lineages]
-        if lineages is not None:
-            _ = abs_probs[lineages]
-        else:
-            lineages = abs_probs.names
-
-        if not len(lineages):
-            raise ValueError("No lineages have been selected.")
-
-        # use `cluster_key` and clusters to subset the data
-        if clusters is not None:
-            if cluster_key not in self.adata.obs.keys():
-                raise KeyError(f"Key `{cluster_key!r}` not found in `adata.obs`.")
-            if isinstance(clusters, str):
-                clusters = [clusters]
 
-            all_clusters = np.array(self.adata.obs[cluster_key].cat.categories)
-            cluster_mask = np.array([name not in all_clusters for name in clusters])
-
-            if any(cluster_mask):
-                raise KeyError(
-                    f"Clusters `{list(np.array(clusters)[cluster_mask])}` not found in "
-                    f"`adata.obs[{cluster_key!r}]`."
-                )
-            subset_mask = np.in1d(self.adata.obs[cluster_key], clusters)
-            adata_comp = self.adata[subset_mask]
-            lin_probs = abs_probs[subset_mask, :]
-        else:
-            adata_comp = self.adata
-            lin_probs = abs_probs
-
-        # check that the layer exists, and that use raw is only used with layer X
-        if layer != "X":
-            if layer not in self.adata.layers:
-                raise KeyError(f"Layer `{layer!r}` not found in `adata.layers`.")
-            if use_raw:
-                raise ValueError("For `use_raw=True`, layer must be 'X'.")
-            data = adata_comp.layers[layer]
-            var_names = adata_comp.var_names
-        else:
-            if use_raw and self.adata.raw is None:
-                logg.warning("No raw attribute set. Using `.X` instead")
-                use_raw = False
-            data = adata_comp.raw.X if use_raw else adata_comp.X
-            var_names = adata_comp.raw.var_names if use_raw else adata_comp.var_names
-
-        start = logg.info(
-            f"Computing correlations for lineages `{lineages}` restricted to clusters `{clusters}` in "
-            f"layer `{layer}` with `use_raw={use_raw}`"
-        )
-
-        lin_probs = lin_probs[lineages]
-        drivers = _correlation_test(
-            data,
-            lin_probs,
-            gene_names=var_names,
-            method=method,
-            n_perms=n_perms,
-            seed=seed,
-            confidence_level=confidence_level,
-            **kwargs,
-        )
-        self._set(A.LIN_DRIVERS, drivers)
-
-        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]
-            self.adata.raw.var[[f"{prefix} {col}" for col in qvals]] = drivers[qvals]
-        else:
-            self.adata.var[[f"{prefix} {col}" for col in corrs]] = drivers[corrs]
-            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 lin_probs.names
-        ]
-
-        logg.info(
-            f"Adding `.{P.LIN_DRIVERS}`\n       "
-            + "\n       ".join(keys_added)
-            + "\n    Finish",
-            time=start,
-        )
+            params = {}
+            for name, param in signature(func).parameters.items():
+                if param.kind in (Parameter.VAR_POSITIONAL, Parameter.VAR_KEYWORD):
+                    continue
+                if name in remove:
+                    continue
+                if name in locs:
+                    params[name] = locs[name]
+
+            return params
+        except AttributeError:
+            # frame can be None
+            return {}
+        except TypeError:
+            return {}
+        finally:
+            del frame
+
+    def _read_params(self, key: str) -> Dict[str, Any]:
+        """
+        Read ``key`` from estimator params in :attr:`adata`.
 
-        if return_drivers:
-            return drivers
+        Usually called in :meth:`_read_adata` during :meth:`from_adata`.
+        """
+        ekey = Key.uns.estimator(self.backward) + "_params"
+        return dict(self.adata.uns.get(ekey, {}).get(key, {}))
 
-    @d.get_sections(base="plot_lineage_drivers", sections=["Parameters"])
     @d.dedent
-    def plot_lineage_drivers(
+    def to_adata(
         self,
-        lineage: str,
-        n_genes: int = 8,
-        ncols: Optional[int] = None,
-        use_raw: bool = False,
-        title_fmt: str = "{gene} qval={qval:.4e}",
-        figsize: Optional[Tuple[float, float]] = None,
-        dpi: Optional[int] = None,
-        save: Optional[Union[str, Path]] = None,
-        **kwargs,
-    ) -> None:
+        keep: Union[Literal["all"], Sequence[Attr_t]] = ("X", "raw"),
+        *,
+        copy: Union[bool, Sequence[Attr_t]] = True,
+    ) -> AnnData:
         """
-        Plot lineage drivers discovered by :meth:`compute_lineage_drivers`.
+        %(to_adata.full_desc)s
 
         Parameters
         ----------
-        lineage
-            Lineage for which to plot the driver genes.
-        n_genes
-            Top most correlated genes to plot.
-        ncols
-            Number of columns.
-        use_raw
-            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.
-        %(plotting)s
-        kwargs
-            Keyword arguments for :func:`scvelo.pl.scatter`.
+        keep
+            Which attributes to keep from the underlying :attr:`adata`. Valid options are:
+
+                - `'all'` - keep all attributes specified in the signature.
+                - :class:`typing.Sequence` - keep only subset of these attributes.
+                - :class:`dict` - the keys correspond the attribute names and values to a subset of keys
+                  which to keep from this attribute. If the values are specified either as `True` or `'all'`,
+                  everything from this attribute will be kept.
+        copy
+            Whether to copy the data. Can be specified on per-attribute basis. Useful for attributes that store arrays.
+            Attributes not specified here will not be copied.
 
         Returns
         -------
-        %(just_plots)s
-        """
-
-        def prepare_format(
-            gene: str,
-            *,
-            pval: Optional[float],
-            qval: Optional[float],
-            corr: Optional[float],
-        ) -> Dict[str, Any]:
-            kwargs = {}
-            if "{gene" in title_fmt:
-                kwargs["gene"] = gene
-            if "{pval" in title_fmt:
-                kwargs["pval"] = float(pval) if pval is not None else np.nan
-            if "{qval" in title_fmt:
-                kwargs["qval"] = float(qval) if qval is not None else np.nan
-            if "{corr" in title_fmt:
-                kwargs["corr"] = float(corr) if corr is not None else np.nan
-
-            return kwargs
-
-        lin_drivers = self._get(P.LIN_DRIVERS)
-        if lin_drivers is None:
-            raise RuntimeError(
-                f"Compute `.{P.LIN_DRIVERS}` first as `.compute_lineage_drivers()`."
-            )
-
-        key = f"{lineage} corr"
-        if key not in lin_drivers:
-            raise KeyError(
-                f"Lineage `{key!r}` not found in `{list(lin_drivers.columns)}`."
-            )
-
-        if n_genes <= 0:
-            raise ValueError(f"Expected `n_genes` to be positive, found `{n_genes}`.")
-
-        kwargs.pop("save", None)
-        genes = lin_drivers.sort_values(by=key, ascending=False).head(n_genes)
-
-        ncols = 4 if ncols is None else ncols
-        nrows = int(np.ceil(len(genes) / ncols))
-
-        fig, axes = plt.subplots(
-            ncols=ncols,
-            nrows=nrows,
-            dpi=dpi,
-            figsize=(ncols * 6, nrows * 4) if figsize is None else figsize,
-        )
-        axes = np.ravel([axes])
+        %(adata)s
+        """  # noqa: D400
 
-        _i = 0
-        for _i, (gene, ax) in enumerate(zip(genes.index, axes)):
-            data = genes.loc[gene]
-            scv.pl.scatter(
-                self.adata,
-                color=gene,
-                ncols=ncols,
-                use_raw=use_raw,
-                ax=ax,
-                show=False,
-                title=title_fmt.format(
-                    **prepare_format(
-                        gene,
-                        pval=data.get(f"{lineage} pval", None),
-                        qval=data.get(f"{lineage} qval", None),
-                        corr=data.get(f"{lineage} corr", None),
+        def handle_attribute(attr: Attr_t, keys: List[str], *, copy: bool) -> None:
+            try:
+                if attr == "X":
+                    adata.X = deepcopy(self.adata.X) if copy else self.adata.X
+                    return
+                if attr == "raw":
+                    adata.raw = (
+                        self.adata.raw.to_adata()
+                        if self.adata.raw is not None
+                        else None
                     )
-                ),
-                **kwargs,
-            )
-
-        for j in range(_i + 1, len(axes)):
-            axes[j].remove()
-
-        if save is not None:
-            save_fig(fig, save)
+                    return
 
+                old = getattr(self.adata, attr)
+                new = getattr(adata, attr)
+                if keys == ["all"]:
+                    keys = list(old.keys())
+
+                # fmt: off
+                if isinstance(new, pd.DataFrame):
+                    old = old[keys]
+                    # avoid duplicates
+                    old = old[old.columns.difference(new.columns)]
+                    setattr(adata, attr, pd.merge(new, old, how="inner", left_index=True, right_index=True, copy=copy))
+                elif isinstance(new, Mapping):
+                    old = {k: old[k] for k in keys}
+                    # old has preference, since it's user supplied
+                    setattr(adata, attr, {**new, **(deepcopy(old) if copy else old)})
+                else:
+                    raise TypeError(f"Expected `adata.{attr}` to be either `Mapping` or `pandas. DataFrame`, "
+                                    f"found `{type(new).__name__}`.")
+                # fmt: on
+            except KeyError:
+                missing = sorted(k for k in keys if k not in old)
+                raise KeyError(
+                    f"Unable to find key(s) `{missing}` in `adata.{attr}`."
+                ) from None
+
+        adata = self._shadow_adata.copy()
+        _adata = self.adata
+        try:
+            # kernel and estimator share the adata
+            self.adata = adata
+            self.kernel.write_to_adata()
+        finally:
+            self.adata = _adata
+        key = Key.uns.estimator(self.backward) + "_params"
+        adata.uns[key] = deepcopy(self.params)
+
+        # fmt: off
+        if isinstance(keep, str):
+            if keep == 'all':
+                keep = ["X", "raw", "layers", "obs", "var", "obsm", "varm", "obsp", "varp", "uns"]
+            else:
+                keep = [keep]
+        if not isinstance(keep, Mapping):
+            keep = {attr: True for attr in keep}
+        if isinstance(copy, bool):
+            copy = {attr: copy for attr in keep}
+        elif isinstance(copy, str):
+            copy = [copy]
+        if not isinstance(copy, Mapping):
+            copy = {attr: True for attr in copy}
+        # fmt: on
+
+        for attr, keys in keep.items():
+            if keys is True:
+                keys = ["all"]
+            elif isinstance(keys, str):
+                keys = [keys]
+            if keys is False or not len(keys):
+                continue
+            handle_attribute(attr, keys=list(keys), copy=copy.get(attr, False))
+
+        return adata
+
+    @classmethod
     @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]:
+    def from_adata(cls, adata: AnnData, obsp_key: str) -> "BaseEstimator":
         """
-        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.
+        %(from_adata.full_desc)s
 
         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`.
+        %(adata)s
+        obsp_key
+            Key in :attr:`anndata.AnnData.obsp` where the transition matrix is stored.
 
         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
+        %(from_adata.returns)s
+        """  # noqa: D400
+        return super().from_adata(adata, obsp_key=obsp_key)
 
-    def _detect_cc_stages(self, rc_labels: Series, p_thresh: float = 1e-15) -> None:
+    @d.dedent
+    def copy(self, *, deep: bool = False) -> "BaseEstimator":
         """
-        Detect cell-cycle driven start or endpoints.
+        Return a copy of self.
 
         Parameters
         ----------
-        rc_labels
-            Approximate recurrent classes.
-        p_thresh
-            P-value threshold for the rank-sum test for the group to be considered cell-cycle driven.
+        deep
+            Whether to return a deep copy or not. If `True`, this also copies the :attr:`adata`.
 
         Returns
         -------
-        None
-            Nothing, but warns if a group is cell-cycle driven.
+        A copy of self.
         """
+        k = deepcopy(self.kernel) if deep else copy_(self.kernel)
+        res = type(self)(k)
+        for k, v in self.__dict__.items():
+            if isinstance(v, Mapping):
+                res.__dict__[k] = deepcopy(v)
+            elif k != "_kernel":
+                res.__dict__[k] = deepcopy(v) if deep else copy_(v)
 
-        # initialize the groups (start or end clusters) and scores
-        groups = rc_labels.cat.categories
-        scores = []
-        if self._G2M_score is not None:
-            scores.append(self._G2M_score)
-        if self._S_score is not None:
-            scores.append(self._S_score)
-
-        # loop over groups and scores
-        for group in groups:
-            mask = rc_labels == group
-            for score in scores:
-                a, b = score[mask], score[~mask]
-                statistic, pvalue = ranksums(a, b)
-                if statistic > 0 and pvalue < p_thresh:
-                    logg.warning(f"Group `{group!r}` appears to be cell-cycle driven")
-                    break
-
-    def _set_categorical_labels(
-        self,
-        attr_key: str,
-        color_key: str,
-        pretty_attr_key: str,
-        categories: Union[Series, Dict[Any, Any]],
-        add_to_existing_error_msg: Optional[str] = None,
-        cluster_key: Optional[str] = None,
-        en_cutoff: Optional[float] = None,
-        p_thresh: Optional[float] = None,
-        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)
-            )
-        if not is_categorical_dtype(categories):
-            raise TypeError(
-                f"Object must be `categorical`, found `{infer_dtype(categories)}`."
-            )
-
-        if add_to_existing:
-            if getattr(self, attr_key) is None:
-                raise RuntimeError(add_to_existing_error_msg)
-            categories = _merge_categorical_series(
-                old=getattr(self, attr_key),
-                new=categories,
-            )
-
-        if cluster_key is not None:
-            logg.debug(f"Creating colors based on `{cluster_key}`")
-
-            # check that we can load the reference series from adata
-            if cluster_key not in self.adata.obs:
-                raise KeyError(
-                    f"Cluster key `{cluster_key!r}` not found in `adata.obs`."
-                )
-            series_query, series_reference = categories, self.adata.obs[cluster_key]
-
-            # load the reference colors if they exist
-            if _colors(cluster_key) in self.adata.uns.keys():
-                colors_reference = _convert_to_hex_colors(
-                    self.adata.uns[_colors(cluster_key)]
-                )
-            else:
-                colors_reference = _create_categorical_colors(
-                    len(series_reference.cat.categories)
-                )
-
-            approx_rcs_names, colors = _map_names_and_colors(
-                series_reference=series_reference,
-                series_query=series_query,
-                colors_reference=colors_reference,
-                en_cutoff=en_cutoff,
-            )
-            setattr(self, color_key, colors)
-            # if approx_rcs_names is categorical, the info is take from .cat.categories
-            categories.cat.categories = approx_rcs_names
-        else:
-            setattr(
-                self,
-                color_key,
-                _create_categorical_colors(len(categories.cat.categories)),
-            )
-
-        if p_thresh is not None:
-            self._detect_cc_stages(categories, p_thresh=p_thresh)
-
-        # write to class and adata
-        if getattr(self, attr_key) is not None:
-            logg.debug(f"Overwriting `.{pretty_attr_key}`")
+        return res
 
-        setattr(self, attr_key, categories)
+    def __copy__(self) -> "BaseEstimator":
+        return self.copy(deep=False)
 
-    def _write_terminal_states(self, time=None) -> None:
-        self.adata.obs[self._term_key] = self._get(P.TERM)
-        self.adata.obs[_probs(self._term_key)] = self._get(P.TERM_PROBS)
+    def __repr__(self) -> str:
+        return f"{self.__class__.__name__}[n={len(self)}, kernel={repr(self.kernel)}]"
 
-        self.adata.uns[_colors(self._term_key)] = self._get(A.TERM_COLORS)
-        self.adata.uns[_lin_names(self._term_key)] = np.array(
-            self._get(P.TERM).cat.categories
-        )
+    def __str__(self) -> str:
+        return f"{self.__class__.__name__}[n={len(self)}, kernel={str(self.kernel)}]"
 
-        logg.info(
-            f"Adding `adata.obs[{_probs(self._term_key)!r}]`\n"
-            f"       `adata.obs[{self._term_key!r}]`\n"
-            f"       `.{P.TERM_PROBS}`\n"
-            f"       `.{P.TERM}`\n"
-            "    Finish",
-            time=time,
-        )
+    @property
+    def params(self) -> Dict[str, Any]:
+        """Estimator parameters."""
+        return self._params
 
     @abstractmethod
-    def _fit_terminal_states(self, *args: Any, **kwargs: Any) -> None:
+    def fit(self, *args: Any, **kwargs: Any) -> "BaseEstimator":
         """
-        High level API helper method called inside :meth:`fit` that should compute the terminal states.
-
-        This method would usually call :meth:`compute_terminal_states` after all the functions that required beforehand.
+        Fit an estimator.
 
         Parameters
         ----------
@@ -1305,95 +473,22 @@ def _fit_terminal_states(self, *args: Any, **kwargs: Any) -> None:
 
         Returns
         -------
-        None
-            Nothing, just sets the terminal states.
-
-        See also
-        --------
-        See :meth:`cellrank.tl.estimators.GPCCA._fit_terminal_states` for an example implementation.
+        Self.
         """
 
-    @d.dedent
-    @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS, ap=P.ABS_PROBS, pd=P.PRIME_DEG)
-    def fit(
-        self,
-        keys: Optional[Sequence] = None,
-        compute_absorption_probabilities: bool = True,
-        **kwargs,
-    ) -> None:
+    @abstractmethod
+    def predict(self, *args: Any, **kwargs: Any) -> None:
         """
-        Run the pipeline.
+        Run a prediction.
 
         Parameters
         ----------
-        keys
-            States for which to compute absorption probabilities.
-        compute_absorption_probabilities
-            Whether to compute absorption probabilities or just %(initial_or_terminal)s states.
+        args
+            Positional arguments.
         kwargs
             Keyword arguments.
 
         Returns
         -------
-        None
-            Nothing, just makes available the following fields:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-                - :attr:`{ap}`
-                - :attr:`{pd}`
+        Nothing.
         """
-        self._fit_terminal_states(**kwargs)
-        if compute_absorption_probabilities:
-            self.compute_absorption_probabilities(keys=keys)
-
-    def _write_absorption_probabilities(
-        self, time: datetime, extra_msg: str = ""
-    ) -> None:
-        self.adata.obsm[self._abs_prob_key] = self._get(P.ABS_PROBS)
-
-        abs_prob = self._get(P.ABS_PROBS)
-
-        self.adata.uns[_lin_names(self._abs_prob_key)] = abs_prob.names
-        self.adata.uns[_colors(self._abs_prob_key)] = abs_prob.colors
-
-        logg.info(
-            f"Adding `adata.obsm[{self._abs_prob_key!r}]`\n"
-            f"{extra_msg}"
-            f"       `.{P.ABS_PROBS}`\n"
-            "    Finish",
-            time=time,
-        )
-
-    def _set(self, n: Union[str, PrettyEnum], v: Any) -> None:
-        setattr(self, n.s if isinstance(n, PrettyEnum) else n, v)
-
-    def _get(self, n: Union[str, PrettyEnum]) -> Any:
-        return getattr(self, n.s if isinstance(n, PrettyEnum) else n)
-
-    def _set_or_debug(
-        self, needle: str, haystack, attr: Optional[Union[str, PrettyEnum]] = None
-    ) -> Optional[Any]:
-        if isinstance(attr, PrettyEnum):
-            attr = attr.s
-        if needle in haystack:
-            if attr is None:
-                return haystack[needle]
-            setattr(self, attr, haystack[needle])
-        elif attr is not None:
-            logg.debug(f"Unable to set attribute `.{attr}`, skipping")
-
-    def copy(self) -> "BaseEstimator":
-        """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():
-            if isinstance(v, dict):
-                res.__dict__[k] = deepcopy(v)
-            elif k != "_kernel":
-                res.__dict__[k] = copy(v)
-
-        return res
-
-    def __copy__(self) -> "BaseEstimator":
-        return self.copy()
diff --git a/cellrank/tl/estimators/_cflare.py b/cellrank/tl/estimators/_cflare.py
deleted file mode 100644
index f1db584b3..000000000
--- a/cellrank/tl/estimators/_cflare.py
+++ /dev/null
@@ -1,319 +0,0 @@
-"""Clustering and Filtering of Left and Right Eigenvectors based on Markov chains."""
-from typing import List, Tuple, Union, Optional, Sequence
-from typing_extensions import Literal
-
-from cellrank import logging as logg
-from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import (
-    _cluster_X,
-    _filter_cells,
-    _complex_warning,
-    _get_connectivities,
-)
-from cellrank.tl.estimators._constants import A, P
-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):
-    """
-    Compute the initial/terminal states of a Markov chain via spectral heuristics.
-
-    This estimator uses the left eigenvectors of the transition matrix to filter to a set of recurrent cells
-    and the right eigenvectors to cluster this set of cells into discrete groups.
-
-    Parameters
-    ----------
-    %(base_estimator.parameters)s
-    """
-
-    @d.dedent
-    @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS)
-    def compute_terminal_states(
-        self,
-        use: Optional[Union[int, Tuple[int], List[int], range]] = None,
-        percentile: Optional[int] = 98,
-        method: Literal["leiden", "means"] = "kmeans",
-        cluster_key: Optional[str] = None,
-        n_clusters_kmeans: Optional[int] = None,
-        n_neighbors: int = 20,
-        resolution: float = 0.1,
-        n_matches_min: Optional[int] = 0,
-        n_neighbors_filtering: int = 15,
-        basis: Optional[str] = None,
-        n_comps: int = 5,
-        scale: bool = False,
-        en_cutoff: Optional[float] = 0.7,
-        p_thresh: float = 1e-15,
-    ) -> None:
-        """
-        Find approximate recurrent classes of the Markov chain.
-
-        Filter to obtain recurrent states in left eigenvectors.
-        Cluster to obtain approximate recurrent classes in right eigenvectors.
-
-        Parameters
-        ----------
-        use
-            Which or how many first eigenvectors to use as features for clustering/filtering.
-            If `None`, use the `eigengap` statistic.
-        percentile
-            Threshold used for filtering out cells which are most likely transient states. Cells which are in the
-            lower ``percentile`` percent of each eigenvector will be removed from the data matrix.
-        method
-            Method to be used for clustering. Valid option are:
-
-                - `'kmeans'` - :class:`sklearn.cluster.KMeans`.
-                - `'leiden'` - :func:`scanpy.tl.leiden`.
-        cluster_key
-            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
-            Number of neighbors in a KNN graph. This is the :math:`K` parameter for that,
-            the number of neighbors for each cell. Only used when ``method = 'leiden'``.
-        resolution
-            Resolution parameter for :func:`scanpy.tl.leiden`. Should be chosen relatively small.
-        n_matches_min
-            Filters out cells which don't have at least n_matches_min neighbors from the same class.
-            This filters out some cells which are transient but have been misassigned.
-        n_neighbors_filtering
-            Parameter for filtering cells. Cells are filtered out if they don't have at least ``n_matches_min``
-            neighbors among their ``n_neighbors_filtering`` nearest cells.
-        basis
-            Key from :attr:`adata` ``.obsm`` to be used as additional features for the clustering.
-        n_comps
-            Number of embedding components to be use when ``basis != None``.
-        scale
-            Scale to z-scores. Consider using this if appending embedding to features.
-        %(en_cutoff_p_thresh)s
-
-        Returns
-        -------
-        None
-            Nothing, but updates the following fields:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-        """
-
-        def _compute_macrostates_prob() -> Series:
-            """Compute a global score of being an approximate recurrent class."""
-
-            # get the truncated eigendecomposition
-            V, evals = eig["V_l"].real[:, use], eig["D"].real[use]
-
-            # shift and scale
-            V_pos = np.abs(V)
-            V_shifted = V_pos - np.min(V_pos, axis=0)
-            V_scaled = V_shifted / np.max(V_shifted, axis=0)
-
-            # check the ranges are correct
-            assert np.allclose(np.min(V_scaled, axis=0), 0), "Lower limit it not zero."
-            assert np.allclose(np.max(V_scaled, axis=0), 1), "Upper limit is not one."
-
-            # further scale by the eigenvalues
-            V_eigs = V_scaled / evals
-
-            # sum over cols and scale
-            c_ = np.sum(V_eigs, axis=1)
-            c = c_ / np.max(c_)
-
-            return Series(c, index=self.adata.obs_names)
-
-        def check_use(use) -> List[int]:
-            if method not in ["kmeans", "leiden"]:
-                raise ValueError(
-                    f"Invalid method `{method!r}`. Valid options are `'leiden'` or `'kmeans'`."
-                )
-
-            if use is None:
-                use = eig["eigengap"] + 1  # add one b/c indexing starts at 0
-            if isinstance(use, int):
-                use = list(range(use))
-            elif not isinstance(use, (tuple, list, range)):
-                raise TypeError(
-                    f"Argument `use` must be either `int`, `tuple`, `list` or `range`, "
-                    f"found `{type(use).__name__!r}`."
-                )
-            else:
-                if not all(map(lambda u: isinstance(u, int), use)):
-                    raise TypeError("Not all values in `use` argument are integers.")
-            use = list(use)
-
-            if len(use) == 0:
-                raise ValueError(
-                    f"Number of eigenvector must be larger than `0`, found `{len(use)}`."
-                )
-
-            muse = max(use)
-            if muse >= eig["V_l"].shape[1] or muse >= eig["V_r"].shape[1]:
-                raise ValueError(
-                    f"Maximum specified eigenvector `{muse}` is larger "
-                    f'than the number of computed eigenvectors `{eig["V_l"].shape[1]}`. '
-                    f"Use `.compute_eigendecomposition(k={muse})` to recompute the eigendecomposition."
-                )
-
-            return use
-
-        eig = self._get(P.EIG)
-        if eig is None:
-            raise RuntimeError(
-                "Compute eigendecomposition first as `.compute_eigendecomposition()`."
-            )
-        use = check_use(use)
-
-        start = logg.info("Computing approximate recurrent classes")
-        # we check for complex values only in the left, that's okay because the complex pattern
-        # will be identical for left and right
-        V_l, V_r = eig["V_l"][:, use], eig["V_r"].real[:, use]
-        V_l = _complex_warning(V_l, use, use_imag=False)
-
-        # compute a rc probability
-        logg.debug("Computing probabilities of approximate recurrent classes")
-        self._set(A.TERM_PROBS, _compute_macrostates_prob())
-
-        # retrieve embedding and concatenate
-        if basis is not None:
-            bkey = f"X_{basis}"
-            if bkey not in self.adata.obsm.keys():
-                raise KeyError(f"Basis key `{bkey!r}` not found in `adata.obsm`")
-
-            X_em = self.adata.obsm[bkey][:, :n_comps]
-            X = np.concatenate([V_r, X_em], axis=1)
-        else:
-            logg.debug("Basis is `None`. Setting X equal to the right eigenvectors")
-            X = V_r
-
-        # filter out cells which are in the lowest q percentile in abs value in each eigenvector
-        if percentile is not None:
-            logg.debug("Filtering out cells according to percentile")
-            if percentile < 0 or percentile > 100:
-                raise ValueError(
-                    f"Percentile must be in interval `[0, 100]`, found `{percentile}`."
-                )
-            cutoffs = np.percentile(np.abs(V_l), percentile, axis=0)
-            ixs = np.sum(np.abs(V_l) < cutoffs, axis=1) < V_l.shape[1]
-            X = X[ixs, :]
-
-        # scale
-        if scale:
-            X = zscore(X, axis=0)
-
-        # cluster X
-        if method == "kmeans" and n_clusters_kmeans is None:
-            n_clusters_kmeans = len(use) + (percentile is None)
-            if X.shape[0] < n_clusters_kmeans:
-                raise ValueError(
-                    f"Filtering resulted in only {X.shape[0]} cell(s), insufficient to cluster into "
-                    f"`{n_clusters_kmeans}` clusters. Consider decreasing the value of `percentile`."
-                )
-
-        logg.debug(
-            f"Using `{use}` eigenvectors, basis `{basis!r}` and method `{method!r}` for clustering"
-        )
-        labels = _cluster_X(
-            X,
-            method=method,
-            n_clusters=n_clusters_kmeans,
-            n_neighbors=n_neighbors,
-            resolution=resolution,
-        )
-
-        # fill in the labels in case we filtered out cells before
-        if percentile is not None:
-            rc_labels = np.repeat(None, self.adata.n_obs)
-            rc_labels[ixs] = labels
-        else:
-            rc_labels = labels
-
-        rc_labels = Series(rc_labels, index=self.adata.obs_names, dtype="category")
-        rc_labels.cat.categories = list(rc_labels.cat.categories.astype("str"))
-
-        # filtering to get rid of some of the left over transient states
-        if n_matches_min > 0:
-            logg.debug(f"Filtering according to `n_matches_min={n_matches_min}`")
-            distances = _get_connectivities(
-                self.adata, mode="distances", n_neighbors=n_neighbors_filtering
-            )
-            rc_labels = _filter_cells(
-                distances, rc_labels=rc_labels, n_matches_min=n_matches_min
-            )
-
-        self.set_terminal_states(
-            labels=rc_labels,
-            cluster_key=cluster_key,
-            en_cutoff=en_cutoff,
-            p_thresh=p_thresh,
-            add_to_existing=False,
-            time=start,
-        )
-
-    def _fit_terminal_states(
-        self,
-        n_lineages: Optional[int] = None,
-        keys: Optional[Sequence[str]] = None,
-        cluster_key: Optional[str] = None,
-        compute_absorption_probabilities: bool = True,
-        **kwargs,
-    ) -> None:
-        self.compute_eigendecomposition(k=20 if n_lineages is None else n_lineages + 1)
-        if n_lineages is None:
-            n_lineages = self._get(P.EIG)["eigengap"] + 1
-
-        self.compute_terminal_states(
-            use=n_lineages,
-            cluster_key=cluster_key,
-            n_clusters_kmeans=n_lineages,
-            method=kwargs.pop("method", "kmeans"),
-            **kwargs,
-        )
-
-    @d.dedent  # because of fit
-    @d.dedent
-    @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS, ap=P.ABS_PROBS, pd=P.PRIME_DEG)
-    def fit(
-        self,
-        n_lineages: Optional[int],
-        keys: Optional[Sequence[str]] = None,
-        cluster_key: Optional[str] = None,
-        compute_absorption_probabilities: bool = True,
-        **kwargs,
-    ):
-        """
-        Run the pipeline, computing the %(initial_or_terminal)s states and optionally the absorption probabilities.
-
-        It is equivalent to running::
-
-            compute_eigendecomposition(...)
-            compute_terminal_states(...)
-            compute_absorption_probabilities(...)
-
-        Parameters
-        ----------
-        %(fit)s
-        kwargs
-            Keyword arguments for :meth:`compute_terminal_states`, such as ``n_cells``.
-
-        Returns
-        -------
-        None
-            Nothing, just makes available the following fields:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-                - :attr:`{ap}`
-                - :attr:`{pd}`
-        """
-        super().fit(
-            n_lineages=n_lineages,
-            keys=keys,
-            cluster_key=cluster_key,
-            compute_absorption_probabilities=compute_absorption_probabilities,
-            **kwargs,
-        )
diff --git a/cellrank/tl/estimators/_constants.py b/cellrank/tl/estimators/_constants.py
deleted file mode 100644
index 2b2d021bb..000000000
--- a/cellrank/tl/estimators/_constants.py
+++ /dev/null
@@ -1,59 +0,0 @@
-"""Attribute, property and function names for estimators."""
-from cellrank.tl._constants import PrettyEnum
-
-META_KEY = "__prop_metadata__"
-
-
-# attributes
-class A(PrettyEnum):
-    """Attribute names of an estimator."""
-
-    EIG = "_eig"
-    SCHUR = "_schur"
-    SCHUR_MAT = "_schur_mat"
-    MACRO = "_macrostates"
-    MACRO_MEMBER = "_macrostates_memberships"
-    MACRO_COLORS = "_macrostates_colors"
-    TERM = "_terminal_states"
-    TERM_PROBS = "_terminal_states_probabilities"
-    TERM_COLORS = "_terminal_colors"
-    TERM_ABS_PROBS = "_terminal_abs_probabilities"  # Lineage object
-    ABS_PROBS = "_absorption_probabilities"
-    PRIME_DEG = "_priming_degree"
-    COARSE_T = "_coarse_T"
-    COARSE_INIT_D = "_coarse_init_dist"
-    COARSE_STAT_D = "_coarse_stat_dist"
-    LIN_ABS_TIMES = "_lin_abs_times"
-    LIN_DRIVERS = "_lin_drivers"
-
-
-# properties
-class P(PrettyEnum):
-    """Property names of an estimator."""
-
-    NO_PROPERTY = "NO_PROPERTY"
-    EMPTY = "EMPTY"
-    EIG = "eigendecomposition"
-    SCHUR = "schur"
-    SCHUR_MAT = "schur_matrix"
-    MACRO = "macrostates"
-    MACRO_MEMBER = "macrostates_memberships"
-    TERM = "terminal_states"
-    TERM_PROBS = "terminal_states_probabilities"
-    ABS_PROBS = "absorption_probabilities"
-    PRIME_DEG = "priming_degree"
-    COARSE_T = "coarse_T"
-    COARSE_INIT_D = "coarse_initial_distribution"
-    COARSE_STAT_D = "coarse_stationary_distribution"
-    LIN_ABS_TIMES = "lineage_absorption_times"
-    LIN_DRIVERS = "lineage_drivers"
-
-
-# functions
-class F(PrettyEnum):
-    """Function formats."""
-
-    NO_FUNC = "NO_FUNC"
-    GENERAL = "{}"
-    COMPUTE = "compute_{}"
-    PLOT = "plot_{}"
diff --git a/cellrank/tl/estimators/_decomposition.py b/cellrank/tl/estimators/_decomposition.py
deleted file mode 100644
index 76bd36015..000000000
--- a/cellrank/tl/estimators/_decomposition.py
+++ /dev/null
@@ -1,580 +0,0 @@
-"""Matrix decomposition module."""
-from typing import Any, Tuple, Union, Mapping, Optional
-
-from abc import ABC
-from pathlib import Path
-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
-
-import matplotlib.pyplot as plt
-from matplotlib.ticker import MultipleLocator, FormatStrFormatter
-from mpl_toolkits.axes_grid1 import make_axes_locatable
-
-EPS = np.finfo(np.float64).eps
-
-
-class Decomposable(KernelHolder, Property, ABC):
-    """Helper class exposes writing the eigendecomposition to :class:`anndata.AnnData` object."""
-
-    def _write_eig_to_adata(
-        self, eig: Mapping[str, Any], start=None, extra_msg: Optional[str] = None
-    ):
-        setattr(self, A.EIG.s, eig)
-        self.adata.uns[f"eig_{self._direction}"] = eig
-
-        msg = f"Adding `adata.uns['eig_{self._direction}']`\n       `.{P.EIG}`"
-        if extra_msg is None:
-            extra_msg = "\n    Finish"
-
-        msg += extra_msg
-
-        logg.info(msg, time=start)
-
-
-class Eigen(VectorPlottable, Decomposable):
-    """Class computing the eigendecomposition."""
-
-    __prop_metadata__ = [
-        Metadata(attr=A.EIG, prop=P.EIG, dtype=Mapping[str, Any], compute_fmt=F.NO_FUNC)
-    ]
-
-    @d.dedent
-    @inject_docs(prop=P.EIG)
-    def compute_eigendecomposition(
-        self,
-        k: int = 20,
-        which: str = "LR",
-        alpha: float = 1,
-        only_evals: bool = False,
-        ncv: Optional[int] = None,
-    ) -> None:
-        """
-        Compute eigendecomposition of transition matrix.
-
-        Uses a sparse implementation, if possible, and only computes the top :math:`k` eigenvectors
-        to speed up the computation. Computes both left and right eigenvectors.
-
-        Parameters
-        ----------
-        k
-            Number of eigenvalues/vectors to compute.
-        %(eigen)s
-        only_evals
-            Compute only eigenvalues.
-        ncv
-            Number of Lanczos vectors generated.
-
-        Returns
-        -------
-        None
-            Nothing, but updates the following field:
-
-                - :attr:`{prop}`
-        """
-
-        def get_top_k_evals():
-            return D[np.flip(np.argsort(D.real))][:k]
-
-        start = logg.info("Computing eigendecomposition of the transition matrix")
-
-        if self.issparse:
-            logg.debug(f"Computing top `{k}` eigenvalues for sparse matrix")
-            D, V_l = eigs(self.transition_matrix.T, k=k, which=which, ncv=ncv)
-            if only_evals:
-                self._write_eig_to_adata(
-                    {
-                        "D": get_top_k_evals(),
-                        "eigengap": _eigengap(get_top_k_evals().real, alpha),
-                        "params": {"which": which, "k": k, "alpha": alpha},
-                    }
-                )
-                return
-            _, V_r = eigs(self.transition_matrix, k=k, which=which, ncv=ncv)
-        else:
-            logg.warning(
-                "This transition matrix is not sparse, computing full eigendecomposition"
-            )
-            D, V_l = np.linalg.eig(self.transition_matrix.T)
-            if only_evals:
-                self._write_eig_to_adata(
-                    {
-                        "D": get_top_k_evals(),
-                        "eigengap": _eigengap(D.real, alpha),
-                        "params": {"which": which, "k": k, "alpha": alpha},
-                    }
-                )
-                return
-            _, V_r = np.linalg.eig(self.transition_matrix)
-
-        # Sort the eigenvalues and eigenvectors and take the real part
-        logg.debug("Sorting eigenvalues by their real part")
-        p = np.flip(np.argsort(D.real))
-        D, V_l, V_r = D[p], V_l[:, p], V_r[:, p]
-        e_gap = _eigengap(D.real, alpha)
-
-        pi = np.abs(V_l[:, 0].real)
-        pi /= np.sum(pi)
-
-        self._write_eig_to_adata(
-            {
-                "D": D,
-                "stationary_dist": pi,
-                "V_l": V_l,
-                "V_r": V_r,
-                "eigengap": e_gap,
-                "params": {"which": which, "k": k, "alpha": alpha},
-            },
-            start=start,
-        )
-
-    @d.dedent
-    def plot_eigendecomposition(self, left: bool = False, *args, **kwargs):
-        """
-        Plot eigenvectors in an embedding.
-
-        Parameters
-        ----------
-        left
-            Whether to plot left or right eigenvectors.
-        %(plot_vectors.parameters)s
-
-        Returns
-        -------
-        %(plot_vectors.returns)s
-        """
-
-        eig = getattr(self, P.EIG.s, None)
-
-        if eig is None:
-            self._plot_vectors(None, P.EIG.s)
-
-        side = "left" if left else "right"
-        D, V = (
-            eig["D"],
-            eig.get(f"V_{side[0]}", None),
-        )
-        if V is None:
-            raise RuntimeError(
-                "Compute eigendecomposition first as `.compute_eigendecomposition(..., only_evals=False)`."
-            )
-
-        # if irreducible, first rigth e-vec should be const.
-        if side == "right":
-            # quick check for irreducibility:
-            if np.sum(np.isclose(D, 1, rtol=1e2 * EPS, atol=1e2 * EPS)) == 1:
-                V[:, 0] = 1.0
-
-        self._plot_vectors(
-            V,
-            P.EIG.s,
-            *args,
-            D=D,
-            **kwargs,
-        )
-
-    @d.dedent
-    def plot_spectrum(
-        self,
-        n: Optional[int] = None,
-        real_only: bool = False,
-        show_eigengap: bool = True,
-        show_all_xticks: bool = True,
-        legend_loc: Optional[str] = None,
-        title: Optional[str] = None,
-        figsize: Optional[Tuple[float, float]] = (5, 5),
-        dpi: int = 100,
-        save: Optional[Union[str, Path]] = None,
-        marker: str = ".",
-        **kwargs,
-    ) -> None:
-        """
-        Plot the top eigenvalues in real or complex plane.
-
-        Parameters
-        ----------
-        n
-            Number of eigenvalues to show. If `None`, show all that have been computed.
-        real_only
-            Whether to plot only the real part of the spectrum.
-        show_eigengap
-            When `real_only=True`, this determines whether to show the inferred eigengap as
-            a dotted line.
-        show_all_xticks
-            When `real_only=True`, this determines whether to show the indices of all eigenvalues
-            on the x-axis.
-        legend_loc
-            Location parameter for the legend.
-        title
-            Title of the figure.
-        %(plotting)s
-        marker
-            Marker symbol used, valid options can be found in :mod:`matplotlib.markers`.
-        kwargs
-            Keyword arguments for :func:`matplotlib.pyplot.scatter`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        eig = getattr(self, P.EIG.s, None)
-        if eig is None:
-            raise RuntimeError(
-                f"Compute `.{P.EIG}` first as `.{F.COMPUTE.fmt(P.EIG)}()`."
-            )
-        if n is None:
-            n = len(eig["D"])
-        elif n <= 0:
-            raise ValueError(f"Expected `n` to be > 0, found `{n}`.")
-
-        if real_only:
-            fig = self._plot_real_spectrum(
-                n,
-                show_eigengap=show_eigengap,
-                show_all_xticks=show_all_xticks,
-                dpi=dpi,
-                figsize=figsize,
-                legend_loc=legend_loc,
-                title=title,
-                marker=marker,
-                **kwargs,
-            )
-        else:
-            fig = self._plot_complex_spectrum(
-                n,
-                dpi=dpi,
-                figsize=figsize,
-                legend_loc=legend_loc,
-                title=title,
-                marker=marker,
-                **kwargs,
-            )
-
-        if save:
-            save_fig(fig, save)
-
-    def _plot_complex_spectrum(
-        self,
-        n: int,
-        dpi: int = 100,
-        figsize: Optional[Tuple[float, float]] = (None, None),
-        legend_loc: Optional[str] = None,
-        title: Optional[str] = None,
-        marker: str = ".",
-        **kwargs,
-    ):
-        # define a function to make the data limits rectangular
-        def adapt_range(min_, max_, range_):
-            return (
-                min_ + (max_ - min_) / 2 - range_ / 2,
-                min_ + (max_ - min_) / 2 + range_ / 2,
-            )
-
-        eig = getattr(self, P.EIG.s)
-        D, params = eig["D"][:n], eig["params"]
-
-        # create fiture and axes
-        fig, ax = plt.subplots(nrows=1, ncols=1, dpi=dpi, figsize=figsize)
-
-        # get the original data ranges
-        lam_x, lam_y = D.real, D.imag
-        x_min, x_max = np.min(lam_x), np.max(lam_x)
-        y_min, y_max = np.min(lam_y), np.max(lam_y)
-        x_range, y_range = x_max - x_min, y_max - y_min
-        final_range = np.max([x_range, y_range]) + 0.05
-
-        x_min_, x_max_ = adapt_range(x_min, x_max, final_range)
-        y_min_, y_max_ = adapt_range(y_min, y_max, final_range)
-
-        # plot the data and the unit circle
-        ax.scatter(D.real, D.imag, marker=marker, label="eigenvalue", **kwargs)
-        t = np.linspace(0, 2 * np.pi, 500)
-        x_circle, y_circle = np.sin(t), np.cos(t)
-        ax.plot(x_circle, y_circle, "k-", label="unit circle")
-
-        # set labels, ranges and legend
-        ax.set_xlabel(r"Re($\lambda$)")
-        ax.set_xlim(x_min_, x_max_)
-
-        ax.set_ylabel(r"Im($\lambda$)")
-        ax.set_ylim(y_min_, y_max_)
-
-        key = "real part" if params["which"] == "LR" else "magnitude"
-        if title is None:
-            title = f"top {n} eigenvalues according to their {key}"
-
-        ax.set_title(title)
-        if legend_loc not in (None, "none"):
-            ax.legend(loc=legend_loc)
-
-        return fig
-
-    def _plot_real_spectrum(
-        self,
-        n: int,
-        show_eigengap: bool = True,
-        show_all_xticks: bool = True,
-        dpi: int = 100,
-        figsize: Optional[Tuple[float, float]] = None,
-        legend_loc: Optional[str] = None,
-        title: Optional[str] = None,
-        marker: str = ".",
-        **kwargs,
-    ):
-        eig = getattr(self, P.EIG.s)
-        D, params = eig["D"][:n], eig["params"]
-
-        D_real, D_imag = D.real, D.imag
-        ixs = np.arange(len(D))
-        mask = D_imag == 0
-
-        # plot the top eigenvalues
-        fig, ax = plt.subplots(nrows=1, ncols=1, dpi=dpi, figsize=figsize)
-        if np.any(mask):
-            ax.scatter(
-                ixs[mask],
-                D_real[mask],
-                marker=marker,
-                label="real eigenvalue",
-                **kwargs,
-            )
-        if np.any(~mask):
-            ax.scatter(
-                ixs[~mask],
-                D_real[~mask],
-                marker=marker,
-                label="complex eigenvalue",
-                **kwargs,
-            )
-
-        # add dashed line for the eigengap, ticks, labels, title and legend
-        if show_eigengap and eig["eigengap"] < n:
-            ax.axvline(eig["eigengap"], label="eigengap", ls="--", lw=2, c="k")
-
-        ax.set_xlabel("index")
-        if show_all_xticks:
-            ax.set_xticks(np.arange(len(D)))
-        else:
-            ax.xaxis.set_major_locator(MultipleLocator(2.0))
-            ax.xaxis.set_major_formatter(FormatStrFormatter("%d"))
-
-        ax.set_ylabel(r"Re($\lambda_i$)")
-
-        key = "real part" if params["which"] == "LR" else "magnitude"
-        if title is None:
-            title = f"real part of top {n} eigenvalues according to their {key}"
-
-        ax.set_title(title)
-        if legend_loc not in (None, "none"):
-            ax.legend(loc=legend_loc)
-
-        return fig
-
-
-class Schur(VectorPlottable, Decomposable):
-    """Class computing the Schur decomposition."""
-
-    __prop_metadata__ = [
-        Metadata(
-            attr=A.SCHUR,
-            prop=P.SCHUR,
-            dtype=np.ndarray,
-            compute_fmt=F.NO_FUNC,
-            doc="Schur vectors.",
-        ),
-        Metadata(attr=A.SCHUR_MAT, prop=P.SCHUR_MAT, dtype=np.ndarray),
-        Metadata(attr=A.EIG, prop=P.EIG, dtype=Mapping[str, Any]),
-        Metadata(attr="_invalid_n_states", prop=P.NO_PROPERTY, dtype=np.ndarray),
-        Metadata(attr="_gpcca", prop=P.NO_PROPERTY),
-    ]
-
-    @d.dedent
-    @inject_docs(schur_vectors=P.SCHUR, schur_matrix=P.SCHUR_MAT, eigendec=P.EIG)
-    def compute_schur(
-        self,
-        n_components: int = 10,
-        initial_distribution: Optional[np.ndarray] = None,
-        method: str = "krylov",
-        which: str = "LR",
-        alpha: float = 1,
-    ):
-        """
-        Compute the Schur decomposition.
-
-        Parameters
-        ----------
-        n_components
-            Number of vectors to compute.
-        initial_distribution
-            Input probability distribution over all cells. If `None`, uniform is chosen.
-        method
-            Method for calculating the Schur vectors. Valid options are: `'krylov'` or `'brandts'`.
-            For benefits of each method, see :class:`pygpcca.GPCCA`.
-
-            The former is an iterative procedure that computes a partial, sorted Schur decomposition for large, sparse
-            matrices whereas the latter computes a full sorted Schur decomposition of a dense matrix.
-        %(eigen)s
-
-        Returns
-        -------
-        None
-            Nothing, but updates the following fields:
-
-                - :attr:`{schur_vectors}`
-                - :attr:`{schur_matrix}`
-                - :attr:`{eigendec}`
-        """
-        if n_components < 2:
-            raise ValueError(
-                f"Number of components must be `>=2`, found `{n_components}`."
-            )
-
-        if method not in ("brandts", "krylov"):
-            raise ValueError(
-                f"Invalid method `{method!r}`. Valid options are: 'brandts', 'krylov'."
-            )
-
-        try:
-            import petsc4py
-            import slepc4py
-        except ImportError:
-            method = "brandts"
-            logg.warning(
-                f"Unable to import `petsc4py` or `slepc4py`. Using `method={method!r}`"
-            )
-
-        if method == "brandts" and issparse(self.transition_matrix):
-            logg.warning("For `method='brandts'`, dense matrix is required. Densifying")
-            tmat = self.transition_matrix.A
-        else:
-            tmat = self.transition_matrix
-
-        self._gpcca = _GPCCA(tmat, eta=initial_distribution, z=which, method=method)
-        start = logg.info("Computing Schur decomposition")
-
-        try:
-            self._gpcca._do_schur_helper(n_components)
-        except ValueError:
-            logg.warning(
-                f"Using `{n_components}` components would split a block of complex conjugates. "
-                f"Increasing `n_components` to `{n_components + 1}`"
-            )
-            self._gpcca._do_schur_helper(n_components + 1)
-
-        # make it available for plotting
-        setattr(self, A.SCHUR.s, self._gpcca._p_X)
-        setattr(self, A.SCHUR_MAT.s, self._gpcca._p_R)
-
-        self._invalid_n_states = np.array(
-            [
-                i
-                for i in range(2, len(self._gpcca._p_eigenvalues))
-                if _check_conj_split(self._gpcca._p_eigenvalues[:i])
-            ]
-        )
-        if len(self._invalid_n_states):
-            logg.info(
-                f"When computing macrostates, choose a number of states NOT in `{list(self._invalid_n_states)}`"
-            )
-
-        self._write_eig_to_adata(
-            {
-                "D": self._gpcca._p_eigenvalues,
-                "eigengap": _eigengap(self._gpcca._p_eigenvalues, alpha),
-                "params": {
-                    "which": which,
-                    "k": len(self._gpcca._p_eigenvalues),
-                    "alpha": alpha,
-                },
-            },
-            start=start,
-            extra_msg=f"\n       `.{P.SCHUR}`\n       `.{P.SCHUR_MAT}`\n    Finish",
-        )
-
-    plot_schur = _delegate(prop_name=P.SCHUR.s)(VectorPlottable._plot_vectors)
-
-    @d.dedent
-    def plot_schur_matrix(
-        self,
-        title: Optional[str] = "schur matrix",
-        cmap: str = "viridis",
-        figsize: Optional[Tuple[float, float]] = None,
-        dpi: Optional[float] = 80,
-        save: Optional[Union[str, Path]] = None,
-        **kwargs,
-    ):
-        """
-        Plot the Schur matrix.
-
-        Parameters
-        ----------
-        title
-            Title of the figure.
-        cmap
-            Colormap to use.
-        %(plotting)s
-        kwargs
-            Keyword arguments for :func:`seaborn.heatmap`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        from seaborn import heatmap
-
-        schur_matrix = getattr(self, P.SCHUR_MAT.s, None)
-
-        if schur_matrix is None:
-            raise RuntimeError(
-                f"Compute Schur matrix first as `.{F.COMPUTE.fmt(P.SCHUR)}()`."
-            )
-
-        fig, ax = plt.subplots(
-            figsize=schur_matrix.shape if figsize is None else figsize, dpi=dpi
-        )
-
-        divider = make_axes_locatable(ax)  # square=True make the colorbar a bit bigger
-        cbar_ax = divider.append_axes("right", size="2%", pad=0.1)
-
-        mask = np.zeros_like(schur_matrix, dtype=np.bool_)
-        mask[np.tril_indices_from(mask, k=-1)] = True
-        mask[~np.isclose(schur_matrix, 0.0)] = False
-
-        vmin, vmax = (
-            np.min(schur_matrix[~mask]),
-            np.max(schur_matrix[~mask]),
-        )
-
-        kwargs["fmt"] = kwargs.get("fmt", "0.2f")
-        heatmap(
-            data=schur_matrix,
-            cmap=cmap,
-            square=True,
-            annot=True,
-            vmin=vmin,
-            vmax=vmax,
-            cbar_ax=cbar_ax,
-            cbar_kws={"ticks": np.linspace(vmin, vmax, 10)},
-            mask=mask,
-            xticklabels=[],
-            yticklabels=[],
-            ax=ax,
-            **kwargs,
-        )
-
-        ax.set_title(title)
-
-        if save is not None:
-            save_fig(fig, path=save)
diff --git a/cellrank/tl/estimators/_gpcca.py b/cellrank/tl/estimators/_gpcca.py
deleted file mode 100644
index 231d1aa58..000000000
--- a/cellrank/tl/estimators/_gpcca.py
+++ /dev/null
@@ -1,1351 +0,0 @@
-"""Generalized Perron Cluster Cluster Analysis :cite:`reuter:18`."""
-
-from typing import Any, Dict, List, Tuple, Union, Mapping, Iterable, Optional, Sequence
-
-from types import MappingProxyType
-from pathlib import Path
-
-from cellrank import logging as logg
-from cellrank.ul._docs import d, inject_docs
-from cellrank.tl._utils import (
-    save_fig,
-    _eigengap,
-    _fuzzy_to_discrete,
-    _series_from_one_hot_matrix,
-)
-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
-from cellrank.tl.estimators._property import Macrostates
-from cellrank.tl.estimators._constants import A, F, P
-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 :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
-    assignment is soft, i.e. each cell is assigned to every macrostate with a certain weight, where weights sum to
-    one per cell. Macrostates are computed by maximizing the 'crispness' which can be thought of as a measure for
-    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 :cite:`reuter:18`.
-
-    Parameters
-    ----------
-    %(base_estimator.parameters)s
-    """  # noqa: E501
-
-    __prop_metadata__ = [
-        Metadata(
-            attr=A.COARSE_T,
-            prop=P.COARSE_T,
-            compute_fmt=F.NO_FUNC,
-            plot_fmt=F.NO_FUNC,
-            dtype=pd.DataFrame,
-            doc="Coarse-grained transition matrix.",
-        ),
-        Metadata(attr=A.TERM_ABS_PROBS, prop=P.NO_PROPERTY, dtype=Lineage),
-        Metadata(attr=A.COARSE_INIT_D, prop=P.COARSE_INIT_D, dtype=pd.Series),
-        Metadata(attr=A.COARSE_STAT_D, prop=P.COARSE_STAT_D, dtype=pd.Series),
-    ]
-
-    def _read_from_adata(self) -> None:
-        super()._read_from_adata()
-        self._reconstruct_lineage(
-            A.TERM_ABS_PROBS,
-            self._term_abs_prob_key,
-        )
-
-    @inject_docs(
-        ms=P.MACRO,
-        msp=P.MACRO_MEMBER,
-        schur=P.SCHUR.s,
-        coarse_T=P.COARSE_T,
-        coarse_stat=P.COARSE_STAT_D,
-    )
-    @d.dedent
-    def compute_macrostates(
-        self,
-        n_states: Optional[
-            Union[int, Tuple[int, int], List[int], Dict[str, int]]
-        ] = None,
-        n_cells: Optional[int] = 30,
-        use_min_chi: bool = False,
-        cluster_key: str = None,
-        en_cutoff: Optional[float] = 0.7,
-        p_thresh: float = 1e-15,
-    ):
-        """
-        Compute the macrostates.
-
-        Parameters
-        ----------
-        n_states
-            Number of macrostates. If `None`, use the `eigengap` heuristic.
-        %(n_cells)s
-        use_min_chi
-            Whether to use :meth:`pygpcca.GPCCA.minChi` to calculate the number of macrostates.
-            If `True`, ``n_states`` corresponds to a closed interval `[min, max]` inside of which the potentially
-            optimal number of macrostates is searched.
-        cluster_key
-            If a key to cluster labels is given, names and colors of the states will be associated with the clusters.
-        %(en_cutoff_p_thresh)s
-
-        Returns
-        -------
-        None
-            Nothing, but updates the following fields:
-
-                - :attr:`{msp}`
-                - :attr:`{ms}`
-                - :attr:`{schur}`
-                - :attr:`{coarse_T}`
-                - :attr:`{coarse_stat}`
-        """
-
-        was_from_eigengap = False
-
-        if use_min_chi:
-            n_states = self._get_n_states_from_minchi(n_states)
-
-        if n_states is None:
-            if self._get(P.EIG) is None:
-                raise RuntimeError(
-                    "Compute eigendecomposition first as `.compute_eigendecomposition()` or `.compute_schur()`."
-                )
-            was_from_eigengap = True
-            n_states = self._get(P.EIG)["eigengap"] + 1
-            logg.info(f"Using `{n_states}` states based on eigengap")
-        elif not isinstance(n_states, int):
-            raise ValueError(
-                f"Expected `n_states` to be an integer when `use_min_chi=False`, "
-                f"found `{type(n_states).__name__!r}`."
-            )
-
-        if n_states <= 0:
-            raise ValueError(
-                f"Expected `n_states` to be positive or `None`, found `{n_states}`."
-            )
-
-        n_states = self._check_states_validity(n_states)
-        if n_states == 1:
-            self._compute_one_macrostate(
-                n_cells=n_cells,
-                cluster_key=cluster_key,
-                p_thresh=p_thresh,
-                en_cutoff=en_cutoff,
-            )
-            return
-
-        if self._gpcca is None:
-            if not was_from_eigengap:
-                raise RuntimeError(
-                    "Compute Schur decomposition first as `.compute_schur()`."
-                )
-
-            logg.warning(
-                f"Number of states `{n_states}` was automatically determined by `eigengap` "
-                "but no Schur decomposition was found. Computing with default parameters"
-            )
-            # this cannot fail if splitting occurs
-            # if it were to split, it's automatically increased in `compute_schur`
-            self.compute_schur(n_states)
-
-        # pre-computed X
-        if self._gpcca._p_X.shape[1] < n_states:
-            logg.warning(
-                f"Requested more macrostates `{n_states}` than available "
-                f"Schur vectors `{self._gpcca._p_X.shape[1]}`. Recomputing the decomposition"
-            )
-
-        start = logg.info(f"Computing `{n_states}` macrostates")
-        try:
-            self._gpcca = self._gpcca.optimize(m=n_states)
-        except ValueError as e:
-            # this is the following case - we have 4 Schur vectors, user requests 5 states, but it splits the conj. ev.
-            # in the try block, Schur decomposition with 5 vectors is computed, but it fails (no way of knowing)
-            # so in this case, we increase it by 1
-            n_states += 1
-            logg.warning(f"{e}\nIncreasing `n_states` to `{n_states}`")
-            self._gpcca = self._gpcca.optimize(m=n_states)
-
-        self._set_macrostates(
-            memberships=self._gpcca.memberships,
-            n_cells=n_cells,
-            cluster_key=cluster_key,
-            p_thresh=p_thresh,
-            en_cutoff=en_cutoff,
-        )
-
-        # cache the results and make sure we don't overwrite
-        self._set(A.SCHUR, self._gpcca._p_X)
-        self._set(A.SCHUR_MAT, self._gpcca._p_R)
-
-        names = self._get(P.MACRO_MEMBER).names
-
-        self._set(
-            A.COARSE_T,
-            pd.DataFrame(
-                self._gpcca.coarse_grained_transition_matrix,
-                index=names,
-                columns=names,
-            ),
-        )
-        self._set(
-            A.COARSE_INIT_D,
-            pd.Series(self._gpcca.coarse_grained_input_distribution, index=names),
-        )
-
-        # careful here, in case computing the stat. dist failed
-        if self._gpcca.coarse_grained_stationary_probability is not None:
-            self._set(
-                A.COARSE_STAT_D,
-                pd.Series(
-                    self._gpcca.coarse_grained_stationary_probability,
-                    index=names,
-                ),
-            )
-            logg.info(
-                f"Adding `.{P.MACRO_MEMBER}`\n"
-                f"       `.{P.MACRO}`\n"
-                f"       `.{P.SCHUR}`\n"
-                f"       `.{P.COARSE_T}`\n"
-                f"       `.{P.COARSE_STAT_D}`\n"
-                f"    Finish",
-                time=start,
-            )
-        else:
-            logg.warning("No stationary distribution found in GPCCA object")
-            logg.info(
-                f"Adding `.{P.MACRO_MEMBER}`\n"
-                f"       `.{P.MACRO}`\n"
-                f"       `.{P.SCHUR}`\n"
-                f"       `.{P.COARSE_T}`\n"
-                f"    Finish",
-                time=start,
-            )
-
-    @d.dedent
-    @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS)
-    def set_terminal_states_from_macrostates(
-        self,
-        names: Optional[Union[Sequence[str], Mapping[str, str], str]] = None,
-        n_cells: int = 30,
-    ):
-        """
-        Manually select terminal states from macrostates.
-
-        Parameters
-        ----------
-        names
-            Names of the macrostates to be marked as terminal. Multiple states can be combined using `','`,
-            such as ``["Alpha, Beta", "Epsilon"]``.  If a :class:`dict`, keys correspond to the names
-            of the macrostates and the values to the new names.  If `None`, select all macrostates.
-        %(n_cells)s
-
-        Returns
-        -------
-        None
-            Nothing, just updates the following fields:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-        """
-
-        if not isinstance(n_cells, int):
-            raise TypeError(
-                f"Expected `n_cells` to be of type `int`, found `{type(n_cells).__name__}`."
-            )
-
-        if n_cells <= 0:
-            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
-
-        probs = self._get(P.MACRO_MEMBER)
-        if probs is None:
-            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
-
-        rename = True
-        if names is None:
-            names = probs.names
-            rename = False
-        if isinstance(names, str):
-            names = [names]
-            rename = False
-        if not isinstance(names, dict):
-            names = {n: n for n in names}
-            rename = False
-
-        if not len(names):
-            raise ValueError("No macrostates have been selected.")
-
-        if not all(isinstance(old, str) for old in names.keys()):
-            raise TypeError("Not all new names are strings.")
-
-        if not all(isinstance(new, (str, int)) for new in names.values()):
-            raise TypeError("Not all macrostates names are strings or integers.")
-
-        # this also checks that the names are correct before renaming
-        macrostates_probs = probs[list(names.keys())]
-
-        # we do this also here because if `rename_terminal_states` fails
-        # invalid states would've been written to this object and nothing to adata
-        new_names = {k: str(v) for k, v in names.items()}
-        names_after_renaming = [new_names.get(n, n) for n in probs.names]
-        if len(set(names_after_renaming)) != probs.shape[1]:
-            raise ValueError(
-                f"After renaming, the names will not be unique: `{names_after_renaming}`."
-            )
-
-        if probs.shape[1] == 1:
-            self._set(A.TERM, self._create_states(probs, n_cells=n_cells))
-            self._set(A.TERM_COLORS, self._get(A.MACRO_COLORS))
-            self._set(
-                A.TERM_PROBS,
-                pd.Series(
-                    probs.X.squeeze() / probs.X.max(), index=self.adata.obs_names
-                ),
-            )
-            self._set(A.TERM_ABS_PROBS, probs)
-            if rename:
-                # access lineage renames join states, e.g. 'Alpha, Beta' becomes 'Alpha or Beta' + whitespace stripping
-                self.rename_terminal_states(
-                    dict(zip(self._get(P.TERM).cat.categories, names.values()))
-                )
-
-            self._write_terminal_states()
-            return
-
-        # compute the aggregated probability of being a initial/terminal state (no matter which)
-        scaled_probs = macrostates_probs.copy()
-        scaled_probs /= scaled_probs.max(0)
-
-        self._set(A.TERM, self._create_states(macrostates_probs, n_cells=n_cells))
-        self._set(
-            A.TERM_PROBS, pd.Series(scaled_probs.X.max(1), index=self.adata.obs_names)
-        )
-        self._set(
-            A.TERM_COLORS,
-            macrostates_probs[list(self._get(P.TERM).cat.categories)].colors,
-        )
-        self._set(A.TERM_ABS_PROBS, scaled_probs)
-        if rename:
-            self.rename_terminal_states(
-                dict(zip(self._get(P.TERM).cat.categories, names.values()))
-            )
-
-        self._write_terminal_states()
-
-    @inject_docs(fs=P.TERM, fsp=P.TERM_PROBS)
-    @d.dedent
-    def compute_terminal_states(
-        self,
-        method: str = "stability",
-        n_cells: int = 30,
-        alpha: Optional[float] = 1,
-        stability_threshold: float = 0.96,
-        n_states: Optional[int] = None,
-    ):
-        """
-        Automatically select terminal states from macrostates.
-
-        Parameters
-        ----------
-        method
-            One of following:
-
-                - `'eigengap'` - select the number of states based on the `eigengap` of the transition matrix.
-                - `'eigengap_coarse'` - select the number of states based on the `eigengap` of the diagonal
-                  of the coarse-grained transition matrix.
-                - `'top_n'` - select top ``n_states`` based on the probability of the diagonal
-                  of the coarse-grained transition matrix.
-                - `'stability'` - select states which have a stability index >= ``stability_threshold``. The stability
-                  index is given by the diagonal elements of the coarse-grained transition matrix.
-        %(n_cells)s
-        alpha
-            Weight given to the deviation of an eigenvalue from one. Used when ``method='eigengap'``
-            or ``method='eigengap_coarse'``.
-        stability_threshold
-            Threshold used when ``method='stability'``.
-        n_states
-            Numer of states used when ``method='top_n'``.
-
-        Returns
-        -------
-        None
-            Nothing, just updates the following fields:
-
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-        """
-
-        if len(self._get(P.MACRO).cat.categories) == 1:
-            logg.warning("Found only one macrostate. Making it the single main state")
-            self.set_terminal_states_from_macrostates(None, n_cells=n_cells)
-            return
-
-        coarse_T = self._get(P.COARSE_T)
-
-        if method == "eigengap":
-            if self._get(P.EIG) is None:
-                raise RuntimeError(
-                    "Compute eigendecomposition first as `.compute_eigendecomposition()`."
-                )
-            n_states = _eigengap(self._get(P.EIG)["D"], alpha=alpha) + 1
-        elif method == "eigengap_coarse":
-            if coarse_T is None:
-                raise RuntimeError(
-                    "Compute macrostates first as `.compute_macrostates()`."
-                )
-            n_states = _eigengap(np.sort(np.diag(coarse_T)[::-1]), alpha=alpha)
-        elif method == "top_n":
-            if n_states is None:
-                raise ValueError(
-                    "Argument `n_states` must be != `None` for `method='top_n'`."
-                )
-            elif n_states <= 0:
-                raise ValueError(
-                    f"Expected `n_states` to be positive, found `{n_states}`."
-                )
-        elif method == "stability":
-            if stability_threshold is None:
-                raise ValueError(
-                    "Argument `stability_threshold` must be != `None` for `method='stability'`."
-                )
-            self_probs = pd.Series(np.diag(coarse_T), index=coarse_T.columns)
-            names = self_probs[self_probs.values >= stability_threshold].index
-            self.set_terminal_states_from_macrostates(names, n_cells=n_cells)
-            return
-        else:
-            raise ValueError(
-                f"Invalid method `{method!r}`. Valid options are `'eigengap', 'eigengap_coarse', "
-                f"'top_n' and 'min_self_prob'`."
-            )
-
-        names = coarse_T.columns[np.argsort(np.diag(coarse_T))][-n_states:]
-        self.set_terminal_states_from_macrostates(names, n_cells=n_cells)
-
-    def compute_gdpt(
-        self, n_components: int = 10, key_added: str = "gdpt_pseudotime", **kwargs
-    ):
-        """
-        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 :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 :attr:`adata` ``.obs[key_added]`` with the computed pseudotime.
-        """
-
-        def _get_dpt_row(e_vals: np.ndarray, e_vecs: np.ndarray, i: int):
-            row = sum(
-                (
-                    np.abs(e_vals[eval_ix])
-                    / (1 - np.abs(e_vals[eval_ix]))
-                    * (e_vecs[i, eval_ix] - e_vecs[:, eval_ix])
-                )
-                ** 2
-                # account for float32 precision
-                for eval_ix in range(0, e_vals.size)
-                if np.abs(e_vals[eval_ix]) < 0.9994
-            )
-
-            return np.sqrt(row)
-
-        if "iroot" not in self.adata.uns.keys():
-            raise KeyError("Key `'iroot'` not found in `adata.uns`.")
-
-        iroot = self.adata.uns["iroot"]
-        if isinstance(iroot, str):
-            iroot = np.where(self.adata.obs_names == iroot)[0]
-            if not len(iroot):
-                raise ValueError(
-                    f"Unable to find cell with name `{self.adata.uns['iroot']!r}` in `adata.obs_names`."
-                )
-            iroot = iroot[0]
-
-        if n_components < 2:
-            raise ValueError(
-                f"Expected number of components >= 2, found `{n_components}`."
-            )
-
-        if self._get(P.SCHUR) is None:
-            logg.warning("No Schur decomposition found. Computing")
-            self.compute_schur(n_components, **kwargs)
-        elif self._get(P.SCHUR_MAT).shape[1] < n_components:
-            logg.warning(
-                f"Requested `{n_components}` components, but only `{self._get(P.SCHUR_MAT).shape[1]}` were found. "
-                f"Recomputing using default values"
-            )
-            self.compute_schur(n_components)
-        else:
-            logg.debug("Using cached Schur decomposition")
-
-        start = logg.info(
-            f"Computing Generalized Diffusion Pseudotime using `n_components={n_components}`"
-        )
-
-        Q, eigenvalues = (
-            self._get(P.SCHUR),
-            self._get(P.EIG)["D"],
-        )
-        # may have to remove some values if too many converged
-        Q, eigenvalues = Q[:, :n_components], eigenvalues[:n_components]
-        D = _get_dpt_row(eigenvalues, Q, i=iroot)
-        pseudotime = D / np.max(D[np.isfinite(D)])
-
-        self.adata.obs[key_added] = pseudotime
-
-        logg.info(f"Adding `{key_added!r}` to `adata.obs`\n    Finish", time=start)
-
-    @d.dedent
-    def plot_coarse_T(
-        self,
-        show_stationary_dist: bool = True,
-        show_initial_dist: bool = False,
-        cmap: Union[str, mcolors.ListedColormap] = "viridis",
-        xtick_rotation: float = 45,
-        annotate: bool = True,
-        show_cbar: bool = True,
-        title: Optional[str] = None,
-        figsize: Tuple[float, float] = (8, 8),
-        dpi: int = 80,
-        save: Optional[Union[Path, str]] = None,
-        text_kwargs: Mapping[str, Any] = MappingProxyType({}),
-        **kwargs,
-    ) -> None:
-        """
-        Plot the coarse-grained transition matrix between macrostates.
-
-        Parameters
-        ----------
-        show_stationary_dist
-            Whether to show the stationary distribution, if present.
-        show_initial_dist
-            Whether to show the initial distribution.
-        cmap
-            Colormap to use.
-        xtick_rotation
-            Rotation of ticks on the x-axis.
-        annotate
-            Whether to display the text on each cell.
-        show_cbar
-            Whether to show colorbar.
-        title
-            Title of the figure.
-        %(plotting)s
-        text_kwargs
-            Keyword arguments for :func:`matplotlib.pyplot.text`.
-        kwargs
-            Keyword arguments for :func:`matplotlib.pyplot.imshow`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        def stylize_dist(
-            ax, data: np.ndarray, xticks_labels: Union[List[str], Tuple[str]] = ()
-        ):
-            _ = ax.imshow(data, aspect="auto", cmap=cmap, norm=norm)
-            for spine in ax.spines.values():
-                spine.set_visible(False)
-
-            if xticks_labels is not None:
-                ax.set_xticks(np.arange(data.shape[1]))
-                ax.set_xticklabels(xticks_labels)
-                plt.setp(
-                    ax.get_xticklabels(),
-                    rotation=xtick_rotation,
-                    ha="right",
-                    rotation_mode="anchor",
-                )
-            else:
-                ax.set_xticks([])
-                ax.tick_params(
-                    which="both", top=False, right=False, bottom=False, left=False
-                )
-
-            ax.set_yticks([])
-
-        def annotate_heatmap(im, valfmt: str = "{x:.2f}"):
-            # modified from matplotlib's site
-
-            data = im.get_array()
-            kw = {"ha": "center", "va": "center"}
-            kw.update(**text_kwargs)
-
-            # Get the formatter in case a string is supplied
-            if isinstance(valfmt, str):
-                valfmt = mpl.ticker.StrMethodFormatter(valfmt)
-
-            # Loop over the data and create a `Text` for each "pixel".
-            # Change the text's color depending on the data.
-            texts = []
-            for i in range(data.shape[0]):
-                for j in range(data.shape[1]):
-                    kw.update(color=_get_black_or_white(im.norm(data[i, j]), cmap))
-                    text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
-                    texts.append(text)
-
-        def annotate_dist_ax(ax, data: np.ndarray, valfmt: str = "{x:.2f}"):
-            if ax is None:
-                return
-
-            if isinstance(valfmt, str):
-                valfmt = mpl.ticker.StrMethodFormatter(valfmt)
-
-            kw = {"ha": "center", "va": "center"}
-            kw.update(**text_kwargs)
-
-            for i, val in enumerate(data):
-                kw.update(color=_get_black_or_white(im.norm(val), cmap))
-                ax.text(
-                    i,
-                    0,
-                    valfmt(val, None),
-                    **kw,
-                )
-
-        coarse_T = self._get(P.COARSE_T)
-        coarse_stat_d = self._get(P.COARSE_STAT_D)
-        coarse_init_d = self._get(P.COARSE_INIT_D)
-
-        if coarse_T is None:
-            raise RuntimeError(
-                "Compute coarse-grained transition matrix first as `.compute_macrostates()` with `n_states > 1`."
-            )
-
-        if show_stationary_dist and coarse_stat_d is None:
-            logg.warning("Coarse stationary distribution is `None`, ignoring")
-            show_stationary_dist = False
-        if show_initial_dist and coarse_init_d is None:
-            logg.warning("Coarse initial distribution is `None`, ignoring")
-            show_initial_dist = False
-
-        hrs, wrs = [1], [1]
-        if show_stationary_dist:
-            hrs += [0.05]
-        if show_initial_dist:
-            hrs += [0.05]
-        if show_cbar:
-            wrs += [0.025]
-
-        dont_show_dist = not show_initial_dist and not show_stationary_dist
-
-        fig = plt.figure(constrained_layout=False, figsize=figsize, dpi=dpi)
-        gs = plt.GridSpec(
-            1 + show_stationary_dist + show_initial_dist,
-            1 + show_cbar,
-            height_ratios=hrs,
-            width_ratios=wrs,
-            wspace=0.05,
-            hspace=0.05,
-        )
-        if isinstance(cmap, str):
-            cmap = plt.get_cmap(cmap)
-
-        ax = fig.add_subplot(gs[0, 0])
-        cax = fig.add_subplot(gs[:1, -1]) if show_cbar else None
-        init_ax, stat_ax = None, None
-
-        labels = list(self.coarse_T.columns)
-
-        tmp = coarse_T
-        if show_initial_dist:
-            tmp = np.c_[tmp, coarse_stat_d]
-        if show_initial_dist:
-            tmp = np.c_[tmp, coarse_init_d]
-
-        minn, maxx = np.nanmin(tmp), np.nanmax(tmp)
-        norm = mpl.colors.Normalize(vmin=minn, vmax=maxx)
-
-        if show_stationary_dist:
-            stat_ax = fig.add_subplot(gs[1, 0])
-            stylize_dist(
-                stat_ax,
-                np.array(coarse_stat_d).reshape(1, -1),
-                xticks_labels=labels if not show_initial_dist else None,
-            )
-            stat_ax.yaxis.set_label_position("right")
-            stat_ax.set_ylabel("stationary dist", rotation=0, ha="left", va="center")
-
-        if show_initial_dist:
-            init_ax = fig.add_subplot(gs[show_stationary_dist + show_initial_dist, 0])
-            stylize_dist(
-                init_ax, np.array(coarse_init_d).reshape(1, -1), xticks_labels=labels
-            )
-
-            init_ax.yaxis.set_label_position("right")
-            init_ax.set_ylabel("initial dist", rotation=0, ha="left", va="center")
-
-        im = ax.imshow(coarse_T, aspect="auto", cmap=cmap, norm=norm, **kwargs)
-        ax.set_title("coarse-grained transition matrix" if title is None else title)
-
-        if cax is not None:
-            _ = mpl.colorbar.ColorbarBase(
-                cax,
-                cmap=cmap,
-                norm=norm,
-                ticks=np.linspace(minn, maxx, 10),
-                format="%0.3f",
-            )
-
-        ax.set_yticks(np.arange(coarse_T.shape[0]))
-        ax.set_yticklabels(labels)
-
-        ax.tick_params(
-            top=False,
-            bottom=dont_show_dist,
-            labeltop=False,
-            labelbottom=dont_show_dist,
-        )
-
-        for spine in ax.spines.values():
-            spine.set_visible(False)
-
-        if dont_show_dist:
-            ax.set_xticks(np.arange(coarse_T.shape[1]))
-            ax.set_xticklabels(labels)
-            plt.setp(
-                ax.get_xticklabels(),
-                rotation=xtick_rotation,
-                ha="right",
-                rotation_mode="anchor",
-            )
-        else:
-            ax.set_xticks([])
-
-        ax.set_yticks(np.arange(coarse_T.shape[0] + 1) - 0.5, minor=True)
-        ax.tick_params(which="minor", bottom=dont_show_dist, left=False, top=False)
-
-        if annotate:
-            annotate_heatmap(im)
-            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,
-        cluster_key: Optional[str],
-        en_cutoff: Optional[float],
-        p_thresh: float,
-    ) -> None:
-        start = logg.warning("For 1 macrostate, stationary distribution is computed")
-
-        eig = self._get(P.EIG)
-        if (
-            eig is not None
-            and "stationary_dist" in eig
-            and eig["params"]["which"] == "LR"
-        ):
-            stationary_dist = eig["stationary_dist"]
-        else:
-            self.compute_eigendecomposition(only_evals=False, which="LR")
-            stationary_dist = self._get(P.EIG)["stationary_dist"]
-
-        self._set_macrostates(
-            memberships=stationary_dist[:, None],
-            n_cells=n_cells,
-            cluster_key=cluster_key,
-            p_thresh=p_thresh,
-            en_cutoff=en_cutoff,
-        )
-        self._set(
-            A.MACRO_MEMBER,
-            Lineage(
-                stationary_dist,
-                names=list(self._get(A.MACRO).cat.categories),
-                colors=self._get(A.MACRO_COLORS),
-            ),
-        )
-
-        # reset all the things
-        for key in (
-            A.ABS_PROBS,
-            A.PRIME_DEG,
-            A.SCHUR,
-            A.SCHUR_MAT,
-            A.COARSE_T,
-            A.COARSE_STAT_D,
-            A.COARSE_STAT_D,
-        ):
-            self._set(key.s, None)
-
-        logg.info(
-            f"Adding `.{P.MACRO_MEMBER}`\n        `.{P.MACRO}`\n    Finish",
-            time=start,
-        )
-
-    def _get_n_states_from_minchi(
-        self, n_states: Union[Tuple[int, int], List[int], Dict[str, int]]
-    ) -> int:
-        if self._gpcca is None:
-            raise RuntimeError(
-                "Compute Schur decomposition first as `.compute_schur()` when `use_min_chi=True`."
-            )
-
-        if not isinstance(n_states, (dict, tuple, list)):
-            raise TypeError(
-                f"Expected `n_states` to be either `dict`, `tuple` or a `list`, "
-                f"found `{type(n_states).__name__}`."
-            )
-        if len(n_states) != 2:
-            raise ValueError(
-                f"Expected `n_states` to be of size `2`, found `{len(n_states)}`."
-            )
-
-        if isinstance(n_states, dict):
-            if "min" not in n_states or "max" not in n_states:
-                raise KeyError(
-                    f"Expected the dictionary to have `'min'` and `'max'` keys, "
-                    f"found `{tuple(n_states.keys())}`."
-                )
-            minn, maxx = n_states["min"], n_states["max"]
-        else:
-            minn, maxx = n_states
-
-        if minn > maxx:
-            logg.debug(f"Swapping minimum and maximum because `{minn}` > `{maxx}`")
-            minn, maxx = maxx, minn
-
-        if minn <= 1:
-            raise ValueError(f"Minimum value must be > `1`, found `{minn}`.")
-        elif minn == 2:
-            logg.warning(
-                "In most cases, 2 clusters will always be optimal. "
-                "If you really expect 2 clusters, use `n_states=2` and `use_minchi=False`. Setting minimum to `3`"
-            )
-            minn = 3
-
-        if minn >= maxx:
-            maxx = minn + 1
-            logg.debug(
-                f"Setting maximum to `{maxx}` because it was <= than minimum `{minn}`"
-            )
-
-        logg.info(f"Calculating minChi within interval `[{minn}, {maxx}]`")
-
-        return int(np.arange(minn, maxx + 1)[np.argmax(self._gpcca.minChi(minn, maxx))])
-
-    @d.dedent
-    def _set_macrostates(
-        self,
-        memberships: np.ndarray,
-        n_cells: Optional[int] = 30,
-        cluster_key: str = "clusters",
-        en_cutoff: Optional[float] = 0.7,
-        p_thresh: float = 1e-15,
-        check_row_sums: bool = True,
-    ) -> None:
-        """
-        Map fuzzy clustering to pre-computed annotations to get names and colors.
-
-        Given the fuzzy clustering, we would like to select the most likely cells from each state and use these to
-        give each state a name and a color by comparing with pre-computed, categorical cluster annotations.
-
-        Parameters
-        ----------
-        memberships
-            Fuzzy clustering.
-        %(n_cells)s
-        cluster_key
-            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 :attr:`adata` ``.obs`` as
-            `'G2M_score'` and `'S_score'`.
-        check_row_sums
-            Check whether rows in `memberships` sum to `1`.
-
-        Returns
-        -------
-        None
-            Writes a :class:`cellrank.tl.Lineage` object which mapped names and colors.
-            Also writes a categorical :class:`pandas.Series`, where top ``n_cells`` cells represent each fuzzy state.
-        """
-
-        if n_cells is None:
-            logg.debug("Setting the macrostates using macrostate assignment")
-
-            # fmt: off
-            max_assignment = np.argmax(memberships, axis=1)
-            _macro_assignment = pd.Series(index=self.adata.obs_names, data=max_assignment, dtype="category")
-            # sometimes, the assignment can have a missing category and the Lineage creation therefore fails
-            # keep it as ints when `n_cells != None`
-            _macro_assignment = _macro_assignment.cat.set_categories(list(range(memberships.shape[1])))
-            macrostates = _macro_assignment.astype(str).astype("category").copy()
-            not_enough_cells = []
-            # fmt: on
-        else:
-            logg.debug("Setting the macrostates using macrostates memberships")
-
-            # select the most likely cells from each macrostate
-            macrostates, not_enough_cells = self._create_states(
-                memberships,
-                n_cells=n_cells,
-                check_row_sums=check_row_sums,
-                return_not_enough_cells=True,
-            )
-            not_enough_cells = not_enough_cells.astype("str")
-
-        # _set_categorical_labels creates the names, we still need to remap the group names
-        orig_cats = macrostates.cat.categories
-        self._set_categorical_labels(
-            attr_key=A.MACRO.v,
-            color_key=A.MACRO_COLORS.v,
-            pretty_attr_key=P.MACRO.v,
-            add_to_existing_error_msg="Compute macrostates first as `.compute_macrostates()`.",
-            categories=macrostates,
-            cluster_key=cluster_key,
-            en_cutoff=en_cutoff,
-            p_thresh=p_thresh,
-            add_to_existing=False,
-        )
-
-        name_mapper = dict(zip(orig_cats, self._get(P.MACRO).cat.categories))
-        _print_insufficient_number_of_cells(
-            [name_mapper.get(n, n) for n in not_enough_cells], n_cells
-        )
-        logg.debug("Setting macrostates memberships based on GPCCA membership vectors")
-
-        self._set(
-            A.MACRO_MEMBER,
-            Lineage(
-                memberships,
-                names=list(macrostates.cat.categories),
-                colors=self._get(A.MACRO_COLORS),
-            ),
-        )
-
-    def _create_states(
-        self,
-        probs: Union[np.ndarray, Lineage],
-        n_cells: int,
-        check_row_sums: bool = False,
-        return_not_enough_cells: bool = False,
-    ) -> pd.Series:
-        if n_cells <= 0:
-            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
-
-        a_discrete, not_enough_cells = _fuzzy_to_discrete(
-            a_fuzzy=probs,
-            n_most_likely=n_cells,
-            remove_overlap=False,
-            raise_threshold=0.2,
-            check_row_sums=check_row_sums,
-        )
-
-        states = _series_from_one_hot_matrix(
-            membership=a_discrete,
-            index=self.adata.obs_names,
-            names=probs.names if isinstance(probs, Lineage) else None,
-        )
-
-        return (states, not_enough_cells) if return_not_enough_cells else states
-
-    def _check_states_validity(self, n_states: int) -> int:
-        if self._invalid_n_states is not None and n_states in self._invalid_n_states:
-            logg.warning(
-                f"Unable to compute macrostates with `n_states={n_states}` because it will "
-                f"split the conjugate eigenvalues. Increasing `n_states` to `{n_states + 1}`"
-            )
-            n_states += 1  # cannot force recomputation of the Schur decomposition
-            assert n_states not in self._invalid_n_states, "Sanity check failed."
-
-        return n_states
-
-    def _fit_terminal_states(
-        self,
-        n_lineages: Optional[int] = None,
-        cluster_key: Optional[str] = None,
-        method: str = "krylov",
-        **kwargs,
-    ) -> None:
-        if n_lineages is None or n_lineages == 1:
-            self.compute_eigendecomposition()
-            if n_lineages is None:
-                n_lineages = self.eigendecomposition["eigengap"] + 1
-
-        if n_lineages > 1:
-            self.compute_schur(n_lineages, method=method)
-
-        try:
-            self.compute_macrostates(
-                n_states=n_lineages, cluster_key=cluster_key, **kwargs
-            )
-        except ValueError:
-            logg.warning(
-                f"Computing `{n_lineages}` macrostates cuts through a block of complex conjugates. "
-                f"Increasing `n_lineages` to {n_lineages + 1}"
-            )
-            self.compute_macrostates(
-                n_states=n_lineages + 1, cluster_key=cluster_key, **kwargs
-            )
-
-        fs_kwargs = {"n_cells": kwargs["n_cells"]} if "n_cells" in kwargs else {}
-
-        if n_lineages is None:
-            self.compute_terminal_states(method="eigengap", **fs_kwargs)
-        else:
-            self.set_terminal_states_from_macrostates(**fs_kwargs)
-
-    @d.dedent  # because of fit
-    @d.dedent
-    @inject_docs(
-        ms=P.MACRO,
-        msp=P.MACRO_MEMBER,
-        fs=P.TERM,
-        fsp=P.TERM_PROBS,
-        ap=P.ABS_PROBS,
-        pd=P.PRIME_DEG,
-    )
-    def fit(
-        self,
-        n_lineages: Optional[int] = None,
-        cluster_key: Optional[str] = None,
-        keys: Optional[Sequence[str]] = None,
-        method: str = "krylov",
-        compute_absorption_probabilities: bool = True,
-        **kwargs,
-    ):
-        """
-        Run the pipeline, computing the macrostates, %(initial_or_terminal)s states \
-        and optionally the absorption probabilities.
-
-        It is equivalent to running::
-
-            if n_lineages is None or n_lineages == 1:
-                compute_eigendecomposition(...)  # get the stationary distribution
-            if n_lineages > 1:
-                compute_schur(...)
-
-            compute_macrostates(...)
-
-            if n_lineages is None:
-                compute_terminal_states(...)
-            else:
-                set_terminal_states_from_macrostates(...)
-
-            if compute_absorption_probabilities:
-                compute_absorption_probabilities(...)
-
-        Parameters
-        ----------
-        %(fit)s
-        method
-            Method to use when computing the Schur decomposition. Valid options are: `'krylov'` or `'brandts'`.
-        compute_absorption_probabilities
-            Whether to compute the absorption probabilities or only the %(initial_or_terminal)s states.
-        kwargs
-            Keyword arguments for :meth:`cellrank.tl.estimators.GPCCA.compute_macrostates`.
-
-        Returns
-        -------
-        None
-            Nothing, just makes available the following fields:
-
-                - :attr:`{msp}`
-                - :attr:`{ms}`
-                - :attr:`{fsp}`
-                - :attr:`{fs}`
-                - :attr:`{ap}`
-                - :attr:`{pd}`
-        """
-
-        super().fit(
-            n_lineages=n_lineages,
-            cluster_key=cluster_key,
-            keys=keys,
-            method=method,
-            compute_absorption_probabilities=compute_absorption_probabilities,
-            **kwargs,
-        )
-
-    @d.dedent
-    def _compute_initial_states(self, n_states: int = 1, n_cells: int = 30) -> None:
-        """
-        Compute initial states from macrostates.
-
-        Parameters
-        ----------
-        n_states
-            Number of initial states.
-        %(n_cells)s
-
-        Returns
-        -------
-        %(set_initial_states_from_macrostates.returns)s
-        """
-
-        if n_states <= 0:
-            raise ValueError(f"Expected `n_states` to be positive, found `{n_states}`.")
-
-        if n_cells <= 0:
-            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
-
-        probs = self._get(P.MACRO_MEMBER)
-
-        if probs is None:
-            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
-
-        if n_states > probs.shape[1]:
-            raise ValueError(
-                f"Requested `{n_states}` initial states, but only `{probs.shape[1]}` macrostates have been computed."
-            )
-
-        if probs.shape[1] == 1:
-            self._set_initial_states_from_macrostates(n_cells=n_cells)
-            return
-
-        stat_dist = self._get(P.COARSE_STAT_D)
-        if stat_dist is None:
-            raise RuntimeError("No coarse-grained stationary distribution found.")
-
-        self._set_initial_states_from_macrostates(
-            stat_dist[np.argsort(stat_dist)][:n_states].index, n_cells=n_cells
-        )
-
-    @d.get_sections(base="set_initial_states_from_macrostates", sections=["Returns"])
-    @d.dedent
-    @inject_docs(
-        key=TermStatesKey.BACKWARD.s, probs_key=_probs(TermStatesKey.BACKWARD.s)
-    )
-    def _set_initial_states_from_macrostates(
-        self,
-        names: Optional[Union[Iterable[str], str]] = None,
-        n_cells: int = 30,
-    ) -> None:
-        """
-        Manually select initial states from macrostates.
-
-        Note that no check is performed to ensure initial and terminal states are distinct.
-
-        Parameters
-        ----------
-        names
-            Names of the macrostates to be marked as initial states. Multiple states can be combined using `','`,
-            such as `["Alpha, Beta", "Epsilon"]`.
-        %(n_cells)s
-
-        Returns
-        -------
-        None
-            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.
-        """
-
-        if not isinstance(n_cells, int):
-            raise TypeError(
-                f"Expected `n_cells` to be of type `int`, found `{type(n_cells).__name__!r}`."
-            )
-
-        if n_cells <= 0:
-            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
-
-        probs = self._get(P.MACRO_MEMBER)
-
-        if probs is None:
-            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
-        elif probs.shape[1] == 1:
-            categorical = self._create_states(probs, n_cells=n_cells)
-            scaled = probs / probs.max()
-        else:
-            if names is None:
-                names = probs.names
-            if isinstance(names, str):
-                names = [names]
-
-            probs = probs[list(names)]
-            categorical = self._create_states(probs, n_cells=n_cells)
-            probs /= probs.max(0)
-
-            # compute the aggregated probability of being a initial/terminal state (no matter which)
-            scaled = probs.X.max(1)
-
-        self._write_initial_states(membership=probs, probs=scaled, cats=categorical)
-
-    def _write_initial_states(
-        self, membership: Lineage, probs: pd.Series, cats: pd.Series, time=None
-    ) -> None:
-        key = TermStatesKey.BACKWARD.s
-
-        self.adata.obs[key] = cats
-        self.adata.obs[_probs(key)] = probs
-
-        self.adata.uns[_colors(key)] = membership.colors
-        self.adata.uns[_lin_names(key)] = membership.names
-
-        logg.info(
-            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
deleted file mode 100644
index bb1feaaa6..000000000
--- a/cellrank/tl/estimators/_property.py
+++ /dev/null
@@ -1,924 +0,0 @@
-"""Base properties used within the estimators."""
-from typing import Any, Dict, List, Tuple, Union, Iterable, Optional
-
-import sys
-from abc import ABC, ABCMeta, abstractmethod
-from inspect import isabstract
-
-import scvelo as scv
-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
-from cellrank.tl._colors import _create_categorical_colors
-from cellrank.tl.kernels import PrecomputedKernel
-from cellrank.tl._lineage import Lineage
-from cellrank.tl._constants import Direction, DirPrefix, DirectionPlot
-from cellrank.tl.kernels._utils import _filter_kwargs
-from cellrank.tl.estimators._utils import (
-    Metadata,
-    Sequence,
-    _create_property,
-    singledispatchmethod,
-    _delegate_method_dispatch,
-)
-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
-    """
-    Check whether class with a given name inside this module is abstract.
-
-    Parameters
-    ----------
-    classname
-        Name of the class.
-
-    Returns
-    -------
-    bool
-        `True` if the class is abstract, otherwise `False`.
-
-    """
-
-    cls = getattr(sys.modules[__name__], classname, None)
-    return cls is not None and isabstract(cls)
-
-
-class PropertyMeta(ABCMeta, type):
-    """Metaclass for all the properties."""
-
-    @staticmethod
-    def update_attributes(md: Metadata, attributedict: Dict[str, Any]) -> Optional[str]:
-        """
-        Update ``attributedict`` with new attribute and property.
-
-        Parameters
-        ----------
-        md
-            Metadata object, containing attribute name, property names, etc.
-        attributedict
-            Dictionary with attributes.
-
-        Returns
-        -------
-        str or None
-            Old property name, if found, otherwise a newly constructed one, or `None` if no property is desired.
-        """
-
-        # TODO: determine whether supporting strings is such a good idea
-        if not isinstance(md, Metadata):
-            raise TypeError(
-                f"Expected a `Metadata` object, found `{type(md).__name__}`."
-            )
-
-        if not isinstance(md.attr, (str, A)):
-            raise TypeError(
-                f"Attribute `{md.attr}` must be of type `A` or `str`, found `{type(md.attr).__name__!r}`."
-            )
-        if not str(md.attr).startswith("_"):
-            raise ValueError(f"Attribute `{md.attr!r}` must start with `'_'`.")
-
-        attributedict[str(md.attr)] = md.default
-
-        if md.prop == P.NO_PROPERTY:
-            return
-
-        if not isinstance(md.prop, (str, P)):
-            raise TypeError(
-                f"Property must be of type `P` or `str`, found `{type(md.attr).__name__!r}`."
-            )
-
-        prop_name = str(md.attr).lstrip("_") if md.prop == P.EMPTY else str(md.prop)
-
-        if not len(prop_name):
-            raise ValueError(f"Property name for attribute `{md.attr}` is empty.")
-        if prop_name.startswith("_"):
-            raise ValueError(
-                f"Property mustn't start with an underscore: `{prop_name!r}`."
-            )
-
-        attributedict[prop_name] = _create_property(
-            str(md.attr), prop_name, doc=md.doc, return_type=md.dtype
-        )
-
-        return prop_name
-
-    def __new__(cls, clsname, superclasses, attributedict):
-        """
-        Create a new instance.
-
-        Parameters
-        ----------
-        clsname
-            Name of class to be constructed.
-        superclasses
-            List of superclasses.
-        attributedict
-            Dictionary of attributes.
-        """
-
-        compute_md, metadata = attributedict.pop(META_KEY, None), []
-
-        if compute_md is None:
-            return super().__new__(cls, clsname, superclasses, attributedict)
-
-        if isinstance(compute_md, str):
-            compute_md = Metadata(attr=compute_md)
-        elif not isinstance(compute_md, (tuple, list)):
-            raise TypeError(
-                f"Expected property metadata to be `list` or `tuple`,"
-                f"found `{type(compute_md).__name__!r}`."
-            )
-        elif len(compute_md) == 0:
-            raise ValueError("No metadata found.")
-        else:
-            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)
-
-        if compute_md.compute_fmt != F.NO_FUNC:
-            if "_compute" in attributedict:
-                attributedict[
-                    str(compute_md.compute_fmt).format(prop_name)
-                ] = attributedict["_compute"]
-
-        if (
-            compute_md.plot_fmt != F.NO_FUNC
-            and VectorPlottable in superclasses
-            and plot_name not in attributedict
-            and not is_abstract(clsname)
-        ):
-            raise TypeError(
-                f"Method `{plot_name}` is not implemented for class `{clsname}`."
-            )
-
-        for md in metadata:
-            PropertyMeta.update_attributes(md, attributedict)
-
-        res = super().__new__(cls, clsname, superclasses, attributedict)
-
-        if compute_md.plot_fmt != F.NO_FUNC and Plottable in res.mro():
-            # _this is intended singledispatchmethod
-            # unfortunately, `_plot` is not always in attributedict, so we can't just check for it
-            # and res._plot is just a regular function
-            # if this gets buggy in the future, consider switching from singlemethoddispatch
-            setattr(
-                res,
-                plot_name,
-                _delegate_method_dispatch(res._plot, "_plot", prop_name, skip=2),
-            )
-
-        return res
-
-
-class Property(ABC, metaclass=PropertyMeta):
-    """Base class for all the properties."""
-
-
-class KernelHolder(ABC):
-    """Base class which holds a :class:`cellrank.tool.kernels._kernel.KernelExpression`."""
-
-    def __init__(
-        self,
-        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
-        key: Optional[str] = None,
-        obsp_key: Optional[str] = None,
-        write_to_adata: bool = True,
-    ):
-        if isinstance(obj, KernelExpression):
-            self._kernel = obj
-        elif isinstance(obj, (np.ndarray, spmatrix)):
-            self._kernel = PrecomputedKernel(obj)
-        elif isinstance(obj, AnnData):
-            if obsp_key is None:
-                raise ValueError(
-                    "Specify `obsp_key=...` when supplying an `AnnData` object."
-                )
-            elif obsp_key not in obj.obsp.keys():
-                raise KeyError(f"Key `{obsp_key!r}` not found in `adata.obsp`.")
-            self._kernel = PrecomputedKernel(obj.obsp[obsp_key], adata=obj)
-        else:
-            raise TypeError(
-                f"Expected an object of type `KernelExpression`, `numpy.ndarray`, `scipy.sparse.spmatrix` "
-                f"or `anndata.AnnData`, got `{type(obj).__name__!r}`."
-            )
-
-        if self.kernel._transition_matrix is None:
-            # access the private attribute to avoid accidentally computing the transition matrix
-            # in principle, it doesn't make a difference, apart from not seeing the message
-            logg.warning("Computing transition matrix using the default parameters")
-            self.kernel.compute_transition_matrix()
-
-        if write_to_adata:
-            self.kernel.write_to_adata(key=key)
-
-    @property
-    def _direction(self):
-        return Direction.BACKWARD if self.kernel.backward else Direction.FORWARD
-
-    @property
-    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
-        """Transition matrix."""
-        return self.kernel.transition_matrix
-
-    @property
-    def issparse(self) -> bool:
-        """Whether the transition matrix is sparse or not."""
-        return issparse(self.transition_matrix)
-
-    @property
-    def kernel(self) -> KernelExpression:
-        """Underlying kernel."""
-        return self._kernel
-
-    @property
-    @d.dedent
-    def adata(self) -> AnnData:
-        """
-        Annotated data object.
-
-        Returns
-        -------
-        %(adata_ret)s
-        """
-        return self.kernel.adata
-
-    def __len__(self):
-        return self.kernel.transition_matrix.shape[0]
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}[n={len(self)}, kernel={repr(self.kernel)}]"
-
-    def __str__(self) -> str:
-        return f"{self.__class__.__name__}[n={len(self)}, kernel={str(self.kernel)}]"
-
-
-class VectorPlottable(KernelHolder, Property):
-    """
-    Injector class which plots vectors.
-
-    To be used in conjunction with:
-
-        - :class:`cellrank.tool.estimators._decomposition.Eig`
-        - :class:`cellrank.tool.estimators._decomposition.Schur`
-    """
-
-    @d.get_sections(base="plot_vectors", sections=["Parameters", "Returns"])
-    @d.dedent
-    def _plot_vectors(
-        self,
-        vectors: Optional[np.ndarray],
-        prop: str,
-        use: Optional[Union[int, Tuple[int], List[int]]] = None,
-        abs_value: bool = False,
-        cluster_key: Optional[str] = None,
-        **kwargs,
-    ) -> None:
-        """
-        Plot vectors in an embedding.
-
-        Parameters
-        ----------
-        use
-            Which or how many vectors are to be plotted.
-        abs_value
-            Whether to take the absolute value before plotting.
-        cluster_key
-            Key in :attr:`adata` ``.obs`` for plotting categorical observations.
-        %(basis)s
-        kwargs
-            Keyword arguments for :func:`scvelo.pl.scatter`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        if prop not in (P.EIG.v, P.SCHUR.v):
-            raise ValueError(
-                f"Invalid kind `{prop!r}`. Valid options are `{P.EIG!r}`, `{P.SCHUR!r}``."
-            )
-        if vectors is None:
-            raise RuntimeError(f"Compute `.{prop}` first as `{F.COMPUTE.fmt(prop)}()`.")
-
-        if prop == P.SCHUR.s:
-            is_schur = True
-            name = "Schur "
-        else:
-            is_schur = False
-            name = "eigen"
-
-        # check whether dimensions are consistent
-        if self.adata.n_obs != vectors.shape[0]:
-            raise ValueError(
-                f"Number of cells ({self.adata.n_obs}) is inconsistent with the first"
-                f"dimension of vectors ({vectors.shape[0]})."
-            )
-
-        if use is None:
-            m = (
-                getattr(self, P.EIG.s).get("eigengap", vectors.shape[1]) + 1
-                if hasattr(self, P.EIG.s) and not is_schur
-                else vectors.shape[1]
-            )
-            use = list(range(is_schur, m))
-        elif isinstance(use, int):
-            use = list(range(is_schur, use + is_schur))
-        elif not isinstance(use, (tuple, list, range)):
-            raise TypeError(
-                f"Argument `use` must be either `int`, `tuple`, `list` or `range`,"
-                f"found `{type(use).__name__}`."
-            )
-        else:
-            if not all(map(lambda u: isinstance(u, int), use)):
-                raise TypeError("Not all values in `use` argument are integers.")
-        use = list(use)
-        if not use:
-            raise ValueError("Nothing to plot.")
-
-        muse = max(use)
-        if muse >= vectors.shape[1]:
-            raise ValueError(
-                f"Maximum specified {name}vector ({muse}) is larger "
-                f"than the number of computed {name}vectors ({vectors.shape[1]})."
-            )
-        V_ = vectors[:, use]
-
-        if is_schur:
-            title = [f"{name}vector {i}" for i in use]
-        else:
-            D = kwargs.pop("D")
-            V_ = _complex_warning(V_, use, use_imag=kwargs.pop("use_imag", False))
-            title = [fr"$\lambda_{i}$={d:.02f}" for i, d in zip(use, D[use])]
-
-        if abs_value:
-            V_ = np.abs(V_)
-
-        color = list(V_.T)
-        if cluster_key is not None:
-            color = [cluster_key] + color
-        cmap = kwargs.pop("cmap", "viridis")
-
-        logg.debug(f"Plotting `{use}` {name}vectors")
-
-        scv.pl.scatter(self.adata, color=color, title=title, cmap=cmap, **kwargs)
-
-
-class Plottable(KernelHolder, Property):
-    """
-    Injector which plots macrostates or terminal states or absorption probabilities.
-
-    To be used in conjunction with:
-
-        - :class:`cellrank.tool.estimators._property.MacroStates`.
-        - :class:`cellrank.tool.estimators._property.TermStates`.
-        - :class:`cellrank.tool.estimators._property.AbsProbs`.
-    """
-
-    @d.get_sections(base="plot_discrete")
-    @d.dedent
-    def _plot_discrete(
-        self,
-        data: pd.Series,
-        prop: str,
-        lineages: Optional[Union[str, Sequence[str]]] = None,
-        cluster_key: Optional[str] = None,
-        same_plot: bool = True,
-        title: Optional[Union[str, List[str]]] = None,
-        **kwargs,
-    ) -> None:
-        """
-        Plot the states for each uncovered lineage.
-
-        Parameters
-        ----------
-        lineages
-            Plot only these lineages. If `None`, plot all lineages.
-        cluster_key
-            Key from :attr:`adata` ``.obs`` for plotting categorical observations.
-        same_plot
-            Whether to plot the lineages on the same plot or separately.
-        title
-            The title of the plot.
-        %(basis)s
-        kwargs
-            Keyword arguments for :func:`scvelo.pl.scatter`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        if data is None:
-            raise RuntimeError(
-                f"Compute `.{prop}` first as `.{F.COMPUTE.fmt(prop)}()`."
-            )
-        if not is_categorical_dtype(data):
-            raise TypeError(
-                f"Expected property `.{prop}` to be categorical, found `{type(data).__name__!r}`."
-            )
-        if prop in (P.ABS_PROBS.s, P.TERM.s):
-            colors = getattr(self, A.TERM_COLORS.v, None)
-        elif prop == P.MACRO.v:
-            colors = getattr(self, A.MACRO_COLORS.v, None)
-        else:
-            logg.debug("No colors found. Creating new ones")
-            colors = _create_categorical_colors(len(data.cat.categories))
-        colors = dict(zip(data.cat.categories, colors))
-
-        if (
-            lineages is not None
-        ):  # these are states per-se, but I want to keep the arg names for dispatch the same
-            if isinstance(lineages, str):
-                lineages = [lineages]
-            for state in lineages:
-                if state not in data.cat.categories:
-                    raise ValueError(
-                        f"Invalid state `{state!r}`. Valid options are `{list(data.cat.categories)}`."
-                    )
-            data = data.copy()
-            to_remove = list(set(data.cat.categories) - set(lineages))
-
-            if len(to_remove) == len(data.cat.categories):
-                raise RuntimeError(
-                    "Nothing to plot because empty subset has been selected."
-                )
-
-            for state in to_remove:
-                data[data == state] = np.nan
-            data = data.cat.remove_categories(to_remove)
-
-        if cluster_key is None:
-            cluster_key = []
-        elif isinstance(cluster_key, str):
-            cluster_key = [cluster_key]
-        if not isinstance(cluster_key, list):
-            cluster_key = list(cluster_key)
-
-        same_plot = same_plot or len(data.cat.categories) == 1
-        kwargs["legend_loc"] = kwargs.get("legend_loc", "on data")
-
-        with RandomKeys(
-            self.adata, None if same_plot else len(data.cat.categories), where="obs"
-        ) as keys:
-            if same_plot:
-                key = keys[0]
-                self.adata.obs[key] = data
-                self.adata.uns[f"{key}_colors"] = [
-                    colors[c] for c in data.cat.categories
-                ]
-
-                if title is None:
-                    title = (
-                        f"{prop.replace('_', ' ')} "
-                        f"({Direction.BACKWARD if self.kernel.backward else Direction.FORWARD})"
-                    )
-                if isinstance(title, str):
-                    title = [title]
-
-                scv.pl.scatter(
-                    self.adata,
-                    title=cluster_key + title,
-                    color=cluster_key + keys,
-                    **_filter_kwargs(scv.pl.scatter, **kwargs),
-                )
-            else:
-                for key, cat in zip(keys, data.cat.categories):
-                    d = data.copy()
-                    d[data != cat] = None
-                    d = d.cat.set_categories([cat])
-
-                    self.adata.obs[key] = d
-                    self.adata.uns[f"{key}_colors"] = [colors[cat]]
-
-                scv.pl.scatter(
-                    self.adata,
-                    color=cluster_key + keys,
-                    title=(
-                        cluster_key
-                        + [
-                            f"{_initial if self.kernel.backward else _terminal} state {c}"
-                            for c in data.cat.categories
-                        ]
-                    )
-                    if title is None
-                    else title,
-                    **_filter_kwargs(scv.pl.scatter, **kwargs),
-                )
-
-    @d.get_sections(base="plot_continuous")
-    @d.dedent
-    def _plot_continuous(
-        self,
-        probs: Optional[Lineage],
-        prop: str,
-        lineages: Optional[Union[str, Iterable[str]]] = None,
-        cluster_key: Optional[str] = None,
-        mode: str = "embedding",
-        time_key: str = "latent_time",
-        title: Optional[str] = None,
-        same_plot: bool = True,
-        cmap: Union[str, mpl.colors.ListedColormap] = cm.viridis,
-        **kwargs,
-    ) -> None:
-        """
-        Plot continuous observations such as macrostates memberships or lineages in an embedding.
-
-        Parameters
-        ----------
-        lineages
-            Plot only these lineages. If `None`, plot all lineages.
-        cluster_key
-            Key from :attr:`adata` ``.obs`` for plotting categorical observations.
-        %(time_mode)s
-        time_key
-            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.
-        same_plot
-            Whether to plot the lineages on the same plot using color gradients when ``mode='embedding'``.
-        cmap
-            Colormap to use.
-        %(basis)s
-        kwargs
-            Keyword arguments for :func:`scvelo.pl.scatter`.
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-
-        if probs is None:
-            raise RuntimeError(
-                f"Compute `.{prop}` first as `.{F.COMPUTE.fmt(prop)}()`."
-            )
-
-        if isinstance(lineages, str):
-            lineages = [lineages]
-
-        if lineages is None:
-            lineages = probs.names
-            A = probs
-        else:
-            A = probs[lineages]
-
-        if not len(lineages):
-            raise RuntimeError(
-                "Nothing to plot because empty subset has been selected."
-            )
-
-        prefix = DirPrefix.BACKWARD if self.kernel.backward else DirPrefix.FORWARD
-        same_plot = same_plot and mode == "embedding"  # set this silently
-
-        A = A.copy()  # the below code modifies stuff inplace
-        X = A.X  # list(A.T) behaves differently, because it's Lineage
-
-        if X.shape[1] == 1:
-            same_plot = (
-                False  # because color_gradients for 1 state is buggy (looks empty)
-            )
-            # this is the case for only 1 recurrent class - all cells have prob. 1 of going there
-            # however, matplotlib's plotting really picks up the slightest differences in the colormap, here we set
-            # everything to one, if applicable
-            if np.allclose(X, 1.0):
-                X = np.ones_like(X)
-
-        for col in X.T:
-            mask = ~np.isclose(col, 1.0)
-            # change the maximum value - the 1 is artificial and obscures the color scaling
-            if np.sum(mask):
-                max_not_one = np.max(col[mask])
-                col[~mask] = max_not_one
-
-        if mode == "time":
-            if time_key not in self.adata.obs.keys():
-                raise KeyError(f"Time key `{time_key!r}` not found in `adata.obs`.")
-            time = self.adata.obs[time_key]
-            if cluster_key is not None:
-                logg.warning(
-                    f"Cluster key `{cluster_key!r}` is ignored when `mode='time'`"
-                )
-            cluster_key = None
-
-        color = list(X.T)
-        if title is None:
-            if same_plot:
-                title = [
-                    f"{prop.replace('_', ' ')} "
-                    f"({DirectionPlot.BACKWARD if self.kernel.backward else Direction.FORWARD})"
-                ]
-            else:
-                title = [f"{prefix} {lin}" for lin in lineages]
-        elif isinstance(title, str):
-            title = [title]
-
-        if isinstance(cluster_key, str):
-            cluster_key = [cluster_key]
-        elif cluster_key is None:
-            cluster_key = []
-        if not isinstance(cluster_key, list):
-            cluster_key = list(cluster_key)
-
-        if not same_plot:
-            color = cluster_key + color
-            title = cluster_key + title
-
-        if mode == "embedding":
-            if same_plot:
-                # to complement: https://github.com/theislab/scvelo/blob/master/scvelo/plotting/scatter.py#L269
-                sorted_idx = np.argsort(X, axis=1)[:, ::-1][:, :2]
-                pairs, cnts = np.unique(
-                    np.sort(sorted_idx, axis=0), axis=0, return_counts=True
-                )
-                absent = set(A.names) - set(A.names[pairs[cnts > 1].flatten()])
-                if absent:
-                    # can't print which because it would be inaccurate, i.e. we would print the legend
-                    # is missing when in fact, it would be visible
-                    # this has been checked exactly according to scVelo's logic
-                    logg.warning("Legend for some lineages may be missing")
-
-                kwargs["color_gradients"] = A
-                if len(cluster_key):
-                    logg.warning(
-                        "Ignoring `cluster_key` when plotting continuous observations in the same plot"
-                    )
-                # kwargs["color"] = cluster_key  this results in a bug, cluster_key data is overwritten, will make a PR
-            else:
-                kwargs["color"] = color
-
-            if probs.shape[1] == 1 and prop in (P.MACRO_MEMBER.s, P.TERM_PROBS.s):
-                if "perc" not in kwargs:
-                    logg.warning(
-                        "Did not detect percentile for stationary distribution. Setting `perc=[0, 95]`"
-                    )
-                    kwargs["perc"] = [0, 95]
-                kwargs["color"] = X
-                kwargs.pop("color_gradients", None)
-
-            scv.pl.scatter(
-                self.adata,
-                title=title,
-                color_map=cmap,
-                **_filter_kwargs(scv.pl.scatter, **kwargs),
-            )
-        elif mode == "time":
-            scv.pl.scatter(
-                self.adata,
-                x=time,
-                color_map=cmap,
-                y=color,
-                title=title,
-                xlabel=[time_key] * len(title),
-                ylabel=["probability"] * len(title),
-                **_filter_kwargs(scv.pl.scatter, **kwargs),
-            )
-        else:
-            raise ValueError(
-                f"Invalid mode `{mode!r}`. Valid options are: `'embedding'` or `'time'`."
-            )
-
-    @singledispatchmethod
-    @d.dedent
-    def _plot(
-        self,
-        data,
-        prop: str,
-        discrete: bool = False,
-        lineages: Optional[Union[str, Sequence[str]]] = None,
-        cluster_key: Optional[str] = None,
-        mode: str = "embedding",
-        time_key: str = "latent_time",
-        title: Optional[str] = None,
-        same_plot: bool = False,
-        cmap: Union[str, mpl.colors.ListedColormap] = "viridis",
-        **kwargs,
-    ) -> None:
-        """
-        Plot discrete states or probabilities in an embedding.
-
-        Parameters
-        ----------
-        discrete
-            Whether to plot in discrete or continuous mode.
-        %(plot_continuous.parameters)s
-
-        Returns
-        -------
-        %(just_plots)s
-        """
-        # we keep all the arguments because from this function, we get the docs + signature
-        raise RuntimeError(f"Compute `.{prop}` first as `.{F.COMPUTE.fmt(prop)}()`.")
-
-    @_plot.register(pd.Series)
-    def _(self, data: pd.Series, prop: str, discrete: bool = False, **kwargs) -> None:
-        if discrete and kwargs.get("mode", "embedding") == "time":
-            logg.warning(
-                "`mode='time'` is implemented in continuous case, plotting in continuous mode"
-            )
-            discrete = False
-
-        if discrete:
-            self._plot_discrete(data, prop, **kwargs)
-        elif prop == P.MACRO.v:  # GPCCA
-            prop = P.MACRO_MEMBER.v
-            self._plot_continuous(getattr(self, prop, None), prop, **kwargs)
-        elif prop == P.TERM.v:
-            probs = getattr(self, A.TERM_ABS_PROBS.s, None)
-            # we have this only in GPCCA
-            if isinstance(probs, Lineage):
-                self._plot_continuous(probs, P.TERM_PROBS.v, **kwargs)
-            else:
-                logg.warning(
-                    f"Unable to plot continuous observations for `{prop!r}`, plotting in discrete mode"
-                )
-                self._plot_discrete(data, prop, **kwargs)
-        else:
-            raise NotImplementedError(
-                f"Unable to plot property `.{prop}` in discrete mode."
-            )
-
-    @_plot.register(Lineage)
-    def _(self, data: Lineage, prop: str, discrete: bool = False, **kwargs) -> None:
-        if discrete and kwargs.get("mode", "embedding") == "time":
-            logg.warning(
-                "`mode='time'` is implemented in continuous case, plotting in continuous mode"
-            )
-            discrete = False
-
-        if not discrete:
-            self._plot_continuous(data, prop, **kwargs)
-        elif prop == P.ABS_PROBS.v:
-            # for discrete and abs. probs, plot the terminal states
-            prop = P.TERM.v
-            self._plot_discrete(getattr(self, prop, None), prop, **kwargs)
-        else:
-            raise NotImplementedError(
-                f"Unable to plot property `.{prop}` in continuous mode."
-            )
-
-
-class Macrostates(Plottable):
-    """Class dealing with macrostates."""
-
-    __prop_metadata__ = [
-        Metadata(attr=A.MACRO, prop=P.MACRO, dtype=pd.Series),
-        Metadata(
-            attr=A.MACRO_MEMBER,
-            prop=P.MACRO_MEMBER,
-            dtype=Lineage,
-        ),
-        Metadata(attr=A.MACRO_COLORS, prop=P.NO_PROPERTY, dtype=np.ndarray),
-    ]
-
-    @abstractmethod
-    def compute_macrostates(self, *args, **kwargs) -> None:  # noqa
-        pass
-
-
-class TerminalStates(Plottable):
-    """Class dealing with terminal states."""
-
-    __prop_metadata__ = [
-        Metadata(attr=A.TERM, prop=P.TERM, dtype=pd.Series, doc="Terminal states."),
-        Metadata(
-            attr=A.TERM_PROBS,
-            prop=P.TERM_PROBS,
-            dtype=pd.Series,
-            doc="Terminal states probabilities.",
-        ),
-        Metadata(attr=A.TERM_COLORS, prop=P.NO_PROPERTY, dtype=np.ndarray),
-    ]
-
-    @abstractmethod
-    def set_terminal_states(self, *args: Any, **kwargs: Any) -> None:  # noqa
-        pass
-
-    @abstractmethod
-    def compute_terminal_states(self, *args: Any, **kwargs: Any) -> None:  # noqa
-        pass
-
-    @abstractmethod
-    def _write_terminal_states(self, *args: Any, **kwargs: Any) -> None:
-        pass
-
-
-class AbsProbs(Plottable):
-    """Class dealing with absorption probabilities."""
-
-    __prop_metadata__ = [
-        Metadata(
-            attr=A.ABS_PROBS,
-            prop=P.ABS_PROBS,
-            dtype=Lineage,
-            doc="Absorption probabilities.",
-        ),
-        Metadata(
-            attr=A.PRIME_DEG,
-            prop=P.PRIME_DEG,
-            dtype=pd.Series,
-            doc="Priming degree.",
-        ),
-        Metadata(attr=A.LIN_ABS_TIMES, prop=P.LIN_ABS_TIMES, dtype=pd.DataFrame),
-    ]
-
-    @abstractmethod
-    def _write_absorption_probabilities(self, *args, **kwargs) -> None:
-        pass
-
-
-class LinDrivers(Plottable):  # noqa
-    __prop_metadata__ = [
-        Metadata(
-            attr=A.LIN_DRIVERS,
-            prop=P.LIN_DRIVERS,
-            dtype=pd.DataFrame,
-            doc="Lineage drivers.",
-            plot_fmt=F.NO_FUNC,
-            # in essence ignore Plottable (could be done by registering DataFrame, but it's ugly
-        )
-    ]
-
-
-class Partitioner(KernelHolder, ABC):
-    """Abstract base class for partitioning transition matrix into sets of recurrent and transient states."""
-
-    def __init__(self, *args, **kwargs):
-        super().__init__(*args, **kwargs)
-
-        self._is_irreducible = None  # no need to refactor these
-        self._rec_classes = None
-        self._trans_classes = None
-
-    def compute_partition(self) -> None:
-        """
-        Compute communication classes for the Markov chain.
-
-        Returns
-        -------
-        None
-            Nothing, but updates the following fields:
-
-                - :attr:`recurrent_classes`
-                - :attr:`transient_classes`
-                - :attr:`is_irreducible`
-        """
-
-        start = logg.info("Computing communication classes")
-        n_states = len(self)
-
-        rec_classes, trans_classes = _partition(self.transition_matrix)
-
-        self._is_irreducible = len(rec_classes) == 1 and len(trans_classes) == 0
-
-        if not self._is_irreducible:
-            self._trans_classes = _make_cat(
-                trans_classes, n_states, self.adata.obs_names
-            )
-            self._rec_classes = _make_cat(rec_classes, n_states, self.adata.obs_names)
-            logg.info(
-                f"Found `{(len(rec_classes))}` recurrent and `{len(trans_classes)}` transient classes\n"
-                f"Adding `.recurrent_classes`\n"
-                f"       `.transient_classes`\n"
-                f"       `.is_irreducible`\n"
-                f"    Finish",
-                time=start,
-            )
-        else:
-            logg.warning(
-                "The transition matrix is irreducible, cannot further partition it\n    Finish",
-                time=start,
-            )
-
-    @property
-    def is_irreducible(self):
-        """Whether the Markov chain is irreducible or not."""
-        return self._is_irreducible
-
-    @property
-    def recurrent_classes(self):
-        """Recurrent classes of the Markov chain."""
-        return self._rec_classes
-
-    @property
-    def transient_classes(self):
-        """Transient classes of the Markov chain."""
-        return self._trans_classes
-
-
-class LineageEstimatorMixin(TerminalStates, AbsProbs, LinDrivers, Pickleable, ABC):
-    """Mixin containing terminal states and absorption probabilities."""
diff --git a/cellrank/tl/estimators/_utils.py b/cellrank/tl/estimators/_utils.py
index d745bac6c..f6efe8fe8 100644
--- a/cellrank/tl/estimators/_utils.py
+++ b/cellrank/tl/estimators/_utils.py
@@ -1,176 +1,61 @@
-"""Utility functions which deal with delegating methods."""
-from typing import *  # noqa
+from typing import Any, Union, TypeVar, Optional, Sequence
 
-import wrapt
-import inspect
-from functools import partial, singledispatch, update_wrapper
-from collections import namedtuple
+from copy import copy, deepcopy
+from collections import Iterable
 
-import cellrank.logging as logg
-from cellrank.tl.estimators._constants import F
+from scanpy import logging as logg
 
-AnnData = TypeVar("AnnData")
-try:
-    Metadata = namedtuple(
-        "Metadata",
-        ["attr", "prop", "dtype", "default", "doc", "compute_fmt", "plot_fmt"],
-        defaults=[None, None, None, None, None, F.COMPUTE, F.PLOT],
-    )
-except TypeError:
-    # Python < 3.7
-    Metadata = namedtuple(
-        "Metadata",
-        ["attr", "prop", "dtype", "default", "doc", "compute_fmt", "plot_fmt"],
-    )
-    Metadata.__new__.__defaults__ = defaults = (
-        None,
-        None,
-        None,
-        None,
-        None,
-        F.COMPUTE,
-        F.PLOT,
-    )
 
-
-# copy of functools.singledispatchmethod (for Python < 3.8)
-class singledispatchmethod:
-    """
-    Single-dispatch generic method descriptor.
-
-    Supports wrapping existing descriptors and handles non-descriptor callables as instance methods.
+class SafeGetter:
     """
+    Context manager that does a simple rollback on :attr:`object.__dict__`.
 
-    def __init__(self, func):
-        if not callable(func) and not hasattr(func, "__get__"):
-            raise TypeError(f"{func!r} is not callable or a descriptor")
-
-        self.dispatcher = singledispatch(func)
-        self.func = func
-
-    def register(self, cls, method=None):
-        """
-        Generic_method.register(cls, func) -> func
-
-        Registers a new implementation for the given *cls* on a *generic_method*.
-        """  # noqa
-        return self.dispatcher.register(cls, func=method)
-
-    def __get__(self, obj, cls=None):
-        def _method(*args, **kwargs):
-            method = self.dispatcher.dispatch(args[0].__class__)
-            return method.__get__(obj, cls)(*args, **kwargs)
-
-        _method.__isabstractmethod__ = self.__isabstractmethod__
-        _method.register = self.register
-        update_wrapper(_method, self.func)
-        return _method
-
-    @property
-    def __isabstractmethod__(self):
-        return getattr(self.func, "__isabstractmethod__", False)
-
-
-def argspec_factory(
-    wrapped: Callable,
-    skip: int = 0,
-    return_type: Optional[Type] = None,
-    fn: Optional[Callable] = None,
-) -> Callable:
-    """
-    Create a dummy adapter function with desired signature.
+    Used in :class:`cellrank.estimators.BaseEstimator` to maintain consistency
+    during deserialization from :class:`anndata.AnnData`.
 
     Parameters
     ----------
-    wrapped
-        Function being wrapped whose signature we wish to replicate.
-    skip
-        Number of arguments from left to skip.
-    return_type
-        Return type of the function.
-    fn
-        If not `None`, take the signature from this function.
-
-    Returns
-    -------
-    :class:`Callable`
-        An adapter with the correct signature.
+    obj
+        Object to watch. Shallow copy of ``obj.__dict__`` is made during initialization, with the exception
+        of the `_shadow_adata` (deep copy), if present.
+    allowed
+        Allow these exception to occur with the context.
+
+    Notes
+    -----
+    :attr:`ok` is `False` iff any exception, regardless of whether it is in ``allowed``, has been encountered.
     """
-    # for locals(), this is a whitelist of types we allow
-    import matplotlib  # this one seems to be missing, so whitelist it
-
-    NoneType = type(None)
-
-    sig = inspect.signature(wrapped if fn is None else fn)
-    if return_type is not None:
-        sig._return_annotation = return_type
-
-    keys = list(sig.parameters.keys())[skip:]
-    sig = sig.replace(parameters=[sig.parameters[k] for k in keys])
-
-    exec(f"def adapter{sig}: pass", globals(), locals())
 
-    return locals()["adapter"]
+    def __init__(self, obj: Any, allowed: Union[type, Sequence[type]] = ()):
+        self._exc_type: Optional[TypeVar[Exception]] = None
+        self._obj = obj
+        if not isinstance(allowed, Iterable):
+            allowed = (allowed,)
+        self._allowed = allowed
+
+        self._dict = copy(obj.__dict__)
+        if "_shadow_adata" in self._dict:
+            # deepcopy `_shadow_adata` (small cost) to easily revert changes
+            # do not run deepcopy on the full `__dict__` (possibly large cost)
+            # silent assumption is that within the context `adata` is not modified
+            self._dict["_shadow_adata"] = deepcopy(self._dict["_shadow_adata"])
+
+    def __enter__(self) -> "SafeGetter":
+        return self
+
+    def __exit__(self, exc_type, exc_val, exc_tb) -> bool:
+        self._exc_type = exc_type
+        # TODO(michalk8): log properly
+        if not self.ok:
+            logg.debug(
+                f"The estimator will not be completely initialized, reason: {exc_type(exc_val)}"
+            )
+            self._obj.__dict__ = self._dict
+            return self._exc_type in self._allowed
+
+        return True
 
-
-def _delegate(
-    *,
-    prop_name: Optional[str] = None,
-    return_type: Optional[Type] = None,
-    skip: int = 2,
-) -> Callable:
-    @wrapt.decorator()
-    def pass_through(wrapped, _instance, args, kwargs):
-        return wrapped(*args, **kwargs)
-
-    if prop_name is None:
-        return pass_through
-
-    adapter = wrapt.adapter_factory(
-        partial(argspec_factory, skip=skip, return_type=return_type)
-    )
-
-    @wrapt.decorator(adapter=adapter)
-    def wrapper(wrapped, instance, args, kwargs):
-        return wrapped(getattr(instance, prop_name), prop_name, *args, **kwargs)
-
-    return wrapper
-
-
-def _delegate_method_dispatch(fn: Callable, attr: str, prop_name: str, skip: int = 2):
-
-    adapter = wrapt.adapter_factory(partial(argspec_factory, fn=fn, skip=skip))
-
-    @wrapt.decorator(adapter=adapter)
-    def delegate(wrapped, instance, args, kwargs):
-        # first, get the correct method based on the dispatcher (can't use wrapped.dispatch..., since it's a function)
-        # then, call the actual function
-        return getattr(instance, attr)(
-            getattr(instance, prop_name), prop_name, *args, **kwargs
-        )
-
-    return delegate(fn)
-
-
-def _create_property(
-    attr_name: str,
-    prop_name: str,
-    doc: Optional[str] = None,
-    return_type: Optional[Type] = None,
-) -> property:
-    def wrapper(self) -> return_type:
-        return getattr(self, attr_name)
-
-    if not doc:
-        doc = f"{prop_name.replace('_', ' ').capitalize()}."
-    wrapper.__doc__ = doc
-
-    return property(wrapper)
-
-
-def _print_insufficient_number_of_cells(groups: Iterable[Any], n_cells: int):
-    if groups:
-        logg.debug(
-            f"The following groups have less than requested number of cells ({n_cells}): "
-            f"`{', '.join(sorted(map(str, groups)))}`"
-        )
+    @property
+    def ok(self) -> bool:
+        """Return `True` if no exception has been raised."""
+        return self._exc_type is None
diff --git a/cellrank/tl/estimators/mixins/__init__.py b/cellrank/tl/estimators/mixins/__init__.py
new file mode 100644
index 000000000..d541f14cd
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/__init__.py
@@ -0,0 +1,3 @@
+from cellrank.tl.estimators.mixins.decomposition import EigenMixin, SchurMixin
+from cellrank.tl.estimators.mixins._lineage_drivers import LinDriversMixin
+from cellrank.tl.estimators.mixins._absorption_probabilities import AbsProbsMixin
diff --git a/cellrank/tl/estimators/mixins/_absorption_probabilities.py b/cellrank/tl/estimators/mixins/_absorption_probabilities.py
new file mode 100644
index 000000000..322301939
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/_absorption_probabilities.py
@@ -0,0 +1,511 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
+from typing_extensions import Literal
+
+from types import MappingProxyType
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.tl._enum import _DEFAULT_BACKEND, Backend_t
+from cellrank.ul._docs import d
+from cellrank.tl._utils import (
+    _pairwise,
+    _process_series,
+    _get_cat_and_null_indices,
+    _calculate_lineage_absorption_time_means,
+)
+from cellrank.tl._lineage import Lineage
+from cellrank.tl._linear_solver import _solve_lin_system
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.estimators.mixins._utils import (
+    BaseProtocol,
+    logger,
+    shadow,
+    register_plotter,
+)
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import issparse, spmatrix
+from pandas.api.types import infer_dtype, is_categorical_dtype
+
+
+class AbsProbsProtocol(BaseProtocol):  # noqa: D101
+    _term_states_colors: np.ndarray
+
+    @property
+    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:  # noqa: D102
+        ...
+
+    @property
+    def terminal_states(self) -> pd.Series:  # noqa: D102
+        ...
+
+    @property
+    def absorption_probabilities(self) -> Optional[Lineage]:  # noqa: D102
+        ...
+
+    @property
+    def absorption_times(self) -> Optional[pd.DataFrame]:  # noqa: D102
+        ...
+
+    @property
+    def priming_degree(self) -> Optional[pd.Series]:  # noqa: D102
+        ...
+
+    def __len__(self) -> int:
+        ...
+
+    def _compute_absorption_probabilities(
+        self,
+        q: Union[np.ndarray, spmatrix],
+        s: Union[np.ndarray, spmatrix],
+        trans_indices: np.ndarray,
+        term_states: np.ndim,
+        solver: str,
+        use_petsc: bool,
+        n_jobs: Optional[int],
+        backend: str,
+        tol: float,
+        show_progress_bar: bool,
+        preconditioner: str,
+    ) -> np.ndarray:
+        ...
+
+    def _ensure_lineage_object(self, attr: str, **kwargs: Any) -> None:
+        ...
+
+    def _write_absorption_probabilities(
+        self,
+        abs_probs: Optional[Lineage],
+        abs_times: Optional[pd.DataFrame],
+    ) -> str:
+        ...
+
+    def _write_lineage_priming(
+        self,
+        priming_degree: Optional[pd.Series],
+    ) -> str:
+        ...
+
+
+def _normalize_abs_times(
+    keys: Sequence[str], time_to_absorption: Any = None
+) -> Dict[Tuple[str, ...], Literal["mean", "var"]]:
+    if time_to_absorption is None:
+        return {}
+
+    if isinstance(time_to_absorption, (str, tuple)):
+        time_to_absorption = [time_to_absorption]
+    if not isinstance(time_to_absorption, dict):
+        time_to_absorption = {ln: "mean" for ln in time_to_absorption}
+
+    res = {}
+    for ln, moment in time_to_absorption.items():
+        if moment not in ("mean", "var"):
+            raise ValueError(
+                f"Moment must be either `'mean'` or `'var'`, found `{moment!r}` in `{ln}`."
+            )
+
+        seen = set()
+        if isinstance(ln, str):
+            ln = tuple(keys) if ln == "all" else (ln,)
+        sorted_ln = tuple(sorted(ln))  # preserve the user order
+
+        if sorted_ln not in seen:
+            seen.add(sorted_ln)
+            for lin in ln:
+                if lin not in keys:
+                    raise ValueError(
+                        f"Invalid absorbing state `{lin!r}` in `{ln}`. "
+                        f"Valid options are `{list(keys)}`."
+                    )
+            res[tuple(ln)] = moment
+
+    return res
+
+
+class AbsProbsMixin:
+    """Mixin that supports computation of absorption probabilities and mean times to absorption."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+        self._absorption_probabilities: Optional[Lineage] = None
+        self._absorption_times: Optional[pd.DataFrame] = None
+        self._priming_degree: Optional[pd.Series] = None
+
+    @property
+    @d.get_summary(base="abs_probs")
+    def absorption_probabilities(self) -> Optional[Lineage]:
+        """Absorption probabilities.
+
+        Informally, given a (finite, discrete) Markov chain with a set of transient states :math:`T` and
+        a set of absorbing states :math:`A`, the absorption probability for cell :math:`i` from :math:`T`
+        to reach cell :math:`j` from :math:`R` is the probability that a random walk initialized in :math:`i`
+        will reach absorbing state :math:`j`.
+
+        In our context, states correspond to cells, in particular, absorbing states correspond to cells in terminal
+        states.
+        """
+        return self._absorption_probabilities
+
+    @property
+    @d.get_summary(base="abs_times")
+    def absorption_times(self) -> Optional[pd.DataFrame]:
+        """Mean and variance of the time until absorption.
+
+        Related to conditional mean first passage times. Corresponds to the expectation of the time until absorption,
+        depending on initialization, and the variance.
+        """
+        return self._absorption_times
+
+    @property
+    @d.get_summary(base="priming_degree")
+    def priming_degree(self) -> Optional[pd.Series]:
+        """Priming degree.
+
+        Given a cell :math:`i` and a set of terminal states, this quantifies how committed vs. naive cell :math:`i` is,
+        i.e. its degree of pluripotency. Low values correspond to naive cells (high degree of pluripotency), high values
+        correspond to committed cells (low degree of pluripotency).
+        """
+        return self._priming_degree
+
+    @d.dedent
+    def compute_absorption_probabilities(
+        self: AbsProbsProtocol,
+        keys: Optional[Sequence[str]] = None,
+        solver: Union[
+            str, Literal["direct", "gmres", "lgmres", "bicgstab", "gcrotmk"]
+        ] = "gmres",
+        use_petsc: bool = True,
+        time_to_absorption: Optional[
+            Union[
+                Literal["all"],
+                Sequence[Union[str, Sequence[str]]],
+                Dict[Union[str, Sequence[str]], Literal["mean", "var"]],
+            ]
+        ] = None,
+        n_jobs: Optional[int] = None,
+        backend: Backend_t = _DEFAULT_BACKEND,
+        show_progress_bar: bool = True,
+        tol: float = 1e-6,
+        preconditioner: Optional[str] = None,
+    ) -> None:
+        """
+        Compute absorption probabilities.
+
+        For each cell, this computes the probability of being absorbed in any of the :attr:`terminal_states`. In
+        particular, this corresponds to the probability that a random walk initialized in transient cell :math:`i`
+        will reach any cell from a fixed transient state before reaching a cell from any other transient state.
+
+        Parameters
+        ----------
+        keys
+            Terminal states for which to compute the absorption probabilities.
+            If `None`, use all states defined in :attr:`terminal_states`.
+        solver
+            Solver to use for the linear problem. Options are `'direct', 'gmres', 'lgmres', 'bicgstab' or 'gcrotmk'`
+            when ``use_petsc = False`` or one of :class:`petsc4py.PETSc.KPS.Type` otherwise.
+
+            Information on the :mod:`scipy` iterative solvers can be found in :func:`scipy.sparse.linalg` or for
+            :mod:`petsc4py` solver `here <https://petsc.org/release/overview/linear_solve_table/>`__.
+        use_petsc
+            Whether to use solvers from :mod:`petsc4py` or :mod:`scipy`. Recommended for large problems.
+            If no installation is found, defaults to :func:`scipy.sparse.linalg.gmres`.
+        time_to_absorption
+            Whether to compute mean time to absorption and its variance to specific absorbing states.
+
+            If a :class:`dict`, can be specified as ``{{'Alpha': 'var', ...}}`` to also compute variance.
+            In case when states are a :class:`tuple`, time to absorption will be computed to the subset of these states,
+            such as ``[('Alpha', 'Beta'), ...]`` or ``{{('Alpha', 'Beta'): 'mean', ...}}``.
+            Can be specified as ``'all'`` to compute it to any absorbing state in ``keys``, which is more efficient
+            than listing all absorbing states explicitly.
+
+            It might be beneficial to disable the progress bar as ``show_progress_bar = False`` because of many solves.
+        n_jobs
+            Number of parallel jobs to use when using an iterative solver.
+        backend
+            Which backend to use for multiprocessing. See :class:`joblib.Parallel` for valid options.
+        show_progress_bar
+            Whether to show progress bar. Only used when ``solver != 'direct'``.
+        tol
+            Convergence tolerance for the iterative solver. The default is fine for most cases, only consider
+            decreasing this for severely ill-conditioned matrices.
+        preconditioner
+            Preconditioner to use, only available when ``use_petsc = True``. For valid options, see
+            `here <https://petsc.org/release/docs/manual/ksp/?highlight=pctype#preconditioners>`__.
+            We recommend the `'ilu'` preconditioner for badly conditioned problems.
+
+        Returns
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`absorption_probabilities` - %(abs_probs.summary)s
+            - :attr:`absorption_times` - %(abs_times.summary)s Only if ``time_to_absorption`` is specified.
+        """
+        if self.terminal_states is None:
+            raise RuntimeError(
+                "Compute terminal states first as `.compute_terminal_states()`."
+            )
+        if keys is not None:
+            keys = sorted(set(keys))
+
+        start = logg.info("Computing absorption probabilities")
+
+        # get the transition matrix
+        if not issparse(self.transition_matrix):
+            logg.warning(
+                "Attempting to solve a potentially large linear system with dense transition matrix"
+            )
+
+        # process the current annotations according to `keys`
+        term_states, colors = _process_series(
+            series=self.terminal_states, keys=keys, colors=self._term_states_colors
+        )
+        # warn in case only one state is left
+        keys = list(term_states.cat.categories)
+        if len(keys) == 1:
+            logg.warning(
+                "There is only `1` terminal state, all cells will have probability `1` of going there"
+            )
+
+        # get indices corresponding to recurrent and transient states
+        rec_indices, trans_indices, lookup_dict = _get_cat_and_null_indices(term_states)
+        if not len(trans_indices):
+            raise RuntimeError("Markov chain is irreducible.")
+
+        # create Q (restriction transient-transient), S (restriction transient-recurrent)
+        q = self.transition_matrix[trans_indices, :][:, trans_indices]
+        s = self.transition_matrix[trans_indices, :][:, rec_indices]
+
+        # take individual solutions and piece them together to get absorption probabilities towards the classes
+        # fmt: off
+        macro_ix_helper = np.cumsum([0] + [len(indices) for indices in lookup_dict.values()])
+        s = np.concatenate([s[:, np.arange(a, b)].sum(axis=1) for a, b in _pairwise(macro_ix_helper)], axis=1)
+        # fmt: on
+
+        abs_probs = self._compute_absorption_probabilities(
+            q,
+            s,
+            trans_indices=trans_indices,
+            term_states=term_states,
+            solver=solver,
+            use_petsc=use_petsc,
+            n_jobs=n_jobs,
+            backend=backend,
+            tol=tol,
+            show_progress_bar=show_progress_bar,
+            preconditioner=preconditioner,
+        )
+        abs_probs = Lineage(abs_probs, names=keys, colors=colors)
+
+        abs_times = None
+        if time_to_absorption is not None:
+            lineages = _normalize_abs_times(keys, time_to_absorption=time_to_absorption)
+            abs_times = _calculate_lineage_absorption_time_means(
+                q,
+                s,
+                trans_indices=trans_indices,
+                ixs=lookup_dict,
+                lineages=lineages,
+                solver=solver,
+                use_petsc=use_petsc,
+                n_jobs=n_jobs,
+                backend=backend,
+                tol=tol,
+                show_progress_bar=show_progress_bar,
+                preconditioner=preconditioner,
+                index=self.adata.obs_names,
+            )
+
+        params = self._create_params(
+            remove=["use_petsc", "n_jobs", "backend", "show_progress_bar"]
+        )
+        self._write_absorption_probabilities(
+            abs_probs, abs_times, params=params, time=start
+        )
+
+    @d.dedent
+    def compute_lineage_priming(
+        self: AbsProbsProtocol,
+        method: Literal["kl_divergence", "entropy"] = "kl_divergence",
+        early_cells: Optional[Union[Mapping[str, Sequence[str]], Sequence[str]]] = None,
+    ) -> pd.Series:
+        """
+        %(lin_pd.full_desc)s
+
+        Parameters
+        ----------
+        %(lin_pd.parameters)s
+            If a :class:`dict`, the key specifies a cluster key in :attr:`anndata.AnnData.obs` and the values
+            specify cluster labels containing early cells.
+
+        Returns
+        -------
+        %(lin_pd.returns)s
+
+        Also updates the following field:
+
+            - :attr:`priming_degree` - %(priming_degree.summary)s
+        """  # noqa: D400
+        abs_probs = self.absorption_probabilities
+        if abs_probs is None:
+            raise RuntimeError(
+                "Compute absorption probabilities first as `.compute_absorption_probabilities()`."
+            )
+        if isinstance(early_cells, dict):
+            if len(early_cells) != 1:
+                raise ValueError(
+                    f"Expected a dictionary with only 1 key, found `{len(early_cells)}`."
+                )
+            key = next(iter(early_cells.keys()))
+            if key not in self.adata.obs:
+                raise KeyError(f"Unable to find clusters in `adata.obs[{key!r}]`.")
+            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])}`."
+                )
+            early_cells = self.adata.obs[key].isin(early_cells[key])
+        elif early_cells is not None:
+            early_cells = np.asarray(early_cells)
+            if not np.issubdtype(early_cells.dtype, np.bool_):
+                early_cells = np.isin(self.adata.obs_names, early_cells)
+
+        values = pd.Series(
+            abs_probs.priming_degree(method, early_cells), index=self.adata.obs_names
+        )
+        self._write_lineage_priming(values)
+
+        return values
+
+    def _compute_absorption_probabilities(
+        self: AbsProbsProtocol,
+        q: Union[np.ndarray, spmatrix],
+        s: Union[np.ndarray, spmatrix],
+        trans_indices: np.ndarray,
+        term_states: np.ndim,
+        solver: str,
+        use_petsc: bool,
+        n_jobs: Optional[int],
+        backend: str,
+        tol: float,
+        show_progress_bar: bool,
+        preconditioner: str,
+    ) -> np.ndarray:
+        _abs_classes = _solve_lin_system(
+            q,
+            s,
+            solver=solver,
+            use_petsc=use_petsc,
+            n_jobs=n_jobs,
+            backend=backend,
+            tol=tol,
+            use_eye=True,
+            show_progress_bar=show_progress_bar,
+            preconditioner=preconditioner,
+        )
+        abs_classes = np.zeros(
+            shape=(len(self), len(term_states.cat.categories)), dtype=np.float64
+        )
+        for col, rec_class in enumerate(term_states.cat.categories):
+            rec_indices = np.where(term_states == rec_class)[0]
+            abs_classes[trans_indices, col] = _abs_classes[:, col]
+            abs_classes[rec_indices, col] = 1.0
+
+        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."
+            )
+
+        return abs_classes
+
+    @logger
+    @shadow
+    def _write_absorption_probabilities(
+        self: AbsProbsProtocol,
+        abs_probs: Optional[Lineage],
+        abs_times: Optional[pd.DataFrame],
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        # fmt: off
+        key1 = Key.obsm.abs_probs(self.backward)
+        self._set("_absorption_probabilities", self.adata.obsm, key=key1, value=abs_probs)
+        key2 = Key.obsm.abs_times(self.backward)
+        self._set("_absorption_times", self.adata.obsm, key=key2, value=abs_times)
+        self._write_lineage_priming(None, log=False)
+        self.params[key1] = dict(params)
+        # fmt: on
+
+        if abs_times is None:
+            return (
+                f"Adding `adata.obsm[{key1!r}]`\n"
+                f"       `.absorption_probabilities`\n"
+                f"    Finish"
+            )
+        return (
+            f"Adding `adata.obsm[{key1!r}]`\n"
+            f"       `adata.obsm[{key2!r}]`\n"
+            f"       `.absorption_probabilities`\n"
+            f"       `.absorption_times`\n"
+            f"    Finish"
+        )
+
+    def _ensure_lineage_object(
+        self: AbsProbsProtocol, attr: str, **kwargs: Any
+    ) -> None:
+        if isinstance(getattr(self, attr), np.ndarray):
+            try:
+                setattr(
+                    self,
+                    attr,
+                    Lineage.from_adata(self.adata, backward=self.backward, **kwargs),
+                )
+            except Exception as e:  # noqa: B902
+                raise RuntimeError(
+                    f"Unable to reconstruct `.absorption_probabilities`. Reason: `{e}`."
+                ) from None
+
+    @logger
+    @shadow
+    def _write_lineage_priming(
+        self: AbsProbsProtocol, priming_degree: Optional[pd.Series]
+    ) -> str:
+        self._priming_degree = priming_degree
+        key = Key.obs.priming_degree(self.backward)
+        self._set("_priming_degree", self.adata.obs, key=key, value=priming_degree)
+
+        return f"Adding `adata.obs[{key!r}]`\n       `.priming_degree`\n    Finish"
+
+    def _read_absorption_probabilities(
+        self: AbsProbsProtocol, anndata: AnnData
+    ) -> bool:
+        # fmt: off
+        with SafeGetter(self, allowed=KeyError) as sg:
+            key1 = Key.obsm.abs_probs(self.backward)
+            self._get("_absorption_probabilities", self.adata.obsm, key=key1, where="obsm", dtype=(np.ndarray, Lineage))
+            self._ensure_lineage_object("_absorption_probabilities", kind="abs_probs")
+            key = Key.obsm.abs_times(self.backward)
+            self._get("_absorption_times", self.adata.obsm, key=key, where="obsm", dtype=pd.DataFrame,
+                      allow_missing=True)
+            key = Key.obs.priming_degree(self.backward)
+            self._get("_priming_degree", self.adata.obs, key=key, where="obs", dtype=pd.Series, allow_missing=True)
+            self.params[key1] = self._read_params(key1)
+        # fmt: on
+
+        return sg.ok
+
+    plot_absorption_probabilities = register_plotter(
+        continuous="absorption_probabilities"
+    )
diff --git a/cellrank/tl/estimators/mixins/_lineage_drivers.py b/cellrank/tl/estimators/mixins/_lineage_drivers.py
new file mode 100644
index 000000000..c2fb6853f
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/_lineage_drivers.py
@@ -0,0 +1,613 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
+from typing_extensions import Literal
+
+from types import MappingProxyType
+from pathlib import Path
+
+import scanpy as sc
+import scvelo as scv
+from anndata import Raw, AnnData
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.ul._docs import d, inject_docs
+from cellrank.tl._utils import RandomKeys, TestMethod, save_fig, _correlation_test
+from cellrank.tl._colors import _create_categorical_colors
+from cellrank.tl._lineage import Lineage
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.estimators.mixins._utils import BaseProtocol, logger, shadow
+from cellrank.tl.estimators.mixins._absorption_probabilities import AbsProbsMixin
+
+import numpy as np
+import pandas as pd
+
+import matplotlib.pyplot as plt
+from matplotlib import rc_context, patheffects
+from matplotlib.axes import Axes
+from matplotlib.patches import ArrowStyle
+
+
+class LinDriversProtocol(BaseProtocol):  # noqa: D101
+    def _write_lineage_drivers(
+        self,
+        drivers: Optional[pd.DataFrame],
+        use_raw: bool,
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> None:
+        ...
+
+    @property
+    def eigendecomposition(self) -> Dict[str, Any]:  # noqa: D102
+        ...
+
+    @property
+    def absorption_probabilities(self) -> Optional[Lineage]:  # noqa: D102
+        ...
+
+    @property
+    def lineage_drivers(self) -> Optional[pd.DataFrame]:  # noqa: D102
+        ...
+
+
+class LinDriversMixin(AbsProbsMixin):
+    """Mixin that computes potential driver genes for lineages."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+        self._lineage_drivers: Optional[pd.DataFrame] = None
+
+    @d.get_full_description(base="lineage_drivers")
+    @d.get_sections(base="lineage_drivers", sections=["Parameters", "Returns"])
+    @d.dedent
+    @inject_docs(tm=TestMethod)
+    def compute_lineage_drivers(
+        self: LinDriversProtocol,
+        lineages: Optional[Union[str, Sequence]] = None,
+        method: Literal["fischer", "perm_test"] = TestMethod.FISCHER,
+        cluster_key: Optional[str] = None,
+        clusters: Optional[Union[str, Sequence]] = None,
+        layer: Optional[str] = None,
+        use_raw: bool = False,
+        confidence_level: float = 0.95,
+        n_perms: int = 1000,
+        seed: Optional[int] = None,
+        **kwargs: Any,
+    ) -> pd.DataFrame:
+        """
+        Compute driver genes per lineage.
+
+        Correlates gene expression with lineage probabilities, for a given lineage and set of clusters.
+        Often, it makes sense to restrict this to a set of clusters which are relevant for the specified lineages.
+
+        Parameters
+        ----------
+        lineages
+            Lineage names from :attr:`absorption_probabilities`. If `None`, use all lineages.
+        method
+            Mode to use when calculating p-values and confidence intervals. Valid options are:
+
+                - `{tm.FISCHER!r}` - use Fischer transformation :cite:`fischer:21`.
+                - `{tm.PERM_TEST!r}` - use permutation test.
+        cluster_key
+            Key from :attr:`anndata.AnnData.obs` to obtain cluster annotations. These are considered for ``clusters``.
+        clusters
+            Restrict the correlations to these clusters.
+        layer
+            Key from :attr:`anndata.AnnData.layers` from which to get the expression.
+            If `None` or `'X'`, use :attr:`anndata.AnnData.X`.
+        use_raw
+            Whether or not to use :attr:`anndata.AnnData.raw` to correlate gene expression.
+        confidence_level
+            Confidence level for the confidence interval calculation. Must be in interval `[0, 1]`.
+        n_perms
+            Number of permutations to use when ``method = {tm.PERM_TEST!r}``.
+        seed
+            Random seed when ``method = {tm.PERM_TEST!r}``.
+        %(parallel)s
+
+        Returns
+        -------
+        %(correlation_test.returns)s
+
+        Also updates the following field:
+
+            - :attr:`lineage_drivers` - the same :class:`pandas.DataFrame` as described above.
+        """
+
+        # check that lineage probs have been computed
+        method = TestMethod(method)
+        abs_probs = self.absorption_probabilities
+        if abs_probs is None:
+            raise RuntimeError(
+                "Compute `.absorption_probabilities` first as `.compute_absorption_probabilities()`."
+            )
+
+        if abs_probs.shape[1] == 1:
+            logg.warning(
+                "There is only 1 lineage present. Using stationary distribution instead"
+            )
+            stat_dist = self.eigendecomposition.get("stationary_dist", None)
+            if stat_dist is None:
+                raise RuntimeError(
+                    "No stationary distribution found in `.eigendecomposition['stationary_dist']`."
+                )
+            abs_probs = Lineage(
+                stat_dist,
+                names=abs_probs.names,
+                colors=abs_probs.colors,
+            )
+
+        # check all lin_keys exist in self.lin_names
+        if isinstance(lineages, str):
+            lineages = [lineages]
+        if lineages is not None:
+            _ = abs_probs[lineages]
+        else:
+            lineages = abs_probs.names
+        lineages = list(lineages)
+
+        if not len(lineages):
+            raise ValueError("No lineages have been selected.")
+
+        # use `cluster_key` and clusters to subset the data
+        if clusters is not None:
+            if cluster_key not in self.adata.obs.keys():
+                raise KeyError(
+                    f"Unable to find clusters in `adata.obs[{cluster_key!r}]`."
+                )
+            if isinstance(clusters, str):
+                clusters = [clusters]
+
+            all_clusters = np.array(self.adata.obs[cluster_key].cat.categories)
+            cluster_mask = np.array([name not in all_clusters for name in clusters])
+
+            if any(cluster_mask):
+                raise KeyError(
+                    f"Clusters `{list(np.array(clusters)[cluster_mask])}` not found in "
+                    f"`adata.obs[{cluster_key!r}]`."
+                )
+            subset_mask = np.in1d(self.adata.obs[cluster_key], clusters)
+            adata_comp = self.adata[subset_mask]
+            lin_probs = abs_probs[subset_mask, :]
+        else:
+            adata_comp = self.adata
+            lin_probs = abs_probs
+
+        # check that the layer exists, and that use raw is only used with layer X
+        if layer in (None, "X"):
+            if use_raw and self.adata.raw is None:
+                logg.warning("No raw attribute set. Using `.X` instead")
+                use_raw = False
+            data = adata_comp.raw.X if use_raw else adata_comp.X
+            var_names = adata_comp.raw.var_names if use_raw else adata_comp.var_names
+        else:
+            if layer not in self.adata.layers:
+                raise KeyError(f"Layer `{layer!r}` not found in `adata.layers`.")
+            if use_raw:
+                logg.warning("If `use_raw=True`, layer must be `None`.")
+                use_raw = False
+            data = adata_comp.layers[layer]
+            var_names = adata_comp.var_names
+
+        start = logg.debug(
+            f"Computing correlations for lineages `{sorted(lineages)}` restricted to clusters `{clusters}` in "
+            f"layer `{'X' if layer is None else layer}` with `use_raw={use_raw}`"
+        )
+
+        lin_probs = lin_probs[lineages]
+        drivers = _correlation_test(
+            data,
+            lin_probs,
+            gene_names=var_names,
+            method=method,
+            n_perms=n_perms,
+            seed=seed,
+            confidence_level=confidence_level,
+            **kwargs,
+        )
+        params = self._create_params()
+        self._write_lineage_drivers(
+            drivers.loc[var_names], use_raw=use_raw, params=params, time=start
+        )
+
+        return drivers
+
+    @d.get_sections(base="plot_lineage_drivers", sections=["Parameters"])
+    @d.dedent
+    def plot_lineage_drivers(
+        self: LinDriversProtocol,
+        lineage: str,
+        n_genes: int = 8,
+        ncols: Optional[int] = None,
+        use_raw: bool = False,
+        title_fmt: str = "{gene} qval={qval:.4e}",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[int] = None,
+        save: Optional[Union[str, Path]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Plot lineage drivers discovered by :meth:`compute_lineage_drivers`.
+
+        Parameters
+        ----------
+        lineage
+            Lineage for which to plot the driver genes.
+        n_genes
+            Top most correlated genes to plot.
+        ncols
+            Number of columns.
+        use_raw
+            Whether to access in :attr:`anndata.AnnData.raw` or not.
+        title_fmt
+            Title format. Can include `{gene}`, `{pval}`, `{qval}` or `{corr}`, which will be substituted
+            with the actual values.
+        %(plotting)s
+        kwargs
+            Keyword arguments for :func:`scvelo.pl.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+
+        def prepare_format(
+            gene: str,
+            *,
+            pval: Optional[float],
+            qval: Optional[float],
+            corr: Optional[float],
+        ) -> Dict[str, Any]:
+            kwargs: Dict[str, Optional[Union[str, float]]] = {}
+            if "{gene" in title_fmt:
+                kwargs["gene"] = gene
+            if "{pval" in title_fmt:
+                kwargs["pval"] = float(pval) if pval is not None else np.nan
+            if "{qval" in title_fmt:
+                kwargs["qval"] = float(qval) if qval is not None else np.nan
+            if "{corr" in title_fmt:
+                kwargs["corr"] = float(corr) if corr is not None else np.nan
+
+            return kwargs
+
+        lin_drivers = self.lineage_drivers
+        if lin_drivers is None:
+            raise RuntimeError(
+                "Compute `.lineage_drivers` first as `.compute_lineage_drivers()`."
+            )
+
+        key = f"{lineage}_corr"
+        if key not in lin_drivers:
+            raise KeyError(
+                f"Lineage `{key!r}` not found in `{list(lin_drivers.columns)}`."
+            )
+
+        if n_genes <= 0:
+            raise ValueError(f"Expected `n_genes` to be positive, found `{n_genes}`.")
+
+        _ = kwargs.pop("save", None)
+        genes = lin_drivers.sort_values(by=key, ascending=False).head(n_genes)
+
+        ncols = 4 if ncols is None else ncols
+        nrows = int(np.ceil(len(genes) / ncols))
+
+        fig, axes = plt.subplots(
+            ncols=ncols,
+            nrows=nrows,
+            dpi=dpi,
+            figsize=(ncols * 6, nrows * 4) if figsize is None else figsize,
+        )
+        axes = np.ravel([axes])
+
+        _i = 0
+        for _i, (gene, ax) in enumerate(zip(genes.index, axes)):
+            data = genes.loc[gene]
+            scv.pl.scatter(
+                self.adata,
+                color=gene,
+                ncols=ncols,
+                use_raw=use_raw,
+                ax=ax,
+                show=False,
+                title=title_fmt.format(
+                    **prepare_format(
+                        gene,
+                        pval=data.get(f"{lineage}_pval", None),
+                        qval=data.get(f"{lineage}_qval", None),
+                        corr=data.get(f"{lineage}_corr", None),
+                    )
+                ),
+                **kwargs,
+            )
+
+        for j in range(_i + 1, len(axes)):
+            axes[j].remove()
+
+        if save is not None:
+            save_fig(fig, save)
+
+    @d.dedent
+    def plot_lineage_drivers_correlation(
+        self: LinDriversProtocol,
+        lineage_x: str,
+        lineage_y: str,
+        color: Optional[str] = None,
+        gene_sets: Optional[Dict[str, Sequence[str]]] = None,
+        gene_sets_colors: Optional[Sequence[str]] = 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:`anndata.AnnData.var` or :attr:`anndata.AnnData.varm`, preferring for the former.
+        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:`anndata.AnnData.raw` or not.
+        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
+        dkey = Key.varm.lineage_drivers(self.backward)
+        # this makes a useless copy, wish there was another way
+        if isinstance(adata, Raw):
+            adata = adata.to_adata()
+
+        if color is not None:
+            if not isinstance(color, str):
+                raise TypeError(
+                    "Expected `color` to be either of type `str` or `None`, "
+                    f"found `{type(color).__name__}`."
+                )
+            if color not in adata.var and color not in adata.varm[dkey]:
+                raise KeyError(
+                    f"Unable to find key `{color!r}` in `adata{'.raw' if use_raw else ''}.var` "
+                    f"or `adata.{'.raw' if use_raw else ''}varm[{dkey!r}]`."
+                )
+
+        # silent assumption: `.compute_lineage_drivers()` always writes to AnnData
+        key1, key2 = f"{lineage_x}_corr", f"{lineage_y}_corr"
+        if key1 not in adata.varm[dkey] or key2 not in adata.varm[dkey]:
+            haystack = "adata.raw.varm" if use_raw else "adata.varm"
+            raise RuntimeError(
+                f"Unable to find correlations in `{haystack}[{key1!r}]` or `{haystack}[{key2!r}]`."
+                f"Compute `.lineage_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), RandomKeys(adata, n=3, where="var") as keys:
+            adata.var[keys[0]] = adata.varm[dkey][key1]
+            adata.var[keys[1]] = adata.varm[dkey][key2]
+            if color is not None:
+                if color in adata.var:
+                    ckey = color
+                else:
+                    adata.var[keys[2]] = adata.varm[dkey][color]
+                    ckey = keys[2]
+
+            ax = sc.pl.scatter(
+                adata,
+                x=keys[0],
+                y=keys[1],
+                color_map=cmap,
+                use_raw=False,
+                show=False,
+                color=None if color is None else ckey,
+                **kwargs,
+            )
+
+            # add some lines to highlight the origin
+            xmin, xmax = np.nanmin(adata.var[keys[0]]), np.nanmax(adata.var[keys[0]])
+            ymin, ymax = np.nanmin(adata.var[keys[1]]), np.nanmax(adata.var[keys[1]])
+            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)
+
+        fig = ax.figure
+        ax.set_xlabel(key1)
+        ax.set_ylabel(key2)
+        ax.set_title(color)
+
+        # 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.absorption_probabilities[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.varm[dkey].loc[value, key1]
+                    y = adata.varm[dkey].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.varm[dkey][key1].values,
+                        y=adata.varm[dkey][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
+
+    @logger
+    @shadow
+    def _write_absorption_probabilities(
+        self: LinDriversProtocol,
+        abs_probs: Optional[Lineage],
+        abs_times: Optional[pd.DataFrame],
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        self._write_lineage_drivers(None, use_raw=False, log=False)
+        try:
+            self._write_lineage_drivers(None, use_raw=True, log=False)
+        except AttributeError:
+            pass
+        return super()._write_absorption_probabilities(
+            abs_probs, abs_times, params=params, log=False
+        )
+
+    @logger
+    @shadow
+    def _write_lineage_drivers(
+        self: LinDriversProtocol,
+        drivers: Optional[pd.DataFrame],
+        use_raw: bool,
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        self._lineage_drivers = drivers
+
+        # fmt: off
+        key = Key.varm.lineage_drivers(self.backward)
+        self._set("_lineage_drivers", self.adata.raw.varm if use_raw else self.adata.varm, key=key, value=drivers)
+        self.params[key] = dict(params)
+        # fmt: on
+
+        return (
+            f"Adding `adata.{'raw.' if use_raw else ''}varm[{key!r}]`\n"
+            f"       `.lineage_drivers`\n"
+            "    Finish"
+        )
+
+    def _read_absorption_probabilities(
+        self: LinDriversProtocol, adata: AnnData
+    ) -> bool:
+        ok = super()._read_absorption_probabilities(adata)
+        if not ok:
+            return False
+
+        # fmt: off
+        key = Key.varm.lineage_drivers(self.backward)
+        with SafeGetter(self, allowed=(KeyError, AttributeError)) as sg:
+            self.params[key] = self._read_params(key)
+            if self.params[key].get("use_raw", False):
+                self._get("_lineage_drivers", self.adata.raw.varm, key=key, where="varm", dtype=pd.DataFrame)
+            else:
+                self._get("_lineage_drivers", self.adata.varm, key=key, where="varm", dtype=pd.DataFrame)
+        # fmt: on
+
+        return sg.ok
+
+    @property
+    @d.dedent
+    def lineage_drivers(self) -> Optional[pd.DataFrame]:
+        """
+        Potential lineage drivers.
+
+        Computes Pearson correlation of each gene with fate probabilities for every terminal state. High Pearson
+        correlation indicates potential lineage drivers. Also computes p-values and confidence intervals.
+
+        Returns
+        -------
+        %(correlation_test.returns)s
+        """
+        return self._lineage_drivers
diff --git a/cellrank/tl/estimators/mixins/_utils.py b/cellrank/tl/estimators/mixins/_utils.py
new file mode 100644
index 000000000..3410c2df1
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/_utils.py
@@ -0,0 +1,427 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Callable, Optional, Sequence
+from typing_extensions import Literal, Protocol
+
+from enum import auto
+from wrapt import decorator
+
+import scvelo as scv
+from anndata import AnnData
+from cellrank import logging as logg
+from scanpy._utils import deprecated_arg_names
+from cellrank.tl._enum import ModeEnum
+from cellrank.ul._docs import d, inject_docs
+from cellrank.tl._utils import RandomKeys, _unique_order_preserving
+from cellrank.tl._colors import _create_categorical_colors
+from cellrank.tl._lineage import Lineage
+
+import numpy as np
+import pandas as pd
+from pandas.api.types import infer_dtype, is_categorical_dtype
+
+
+class PlotMode(ModeEnum):  # noqa: D101
+    EMBEDDING = auto()
+    TIME = auto()
+
+
+class BaseProtocol(Protocol):  # noqa: D101
+    @property
+    def adata(self) -> AnnData:  # noqa: D102
+        ...
+
+    @property
+    def backward(self) -> bool:  # noqa: D102
+        ...
+
+    @property
+    def params(self) -> Dict[str, Any]:  # noqa: D102
+        ...
+
+    def _set(
+        self,
+        attr: Optional[str] = None,
+        obj: Optional[Union[pd.DataFrame, Mapping[str, Any]]] = None,
+        key: Optional[str] = None,
+        value: Optional[
+            Union[np.ndarray, pd.Series, pd.DataFrame, Lineage, AnnData, Dict[str, Any]]
+        ] = None,
+        copy: bool = True,
+        shadow_only: bool = False,
+    ) -> None:
+        ...
+
+    def _get(
+        self,
+        attr: str,
+        obj: Union[pd.DataFrame, Mapping[str, Any]],
+        key: str,
+        where: Optional[Literal["obs", "obsm", "var", "varm", "uns"]] = None,
+        dtype: Optional[Union[type, Tuple[type, ...]]] = None,
+        copy: bool = True,
+        allow_missing: bool = False,
+    ) -> None:
+        ...
+
+    def _create_params(
+        self,
+        locs: Optional[Mapping[str, Any]] = None,
+        func: Optional[Callable] = None,
+        remove: Sequence[str] = (),
+    ) -> Dict[str, Any]:
+        ...
+
+    def _read_params(self, key: str) -> Dict[str, Any]:
+        ...
+
+
+class PlotterProtocol:  # noqa: D101
+    @property
+    def adata(self) -> AnnData:  # noqa: D102
+        ...
+
+
+@decorator()
+def logger(
+    wrapped: Callable[..., str], instance: Any, args: Any, kwargs: Dict[str, Any]
+) -> str:
+    """Handle logging for :class:`anndata.AnnData` writing functions of :class:`cellrank.estimators.BaseEstimator`."""
+    log, time = kwargs.pop("log", True), kwargs.pop("time", None)
+    msg = wrapped(*args, **kwargs)
+
+    if log:
+        logg.info(msg, time=time)
+
+    return msg
+
+
+@decorator()
+def shadow(
+    wrapped: Callable[..., str], instance: Any, args: Any, kwargs: Mapping[str, Any]
+) -> str:
+    """
+    Duplicate function call with shadowed :class:`anndata.AnnData.object`.
+
+    Used to easily create :class:`anndata.AnnData` serialization object in :class:`cellrank.estimators.BaseEstimator`.
+    """
+    res = wrapped(*args, **kwargs)
+    with instance._shadow:
+        try:
+            # don't care what the "shadowed" function returns
+            _ = wrapped(*args, **kwargs)
+        except Exception as e:  # noqa: B902
+            logg.error(
+                f"Unable to duplicate function call using shadow `anndata.AnnData` object. Reason: `{e}`"
+            )
+
+    return res
+
+
+def _plot_discrete(
+    self: PlotterProtocol,
+    _data: pd.Series,
+    _colors: Optional[np.ndarray] = None,
+    _title: Optional[str] = None,
+    states: Optional[Union[str, Sequence[str]]] = None,
+    color: Optional[str] = None,
+    title: Optional[Union[str, Sequence[str]]] = None,
+    same_plot: bool = True,
+    cmap: str = "viridis",
+    **kwargs: Any,
+) -> None:
+    if not isinstance(_data, pd.Series):
+        raise TypeError(
+            f"Expected `data` to be of type `pandas.Series`, found `{type(_data).__name__}`."
+        )
+    if not is_categorical_dtype(_data):
+        raise TypeError(
+            f"Expected `data` to be `categorical`, found `{infer_dtype(_data)}`."
+        )
+
+    names = list(_data.cat.categories)
+    if _colors is None:
+        _colors = _create_categorical_colors(len(names))
+    if len(_colors) != len(names):
+        raise ValueError(
+            f"Expected `colors` to be of length `{len(names)}`, found `{len(_colors)}`."
+        )
+    color_mapper = dict(zip(names, _colors))
+
+    states = _unique_order_preserving(states or names)
+    if not len(states):
+        raise ValueError("No states have been selected.")
+
+    for name in states:
+        if name not in names:
+            raise ValueError(
+                f"Invalid name `{name!r}`. Valid options are: `{sorted(names)}`."
+            )
+    _data = _data.cat.set_categories(states)
+
+    color = [] if color is None else (color,) if isinstance(color, str) else color
+    color = _unique_order_preserving(color)
+
+    same_plot = same_plot or len(names) == 1
+    kwargs.setdefault("legend_loc", "on data")
+    kwargs["color_map"] = cmap
+
+    # fmt: off
+    with RandomKeys(self.adata, n=1 if same_plot else len(states), where="obs") as keys:
+        if same_plot:
+            self.adata.obs[keys[0]] = _data
+            self.adata.uns[f"{keys[0]}_colors"] = [color_mapper[name] for name in states]
+            title = _title if title is None else title
+        else:
+            for key, cat in zip(keys, states):
+                self.adata.obs[key] = _data.cat.set_categories([cat])
+                self.adata.uns[f"{key}_colors"] = [color_mapper[cat]]
+            title = [f"{_title} {name}" for name in states] if title is None else title
+
+        if isinstance(title, str):
+            title = [title]
+
+        scv.pl.scatter(
+            self.adata,
+            color=color + keys,
+            title=color + title,
+            **kwargs,
+        )
+    # fmt: on
+
+
+def _plot_continuous(
+    self: PlotterProtocol,
+    _data: Lineage,
+    _colors: Optional[np.ndarray] = None,
+    _title: Optional[str] = None,
+    states: Optional[Union[str, Sequence[str]]] = None,
+    color: Optional[str] = None,
+    mode: Literal["embedding", "time"] = PlotMode.EMBEDDING,
+    time_key: str = "latent_time",
+    title: Optional[Union[str, Sequence[str]]] = None,
+    same_plot: bool = True,
+    cmap: str = "viridis",
+    **kwargs: Any,
+) -> None:
+    mode = PlotMode(mode)
+    if not isinstance(_data, Lineage):
+        raise TypeError(
+            f"Expected data to be of type `Lineage`, found `{type(_data).__name__}`."
+        )
+
+    if states is None:
+        states = _data.names
+    if not len(states):
+        raise ValueError("No lineages have been selected.")
+    is_singleton = _data.shape[1] == 1
+    _data = _data[states].copy()
+
+    if mode == "time" and same_plot:
+        logg.warning(
+            "Invalid combination `mode='time'` and `same_plot=True`. Using `same_plot=False`"
+        )
+        same_plot = False
+
+    _data_X = _data.X  # list(_data.T) behaves differently than a numpy.array
+    if _data_X.shape[1] == 1:
+        same_plot = False
+        if np.allclose(_data_X, 1.0):
+            # matplotlib shows even tiny perturbations in the colormap
+            _data_X = np.ones_like(_data_X)
+
+    for col in _data_X.T:
+        mask = ~np.isclose(col, 1.0)
+        # change the maximum value - the 1 is artificial and obscures the color scaling
+        if np.any(mask):
+            col[~mask] = np.nanmax(col[mask])
+
+    # fmt: off
+    color = [] if color is None else (color,) if isinstance(color, str) else color
+    color = _unique_order_preserving(color)
+
+    if mode == PlotMode.TIME:
+        kwargs.setdefault("legend_loc", "best")
+        if title is None:
+            title = [f"{_title} {state}" for state in states]
+        if time_key not in self.adata.obs:
+            raise KeyError(f"Unable to find pseudotime in `adata.obs[{time_key!r}]`.")
+        if len(color) and len(color) not in (1, _data_X.shape[1]):
+            raise ValueError(f"Expected `color` to be of length `1` or `{_data_X.shape[1]}`, "
+                             f"found `{len(color)}`.")
+        kwargs["x"] = self.adata.obs[time_key]
+        kwargs["y"] = list(_data_X.T)
+        kwargs["color"] = color if len(color) else None
+        kwargs["xlabel"] = [time_key] * len(states)
+        kwargs["ylabel"] = ["probability"] * len(states)
+    elif mode == PlotMode.EMBEDDING:
+        kwargs.setdefault("legend_loc", "on data")
+        if same_plot:
+            if color:
+                # https://github.com/theislab/scvelo/issues/673
+                logg.warning("Ignoring `color` when `mode='embedding'` and `same_plot=True`")
+            title = [_title] if title is None else title
+            kwargs["color_gradients"] = _data
+        else:
+            title = [f"{_title} {state}" for state in states] if title is None else title
+            if isinstance(title, str):
+                title = [title]
+            title = color + title
+            kwargs["color"] = color + list(_data_X.T)
+    else:
+        raise NotImplementedError(f"Mode `{mode}` is not yet implemented.")
+    # fmt: on
+
+    # e.g. a stationary distribution
+    if is_singleton and not np.allclose(_data_X, 1.0):
+        kwargs.setdefault("perc", [0, 95])
+        _ = kwargs.pop("color_gradients", None)
+
+    scv.pl.scatter(
+        self.adata,
+        title=title,
+        color_map=cmap,
+        **kwargs,
+    )
+
+
+# TODO(michalk8): in 2.0, remove %(time_mode)s and the deprecation
+@d.dedent
+@inject_docs(m=PlotMode)
+@deprecated_arg_names({"cluster_key": "color"})
+def _plot_dispatcher(
+    self: PlotterProtocol,
+    states: Optional[Union[str, Sequence[str]]] = None,
+    color: Optional[str] = None,
+    discrete: bool = False,
+    mode: Literal["embedding", "time"] = PlotMode.EMBEDDING,
+    time_key: str = "latent_time",
+    same_plot: bool = True,
+    title: Optional[Union[str, Sequence[str]]] = None,
+    cmap: str = "viridis",
+    **kwargs: Any,
+) -> None:
+    """
+    Plot continuous or categorical observations in an embedding or along pseudotime.
+
+    Parameters
+    ----------
+    states
+        States to plot.
+    color
+        Key in :attr:`anndata.AnnData.obs`.
+    discrete
+        Whether to plot the data as continuous or discrete observations.
+        If the data cannot be plotted as continuous observations, it will be plotted as discrete.
+    %(time_mode)s
+    time_key
+        Key in :attr:`anndata.AnnData.obs` where pseudotime is stored. Only used when ``mode = {m.TIME!r}``.
+    title
+        Title of the plot(s).
+    same_plot
+        Whether to plot the data on the same plot or not. Only use when ``mode = {m.EMBEDDING!r}``.
+        If `True` and ``discrete = False``, ``color`` is ignored.
+    cmap
+        Colormap for continuous data.
+    kwargs
+        Keyword arguments for :func:`scvelo.pl.scatter`.
+
+    Returns
+    -------
+    %(just_plots)s
+    """
+    if discrete:
+        return _plot_discrete(
+            self,
+            states=states,
+            color=color,
+            same_plot=same_plot,
+            title=title,
+            cmap=cmap,
+            **kwargs,
+        )
+
+    return _plot_continuous(
+        self,
+        states=states,
+        color=color,
+        mode=mode,
+        time_key=time_key,
+        same_plot=same_plot,
+        title=title,
+        cmap=cmap,
+        **kwargs,
+    )
+
+
+def register_plotter(
+    *,
+    discrete: Optional[str] = None,
+    continuous: Optional[str] = None,
+    colors: Optional[str] = None,
+):
+    """
+    Register a plotting function.
+
+    Used by the subclasses of :class`cellrank.estimators.BaseEstimator`.
+
+    Parameters
+    ----------
+    discrete
+        Attribute which stores categorical :class:`pandas.Series`.
+    continuous
+        Attribute which stores :class:`cellrank.tl.Lineage` containing continuous observations, such as
+        macrostates memberships or absorption probabilities.
+    colors
+        Attribute where color is stored. Useful only when specifying ``discrete``,
+        since :class:`cellrank.tl.Lineage` already contains color information.
+
+    Returns
+    -------
+    The plotting function.
+    """
+
+    @decorator()
+    def wrapper(
+        wrapped: Callable[..., None],
+        instance: Protocol,
+        args: Any,
+        kwargs: Dict[str, Any],
+    ) -> None:
+        # fmt: off
+        disc: Optional[bool] = kwargs.pop("discrete", None)
+        # prefer `continuous`
+        disc = continuous is None if disc is None else disc
+
+        if disc and discrete is None:
+            logg.warning(f"Unable to plot `.{continuous}` in `discrete` mode. Using `continuous` mode")
+            disc = False
+        if not disc and continuous is None:
+            logg.warning(f"Unable to plot `.{discrete}` in `continuous` mode. Using `discrete` mode")
+            disc = True
+        # fmt: on
+
+        attr = discrete if disc else continuous
+        data = getattr(instance, attr, None)
+        if data is None:
+            # prefer `continuous` for a prettier format
+            attr = continuous if continuous is not None else discrete
+            raise RuntimeError(f"Compute `.{attr}` first as `.compute_{attr}()`.")
+
+        if colors is None:
+            # extract colors from Lineage object, if present
+            _colors = getattr(getattr(instance, continuous, None), "colors", None)
+        else:
+            _colors = getattr(instance, colors, None)
+
+        return wrapped(
+            *args,
+            _data=data,
+            _colors=_colors,
+            _title=attr,
+            discrete=disc,
+            **kwargs,
+        )
+
+    if discrete is None and continuous is None:
+        raise ValueError("At least 1 of `discrete` or `continuous` must be set.")
+
+    return wrapper(_plot_dispatcher)
diff --git a/cellrank/tl/estimators/mixins/decomposition/__init__.py b/cellrank/tl/estimators/mixins/decomposition/__init__.py
new file mode 100644
index 000000000..4cd595047
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/decomposition/__init__.py
@@ -0,0 +1,2 @@
+from cellrank.tl.estimators.mixins.decomposition._eigen import EigenMixin
+from cellrank.tl.estimators.mixins.decomposition._schur import SchurMixin
diff --git a/cellrank/tl/estimators/mixins/decomposition/_eigen.py b/cellrank/tl/estimators/mixins/decomposition/_eigen.py
new file mode 100644
index 000000000..46f6c6b72
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/decomposition/_eigen.py
@@ -0,0 +1,380 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional
+from typing_extensions import Literal
+
+from types import MappingProxyType
+from pathlib import Path
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.ul._docs import d
+from cellrank.tl._utils import save_fig, _eigengap
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.estimators.mixins._utils import BaseProtocol, logger, shadow
+
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+from scipy.sparse.linalg import eigs
+
+import matplotlib.pyplot as plt
+from matplotlib.ticker import MultipleLocator, FormatStrFormatter
+
+EPS = np.finfo(np.float64).eps
+
+
+class EigenProtocol(BaseProtocol):  # noqa: D101
+    @property
+    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:  # noqa: D102
+        ...
+
+    def _write_eigendecomposition(
+        self, decomp: Dict[str, Any], params: Mapping[str, Any] = MappingProxyType({})
+    ) -> str:
+        ...
+
+
+class EigenMixin:
+    """Mixin that computes Schur decomposition."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+        self._eigendecomposition = None
+
+    @property
+    @d.get_summary(base="eigen")
+    def eigendecomposition(self) -> Optional[Dict[str, Any]]:
+        """Eigendecomposition of :attr:`transition_matrix`.
+
+        For non-symmetric real matrices, left and right eigenvectors will in general be different and complex. We
+        compute both left and right eigenvectors.
+
+        Returns
+        -------
+        A dictionary with the following keys:
+
+            - `'D'` - the eigenvalues.
+            - `'eigengap'` - the eigengap.
+            - `'params'` - parameters used for the computation.
+            - `'V_l'` - left eigenvectors (optional).
+            - `'V_r'` - right eigenvectors (optional).
+            - `'stationary_dist'` - stationary distribution of :attr:`transition_matrix`, if present.
+        """
+        return self._eigendecomposition
+
+    @d.dedent
+    def compute_eigendecomposition(
+        self: EigenProtocol,
+        k: int = 20,
+        which: Literal["LR", "LM"] = "LR",
+        alpha: float = 1.0,
+        only_evals: bool = False,
+        ncv: Optional[int] = None,
+    ) -> None:
+        """
+        Compute eigendecomposition of :attr:`transition_matrix`.
+
+        Uses a sparse implementation, if possible, and only computes the top :math:`k` eigenvectors
+        to speed up the computation. Computes both left and right eigenvectors.
+
+        Parameters
+        ----------
+        k
+            Number of eigenvectors or eigenvalues to compute.
+        %(eigen)s
+        only_evals
+            Whether to compute only eigenvalues.
+        ncv
+            Number of Lanczos vectors generated.
+
+        Returns
+        -------
+        Nothing, just updates the following field:
+
+            - :attr:`eigendecomposition` - %(eigen.summary)s
+        """
+
+        def get_top_k_evals() -> np.ndarray:
+            return D[np.flip(np.argsort(D.real))][:k]
+
+        start = logg.info("Computing eigendecomposition of the transition matrix")
+
+        if issparse(self.transition_matrix):
+            logg.debug(f"Computing top `{k}` eigenvalues of a sparse matrix")
+            D, V_l = eigs(self.transition_matrix.T, k=k, which=which, ncv=ncv)
+            if only_evals:
+                self._write_eigendecomposition(
+                    {
+                        "D": get_top_k_evals(),
+                        "eigengap": _eigengap(get_top_k_evals().real, alpha),
+                        "params": {"which": which, "k": k, "alpha": alpha},
+                    },
+                    time=start,
+                )
+                return
+            _, V_r = eigs(self.transition_matrix, k=k, which=which, ncv=ncv)
+        else:
+            logg.warning(
+                "This transition matrix is not sparse, computing full eigendecomposition"
+            )
+            D, V_l = np.linalg.eig(self.transition_matrix.T)
+            if only_evals:
+                self._write_eigendecomposition(
+                    {
+                        "D": get_top_k_evals(),
+                        "eigengap": _eigengap(D.real, alpha),
+                        "params": {"which": which, "k": k, "alpha": alpha},
+                    },
+                    time=start,
+                )
+                return
+            _, V_r = np.linalg.eig(self.transition_matrix)
+
+        # Sort the eigenvalues and eigenvectors and take the real part
+        logg.debug("Sorting eigenvalues by their real part")
+        p = np.flip(np.argsort(D.real))
+        D, V_l, V_r = D[p], V_l[:, p], V_r[:, p]
+        e_gap = _eigengap(D.real, alpha)
+
+        pi = np.abs(V_l[:, 0].real)
+        pi /= np.sum(pi)
+
+        self._write_eigendecomposition(
+            {
+                "D": D,
+                "stationary_dist": pi,
+                "V_l": V_l,
+                "V_r": V_r,
+                "eigengap": e_gap,
+                "params": {"which": which, "k": k, "alpha": alpha},
+            },
+            time=start,
+        )
+
+    @d.dedent
+    def plot_spectrum(
+        self,
+        n: Optional[int] = None,
+        real_only: bool = False,
+        show_eigengap: bool = True,
+        show_all_xticks: bool = True,
+        legend_loc: Optional[str] = None,
+        title: Optional[str] = None,
+        marker: str = ".",
+        figsize: Optional[Tuple[float, float]] = (5, 5),
+        dpi: int = 100,
+        save: Optional[Union[str, Path]] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Plot the top eigenvalues in real or complex plane.
+
+        Parameters
+        ----------
+        n
+            Number of eigenvalues to show. If `None`, show all that have been computed.
+        real_only
+            Whether to plot only the real part of the spectrum.
+        show_eigengap
+            When ``real_only = True``, this determines whether to show the inferred eigengap as
+            a dotted line.
+        show_all_xticks
+            When ``real_only = True``, this determines whether to show the indices of all eigenvalues on the x-axis.
+        legend_loc
+            Location parameter for the legend.
+        title
+            Title of the figure.
+        marker
+            Marker symbol used, valid options can be found in :mod:`matplotlib.markers`.
+        %(plotting)s
+        kwargs
+            Keyword arguments for :func:`matplotlib.pyplot.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+
+        eig = self.eigendecomposition
+        if eig is None:
+            raise RuntimeError(
+                "Compute eigendecomposition first as `.compute_eigendecomposition()`."
+            )
+
+        if n is None:
+            n = len(eig["D"])
+        if n <= 0:
+            raise ValueError(f"Expected `n` to be > 0, found `{n}`.")
+
+        if real_only:
+            fig = self._plot_real_spectrum(
+                n,
+                show_eigengap=show_eigengap,
+                show_all_xticks=show_all_xticks,
+                dpi=dpi,
+                figsize=figsize,
+                legend_loc=legend_loc,
+                title=title,
+                marker=marker,
+                **kwargs,
+            )
+        else:
+            fig = self._plot_complex_spectrum(
+                n,
+                dpi=dpi,
+                figsize=figsize,
+                legend_loc=legend_loc,
+                title=title,
+                marker=marker,
+                **kwargs,
+            )
+
+        if save:
+            save_fig(fig, save)
+
+    def _plot_complex_spectrum(
+        self,
+        n: int,
+        dpi: int = 100,
+        figsize: Optional[Tuple[float, float]] = None,
+        legend_loc: Optional[str] = None,
+        title: Optional[str] = None,
+        marker: str = ".",
+        **kwargs: Any,
+    ):
+        # define a function to make the data limits rectangular
+        def adapt_range(min_, max_, range_):
+            return (
+                min_ + (max_ - min_) / 2 - range_ / 2,
+                min_ + (max_ - min_) / 2 + range_ / 2,
+            )
+
+        eig = self.eigendecomposition
+        D, params = eig["D"][:n], eig["params"]
+
+        # create figure and axes
+        fig, ax = plt.subplots(nrows=1, ncols=1, dpi=dpi, figsize=figsize)
+
+        # get the original data ranges
+        lam_x, lam_y = D.real, D.imag
+        x_min, x_max = np.min(lam_x), np.max(lam_x)
+        y_min, y_max = np.min(lam_y), np.max(lam_y)
+        x_range, y_range = x_max - x_min, y_max - y_min
+        final_range = np.max([x_range, y_range]) + 0.05
+
+        x_min_, x_max_ = adapt_range(x_min, x_max, final_range)
+        y_min_, y_max_ = adapt_range(y_min, y_max, final_range)
+
+        # plot the data and the unit circle
+        ax.scatter(D.real, D.imag, marker=marker, label="eigenvalue", **kwargs)
+        t = np.linspace(0, 2 * np.pi, 500)
+        x_circle, y_circle = np.sin(t), np.cos(t)
+        ax.plot(x_circle, y_circle, "k-", label="unit circle")
+
+        # set labels, ranges and legend
+        ax.set_xlabel(r"Re($\lambda$)")
+        ax.set_xlim(x_min_, x_max_)
+
+        ax.set_ylabel(r"Im($\lambda$)")
+        ax.set_ylim(y_min_, y_max_)
+
+        key = "real part" if params["which"] == "LR" else "magnitude"
+        if title is None:
+            title = f"top {n} eigenvalues according to their {key}"
+
+        ax.set_title(title)
+        if legend_loc not in (None, "none"):
+            ax.legend(loc=legend_loc)
+
+        return fig
+
+    def _plot_real_spectrum(
+        self,
+        n: int,
+        show_eigengap: bool = True,
+        show_all_xticks: bool = True,
+        dpi: int = 100,
+        figsize: Optional[Tuple[float, float]] = None,
+        legend_loc: Optional[str] = None,
+        title: Optional[str] = None,
+        marker: str = ".",
+        **kwargs,
+    ):
+        eig = self.eigendecomposition
+        D, params = eig["D"][:n], eig["params"]
+
+        D_real, D_imag = D.real, D.imag
+        ixs = np.arange(len(D))
+        mask = D_imag == 0
+
+        # plot the top eigenvalues
+        fig, ax = plt.subplots(nrows=1, ncols=1, dpi=dpi, figsize=figsize)
+        if np.any(mask):
+            ax.scatter(
+                ixs[mask],
+                D_real[mask],
+                marker=marker,
+                label="real eigenvalue",
+                **kwargs,
+            )
+        if np.any(~mask):
+            ax.scatter(
+                ixs[~mask],
+                D_real[~mask],
+                marker=marker,
+                label="complex eigenvalue",
+                **kwargs,
+            )
+
+        # add dashed line for the eigengap, ticks, labels, title and legend
+        if show_eigengap and eig["eigengap"] < n:
+            ax.axvline(eig["eigengap"], label="eigengap", ls="--", lw=2, c="k")
+
+        ax.set_xlabel("index")
+        if show_all_xticks:
+            ax.set_xticks(np.arange(len(D)))
+        else:
+            ax.xaxis.set_major_locator(MultipleLocator(2.0))
+            ax.xaxis.set_major_formatter(FormatStrFormatter("%d"))
+
+        ax.set_ylabel(r"Re($\lambda_i$)")
+
+        key = "real part" if params["which"] == "LR" else "magnitude"
+        if title is None:
+            title = f"real part of top {n} eigenvalues according to their {key}"
+
+        ax.set_title(title)
+        if legend_loc not in (None, "none"):
+            ax.legend(loc=legend_loc)
+
+        return fig
+
+    @logger
+    @shadow
+    def _write_eigendecomposition(
+        self: EigenProtocol,
+        decomp: Dict[str, Any],
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        key = Key.uns.eigen(self.backward)
+        self._set("_eigendecomposition", self.adata.uns, key=key, value=decomp)
+        params = params or decomp.get("params", {})
+        self.params[key] = dict(params)
+
+        return (
+            f"Adding `adata.uns[{key!r}]`\n"
+            f"       `.eigendecomposition`\n"
+            "    Finish"
+        )
+
+    def _read_eigendecomposition(
+        self: EigenProtocol, adata: AnnData, allow_missing: bool = True
+    ) -> bool:
+        # fmt: off
+        key = Key.uns.eigen(self.backward)
+        with SafeGetter(self, allowed=KeyError) as sg:
+            self._get("_eigendecomposition", adata.uns, key=key, where="uns", dtype=dict, allow_missing=allow_missing)
+            self.params[key] = self._read_params(key)
+        # fmt: on
+
+        return sg.ok
diff --git a/cellrank/tl/estimators/mixins/decomposition/_schur.py b/cellrank/tl/estimators/mixins/decomposition/_schur.py
new file mode 100644
index 000000000..fed9d02b6
--- /dev/null
+++ b/cellrank/tl/estimators/mixins/decomposition/_schur.py
@@ -0,0 +1,313 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional
+from typing_extensions import Literal
+
+from types import MappingProxyType
+from pathlib import Path
+from pygpcca import GPCCA
+from pygpcca._sorted_schur import _check_conj_split
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.ul._docs import d
+from cellrank.tl._utils import save_fig, _eigengap
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.estimators.mixins._utils import BaseProtocol, logger, shadow
+
+import numpy as np
+from scipy.sparse import issparse, spmatrix
+
+import matplotlib.pyplot as plt
+from seaborn import heatmap
+from mpl_toolkits.axes_grid1 import make_axes_locatable
+
+EPS = np.finfo(np.float64).eps
+
+
+class SchurProtocol(BaseProtocol):  # noqa: D101
+    def transition_matrix(self) -> Union[np.ndarray, spmatrix]:  # noqa: D102
+        ...
+
+    def _write_schur_decomposition(
+        self,
+        decomp: Dict[str, Any],
+        vectors: np.ndarray,
+        matrix: np.ndarray,
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        ...
+
+
+class SchurMixin:
+    """Mixin that computes Schur decomposition."""
+
+    def __init__(self, **kwargs: Any):
+        super().__init__(**kwargs)
+
+        self._gpcca: Optional[GPCCA] = None
+        self._invalid_n_states: Optional[np.ndarray] = None
+
+        self._schur_vectors: Optional[np.ndarray] = None
+        self._schur_matrix: Optional[np.ndarray] = None
+        self._eigendecomposition: Optional[Dict[str, Any]] = None
+
+    @property
+    @d.get_summary(base="schur_vectors")
+    @d.get_extended_summary(base="schur_vectors")
+    def schur_vectors(self) -> Optional[np.ndarray]:
+        """Real Schur vectors of the transition matrix.
+
+        The real Schur decomposition is a generalization of the Eigendecomposition and can be computed for any
+        real-valued, square matrix :math:`A`. It is given by :math:`A = Q R Q^T`, where :math:`Q` contains the
+        real Schur vectors and :math:`R` is the Schur matrix. :math:`Q` is orthogonal and :math:`R` is quasi-upper
+        triangular with 1x1 and 2x2 blocks on the diagonal.
+        If PETSc and SLEPc are installed, only the leading Schur vectors are computed.
+        """
+        return self._schur_vectors
+
+    @property
+    @d.get_summary(base="schur_matrix")
+    @d.dedent
+    def schur_matrix(self) -> Optional[np.ndarray]:
+        """Schur matrix.
+
+        %(schur_vectors.summary_ext)s
+        """
+        return self._schur_matrix
+
+    @property
+    @d.get_summary(base="eigen")
+    def eigendecomposition(self) -> Optional[Dict[str, Any]]:
+        """Eigendecomposition of :attr:`transition_matrix`.
+
+        For non-symmetric real matrices, left and right eigenvectors will in general be different and complex. We
+        compute both left and right eigenvectors.
+
+        Returns
+        -------
+        A dictionary with the following keys:
+
+            - `'D'` - the eigenvalues.
+            - `'eigengap'` - the eigengap.
+            - `'params'` - parameters used for the computation.
+            - `'V_l'` - left eigenvectors (optional).
+            - `'V_r'` - right eigenvectors (optional).
+            - `'stationary_dist'` - stationary distribution of :attr:`transition_matrix`, if present.
+        """
+        return self._eigendecomposition
+
+    @d.dedent
+    def compute_schur(
+        self: SchurProtocol,
+        n_components: int = 10,
+        initial_distribution: Optional[np.ndarray] = None,
+        method: Literal["krylov", "brandts"] = "krylov",
+        which: Literal["LR", "LM"] = "LR",
+        alpha: float = 1.0,
+    ):
+        """
+        Compute Schur decomposition.
+
+        Parameters
+        ----------
+        n_components
+            Number of Schur vectors to compute.
+        initial_distribution
+            Input distribution over all cells. If `None`, uniform distribution is used.
+        method
+            Method for calculating the Schur vectors. Valid options are:
+
+                - `'krylov'` - an iterative procedure that computes a partial, sorted Schur decomposition for
+                  large, sparse matrices.
+                - `'brandts'` - full sorted Schur decomposition of a dense matrix.
+
+            For benefits of each method, see :class:`pygpcca.GPCCA`.
+        %(eigen)s
+
+        Returns
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`schur_vectors` - %(schur_vectors.summary)s
+            - :attr:`schur_matrix` -  %(schur_matrix.summary)s
+            - :attr:`eigendecomposition` - %(eigen.summary)s
+        """
+        if n_components < 2:
+            logg.warning(
+                f"Number of Schur vectors `>=2`, but only `{n_components}` "
+                f"were requested. Using `n_components=2`"
+            )
+            n_components = 2
+
+        if method not in ("brandts", "krylov"):
+            raise ValueError(
+                f"Invalid method `{method!r}`. Valid options are:`'brandts'` or `'krylov'`."
+            )
+
+        try:
+            import petsc4py
+            import slepc4py
+        except ImportError:
+            method = "brandts"
+            logg.warning(
+                f"Unable to import `petsc4py` or `slepc4py`. Using `method={method!r}`"
+            )
+
+        tmat = self.transition_matrix
+        if method == "brandts" and issparse(self.transition_matrix):
+            logg.warning("For `method='brandts'`, dense matrix is required. Densifying")
+            tmat = tmat.A
+
+        self._gpcca = GPCCA(tmat, eta=initial_distribution, z=which, method=method)
+        start = logg.info("Computing Schur decomposition")
+
+        try:
+            self._gpcca._do_schur_helper(n_components)
+        except ValueError as e:
+            if "will split complex conjugate eigenvalues" not in str(e):
+                raise
+            logg.warning(
+                f"Using `{n_components}` components would split a block of complex conjugate eigenvalues. "
+                f"Using `n_components={n_components + 1}`"
+            )
+            self._gpcca._do_schur_helper(n_components + 1)
+
+        self._invalid_n_states = np.array(
+            [
+                i
+                for i in range(2, len(self._gpcca._p_eigenvalues))
+                if _check_conj_split(self._gpcca._p_eigenvalues[:i])
+            ]
+        )
+        if len(self._invalid_n_states):
+            logg.info(
+                f"When computing macrostates, choose a number of states NOT in `{list(self._invalid_n_states)}`"
+            )
+
+        self._write_schur_decomposition(
+            {
+                "D": self._gpcca._p_eigenvalues,
+                "eigengap": _eigengap(self._gpcca._p_eigenvalues, alpha),
+                "params": {
+                    "which": which,
+                    "k": len(self._gpcca._p_eigenvalues),
+                    "alpha": alpha,
+                },
+            },
+            vectors=self._gpcca._p_X,
+            matrix=self._gpcca._p_R,
+            params=self._create_params(),
+            time=start,
+        )
+
+    @d.dedent
+    def plot_schur_matrix(
+        self,
+        title: Optional[str] = "schur matrix",
+        cmap: str = "viridis",
+        figsize: Optional[Tuple[float, float]] = None,
+        dpi: Optional[float] = 80,
+        save: Optional[Union[str, Path]] = None,
+        **kwargs: Any,
+    ):
+        """
+        Plot the Schur matrix.
+
+        Parameters
+        ----------
+        title
+            Title of the figure.
+        cmap
+            Colormap to use.
+        %(plotting)s
+        kwargs
+            Keyword arguments for :func:`seaborn.heatmap`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+        schur_matrix = self.schur_matrix
+        if schur_matrix is None:
+            raise RuntimeError("Compute Schur matrix first as `.compute_schur()`.")
+
+        fig, ax = plt.subplots(
+            figsize=schur_matrix.shape if figsize is None else figsize, dpi=dpi
+        )
+
+        divider = make_axes_locatable(ax)  # square=True make the colorbar a bit bigger
+        cbar_ax = divider.append_axes("right", size="2%", pad=0.1)
+
+        mask = np.zeros_like(schur_matrix, dtype=np.bool_)
+        mask[np.tril_indices_from(mask, k=-1)] = True
+        mask[~np.isclose(schur_matrix, 0.0)] = False
+
+        vmin, vmax = (
+            np.min(schur_matrix[~mask]),
+            np.max(schur_matrix[~mask]),
+        )
+
+        kwargs["fmt"] = kwargs.get("fmt", "0.2f")
+        heatmap(
+            data=schur_matrix,
+            cmap=cmap,
+            square=True,
+            annot=True,
+            vmin=vmin,
+            vmax=vmax,
+            cbar_ax=cbar_ax,
+            cbar_kws={"ticks": np.linspace(vmin, vmax, 10)},
+            mask=mask,
+            xticklabels=[],
+            yticklabels=[],
+            ax=ax,
+            **kwargs,
+        )
+        ax.set_title(title)
+
+        if save is not None:
+            save_fig(fig, path=save)
+
+    @logger
+    @shadow
+    def _write_schur_decomposition(
+        self: SchurProtocol,
+        decomp: Dict[str, Any],
+        vectors: np.ndarray,
+        matrix: np.ndarray,
+        params: Mapping[str, Any] = MappingProxyType({}),
+    ) -> str:
+        key = Key.uns.schur_matrix(self.backward)
+        self._set("_schur_matrix", self.adata.uns, key=key, value=matrix)
+        key = Key.obsm.schur_vectors(self.backward)
+        self._set("_schur_vectors", self.adata.obsm, key=key, value=vectors)
+        key = Key.uns.eigen(self.backward)
+        self._set("_eigendecomposition", self.adata.uns, key=key, value=decomp)
+        self.params[f"schur_decomposition_{Key.backward(self.backward)}"] = dict(params)
+
+        return (
+            f"Adding `adata.uns[{key!r}]`\n"
+            f"       `.schur_vectors`\n"
+            f"       `.schur_matrix`\n"
+            f"       `.eigendecomposition`\n"
+            "    Finish"
+        )
+
+    def _read_schur_decomposition(
+        self: SchurProtocol, adata: AnnData, allow_missing: bool = True
+    ) -> bool:
+        # fmt: off
+        key = Key.uns.eigen(self.backward)
+        with SafeGetter(self, allowed=KeyError) as sg:
+            self._get("_eigendecomposition", adata.uns, key=key, where="uns", dtype=dict, allow_missing=allow_missing)
+            key = Key.obsm.schur_vectors(self.backward)
+            self._get("_schur_vectors", self.adata.obsm, key=key, where="obsm", dtype=np.ndarray,
+                      allow_missing=allow_missing)
+            key = Key.uns.schur_matrix(self.backward)
+            self._get("_schur_matrix", self.adata.uns, key=key, where="uns", dtype=np.ndarray,
+                      allow_missing=allow_missing)
+            key = f"schur_decomposition_{Key.backward(self.backward)}"
+            self.params[key] = self._read_params(key)
+        # fmt: on
+
+        return sg.ok
diff --git a/cellrank/tl/estimators/terminal_states/__init__.py b/cellrank/tl/estimators/terminal_states/__init__.py
new file mode 100644
index 000000000..4f413d637
--- /dev/null
+++ b/cellrank/tl/estimators/terminal_states/__init__.py
@@ -0,0 +1,5 @@
+from cellrank.tl.estimators.terminal_states._gpcca import GPCCA
+from cellrank.tl.estimators.terminal_states._cflare import CFLARE
+from cellrank.tl.estimators.terminal_states._term_states_estimator import (
+    TermStatesEstimator,
+)
diff --git a/cellrank/tl/estimators/terminal_states/_cflare.py b/cellrank/tl/estimators/terminal_states/_cflare.py
new file mode 100644
index 000000000..c0dc01cd4
--- /dev/null
+++ b/cellrank/tl/estimators/terminal_states/_cflare.py
@@ -0,0 +1,264 @@
+from typing import Any, Dict, List, Union, Optional, Sequence
+from typing_extensions import Literal
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank.ul._docs import d
+from cellrank.tl._utils import (
+    _cluster_X,
+    _filter_cells,
+    _complex_warning,
+    _get_connectivities,
+)
+from cellrank.tl.kernels._utils import _get_basis
+from cellrank.tl.estimators.mixins import EigenMixin, LinDriversMixin
+from cellrank.tl.estimators.terminal_states._term_states_estimator import (
+    TermStatesEstimator,
+)
+
+import numpy as np
+import pandas as pd
+from scipy.stats import zscore
+
+
+@d.dedent
+class CFLARE(TermStatesEstimator, LinDriversMixin, EigenMixin):
+    """
+    Compute the initial/terminal states of a Markov chain via spectral heuristics.
+
+    This estimator uses the left eigenvectors of the transition matrix to filter to a set of recurrent cells
+    and the right eigenvectors to cluster this set of cells into discrete groups.
+
+    Parameters
+    ----------
+    %(base_estimator.parameters)s
+    """
+
+    @d.dedent
+    def fit(self, k: int = 20, **kwargs: Any) -> "TermStatesEstimator":
+        """
+        %(tse_fit.full_desc)s
+
+        Parameters
+        ----------
+        k
+            Number of eigenvectors to compute.
+        kwargs
+            Keyword arguments for :meth:`compute_eigendecomposition`.
+
+        Returns
+        -------
+        Self.
+        Also modifies the following field:
+
+            - :attr:`eigendecomposition` - %(eigen.summary)s
+        """  # noqa: D400
+        self.compute_eigendecomposition(k=k, only_evals=False, **kwargs)
+        return self
+
+    @d.dedent
+    def predict(
+        self,
+        use: Optional[Union[int, Sequence[int]]] = None,
+        percentile: Optional[int] = 98,
+        method: Literal["leiden", "means"] = "leiden",
+        cluster_key: Optional[str] = None,
+        n_clusters_kmeans: Optional[int] = None,
+        n_neighbors: int = 20,
+        resolution: float = 0.1,
+        n_matches_min: int = 0,
+        n_neighbors_filtering: int = 15,
+        basis: Optional[str] = None,
+        n_comps: int = 5,
+        scale: Optional[bool] = None,
+    ) -> None:
+        """
+        Find approximate recurrent classes of the Markov chain.
+
+        Filter to obtain recurrent states in left eigenvectors.
+        Cluster to obtain approximate recurrent classes in right eigenvectors.
+
+        Parameters
+        ----------
+        use
+            Which or how many first eigenvectors to use as features for filtering and clustering.
+            If `None`, use the *eigengap* statistic.
+        percentile
+            Threshold used for filtering out cells which are most likely transient states. Cells which are in the
+            lower ``percentile`` percent of each eigenvector will be removed from the data matrix.
+        method
+            Method to be used for clustering. Valid option are:
+
+                - `'kmeans'` - :class:`sklearn.cluster.KMeans`.
+                - `'leiden'` - :func:`scanpy.tl.leiden`.
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` in order to associate names and colors with :attr:`terminal_states`.
+        n_clusters_kmeans
+            If `None`, this is set to ``use + 1``.
+        n_neighbors
+            Number of neighbors in a KNN graph. This is the :math:`K` parameter for that,
+            the number of neighbors for each cell. Only used when ``method = 'leiden'``.
+        resolution
+            Resolution parameter for :func:`scanpy.tl.leiden`. Should be chosen relatively small.
+        n_matches_min
+            Filters out cells which don't have at least ``n_matches_min`` neighbors from the same category.
+            This filters out some cells which are transient but have been misassigned.
+        n_neighbors_filtering
+            Parameter for filtering cells. Cells are filtered out if they don't have at least ``n_matches_min``
+            neighbors among their ``n_neighbors_filtering`` nearest cells.
+        basis
+            Key from :attr:`anndata.AnnData.obsm` as additional features for clustering.
+            If `None`, use only the right eigenvectors.
+        n_comps
+            Number of embedding components to be use when ``basis != None``.
+        scale
+            Scale the values to z-scores. If `None`, scale the values if ``basis != None``.
+
+        Returns
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+            - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
+        """
+
+        def convert_use(
+            use: Optional[Union[int, Sequence[int], np.ndarray]]
+        ) -> List[int]:
+            if method not in ["kmeans", "leiden"]:
+                raise ValueError(
+                    f"Invalid method `{method!r}`. Valid options are `leiden` or `kmeans`."
+                )
+
+            if use is None:
+                use = eig["eigengap"] + 1  # add one b/c indexing starts at 0
+            if isinstance(use, int):
+                use = list(range(use))
+            elif not isinstance(use, (np.ndarray, Sequence)):
+                raise TypeError(
+                    f"Expected `use` to be `int` or a `Sequence`, found `{type(use).__name__}`."
+                )
+            use = list(use)
+
+            if not use:
+                raise ValueError("No eigenvectors have been selected.")
+
+            muse, mevecs = max(use), max(eig["V_l"].shape[1], eig["V_r"].shape[1])
+            if muse >= mevecs:
+                raise ValueError(
+                    f"Maximum specified eigenvector `{muse}` is larger than the number of eigenvectors `{mevecs}`. "
+                    f"Use `.compute_eigendecomposition(k={muse})` to recompute the eigendecomposition."
+                )
+
+            return use
+
+        eig = self.eigendecomposition
+        if eig is None:
+            raise RuntimeError(
+                "Compute eigendecomposition first as `.compute_eigendecomposition()`."
+            )
+        use = convert_use(use)
+
+        start = logg.info("Computing terminal states")
+        # we check for complex values only in the left, that's okay because the complex pattern
+        # will be identical for left and right
+        V_l, V_r = eig["V_l"][:, use], eig["V_r"].real[:, use]
+        V_l = _complex_warning(V_l, use, use_imag=False)
+
+        # retrieve embedding and concatenate
+        if basis is not None:
+            X_em = _get_basis(self.adata, basis=basis)[:, :n_comps]
+            X = np.concatenate([V_r, X_em], axis=1)
+        else:
+            X = V_r
+
+        # filter out cells which are in the lowest q percentile in abs value in each eigenvector
+        if percentile is not None:
+            logg.debug("Filtering out cells according to percentile")
+            if not (0 <= percentile <= 100):
+                raise ValueError(
+                    f"Expected `percentile` to be in interval `[0, 100]`, found `{percentile}`."
+                )
+            cutoffs = np.percentile(np.abs(V_l), percentile, axis=0)
+            ixs = np.any(cutoffs <= np.abs(V_l), axis=1)
+            X = X[ixs, :]
+
+        if scale is None:
+            scale = basis is not None
+        if scale:
+            X = zscore(X, axis=0)
+
+        # cluster X
+        if method == "kmeans" and n_clusters_kmeans is None:
+            n_clusters_kmeans = len(use) + (percentile is None)
+            if X.shape[0] < n_clusters_kmeans:
+                raise ValueError(
+                    f"Filtering resulted in only {X.shape[0]} cell(s), insufficient to cluster into "
+                    f"`{n_clusters_kmeans}` clusters. Consider decreasing the value of `percentile`."
+                )
+
+        # fmt: off
+        logg.debug(f"Using `{use}` eigenvectors, basis `{basis!r}` and method `{method!r}` for clustering")
+        clusters = _cluster_X(
+            X,
+            method=method,
+            n_clusters=n_clusters_kmeans,
+            n_neighbors=n_neighbors,
+            resolution=resolution,
+        )
+
+        # fill in the labels in case we filtered out cells before
+        if percentile is not None:
+            labels = np.repeat(None, len(self))
+            labels[ixs] = clusters
+        else:
+            labels = clusters
+        labels = pd.Series(labels, index=self.adata.obs_names, dtype="category")
+        labels = labels.cat.rename_categories({c: str(c) for c in labels.cat.categories})
+
+        # filtering to get rid of some of the left over transient states
+        if n_matches_min > 0:
+            logg.debug(f"Filtering according to `n_matches_min={n_matches_min}`")
+            distances = _get_connectivities(self.adata, mode="distances", n_neighbors=n_neighbors_filtering)
+            labels = _filter_cells(distances, rc_labels=labels, n_matches_min=n_matches_min)
+        # fmt: on
+
+        self.set_terminal_states(
+            labels=labels,
+            cluster_key=cluster_key,
+            probs=self._compute_term_states_probs(eig, use),
+            params=self._create_params(),
+            time=start,
+        )
+
+    def _compute_term_states_probs(
+        self, eig: Dict[str, Any], use: List[int]
+    ) -> pd.Series:
+        # get the truncated eigendecomposition
+        V, evals = eig["V_l"].real[:, use], eig["D"].real[use]
+
+        # shift and scale
+        V_pos = np.abs(V)
+        V_shifted = V_pos - np.min(V_pos, axis=0)
+        V_scaled = V_shifted / np.max(V_shifted, axis=0)
+
+        # check the ranges are correct
+        np.testing.assert_allclose(np.min(V_scaled, axis=0), 0)
+        np.testing.assert_allclose(np.max(V_scaled, axis=0), 1)
+
+        # further scale by the eigenvalues
+        V_eigs = V_scaled / evals
+
+        # sum over cols and scale
+        c = np.sum(V_eigs, axis=1)
+        c /= np.max(c)
+
+        return pd.Series(c, index=self.adata.obs_names)
+
+    def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
+        ok = super()._read_from_adata(adata, **kwargs)
+        return (
+            ok
+            and self._read_eigendecomposition(adata, allow_missing=False)
+            and self._read_absorption_probabilities(adata)
+        )
diff --git a/cellrank/tl/estimators/terminal_states/_gpcca.py b/cellrank/tl/estimators/terminal_states/_gpcca.py
new file mode 100644
index 000000000..727521f88
--- /dev/null
+++ b/cellrank/tl/estimators/terminal_states/_gpcca.py
@@ -0,0 +1,1192 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
+from typing_extensions import Literal
+
+from enum import auto
+from types import MappingProxyType
+from pathlib import Path
+from datetime import datetime
+
+from anndata import AnnData
+from cellrank import logging as logg
+from cellrank._key import Key
+from cellrank.tl._enum import ModeEnum
+from cellrank.ul._docs import d
+from cellrank.tl._utils import (
+    save_fig,
+    _eigengap,
+    _fuzzy_to_discrete,
+    _series_from_one_hot_matrix,
+)
+from cellrank.tl._colors import _get_black_or_white, _create_categorical_colors
+from cellrank.tl._lineage import Lineage
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.estimators.mixins import EigenMixin, SchurMixin, LinDriversMixin
+from cellrank.tl.kernels._base_kernel import KernelExpression
+from cellrank.tl.estimators.mixins._utils import logger, shadow, register_plotter
+from cellrank.tl.estimators.terminal_states._term_states_estimator import (
+    TermStatesEstimator,
+)
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import spmatrix
+from pandas.api.types import infer_dtype, is_categorical_dtype
+
+import matplotlib.pyplot as plt
+from matplotlib.axes import Axes
+from matplotlib.colors import Normalize, ListedColormap
+from matplotlib.ticker import StrMethodFormatter
+from matplotlib.colorbar import ColorbarBase
+
+
+class TermStatesMethod(ModeEnum):  # noqa: D101
+    EIGENGAP = auto()
+    EIGENGAP_COARSE = auto()
+    TOP_N = auto()
+    STABILITY = auto()
+
+
+@d.dedent
+class GPCCA(TermStatesEstimator, LinDriversMixin, SchurMixin, EigenMixin):
+    """
+    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
+    assignment is soft, i.e. each cell is assigned to every macrostate with a certain weight, where weights sum to
+    one per cell. Macrostates are computed by maximizing the 'crispness' which can be thought of as a measure for
+    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 :cite:`reuter:18`.
+
+    Parameters
+    ----------
+    %(base_estimator.parameters)s
+    """
+
+    def __init__(
+        self,
+        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
+        obsp_key: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(obj=obj, obsp_key=obsp_key, **kwargs)
+
+        self._coarse_init_dist: Optional[pd.Series] = None
+        self._coarse_stat_dist: Optional[pd.Series] = None
+        self._coarse_tmat: Optional[pd.DataFrame] = None
+
+        self._macrostates: Optional[pd.Series] = None
+        self._macrostates_memberships: Optional[Lineage] = None
+        self._macrostates_colors: Optional[np.ndarray] = None
+
+        self._term_states_memberships: Optional[Lineage] = None
+
+    @property
+    @d.get_summary(base="gpcca_macro")
+    def macrostates(self) -> Optional[pd.Series]:
+        """Macrostates of the transition matrix."""
+        return self._macrostates
+
+    @property
+    @d.get_summary(base="gpcca_macro_memberships")
+    def macrostates_memberships(self) -> Optional[Lineage]:
+        """Macrostate membership matrix.
+
+        Soft assignment of microstates (cells) to macrostates.
+        """
+        return self._macrostates_memberships
+
+    @property
+    @d.get_summary(base="gpcca_term_states_memberships")
+    def terminal_states_memberships(self) -> Optional[Lineage]:
+        """Terminal state membership matrix.
+
+        Soft assignment of cells to terminal states.
+        """
+        return self._term_states_memberships
+
+    @property
+    @d.get_summary(base="gpcca_coarse_init")
+    def coarse_initial_distribution(self) -> Optional[pd.Series]:
+        """Coarse-grained initial distribution."""
+        return self._coarse_init_dist
+
+    @property
+    @d.get_summary(base="gpcca_coarse_stat")
+    def coarse_stationary_distribution(self) -> Optional[pd.Series]:
+        """Coarse-grained stationary distribution."""
+        return self._coarse_stat_dist
+
+    @property
+    @d.get_summary(base="gpcca_coarse_tmat")
+    def coarse_T(self) -> Optional[pd.DataFrame]:
+        """Coarse-grained transition matrix."""
+        return self._coarse_tmat
+
+    @d.get_sections(base="gpcca_compute_macro", sections=["Parameters", "Returns"])
+    @d.dedent
+    def compute_macrostates(
+        self,
+        n_states: Optional[Union[int, Sequence[int]]] = None,
+        n_cells: Optional[int] = 30,
+        cluster_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Compute the macrostates.
+
+        Parameters
+        ----------
+        n_states
+            Number of macrostates. If a :class:`typing.Sequence`, use the *minChi* criterion :cite:`reuter:18`.
+            If `None`, use the *eigengap* heuristic.
+        %(n_cells)s
+        cluster_key
+            If a key to cluster labels is given, names and colors of the states will be associated with the clusters.
+        kwargs
+            Keyword arguments for :meth:`compute_schur`.
+
+        Returns
+        -------
+        Updates the following fields:
+
+            - :attr:`macrostates` - %(gpcca_macro.summary)s
+            - :attr:`macrostates_memberships` - %(gpcca_macro_memberships.summary)s
+            - :attr:`coarse_T` - %(gpcca_coarse_tmat.summary)s
+            - :attr:`coarse_initial_distribution` - %(gpcca_coarse_init.summary)s
+            - :attr:`coarse_stationary_distribution` - %(gpcca_coarse_stat.summary)s
+            - :attr:`schur_vectors` - %(schur_vectors.summary)s
+            - :attr:`schur_matrix` - %(schur_matrix.summary)s
+            - :attr:`eigendecomposition` - %(eigen.summary)s
+        """
+
+        n_states = self._n_states(n_states)
+        if n_states == 1:
+            self._compute_one_macrostate(
+                n_cells=n_cells,
+                cluster_key=cluster_key,
+            )
+            return
+
+        if self._gpcca is None or kwargs:
+            self.compute_schur(n_states, **kwargs)
+        n_states = self._validate_n_states(n_states)
+
+        if self._gpcca._p_X.shape[1] < n_states:
+            # precomputed X
+            logg.warning(
+                f"Requested more macrostates `{n_states}` than available "
+                f"Schur vectors `{self._gpcca._p_X.shape[1]}`. Recomputing the decomposition"
+            )
+
+        start = logg.info(f"Computing `{n_states}` macrostates")
+        try:
+            self._gpcca = self._gpcca.optimize(m=n_states)
+        except ValueError as e:
+            if "will split complex conjugate eigenvalues" not in str(e):
+                raise
+            # this is the following case - we have 4 Schur vectors, user requests 5 states, but it splits the conj. ev.
+            # in the try block, Schur decomposition with 5 vectors is computed, but it fails (no way of knowing)
+            # so in this case, we increase it by 1
+            logg.warning(
+                f"Unable to compute macrostates with `n_states={n_states}` because it will "
+                f"split complex conjugate eigenvalues. Using `n_states={n_states + 1}`"
+            )
+            self._gpcca = self._gpcca.optimize(m=n_states + 1)
+
+        self._set_macrostates(
+            memberships=self._gpcca.memberships,
+            n_cells=n_cells,
+            cluster_key=cluster_key,
+            params=self._create_params(),
+            time=start,
+        )
+
+    @d.dedent
+    def predict(
+        self,
+        method: Literal[
+            "stability", "top_n", "eigengap", "eigengap_coarse"
+        ] = TermStatesMethod.STABILITY,
+        n_cells: int = 30,
+        alpha: Optional[float] = 1,
+        stability_threshold: float = 0.96,
+        n_states: Optional[int] = None,
+    ) -> None:
+        """
+        Automatically select terminal states from macrostates.
+
+        Parameters
+        ----------
+        method
+            How to select the terminal states. Valid option are:
+
+                - `'eigengap'` - select the number of states based on the *eigengap* of :attr:`transition_matrix`.
+                - `'eigengap_coarse'` - select the number of states based on the *eigengap* of the diagonal
+                  of :attr:`coarse_T`.
+                - `'top_n'` - select top ``n_states`` based on the probability of the diagonal of :attr:`coarse_T`.
+                - `'stability'` - select states which have a stability >= ``stability_threshold``.
+                  The stability is given by the diagonal elements of :attr:`coarse_T`.
+        %(n_cells)s
+        alpha
+            Weight given to the deviation of an eigenvalue from one.
+            Only used when ``method = 'eigengap'`` or ``method = 'eigengap_coarse'``.
+        stability_threshold
+            Threshold used when ``method = 'stability'``.
+        n_states
+            Number of states used when ``method = 'top_n'``.
+
+        Returns
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+            - :attr:`terminal_states_memberships` - %(gpcca_term_states_memberships.summary)s
+            - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
+        """
+        if self.macrostates is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+
+        # fmt: off
+        if len(self._macrostates.cat.categories) == 1:
+            logg.warning("Found only one macrostate. Making it the single terminal state")
+            self.set_terminal_states_from_macrostates(None, n_cells=n_cells, params=self._create_params())
+            return
+
+        method = TermStatesMethod(method)
+        eig = self.eigendecomposition
+        coarse_T = self.coarse_T
+
+        if method == TermStatesMethod.EIGENGAP:
+            if eig is None:
+                raise RuntimeError("Compute eigendecomposition first as `.compute_eigendecomposition()`.")
+            n_states = _eigengap(eig["D"], alpha=alpha) + 1
+        elif method == TermStatesMethod.EIGENGAP_COARSE:
+            if coarse_T is None:
+                raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+            n_states = _eigengap(np.sort(np.diag(coarse_T)[::-1]), alpha=alpha)
+        elif method == TermStatesMethod.TOP_N:
+            if n_states is None:
+                raise ValueError("Expected `n_states != None` for `method='top_n'`.")
+            elif n_states <= 0:
+                raise ValueError(f"Expected `n_states` to be positive, found `{n_states}`.")
+        elif method == TermStatesMethod.STABILITY:
+            if stability_threshold is None:
+                raise ValueError("Expected `stability_threshold != None` for `method='stability'`.")
+            stability = pd.Series(np.diag(coarse_T), index=coarse_T.columns)
+            names = stability[stability.values >= stability_threshold].index
+            self.set_terminal_states_from_macrostates(names, n_cells=n_cells, params=self._create_params())
+            return
+        else:
+            raise NotImplementedError(f"Method `{method}` is not yet implemented.")
+        # fmt: on
+
+        names = coarse_T.columns[np.argsort(np.diag(coarse_T))][-n_states:]
+        self.set_terminal_states_from_macrostates(
+            names, n_cells=n_cells, params=self._create_params()
+        )
+
+        return
+
+    @d.dedent
+    def set_terminal_states_from_macrostates(
+        self,
+        names: Optional[Union[str, Sequence[str], Mapping[str, str]]] = None,
+        n_cells: int = 30,
+        **kwargs: Any,
+    ):
+        """
+        Manually select terminal states from macrostates.
+
+        Parameters
+        ----------
+        names
+            Names of the macrostates to be marked as terminal. Multiple states can be combined using `','`,
+            such as ``["Alpha, Beta", "Epsilon"]``.  If a :class:`dict`, keys correspond to the names
+            of the macrostates and the values to the new names.  If `None`, select all macrostates.
+        %(n_cells)s
+
+        Returns
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+            - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
+            - :attr:`terminal_states_probabilities_memberships` - %(gpcca_term_states_memberships.summary)s
+        """
+        if n_cells <= 0:
+            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
+
+        memberships = self.macrostates_memberships
+        if memberships is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+
+        rename = True
+        if names is None:
+            names = memberships.names
+            rename = False
+        if isinstance(names, str):
+            names = [names]
+            rename = False
+        if not isinstance(names, dict):
+            names = {n: n for n in names}
+            rename = False
+        if not len(names):
+            raise ValueError("No macrostates have been selected.")
+
+        # we do this also here because if `rename_terminal_states` fails
+        # invalid states would've been written to this object and nothing to adata
+        names = {str(k): str(v) for k, v in names.items()}
+        names_after_renaming = {names.get(n, n) for n in memberships.names}
+        if len(names_after_renaming) != memberships.shape[1]:
+            raise ValueError(
+                f"After renaming, terminal state names will no longer be unique: `{names_after_renaming}`."
+            )
+
+        # this also checks that the names are correct before renaming
+        is_singleton = memberships.shape[1] == 1
+        memberships = memberships[list(names.keys())].copy()
+
+        states = self._create_states(memberships, n_cells=n_cells, check_row_sums=False)
+        if is_singleton:
+            colors = self._macrostates_colors.copy()
+            probs = memberships.X.squeeze() / memberships.X.max()
+        else:
+            colors = memberships[list(states.cat.categories)].colors
+            probs = (memberships.X / memberships.X.max(0)).max(1)
+        probs = pd.Series(probs, index=self.adata.obs_names)
+
+        self._write_terminal_states(
+            states, colors, probs, memberships, params=kwargs.pop("params", {})
+        )
+        if rename:
+            # TODO(michalk8): in a future PR, remove this behavior in Lineage
+            # access lineage renames join states, e.g. 'Alpha, Beta' becomes 'Alpha or Beta' + whitespace stripping
+            self.rename_terminal_states(
+                dict(zip(self.terminal_states.cat.categories, names.values()))
+            )
+
+    @d.dedent
+    def rename_terminal_states(self, new_names: Mapping[str, str]) -> None:
+        """
+        %(tse_rename_term_states.full_desc)s
+
+        Parameters
+        ----------
+        %(tse_rename_term_states.parameters)s
+
+        Returns
+        -------
+        %(tse_rename_term_states.returns)s
+            - :attr:`terminal_states_memberships` - %(gpcca_term_states_memberships.summary)s
+        """  # noqa: D400
+        term_states_memberships = self.terminal_states_memberships
+        super().rename_terminal_states(new_names)
+
+        # fmt: off
+        new_names = {str(k): str(v) for k, v in new_names.items()}
+        term_states_memberships.names = [new_names.get(n, n) for n in term_states_memberships.names]
+        self._set("_term_states_memberships", value=term_states_memberships, shadow_only=True)
+        # fmt: on
+
+        with self._shadow:
+            key = Key.obsm.memberships(Key.obs.macrostates(self.backward))
+            self._set(obj=self.adata.obsm, key=key, value=term_states_memberships)
+
+    @d.dedent
+    def fit(
+        self,
+        n_states: Optional[Union[int, Sequence[int]]] = None,
+        n_cells: Optional[int] = 30,
+        cluster_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> "GPCCA":
+        """
+        %(tse_fit.full_desc)s
+
+        Parameters
+        ----------
+        %(gpcca_compute_macro.parameters)s
+
+        Returns
+        -------
+        %(gpcca_compute_macro.returns)s
+        """  # noqa: D400
+        if n_states is None:
+            self.compute_eigendecomposition()
+            n_states = self.eigendecomposition["eigengap"] + 1
+        if isinstance(n_states, int) and n_states == 1:
+            self.compute_eigendecomposition()
+
+        self.compute_macrostates(n_states=n_states, cluster_key=cluster_key, **kwargs)
+
+        return self
+
+    @d.dedent
+    def plot_coarse_T(
+        self,
+        show_stationary_dist: bool = True,
+        show_initial_dist: bool = False,
+        cmap: Union[str, ListedColormap] = "viridis",
+        xtick_rotation: float = 45,
+        annotate: bool = True,
+        show_cbar: bool = True,
+        title: Optional[str] = None,
+        figsize: Tuple[float, float] = (8, 8),
+        dpi: int = 80,
+        save: Optional[Union[Path, str]] = None,
+        text_kwargs: Mapping[str, Any] = MappingProxyType({}),
+        **kwargs: Any,
+    ) -> None:
+        """
+        Plot the coarse-grained transition matrix between macrostates.
+
+        Parameters
+        ----------
+        show_stationary_dist
+            Whether to show :attr:`coarse_stationary_distribution`, if present.
+        show_initial_dist
+            Whether to show :attr:`coarse_initial_distribution`.
+        cmap
+            Colormap to use.
+        xtick_rotation
+            Rotation of ticks on the x-axis.
+        annotate
+            Whether to display the text on each cell.
+        show_cbar
+            Whether to show colorbar.
+        title
+            Title of the figure.
+        %(plotting)s
+        text_kwargs
+            Keyword arguments for :func:`matplotlib.pyplot.text`.
+        kwargs
+            Keyword arguments for :func:`matplotlib.pyplot.imshow`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+
+        def stylize_dist(
+            ax: Axes, data: np.ndarray, xticks_labels: Sequence[str] = ()
+        ) -> None:
+            _ = ax.imshow(data, aspect="auto", cmap=cmap, norm=norm)
+            for spine in ax.spines.values():
+                spine.set_visible(False)
+
+            if xticks_labels is not None:
+                ax.set_xticks(np.arange(data.shape[1]))
+                ax.set_xticklabels(xticks_labels)
+                plt.setp(
+                    ax.get_xticklabels(),
+                    rotation=xtick_rotation,
+                    ha="right",
+                    rotation_mode="anchor",
+                )
+            else:
+                ax.set_xticks([])
+                ax.tick_params(
+                    which="both", top=False, right=False, bottom=False, left=False
+                )
+
+            ax.set_yticks([])
+
+        def annotate_heatmap(im, valfmt: str = "{x:.2f}") -> None:
+            # modified from matplotlib's site
+
+            data = im.get_array()
+            kw = {"ha": "center", "va": "center"}
+            kw.update(**text_kwargs)
+
+            # Get the formatter in case a string is supplied
+            if isinstance(valfmt, str):
+                valfmt = StrMethodFormatter(valfmt)
+
+            # Loop over the data and create a `Text` for each "pixel".
+            # Change the text's color depending on the data.
+            texts = []
+            for i in range(data.shape[0]):
+                for j in range(data.shape[1]):
+                    kw.update(color=_get_black_or_white(im.norm(data[i, j]), cmap))
+                    text = im.axes.text(j, i, valfmt(data[i, j], None), **kw)
+                    texts.append(text)
+
+        def annotate_dist_ax(ax, data: np.ndarray, valfmt: str = "{x:.2f}"):
+            if ax is None:
+                return
+
+            if isinstance(valfmt, str):
+                valfmt = StrMethodFormatter(valfmt)
+
+            kw = {"ha": "center", "va": "center"}
+            kw.update(**text_kwargs)
+
+            for i, val in enumerate(data):
+                kw.update(color=_get_black_or_white(im.norm(val), cmap))
+                ax.text(
+                    i,
+                    0,
+                    valfmt(val, None),
+                    **kw,
+                )
+
+        coarse_T = self.coarse_T
+        coarse_init_d = self.coarse_initial_distribution
+        coarse_stat_d = self.coarse_stationary_distribution
+
+        if coarse_T is None:
+            raise RuntimeError(
+                "Compute coarse-grained transition matrix first as `.compute_macrostates()` with `n_states > 1`."
+            )
+
+        if show_stationary_dist and coarse_stat_d is None:
+            logg.warning("Coarse stationary distribution is `None`, ignoring")
+            show_stationary_dist = False
+        if show_initial_dist and coarse_init_d is None:
+            logg.warning("Coarse initial distribution is `None`, ignoring")
+            show_initial_dist = False
+
+        hrs, wrs = [1], [1]
+        if show_stationary_dist:
+            hrs += [0.05]
+        if show_initial_dist:
+            hrs += [0.05]
+        if show_cbar:
+            wrs += [0.025]
+
+        dont_show_dist = not show_initial_dist and not show_stationary_dist
+
+        fig = plt.figure(constrained_layout=False, figsize=figsize, dpi=dpi)
+        gs = plt.GridSpec(
+            1 + show_stationary_dist + show_initial_dist,
+            1 + show_cbar,
+            height_ratios=hrs,
+            width_ratios=wrs,
+            wspace=0.05,
+            hspace=0.05,
+        )
+        if isinstance(cmap, str):
+            cmap = plt.get_cmap(cmap)
+
+        ax = fig.add_subplot(gs[0, 0])
+        cax = fig.add_subplot(gs[:1, -1]) if show_cbar else None
+        init_ax, stat_ax = None, None
+
+        labels = list(self.coarse_T.columns)
+
+        tmp = coarse_T
+        if show_initial_dist:
+            tmp = np.c_[tmp, coarse_stat_d]
+        if show_initial_dist:
+            tmp = np.c_[tmp, coarse_init_d]
+
+        minn, maxx = np.nanmin(tmp), np.nanmax(tmp)
+        norm = Normalize(vmin=minn, vmax=maxx)
+
+        if show_stationary_dist:
+            stat_ax = fig.add_subplot(gs[1, 0])
+            stylize_dist(
+                stat_ax,
+                np.array(coarse_stat_d).reshape(1, -1),
+                xticks_labels=labels if not show_initial_dist else None,
+            )
+            stat_ax.yaxis.set_label_position("right")
+            stat_ax.set_ylabel("stationary dist", rotation=0, ha="left", va="center")
+
+        if show_initial_dist:
+            init_ax = fig.add_subplot(gs[show_stationary_dist + show_initial_dist, 0])
+            stylize_dist(
+                init_ax, np.array(coarse_init_d).reshape(1, -1), xticks_labels=labels
+            )
+
+            init_ax.yaxis.set_label_position("right")
+            init_ax.set_ylabel("initial dist", rotation=0, ha="left", va="center")
+
+        im = ax.imshow(coarse_T, aspect="auto", cmap=cmap, norm=norm, **kwargs)
+        ax.set_title("coarse-grained transition matrix" if title is None else title)
+
+        if cax is not None:
+            _ = ColorbarBase(
+                cax,
+                cmap=cmap,
+                norm=norm,
+                ticks=np.linspace(minn, maxx, 10),
+                format="%0.3f",
+            )
+
+        ax.set_yticks(np.arange(coarse_T.shape[0]))
+        ax.set_yticklabels(labels)
+
+        ax.tick_params(
+            top=False,
+            bottom=dont_show_dist,
+            labeltop=False,
+            labelbottom=dont_show_dist,
+        )
+
+        for spine in ax.spines.values():
+            spine.set_visible(False)
+
+        if dont_show_dist:
+            ax.set_xticks(np.arange(coarse_T.shape[1]))
+            ax.set_xticklabels(labels)
+            plt.setp(
+                ax.get_xticklabels(),
+                rotation=xtick_rotation,
+                ha="right",
+                rotation_mode="anchor",
+            )
+        else:
+            ax.set_xticks([])
+
+        ax.set_yticks(np.arange(coarse_T.shape[0] + 1) - 0.5, minor=True)
+        ax.tick_params(which="minor", bottom=dont_show_dist, left=False, top=False)
+
+        if annotate:
+            annotate_heatmap(im)
+            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.values)
+
+        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:`anndata.AnnData.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.macrostates
+        if macrostates is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+        if key not in self.adata.obs:
+            raise KeyError(f"Data not found in `adata.obs[{key!r}]`.")
+        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 _n_states(self, n_states: Optional[Union[int, Sequence[int]]]) -> int:
+        if n_states is None:
+            if self.eigendecomposition is None:
+                raise RuntimeError(
+                    "Compute eigendecomposition first as `.compute_eigendecomposition()` or "
+                    "supply `n_states != None`."
+                )
+            return self.eigendecomposition["eigengap"] + 1
+
+        # fmt: off
+        if isinstance(n_states, int):
+            if n_states <= 0:
+                raise ValueError(f"Expected `n_states` to be positive, found `{n_states}`.")
+            return n_states
+
+        if self._gpcca is None:
+            raise RuntimeError("Compute Schur decomposition first as `.compute_schur()`.")
+
+        if not isinstance(n_states, Sequence):
+            raise TypeError(f"Expected `n_states` to be a `Sequence`, found `{type(n_states).__name__!r}`.")
+        if len(n_states) != 2:
+            raise ValueError(f"Expected `n_states` to be of size `2`, found `{len(n_states)}`.")
+
+        minn, maxx = sorted(n_states)
+        if minn <= 1:
+            minn = 2
+            logg.warning(f"Minimum value must be larger than `1`, found `{minn}`. Setting `min={minn}`")
+        if minn == 2:
+            minn = 3
+            logg.warning(
+                f"In most cases, 2 clusters will always be optimal. "
+                f"If you really expect 2 clusters, use `n_states=2`. Setting `min={minn}`"
+            )
+        # fmt: on
+        maxx = max(minn + 1, maxx)
+
+        logg.info(f"Calculating minChi criterion in interval `[{minn}, {maxx}]`")
+
+        return int(np.arange(minn, maxx + 1)[np.argmax(self._gpcca.minChi(minn, maxx))])
+
+    def _create_states(
+        self,
+        probs: Union[np.ndarray, Lineage],
+        n_cells: int,
+        check_row_sums: bool = False,
+        return_not_enough_cells: bool = False,
+    ) -> Union[pd.Series, Tuple[pd.Series, np.ndarray]]:
+        if n_cells <= 0:
+            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
+
+        discrete, not_enough_cells = _fuzzy_to_discrete(
+            a_fuzzy=probs,
+            n_most_likely=n_cells,
+            remove_overlap=False,
+            raise_threshold=0.2,
+            check_row_sums=check_row_sums,
+        )
+
+        states = _series_from_one_hot_matrix(
+            membership=discrete,
+            index=self.adata.obs_names,
+            names=probs.names if isinstance(probs, Lineage) else None,
+        )
+
+        return (states, not_enough_cells) if return_not_enough_cells else states
+
+    def _validate_n_states(self, n_states: int) -> int:
+        if self._invalid_n_states is not None and n_states in self._invalid_n_states:
+            logg.warning(
+                f"Unable to compute macrostates with `n_states={n_states}` because it will "
+                f"split complex conjugate eigenvalues. Using `n_states={n_states + 1}`"
+            )
+            n_states += 1  # cannot force recomputation of the Schur decomposition
+            assert n_states not in self._invalid_n_states, "Sanity check failed."
+
+        return n_states
+
+    def _compute_one_macrostate(
+        self,
+        n_cells: Optional[int],
+        cluster_key: Optional[str],
+    ) -> None:
+        start = logg.info("For 1 macrostate, stationary distribution is computed")
+
+        eig = self.eigendecomposition
+        if (
+            eig is not None
+            and "stationary_dist" in eig
+            and eig["params"]["which"] == "LR"
+        ):
+            stationary_dist = eig["stationary_dist"]
+        else:
+            self.compute_eigendecomposition(only_evals=False, which="LR")
+            stationary_dist = self.eigendecomposition["stationary_dist"]
+
+        self._set_macrostates(
+            memberships=stationary_dist[:, None],
+            n_cells=n_cells,
+            cluster_key=cluster_key,
+            check_row_sums=False,
+            time=start,
+        )
+
+    @d.dedent
+    def _set_macrostates(
+        self,
+        memberships: np.ndarray,
+        n_cells: Optional[int] = 30,
+        cluster_key: str = "clusters",
+        check_row_sums: bool = True,
+        time: Optional[datetime] = None,
+        params: Dict[str, Any] = MappingProxyType({}),
+    ) -> None:
+        """
+        Map fuzzy clustering to pre-computed annotations to get names and colors.
+
+        Given the fuzzy clustering, we would like to select the most likely cells from each state and use these to
+        give each state a name and a color by comparing with pre-computed, categorical cluster annotations.
+
+        Parameters
+        ----------
+        memberships
+            Fuzzy clustering.
+        %(n_cells)s
+        cluster_key
+            Key from :attr:`anndata.AnnData.obs` to get reference cluster annotations.
+        check_row_sums
+            Check whether rows in `memberships` sum to `1`.
+        time
+            Start time of macrostates computation.
+        params
+            Parameters used in macrostates computation.
+
+        Returns
+        -------
+        Nothing, just update the field as described in :meth:`compute_macrostates`.
+        """
+
+        if n_cells is None:
+            # fmt: off
+            logg.debug("Setting the macrostates using macrostate assignment")
+            assignment = pd.Series(np.argmax(memberships, axis=1).astype(str), dtype="category")
+            # sometimes, a category can be missing
+            assignment = assignment.cat.reorder_categories([str(i) for i in range(memberships.shape[1])])
+            not_enough_cells = []
+            # fmt: on
+        else:
+            logg.debug("Setting the macrostates using macrostates memberships")
+
+            # select the most likely cells from each macrostate
+            assignment, not_enough_cells = self._create_states(
+                memberships,
+                n_cells=n_cells,
+                check_row_sums=check_row_sums,
+                return_not_enough_cells=True,
+            )
+
+        # remove previous fields
+        self._write_terminal_states(None, None, None, None, log=False)
+
+        # fmt: off
+        assignment, colors = self._set_categorical_labels(assignment, cluster_key=cluster_key)
+        memberships = Lineage(memberships, names=list(assignment.cat.categories), colors=colors)
+        # fmt: on
+
+        groups = pd.DataFrame(assignment).groupby(0).size()
+        groups = groups[groups != n_cells].to_dict()
+        if len(groups):
+            logg.warning(
+                f"The following terminal states have different number "
+                f"of cells than requested ({n_cells}): {groups}"
+            )
+
+        self._write_macrostates(
+            assignment, colors, memberships, time=time, params=params
+        )
+
+    @logger
+    @shadow
+    def _write_macrostates(
+        self,
+        macrostates: pd.Series,
+        colors: np.ndarray,
+        memberships: Lineage,
+        params: Dict[str, Any] = MappingProxyType({}),
+    ) -> str:
+        # fmt: off
+        names = list(macrostates.cat.categories)
+
+        key = Key.obs.macrostates(self.backward)
+        self._set("_macrostates", obj=self.adata.obs, key=key, value=macrostates, shadow_only=True)
+        ckey = Key.uns.colors(key)
+        self._set("_macrostates_colors", obj=self.adata.uns, key=ckey, value=colors, shadow_only=True)
+        mkey = Key.obsm.memberships(key)
+        self._set("_macrostates_memberships", obj=self.adata.obsm, key=mkey, value=memberships, shadow_only=True)
+        self.params[key] = dict(params)
+
+        if len(names) > 1:
+            # not using stationary distribution
+            g = self._gpcca
+            tmat = pd.DataFrame(g.coarse_grained_transition_matrix, index=names, columns=names)
+            init_dist = pd.Series(g.coarse_grained_input_distribution, index=names)
+            stat_dist = pd.Series(g.coarse_grained_stationary_probability, index=names)
+            dists = pd.DataFrame({"coarse_init_dist": init_dist})
+            if stat_dist is not None:
+                dists["coarse_stat_dist"] = pd.Series(stat_dist, index=names)
+
+            key = Key.obsm.schur_vectors(self.backward)
+            self._set("_schur_vectors", obj=self.adata.obsm, key=key, value=g._p_X, shadow_only=True)
+            key = Key.uns.schur_matrix(self.backward)
+            self._set("_schur_matrix", obj=self.adata.uns, key=key, value=g._p_R, shadow_only=True)
+            self._set("_coarse_tmat", value=tmat, shadow_only=True)
+            self._set("_coarse_init_dist", value=init_dist, shadow_only=True)
+            self._set("_coarse_stat_dist", value=stat_dist, shadow_only=True)
+            self._set(obj=self.adata.uns, key=Key.uns.coarse(self.backward), value=AnnData(tmat, obs=dists))
+        else:
+            for attr in ["_schur_vectors", "_schur_matrix", "_coarse_tmat", "_coarse_init_dist", "_coarse_stat_dist"]:
+                self._set(attr, value=None, shadow_only=True)
+            self._set(obj=self.adata.uns, key=Key.uns.coarse(self.backward), value=None)
+        # fmt: on
+
+        return (
+            "Adding `.macrostates`\n"
+            "       `.macrostates_memberships`\n"
+            "       `.coarse_T`\n"
+            "       `.coarse_initial_distribution\n"
+            "       `.coarse_stationary_distribution`\n"
+            "       `.schur_vectors`\n"
+            "       `.schur_matrix`\n"
+            "       `.eigendecomposition`\n"
+            "    Finish"
+        )
+
+    @logger
+    @shadow
+    def _write_terminal_states(
+        self,
+        states: Optional[pd.Series],
+        colors: Optional[np.ndarray],
+        probs: Optional[pd.Series] = None,
+        memberships: Optional[Lineage] = None,
+        params: Dict[str, Any] = MappingProxyType({}),
+    ) -> str:
+        # fmt: off
+        msg = super()._write_terminal_states(states, colors, probs, params=params, log=False)
+        msg = "\n".join(msg.split("\n")[:-1])
+        msg += "\n       `.terminal_states_memberships\n    Finish`"
+
+        self._write_absorption_probabilities(None, None, log=False)
+        key = Key.obsm.memberships(Key.obs.term_states(self.backward))
+        self._set("_term_states_memberships", obj=self.adata.obsm, key=key, value=memberships)
+        # fmt: on
+
+        return msg
+
+    def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
+        _ = self._read_eigendecomposition(adata, allow_missing=True)
+        # TODO(michalk8): reintroduce in 2.0
+        ok = self._read_schur_decomposition(adata, allow_missing=True)
+        if not ok:
+            return False
+
+        # fmt: off
+        with SafeGetter(self, allowed=KeyError) as sg:
+            key = Key.obs.macrostates(self.backward)
+            # TODO(michalk8): in the future, be more stringent and ensure the categories match the macro memberships
+            self._get("_macrostates", self.adata.obs, key=key, where="obs", dtype=pd.Series)
+            ckey = Key.uns.colors(key)
+            self._get("_macrostates_colors", self.adata.uns, key=ckey, where="uns", dtype=(list, tuple, np.ndarray))
+            mkey = Key.obsm.memberships(key)
+            self._get("_macrostates_memberships", self.adata.obsm, key=mkey, where="obsm", dtype=(Lineage, np.ndarray))
+            self._ensure_lineage_object("_macrostates_memberships", kind="macrostates")
+
+            self._macrostates_colors = self.macrostates_memberships.colors.copy()
+            self.params[key] = self._read_params(key)
+
+            # TODO(michalk8): allow missing?
+            tmat: AnnData = self.adata.uns[Key.uns.coarse(self.backward)]
+            if not isinstance(tmat, AnnData):
+                raise TypeError(f"Expected coarse-grained transition matrix to be stored "
+                                f"as `anndata.AnnData`, found `{type(tmat).__name__}`.")
+            tmat = tmat.copy()
+            names = tmat.obs_names
+
+            self._coarse_tmat = pd.DataFrame(tmat.X, index=names, columns=names)
+            self._coarse_init_dist = tmat.obs["coarse_init_dist"]
+            self._coarse_stat_dist = tmat.obs.get("coarse_stat_dist", None)
+
+            self._set(obj=self._shadow_adata.uns, key=Key.uns.coarse(self.backward), value=tmat)
+
+        # TODO(michalk8): reintroduce this in 2.0 - this is done for high-level plotting of init/term states only
+        # if not sg.ok:
+        #    return False
+
+        if not super()._read_from_adata(adata, **kwargs):
+            return False
+
+        with SafeGetter(self, allowed=KeyError) as sg:
+            key = Key.obsm.memberships(Key.obs.term_states(self.backward))
+            self._get("_term_states_memberships", self.adata.obsm, key=key, where="obsm", dtype=(np.ndarray, Lineage))
+            self._ensure_lineage_object("_term_states_memberships", kind="term_states")
+        # fmt: on
+
+        return sg.ok and self._read_absorption_probabilities(adata)
+
+    plot_macrostates = register_plotter(
+        discrete="macrostates", continuous="macrostates_memberships"
+    )
+    plot_terminal_states = register_plotter(
+        discrete="terminal_states", continuous="terminal_states_memberships"
+    )
+
+    @d.dedent
+    def _compute_initial_states(self, n_states: int = 1, n_cells: int = 30) -> None:
+        """
+        Compute initial states from macrostates using :attr:`coarse_stationary_distribution`.
+
+        Parameters
+        ----------
+        n_states
+            Number of initial states.
+        %(n_cells)s
+
+        Returns
+        -------
+        %(set_initial_states_from_macrostates.returns)s
+        """
+
+        if n_states <= 0:
+            raise ValueError(f"Expected `n_states` to be positive, found `{n_states}`.")
+
+        if n_cells <= 0:
+            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
+
+        probs = self.macrostates_memberships
+        if probs is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+
+        if n_states > probs.shape[1]:
+            raise ValueError(
+                f"Requested `{n_states}` initial states, but only `{probs.shape[1]}` macrostates have been computed."
+            )
+
+        if probs.shape[1] == 1:
+            self._set_initial_states_from_macrostates(n_cells=n_cells)
+            return
+
+        stat_dist = self.coarse_stationary_distribution
+        if stat_dist is None:
+            raise RuntimeError("No coarse-grained stationary distribution found.")
+
+        self._set_initial_states_from_macrostates(
+            stat_dist[np.argsort(stat_dist)][:n_states].index, n_cells=n_cells
+        )
+
+    @d.get_sections(base="set_initial_states_from_macrostates", sections=["Returns"])
+    @d.dedent
+    def _set_initial_states_from_macrostates(
+        self,
+        names: Optional[Union[str, Sequence[str]]] = None,
+        n_cells: int = 30,
+    ) -> None:
+        """
+        Manually select initial states from macrostates.
+
+        Note that no check is performed to ensure initial and terminal states are distinct.
+
+        Parameters
+        ----------
+        names
+            Names of the macrostates to be marked as initial states. Multiple states can be combined using `','`,
+            such as `["Alpha, Beta", "Epsilon"]`.
+        %(n_cells)s
+
+        Returns
+        -------
+        Nothing, just modifies :attr:`anndata.AnnData.obs`. The actual keys depend of :attr:`backward`.
+        """
+
+        if not isinstance(n_cells, int):
+            raise TypeError(
+                f"Expected `n_cells` to be of type `int`, found `{type(n_cells).__name__}`."
+            )
+
+        if n_cells <= 0:
+            raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
+
+        probs = self.macrostates_memberships
+        if probs is None:
+            raise RuntimeError("Compute macrostates first as `.compute_macrostates()`.")
+        elif probs.shape[1] == 1:
+            categorical = self._create_states(probs, n_cells=n_cells)
+            scaled = probs / probs.max()
+        else:
+            if names is None:
+                names = probs.names
+            if isinstance(names, str):
+                names = [names]
+
+            probs = probs[list(names)]
+            categorical = self._create_states(probs, n_cells=n_cells)
+            probs /= probs.max(0)
+
+            # compute the aggregated probability of being a initial/terminal state (no matter which)
+            scaled = probs.X.max(1)
+
+        self._write_initial_states(membership=probs, probs=scaled, cats=categorical)
+
+    def _write_initial_states(
+        self, membership: Lineage, probs: pd.Series, cats: pd.Series, time=None
+    ) -> None:
+        key = Key.obs.term_states(not self.backward)
+
+        self.adata.obs[key] = cats
+        self.adata.obs[Key.obs.probs(key)] = probs
+
+        self.adata.uns[Key.uns.colors(key)] = membership.colors
+
+        logg.info(
+            f"Adding `adata.obs[{key!r}]`\n       `adata.obs[{Key.uns.colors(key)!r}]`\n",
+            time=time,
+        )
diff --git a/cellrank/tl/estimators/terminal_states/_term_states_estimator.py b/cellrank/tl/estimators/terminal_states/_term_states_estimator.py
new file mode 100644
index 000000000..36f9035e9
--- /dev/null
+++ b/cellrank/tl/estimators/terminal_states/_term_states_estimator.py
@@ -0,0 +1,306 @@
+from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
+
+from abc import ABC
+from types import MappingProxyType
+
+from anndata import AnnData
+from cellrank._key import Key
+from cellrank.ul._docs import d
+from cellrank.tl._utils import _merge_categorical_series, _convert_to_categorical_series
+from cellrank.tl._colors import (
+    _map_names_and_colors,
+    _convert_to_hex_colors,
+    _create_categorical_colors,
+)
+from cellrank.tl.estimators import BaseEstimator
+from cellrank.tl.estimators._utils import SafeGetter
+from cellrank.tl.kernels._base_kernel import KernelExpression
+from cellrank.tl.estimators.mixins._utils import logger, shadow, register_plotter
+
+import numpy as np
+import pandas as pd
+from scipy.sparse import spmatrix
+from pandas.api.types import infer_dtype, is_categorical_dtype
+
+from matplotlib.colors import to_hex
+
+
+@d.dedent
+class TermStatesEstimator(BaseEstimator, ABC):
+    """
+    Base class for all estimators predicting terminal states.
+
+    Parameters
+    ----------
+    %(base_estimator.parameters)s
+    """
+
+    def __init__(
+        self,
+        obj: Union[AnnData, np.ndarray, spmatrix, KernelExpression],
+        obsp_key: Optional[str] = None,
+        **kwargs: Any,
+    ):
+        super().__init__(obj=obj, obsp_key=obsp_key, **kwargs)
+
+        self._term_states: Optional[pd.Series] = None
+        self._term_states_probs: Optional[pd.Series] = None
+        self._term_states_colors: Optional[np.ndarray] = None
+
+    @property
+    @d.get_summary(base="tse_term_states")
+    def terminal_states(self) -> Optional[pd.Series]:
+        """Categorical annotation of terminal states.
+
+        By default, all cells in transient cells will be labelled as `NaN`.
+        """
+        return self._term_states
+
+    @property
+    @d.get_summary(base="tse_term_states_probs")
+    def terminal_states_probabilities(self) -> Optional[pd.Series]:
+        """Aggregated probability of cells to be in terminal states."""  # noqa: D401
+        return self._term_states_probs
+
+    @d.dedent
+    def set_terminal_states(
+        self,
+        labels: Union[pd.Series, Dict[str, Sequence[Any]]],
+        cluster_key: Optional[str] = None,
+        add_to_existing: bool = False,
+        **kwargs: Any,
+    ) -> None:
+        """
+        Manually define terminal states.
+
+        Parameters
+        ----------
+        labels
+            Defines the terminal states. Valid options are:
+
+                - categorical :class:`pandas.Series` where each category corresponds to a 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.AnnData.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:`anndata.AnnData.obs`.
+
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` in order to associate names and colors with :attr:`terminal_states`.
+            Each terminal state will be given the name and color corresponding to the cluster it mostly overlaps with.
+        add_to_existing
+            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
+        -------
+        Nothing, just updates the following fields:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+            - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
+        """
+        if add_to_existing:
+            existing = self.terminal_states
+            if existing is None:
+                raise RuntimeError(
+                    "Compute terminal states first as `.compute_terminal_states()` or "
+                    "set them manually as `.set_terminal_states()`."
+                )
+        else:
+            existing = None
+
+        states, colors = self._set_categorical_labels(
+            categories=labels,
+            cluster_key=cluster_key,
+            existing=existing,
+        )
+        self._write_terminal_states(
+            states,
+            colors,
+            **kwargs,
+        )
+
+    @d.get_sections(base="tse_rename_term_states", sections=["Parameters", "Returns"])
+    @d.get_full_description(base="tse_rename_term_states")
+    @d.dedent
+    def rename_terminal_states(self, new_names: Mapping[str, str]) -> None:
+        """
+        Rename categories in :attr:`terminal_states`.
+
+        Parameters
+        ----------
+        new_names
+            Mapping where keys corresponds to the old names and the values to the new names.
+            The new names must be unique.
+
+        Returns
+        -------
+        Nothing, just updates the names of:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+        """
+
+        term_states = self.terminal_states
+        if term_states is None:
+            raise RuntimeError(
+                "Compute terminal states first as `.compute_terminal_states()` or "
+                "set them manually as `.set_terminal_states()`."
+            )
+
+        # fmt: off
+        if not isinstance(new_names, Mapping):
+            raise TypeError(f"Expected new names to be a `Mapping`, found `{type(new_names).__name__}`.")
+        if not len(new_names):
+            return
+
+        old_names = term_states.cat.categories
+        new_names = {str(k): str(v) for k, v in new_names.items()}
+        mask = np.isin(list(new_names.keys()), old_names)
+        if not np.all(mask):
+            invalid = sorted(np.array(list(new_names.keys()))[~mask])
+            raise ValueError(f"Invalid terminal states names: `{invalid}`. Valid names are: `{sorted(old_names)}`")
+
+        names_after_renaming = [new_names.get(n, n) for n in old_names]
+        if len(set(names_after_renaming)) != len(old_names):
+            raise ValueError(f"After renaming, terminal states will no longer unique: `{names_after_renaming}`.")
+        # fmt: on
+
+        # GPCCA; alt. is to subclass
+        self._term_states = term_states.cat.rename_categories(new_names)
+        memberships = getattr(self, "terminal_states_memberships", None)
+        if memberships is not None:
+            memberships.names = [new_names.get(n, n) for n in memberships.names]
+
+        self._write_terminal_states(
+            self.terminal_states,
+            self._term_states_colors,
+            self.terminal_states_probabilities,
+            log=False,
+        )
+
+    def _set_categorical_labels(
+        self,
+        categories: Union[pd.Series, Dict[str, Any]],
+        cluster_key: Optional[str] = None,
+        existing: Optional[pd.Series] = None,
+    ) -> Tuple[pd.Series, np.ndarray]:
+        # fmt: off
+        if isinstance(categories, dict):
+            key = next(iter(categories.keys()))
+            if len(categories) == 1 and is_categorical_dtype(self.adata.obs.get(key, None)):
+                vals = categories[key]
+                if isinstance(vals, str) or not isinstance(vals, Sequence):
+                    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))
+        if not is_categorical_dtype(categories):
+            raise TypeError(f"Expected object to be `categorical`, found `{infer_dtype(categories)}`.")
+
+        if existing is not None:
+            categories = _merge_categorical_series(old=existing, new=categories)
+
+        if cluster_key is not None:
+            # check that we can load the reference series from adata
+            if cluster_key not in self.adata.obs:
+                raise KeyError(f"Unable to find clusters in `adata.obs[{cluster_key!r}]`.")
+            series_query, series_reference = categories, self.adata.obs[cluster_key]
+
+            # load the reference colors if they exist
+            if Key.uns.colors(cluster_key) in self.adata.uns:
+                colors_reference = _convert_to_hex_colors(self.adata.uns[Key.uns.colors(cluster_key)])
+            else:
+                colors_reference = _create_categorical_colors(len(series_reference.cat.categories))
+
+            names, colors = _map_names_and_colors(
+                series_reference=series_reference,
+                series_query=series_query,
+                colors_reference=colors_reference,
+            )
+            categories.cat.categories = names
+        else:
+            colors = _create_categorical_colors(len(categories.cat.categories))
+
+        return categories, colors
+        # fmt: on
+
+    @logger
+    @shadow
+    def _write_terminal_states(
+        self,
+        states: Optional[pd.Series],
+        colors: Optional[np.ndarray],
+        probs: Optional[pd.Series] = None,
+        params: Dict[str, Any] = MappingProxyType({}),
+    ) -> str:
+        # fmt: off
+        key = Key.obs.term_states(self.backward)
+        self._set("_term_states", self.adata.obs, key=key, value=states)
+        self._set("_term_states_probs", self.adata.obs, key=Key.obs.probs(key), value=probs)
+        self._set("_term_states_colors", self.adata.uns, key=Key.uns.colors(key), value=colors)
+        self.params[key] = dict(params)
+        # fmt: on
+
+        return (
+            f"Adding `adata.obs[{key!r}]`\n"
+            f"       `adata.obs[{Key.obs.probs(key)!r}]`\n"
+            f"       `.terminal_states`\n"
+            f"       `.terminal_states_probabilities`\n"
+            "    Finish"
+        )
+
+    def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
+        ok = super()._read_from_adata(adata, **kwargs)
+        if not ok:
+            return False
+
+        # fmt: off
+        key = Key.obs.term_states(self.backward)
+        with SafeGetter(self, allowed=KeyError) as sg:
+            self._get("_term_states", self.adata.obs, key=key, where="obs", dtype=pd.Series)
+            self._get("_term_states_probs", self.adata.obs, key=Key.obs.probs(key), where="obs", dtype=pd.Series)
+            self._get("_term_states_colors", self.adata.uns, key=Key.uns.colors(key), where="uns",
+                      dtype=(list, tuple, np.ndarray))
+            self._term_states_colors = np.asarray([to_hex(c) for c in self._term_states_colors])
+            self.params[key] = self._read_params(key)
+        # fmt: on
+
+        # TODO(michalk8): log
+
+        return sg.ok
+
+    @d.get_full_description(base="tse_fit")
+    def fit(self, *args: Any, **kwargs: Any) -> "TermStatesEstimator":
+        """Prepare self for terminal states prediction."""
+        return self
+
+    @d.dedent
+    def compute_terminal_states(self, *args: Any, **kwargs: Any) -> None:
+        """
+        Compute terminal states of the process.
+
+        This is an alias for :meth:`predict`.
+
+        Parameters
+        ----------
+        args
+            Positional arguments.
+        kwargs
+            Keyword arguments arguments.
+
+        Return
+        ------
+        Nothing, just updates the following fields:
+
+            - :attr:`terminal_states` - %(tse_term_states.summary)s
+            - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
+        """
+        return self.predict(*args, **kwargs)
+
+    plot_terminal_states = register_plotter(
+        discrete="terminal_states", colors="_term_states_colors"
+    )
diff --git a/cellrank/tl/kernels/_base_kernel.py b/cellrank/tl/kernels/_base_kernel.py
index 93c09cc87..9a0889c29 100644
--- a/cellrank/tl/kernels/_base_kernel.py
+++ b/cellrank/tl/kernels/_base_kernel.py
@@ -30,9 +30,8 @@
     _get_neighs,
     _irreducible,
 )
-from cellrank.ul._utils import Pickleable
 from scvelo.plotting.utils import default_size, plot_outline
-from cellrank.tl._constants import Direction, _transition
+from cellrank.tl._mixins._io import IOMixin
 from cellrank.tl.kernels._utils import _get_basis, _filter_kwargs
 from cellrank.tl.kernels._tmat_flow import FlowPlotter
 from cellrank.tl.kernels._random_walk import RandomWalk
@@ -66,7 +65,7 @@
 ]
 
 
-class KernelExpression(Pickleable, ABC):
+class KernelExpression(IOMixin, ABC):
     """Base class for all kernels and kernel expressions."""
 
     def __init__(
@@ -77,7 +76,7 @@ def __init__(
     ):
         self._op_name = op_name
         self._transition_matrix = None
-        self._direction = Direction.BACKWARD if backward else Direction.FORWARD
+        self._backward = backward
         self._compute_cond_num = compute_cond_num
         self._cond_num = None
         self._params = {}
@@ -108,25 +107,23 @@ def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
     @property
     def backward(self) -> bool:
         """Direction of the process."""
-        return self._direction == Direction.BACKWARD
+        return self._backward
 
     @property
     @abstractmethod
-    @d.dedent
     def adata(self) -> AnnData:
-        """
-        Annotated data object.
-
-        Returns
-        -------
-        %(adata_ret)s
-        """
+        """Annotated data object."""
 
     @adata.setter
     @abstractmethod
     def adata(self, value: AnnData) -> None:
         pass
 
+    @property
+    @abstractmethod
+    def shape(self) -> Tuple[int, int]:
+        """`(n_cells, n_cells)`."""
+
     @property
     def params(self) -> Dict[str, Any]:
         """Parameters which are used to compute the transition matrix."""
@@ -196,23 +193,21 @@ def write_to_adata(self, key: Optional[str] = None) -> None:
         Parameters
         ----------
         key
-            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'`.
+            Key used when writing transition matrix to :attr:`adata`. If `None`, determine the key automatically.
 
         Returns
         -------
         None
             %(write_to_adata)s
         """
+        from cellrank._key import Key
 
         if self._transition_matrix is None:
             raise ValueError(
                 "Compute transition matrix first as `.compute_transition_matrix()`."
             )
 
-        if key is None:
-            key = _transition(self._direction)
-
+        key = Key.uns.kernel(self.backward, key=key)
         # retain the embedding info
         self.adata.uns[f"{key}_params"] = {
             **self.adata.uns.get(f"{key}_params", {}),
@@ -227,7 +222,7 @@ def copy(self) -> "KernelExpression":
     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)}`")
+            logg.debug("Computing condition number")
             self._cond_num = np.linalg.cond(
                 self._transition_matrix.toarray()
                 if issparse(self._transition_matrix)
@@ -235,7 +230,7 @@ def _maybe_compute_cond_num(self) -> None:
             )
             if self._cond_num > _cond_num_tolerance:
                 logg.warning(
-                    f"`{repr(self)}` may be ill-conditioned, its condition number is `{self._cond_num:.2e}`"
+                    f"Transition matrix may be ill-conditioned, its condition number is `{self._cond_num:.2e}`"
                 )
             else:
                 logg.info(f"Condition number is `{self._cond_num:.2e}`")
@@ -278,9 +273,9 @@ def compute_projection(
         Parameters
         ----------
         basis
-            Basis in :attr:`adata` ``.obsm`` for which to compute the projection.
+            Basis in :attr:`anndata.AnnData.obsm` for which to compute the projection.
         key_added
-            If not `None` and ``copy=False``, save the result to :attr:`adata` ``.obsm['{key_added}']``.
+            If not `None` and ``copy = False``, save the result to :attr:`anndata.AnnData.obsm` ``['{key_added}']``.
             Otherwise, save the result to `'T_fwd_{basis}'` or `T_bwd_{basis}`, depending on the direction.
         copy
             Whether to return the projection or modify :attr:`adata` inplace.
@@ -291,6 +286,7 @@ def compute_projection(
         Otherwise, it modifies :attr:`anndata.AnnData.obsm` with a key based on ``key_added``.
         """
         # modified from: https://github.com/theislab/scvelo/blob/master/scvelo/tools/velocity_embedding.py
+        from cellrank._key import Key
         from scvelo.tools.velocity_embedding import quiver_autoscale
 
         if self._transition_matrix is None:
@@ -334,7 +330,7 @@ def compute_projection(
         if copy:
             return T_emb
 
-        key = _transition(self._direction) if key_added is None else key_added
+        key = Key.uns.kernel(self.backward, key=key_added)
         ukey = f"{key}_params"
 
         embs = self.adata.uns.get(ukey, {}).get("embeddings", [])
@@ -392,7 +388,7 @@ def plot_random_walks(
         stop_ixs
             Cells which when hit, the random walk is terminated. If `None`, terminate after ``max_iters``.
             %(rw_ixs)s
-            For example ``{'clusters': ['Alpha', 'Beta']}`` and ``successive_hits=3`` means that the random walk will
+            For example ``{'clusters': ['Alpha', 'Beta']}`` and ``successive_hits = 3`` means that the random walk will
             stop prematurely after cells in the above specified clusters have been visited successively 3 times in a
             row.
         basis
@@ -554,11 +550,11 @@ def plot_single_flow(
         Parameters
         ----------
         cluster
-            Cluster for which to visualize outgoing compute_flow.
+            Cluster for which to visualize outgoing flow.
         cluster_key
-            Key in :attr:`adata` ``.obs`` where clustering is stored.
+            Key in :attr:`anndata.AnnData.obs` where clustering is stored.
         time_key
-            Key in :attr:`adata` ``.obs`` where experimental time is stored.
+            Key in :attr:`anndata.AnnData.obs` where experimental time is stored.
         clusters
             Visualize flow only for these clusters. If `None`, use all clusters.
         time_points
@@ -571,7 +567,7 @@ def plot_single_flow(
         Returns
         -------
         :class:`matplotlib.pyplot.Axes`
-            The axis object if ``show=False``.
+            The axis object if ``show = False``.
         %(just_plots)s
 
         Notes
@@ -753,7 +749,7 @@ def __invert__(self: "KernelExpression") -> "KernelExpression":
         # mustn't return a copy because transition matrix
         self._transition_matrix = None
         self._params = {}
-        self._direction = Direction.FORWARD if self.backward else Direction.BACKWARD
+        self._backward = not self.backward
 
         return self
 
@@ -785,6 +781,7 @@ def __init__(
             op_name is None
         ), "Unary kernel does not support any kind operation associated with it."
         self._adata = adata
+        self._n_obs = adata.n_obs
 
     @property
     @d.dedent
@@ -798,19 +795,27 @@ def adata(self) -> AnnData:
         """
         return self._adata
 
+    @property
+    def shape(self) -> Tuple[int, int]:
+        """`(n_cells, n_cells)`."""
+        return self._n_obs, self._n_obs
+
     @adata.setter
-    def adata(self, _adata: AnnData) -> None:
-        if not isinstance(_adata, AnnData):
+    def adata(self, adata: Optional[AnnData]) -> None:
+        if adata is None:
+            self._adata = None
+            return
+        if not isinstance(adata, AnnData):
             raise TypeError(
-                f"Expected argument of type `anndata.AnnData`, found `{type(_adata).__name__!r}`."
+                f"Expected argument of type `anndata.AnnData`, found `{type(adata).__name__!r}`."
             )
-        # otherwise, we'd have to reread bunch of attributes - it's better to initialize new object
-        if _adata.shape != self.adata.shape:
+        shape = (adata.n_obs, adata.n_obs)
+        if self.shape != shape:
             raise ValueError(
-                f"Expected the new object to have same shape as previous object `{self.adata.shape}`, "
-                f"found `{_adata.shape}`."
+                f"Expected the new object to have same shape as previous object `{self.shape}`, "
+                f"found `{shape}`."
             )
-        self._adata = _adata
+        self._adata = adata
 
     def __repr__(self):
         return f"{'~' if self.backward and self._parent is None else ''}<{self.__class__.__name__}>"
@@ -847,6 +852,11 @@ def __init__(self, kexprs: List[KernelExpression], op_name: Optional[str] = None
         for kexprs in self._kexprs:
             kexprs._parent = self
 
+    @property
+    def shape(self) -> Tuple[int, int]:
+        """`(n_cells, n_cells)`."""
+        return self.kernels[0].shape
+
     def _maybe_recalculate_constants(self, type_: Type):
         if type_ == Constant:
             accessor = "transition_matrix"
@@ -854,7 +864,7 @@ def _maybe_recalculate_constants(self, type_: Type):
             accessor = "_value"
         else:
             raise RuntimeError(
-                f"Unable to determine accessor for type `{type.__name__!r}`."
+                f"Unable to determine accessor for type `{type(type_).__name__}`."
             )
 
         constants = [_is_bin_mult(k, type_) for k in self]
@@ -890,8 +900,9 @@ def adata(self) -> AnnData:
         return self._kexprs[0].adata
 
     @adata.setter
-    def adata(self, _adata: AnnData) -> None:
-        self._kexprs[0].adata = _adata
+    def adata(self, adata: Optional[AnnData]) -> None:
+        for kexpr in self._kexprs:
+            kexpr.adata = adata
 
     def __invert__(self) -> "NaryKernelExpression":
         super().__invert__()
@@ -1119,7 +1130,7 @@ def copy(self) -> "Constant":
 
     def __invert__(self) -> "Constant":
         # do not call parent's invert, since it removes the transition matrix
-        self._direction = Direction.FORWARD if self.backward else Direction.BACKWARD
+        self._backward = not self.backward
         return self
 
     def __repr__(self) -> str:
@@ -1176,7 +1187,7 @@ def compute_transition_matrix(self, *args, **kwargs) -> "ConstantMatrix":
 
     def __invert__(self) -> "ConstantMatrix":
         # do not call parent's invert, since it removes the transition matrix
-        self._direction = Direction.FORWARD if self.backward else Direction.BACKWARD
+        self._backward = not self.backward
         return self
 
     def __repr__(self) -> str:
diff --git a/cellrank/tl/kernels/_connectivity_kernel.py b/cellrank/tl/kernels/_connectivity_kernel.py
index 75d80eb7d..766ae26e1 100644
--- a/cellrank/tl/kernels/_connectivity_kernel.py
+++ b/cellrank/tl/kernels/_connectivity_kernel.py
@@ -1,4 +1,3 @@
-"""Connectivity kernel module."""
 from copy import copy
 
 from anndata import AnnData
diff --git a/cellrank/tl/kernels/_cytotrace_kernel.py b/cellrank/tl/kernels/_cytotrace_kernel.py
index ddb395d55..fffb6629b 100644
--- a/cellrank/tl/kernels/_cytotrace_kernel.py
+++ b/cellrank/tl/kernels/_cytotrace_kernel.py
@@ -1,11 +1,14 @@
 from typing import Any
 from typing_extensions import Literal
 
+from enum import auto
+
 from anndata import AnnData
 from cellrank import logging as logg
-from cellrank.ul._docs import d
+from cellrank._key import Key
+from cellrank.tl._enum import ModeEnum
+from cellrank.ul._docs import d, inject_docs
 from cellrank.tl._utils import _correlation_test_helper
-from cellrank.tl._constants import ModeEnum
 from cellrank.tl.kernels._pseudotime_kernel import PseudotimeKernel
 
 import numpy as np
@@ -13,15 +16,11 @@
 from scipy.sparse import issparse
 
 
-def _ct(key: str) -> str:
-    return f"ct_{key}"
-
-
-class CytoTRACEAggregation(ModeEnum):  # noqa
-    MEAN = "mean"
-    MEDIAN = "median"
-    GMEAN = "gmean"
-    HMEAN = "hmean"
+class CytoTRACEAggregation(ModeEnum):  # noqa: D101
+    MEAN = auto()
+    MEDIAN = auto()
+    GMEAN = auto()
+    HMEAN = auto()
 
 
 @d.dedent
@@ -85,7 +84,9 @@ def __init__(
         adata: AnnData,
         backward: bool = False,
         layer: str = "Ms",
-        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        aggregation: Literal[
+            "mean", "median", "hmean", "gmean"
+        ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = False,
         compute_cond_num: bool = False,
         check_connectivity: bool = False,
@@ -94,7 +95,7 @@ def __init__(
         super().__init__(
             adata,
             backward=backward,
-            time_key=_ct("pseudotime"),
+            time_key=Key.cytotrace("pseudotime"),
             compute_cond_num=compute_cond_num,
             check_connectivity=check_connectivity,
             layer=layer,
@@ -102,13 +103,15 @@ def __init__(
             use_raw=use_raw,
             **kwargs,
         )
-        self._time_key = _ct("pseudotime")  # quirk or PT kernel
+        self._time_key = Key.cytotrace("pseudotime")  # quirk or PT kernel
 
     def _read_from_adata(
         self,
         time_key: str,
         layer: str = "Ms",
-        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        aggregation: Literal[
+            "mean", "median", "hmean", "gmean"
+        ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = True,
         **kwargs: Any,
     ) -> None:
@@ -117,10 +120,13 @@ def _read_from_adata(
         super()._read_from_adata(time_key=time_key, **kwargs)
 
     @d.get_sections(base="cytotrace", sections=["Parameters"])
+    @inject_docs(ct=CytoTRACEAggregation)
     def compute_cytotrace(
         self,
         layer: str = "Ms",
-        aggregation: Literal["mean", "median", "hmean", "gmean"] = "mean",
+        aggregation: Literal[
+            "mean", "median", "hmean", "gmean"
+        ] = CytoTRACEAggregation.MEAN,
         use_raw: bool = False,
     ) -> None:
         """
@@ -139,10 +145,10 @@ def compute_cytotrace(
         aggregation
             How to aggregate expression of the top-correlating genes. Valid options are:
 
-                - `'mean'`: arithmetic mean.
-                - `'median'`: median.
-                - `'gmean'`: geometric mean.
-                - `'hmean'`: harmonic mean.
+                - `{ct.MEAN!r}` - arithmetic mean.
+                - `{ct.MEDIAN!r}` - median.
+                - `{ct.HMEAN!r}` - harmonic mean.
+                - `{ct.GMEAN!r}` - geometric mean.
 
         use_raw
             Whether to use the :attr:`anndata.AnnData.raw` to compute the number of genes expressed per cell
@@ -152,14 +158,14 @@ def compute_cytotrace(
         -------
         Nothing, just modifies :attr:`anndata.AnnData.obs` with the following keys:
 
-            - `'ct_score'`: the normalized CytoTRACE score.
-            - `'ct_pseudotime'`: associated pseudotime, essentially `1 - CytoTRACE score`.
-            - `'ct_num_exp_genes'`: the number of genes expressed per cell, basis of the CytoTRACE score.
+            - `'ct_score'` - the normalized CytoTRACE score.
+            - `'ct_pseudotime'` - associated pseudotime, essentially `1 - CytoTRACE score`.
+            - `'ct_num_exp_genes'` - the number of genes expressed per cell, basis of the CytoTRACE score.
 
         It also modifies :attr:`anndata.AnnData.var` with the following keys:
 
-            - `'ct_gene_corr'`: the correlation as specified above.
-            - `'ct_correlates'`: indication of the genes used to compute the CytoTRACE score, i.e. the ones that
+            - `'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
@@ -196,7 +202,7 @@ def compute_cytotrace(
             f"Computing number of genes expressed per cell with `use_raw={use_raw}`"
         )
         num_exp_genes = np.array((adata_mraw.X > 0).sum(axis=1)).reshape(-1)
-        self.adata.obs[_ct("num_exp_genes")] = num_exp_genes
+        self.adata.obs[Key.cytotrace("num_exp_genes")] = num_exp_genes
 
         # fmt: off
         # compute correlation with all genes
@@ -205,14 +211,14 @@ def compute_cytotrace(
 
         # annotate the top 200 genes in terms of correlation
         logg.debug("Finding the top `200` most correlated genes")
-        self.adata.var[_ct("gene_corr")] = gene_corr
-        top_200 = self.adata.var.sort_values(by=_ct("gene_corr"), ascending=False).index[:200]
-        self.adata.var[_ct("correlates")] = False
-        self.adata.var.loc[top_200, _ct("correlates")] = True
+        self.adata.var[Key.cytotrace("gene_corr")] = gene_corr
+        top_200 = self.adata.var.sort_values(by=Key.cytotrace("gene_corr"), ascending=False).index[:200]
+        self.adata.var[Key.cytotrace("correlates")] = False
+        self.adata.var.loc[top_200, Key.cytotrace("correlates")] = True
 
         # compute mean/median over top 200 genes, aggregate over genes and shift to [0, 1] range
         logg.debug(f"Aggregating imputed gene expression using aggregation `{aggregation}` in layer `{layer}`")
-        corr_mask = self.adata.var[_ct("correlates")]
+        corr_mask = self.adata.var[Key.cytotrace("correlates")]
         imputed_exp = self.adata[:, corr_mask].X if layer == "X" else self.adata[:, corr_mask].layers[layer]
         if issparse(imputed_exp):
             imputed_exp = imputed_exp.A
@@ -233,22 +239,22 @@ def compute_cytotrace(
         # scale to 0-1 range
         cytotrace_score -= np.min(cytotrace_score)
         cytotrace_score /= np.max(cytotrace_score)
-        self.adata.obs[_ct("score")] = cytotrace_score
-        self.adata.obs[_ct("pseudotime")] = 1 - cytotrace_score
+        self.adata.obs[Key.cytotrace("score")] = cytotrace_score
+        self.adata.obs[Key.cytotrace("pseudotime")] = 1 - cytotrace_score
 
-        self.adata.uns[_ct("params")] = {
-            "aggregation": aggregation.s,
+        self.adata.uns[Key.cytotrace("params")] = {
+            "aggregation": aggregation,
             "layer": layer,
             "use_raw": use_raw,
         }
 
         logg.info(
-            f"Adding `adata.obs[{_ct('score')!r}]`\n"
-            f"       `adata.obs[{_ct('pseudotime')!r}]`\n"
-            f"       `adata.obs[{_ct('num_exp_genes')!r}]`\n"
-            f"       `adata.var[{_ct('gene_corr')!r}]`\n"
-            f"       `adata.var[{_ct('correlates')!r}]`\n"
-            f"       `adata.uns[{_ct('params')!r}]`\n"
+            f"Adding `adata.obs[{Key.cytotrace('score')!r}]`\n"
+            f"       `adata.obs[{Key.cytotrace('pseudotime')!r}]`\n"
+            f"       `adata.obs[{Key.cytotrace('num_exp_genes')!r}]`\n"
+            f"       `adata.var[{Key.cytotrace('gene_corr')!r}]`\n"
+            f"       `adata.var[{Key.cytotrace('correlates')!r}]`\n"
+            f"       `adata.uns[{Key.cytotrace('params')!r}]`\n"
             f"    Finish",
             time=start,
         )
diff --git a/cellrank/tl/kernels/_exp_time_kernel.py b/cellrank/tl/kernels/_exp_time_kernel.py
index ccc3a27c3..fc0751fdb 100644
--- a/cellrank/tl/kernels/_exp_time_kernel.py
+++ b/cellrank/tl/kernels/_exp_time_kernel.py
@@ -6,10 +6,10 @@
 from types import MappingProxyType
 
 import scanpy as sc
+from cellrank.tl._enum import ModeEnum
 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
 
diff --git a/cellrank/tl/kernels/_precomputed_kernel.py b/cellrank/tl/kernels/_precomputed_kernel.py
index c28e9d7ed..a381c43dd 100644
--- a/cellrank/tl/kernels/_precomputed_kernel.py
+++ b/cellrank/tl/kernels/_precomputed_kernel.py
@@ -1,18 +1,15 @@
-"""Precomputed kernel module."""
-from typing import Any, Union, Optional
+from typing import Any, Dict, Union, Optional
 
 from copy import copy
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank._key import Key
 from cellrank.ul._docs import d
 from cellrank.ul._utils import _read_graph_data
-from cellrank.tl.kernels import Kernel
-from cellrank.tl._constants import Direction, _transition
-from cellrank.tl.kernels._base_kernel import _RTOL, KernelExpression
+from cellrank.tl.kernels._base_kernel import _RTOL, Kernel, KernelExpression
 
 import numpy as np
-from scipy.sparse import eye as speye
 from scipy.sparse import spmatrix, csr_matrix
 
 
@@ -24,7 +21,7 @@ class PrecomputedKernel(Kernel):
     Parameters
     ----------
     transition_matrix
-        Row-normalized transition matrix or a key in :attr:`adata` ``.obsp``
+        Row-normalized transition matrix or a key in :attr:`anndata.AnnData.obsp`.
         or a :class:`cellrank.tl.kernels.KernelExpression` with a precomputed transition matrix.
         If `None`, try to determine the key based on ``backward``.
     %(adata)s
@@ -45,25 +42,21 @@ def __init__(
         compute_cond_num: bool = False,
         **kwargs: Any,
     ):
-        from anndata import AnnData as _AnnData
-
-        self._origin = "'array'"
-        params = {}
+        origin = "'array'"
+        params: Dict[str, Any] = {}
 
         if transition_matrix is None:
-            transition_matrix = _transition(
-                Direction.BACKWARD if backward else Direction.FORWARD
-            )
-            logg.debug(f"Setting transition matrix key to `{transition_matrix!r}`")
+            transition_matrix = Key.uns.kernel(backward)
+            logg.info(f"Accessing `adata.obsp[{transition_matrix!r}]`")
 
         if isinstance(transition_matrix, str):
             if adata is None:
                 raise ValueError(
                     "When `transition_matrix` specifies a key to `adata.obsp`, `adata` cannot be None."
                 )
-            self._origin = f"adata.obsp[{transition_matrix!r}]"
+            origin = f"adata.obsp[{transition_matrix!r}]"
+            backward = Key.uns.kernel(bwd=True) == transition_matrix
             transition_matrix = _read_graph_data(adata, transition_matrix)
-
         elif isinstance(transition_matrix, KernelExpression):
             if transition_matrix._transition_matrix is None:
                 raise ValueError(
@@ -74,8 +67,8 @@ def __init__(
                     "Ignoring supplied `adata` object because it differs from the kernel's `adata` object."
                 )
 
-            # use `str` because it captures the params
-            self._origin = str(transition_matrix).strip("~<>")
+            # use `str` rather than `repr` because it captures the parameters
+            origin = str(transition_matrix).strip("~<>")
             params = transition_matrix.params.copy()
             backward = transition_matrix.backward
             adata = transition_matrix.adata
@@ -84,7 +77,7 @@ def __init__(
         if not isinstance(transition_matrix, (np.ndarray, spmatrix)):
             raise TypeError(
                 f"Expected transition matrix to be of type `numpy.ndarray` or `scipy.sparse.spmatrix`, "
-                f"found `{type(transition_matrix).__name__!r}`."
+                f"found `{type(transition_matrix).__name__}`."
             )
 
         if transition_matrix.shape[0] != transition_matrix.shape[1]:
@@ -93,11 +86,11 @@ def __init__(
             )
 
         if not np.allclose(np.sum(transition_matrix, axis=1), 1.0, rtol=_RTOL):
-            raise ValueError("Not a valid transition matrix, not all rows sum to 1")
+            raise ValueError("Not a valid transition matrix, not all rows sum to 1.")
 
         if adata is None:
             logg.warning("Creating empty `AnnData` object")
-            adata = _AnnData(
+            adata = AnnData(
                 csr_matrix((transition_matrix.shape[0], 1), dtype=np.float32)
             )
 
@@ -105,9 +98,10 @@ def __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()
+        self._params = params
+        self._origin = origin
 
     def _read_from_adata(self, **kwargs: Any) -> None:
         self._conn = None
@@ -127,13 +121,15 @@ def copy(self) -> "PrecomputedKernel":
 
         return pk
 
-    def compute_transition_matrix(self, *args, **kwargs) -> "PrecomputedKernel":
+    def compute_transition_matrix(
+        self, *args: Any, **kwargs: Any
+    ) -> "PrecomputedKernel":
         """Return self."""
         return self
 
     def __invert__(self) -> "PrecomputedKernel":
         # do not call parent's invert, since it removes the transition matrix
-        self._direction = Direction.FORWARD if self.backward else Direction.BACKWARD
+        self._backward = not self.backward
         return self
 
     def __repr__(self):
@@ -144,27 +140,3 @@ def __repr__(self):
 
     def __str__(self):
         return repr(self)
-
-
-@d.dedent
-class DummyKernel(PrecomputedKernel):
-    """
-    Kernel with 1s on the diagonal.
-
-    Parameters
-    ----------
-    %(adata)s
-    %(backward)s
-    """
-
-    def __init__(
-        self,
-        adata: AnnData,
-        backward: bool = False,
-    ):
-        super().__init__(
-            speye(adata.n_obs, format="csr"),
-            adata,
-            backward=backward,
-            compute_cond_num=False,
-        )
diff --git a/cellrank/tl/kernels/_pseudotime_kernel.py b/cellrank/tl/kernels/_pseudotime_kernel.py
index ed904ca61..31d3ee64a 100644
--- a/cellrank/tl/kernels/_pseudotime_kernel.py
+++ b/cellrank/tl/kernels/_pseudotime_kernel.py
@@ -2,13 +2,14 @@
 from typing_extensions import Literal
 
 from copy import copy
+from enum import auto
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum, Backend_t
 from cellrank.ul._docs import d
 from cellrank.tl._utils import _connected
 from cellrank.tl.kernels import Kernel
-from cellrank.tl._constants import ThresholdScheme
 from cellrank.tl.kernels._base_kernel import _dtype
 from cellrank.tl.kernels._pseudotime_schemes import (
     ThresholdSchemeABC,
@@ -20,6 +21,11 @@
 import numpy as np
 
 
+class ThresholdScheme(ModeEnum):  # noqa: D101
+    SOFT = auto()
+    HARD = auto()
+
+
 @d.dedent
 class PseudotimeKernel(Kernel):
     """
@@ -81,7 +87,7 @@ def compute_transition_matrix(
         nu: float = 0.5,
         check_irreducibility: bool = False,
         n_jobs: Optional[int] = None,
-        backend: str = "loky",
+        backend: Backend_t = _DEFAULT_BACKEND,
         show_progress_bar: bool = True,
         **kwargs: Any,
     ) -> "PseudotimeKernel":
@@ -96,20 +102,19 @@ def compute_transition_matrix(
         threshold_scheme
             Which method to use when biasing the graph. Valid options are:
 
-                - `'hard'`: based on *Palantir* :cite:`setty:19` which removes some edges that point against
+                - `'hard'` - based on *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.
-                - `'soft'`: based on *VIA* :cite:`stassen:21` which downweights edges that points against the direction
+                - `'soft'` - based on *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.
-                - :class:`callable`: any function conforming to the signature of
+                - :class:`callable` - any function conforming to the signature of
                   :func:`cellrank.tl.kernels.ThresholdSchemeABC.__call__`.
         frac_to_keep
             The `frac_to_keep` * number of the closest neighbors (according to graph connectivities) are kept, no matter
             whether they lie in the pseudotemporal past or future. This is done to ensure that the graph remains
-            connected. Only used when `threshold_scheme='hard'`. `frac_to_keep` needs to fall within the
-            interval `[0, 1]`.
+            connected. Only used when ``threshold_scheme = 'hard'``. Needs to fall within the interval `[0, 1]`.
         %(soft_scheme_kernel)s
         check_irreducibility
             Optional check for irreducibility of the final transition matrix.
diff --git a/cellrank/tl/kernels/_random_walk.py b/cellrank/tl/kernels/_random_walk.py
index 39b14e0ca..865977436 100644
--- a/cellrank/tl/kernels/_random_walk.py
+++ b/cellrank/tl/kernels/_random_walk.py
@@ -63,11 +63,8 @@ def _sample(self, ix: int, *, rs: np.random.RandomState) -> int:
         )
 
     def _max_iter(self, max_iter: Union[int, float]) -> int:
-        max_iter = (
-            int(np.ceil(max_iter * len(self._ixs)))
-            if isinstance(max_iter, float)
-            else max_iter
-        )
+        if isinstance(max_iter, float):
+            max_iter = int(np.ceil(max_iter * len(self._ixs)))
         if max_iter <= 1:
             raise ValueError(
                 f"Expected number of iterations to be > 1, found `{max_iter}`."
diff --git a/cellrank/tl/kernels/_utils.py b/cellrank/tl/kernels/_utils.py
index 70bab41e9..6a0dd10b8 100644
--- a/cellrank/tl/kernels/_utils.py
+++ b/cellrank/tl/kernels/_utils.py
@@ -1,4 +1,3 @@
-"""Utility functions for kernels, mostly VelocityKernel."""
 from typing import Tuple, Callable, Optional
 
 from inspect import signature
diff --git a/cellrank/tl/kernels/_velocity_kernel.py b/cellrank/tl/kernels/_velocity_kernel.py
index 0182a50c4..e93182f44 100644
--- a/cellrank/tl/kernels/_velocity_kernel.py
+++ b/cellrank/tl/kernels/_velocity_kernel.py
@@ -1,15 +1,16 @@
-"""Velocity kernel module."""
 from typing import Any, Union, Callable, Iterable, Optional
+from typing_extensions import Literal
 
 from copy import copy
+from enum import auto
 from math import fsum
 
 from anndata import AnnData
 from cellrank import logging as logg
+from cellrank.tl._enum import _DEFAULT_BACKEND, ModeEnum
 from cellrank.ul._docs import d, inject_docs
 from cellrank.ul._utils import valuedispatch
 from cellrank.tl.kernels import Kernel
-from cellrank.tl._constants import _DEFAULT_BACKEND, ModeEnum
 from cellrank.ul._parallelize import parallelize
 from cellrank.tl.kernels._utils import (
     prange,
@@ -28,16 +29,16 @@
 from scipy.sparse import issparse, csr_matrix
 
 
-class VelocityMode(ModeEnum):  # noqa
-    DETERMINISTIC = "deterministic"
-    STOCHASTIC = "stochastic"
-    SAMPLING = "sampling"
-    MONTE_CARLO = "monte_carlo"
+class VelocityMode(ModeEnum):  # noqa: D101
+    DETERMINISTIC = auto()
+    STOCHASTIC = auto()
+    SAMPLING = auto()
+    MONTE_CARLO = auto()
 
 
-class BackwardMode(ModeEnum):  # noqa
-    TRANSPOSE = "transpose"
-    NEGATE = "negate"
+class BackwardMode(ModeEnum):  # noqa: D101
+    TRANSPOSE = auto()
+    NEGATE = auto()
 
 
 @d.dedent
@@ -56,12 +57,12 @@ class VelocityKernel(Kernel):
     %(adata)s
     %(backward)s
     vkey
-        Key in :attr:`adata` ``.layers`` where velocities are stored.
+        Key in :attr:`anndata.AnnData.layers` where velocities are stored.
     xkey
-        Key in :attr:`adata` ``.layers`` where expected gene expression counts are stored.
+        Key in :attr:`anndata.AnnData.layers` where expected gene expression counts are stored.
     gene_subset
         List of genes to be used to compute transition probabilities.
-        By default, genes from :attr:`adata` ``.var['velocity_genes']`` are used.
+        By default, genes from :attr:`anndata.AnnData.var` ``['velocity_genes']`` are used.
     %(cond_num)s
     check_connectivity
         Check whether the underlying KNN graph is connected.
@@ -101,7 +102,7 @@ def _read_from_adata(self, **kwargs: Any) -> None:
         # check whether velocities have been computed
         vkey = kwargs.pop("vkey", "velocity")
         if vkey not in self.adata.layers.keys():
-            raise KeyError("Compute RNA velocity first as `scv.tl.velocity()`.")
+            raise KeyError("Compute RNA velocity first as `scvelo.tl.velocity()`.")
 
         # restrict genes to a subset, i.e. velocity genes or user provided list
         gene_subset = kwargs.pop("gene_subset", None)
@@ -151,14 +152,18 @@ def _read_from_adata(self, **kwargs: Any) -> None:
     @d.dedent
     def compute_transition_matrix(
         self,
-        mode: str = VelocityMode.DETERMINISTIC.s,
-        backward_mode: str = BackwardMode.TRANSPOSE.s,
-        scheme: Union[str, Callable] = Scheme.CORRELATION.s,
+        mode: Literal[
+            "deterministic", "stochastic", "sampling", "monte_carlo"
+        ] = VelocityMode.DETERMINISTIC,
+        backward_mode: Literal["transpose", "negate"] = BackwardMode.TRANSPOSE,
+        scheme: Union[
+            Literal["dot_product", "cosine", "correlation"], Callable
+        ] = Scheme.CORRELATION,
         softmax_scale: Optional[float] = None,
         n_samples: int = 1000,
         seed: Optional[int] = None,
         check_irreducibility: bool = False,
-        **kwargs,
+        **kwargs: Any,
     ) -> "VelocityKernel":
         """
         Compute transition matrix based on velocity directions on the local manifold.
@@ -173,7 +178,7 @@ def compute_transition_matrix(
         %(softmax_scale)s
         %(velocity_scheme)s
         n_samples
-            Number of bootstrap samples when ``mode={m.MONTE_CARLO.s!r}``.
+            Number of bootstrap samples when ``mode = {m.MONTE_CARLO!r}``.
         seed
             Set the seed for random state when the method requires ``n_samples``.
         check_irreducibility
@@ -201,8 +206,8 @@ def compute_transition_matrix(
 
         if self.backward and mode != VelocityMode.DETERMINISTIC:
             logg.warning(
-                f"Mode `{mode.s!r}` is currently not supported for the backward process. "
-                f"Defaulting to mode `{VelocityMode.DETERMINISTIC.s!r}`"
+                f"Mode `{mode!r}` is currently not supported for the backward process. "
+                f"Defaulting to mode `{VelocityMode.DETERMINISTIC!r}`"
             )
             mode = VelocityMode.DETERMINISTIC
 
@@ -216,21 +221,21 @@ def compute_transition_matrix(
                 logg.warning(
                     f"Unable to detect a method for Hessian computation. If using predefined functions, consider "
                     f"installing `jax` as `pip install jax jaxlib`\n"
-                    f"Defaulting to `mode={VelocityMode.MONTE_CARLO.s!r}` and `n_samples={n_samples}`"
+                    f"Defaulting to `mode={VelocityMode.MONTE_CARLO!r}` and `n_samples={n_samples}`"
                 )
                 mode = VelocityMode.MONTE_CARLO
 
         start = logg.info(
-            f"Computing transition matrix based on logits using `{mode.s!r}` mode"
+            f"Computing transition matrix based on logits using `{mode!r}` mode"
         )
 
         if seed is None:
             seed = np.random.randint(0, 2 ** 16)
 
         # fmt: off
-        params = {"softmax_scale": softmax_scale, "mode": mode.s, "seed": seed, "scheme": str(scheme)}
+        params = {"softmax_scale": softmax_scale, "mode": mode, "seed": seed, "scheme": str(scheme)}
         if self.backward:
-            params["bwd_mode"] = backward_mode.s
+            params["bwd_mode"] = backward_mode
         if self._reuse_cache(params, time=start):
             return self
 
@@ -248,14 +253,14 @@ def compute_transition_matrix(
         if mode != VelocityMode.STOCHASTIC and backend == "multiprocessing":
             # this is because on jitting and pickling (cloudpickle, used by loky, handles it correctly)
             logg.warning(
-                f"Multiprocessing backend is supported only for mode `{VelocityMode.STOCHASTIC.s!r}`. "
-                f"Defaulting to `{_DEFAULT_BACKEND}`"
+                f"Multiprocessing backend is supported only for mode `{VelocityMode.STOCHASTIC!r}`. "
+                f"Defaulting to `{_DEFAULT_BACKEND!r}`"
             )
             backend = _DEFAULT_BACKEND
 
         if softmax_scale is None:
             logg.info(
-                f"Estimating `softmax_scale` using `{VelocityMode.DETERMINISTIC.s!r}` mode"
+                f"Estimating `softmax_scale` using `{VelocityMode.DETERMINISTIC!r}` mode"
             )
             _, cmat = _dispatch_computation(
                 VelocityMode.DETERMINISTIC,
diff --git a/cellrank/tl/kernels/_velocity_schemes.py b/cellrank/tl/kernels/_velocity_schemes.py
index 78ac7c550..597c04760 100644
--- a/cellrank/tl/kernels/_velocity_schemes.py
+++ b/cellrank/tl/kernels/_velocity_schemes.py
@@ -1,11 +1,12 @@
 from typing import Any, Tuple
 
 from abc import ABC, abstractmethod
+from enum import auto
 from functools import partial
 
+from cellrank.tl._enum import ModeEnum
 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
@@ -13,9 +14,9 @@
 
 
 class Scheme(ModeEnum):  # noqa: D101
-    DOT_PRODUCT = "dot_product"
-    COSINE = "cosine"
-    CORRELATION = "correlation"
+    DOT_PRODUCT = auto()
+    COSINE = auto()
+    CORRELATION = auto()
 
 
 try:
diff --git a/cellrank/ul/_docs.py b/cellrank/ul/_docs.py
index 6cd7e9c33..100eca49a 100644
--- a/cellrank/ul/_docs.py
+++ b/cellrank/ul/_docs.py
@@ -1,4 +1,4 @@
-"""Module for documentation helper function."""
+from typing import Any
 
 from docrep import DocstringProcessor
 from textwrap import dedent
@@ -36,7 +36,7 @@
     Nothing, just plots the figure. Optionally saves it based on ``save``."""
 _plots_or_returns_models = """\
 None
-    If ``return_models=False``, just plots the figure and optionally saves it based on ``save``.
+    If ``return_models = False``, just plots the figure and optionally saves it based on ``save``.
 Dict[str, Dict[str, :class:`cellrank.ul.models.BaseModel`]]
     Otherwise returns the fitted models as ``{'gene_1': {'lineage_1': <model_11>, ...}, ...}``.
     Models which have failed will be instances of :class:`cellrank.ul.models.FailedModel`."""
@@ -45,9 +45,12 @@
     Direction of the process."""
 _eigen = """\
 which
-    Eigenvalues are in general complex. `'LR'` - largest real part, `'LM'` - largest magnitude.
+    How to sort the eigenvalues. Valid option are:
+
+        - `'LR'` - the largest real part.
+        - `'LM'` - the largest magnitude.
 alpha
-    Used to compute the `eigengap`. ``alpha`` is the weight given to the deviation of an eigenvalue from one."""
+    Used to compute the *eigengap*. ``alpha`` is the weight given to the deviation of an eigenvalue from one."""
 _n_cells = """\
 n_cells
     Number of most likely cells from each macrostate to select."""
@@ -78,22 +81,22 @@
 mode
     How to compute transition probabilities. Valid options are:
 
-        - `{m.DETERMINISTIC.s!r}` - deterministic computation that doesn't propagate uncertainty.
-        - `{m.MONTE_CARLO.s!r}` - Monte Carlo average of randomly sampled velocity vectors.
-        - `{m.STOCHASTIC.s!r}` - second order approximation, only available when :mod:`jax` is installed.
-        - `{m.SAMPLING.s!r}` - sample 1 transition matrix from the velocity distribution."""
+        - `{m.DETERMINISTIC!r}` - deterministic computation that doesn't propagate uncertainty.
+        - `{m.MONTE_CARLO!r}` - Monte Carlo average of randomly sampled velocity vectors.
+        - `{m.STOCHASTIC!r}` - second order approximation, only available when :mod:`jax` is installed.
+        - `{m.SAMPLING!r}` - sample 1 transition matrix from the velocity distribution."""
 _velocity_backward_mode = """\
 backward_mode
-    Only matters if initialized as :attr:`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."""
+        - `{b.TRANSPOSE!r}` - compute transitions from neighboring cells :math:`j` to cell :math:`i`.
+        - `{b.NEGATE!r}` - negate the velocity vector."""
 _velocity_backward_mode_high_lvl = """\
 backward_mode
     How to compute the backward transitions. Valid options are:
 
-        - `{b.TRANSPOSE.s!r}` - compute transitions from neighboring cells `j` to cell `i`.
-        - `{b.NEGATE.s!r}` - negate the velocity vector."""
+        - `{b.TRANSPOSE!r}` - compute transitions from neighboring cells :math:`j` to cell :math:`i`.
+        - `{b.NEGATE!r}` - negate the velocity vector."""
 _copy = """Return a copy of self."""
 _initial = "initial"
 _terminal = "terminal"
@@ -104,17 +107,17 @@
     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``."""
+    Genes in :attr:`anndata.AnnData.var_names` or in :attr:`anndata.AnnData.raw.var_names`, if ``use_raw = True``."""
 _softmax_scale = """\
 softmax_scale
     Scaling parameter for the softmax. If `None`, it will be estimated using ``1 / median(correlations)``.
-    The idea behind this is to scale the softmax to counteract everything tending to orthogonality in high dimensions."""  # noqa
+    The idea behind this is to scale the softmax to counter everything tending to orthogonality in high dimensions."""
 _time_mode = """\
 mode
-    Can be either `'embedding'` or `'time'`:
+    Valid options are:
 
-        - `'embedding'` - plot the embedding while coloring in the absorption probabilities.
-        - `'time'` - plot the pseudotime on x-axis and the absorption probabilities on y-axis."""
+        - `'embedding'` - plot the embedding while coloring in continuous or categorical observations.
+        - `'time'` - plot the pseudotime on x-axis and the probabilities/memberships on y-axis."""
 _write_to_adata = """\
 Updates the :attr:`adata` with the following fields:
 
@@ -132,14 +135,14 @@
     If `True`, return the fitted models for each gene in ``genes`` and lineage in ``lineages``."""
 _basis = """\
 basis
-    Basis to use when ``mode='embedding'``. If `None`, use `'umap'`."""
+    Basis to use when ``mode = 'embedding'``. If `None`, use `'umap'`."""
 _velocity_scheme = """\
 scheme
     Similarity scheme between cells as described in :cite:`li:20`. Can be one of the following:
 
-        - `{s.DOT_PRODUCT.s!r}`: :class:`cellrank.tl.kernels.DotProductScheme`.
-        - `{s.COSINE.s!r}`: :class:`cellrank.tl.kernels.CosineScheme`.
-        - `{s.CORRELATION.s!r}`: :class:`cellrank.tl.kernels.CorrelationScheme`.
+        - `{s.DOT_PRODUCT!r}` - :class:`cellrank.tl.kernels.DotProductScheme`.
+        - `{s.COSINE!r}` - :class:`cellrank.tl.kernels.CosineScheme`.
+        - `{s.CORRELATION!r}` - :class:`cellrank.tl.kernels.CorrelationScheme`.
 
     Alternatively, any function can be passed as long as it follows the signature of
     :meth:`cellrank.tl.kernels.SimilaritySchemeABC.__call__`."""
@@ -155,17 +158,17 @@
 _rw_ixs = """\
 Can be specified as:
 
-        - :class:`dict`: dictionary with 1 key in :attr:`anndata.AnnData.obs` with values corresponding
+        - :class:`dict` - dictionary with 1 key in :attr:`anndata.AnnData.obs` with values corresponding
           to either 1 or more clusters (if the column is categorical) or a :class:`tuple` specifying
           `[min, max]` interval from which to select the indices.
-        - :class:`typing.Sequence`: sequence of cell ids in :attr:`anndata.AnnData.obs_names`.
+        - :class:`typing.Sequence` - sequence of cell ids in :attr:`anndata.AnnData.obs_names`.
 """
 _gene_symbols = """\
 gene_symbols
     Key in :attr:`anndata.AnnData.var` to use instead of :attr:`anndata.AnnData.var_names`."""
 
 
-def inject_docs(**kwargs):  # noqa
+def inject_docs(**kwargs: Any):  # noqa
     def decorator(obj):
         obj.__doc__ = dedent(obj.__doc__).format(**kwargs)
         return obj
@@ -214,7 +217,7 @@ def decorator2(obj):
     cond_num=_cond_num,
     soft_scheme=_soft_scheme_fmt.format("", "", ""),
     soft_scheme_kernel=_soft_scheme_fmt.format(
-        *([" Only used when `threshold_scheme='soft'`."] * 3)
+        *([" Only used when ``threshold_scheme = 'soft'``."] * 3)
     ),
     rw_ixs=_rw_ixs,
     gene_symbols=_gene_symbols,
diff --git a/cellrank/ul/_utils.py b/cellrank/ul/_utils.py
index c8de7c5ac..9c5cf6559 100644
--- a/cellrank/ul/_utils.py
+++ b/cellrank/ul/_utils.py
@@ -1,10 +1,7 @@
-"""General utility functions module."""
 from typing import Any, Dict, List, Tuple, Union, Callable, Iterable, Optional
 
 import wrapt
-import pickle
 from types import MappingProxyType
-from pathlib import Path
 from functools import wraps, update_wrapper
 from contextlib import contextmanager
 from multiprocessing import cpu_count
@@ -18,60 +15,6 @@
 from scipy.sparse import issparse, spmatrix
 
 
-class Pickleable:
-    """Class which allows serialization and deserialization using :mod:pickle."""
-
-    @d.get_full_description(base="pickleable")
-    @d.get_sections(base="pickleable", sections=["Parameters", "Returns"])
-    def write(self, fname: Union[str, Path], ext: Optional[str] = "pickle") -> None:
-        """
-        Serialize self to a file.
-
-        Parameters
-        ----------
-        fname
-            Filename where to save the object.
-        ext
-            Filename extension to use. If `None`, don't append any extension.
-
-        Returns
-        -------
-        None
-            Nothing, just writes itself to a file using :mod:`pickle`.
-        """
-
-        fname = str(fname)
-        if ext is not None:
-            if not ext.startswith("."):
-                ext = "." + ext
-            if not fname.endswith(ext):
-                fname += ext
-
-        logg.info(f"Writing `{self}` to `{fname}`")
-
-        with open(fname, "wb") as fout:
-            pickle.dump(self, fout)
-
-    @staticmethod
-    def read(fname: Union[str, Path]) -> Any:
-        """
-        Deserialize self from a file.
-
-        Parameters
-        ----------
-        fname
-            Filename from which to read the object.
-
-        Returns
-        -------
-        :class:`typing.Any`
-            The deserialized object.
-        """
-
-        with open(fname, "rb") as fin:
-            return pickle.load(fin)
-
-
 def _check_collection(
     adata: AnnData,
     needles: Iterable[str],
@@ -225,16 +168,14 @@ def _read_graph_data(adata: AnnData, key: str) -> Union[np.ndarray, spmatrix]:
     adata
         Annotated data object.
     key
-        Key in ``adata.obsp``.
+        Key in :attr:`anndata.AnnData.obsp`.
 
     Returns
     -------
     :class:`numpy.ndarray` or :class:`scipy.sparse.spmatrix`
         The graph data.
     """
-
-    logg.debug(f"Reading key `{key!r}` from `adata.obsp`")
-    if key in adata.obsp.keys():
+    if key in adata.obsp:
         return adata.obsp[key]
 
     raise KeyError(f"Unable to find data in `adata.obsp[{key!r}]`.")
diff --git a/cellrank/ul/models/_base_model.py b/cellrank/ul/models/_base_model.py
index 9da90da99..a63840739 100644
--- a/cellrank/ul/models/_base_model.py
+++ b/cellrank/ul/models/_base_model.py
@@ -1,16 +1,4 @@
-"""Base class for all models."""
-from typing import (
-    Any,
-    Dict,
-    List,
-    Tuple,
-    Union,
-    Mapping,
-    TypeVar,
-    Callable,
-    Optional,
-    Sequence,
-)
+from typing import Any, Dict, List, Tuple, Union, Mapping, Callable, Optional, Sequence
 
 import re
 import wrapt
@@ -18,15 +6,18 @@
 from abc import ABC, ABCMeta, abstractmethod
 from copy import copy as _copy
 from copy import deepcopy
+from enum import auto
 from collections import defaultdict
 
+from anndata import AnnData
 from cellrank import logging as logg
-from cellrank.tl import Lineage
+from cellrank.tl._enum import ModeEnum
 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 cellrank.ul._utils import _minmax, valuedispatch, _densify_squeeze
+from cellrank.tl._lineage import Lineage
 from scanpy.plotting._utils import add_colors_for_categorical_sample_annotation
+from cellrank.tl._mixins._io import IOMixin
 
 import numpy as np
 from scipy.sparse import spmatrix
@@ -41,7 +32,6 @@
 from matplotlib.colors import to_rgb, is_color_like
 from mpl_toolkits.axes_grid1 import make_axes_locatable
 
-AnnData = TypeVar("AnnData")
 _dup_spaces = re.compile(r" +")  # used on repr for underlying model's repr
 ArrayLike = Union[np.ndarray, spmatrix, List, Tuple]
 
@@ -50,19 +40,19 @@ class UnknownModelError(RuntimeError):  # noqa
     pass
 
 
-class FailedReturnType(ModeEnum):  # noqa
-    PREPARE = "prepare"
-    FIT = "fit"
-    PREDICT = "predict"
-    CONFIDENCE_INTERVAL = "confidence_interval"
-    DEFAULT_CONFIDENCE_INTERVAL = "default_confidence_interval"
-    PLOT = "plot"
+class FailedReturnType(ModeEnum):  # noqa: D101
+    PREPARE = auto()
+    FIT = auto()
+    PREDICT = auto()
+    CONFIDENCE_INTERVAL = auto()
+    DEFAULT_CONFIDENCE_INTERVAL = auto()
+    PLOT = auto()
 
 
 class ColorType(ModeEnum):  # noqa: D101
-    CONT = "cont"
-    CAT = "cat"
-    STR = "str"
+    CONT = auto()
+    CAT = auto()
+    STR = auto()
 
 
 def _handle_exception(return_type: FailedReturnType, func: Callable) -> Callable:
@@ -142,7 +132,9 @@ def _(func: Callable):
 class BaseModelMeta(ABCMeta):
     """Metaclass for all base models."""
 
-    def __new__(cls, clsname, superclasses, attributedict):
+    def __new__(
+        cls, clsname: str, superclasses: Tuple[type, ...], attributedict: Dict[str, Any]
+    ):
         """
         Create a new instance.
 
@@ -160,8 +152,8 @@ def __new__(cls, clsname, superclasses, attributedict):
         for fun_name in list(FailedReturnType):
             setattr(
                 obj,
-                fun_name.s,
-                _handle_exception(FailedReturnType(fun_name), getattr(obj, fun_name.s)),
+                fun_name,
+                _handle_exception(FailedReturnType(fun_name), getattr(obj, fun_name)),
             )
 
         return obj
@@ -169,7 +161,7 @@ def __new__(cls, clsname, superclasses, attributedict):
 
 @d.get_sections(base="base_model", sections=["Parameters"])
 @d.dedent
-class BaseModel(Pickleable, ABC, metaclass=BaseModelMeta):
+class BaseModel(IOMixin, ABC, metaclass=BaseModelMeta):
     """
     Base class for all model classes.
 
@@ -182,18 +174,18 @@ class BaseModel(Pickleable, ABC, metaclass=BaseModelMeta):
 
     def __init__(
         self,
-        adata: AnnData,
+        adata: Optional[AnnData],
         model: Any,
     ):
-        from anndata import AnnData as _AnnData
-
-        if not isinstance(adata, _AnnData) and not isinstance(self, FittedModel):
+        if not isinstance(adata, AnnData) and not isinstance(self, FittedModel):
             # FittedModel doesn't need it
             raise TypeError(
-                f"Expected `adata` to be of type `anndata.AnnData`, found `{type(adata).__name__!r}`."
+                f"Expected `adata` to be of type `anndata.AnnData`, found `{type(adata).__name__}`."
             )
+        super().__init__()
 
         self._adata = adata
+        self._n_obs = 0 if adata is None else adata.n_obs
         self._model = model
         self._gene = None
         self._use_raw = False
@@ -239,9 +231,18 @@ def adata(self) -> AnnData:
         """
         return self._adata
 
+    @adata.setter
+    def adata(self, adata: Optional[AnnData]) -> None:
+        self._adata = adata
+
+    @property
+    def shape(self) -> Tuple[int]:
+        """Number of cells in :attr:`adata`."""  # noqa: D401
+        return (self._n_obs,)
+
     @property
     def model(self) -> Any:
-        """The underlying model."""  # noqa
+        """Underlying model."""
         return self._model
 
     @property
@@ -319,7 +320,7 @@ def prepare(
         lineage: Optional[str],
         backward: bool = False,
         time_range: Optional[Union[float, Tuple[float, float]]] = None,
-        data_key: str = "X",
+        data_key: Optional[str] = "X",
         time_key: str = "latent_time",
         use_raw: bool = False,
         threshold: Optional[float] = None,
@@ -333,18 +334,19 @@ def prepare(
         Parameters
         ----------
         gene
-            Gene in :attr:`adata` ``.var_names`` or in :attr:`adata` ``.raw.var_names``.
+            Gene in :attr:`anndata.AnnData.var_names`.
         lineage
-            Name of a lineage in :attr:`adata` ``.obsm['{lineage_key}']``. If `None`, all weights will be set to `1`.
+            Name of a lineage in :attr:`anndata.AnnData.obsm` ``['{lineage_key}']``.
+            If `None`, all weights will be set to `1`.
         %(backward)s
         %(time_range)s
         data_key
-            Key in :attr:`adata` ``.layers`` or `'X'` for :attr:`adata` ``.X``.
-            If ``use_raw=True``, it's always set to `'X'`.
+            Key in :attr:`anndata.AnnData.layers` or `'X'` for :attr:`anndata.AnnData.X`.
+            If ``use_raw = True``, it's always set to `'X'`.
         time_key
-            Key in :attr:`adata` ``.obs`` where the pseudotime is stored.
+            Key in :attr:`anndata.AnnData.obs` where the pseudotime is stored.
         use_raw
-            Whether to access :attr:`adata` ``.raw`` or not.
+            Whether to access :attr:`anndata.AnnData.raw`.
         threshold
             Consider only cells with weights > ``threshold`` when estimating the test endpoint.
             If `None`, use the median of the weights.
@@ -358,36 +360,36 @@ def prepare(
 
         Returns
         -------
-        None
-            Nothing, but updates the following fields:
+        Nothing, but updates the following fields:
 
-                - :attr:`x` - %(base_model_x.summary)s
-                - :attr:`y` - %(base_model_y.summary)s
-                - :attr:`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
 
-                - :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
+            - :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
 
-                - :attr:`x_test` - %(base_model_x_test.summary)s
+            - :attr:`x_test` - %(base_model_x_test.summary)s
 
-                - :attr:`prepared` - %(base_model_prepared.summary)s
+            - :attr:`prepared` - %(base_model_prepared.summary)s
         """
+        from cellrank._key import Key
 
-        self._use_raw = use_raw
         if use_raw:
             if self.adata.raw is None:
                 raise AttributeError("AnnData object has no attribute `.raw`.")
             if data_key != "X":
                 data_key = "X"
+        self._use_raw = use_raw
 
-        if data_key not in ["X", "obs"] + list(self.adata.layers.keys()):
+        if data_key not in ["X", "obs", None] + list(self.adata.layers.keys()):
             raise KeyError(
                 f"Data key must be a key of `adata.layers`: `{list(self.adata.layers.keys())}`, "
                 f"`adata.X` or `adata.obs`."
             )
 
-        lineage_key = str(AbsProbKey.BACKWARD if backward else AbsProbKey.FORWARD)
+        lineage_key = Key.obsm.abs_probs(backward)
         if lineage_key not in self.adata.obsm:
             raise KeyError(f"Lineage key `{lineage_key!r}` not found in `adata.obsm`.")
         if not isinstance(self.adata.obsm[lineage_key], Lineage):
@@ -404,7 +406,7 @@ def prepare(
 
         if data_key == "obs":
             if gene not in self.adata.obs:
-                raise KeyError(f"Unable to find key `{gene!r}` in `adata.obs`.")
+                raise KeyError(f"Unable to find data in `adata.obs[{gene!r}]`.")
         else:
             if use_raw and gene not in self.adata.raw.var_names:
                 raise KeyError(f"Gene `{gene!r}` not found in `adata.raw.var_names`.")
@@ -451,7 +453,7 @@ def prepare(
         adata = self.adata.raw.to_adata() if use_raw else self.adata
         gene_ix = np.where(adata.var_names == gene)[0]
 
-        if data_key == "X":
+        if data_key in ("X", None):
             y = adata.X[:, gene_ix]
             self._data_key = None
         elif data_key == "obs":
diff --git a/cellrank/ul/models/_gamr_model.py b/cellrank/ul/models/_gamr_model.py
index 393ad2677..a3d2dda64 100644
--- a/cellrank/ul/models/_gamr_model.py
+++ b/cellrank/ul/models/_gamr_model.py
@@ -1,11 +1,13 @@
 """Module containing all models interfacing R's mgcv package."""
 from typing import Any, Tuple, Union, Optional
+from typing_extensions import Literal
 
 from copy import copy, deepcopy
+from enum import auto
 
+from cellrank.tl._enum import ModeEnum
 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
 
@@ -17,9 +19,9 @@
 _r_lib_name = None
 
 
-class KnotLocs(ModeEnum):  # noqa
-    AUTO = "auto"
-    DENSITY = "density"
+class KnotLocs(ModeEnum):  # noqa: D101
+    AUTO = auto()
+    DENSITY = auto()
 
 
 @inject_docs(key=_OFFSET_KEY, kloc=KnotLocs)
@@ -43,8 +45,8 @@ class GAMR(BaseModel):
     knotlocs
         Position of the knots. Can be one of the following:
 
-            - `{kloc.AUTO.s!r}` - let `mgcv` handle the knot positions.
-            - `{kloc.DENSITY.s!r}` - position the knots based on the density of the pseudotime.
+            - `{kloc.AUTO!r}` - let `mgcv` handle the knot positions.
+            - `{kloc.DENSITY!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 :cite:`robinson:10`. The values are saved in :attr:`adata` ``.obs[{key!r}]``.
@@ -62,8 +64,8 @@ def __init__(
         n_knots: int = 5,
         distribution: str = "gaussian",
         basis: str = "cr",
-        knotlocs: str = KnotLocs.AUTO.s,
-        offset: Optional[Union[np.ndarray, str]] = "default",
+        knotlocs: Literal["auto", "density"] = KnotLocs.AUTO,
+        offset: Optional[Union[np.ndarray, Literal["default"]]] = "default",
         smoothing_penalty: float = 1.0,
         **kwargs,
     ):
@@ -334,7 +336,7 @@ def copy(self) -> "GAMR":
             self._n_knots,
             distribution=self._family,
             offset=self._offset,
-            knotlocs=self._knotslocs.s,
+            knotlocs=self._knotslocs,
             perform_import_check=False,
         )
         self._shallowcopy_attributes(
diff --git a/cellrank/ul/models/_pygam_model.py b/cellrank/ul/models/_pygam_model.py
index 4adc81560..2cef19c80 100644
--- a/cellrank/ul/models/_pygam_model.py
+++ b/cellrank/ul/models/_pygam_model.py
@@ -1,15 +1,16 @@
-"""Module containing :mod:`pygam` model implementation."""
-from typing import Union, Mapping, Optional
+from typing import Any, Union, Mapping, Optional
+from typing_extensions import Literal
 
 from copy import copy as _copy
 from copy import deepcopy
+from enum import auto
 from types import MappingProxyType
 from collections import defaultdict
 
 from cellrank import logging as logg
+from cellrank.tl._enum import ModeEnum
 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
 
@@ -25,25 +26,22 @@
     s,
 )
 
-_r_lib = None
-_r_lib_name = None
 
-
-class GamLinkFunction(ModeEnum):  # noqa
-    IDENTITY = "identity"
-    LOGIT = "logit"
-    INV = "inverse"
-    LOG = "log"
+class GamLinkFunction(ModeEnum):  # noqa: D101
+    IDENTITY = auto()
+    LOGIT = auto()
+    INVERSE = auto()
+    LOG = auto()
     INV_SQUARED = "inverse-squared"
 
 
-class GamDistribution(ModeEnum):  # noqa
-    NORMAL = "normal"
-    BINOMIAL = "binomial"
-    POISSON = "poisson"
-    GAMMA = "gamma"
-    GAUSS = "gaussian"
-    INV_GAUSS = "inv_gauss"
+class GamDistribution(ModeEnum):  # noqa: D101
+    NORMAL = auto()
+    BINOMIAL = auto()
+    POISSON = auto()
+    GAMMA = auto()
+    GAUSSIAN = auto()
+    INV_GAUSS = auto()
 
 
 _gams = defaultdict(
@@ -95,13 +93,15 @@ def __init__(
         adata: AnnData,
         n_knots: Optional[int] = 6,
         spline_order: int = 3,
-        distribution: str = "gamma",
-        link: str = "log",
+        distribution: Literal[
+            "normal", "binomial", "poisson", "gamma", "gaussian", "inv_gauss"
+        ] = "gamma",
+        link: Literal["identity", "logit", "inverse", "log", "inverse-squared"] = "log",
         max_iter: int = 2000,
         expectile: Optional[float] = None,
-        grid: Optional[Union[str, Mapping]] = None,
-        spline_kwargs: Mapping = MappingProxyType({}),
-        **kwargs,
+        grid: Optional[Union[str, Mapping[str, Any]]] = None,
+        spline_kwargs: Mapping[str, Any] = MappingProxyType({}),
+        **kwargs: Any,
     ):
         term = s(
             0,
@@ -112,7 +112,7 @@ def __init__(
         )
         link = GamLinkFunction(link)
         distribution = GamDistribution(distribution)
-        if distribution == GamDistribution.GAUSS:
+        if distribution == GamDistribution.GAUSSIAN:
             distribution = GamDistribution.NORMAL
 
         if expectile is not None:
@@ -139,8 +139,8 @@ def __init__(
                     f"for `{type(gam).__name__!r}`."
                 )
 
-            filtered_kwargs["link"] = link.s
-            filtered_kwargs["distribution"] = distribution.s
+            filtered_kwargs["link"] = link
+            filtered_kwargs["distribution"] = distribution
 
             model = gam(
                 term,
diff --git a/cellrank/ul/models/_utils.py b/cellrank/ul/models/_utils.py
index e1f9885ea..f50b872f5 100644
--- a/cellrank/ul/models/_utils.py
+++ b/cellrank/ul/models/_utils.py
@@ -1,10 +1,12 @@
 from typing import Union, Optional, Sequence
 
+from enum import auto
+
 import cellrank.logging as logg
 from anndata import AnnData
+from cellrank.tl._enum import ModeEnum
 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
@@ -15,11 +17,11 @@
 _OFFSET_KEY = "cellrank_offset"
 
 
-class NormMode(ModeEnum):  # noqa
-    TMM = "tmm"
-    RLE = "rle"
-    UPPER_QUANT = "upper_quant"
-    NONE = "none"
+class NormMode(ModeEnum):  # noqa: D101
+    TMM = auto()
+    RLE = auto()
+    UPPER_QUANT = auto()
+    NONE = auto()
 
 
 @d.dedent
@@ -34,7 +36,8 @@ def _extract_data(
     data
         Annotated data object or an array.
     layer
-        Key in ``adata.layers`` accessed when ``use_raw=False``. Only used when ``data`` is :class:`anndata.AnnData`.
+        Key in :attr:`anndata.AnnData.layers` accessed when ``use_raw = False``.
+        Only used when ``data`` is :class:`anndata.AnnData`.
     use_raw
         Whether to access ``adata.raw``. Only used when ``data`` is :class:`anndata.AnnData`.
 
@@ -114,7 +117,7 @@ def _rankdata(a: np.ndarray, method: str = "average") -> np.ndarray:
 @inject_docs(m=NormMode)
 def _calculate_norm_factors(
     data: Union[AnnData, np.ndarray, spmatrix],
-    method: str = NormMode.TMM.s,
+    method: str = NormMode.TMM,
     layer: Optional[str] = None,
     use_raw: bool = True,
     library_size: Optional[np.ndarray] = None,
@@ -126,7 +129,7 @@ def _calculate_norm_factors(
     perc: float = 0.75,
 ) -> np.ndarray:
     """
-    Calculate normalization factors according to
+    Calculate normalization factors according to \
     `edgeR <https://www.rdocumentation.org/packages/edgeR/versions/3.14.0/topics/calcNormFactors>`__.
 
     Parameters
@@ -134,30 +137,31 @@ def _calculate_norm_factors(
     data
         Data of shape `(n_cells, n_genes)` containing e.g. the counts.
     method
-        One of following:
+        Which method to use. Valid options are:
 
-            - 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.
+            - `{m.TMM!r}` - weighted trimmed mean of M-values from :cite:`robinson:10`.
+            - `{m.RLE!r}` - relative log expression from [Anders10]_.
+            - `{m.UPPER_QUANT!r}` - upper-quartile normalization method from [Bullard10]_.
+            - `{m.NONE!r}` - all the factors are set to 1.
     layer
-        Layer in ``adata.layers`` or `None` for ``adata.X``. Only used when ``use_raw=False``.
+        Layer in :attr:`anndata.AnnData.layers` or `None` for :attr:`anndata.AnnData.X`.
+        Only used when ``use_raw = False``.
     use_raw
-        Whether to access ``adata.raw``.
+        Whether to access :attr:`anndata.AnnData.raw` or not.
     library_size
         Library size. If `None`, it will be set as the sum of values for each cell.
     ref_ix
         Index of a reference cell. If `None`, it will be determined automatically.
     logratio_trim
-        Amount of trim to use on log-ratios ('M' values) when ``method={m.TMM.s!r}``.
+        Amount of trim to use on log-ratios ('M' values) when ``method={m.TMM!r}``.
     sum_trim
-        Amount of trim to use on the combined absolute levels ('A' values) when ``method={m.TMM.s!r}``.
+        Amount of trim to use on the combined absolute levels ('A' values) when ``method = {m.TMM!r}``.
     weight
-        Whether to compute asymptotic binomial precision weights when ``method={m.TMM.s!r}``.
+        Whether to compute asymptotic binomial precision weights when ``method = {m.TMM!r}``.
     a_cutoff
-        Cutoff on 'A' values to use before trimming when ``method={m.TMM.s!r}``.
+        Cutoff on 'A' values to use before trimming when ``method = {m.TMM!r}``.
     perc
-        Percentile of the counts that is aligned when ``method={m.UPPER_QUANT.s!r}``.
+        Percentile of the counts that is aligned when ``method = {m.UPPER_QUANT!r}``.
 
     Returns
     -------
@@ -172,7 +176,7 @@ def _calculate_norm_factors(
     .. [Bullard10] Bullard, J. H. *et al.* (2010),
         *Evaluation of statistical methods for normalization and differential expression in mRNA-Seq experiments*,
         `BMC Bioinformatics <https://doi.org/10.1186/1471-2105-11-94>`__.
-    """  # noqa
+    """
 
     method = NormMode(method)
 
diff --git a/examples/estimators/compute_abs_probs.py b/examples/estimators/compute_abs_probs.py
index bb090b6ca..a7f3bdfcb 100644
--- a/examples/estimators/compute_abs_probs.py
+++ b/examples/estimators/compute_abs_probs.py
@@ -58,11 +58,9 @@
 # absorption to all terminal states. To compute the mean and the variance only for the `"Alpha"` absorbing state,
 # one specify the following ``time_to_absorption={"Alpha": "var"}``.
 g.compute_absorption_probabilities(time_to_absorption="all")
-g.lineage_absorption_times
+g.absorption_times
 
 # %%
 # Lastly, we plot the above computed time.
-adata.obs["mean_time_to_absorption"] = g.lineage_absorption_times[
-    "Alpha, Beta, Epsilon mean"
-]
+adata.obs["mean_time_to_absorption"] = g.absorption_times["Alpha, Beta, Epsilon mean"]
 scv.pl.scatter(adata, color="mean_time_to_absorption")
diff --git a/examples/estimators/compute_lineage_drivers.py b/examples/estimators/compute_lineage_drivers.py
index 2fc9b77aa..935e77b11 100644
--- a/examples/estimators/compute_lineage_drivers.py
+++ b/examples/estimators/compute_lineage_drivers.py
@@ -39,7 +39,7 @@
 # parameters.
 # We also compute the corrected p-values (qval) and the 95% confidence intervals for the correlations.
 g.compute_lineage_drivers(lineages="Alpha")
-g.lineage_drivers.sort_values("Alpha corr", ascending=False)
+g.lineage_drivers.sort_values("Alpha_corr", ascending=False)
 
 # %%
 # Lastly, we plot the top 3 potential driver genes for the `"Alpha"` lineage.
diff --git a/examples/plotting/plot_heatmap.py b/examples/plotting/plot_heatmap.py
index 0e5fcc4c8..3dd0442fb 100644
--- a/examples/plotting/plot_heatmap.py
+++ b/examples/plotting/plot_heatmap.py
@@ -28,7 +28,7 @@
 model = cr.ul.models.GAM(adata)
 
 # %%
-# We can now plot the heatmap. By default, the model is fitted for all specified genes and all lineages. Paramterer
+# We can now plot the heatmap. By default, the model is fitted for all specified genes and all lineages. Parameter
 # ``show_absorption_probabilities`` is used to create a bar on top of the heatmap.
 #
 # Apart from the default gene sorting, we can use hierarchical clustering to cluster the genes by specifying
diff --git a/setup.py b/setup.py
index b0ad3137f..fea48cd76 100644
--- a/setup.py
+++ b/setup.py
@@ -69,13 +69,12 @@
         ),
         zip_safe=False,
         packages=find_packages(),
-        python_requires=">=3.6",
+        python_requires=">=3.7",
         platforms=["Linux", "MacOs", "Windows"],
         keywords=sorted(
             [
                 "single-cell",
                 "bio-informatics",
-                "single-cell",
                 "RNA velocity",
                 "Markov chain",
                 "GPCCA",
@@ -87,9 +86,9 @@
             "Intended Audience :: Science/Research",
             "Natural Language :: English",
             "Framework :: Jupyter",
+            "Operating System :: POSIX :: Linux",
             "Operating System :: MacOS :: MacOS X",
             "Operating System :: Microsoft :: Windows",
-            "Operating System :: POSIX :: Linux",
             "Typing :: Typed",
             "Programming Language :: Python :: 3",
             "Programming Language :: Python :: 3.7",
diff --git a/tox.ini b/tox.ini
index 614d79c22..79f307c9f 100644
--- a/tox.ini
+++ b/tox.ini
@@ -3,6 +3,7 @@ per-file-ignores =
     */*.py: F401,D104,D100
     examples/*: D400,D205,D100
     docs/source/index.rst: D000
+    cellrank/_key.py: D101,D102,D106
 max_line_length = 120
 filename = *.py
 # D202 No blank lines allowed after function docstring
@@ -13,6 +14,7 @@ filename = *.py
 # E203 whitespace before ':'
 # F405 ... may be undefined, or defined from star imports: ...
 # RST306 Unknown target name
+# D106: Missing docstring in public nested class
 ignore = D202,D107,B008,W503,D105,E203,F405,RST306
 exclude =
     .git,


Test patch:
diff --git a/tests/_ground_truth_adatas/adata_100.h5ad b/tests/_ground_truth_adatas/adata_100.h5ad
index 97c3b61ce..a1b391b87 100644
Binary files a/tests/_ground_truth_adatas/adata_100.h5ad and b/tests/_ground_truth_adatas/adata_100.h5ad differ
diff --git a/tests/_ground_truth_adatas/adata_200.h5ad b/tests/_ground_truth_adatas/adata_200.h5ad
index 0ea52c7a9..d1d52a309 100644
Binary files a/tests/_ground_truth_adatas/adata_200.h5ad and b/tests/_ground_truth_adatas/adata_200.h5ad differ
diff --git a/tests/_ground_truth_adatas/adata_50.h5ad b/tests/_ground_truth_adatas/adata_50.h5ad
index 89e4dee18..60ab43550 100644
Binary files a/tests/_ground_truth_adatas/adata_50.h5ad and b/tests/_ground_truth_adatas/adata_50.h5ad differ
diff --git a/tests/_ground_truth_figures/driver_corr_adjust_text.png b/tests/_ground_truth_figures/driver_corr_adjust_text.png
index 6e841088f..95caa6f8b 100644
Binary files a/tests/_ground_truth_figures/driver_corr_adjust_text.png and b/tests/_ground_truth_figures/driver_corr_adjust_text.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_cmap.png b/tests/_ground_truth_figures/driver_corr_cmap.png
index 94e11098a..cab76756c 100644
Binary files a/tests/_ground_truth_figures/driver_corr_cmap.png and b/tests/_ground_truth_figures/driver_corr_cmap.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_color.png b/tests/_ground_truth_figures/driver_corr_color.png
index aeb3753df..a6385b72a 100644
Binary files a/tests/_ground_truth_figures/driver_corr_color.png and b/tests/_ground_truth_figures/driver_corr_color.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_fontsize.png b/tests/_ground_truth_figures/driver_corr_fontsize.png
index 8226ecfe5..384dd2b5b 100644
Binary files a/tests/_ground_truth_figures/driver_corr_fontsize.png and b/tests/_ground_truth_figures/driver_corr_fontsize.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets.png b/tests/_ground_truth_figures/driver_corr_gene_sets.png
index 444576d3a..79c92609d 100644
Binary files a/tests/_ground_truth_figures/driver_corr_gene_sets.png and b/tests/_ground_truth_figures/driver_corr_gene_sets.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png
index 23461b073..f24953d5e 100644
Binary files a/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png and b/tests/_ground_truth_figures/driver_corr_gene_sets_colors.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_legend_loc.png b/tests/_ground_truth_figures/driver_corr_legend_loc.png
index 555b5744c..a16c49400 100644
Binary files a/tests/_ground_truth_figures/driver_corr_legend_loc.png and b/tests/_ground_truth_figures/driver_corr_legend_loc.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_return_ax.png b/tests/_ground_truth_figures/driver_corr_return_ax.png
index 0cbdddfa9..19beb887b 100644
Binary files a/tests/_ground_truth_figures/driver_corr_return_ax.png and b/tests/_ground_truth_figures/driver_corr_return_ax.png differ
diff --git a/tests/_ground_truth_figures/driver_corr_use_raw.png b/tests/_ground_truth_figures/driver_corr_use_raw.png
index d09bb54d8..dc9642800 100644
Binary files a/tests/_ground_truth_figures/driver_corr_use_raw.png and b/tests/_ground_truth_figures/driver_corr_use_raw.png differ
diff --git a/tests/_ground_truth_figures/scvelo_eigendecomposition_clusters.png b/tests/_ground_truth_figures/scvelo_eigendecomposition_clusters.png
deleted file mode 100644
index ced07307d..000000000
Binary files a/tests/_ground_truth_figures/scvelo_eigendecomposition_clusters.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_eigendecomposition_left.png b/tests/_ground_truth_figures/scvelo_eigendecomposition_left.png
deleted file mode 100644
index 053e9dbd7..000000000
Binary files a/tests/_ground_truth_figures/scvelo_eigendecomposition_left.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_eigendecomposition_right.png b/tests/_ground_truth_figures/scvelo_eigendecomposition_right.png
deleted file mode 100644
index 053e9dbd7..000000000
Binary files a/tests/_ground_truth_figures/scvelo_eigendecomposition_right.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_eigendecomposition_use_2.png b/tests/_ground_truth_figures/scvelo_eigendecomposition_use_2.png
deleted file mode 100644
index 6a8e690f9..000000000
Binary files a/tests/_ground_truth_figures/scvelo_eigendecomposition_use_2.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_schur_cluster_key.png b/tests/_ground_truth_figures/scvelo_gpcca_schur_cluster_key.png
deleted file mode 100644
index ea5b2bdf7..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_schur_cluster_key.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb.png b/tests/_ground_truth_figures/scvelo_gpcca_schur_emb.png
deleted file mode 100644
index 9c47d7eda..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_abs.png b/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_abs.png
deleted file mode 100644
index 0bfb1f1bf..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_abs.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_use_2.png b/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_use_2.png
deleted file mode 100644
index 1647269b9..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_schur_emb_use_2.png and /dev/null differ
diff --git a/tests/_helpers.py b/tests/_helpers.py
index e50004a49..2c8ac52b3 100644
--- a/tests/_helpers.py
+++ b/tests/_helpers.py
@@ -3,19 +3,19 @@
 import os
 import pytest
 from PIL import Image
-from sys import version_info
 from pathlib import Path
 
 import scanpy as sc
 import scvelo as scv
 import cellrank as cr
 from anndata import AnnData
-from cellrank.tl.kernels import VelocityKernel, CytoTRACEKernel, ConnectivityKernel
+from cellrank.tl.kernels import VelocityKernel, PrecomputedKernel, ConnectivityKernel
 
 import numpy as np
 import pandas as pd
 from sklearn.svm import SVR
 from scipy.sparse import spdiags, issparse, csr_matrix
+from pandas.testing import assert_frame_equal, assert_series_equal
 
 
 def _jax_not_installed() -> bool:
@@ -121,12 +121,12 @@ def create_kernels(
     connectivity_variances: Optional[str] = None,
 ) -> Tuple[VelocityKernel, ConnectivityKernel]:
     vk = VelocityKernel(adata)
-    vk._mat_scaler = adata.uns.get(
+    vk._mat_scaler = adata.obsp.get(
         velocity_variances, np.random.normal(size=(adata.n_obs, adata.n_obs))
     )
 
     ck = ConnectivityKernel(adata)
-    ck._mat_scaler = adata.uns.get(
+    ck._mat_scaler = adata.obsp.get(
         connectivity_variances, np.random.normal(size=(adata.n_obs, adata.n_obs))
     )
 
@@ -209,11 +209,12 @@ def assert_models_equal(
     deepcopy: bool = True,
 ) -> None:
     assert actual is not expected
-    if not pickled:
-        assert actual.adata is expected.adata
-    else:
+    if pickled:
         assert actual.adata is not expected.adata
+    else:
+        assert actual.adata is expected.adata
 
+    assert actual.shape == expected.shape
     assert actual.adata.shape == expected.adata.shape
 
     assert expected.__dict__.keys() == actual.__dict__.keys()
@@ -259,18 +260,33 @@ def assert_estimators_equal(
     expected: cr.tl.estimators.BaseEstimator,
     actual: cr.tl.estimators.BaseEstimator,
     copy: bool = False,
+    deep: bool = False,
+    from_adata: bool = False,
 ) -> None:
+    def check_arrays(x, y):
+        if isinstance(x, (np.ndarray, list, tuple)):
+            try:
+                np.testing.assert_array_compare(np.array_equal, x, y, equal_nan=True)
+            except:
+                np.testing.assert_array_compare(np.allclose, x, y, equal_nan=True)
+        elif isinstance(x, pd.Series):
+            assert_series_equal(x, y, check_names=False)
+        elif isinstance(x, pd.DataFrame):
+            assert_frame_equal(x, y, check_dtype=False)
+
     assert actual is not expected
-    assert actual.adata is not expected.adata
+    if copy:
+        if deep:
+            assert actual.adata is not expected.adata
+        else:
+            assert actual.adata is expected.adata
+    else:
+        assert actual.adata is not expected.adata
     assert actual.kernel is not expected.kernel
-    if copy or version_info[:2] > (3, 6):
-        # pickling of Enums doesn't work in Python3.6
-        assert isinstance(actual.kernel, type(expected.kernel)), (
-            type(actual.kernel),
-            type(expected.kernel),
-        )
+    if from_adata:
+        assert isinstance(actual.kernel, PrecomputedKernel)
     else:
-        assert isinstance(actual.kernel, cr.tl.kernels.PrecomputedKernel)
+        assert isinstance(actual.kernel, type(expected.kernel))
 
     assert actual.adata.shape == expected.adata.shape
     assert actual.adata is actual.kernel.adata
@@ -280,33 +296,52 @@ def assert_estimators_equal(
         actual.transition_matrix.A, expected.transition_matrix.A
     )
 
-    assert expected.__dict__.keys() == actual.__dict__.keys()
+    k1 = sorted(expected.__dict__.keys())
+    k2 = sorted(actual.__dict__.keys())
+    np.testing.assert_array_equal(k1, k2)
 
     for attr in expected.__dict__.keys():
-        val2, val1 = getattr(actual, attr), getattr(expected, attr)
-        if isinstance(val1, cr.tl.Lineage):
-            assert val2 is not val1, attr
-            assert_array_nan_equal(val2.X, val1.X)
-        elif isinstance(val1, (np.ndarray, pd.Series, pd.DataFrame)):
-            assert val2 is not val1, attr
+        if attr == "_invalid_n_states" and from_adata:
+            continue
+        actual_val, expected_val = getattr(actual, attr), getattr(expected, attr)
+        if isinstance(actual_val, cr.tl.Lineage):
+            assert actual_val is not expected_val, attr
+            assert_array_nan_equal(actual_val.X, expected_val.X)
+        elif isinstance(actual_val, (np.ndarray, pd.Series, pd.DataFrame, list, tuple)):
+            assert actual_val is not expected_val, attr
+            check_arrays(actual_val, expected_val)
+        elif isinstance(actual_val, dict):
+            if from_adata:
+                # _params can sometimes contain extra empty dict if initialized from `adata`
+                for k in set(actual_val.keys()) | set(expected_val.keys()):
+                    v2, v1 = actual_val.get(k, {}), expected_val.get(k, {})
+                    if isinstance(v1, (np.ndarray, pd.Series, pd.DataFrame)):
+                        check_arrays(v2, v1)
+                    else:
+                        assert v2 == v1, (v2, v1, attr, k)
+            else:
+                np.testing.assert_array_equal(
+                    sorted(actual_val.keys()), sorted(expected_val.keys())
+                )
+                for k in sorted(actual_val.keys()):
+                    v2, v1 = actual_val[k], expected_val[k]
+                    if isinstance(v1, (np.ndarray, pd.Series, pd.DataFrame)):
+                        check_arrays(v2, v1)
+                    else:
+                        assert v2 == v1, (v2, v1, attr, k)
+        elif attr not in ("_kernel", "_gpcca", "_adata", "_shadow_adata"):
+            assert actual_val == expected_val, (actual_val, expected_val, attr)
+        else:
             try:
-                # can be array of strings, can't get NaN
-                assert_array_nan_equal(val2, val1)
-            except:
-                np.testing.assert_array_equal(val2, val1)
-        elif isinstance(val1, dict):
-            assert val2.keys() == val1.keys()
-            for v2, v1 in zip(val2.values(), val1.values()):
-                if isinstance(v2, np.ndarray):
-                    assert v2 is not v1, attr
-                    np.testing.assert_array_equal(v2, v1)
-                else:
-                    assert v2 == v1, (v2, v1, attr)
-        elif attr not in ("_kernel", "_gpcca"):
-            assert val2 == val1, (val2, val1, attr)
-        elif copy or version_info[:2] > (3, 6):
-            # we can compare the kernel types, but for 3.6, it's saved as Precomputed
-            assert isinstance(val2, type(val1)), (val2, val1, attr)
+                assert isinstance(actual_val, type(expected_val)), (
+                    actual_val,
+                    expected_val,
+                    attr,
+                )
+            except AssertionError:
+                # objects initialized from `adata` don't have `_gpcca`
+                if attr != "_gpcca" and not from_adata:
+                    raise
 
 
 def random_transition_matrix(n: int) -> np.ndarray:
@@ -368,8 +403,8 @@ def _create_dummy_adata(n_obs: int) -> AnnData:
     sc.tl.dpt(adata)
 
     if "velocity_graph" in adata.uns:
-        adata.obsp["velocity_graph"] = adata.uns["velocity_graph"]
-        adata.obsp["velocity_graph_neg"] = adata.uns["velocity_graph_neg"]
+        adata.obsp["velocity_graph"] = adata.uns.pop("velocity_graph")
+        adata.obsp["velocity_graph_neg"] = adata.uns.pop("velocity_graph_neg")
     adata.obsp["connectivity_variances"] = np.ones((n_obs, n_obs), dtype=np.float64)
     adata.obsp["velocity_variances"] = np.ones((n_obs, n_obs), dtype=np.float64)
 
diff --git a/tests/conftest.py b/tests/conftest.py
index 03def0584..0624423ae 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -17,7 +17,6 @@
 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
@@ -63,18 +62,14 @@ def _create_cflare(*, backward: bool = False) -> Tuple[AnnData, CFLARE]:
     final_kernel = 0.8 * vk + 0.2 * ck
 
     mc = CFLARE(final_kernel)
+    final_kernel.write_to_adata()
 
-    mc.compute_partition()
     mc.compute_eigendecomposition()
-    mc.compute_terminal_states(use=2)
+    mc.compute_terminal_states(use=2, method="kmeans")
     mc.compute_absorption_probabilities(use_petsc=False)
     mc.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
     assert adata is mc.adata
-    if backward:
-        assert str(AbsProbKey.BACKWARD) in adata.obsm
-    else:
-        assert str(AbsProbKey.FORWARD) in adata.obsm
     np.testing.assert_allclose(mc.absorption_probabilities.X.sum(1), 1.0, rtol=1e-6)
 
     return adata, mc
@@ -92,20 +87,16 @@ def _create_gpcca(*, backward: bool = False) -> Tuple[AnnData, GPCCA]:
     final_kernel = 0.8 * vk + 0.2 * ck
 
     mc = GPCCA(final_kernel)
+    final_kernel.write_to_adata()
 
-    mc.compute_partition()
     mc.compute_eigendecomposition()
     mc.compute_schur(method="krylov")
     mc.compute_macrostates(n_states=2)
     mc.set_terminal_states_from_macrostates()
-    mc.compute_absorption_probabilities()
+    mc.compute_absorption_probabilities(use_petsc=False)
     mc.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
     assert adata is mc.adata
-    if backward:
-        assert str(AbsProbKey.BACKWARD) in adata.obsm
-    else:
-        assert str(AbsProbKey.FORWARD) in adata.obsm
     np.testing.assert_allclose(mc.absorption_probabilities.X.sum(1), 1.0, rtol=1e-6)
 
     return adata, mc
@@ -117,7 +108,7 @@ def _create_gamr_model(_adata: AnnData) -> Optional[GAMR]:
         m.prepare(_adata.var_names[0], "0").fit()
         m.predict(level=0.95)
         return m
-    except:
+    except Exception:
         return None
 
 
diff --git a/tests/test_cflare.py b/tests/test_cflare.py
index 50b4e7494..964ddf60b 100644
--- a/tests/test_cflare.py
+++ b/tests/test_cflare.py
@@ -3,23 +3,11 @@
 import os
 import pytest
 from _helpers import assert_estimators_equal
-from tempfile import TemporaryDirectory
 
 import cellrank as cr
-import cellrank.tl.kernels._precomputed_kernel
 from anndata import AnnData
+from cellrank._key import Key
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
-from cellrank.tl._constants import (
-    Direction,
-    DirPrefix,
-    AbsProbKey,
-    TermStatesKey,
-    _pd,
-    _probs,
-    _colors,
-    _lin_names,
-)
-from cellrank.tl.estimators._constants import A, P
 
 import numpy as np
 import pandas as pd
@@ -29,24 +17,6 @@
 
 
 class TestCFLARE:
-    def test_compute_partition(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.CFLARE(terminal_kernel)
-        mc.compute_partition()
-
-        assert isinstance(mc.is_irreducible, bool)
-        if not mc.is_irreducible:
-            assert isinstance(mc.recurrent_classes, list)
-            assert isinstance(mc.transient_classes, list)
-            assert f"{TermStatesKey.FORWARD}_rec_classes" in mc.adata.obs
-            assert f"{TermStatesKey.FORWARD}_trans_classes" in mc.adata.obs
-        else:
-            assert mc.recurrent_classes is None
-            assert mc.transient_classes is None
-
     def test_compute_eigendecomposition(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -56,7 +26,7 @@ def test_compute_eigendecomposition(self, adata_large: AnnData):
         mc.compute_eigendecomposition(k=2)
 
         assert isinstance(mc.eigendecomposition, dict)
-        assert set(mc._get(P.EIG).keys()) == {
+        assert set(mc.eigendecomposition.keys()) == {
             "D",
             "V_l",
             "V_r",
@@ -64,7 +34,7 @@ def test_compute_eigendecomposition(self, adata_large: AnnData):
             "params",
             "stationary_dist",
         }
-        assert f"eig_{Direction.FORWARD}" in mc.adata.uns.keys()
+        assert Key.uns.eigen(mc.backward) in mc.adata.uns
 
     def test_compute_terminal_states_no_eig(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -94,12 +64,13 @@ def test_compute_approx_normal_run(self, adata_large: AnnData):
         mc.compute_eigendecomposition(k=5)
         mc.compute_terminal_states(use=2)
 
-        assert is_categorical_dtype(mc._get(P.TERM))
-        assert mc._get(P.TERM_PROBS) is not None
+        assert is_categorical_dtype(mc.terminal_states)
+        assert mc.terminal_states_probabilities is not None
 
-        assert TermStatesKey.FORWARD.s in mc.adata.obs.keys()
-        assert _probs(TermStatesKey.FORWARD) in mc.adata.obs.keys()
-        assert _colors(TermStatesKey.FORWARD) in mc.adata.uns.keys()
+        key = Key.obs.term_states(mc.backward)
+        assert key in mc.adata.obs
+        assert Key.obs.probs(key) in mc.adata.obs
+        assert Key.uns.colors(key) in mc.adata.uns
 
     def test_rename_terminal_states_no_terminal_states(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -129,11 +100,11 @@ def test_rename_terminal_states_invalid_new_name(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.compute_terminal_states(use=2, method="kmeans")
         with pytest.raises(ValueError):
             mc.rename_terminal_states({"0": "1"})
 
-    def test_rename_terminal_states_try_joinining_states(self, adata_large: AnnData):
+    def test_rename_terminal_states_try_joining_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
@@ -153,10 +124,10 @@ def test_rename_terminal_states_empty_mapping(self, adata_large: AnnData):
         mc.compute_eigendecomposition(k=5)
         mc.compute_terminal_states(use=2)
 
-        orig_cats = list(mc._get(P.TERM).cat.categories)
+        orig_cats = list(mc.terminal_states.cat.categories)
         mc.rename_terminal_states({})
 
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, orig_cats)
+        np.testing.assert_array_equal(mc.terminal_states.cat.categories, orig_cats)
 
     def test_rename_terminal_states_normal_run(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -165,63 +136,13 @@ def test_rename_terminal_states_normal_run(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
-
-        mc.rename_terminal_states({"0": "foo", "1": "bar"})
-
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, ["foo", "bar"])
-
-    def test_rename_terminal_states_non_string_new_names(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.CFLARE(terminal_kernel)
-        mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
-
-        mc.rename_terminal_states({"0": 42, "1": object})
-
-        np.testing.assert_array_equal(
-            mc._get(P.TERM).cat.categories, ["42", str(object)]
-        )
-
-    def test_rename_terminal_states_update_adata(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.CFLARE(terminal_kernel)
-        mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
-
+        mc.compute_terminal_states(use=2, method="kmeans")
         mc.rename_terminal_states({"0": "foo", "1": "bar"})
 
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, ["foo", "bar"])
-        np.testing.assert_array_equal(
-            mc.adata.obs[TermStatesKey.FORWARD.s].cat.categories, ["foo", "bar"]
-        )
-        np.testing.assert_array_equal(
-            mc.adata.uns[_lin_names(TermStatesKey.FORWARD.s)], ["foo", "bar"]
-        )
-
-    def test_rename_terminal_states_dont_update_adata(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.CFLARE(terminal_kernel)
-        mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
-
-        mc.rename_terminal_states({"0": "foo", "1": "bar"}, update_adata=False)
-
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, ["foo", "bar"])
-        np.testing.assert_array_equal(
-            mc.adata.obs[TermStatesKey.FORWARD.s].cat.categories, ["0", "1"]
-        )
+        np.testing.assert_array_equal(mc.terminal_states.cat.categories, ["foo", "bar"])
         np.testing.assert_array_equal(
-            mc.adata.uns[_lin_names(TermStatesKey.FORWARD.s)], ["0", "1"]
+            mc.adata.obs[Key.obs.term_states(mc.backward)].cat.categories,
+            ["foo", "bar"],
         )
 
     def test_compute_absorption_probabilities_no_args(self, adata_large: AnnData):
@@ -240,35 +161,34 @@ def test_compute_absorption_probabilities_normal_run(self, adata_large: AnnData)
 
         mc = cr.tl.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.compute_terminal_states(use=2, method="kmeans")
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
-        assert isinstance(mc._get(P.PRIME_DEG), pd.Series)
-        assert _pd(AbsProbKey.FORWARD) in mc.adata.obs.keys()
-        np.testing.assert_array_equal(
-            mc._get(P.PRIME_DEG), mc.adata.obs[_pd(AbsProbKey.FORWARD)]
-        )
+        key = Key.obs.priming_degree(mc.backward)
+        assert isinstance(mc.priming_degree, pd.Series)
+        assert key in mc.adata.obs
+        np.testing.assert_array_equal(mc.priming_degree, mc.adata.obs[key])
 
-        assert isinstance(mc._get(P.ABS_PROBS), cr.tl.Lineage)
-        assert mc._get(P.ABS_PROBS).shape == (mc.adata.n_obs, 2)
-        assert f"{AbsProbKey.FORWARD}" in mc.adata.obsm.keys()
-        np.testing.assert_array_equal(
-            mc._get(P.ABS_PROBS).X, mc.adata.obsm[f"{AbsProbKey.FORWARD}"]
-        )
+        key = Key.obsm.abs_probs(mc.backward)
+        assert isinstance(mc.absorption_probabilities, cr.tl.Lineage)
+        assert mc.absorption_probabilities.shape == (mc.adata.n_obs, 2)
+        assert key in mc.adata.obsm
+        assert isinstance(mc.adata.obsm[key], cr.tl.Lineage)
+        np.testing.assert_array_equal(mc.absorption_probabilities.X, mc.adata.obsm[key])
 
-        assert _lin_names(AbsProbKey.FORWARD) in mc.adata.uns.keys()
         np.testing.assert_array_equal(
-            mc._get(P.ABS_PROBS).names,
-            mc.adata.uns[_lin_names(AbsProbKey.FORWARD)],
+            mc.absorption_probabilities.names,
+            mc.adata.obs[Key.obs.term_states(mc.backward)].cat.categories,
         )
 
-        assert _colors(AbsProbKey.FORWARD) in mc.adata.uns.keys()
+        key = Key.uns.colors(Key.obs.term_states(mc.backward))
+        assert key in mc.adata.uns
         np.testing.assert_array_equal(
-            mc._get(P.ABS_PROBS).colors,
-            mc.adata.uns[_colors(AbsProbKey.FORWARD)],
+            mc.absorption_probabilities.colors, mc.adata.uns[key]
         )
-        np.testing.assert_allclose(mc._get(P.ABS_PROBS).X.sum(1), 1, rtol=1e-6)
+        np.testing.assert_array_equal(mc._term_states_colors, mc.adata.uns[key])
+        np.testing.assert_allclose(mc.absorption_probabilities.X.sum(1), 1, rtol=1e-6)
 
     def test_compute_absorption_probabilities_solver(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -282,11 +202,11 @@ def test_compute_absorption_probabilities_solver(self, adata_large: AnnData):
 
         # compute lin probs using direct solver
         mc.compute_absorption_probabilities(solver="direct")
-        l_direct = mc._get(P.ABS_PROBS).copy()
+        l_direct = mc.absorption_probabilities.copy()
 
         # compute lin probs using iterative solver
         mc.compute_absorption_probabilities(solver="gmres", tol=tol)
-        l_iterative = mc._get(P.ABS_PROBS).copy()
+        l_iterative = mc.absorption_probabilities.copy()
 
         assert not np.shares_memory(l_direct.X, l_iterative.X)  # sanity check
         np.testing.assert_allclose(l_direct.X, l_iterative.X, rtol=0, atol=tol)
@@ -299,15 +219,15 @@ def test_compute_absorption_probabilities_solver_petsc(self, adata_large: AnnDat
 
         mc = cr.tl.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.compute_terminal_states(use=2, method="kmeans")
 
         # compute lin probs using direct solver
         mc.compute_absorption_probabilities(solver="gmres", use_petsc=False, tol=tol)
-        l_iter = mc._get(P.ABS_PROBS).copy()
+        l_iter = mc.absorption_probabilities.copy()
 
         # compute lin probs using petsc iterative solver
         mc.compute_absorption_probabilities(solver="gmres", use_petsc=True, tol=tol)
-        l_iter_petsc = mc._get(P.ABS_PROBS).copy()
+        l_iter_petsc = mc.absorption_probabilities.copy()
 
         assert not np.shares_memory(l_iter.X, l_iter_petsc.X)  # sanity check
         np.testing.assert_allclose(l_iter.X, l_iter_petsc.X, rtol=0, atol=tol)
@@ -331,7 +251,7 @@ def test_compute_absorption_probabilities_lineage_absorption_mean(
             tol=tol,
             time_to_absorption="0",
         )
-        at = mc._get(P.LIN_ABS_TIMES)
+        at = mc.absorption_times
 
         assert isinstance(at, pd.DataFrame)
         np.testing.assert_array_equal(at.index, adata_large.obs_names)
@@ -353,7 +273,7 @@ def test_compute_absorption_probabilities_lineage_absorption_var(
         mc.compute_absorption_probabilities(
             solver="gmres", use_petsc=False, tol=tol, time_to_absorption={"0": "var"}
         )
-        at = mc._get(P.LIN_ABS_TIMES)
+        at = mc.absorption_times
 
         assert isinstance(at, pd.DataFrame)
         np.testing.assert_array_equal(at.index, adata_large.obs_names)
@@ -375,8 +295,8 @@ def test_compute_absorption_probabilities_lineage_absorption_all(
         mc.compute_absorption_probabilities(
             solver="gmres", use_petsc=False, tol=tol, time_to_absorption={"all": "var"}
         )
-        name = ", ".join(mc._get(P.ABS_PROBS).names)
-        at = mc._get(P.LIN_ABS_TIMES)
+        name = ", ".join(mc.absorption_probabilities.names)
+        at = mc.absorption_times
 
         assert isinstance(at, pd.DataFrame)
         np.testing.assert_array_equal(at.index, adata_large.obs_names)
@@ -426,16 +346,17 @@ def test_compute_lineage_drivers_normal_run(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.compute_terminal_states(use=2, method="kmeans")
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
+        key = Key.varm.lineage_drivers(False)
         for lineage in ["0", "1"]:
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} corr"] >= -1.0)
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} corr"] <= 1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_corr"] >= -1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_corr"] <= 1.0)
 
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} qval"] >= 0)
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} qval"] <= 1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_qval"] >= 0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_qval"] <= 1.0)
 
     def test_plot_lineage_drivers_not_computed(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -497,25 +418,24 @@ def test_compute_absorption_probabilities_keys_colors(self, adata_large: AnnData
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc_fwd = cr.tl.estimators.CFLARE(terminal_kernel)
-        mc_fwd.compute_partition()
         mc_fwd.compute_eigendecomposition()
-        mc_fwd.compute_terminal_states(use=3)
+        mc_fwd.compute_terminal_states(use=3, method="kmeans")
 
         arcs = ["0", "2"]
         arc_colors = [
             c
             for arc, c in zip(
-                mc_fwd._get(P.TERM).cat.categories, mc_fwd._get(A.TERM_COLORS)
+                mc_fwd.terminal_states.cat.categories, mc_fwd._term_states_colors
             )
             if arc in arcs
         ]
 
         mc_fwd.compute_absorption_probabilities(keys=arcs)
-        lin_colors = mc_fwd._get(P.ABS_PROBS)[arcs].colors
+        lin_colors = mc_fwd.absorption_probabilities[arcs].colors
 
         np.testing.assert_array_equal(arc_colors, lin_colors)
 
-    def compare_absorption_probabilites_with_reference(self):
+    def test_compare_absorption_probabilities_with_reference(self):
         # define a reference transition matrix. This is an absorbing MC with 2 absorbing states
         transition_matrix = np.array(
             [
@@ -550,23 +470,21 @@ def compare_absorption_probabilites_with_reference(self):
             ]
         )
 
-        # initialise a pre-computed kernel and CFLARE estimator object
-        c = cr.tl.estimators.CFLARE(
-            cellrank.tl.kernels._precomputed_kernel.PrecomputedKernel(transition_matrix)
-        )
+        c = cr.tl.estimators.CFLARE(cr.tl.kernels.PrecomputedKernel(transition_matrix))
 
-        # define the set of macrostates
-        state_annotation = pd.Series(index=range(p.shape[0]))
+        state_annotation = pd.Series(index=range(len(c)))
         state_annotation[7] = "terminal_1"
         state_annotation[10] = "terminal_2"
         state_annotation = state_annotation.astype("category")
-        c._set(A.TERM, state_annotation)
+        c._term_states = state_annotation
+        c._term_states_colors = np.array(["#000000", "#ffffff"])
 
         # compute absorption probabilities
         c.compute_absorption_probabilities()
-        absorption_probabilities_query = c.get(P.ABS_PROBS)[state_annotation.isna()]
+        absorption_probabilities_query = c.absorption_probabilities[
+            state_annotation.isna()
+        ]
 
-        # check whether these two agree
         np.allclose(absorption_probabilities_query, absorption_probabilities_reference)
 
     def test_manual_approx_rc_set(self, adata_large):
@@ -576,18 +494,18 @@ def test_manual_approx_rc_set(self, adata_large):
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc_fwd = cr.tl.estimators.CFLARE(terminal_kernel)
-        mc_fwd.compute_partition()
         mc_fwd.compute_eigendecomposition()
+        key = Key.obs.term_states(mc_fwd.backward)
 
         mc_fwd.compute_terminal_states(use=3)
-        original = np.array(adata.obs[f"{TermStatesKey.FORWARD}"].copy())
+        original = np.array(adata.obs[key].copy())
         zero_mask = original == "0"
 
         cells = list(adata[zero_mask].obs_names)
         mc_fwd.set_terminal_states({"foo": cells})
 
-        assert (adata.obs[f"{TermStatesKey.FORWARD}"][zero_mask] == "foo").all()
-        assert pd.isna(adata.obs[f"{TermStatesKey.FORWARD}"][~zero_mask]).all()
+        assert (adata.obs[key][zero_mask] == "foo").all()
+        assert pd.isna(adata.obs[key][~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)
@@ -605,7 +523,7 @@ def test_abs_probs_do_not_sum_to_1(self, adata_large: AnnData, mocker):
         abs_prob[:, 0] = 1.0
         abs_prob[0, 0] = 1.01
         mocker.patch(
-            "cellrank.tl.estimators._base_estimator._solve_lin_system",
+            "cellrank.tl.estimators.mixins._absorption_probabilities._solve_lin_system",
             return_value=abs_prob,
         )
 
@@ -631,7 +549,7 @@ def test_abs_probs_negative(self, adata_large: AnnData, mocker):
         abs_prob[0, 0] = -0.5
         abs_prob[0, 1] = -1.5
         mocker.patch(
-            "cellrank.tl.estimators._base_estimator._solve_lin_system",
+            "cellrank.tl.estimators.mixins._absorption_probabilities._solve_lin_system",
             return_value=abs_prob,
         )
 
@@ -640,17 +558,21 @@ def test_abs_probs_negative(self, adata_large: AnnData, mocker):
 
 
 class TestCFLAREIO:
-    def test_copy(self, adata_cflare_fwd: Tuple[AnnData, cr.tl.estimators.CFLARE]):
+    @pytest.mark.parametrize("deep", [False, True])
+    def test_copy(
+        self, adata_cflare_fwd: Tuple[AnnData, cr.tl.estimators.CFLARE], deep: bool
+    ):
         _, mc1 = adata_cflare_fwd
-        mc2 = mc1.copy()
+        mc2 = mc1.copy(deep=deep)
 
-        assert_estimators_equal(mc1, mc2, copy=True)
+        assert_estimators_equal(mc1, mc2, copy=True, deep=deep)
 
-    def test_read(self, adata_cflare_fwd: Tuple[AnnData, cr.tl.estimators.CFLARE]):
+    def test_read(
+        self, adata_cflare_fwd: Tuple[AnnData, cr.tl.estimators.CFLARE], tmpdir
+    ):
         _, mc1 = adata_cflare_fwd
 
-        with TemporaryDirectory() as tmpdir:
-            mc1.write(os.path.join(tmpdir, "foo"))
-            mc2 = cr.tl.estimators.CFLARE.read(os.path.join(tmpdir, "foo.pickle"))
+        mc1.write(os.path.join(tmpdir, "foo.pickle"))
+        mc2 = cr.tl.estimators.CFLARE.read(os.path.join(tmpdir, "foo.pickle"))
 
         assert_estimators_equal(mc1, mc2)
diff --git a/tests/test_external.py b/tests/test_external.py
index 8917736a0..34f2dc9b0 100644
--- a/tests/test_external.py
+++ b/tests/test_external.py
@@ -215,7 +215,7 @@ def test_birth_death_process(self, adata_large: AnnData, key_added: Optional[str
     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,
+            last_time_point=ltp,
             conn_kwargs={"n_neighbors": 11},
             threshold=None,
         )
diff --git a/tests/test_gpcca.py b/tests/test_gpcca.py
index 3446eeefc..591f452f8 100644
--- a/tests/test_gpcca.py
+++ b/tests/test_gpcca.py
@@ -1,138 +1,323 @@
-from typing import List, Tuple, Union
+from typing import List, Tuple, Union, Optional, Sequence
 
 import os
 import pytest
 from copy import deepcopy
+from enum import Enum
 from _helpers import assert_array_nan_equal, assert_estimators_equal
-from tempfile import TemporaryDirectory
 
 import cellrank as cr
 from anndata import AnnData
+from cellrank.tl import Lineage
+from cellrank._key import Key
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
-from cellrank.tl._constants import (
-    Direction,
-    DirPrefix,
-    AbsProbKey,
-    TermStatesKey,
-    _pd,
-    _probs,
-    _colors,
-    _lin_names,
-)
-from cellrank.tl.estimators._constants import A, P
 
 import numpy as np
 import pandas as pd
-from pandas.testing import assert_series_equal
+from scipy.sparse import issparse
+from pandas.testing import assert_frame_equal, assert_series_equal
+
+
+# fmt: off
+class State(str, Enum):
+    SCHUR = "schur"
+    MACRO = "macro"
+    TERM = "term"
+    ABS = "abs"
+    LIN = "lin"
+
+    @property
+    def prev(self) -> Optional["State"]:
+        if self.value == "schur":
+            return None
+        if self.value == "macro":
+            return State.SCHUR
+        if self.value == "term":
+            return State.MACRO
+        if self.value == "abs":
+            return State.TERM
+        if self.value == "lin":
+            return State.ABS
+        return None
+
+    @property
+    def next(self) -> Optional["State"]:
+        if self.value == "schur":
+            return State.MACRO
+        if self.value == "macro":
+            return State.TERM
+        if self.value == "term":
+            return State.ABS
+        if self.value == "abs":
+            return State.LIN
+        if self.value == "lin":
+            return None
+        return None
+
+    @property
+    def key(self) -> Optional[str]:
+        bwd = False
+        if self.value == "schur":
+            return f"schur_decomposition_{Key.backward(bwd)}"
+        if self.value == "macro":
+            return Key.obs.macrostates(bwd)
+        if self.value == "term":
+            return Key.obs.term_states(bwd)
+        if self.value == "abs":
+            return Key.obsm.abs_probs(bwd)
+        if self.value == "lin":
+            return Key.varm.lineage_drivers(bwd)
+        return None
+
+    @property
+    def attr_keys(self) -> Optional[Sequence[Tuple[str, str]]]:
+        bwd = False
+        if self.value == "schur":
+            key1 = Key.uns.eigen(bwd)
+            key2 = Key.obsm.schur_vectors(bwd)
+            key3 = Key.uns.schur_matrix(bwd)
+            return ("uns", key1), ("obsm", key2), ("uns", key3)
+        if self.value == "macro":
+            key1 = Key.obs.macrostates(bwd)
+            key2 = Key.uns.colors(key1)
+            key3 = Key.obsm.memberships(key1)
+            return ("obs", key1), ("uns", key2), ("obsm", key3)
+        if self.value == "term":
+            key1 = Key.obs.term_states(bwd)
+            key2 = Key.obs.probs(key1)
+            key3 = Key.uns.colors(key1)
+            return ("obs", key1), ("obs", key2), ("uns", key3)
+        if self.value == "abs":
+            key1 = Key.obsm.abs_probs(bwd)
+            key2 = Key.obsm.abs_times(bwd)
+            key3 = Key.obs.priming_degree(bwd)
+            return ("obsm", key1), ("obsm", key2), ("obs", key3)
+        if self.value == "lin":
+            key1 = Key.varm.lineage_drivers(bwd)
+            return ("varm", key1),
+        return None
+
+    @property
+    def attrs(self) -> Optional[Sequence[Tuple[str, type]]]:
+        if self.value == "schur":
+            return ("_eigendecomposition", dict), ("_schur_vectors", np.ndarray), ("_schur_matrix", np.ndarray)
+        if self.value == "macro":
+            return (
+                ("_macrostates", pd.Series), ("_macrostates_colors", np.ndarray),
+                ("_macrostates_memberships", Lineage), ("_coarse_tmat", pd.DataFrame),
+                ("_coarse_init_dist", pd.Series), ("_coarse_stat_dist", pd.Series)
+            )
+        if self.value == "term":
+            return ("_term_states", pd.Series), ("_term_states_probs", pd.Series), ("_term_states_colors", np.ndarray)
+        if self.value == "abs":
+            return (
+                ("_absorption_probabilities", Lineage), ("_absorption_times", pd.DataFrame),
+                ("_priming_degree", pd.Series)
+            )
+        if self.value == "lin":
+            return ("_lineage_drivers", pd.DataFrame),
+        return None
+
+# fmt: on
+
+
+def shares_mem(x, y) -> bool:
+    if isinstance(x, Lineage):
+        return np.shares_memory(x.X, y.X)
+    if issparse(x):
+        return np.shares_memory(x.data, y.data)
+    return np.shares_memory(x, y)
 
 
 def _check_eigdecomposition(mc: cr.tl.estimators.GPCCA) -> None:
-    assert isinstance(mc._get(P.EIG), dict)
-    assert set(mc._get(P.EIG).keys()) == {
+    assert isinstance(mc.eigendecomposition, dict)
+    assert set(mc.eigendecomposition.keys()) == {
         "D",
         "eigengap",
         "params",
     }
-    assert "V_l" not in mc._get(P.EIG)
-    assert "V_r" not in mc._get(P.EIG)
-    assert "stationary_dist" not in mc._get(P.EIG)
+    assert "V_l" not in mc.eigendecomposition
+    assert "V_r" not in mc.eigendecomposition
+    assert "stationary_dist" not in mc.eigendecomposition
 
-    assert f"eig_{Direction.FORWARD}" in mc.adata.uns.keys()
+    assert Key.uns.eigen(mc.backward) in mc.adata.uns
 
 
 def _check_compute_macro(mc: cr.tl.estimators.GPCCA) -> None:
-    assert isinstance(mc._get(P.MACRO), pd.Series)
-    assert len(mc._get(A.MACRO_COLORS)) == len(mc._get(P.MACRO).cat.categories)
-
-    if "stationary_dist" in mc._get(P.EIG):  # one state
-        assert isinstance(mc._get(P.MACRO_MEMBER), cr.tl.Lineage)
-        assert mc._get(P.MACRO_MEMBER).shape[1] == 1
-        np.testing.assert_array_almost_equal(mc._get(P.MACRO_MEMBER).X.sum(), 1.0)
-
-        assert mc._get(P.COARSE_INIT_D) is None
-        assert mc._get(P.SCHUR_MAT) is None
-        assert mc._get(P.COARSE_STAT_D) is None
-        assert mc._get(P.COARSE_T) is None
-        assert mc._get(P.SCHUR) is None
+    assert isinstance(mc.macrostates, pd.Series)
+    assert len(mc._macrostates_colors) == len(mc.macrostates.cat.categories)
+
+    if "stationary_dist" in mc.eigendecomposition:  # one state
+        assert isinstance(mc.macrostates_memberships, cr.tl.Lineage)
+        assert mc.macrostates_memberships.shape[1] == 1
+        np.testing.assert_allclose(mc.macrostates_memberships.X.sum(), 1.0)
+
+        assert mc.schur_matrix is None
+        assert mc.schur_vectors is None
+        assert mc.coarse_initial_distribution is None
+        assert mc.coarse_stationary_distribution is None
+        assert mc.coarse_T is None
     else:
-        assert isinstance(mc._get(P.MACRO_MEMBER), cr.tl.Lineage)
-        if mc._get(P.MACRO_MEMBER).shape[1] > 1:
-            np.testing.assert_array_almost_equal(mc._get(P.MACRO_MEMBER).sum(1), 1.0)
-
-        assert isinstance(mc._get(P.COARSE_INIT_D), pd.Series)
-        assert isinstance(mc._get(P.SCHUR_MAT), np.ndarray)
-        assert mc._get(P.COARSE_STAT_D) is None or isinstance(
-            mc._get(P.COARSE_STAT_D), pd.Series
-        )
-        assert isinstance(mc._get(P.COARSE_T), pd.DataFrame)
-        assert isinstance(mc._get(P.SCHUR), np.ndarray)
-
-
-def _check_renaming_no_write_terminal(mc: cr.tl.estimators.GPCCA) -> None:
-    assert mc._get(P.TERM) is None
-    assert mc._get(P.TERM_PROBS) is None
-    assert mc._get(A.TERM_ABS_PROBS) is None
-
-    assert TermStatesKey.FORWARD.s not in mc.adata.obs
-    assert _probs(TermStatesKey.FORWARD.s) not in mc.adata.obs
-    assert _colors(TermStatesKey.FORWARD.s) not in mc.adata.uns
-    assert _lin_names(TermStatesKey.FORWARD.s) not in mc.adata.uns
-
-
-def _check_abs_probs(mc: cr.tl.estimators.GPCCA, has_main_states: bool = True):
-    if has_main_states:
-        assert isinstance(mc._get(P.TERM), pd.Series)
-        assert_array_nan_equal(
-            mc.adata.obs[str(TermStatesKey.FORWARD)], mc._get(P.TERM)
-        )
+        assert isinstance(mc.macrostates_memberships, cr.tl.Lineage)
+        if mc.macrostates_memberships.shape[1] > 1:
+            np.testing.assert_allclose(mc.macrostates_memberships.sum(1), 1.0)
+
+        assert isinstance(mc.schur_matrix, np.ndarray)
+        assert isinstance(mc.schur_vectors, np.ndarray)
+        assert isinstance(mc.coarse_initial_distribution, pd.Series)
+        assert isinstance(mc.coarse_T, pd.DataFrame)
+        np.testing.assert_array_equal(mc.coarse_T.index, mc.coarse_T.columns)
         np.testing.assert_array_equal(
-            mc.adata.uns[_colors(TermStatesKey.FORWARD)],
-            mc._get(A.TERM_ABS_PROBS)[list(mc._get(P.TERM).cat.categories)].colors,
+            mc.coarse_T.index, mc.coarse_stationary_distribution.index
         )
+        if mc.coarse_stationary_distribution is not None:
+            assert isinstance(mc.coarse_stationary_distribution, pd.Series)
+            np.testing.assert_array_equal(
+                mc.coarse_T.index, mc.coarse_stationary_distribution.index
+            )
 
-    assert isinstance(mc._get(P.PRIME_DEG), pd.Series)
-    assert isinstance(mc._get(P.ABS_PROBS), cr.tl.Lineage)
-    np.testing.assert_array_almost_equal(mc._get(P.ABS_PROBS).sum(1), 1.0)
-
-    np.testing.assert_array_equal(
-        mc.adata.obsm[str(AbsProbKey.FORWARD)], mc._get(P.ABS_PROBS).X
-    )
-    np.testing.assert_array_equal(
-        mc.adata.uns[_lin_names(AbsProbKey.FORWARD)], mc._get(P.ABS_PROBS).names
-    )
-    np.testing.assert_array_equal(
-        mc.adata.uns[_colors(AbsProbKey.FORWARD)], mc._get(P.ABS_PROBS).colors
-    )
-
-    np.testing.assert_array_equal(
-        mc.adata.obs[_pd(AbsProbKey.FORWARD)], mc._get(P.PRIME_DEG)
-    )
 
-    assert_array_nan_equal(mc.adata.obs[TermStatesKey.FORWARD.s], mc._get(P.TERM))
-    np.testing.assert_array_equal(
-        mc.adata.obs[_probs(TermStatesKey.FORWARD)], mc._get(P.TERM_PROBS)
+def _check_renaming_no_write_terminal(mc: cr.tl.estimators.GPCCA) -> None:
+    assert mc.terminal_states is None
+    assert mc.terminal_states_probabilities is None
+    assert mc.terminal_states_memberships is None
+
+    key = Key.obs.term_states(mc.backward)
+    assert key not in mc.adata.obs
+    assert Key.obs.probs(key) not in mc.adata.obs
+    assert Key.uns.colors(key) not in mc.adata.uns
+
+
+def _check_abs_probs(mc: cr.tl.estimators.GPCCA) -> None:
+    # fmt: off
+    # macrostates
+    assert isinstance(mc.macrostates, pd.Series)
+    assert isinstance(mc.macrostates_memberships, cr.tl.Lineage)
+    np.testing.assert_array_equal(mc._macrostates_colors, mc.macrostates_memberships.colors)
+
+    # term states
+    key = Key.obs.term_states(mc.backward)
+    assert isinstance(mc.terminal_states, pd.Series)
+    # TODO(michalk8): assert series equal from pandas
+    assert_array_nan_equal(mc.adata.obs[key], mc.terminal_states)
+    np.testing.assert_array_equal(mc.adata.uns[Key.uns.colors(key)], mc.absorption_probabilities.colors)
+    np.testing.assert_array_equal(mc.adata.uns[Key.uns.colors(key)], mc._term_states_colors)
+    assert isinstance(mc.terminal_states_probabilities, pd.Series)
+    np.testing.assert_array_equal(mc.adata.obs[Key.obs.probs(key)], mc.terminal_states_probabilities)
+
+    # abs probs
+    key = Key.obsm.abs_probs(mc.backward)
+    assert isinstance(mc.absorption_probabilities, cr.tl.Lineage)
+    np.testing.assert_array_almost_equal(mc.absorption_probabilities.sum(1), 1.0)
+    assert isinstance(mc.adata.obsm[key], cr.tl.Lineage)
+    np.testing.assert_array_equal(mc.adata.obsm[key], mc.absorption_probabilities.X)
+
+    # priming
+    key = Key.obs.priming_degree(mc.backward)
+    assert isinstance(mc.priming_degree, pd.Series)
+    np.testing.assert_array_equal(mc.adata.obs[key], mc.priming_degree)
+    # fmt: on
+
+
+def _fit_gpcca(adata, state: str, backward: bool = False) -> cr.estimators.GPCCA:
+    vk = VelocityKernel(adata, backward=backward).compute_transition_matrix(
+        softmax_scale=4
     )
+    ck = ConnectivityKernel(adata, backward=backward).compute_transition_matrix()
+    terminal_kernel = 0.8 * vk + 0.2 * ck
+
+    mc = cr.tl.estimators.GPCCA(terminal_kernel)
+    mc.compute_eigendecomposition()
+    mc.compute_schur(n_components=10, method="krylov")
+    if state == State.SCHUR:
+        return mc
+    mc.compute_macrostates(n_states=2)
+    if state == State.MACRO:
+        return mc
+    mc.set_terminal_states_from_macrostates()
+    if state == State.TERM:
+        return mc
+    mc.compute_absorption_probabilities(time_to_absorption="all")
+    mc.compute_lineage_priming()
+    if state == State.ABS:
+        return mc
+    mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
+    if state == State.LIN:
+        return mc
+
+    raise NotImplementedError(state)
+
+
+def _assert_params(
+    g: cr.tl.estimators.GPCCA,
+    state: Optional[State],
+    fwd: bool = False,
+    init: bool = True,
+) -> None:
+    if state is None:
+        return
+
+    key = state.key
+    if fwd:
+        if init:
+            assert isinstance(g.params[key], dict), state
+        else:
+            assert g.params.get(key, {}) == {}, state
+        _assert_params(g, state.next, fwd=True, init=False)
+    else:
+        assert isinstance(g.params[key], dict), state
+        _assert_params(g, state.prev, fwd=False)
 
 
-class TestGPCCA:
-    def test_compute_partition(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_partition()
+def _assert_adata(
+    adata: AnnData, state: Optional[State], fwd: bool = False, init: bool = True
+) -> None:
+    if state is None:
+        return
 
-        assert isinstance(mc.is_irreducible, bool)
-        if not mc.is_irreducible:
-            assert isinstance(mc.recurrent_classes, list)
-            assert isinstance(mc.transient_classes, list)
-            assert f"{TermStatesKey.FORWARD}_rec_classes" in mc.adata.obs
-            assert f"{TermStatesKey.FORWARD}_trans_classes" in mc.adata.obs
+    if fwd:
+        if init:
+            for attr, key in state.attr_keys:
+                obj = getattr(adata, attr)
+                assert key in obj, sorted(obj.keys())
+        else:
+            for attr, key in state.attr_keys:
+                obj = getattr(adata, attr)
+                assert key not in obj, sorted(obj.keys())
+        _assert_adata(adata, state.next, fwd=True, init=False)
+    else:
+        for attr, key in state.attr_keys:
+            obj = getattr(adata, attr)
+            assert key in obj, sorted(obj.keys())
+        _assert_adata(adata, state.prev, fwd=False)
+
+
+def _assert_gpcca_attrs(
+    g: cr.tl.estimators.GPCCA,
+    state: Optional[State] = None,
+    fwd: bool = False,
+    init: bool = True,
+) -> None:
+    if state is None:
+        return
+    if fwd:
+        if init:
+            for attr, dtype in state.attrs:
+                obj = getattr(g, attr)
+                assert isinstance(obj, dtype)
         else:
-            assert mc.recurrent_classes is None
-            assert mc.transient_classes is None
+            for attr, dtype in state.attrs:
+                obj = getattr(g, attr)
+                assert obj is None, attr
+        _assert_gpcca_attrs(g, state.next, fwd=True, init=False)
+    else:
+        for attr, dtype in state.attrs:
+            obj = getattr(g, attr)
+            assert isinstance(obj, dtype), attr
+        _assert_gpcca_attrs(g, state.prev, fwd=False)
 
+
+class TestGPCCA:
     def test_compute_eigendecomposition(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -149,8 +334,9 @@ def test_compute_schur_invalid_n_comps(self, adata_large: AnnData):
         terminal_kernel = 0.8 * vk + 0.2 * ck
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
-        with pytest.raises(ValueError):
-            mc.compute_schur(n_components=1, method="krylov")
+        mc.compute_schur(n_components=1, method="krylov")
+
+        assert mc.schur_vectors.shape[1] == 2
 
     def test_compute_schur_invalid_method(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -191,13 +377,13 @@ def test_compute_schur_write_eigvals_similar_to_orig_eigdecomp(
         mc.compute_eigendecomposition(k=10, only_evals=True)
 
         _check_eigdecomposition(mc)
-        orig_ed = deepcopy(mc._get(P.EIG))
+        orig_ed = deepcopy(mc.eigendecomposition)
 
-        mc._set(A.EIG, None)
+        mc._eigendecomposition = None
         mc.compute_schur(n_components=10, method="krylov")
 
         _check_eigdecomposition(mc)
-        schur_ed = mc._get(P.EIG)
+        schur_ed = mc.eigendecomposition
 
         assert orig_ed.keys() == schur_ed.keys()
         assert orig_ed["eigengap"] == schur_ed["eigengap"]
@@ -255,8 +441,9 @@ def test_compute_macrostates_min_chi_too_low_min(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
-        with pytest.raises(ValueError):
-            mc.compute_macrostates(n_states=[1, 4], use_min_chi=True)
+        mc.compute_macrostates(n_states=[1, 1])
+
+        _check_compute_macro(mc)
 
     def test_compute_macrostates_min_chi_inverted_range(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -265,7 +452,7 @@ def test_compute_macrostates_min_chi_inverted_range(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
-        mc.compute_macrostates(n_states=[4, 2], use_min_chi=True)
+        mc.compute_macrostates(n_states=[4, 2])
 
         _check_compute_macro(mc)
 
@@ -276,20 +463,10 @@ def test_compute_macrostates_min_chi_range_same_values(self, adata_large: AnnDat
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
-        mc.compute_macrostates(n_states=[2, 2], use_min_chi=True)
+        mc.compute_macrostates(n_states=[2, 2])
 
         _check_compute_macro(mc)
 
-    def test_compute_macrostates_min_chi_dict_wrong_keys(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")
-        with pytest.raises(KeyError):
-            mc.compute_macrostates(n_states={"foo": 2, "max": 3}, use_min_chi=True)
-
     def test_compute_macrostates_min_chi_normal_run(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -297,7 +474,7 @@ def test_compute_macrostates_min_chi_normal_run(self, adata_large: AnnData):
 
         mc = cr.tl.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
-        mc.compute_macrostates(n_states=[2, 4], use_min_chi=True)
+        mc.compute_macrostates(n_states=[2, 4])
 
         _check_compute_macro(mc)
 
@@ -321,8 +498,8 @@ def test_compute_macro_cache(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2)
 
-        assert mc._get(P.SCHUR).shape[1] == 11
-        assert mc._get(P.SCHUR_MAT).shape == (11, 11)
+        assert mc.schur_vectors.shape[1] == 11
+        np.testing.assert_array_equal(mc.schur_matrix.shape, [11, 11])
 
     def test_set_initial_states_from_forward(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
@@ -336,15 +513,15 @@ def test_set_initial_states_from_forward(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2, n_cells=5)
         obsm_keys = set(mc.adata.obsm.keys())
+        # TODO
         mc._set_initial_states_from_macrostates("0")
 
-        key = TermStatesKey.BACKWARD.s
+        key = Key.obs.term_states(not mc.backward)
 
         assert key in mc.adata.obs
         np.testing.assert_array_equal(mc.adata.obs[key].cat.categories, ["0"])
-        assert _probs(key) in mc.adata.obs
-        assert _colors(key) in mc.adata.uns
-        assert _lin_names(key) in mc.adata.uns
+        assert Key.obs.probs(key) in mc.adata.obs
+        assert Key.uns.colors(key) in mc.adata.uns
 
         # make sure that we don't write anything there - it's useless
         assert set(mc.adata.obsm.keys()) == obsm_keys
@@ -360,6 +537,7 @@ def test_compute_initial_states_from_forward_no_macro(self, adata_large: AnnData
         mc.compute_schur(n_components=10, method="krylov")
 
         with pytest.raises(RuntimeError):
+            # TODO
             mc._compute_initial_states(1)
 
     def test_compute_initial_states_from_forward_too_many_states(
@@ -424,17 +602,17 @@ def test_compute_initial_states_from_forward_normal_run(self, adata_large: AnnDa
 
         mc.compute_macrostates(n_states=2, n_cells=5)
         obsm_keys = set(mc.adata.obsm.keys())
-        expected = mc._get(P.COARSE_STAT_D).index[np.argmin(mc._get(P.COARSE_STAT_D))]
+        csd = mc.coarse_stationary_distribution
+        expected = csd.index[np.argmin(csd)]
 
         mc._compute_initial_states(1)
 
-        key = TermStatesKey.BACKWARD.s
+        key = Key.obs.term_states(not mc.backward)
 
         assert key in mc.adata.obs
         np.testing.assert_array_equal(mc.adata.obs[key].cat.categories, [expected])
-        assert _probs(key) in mc.adata.obs
-        assert _colors(key) in mc.adata.uns
-        assert _lin_names(key) in mc.adata.uns
+        assert Key.obs.probs(key) in mc.adata.obs
+        assert Key.uns.colors(key) in mc.adata.uns
 
         # make sure that we don't write anything there - it's useless
         assert set(mc.adata.obsm.keys()) == obsm_keys
@@ -516,8 +694,8 @@ def test_set_terminal_states_from_macrostates_rename_states(
 
         _check_abs_probs(mc)
 
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, ["foo"])
-        np.testing.assert_array_equal(mc._get(A.TERM_ABS_PROBS).names, ["foo"])
+        np.testing.assert_array_equal(mc.terminal_states.cat.categories, ["foo"])
+        np.testing.assert_array_equal(mc.terminal_states_memberships.names, ["foo"])
 
     def test_set_terminal_states_from_macrostates_join_and_rename_states(
         self, adata_large: AnnData
@@ -536,8 +714,10 @@ def test_set_terminal_states_from_macrostates_join_and_rename_states(
 
         _check_abs_probs(mc)
 
-        np.testing.assert_array_equal(mc._get(P.TERM).cat.categories, ["foo", "bar"])
-        np.testing.assert_array_equal(mc._get(A.TERM_ABS_PROBS).names, ["foo", "bar"])
+        np.testing.assert_array_equal(mc.terminal_states.cat.categories, ["foo", "bar"])
+        np.testing.assert_array_equal(
+            mc.terminal_states_memberships.names, ["foo", "bar"]
+        )
 
     def test_set_terminal_states_from_macrostates_no_cells(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -682,11 +862,13 @@ def test_compute_terminal_states_stability(self, adata_large: AnnData):
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
-        coarse_T = mc._get(P.COARSE_T)
+        coarse_T = mc.coarse_T
         self_probs = pd.Series(np.diag(coarse_T), index=coarse_T.columns)
         names = self_probs[self_probs.values >= thresh].index
 
-        np.testing.assert_array_equal(set(names), set(mc._get(P.TERM).cat.categories))
+        np.testing.assert_array_equal(
+            set(names), set(mc.terminal_states.cat.categories)
+        )
         _check_abs_probs(mc)
 
     def test_compute_terminal_states_no_selected(self, adata_large: AnnData):
@@ -730,62 +912,6 @@ def test_compute_terminal_states_default(self, adata_large: AnnData):
 
         _check_abs_probs(mc)
 
-    def test_compute_gdpt_no_schur(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_gdpt(method="krylov")
-
-        assert "gdpt_pseudotime" in mc.adata.obs
-
-    def test_compute_gdpt_no_iroot(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.adata.uns.pop("iroot", None)
-
-        with pytest.raises(KeyError):
-            mc.compute_gdpt()
-
-    def test_compute_gdpt_invalid_n_comps(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)
-
-        with pytest.raises(ValueError):
-            mc.compute_gdpt(n_components=1)
-
-    def test_compute_gdpt_cellname_key_added(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_gdpt(key_added="foobar")
-
-        assert "foobar" in mc.adata.obs
-
-    def test_compute_gdpt_cellname_iroot(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.adata.uns["iroot"] = mc.adata.obs_names[0]
-
-        mc.compute_gdpt()
-
-        assert "gdpt_pseudotime" in mc.adata.obs
-
     def test_compute_lineage_drivers_invalid_lineages(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -828,12 +954,13 @@ def test_compute_lineage_drivers_normal_run(self, adata_large: AnnData):
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
+        key = Key.varm.lineage_drivers(False)
         for lineage in ["0", "1"]:
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} corr"] >= -1.0)
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} corr"] <= 1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_corr"] >= -1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_corr"] <= 1.0)
 
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} qval"] >= 0)
-            assert np.all(mc.adata.var[f"{DirPrefix.FORWARD} {lineage} qval"] <= 1.0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_qval"] >= 0)
+            assert np.all(mc.adata.varm[key][f"{lineage}_qval"] <= 1.0)
 
     def test_plot_lineage_drivers_not_computed(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -915,10 +1042,9 @@ def test_compute_priming_clusters(self, adata_large: AnnData):
 
         assert_series_equal(deg1, deg2)
         # because passing it to a dataframe changes its name
-        assert_series_equal(
-            adata_large.obs[_pd(mc._abs_prob_key)], deg1, check_names=False
-        )
-        assert_series_equal(mc._get(A.PRIME_DEG), deg1)
+        key = Key.obs.priming_degree(mc.backward)
+        assert_series_equal(adata_large.obs[key], deg1, check_names=False)
+        assert_series_equal(mc.priming_degree, deg1)
 
     def test_recompute_terminal_states_different_n_states(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -930,59 +1056,236 @@ def test_recompute_terminal_states_different_n_states(self, adata_large: AnnData
         mc.compute_macrostates(n_states=2)
         mc.set_terminal_states_from_macrostates()
 
+        key = Key.obsm.memberships(Key.obs.term_states(mc.backward))
         assert len(mc.terminal_states.cat.categories) == 2
-        assert adata_large.obsm[mc._term_abs_prob_key].shape == (adata_large.n_obs, 2)
+        assert adata_large.obsm[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
+        assert mc.terminal_states_memberships is None
+        assert key not in adata_large.obsm
+
+    def test_fit_predict(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
+
+        g = cr.tl.estimators.GPCCA(terminal_kernel)
+        tmp = g.fit(n_states=2)
+
+        assert tmp is g
+        _assert_params(g, State.MACRO, fwd=False)
+        _assert_params(g, State.MACRO, fwd=True)
+
+        res = g.predict()
+        assert res is None
+        _assert_params(g, State.TERM, fwd=False)
+        _assert_params(g, State.TERM, fwd=True)
+
+    @pytest.mark.parametrize("state", list(State))
+    def test_params(self, adata_large: AnnData, state: State):
+        g = _fit_gpcca(adata_large, state)
+        _assert_params(g, state, fwd=False)
+        _assert_params(g, state, fwd=True)
+
+
+class TestGPCCASerialization:
+    @pytest.mark.parametrize("state", list(State))
+    def test_to_adata(self, adata_large: AnnData, state: State):
+        g = _fit_gpcca(adata_large, state)
+        adata = g.to_adata()
+        _assert_adata(adata, state, fwd=False)
+        _assert_adata(adata, state, fwd=True)
+
+    @pytest.mark.parametrize("copy", [False, True])
+    @pytest.mark.parametrize("keep", ["X", ("obs", "obsm"), ("layers")])
+    def test_to_adata_keep(
+        self, adata_large: AnnData, keep: Union[str, Sequence[str]], copy: bool
+    ):
+        g = _fit_gpcca(adata_large, State.MACRO)
+        adata = g.to_adata(keep=keep, copy=copy)
+        if "X" in keep:
+            res = shares_mem(adata.X, g.adata.X)
+            assert not res if copy else res
+        if "obs" in keep:
+            # we don't save macrostates in the underlying object, only during serialization
+            key = Key.obs.macrostates(g.backward)
+            cols = [c for c in adata.obs.columns if c != key]
+            obs = adata.obs[cols]
+            np.testing.assert_array_equal(obs.shape, g.adata.obs.shape)
+            assert_frame_equal(obs, g.adata.obs[cols])
+        if "obsm" in keep:
+            for key in g.adata.obsm.keys():
+                res = shares_mem(adata.obsm[key], g.adata.obsm[key])
+                assert not res if copy else res, key
+        if "layers" in keep:
+            for key in g.adata.layers.keys():
+                res = shares_mem(adata.layers[key], g.adata.layers[key])
+                assert not res if copy else res, key
+
+    @pytest.mark.parametrize("copy", [False, True])
+    def test_to_adata_copy_finegrained(self, adata_large: AnnData, copy):
+        g = _fit_gpcca(adata_large, State.MACRO)
+        adata = g.to_adata(keep=["obsp", "varm", "var"], copy=["obsp"])
+
+        assert_frame_equal(adata.var, g.adata.var)
+        for key in g.adata.obsp.keys():
+            assert not shares_mem(adata.obsp[key], g.adata.obsp[key]), key
+        for key in g.adata.varm.keys():
+            assert shares_mem(adata.varm[key], g.adata.varm[key]), key
+
+    @pytest.mark.parametrize("state", list(State))
+    def test_from_adata(self, adata_large: AnnData, state: State):
+        g1 = _fit_gpcca(adata_large, state)
+        adata = g1.to_adata()
+        g2 = cr.tl.estimators.GPCCA.from_adata(adata, obsp_key="T_fwd")
+
+        _assert_adata(g2.adata, state, fwd=False)
+        _assert_adata(g2.adata, state, fwd=True)
+        assert_estimators_equal(g1, g2, from_adata=True)
+
+    @pytest.mark.parametrize("state", list(State))
+    def test_from_adata_incomplete(self, adata_large: AnnData, state: State):
+        if state in (State.SCHUR, State.MACRO):
+            pytest.skip("See the reintroduce TODO in GPCCA")
+        g_orig = _fit_gpcca(adata_large, State.LIN)
+        adata = g_orig.to_adata()
+
+        for attr, key in state.attr_keys:
+            obj = getattr(adata, attr)
+            del obj[key]
+
+        g = cr.tl.estimators.GPCCA.from_adata(adata, obsp_key="T_fwd")
+        _assert_gpcca_attrs(g, state, fwd=True, init=False)
+        _assert_gpcca_attrs(g, state.prev, fwd=False)
+        _assert_adata(g._shadow_adata, state, fwd=True, init=False)
+        _assert_adata(g._shadow_adata, state.prev, fwd=False)
+
+    @pytest.mark.parametrize("bwd", [False, True])
+    def test_reconstruct_lineage(self, adata_large: AnnData, bwd: bool):
+        g_orig = _fit_gpcca(adata_large, State.LIN, backward=bwd)
+        adata = g_orig.to_adata()
+        bdata = adata.copy()
+        keys = [
+            Key.obsm.memberships(Key.obs.macrostates(bwd)),
+            Key.obsm.memberships(Key.obs.term_states(bwd)),
+            Key.obsm.abs_probs(bwd),
+        ]
+        attrs = [
+            "macrostates_memberships",
+            "terminal_states_memberships",
+            "absorption_probabilities",
+        ]
+
+        for key in keys:
+            assert isinstance(adata.obsm[key], Lineage)
+            adata.obsm[key] = np.array(adata.obsm[key])
+            assert isinstance(adata.obsm[key], np.ndarray)
+
+        g = cr.tl.estimators.GPCCA.from_adata(adata, obsp_key=Key.uns.kernel(bwd))
+        for attr, key in zip(attrs, keys):
+            obj = getattr(g, attr)
+            assert isinstance(obj, Lineage)
+            np.testing.assert_allclose(obj.X, bdata.obsm[key].X)
+            np.testing.assert_array_equal(obj.names, bdata.obsm[key].names)
+            np.testing.assert_array_equal(obj.colors, bdata.obsm[key].colors)
+
+        assert_estimators_equal(g_orig, g, from_adata=True)
 
 
 class TestGPCCAIO:
-    def test_copy(self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]):
+    @pytest.mark.parametrize("deep", [False, True])
+    def test_copy(
+        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA], deep: bool
+    ):
         _, mc1 = adata_gpcca_fwd
-        mc2 = mc1.copy()
+        mc2 = mc1.copy(deep=deep)
 
-        assert_estimators_equal(mc1, mc2, copy=True)
+        assert_estimators_equal(mc1, mc2, copy=True, deep=deep)
 
-    def test_write_ext(self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]):
+    def test_write_ext(
+        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA], tmpdir
+    ):
         _, mc = adata_gpcca_fwd
 
-        with TemporaryDirectory() as tmpdir:
-            fname = "foo"
-            mc.write(os.path.join(tmpdir, fname), ext="bar")
+        fname = "foo"
+        mc.write(os.path.join(tmpdir, fname), ext="bar")
 
-            assert os.path.isfile(os.path.join(tmpdir, f"foo.bar"))
+        assert os.path.isfile(os.path.join(tmpdir, f"foo.bar"))
 
     def test_write_no_ext(
-        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]
+        self,
+        adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA],
+        tmpdir,
     ):
         _, mc = adata_gpcca_fwd
+        fname = "foo"
+        mc.write(os.path.join(tmpdir, fname), ext=None)
 
-        with TemporaryDirectory() as tmpdir:
-            fname = "foo"
-            mc.write(os.path.join(tmpdir, fname), ext=None)
-
-            assert os.path.isfile(os.path.join(tmpdir, f"foo"))
+        assert os.path.isfile(os.path.join(tmpdir, f"foo"))
 
     def test_write_ext_with_dot(
-        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]
+        self,
+        adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA],
+        tmpdir,
     ):
         _, mc = adata_gpcca_fwd
 
-        with TemporaryDirectory() as tmpdir:
-            fname = "foo"
-            mc.write(os.path.join(tmpdir, fname), ext=".bar")
+        fname = "foo"
+        mc.write(os.path.join(tmpdir, fname), ext=".bar")
 
-            assert os.path.isfile(os.path.join(tmpdir, f"foo.bar"))
+        assert os.path.isfile(os.path.join(tmpdir, f"foo.bar"))
 
-    def test_read(self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA]):
+    def test_read(
+        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA], tmpdir
+    ):
         _, mc1 = adata_gpcca_fwd
 
-        with TemporaryDirectory() as tmpdir:
-            mc1.write(os.path.join(tmpdir, "foo"))
-            mc2 = cr.tl.estimators.GPCCA.read(os.path.join(tmpdir, "foo.pickle"))
+        mc1.write(os.path.join(tmpdir, "foo"))
+        mc2 = cr.tl.estimators.GPCCA.read(os.path.join(tmpdir, "foo.pickle"))
+
+        assert_estimators_equal(mc1, mc2)
+
+    @pytest.mark.parametrize("copy", [False, True])
+    def test_write_no_adata(
+        self,
+        adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA],
+        copy: bool,
+        tmpdir,
+    ):
+        adata, mc1 = adata_gpcca_fwd
+
+        mc1.write(os.path.join(tmpdir, "foo"), write_adata=False)
+        mc2 = cr.tl.estimators.GPCCA.read(
+            os.path.join(tmpdir, "foo.pickle"), adata=adata, copy=copy
+        )
 
+        if copy:
+            assert adata is not mc2.adata
+        else:
+            assert adata is mc2.adata
         assert_estimators_equal(mc1, mc2)
+
+    def test_write_no_adata_read_none_supplied(
+        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA], tmpdir
+    ):
+        _, mc1 = adata_gpcca_fwd
+
+        mc1.write(os.path.join(tmpdir, "foo"), write_adata=False)
+        with pytest.raises(TypeError, match="This object was saved without"):
+            _ = cr.tl.estimators.GPCCA.read(
+                os.path.join(tmpdir, "foo.pickle"), adata=None
+            )
+
+    def test_write_no_adata_read_wrong_length(
+        self, adata_gpcca_fwd: Tuple[AnnData, cr.tl.estimators.GPCCA], tmpdir
+    ):
+        _, mc1 = adata_gpcca_fwd
+        adata = AnnData(np.random.normal(size=(len(mc1) + 1, 1)))
+
+        mc1.write(os.path.join(tmpdir, "foo"), write_adata=False)
+        with pytest.raises(ValueError, match="Expected `adata` to be of length"):
+            _ = cr.tl.estimators.GPCCA.read(
+                os.path.join(tmpdir, "foo.pickle"), adata=adata
+            )
diff --git a/tests/test_kernels.py b/tests/test_kernels.py
index 78c1cda0d..144f83e91 100644
--- a/tests/test_kernels.py
+++ b/tests/test_kernels.py
@@ -1,7 +1,9 @@
 from typing import Tuple, Callable, Optional
 
+import pickle
 import pytest
 from copy import copy
+from pathlib import Path
 from _helpers import (
     bias_knn,
     create_kernels,
@@ -14,6 +16,7 @@
 import cellrank as cr
 from scanpy import Neighbors
 from anndata import AnnData
+from cellrank._key import Key
 from cellrank.tl._utils import _normalize
 from cellrank.ul._utils import _get_neighs, _get_neighs_params
 from cellrank.tl.kernels import (
@@ -23,7 +26,6 @@
     PrecomputedKernel,
     ConnectivityKernel,
 )
-from cellrank.tl._constants import _transition
 from cellrank.tl.kernels._base_kernel import (
     Kernel,
     Constant,
@@ -32,7 +34,7 @@
     _dtype,
     _is_bin_mult,
 )
-from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEAggregation, _ct
+from cellrank.tl.kernels._cytotrace_kernel import CytoTRACEAggregation
 
 import numpy as np
 from scipy.sparse import eye as speye
@@ -201,7 +203,7 @@ def test_negative_constant(self, adata: AnnData):
 
     def test_invalid_constant(self, adata: AnnData):
         with pytest.raises(TypeError):
-            _ = Constant(None, None)
+            _ = Constant(adata, None)
 
     def test_inversion(self, adata: AnnData):
         c = ConnectivityKernel(adata, backward=False)
@@ -815,10 +817,10 @@ def test_addition_3_kernels(self, adata: AnnData):
         np.testing.assert_allclose(k.transition_matrix.A, expected)
 
     def test_addition_adaptive(self, adata: AnnData):
-        adata.uns["velocity_variances"] = vv = np.random.random(
+        adata.obsp["velocity_variances"] = vv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
-        adata.uns["connectivity_variances"] = cv = np.random.random(
+        adata.obsp["connectivity_variances"] = cv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
         vk, ck = create_kernels(
@@ -837,10 +839,10 @@ def test_addition_adaptive(self, adata: AnnData):
     def test_addition_adaptive_constants(self, adata: AnnData):
         a, b = np.random.uniform(0, 10, 2)
         s = a + b
-        adata.uns["velocity_variances"] = vv = np.random.random(
+        adata.obsp["velocity_variances"] = vv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
-        adata.uns["connectivity_variances"] = cv = np.random.random(
+        adata.obsp["connectivity_variances"] = cv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
         vk, ck = create_kernels(
@@ -859,10 +861,10 @@ def test_addition_adaptive_constants(self, adata: AnnData):
     def test_addition_adaptive_wrong_variances(self, adata: AnnData):
         a, b = np.random.uniform(0, 10, 2)
         s = a + b
-        adata.uns["velocity_variances"] = np.random.random(
+        adata.obsp["velocity_variances"] = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
-        adata.uns["connectivity_variances"] = np.random.random(
+        adata.obsp["connectivity_variances"] = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
         vk, ck = create_kernels(
@@ -881,10 +883,10 @@ def test_addition_adaptive_wrong_variances(self, adata: AnnData):
     def test_addition_adaptive_4_kernels(self, adata: AnnData):
         a, b, c, d = np.random.uniform(0, 10, 4)
         s = a + b + c + d
-        adata.uns["velocity_variances"] = vv = np.random.random(
+        adata.obsp["velocity_variances"] = vv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
-        adata.uns["connectivity_variances"] = cv = np.random.random(
+        adata.obsp["connectivity_variances"] = cv = np.random.random(
             size=(adata.n_obs, adata.n_obs)
         )
         vk, ck = create_kernels(
@@ -1319,7 +1321,7 @@ def test_write_to_adata(self, adata: AnnData, write_first: bool):
             ck.compute_projection(basis="umap")
             ck.write_to_adata()
 
-        assert adata.uns[_transition(ck._direction) + "_params"] == {
+        assert adata.uns[Key.uns.kernel(ck.backward) + "_params"] == {
             "params": ck.params,
             "embeddings": ["umap"],
         }
@@ -1329,7 +1331,7 @@ def test_key_added(self, adata: AnnData, key_added: Optional[str]):
         ck = cr.tl.kernels.ConnectivityKernel(adata).compute_transition_matrix()
         ck.compute_projection(basis="umap", copy=False, key_added=key_added)
 
-        key = key_added if key_added is not None else _transition(ck._direction)
+        key = Key.uns.kernel(ck.backward, key=key_added)
         ukey = f"{key}_params"
         key = f"{key}_umap"
 
@@ -1346,7 +1348,7 @@ def test_copy(self, adata: AnnData, copy: bool):
             np.testing.assert_array_equal(res.shape, adata.obsm["X_umap"].shape)
         else:
             assert res is None
-            key = _transition(ck._direction) + "_umap"
+            key = Key.uns.kernel(ck.backward) + "_umap"
             np.testing.assert_array_equal(
                 adata.obsm[key].shape, adata.obsm["X_umap"].shape
             )
@@ -1415,44 +1417,46 @@ def test_layer(self, adata: AnnData, layer: str):
                 _ = CytoTRACEKernel(adata, layer=layer)
         else:
             _ = CytoTRACEKernel(adata, layer=layer)
-            assert adata.uns[_ct("params")]["layer"] == layer
+            assert adata.uns[Key.cytotrace("params")]["layer"] == layer
 
     @pytest.mark.parametrize("agg", list(CytoTRACEAggregation))
     def test_aggregation(self, adata: AnnData, agg: CytoTRACEAggregation):
         _ = CytoTRACEKernel(adata, aggregation=agg)
-        assert adata.uns[_ct("params")]["aggregation"] == agg.s
+        assert adata.uns[Key.cytotrace("params")]["aggregation"] == agg
 
     @pytest.mark.parametrize("use_raw", [False, True])
     def test_raw(self, adata: AnnData, use_raw: bool):
         _ = CytoTRACEKernel(adata, use_raw=use_raw)
-        assert adata.uns[_ct("params")]["use_raw"] == (
+        assert adata.uns[Key.cytotrace("params")]["use_raw"] == (
             adata.raw.n_vars == adata.n_vars if use_raw else False
         )
 
     def test_correct_class(self, adata: AnnData):
         k = CytoTRACEKernel(adata)
         assert isinstance(k, PseudotimeKernel)
-        assert k._time_key == _ct("pseudotime")
+        assert k._time_key == Key.cytotrace("pseudotime")
 
     def test_writes_params(self, adata: AnnData):
         k = CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
 
-        assert adata.uns[_ct("params")] == {
+        assert adata.uns[Key.cytotrace("params")] == {
             "layer": "X",
             "aggregation": "mean",
             "use_raw": False,
         }
 
-        assert np.all(adata.var[_ct("gene_corr")] <= 1.0)
-        assert np.all(-1 <= adata.var[_ct("gene_corr")])
-        assert is_bool_dtype(adata.var[_ct("correlates")])
-        assert adata.var[_ct("correlates")].sum() == min(200, adata.n_vars)
+        assert np.all(adata.var[Key.cytotrace("gene_corr")] <= 1.0)
+        assert np.all(-1 <= adata.var[Key.cytotrace("gene_corr")])
+        assert is_bool_dtype(adata.var[Key.cytotrace("correlates")])
+        assert adata.var[Key.cytotrace("correlates")].sum() == min(200, adata.n_vars)
 
-        assert _ct("score") in adata.obs
-        assert _ct("pseudotime") in adata.obs
-        assert _ct("num_exp_genes") in adata.obs
-        assert is_integer_dtype(adata.obs[_ct("num_exp_genes")])
-        np.testing.assert_array_equal(k.pseudotime, adata.obs[_ct("pseudotime")].values)
+        assert Key.cytotrace("score") in adata.obs
+        assert Key.cytotrace("pseudotime") in adata.obs
+        assert Key.cytotrace("num_exp_genes") in adata.obs
+        assert is_integer_dtype(adata.obs[Key.cytotrace("num_exp_genes")])
+        np.testing.assert_array_equal(
+            k.pseudotime, adata.obs[Key.cytotrace("pseudotime")].values
+        )
         np.testing.assert_array_equal(k.pseudotime.min(), 0.0)
         np.testing.assert_array_equal(k.pseudotime.max(), 1.0)
 
@@ -1465,7 +1469,7 @@ def test_compute_transition_matrix(self, adata: AnnData):
     def test_inversion(self, adata: AnnData):
         k = ~CytoTRACEKernel(adata, use_raw=False, layer="X", aggregation="mean")
 
-        pt = adata.obs[_ct("pseudotime")].values
+        pt = adata.obs[Key.cytotrace("pseudotime")].values
         np.testing.assert_array_equal(np.max(pt) - pt, k.pseudotime)
 
 
@@ -1526,3 +1530,27 @@ def test_remove_empty_clusters_none_remain(self, kernel: Kernel):
                 min_flow=np.inf,
                 remove_empty_clusters=True,
             )
+
+
+class TestKernelIO:
+    @pytest.mark.parametrize("copy", [False, True])
+    @pytest.mark.parametrize("write_adata", [False, True])
+    def test_read_write(self, kernel: Kernel, tmpdir, write_adata: bool, copy: bool):
+        path = Path(tmpdir) / "kernel.pickle"
+        kernel.write(path, write_adata=write_adata)
+
+        if write_adata:
+            k: Kernel = type(kernel).read(path)
+            assert k.adata is not None
+        else:
+            with open(path, "rb") as fin:
+                k: Kernel = pickle.load(fin)
+                assert k.adata is None
+                assert k.shape == (kernel.adata.n_obs, kernel.adata.n_obs)
+            k: Kernel = type(kernel).read(path, adata=kernel.adata, copy=copy)
+            if copy:
+                assert k.adata is not kernel.adata
+            else:
+                assert k.adata is kernel.adata
+
+        np.testing.assert_array_equal(k.transition_matrix.A, kernel.transition_matrix.A)
diff --git a/tests/test_lineage.py b/tests/test_lineage.py
index b4600cc40..d01a690b3 100644
--- a/tests/test_lineage.py
+++ b/tests/test_lineage.py
@@ -5,10 +5,10 @@
 from collections import defaultdict
 from html.parser import HTMLParser
 
+from cellrank import Lin
 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
@@ -193,7 +193,7 @@ def test_non_unique_names_conversion(self):
 class TestLineageAccessor:
     def test_too_large_tuple(self, lineage: Lineage):
         with pytest.raises(ValueError):
-            lineage[0, 0, 0]
+            _ = lineage[0, 0, 0]
 
     def test_none(self, lineage: Lineage):
         y = lineage[None, None]
@@ -301,7 +301,7 @@ def test_column_invalid_name(self):
         )
 
         with pytest.raises(KeyError):
-            l["quux"]
+            _ = l["quux"]
 
     def test_row_subset_with_ints(self):
         x = np.random.random((10, 3))
@@ -336,7 +336,7 @@ def test_column_subset_boolean_invalid_dim(self):
         )
 
         with pytest.raises(IndexError):
-            y = l[:, [True]]
+            _ = l[:, [True]]
 
     def test_row_subset_with_mask(self):
         x = np.random.random((10, 3))
@@ -917,7 +917,7 @@ def test_copy(self, lineage: Lineage):
     def test_simple_access(self, lineage: Lineage):
         y = lineage.T["foo"]
         with pytest.raises(TypeError):
-            lineage.T[:, "foo"]
+            _ = lineage.T[:, "foo"]
 
         assert y.shape == (1, lineage.shape[0])
         np.testing.assert_array_equal(y.T, lineage["foo"])
@@ -1106,10 +1106,10 @@ def test_pickle_transposed(self, lineage: Lineage):
 class TestPriming:
     def test_invalid_method(self, lineage: Lineage):
         with pytest.raises(ValueError, match="foobar"):
-            lineage.priming_degree("foobar")
+            _ = lineage.priming_degree("foobar")
 
     @pytest.mark.parametrize("method", list(PrimingDegree))
-    def test_priming_degree(self, lineage: Lineage, method: str):
+    def test_priming_degree(self, lineage: Lineage, method: PrimingDegree):
         deg = lineage.priming_degree(method=method)
 
         assert isinstance(deg, np.ndarray)
@@ -1120,26 +1120,12 @@ def test_priming_degree(self, lineage: Lineage, method: str):
 
     def test_early_cells_empty(self, lineage: Lineage):
         with pytest.raises(ValueError, match="No early cells have been specified."):
-            mask = np.zeros(
-                (
-                    len(
-                        lineage,
-                    )
-                ),
-                dtype=np.bool_,
-            )
+            mask = np.zeros(shape=(len(lineage),), dtype=bool)
             lineage.priming_degree("kl_divergence", early_cells=mask)
 
     def test_early_cells(self, lineage: Lineage):
         deg1 = lineage.priming_degree("kl_divergence", early_cells=[0, 1, 2])
-        mask = np.zeros(
-            (
-                len(
-                    lineage,
-                )
-            ),
-            dtype=np.bool_,
-        )
+        mask = np.zeros(shape=(len(lineage),), dtype=bool)
         mask[:3] = True
         deg2 = lineage.priming_degree("kl_divergence", early_cells=mask)
 
diff --git a/tests/test_linear_solver.py b/tests/test_linear_solver.py
index a6fd97e7e..85f77fea9 100644
--- a/tests/test_linear_solver.py
+++ b/tests/test_linear_solver.py
@@ -1,8 +1,8 @@
 import pytest
 
 from cellrank.tl._linear_solver import (
-    _petsc_mat_solve,
     _solve_lin_system,
+    _petsc_direct_solve,
     _create_petsc_matrix,
 )
 
@@ -51,7 +51,7 @@ def test_invalid_solver(self):
         with pytest.raises(ValueError):
             _solve_lin_system(A, B, solver="foobar", use_petsc=False)
 
-    @pytest.mark.parametrize("seed, sparse", zip(range(10), [False] * 5 + [True] * 5))
+    @pytest.mark.parametrize("seed,sparse", zip(range(10), [False] * 5 + [True] * 5))
     def test_gmres(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
 
@@ -64,6 +64,7 @@ def test_gmres(self, seed: int, sparse: bool):
             show_progress_bar=False,
             tol=1e-6,
         )
+        assert sol.ndim == 2
 
         if sparse:
             A = A.A
@@ -72,7 +73,7 @@ def test_gmres(self, seed: int, sparse: bool):
         np.testing.assert_allclose(A @ sol, B, rtol=1e-6, atol=1e-10)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
     )
     def test_direct_solver_dense(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
@@ -86,6 +87,7 @@ def test_direct_solver_dense(self, seed: int, sparse: bool):
             show_progress_bar=False,
             tol=1e-6,
         )
+        assert sol.ndim == 2
 
         if sparse:
             A = A.A
@@ -94,7 +96,7 @@ def test_direct_solver_dense(self, seed: int, sparse: bool):
         np.testing.assert_allclose(A @ sol, B, rtol=1e-6, atol=1e-10)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(30, 40), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(30, 40), [False] * 5 + [True] * 5)
     )
     def test_eye(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
@@ -108,6 +110,7 @@ def test_eye(self, seed: int, sparse: bool):
             show_progress_bar=False,
             tol=1e-6,
         )
+        assert sol.ndim == 2
 
         if sparse:
             A = A.A
@@ -213,6 +216,8 @@ def test_petsc_scipy_matches(self, seed: int, solver: str, sparse: bool):
             show_progress_bar=False,
             tol=1e-6,
         )
+        assert sol_petsc.ndim == 2
+        assert sol_scipy.ndim == 2
         if sparse:
             A = A.A
             B = B.A
@@ -222,7 +227,7 @@ def test_petsc_scipy_matches(self, seed: int, solver: str, sparse: bool):
         np.testing.assert_allclose(A @ sol_petsc, B, rtol=1e-6, atol=1e-8)
         np.testing.assert_allclose(sol_petsc, sol_scipy, rtol=1e-6, atol=1e-8)
 
-    @pytest.mark.parametrize("seed, sparse", zip(range(10), [False] * 5 + [True] * 5))
+    @pytest.mark.parametrize("seed,sparse", zip(range(10), [False] * 5 + [True] * 5))
     def test_gmres(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
 
@@ -236,6 +241,7 @@ def test_gmres(self, seed: int, sparse: bool):
             preconditioner="lu",
             tol=1e-6,
         )
+        assert sol.ndim == 2
 
         if sparse:
             A = A.A
@@ -245,7 +251,7 @@ def test_gmres(self, seed: int, sparse: bool):
         np.testing.assert_allclose(A @ sol, B, rtol=1e-6, atol=1e-8)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
     )
     def test_direct_solver(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
@@ -259,6 +265,7 @@ def test_direct_solver(self, seed: int, sparse: bool):
             show_progress_bar=False,
             tol=1e-6,
         )
+        assert sol.ndim == 2
 
         if sparse:
             A = A.A
@@ -268,13 +275,14 @@ def test_direct_solver(self, seed: int, sparse: bool):
         np.testing.assert_allclose(A @ sol, B, rtol=1e-6, atol=1e-10)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
     )
     def test_mat_solve(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
         A = (speye(*A.shape) if sparse else np.eye(*A.shape)) - A
 
-        X = _petsc_mat_solve(A, B, tol=1e-8)
+        X = _petsc_direct_solve(A, B, tol=1e-8)
+        assert X.ndim == 2
         if sparse:
             A = A.A
             B = B.A
@@ -282,28 +290,31 @@ def test_mat_solve(self, seed: int, sparse: bool):
         np.testing.assert_allclose(A @ X, B, rtol=1e-6)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
     )
     def test_mat_solve_1_dim_b(self, seed: int, sparse: bool):
         A, B = _create_a_b_matrices(seed, sparse)
         A = (speye(*A.shape) if sparse else np.eye(*A.shape)) - A
         B = B[:, 0]
 
-        X = _petsc_mat_solve(A, B, tol=1e-8)
+        X = _petsc_direct_solve(A, B, tol=1e-8)
+        assert X.ndim == 2
+        assert X.shape[1] == 1
         if sparse:
             A = A.A
             B = B.A
 
-        np.testing.assert_allclose(A @ X, B.squeeze(), rtol=1e-6)
+        np.testing.assert_allclose(((A @ X).squeeze()), B.squeeze(), rtol=1e-6)
 
     @pytest.mark.parametrize(
-        "seed, sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
+        "seed,sparse", zip(range(10, 20), [False] * 5 + [True] * 5)
     )
     def test_mat_solve_inversion(self, seed: int, sparse: bool):
         A, _ = _create_a_b_matrices(seed, sparse)
         A = (speye(*A.shape) if sparse else np.eye(*A.shape)) - A
 
-        X = _petsc_mat_solve(A, tol=1e-8)
+        X = _petsc_direct_solve(A, tol=1e-8)
+        assert X.ndim == 2
         if sparse:
             A = A.A
 
diff --git a/tests/test_model.py b/tests/test_model.py
index ad324d479..6e8c33a96 100644
--- a/tests/test_model.py
+++ b/tests/test_model.py
@@ -2,13 +2,14 @@
 import pytest
 from io import BytesIO
 from copy import copy, deepcopy
+from pathlib import Path
 from _helpers import gamr_skip, create_model, assert_models_equal
 from itertools import product
 
 from anndata import AnnData
 from cellrank.tl import Lineage
+from cellrank._key import Key
 from cellrank.ul.models import GAM, GAMR, FittedModel, SKLearnModel
-from cellrank.tl._constants import AbsProbKey
 from cellrank.ul.models._utils import (
     _OFFSET_KEY,
     NormMode,
@@ -101,7 +102,7 @@ def test_confidence_interval(self, adata_cflare):
         assert ci is model.conf_int
 
     def test_model_1_lineage(self, adata_cflare):
-        adata_cflare.obsm[AbsProbKey.FORWARD.s] = Lineage(
+        adata_cflare.obsm[Key.obsm.abs_probs(False)] = Lineage(
             np.ones((adata_cflare.n_obs, 1)), names=["foo"]
         )
         model = create_model(adata_cflare)
@@ -132,12 +133,12 @@ def test_prepare_resets_fields(self, adata_cflare: AnnData):
 class TestUtils:
     def test_extract_data_wrong_type(self):
         with pytest.raises(TypeError):
-            _extract_data(None)
+            _ = _extract_data(None)
 
     def test_extract_data_raw_None(self, adata: AnnData):
         adata = AnnData(adata.X, raw=None)
         with pytest.raises(ValueError):
-            _extract_data(adata, use_raw=True)
+            _ = _extract_data(adata, use_raw=True)
 
     def test_extract_data_invalid_layer(self, adata: AnnData):
         with pytest.raises(KeyError):
@@ -450,7 +451,7 @@ def test_invalid_link_function(self, adata: AnnData):
 
     def test_invalid_grid_type(self, adata: AnnData):
         with pytest.raises(TypeError):
-            GAM(adata, grid=1311)
+            _ = GAM(adata, grid=1311)
 
     def test_default_grid(self, adata_cflare: AnnData):
         g = GAM(adata_cflare, grid="default")
@@ -626,6 +627,30 @@ def test_pickling_pygam(self, pygam_model: GAM):
 
         assert_models_equal(pygam_model, actual_model, pickled=True)
 
+    @pytest.mark.parametrize("copy", [False, True])
+    @pytest.mark.parametrize("write_adata", [False, True])
+    def test_read_write(
+        self, sklearn_model: SKLearnModel, tmpdir, write_adata: bool, copy: bool
+    ):
+        path = Path(tmpdir) / "model.pickle"
+        sklearn_model.write(path, write_adata=write_adata)
+
+        if write_adata:
+            model = SKLearnModel.read(path)
+            assert model.adata is not None
+        else:
+            with open(path, "rb") as fin:
+                model: SKLearnModel = pickle.load(fin)
+                assert model.adata is None
+                assert model.shape == (sklearn_model.adata.n_obs,)
+            model = SKLearnModel.read(path, adata=sklearn_model.adata, copy=copy)
+            if copy:
+                assert model.adata is not sklearn_model.adata
+            else:
+                assert model.adata is sklearn_model.adata
+            model.adata = model.adata.copy()
+        assert_models_equal(sklearn_model, model, pickled=True)
+
 
 class TestFittedModel:
     def test_wrong_xt_yt_shape(self):
@@ -775,7 +800,7 @@ def test_normal_run(self):
             fm.default_confidence_interval(), [[4, 5], [6, 7]]
         )
 
-    def from_model_wrong_type(self, adata_cflare):
+    def test_from_model_wrong_type(self, adata_cflare):
         m = create_model(adata_cflare)
         with pytest.raises(TypeError):
             FittedModel.from_model(m.model)
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 244f54327..fbc02566f 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -2,73 +2,39 @@
 
 import cellrank as cr
 from anndata import AnnData
+from cellrank._key import Key
 from cellrank.tl.kernels import VelocityKernel, ConnectivityKernel
-from cellrank.tl._constants import (
-    Direction,
-    DirPrefix,
-    AbsProbKey,
-    TermStatesKey,
-    _probs,
-    _colors,
-    _lin_names,
-    _transition,
-)
-from cellrank.tl.estimators._constants import P
 
 import numpy as np
+import pandas as pd
 from pandas.api.types import is_categorical_dtype
 
 
-def _assert_has_all_keys(adata: AnnData, direction: Direction):
-    assert _transition(direction) in adata.obsp.keys()
-    # check if it's not a dummy transition matrix
-    assert not np.all(np.isclose(np.diag(adata.obsp[_transition(direction)].A), 1.0))
-    assert f"{_transition(direction)}_params" in adata.uns.keys()
+def _assert_has_all_keys(adata: AnnData, bwd: bool = False) -> None:
+    # fmt: off
+    # term states
+    key = Key.obs.term_states(bwd)
+    assert is_categorical_dtype(adata.obs[key])
+    assert Key.obs.probs(key) in adata.obs
+    assert Key.uns.colors(key) in adata.uns
 
-    if direction == Direction.FORWARD:
-        assert str(AbsProbKey.FORWARD) in adata.obsm
-        assert isinstance(adata.obsm[str(AbsProbKey.FORWARD)], cr.tl.Lineage)
+    # lineages
+    abs_probs = adata.obsm[Key.obsm.abs_probs(bwd)]
+    assert isinstance(abs_probs, cr.tl.Lineage)
+    np.testing.assert_array_equal(abs_probs.names, adata.obs[key].cat.categories)
+    np.testing.assert_array_equal(abs_probs.colors, adata.uns[Key.uns.colors(key)])
+    np.testing.assert_allclose(abs_probs.X.sum(1), 1.0, rtol=1e-3)
 
-        assert _colors(AbsProbKey.FORWARD) in adata.uns.keys()
-        assert _lin_names(AbsProbKey.FORWARD) in adata.uns.keys()
-
-        assert str(TermStatesKey.FORWARD) in adata.obs
-        assert is_categorical_dtype(adata.obs[str(TermStatesKey.FORWARD)])
-
-        assert _probs(TermStatesKey.FORWARD) in adata.obs
-
-        # check the correlations with all lineages have been computed
-        lin_probs = adata.obsm[str(AbsProbKey.FORWARD)]
-        np.in1d(
-            [f"{str(DirPrefix.FORWARD)} {key}" for key in lin_probs.names],
-            adata.var.keys(),
-        ).all()
-
-    else:
-        assert str(AbsProbKey.BACKWARD) in adata.obsm
-        assert isinstance(adata.obsm[str(AbsProbKey.BACKWARD)], cr.tl.Lineage)
-
-        assert _colors(AbsProbKey.BACKWARD) in adata.uns.keys()
-        assert _lin_names(AbsProbKey.BACKWARD) in adata.uns.keys()
-
-        assert str(TermStatesKey.BACKWARD) in adata.obs
-        assert is_categorical_dtype(adata.obs[str(TermStatesKey.BACKWARD)])
-
-        assert _probs(TermStatesKey.BACKWARD) in adata.obs
-
-        # check the correlations with all lineages have been computed
-        lin_probs = adata.obsm[str(AbsProbKey.BACKWARD)]
-        np.in1d(
-            [f"{str(DirPrefix.BACKWARD)} {key}" for key in lin_probs.names],
-            adata.var.keys(),
-        ).all()
+    # drivers
+    assert isinstance(adata.varm[Key.varm.lineage_drivers(bwd)], pd.DataFrame)
+    # fmt: on
 
 
 class TestHighLevelPipeline:
     def test_plot_states_not_computed(self, adata: AnnData):
-        with pytest.raises(RuntimeError):
+        with pytest.raises(KeyError):
             cr.pl.initial_states(adata)
-        with pytest.raises(RuntimeError):
+        with pytest.raises(KeyError):
             cr.pl.terminal_states(adata)
 
     def test_write_transition_matrix(self, adata: AnnData):
@@ -81,7 +47,7 @@ def test_write_transition_matrix(self, adata: AnnData):
     def test_states_no_precomputed_transition_matrix(self, adata: AnnData):
         cr.tl.terminal_states(adata, key="foo")
 
-        assert str(_transition(Direction.FORWARD)) in adata.obsp
+        np.testing.assert_allclose(adata.obsp[Key.uns.kernel(False)].A.sum(1), 1.0)
 
     def test_states_use_precomputed_transition_matrix(self, adata: AnnData):
         cr.tl.transition_matrix(adata, key="foo")
@@ -104,10 +70,10 @@ def test_fwd_pipeline_cflare(self, adata: AnnData):
         cr.pl.lineages(adata)
         cr.tl.lineage_drivers(adata, use_raw=False)
 
-        ln = adata.obsm[str(AbsProbKey.FORWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(False)].names[0]
         cr.pl.lineage_drivers(adata, ln, use_raw=False, backward=False)
 
-        _assert_has_all_keys(adata, Direction.FORWARD)
+        _assert_has_all_keys(adata)
 
     def test_fwd_pipeline_invalid_raw_requested(self, adata: AnnData):
         cr.tl.terminal_states(
@@ -122,7 +88,7 @@ def test_fwd_pipeline_invalid_raw_requested(self, adata: AnnData):
         cr.pl.lineages(adata)
         cr.tl.lineage_drivers(adata, use_raw=False)
 
-        ln = adata.obsm[str(AbsProbKey.FORWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(False)].names[0]
         with pytest.raises(RuntimeError):
             cr.pl.lineage_drivers(adata, ln, use_raw=True, backward=False)
 
@@ -139,10 +105,10 @@ def test_bwd_pipeline_cflare(self, adata: AnnData):
         cr.pl.lineages(adata, backward=True)
         cr.tl.lineage_drivers(adata, use_raw=False, backward=True)
 
-        ln = adata.obsm[str(AbsProbKey.BACKWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(True)].names[0]
         cr.pl.lineage_drivers(adata, ln, use_raw=False, backward=True)
 
-        _assert_has_all_keys(adata, Direction.BACKWARD)
+        _assert_has_all_keys(adata, bwd=True)
 
     def test_fwd_pipeline_gpcca(self, adata: AnnData):
         cr.tl.terminal_states(
@@ -156,10 +122,10 @@ def test_fwd_pipeline_gpcca(self, adata: AnnData):
         cr.tl.lineages(adata)
         cr.pl.lineages(adata)
         cr.tl.lineage_drivers(adata, use_raw=False)
-        ln = adata.obsm[str(AbsProbKey.FORWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(False)].names[0]
         cr.pl.lineage_drivers(adata, ln, use_raw=False, backward=False)
 
-        _assert_has_all_keys(adata, Direction.FORWARD)
+        _assert_has_all_keys(adata)
 
     def test_fwd_pipeline_gpcca_invalid_raw_requested(self, adata: AnnData):
         cr.tl.terminal_states(
@@ -173,7 +139,7 @@ def test_fwd_pipeline_gpcca_invalid_raw_requested(self, adata: AnnData):
         cr.tl.lineages(adata)
         cr.pl.lineages(adata)
         cr.tl.lineage_drivers(adata, use_raw=False)
-        ln = adata.obsm[str(AbsProbKey.FORWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(False)].names[0]
         with pytest.raises(RuntimeError):
             cr.pl.lineage_drivers(adata, ln, use_raw=True, backward=False)
 
@@ -189,10 +155,10 @@ def test_bwd_pipeline_gpcca(self, adata: AnnData):
         cr.tl.lineages(adata, backward=True)
         cr.pl.lineages(adata, backward=True)
         cr.tl.lineage_drivers(adata, use_raw=False, backward=True)
-        ln = adata.obsm[str(AbsProbKey.BACKWARD)].names[0]
+        ln = adata.obsm[Key.obsm.abs_probs(True)].names[0]
         cr.pl.lineage_drivers(adata, ln, use_raw=False, backward=True)
 
-        _assert_has_all_keys(adata, Direction.BACKWARD)
+        _assert_has_all_keys(adata, bwd=True)
 
     def test_multiple_read_write_diff(self, adata: AnnData):
         for n_states in range(2, 5):
@@ -210,9 +176,10 @@ def test_multiple_read_write_diff(self, adata: AnnData):
 
             cr.tl.lineages(adata, backward=False)
             cr.pl.lineages(adata)
+            cr.tl.lineage_drivers(adata)
 
-            _assert_has_all_keys(adata, Direction.FORWARD)
-            probs = adata.obsm[str(AbsProbKey.FORWARD)]
+            _assert_has_all_keys(adata)
+            probs = adata.obsm[Key.obsm.abs_probs(False)]
 
             assert probs.shape[1] == n_states
             for name in probs.names:
@@ -234,28 +201,24 @@ def test_multiple_read_write_same(self, adata: AnnData):
 
         e.set_terminal_states({"foo": adata.obs_names[:20]})
 
-        e = cr.tl.estimators.GPCCA(adata, read_from_adata=True, obsp_key="T_fwd")
+        e = cr.tl.estimators.GPCCA(adata, 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):
+    def test_fwd_pipeline_cflare(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata).compute_transition_matrix()
         final_kernel = 0.8 * vk + 0.2 * ck
 
         estimator_fwd = cr.tl.estimators.CFLARE(final_kernel)
 
-        estimator_fwd.compute_partition()
-
         estimator_fwd.compute_eigendecomposition()
         estimator_fwd.plot_spectrum()
         estimator_fwd.plot_spectrum(real_only=True)
-        estimator_fwd.plot_eigendecomposition()
-        estimator_fwd.plot_eigendecomposition(left=False)
 
-        estimator_fwd.compute_terminal_states(use=1)
+        estimator_fwd.compute_terminal_states(use=1, method="leiden")
         estimator_fwd.plot_terminal_states()
 
         estimator_fwd.compute_absorption_probabilities()
@@ -263,9 +226,9 @@ def test_fwd_pipelne_cflare(self, adata: AnnData):
 
         estimator_fwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.FORWARD)
+        _assert_has_all_keys(adata)
 
-    def test_bwd_pipelne_cflare(self, adata: AnnData):
+    def test_bwd_pipeline_cflare(self, adata: AnnData):
         vk = VelocityKernel(adata, backward=True).compute_transition_matrix(
             softmax_scale=4
         )
@@ -274,15 +237,11 @@ def test_bwd_pipelne_cflare(self, adata: AnnData):
 
         estimator_bwd = cr.tl.estimators.CFLARE(final_kernel)
 
-        estimator_bwd.compute_partition()
-
         estimator_bwd.compute_eigendecomposition()
         estimator_bwd.plot_spectrum()
         estimator_bwd.plot_spectrum(real_only=True)
-        estimator_bwd.plot_eigendecomposition()
-        estimator_bwd.plot_eigendecomposition(left=False)
 
-        estimator_bwd.compute_terminal_states(use=1)
+        estimator_bwd.compute_terminal_states(use=1, method="kmeans")
         estimator_bwd.plot_terminal_states()
 
         estimator_bwd.compute_absorption_probabilities()
@@ -290,23 +249,20 @@ def test_bwd_pipelne_cflare(self, adata: AnnData):
 
         estimator_bwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.BACKWARD)
+        _assert_has_all_keys(adata, bwd=True)
 
-    def test_fwd_pipelne_gpcca(self, adata: AnnData):
+    def test_fwd_pipeline_gpcca(self, adata: AnnData):
         vk = VelocityKernel(adata).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata).compute_transition_matrix()
         final_kernel = 0.8 * vk + 0.2 * ck
 
         estimator_fwd = cr.tl.estimators.GPCCA(final_kernel)
 
-        estimator_fwd.compute_partition()
-
         estimator_fwd.compute_eigendecomposition()
         estimator_fwd.plot_spectrum()
         estimator_fwd.plot_spectrum(real_only=True)
 
         estimator_fwd.compute_schur(5, method="brandts")
-        estimator_fwd.plot_schur()
 
         estimator_fwd.compute_macrostates(3, n_cells=10)
         estimator_fwd.plot_macrostates()
@@ -320,21 +276,21 @@ def test_fwd_pipelne_gpcca(self, adata: AnnData):
         estimator_fwd.compute_absorption_probabilities()
         estimator_fwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.FORWARD)
+        _assert_has_all_keys(adata)
 
         # select a subset of states
         estimator_fwd.set_terminal_states_from_macrostates(
             n_cells=16,
-            names=estimator_fwd._get(P.MACRO).cat.categories[:2],
+            names=estimator_fwd.macrostates.cat.categories[:2],
         )
         estimator_fwd.plot_terminal_states()
 
         estimator_fwd.compute_absorption_probabilities()
         estimator_fwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.FORWARD)
+        _assert_has_all_keys(adata)
 
-    def test_bwd_pipelne_gpcca(self, adata: AnnData):
+    def test_bwd_pipeline_gpcca(self, adata: AnnData):
         vk = VelocityKernel(adata, backward=True).compute_transition_matrix(
             softmax_scale=4
         )
@@ -348,7 +304,6 @@ def test_bwd_pipelne_gpcca(self, adata: AnnData):
         estimator_bwd.plot_spectrum(real_only=True)
 
         estimator_bwd.compute_schur(5, method="brandts")
-        estimator_bwd.plot_schur()
 
         estimator_bwd.compute_macrostates(3, n_cells=16)
         estimator_bwd.plot_macrostates()
@@ -362,16 +317,16 @@ def test_bwd_pipelne_gpcca(self, adata: AnnData):
         estimator_bwd.compute_absorption_probabilities()
         estimator_bwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.BACKWARD)
+        _assert_has_all_keys(adata, bwd=True)
 
         # select a subset of states
         estimator_bwd.set_terminal_states_from_macrostates(
             n_cells=16,
-            names=estimator_bwd._get(P.MACRO).cat.categories[:2],
+            names=estimator_bwd.macrostates.cat.categories[:2],
         )
         estimator_bwd.plot_terminal_states()
 
         estimator_bwd.compute_absorption_probabilities()
         estimator_bwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
-        _assert_has_all_keys(adata, Direction.BACKWARD)
+        _assert_has_all_keys(adata, bwd=True)
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 54adcf5c9..0f4876f8b 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -15,9 +15,9 @@
 import cellrank as cr
 from anndata import AnnData
 from cellrank.tl import Lineage
+from cellrank._key import Key
 from cellrank.ul.models import GAMR
 from cellrank.tl.kernels import VelocityKernel, PseudotimeKernel, ConnectivityKernel
-from cellrank.tl._constants import AbsProbKey, _pd
 from cellrank.tl.estimators import GPCCA, CFLARE
 
 import numpy as np
@@ -1229,7 +1229,7 @@ def test_heatmap_lineages_return_genes(self, adata_cflare: AnnData):
 
         assert isinstance(df, pd.DataFrame)
         np.testing.assert_array_equal(
-            df.columns, adata_cflare.obsm[AbsProbKey.FORWARD.s].names
+            df.columns, adata_cflare.obsm[Key.obsm.abs_probs(False)].names
         )
         assert len(df) == 10
         assert set(df.iloc[:, 0].values) == set(GENES[:10])
@@ -1249,7 +1249,7 @@ def test_heatmap_lineages_return_models(self, adata_cflare: AnnData):
         models = pd.DataFrame(models).T
         np.testing.assert_array_equal(models.index, GENES[:10])
         np.testing.assert_array_equal(
-            models.columns, adata_cflare.obsm[AbsProbKey.FORWARD.s].names
+            models.columns, adata_cflare.obsm[Key.obsm.abs_probs(False)].names
         )
         assert np.all(models.astype(bool))
 
@@ -1266,7 +1266,7 @@ def test_heatmap_lineages_return_models_and_genes(self, adata_cflare: AnnData):
             dpi=DPI,
         )
 
-        lnames = adata_cflare.obsm[AbsProbKey.FORWARD.s].names
+        lnames = adata_cflare.obsm[Key.obsm.abs_probs(False)].names
 
         models = pd.DataFrame(models).T
         np.testing.assert_array_equal(models.index, GENES[:10])
@@ -1293,7 +1293,7 @@ def test_heatmap_lineages_return_genes_large_number(self, adata_cflare: AnnData)
 
         assert isinstance(df, pd.DataFrame)
         np.testing.assert_array_equal(
-            df.columns, adata_cflare.obsm[AbsProbKey.FORWARD.s].names
+            df.columns, adata_cflare.obsm[Key.obsm.abs_probs(False)].names
         )
         assert len(df) == len(genes)
         assert set(df.iloc[:, 0].values) == set(genes)
@@ -1313,7 +1313,7 @@ def test_heatmap_lineages_return_genes_same_order(self, adata_cflare: AnnData):
 
         assert isinstance(df, pd.DataFrame)
         np.testing.assert_array_equal(
-            df.columns, adata_cflare.obsm[AbsProbKey.FORWARD.s].names
+            df.columns, adata_cflare.obsm[Key.obsm.abs_probs(False)].names
         )
         assert len(df) == 10
         assert set(df.iloc[:, 0].values) == set(GENES[:10])
@@ -2357,29 +2357,13 @@ def test_mc_spectrum_evals(self, mc: CFLARE, fpath: str):
     def test_mc_spectrum_evals_complex(self, mc: CFLARE, fpath: str):
         mc.plot_spectrum(2, real_only=False, dpi=DPI, save=fpath)
 
-    @compare(kind="cflare")
-    def test_scvelo_eigendecomposition_clusters(self, mc: CFLARE, fpath: str):
-        mc.plot_eigendecomposition(cluster_key="clusters", dpi=DPI, save=fpath)
-
-    @compare(kind="cflare")
-    def test_scvelo_eigendecomposition_left(self, mc: CFLARE, fpath: str):
-        mc.plot_eigendecomposition(dpi=DPI, save=fpath)
-
-    @compare(kind="cflare")
-    def test_scvelo_eigendecomposition_right(self, mc: CFLARE, fpath: str):
-        mc.plot_eigendecomposition(left=False, dpi=DPI, save=fpath)
-
-    @compare(kind="cflare")
-    def test_scvelo_eigendecomposition_use_2(self, mc: CFLARE, fpath: str):
-        mc.plot_eigendecomposition(use=2, dpi=DPI, save=fpath)
-
     @compare(kind="cflare")
     def test_scvelo_final_states(self, mc: CFLARE, fpath: str):
         mc.plot_terminal_states(dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_final_states_clusters(self, mc: CFLARE, fpath: str):
-        mc.plot_terminal_states(cluster_key="clusters", dpi=DPI, save=fpath)
+        mc.plot_terminal_states(color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs(self, mc: CFLARE, fpath: str):
@@ -2387,7 +2371,7 @@ def test_scvelo_lin_probs(self, mc: CFLARE, fpath: str):
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs_clusters(self, mc: CFLARE, fpath: str):
-        mc.plot_absorption_probabilities(cluster_key="clusters", dpi=DPI, save=fpath)
+        mc.plot_absorption_probabilities(color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs_cmap(self, mc: CFLARE, fpath: str):
@@ -2395,7 +2379,7 @@ def test_scvelo_lin_probs_cmap(self, mc: CFLARE, fpath: str):
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs_lineages(self, mc: CFLARE, fpath: str):
-        mc.plot_absorption_probabilities(lineages=["0"], dpi=DPI, save=fpath)
+        mc.plot_absorption_probabilities(states=["0"], dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs_time(self, mc: CFLARE, fpath: str):
@@ -2439,22 +2423,6 @@ def test_gpcca_schur_matrix_title(self, mc: GPCCA, fpath: str):
     def test_gpcca_schur_matrix_cmap(self, mc: GPCCA, fpath: str):
         mc.plot_schur_matrix(cmap=cm.inferno, dpi=DPI, save=fpath)
 
-    @compare(kind="gpcca")
-    def test_scvelo_gpcca_schur_emb(self, mc: GPCCA, fpath: str):
-        mc.plot_schur(dpi=DPI, save=fpath)
-
-    @compare(kind="gpcca")
-    def test_scvelo_gpcca_schur_emb_use_2(self, mc: GPCCA, fpath: str):
-        mc.plot_schur(use=1, dpi=DPI, save=fpath)
-
-    @compare(kind="gpcca")
-    def test_scvelo_gpcca_schur_emb_abs(self, mc: GPCCA, fpath: str):
-        mc.plot_schur(abs_value=True, dpi=DPI, save=fpath)
-
-    @compare(kind="gpcca")
-    def test_scvelo_gpcca_schur_cluster_key(self, mc: GPCCA, fpath: str):
-        mc.plot_schur(cluster_key="clusters", dpi=DPI, save=fpath)
-
     @compare(kind="gpcca")
     def test_gpcca_coarse_T(self, mc: GPCCA, fpath: str):
         mc.plot_coarse_T(
@@ -2501,7 +2469,7 @@ def test_scvelo_gpcca_meta_states(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_lineages(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(lineages=["0"], dpi=DPI, save=fpath)
+        mc.plot_macrostates(states=["0"], dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_discrete(self, mc: GPCCA, fpath: str):
@@ -2509,7 +2477,7 @@ def test_scvelo_gpcca_meta_states_discrete(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_cluster_key(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(cluster_key="clusters", dpi=DPI, save=fpath)
+        mc.plot_macrostates(color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_no_same_plot(self, mc: GPCCA, fpath: str):
@@ -2533,7 +2501,7 @@ def test_scvelo_gpcca_final_states(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_lineages(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(lineages=["0"], dpi=DPI, save=fpath)
+        mc.plot_terminal_states(states=["0"], dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_discrete(self, mc: GPCCA, fpath: str):
@@ -2541,7 +2509,7 @@ def test_scvelo_gpcca_final_states_discrete(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_cluster_key(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(cluster_key="clusters", dpi=DPI, save=fpath)
+        mc.plot_terminal_states(color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_no_same_plot(self, mc: GPCCA, fpath: str):
@@ -2562,13 +2530,13 @@ def test_scvelo_gpcca_final_states_time(self, mc: GPCCA, fpath: str):
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_disc_same(self, mc: GPCCA, fpath: str):
         mc.plot_absorption_probabilities(
-            cluster_key="clusters", discrete=True, same_plot=True, dpi=DPI, save=fpath
+            color="clusters", discrete=True, same_plot=True, dpi=DPI, save=fpath
         )
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_disc_not_same(self, mc: GPCCA, fpath: str):
         mc.plot_absorption_probabilities(
-            cluster_key="clusters", discrete=True, same_plot=False, dpi=DPI, save=fpath
+            color="clusters", discrete=True, same_plot=False, dpi=DPI, save=fpath
         )
 
     @compare(kind="gpcca")
@@ -2580,13 +2548,13 @@ def test_scvelo_gpcca_abs_probs_cont_same_no_clusters(self, mc: GPCCA, fpath: st
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_cont_same_clusters(self, mc: GPCCA, fpath: str):
         mc.plot_absorption_probabilities(
-            cluster_key="clusters", discrete=False, same_plot=True, dpi=DPI, save=fpath
+            color="clusters", discrete=False, same_plot=True, dpi=DPI, save=fpath
         )
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_cont_not_same(self, mc: GPCCA, fpath: str):
         mc.plot_absorption_probabilities(
-            cluster_key="clusters", discrete=False, same_plot=False, dpi=DPI, save=fpath
+            color="clusters", discrete=False, same_plot=False, dpi=DPI, save=fpath
         )
 
     @compare(kind="gpcca")
@@ -2622,7 +2590,7 @@ def test_scvelo_lineages_cmap(self, adata: AnnData, fpath: str):
 
     @compare()
     def test_scvelo_lineages_subset(self, adata: AnnData, fpath: str):
-        cr.pl.lineages(adata, cluster_key="clusters", dpi=DPI, save=fpath)
+        cr.pl.lineages(adata, color="clusters", dpi=DPI, save=fpath)
 
 
 class TestHighLvlStates:
@@ -2867,7 +2835,7 @@ def test_model_show_lin_prob_cells_lineage_ci(self, adata: AnnData, fpath: str):
 
     @compare()
     def test_model_1_lineage(self, adata: AnnData, fpath: str):
-        adata.obsm[AbsProbKey.FORWARD.s] = Lineage(
+        adata.obsm[Key.obsm.abs_probs(False)] = Lineage(
             np.ones((adata.n_obs, 1)), names=["foo"]
         )
         model = create_model(adata)
@@ -3112,7 +3080,7 @@ def test_fitted_heatmap(self, adata: AnnData, fpath: str):
 class TestCircularProjection:
     def test_proj_too_few_lineages(self, adata_gpcca_fwd):
         adata, _ = adata_gpcca_fwd
-        lineages = adata.obsm[AbsProbKey.FORWARD.s].names[:2]
+        lineages = adata.obsm[Key.obsm.abs_probs(False)].names[:2]
 
         with pytest.raises(ValueError, match=r"Expected at least `3` lineages"):
             cr.pl.circular_projection(
@@ -3168,13 +3136,13 @@ def test_proj_ncols(self, adata: AnnData, fpath: str):
     @compare()
     def test_proj_labelrot(self, adata: AnnData, fpath: str):
         cr.pl.circular_projection(
-            adata, keys="clusters", labelrot="default", dpi=DPI, save=fpath
+            adata, keys="clusters", label_rot="default", dpi=DPI, save=fpath
         )
 
     @compare()
     def test_proj_labeldistance(self, adata: AnnData, fpath: str):
         cr.pl.circular_projection(
-            adata, keys="clusters", labeldistance=1.5, dpi=DPI, save=fpath
+            adata, keys="clusters", label_distance=1.5, dpi=DPI, save=fpath
         )
 
     @compare()
@@ -3195,7 +3163,7 @@ def test_proj_extra_keys(self, adata: AnnData, fpath: str):
             adata, keys=["kl_divergence", "entropy"], dpi=DPI, save=fpath
         )
 
-        apk = AbsProbKey.FORWARD.s
+        apk = Key.obsm.abs_probs(False)
         assert f"{apk}_kl_divergence" in adata.obs
         assert f"{apk}_entropy" in adata.obs
 
@@ -3410,7 +3378,7 @@ def test_driver_corr(self, mc: GPCCA, fpath: str):
     @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"
+            "0", "1", dpi=DPI, save=fpath, color="2_corr"
         )
 
     @compare(kind="gpcca")
@@ -3445,13 +3413,13 @@ def test_driver_corr_legend_loc(self, mc: GPCCA, fpath: str):
     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"
+            "0", "1", dpi=DPI, save=fpath, use_raw=True, color="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"
+            "0", "1", dpi=DPI, save=fpath, color="1_qval", cmap="inferno"
         )
 
     @compare(kind="gpcca")
diff --git a/tests/test_random_walk.py b/tests/test_random_walk.py
index 5d48d2cbd..d39760cd4 100644
--- a/tests/test_random_walk.py
+++ b/tests/test_random_walk.py
@@ -60,20 +60,21 @@ def test_stop_ixs(self, test_matrix_1: np.ndarray):
             max_iter=1000, seed=42
         )
 
-        assert len(res) <= 1000
+        assert len(res) <= 1001
         assert res[-1] == 1
         assert res[-2] != 1
 
-    def successive_hits(self, test_matrix_1: np.ndarray):
-        res = RandomWalk(test_matrix_1, stop_ixs=[1]).simulate_one(
-            max_iter=1000, seed=42, successive_hits=2
+    def test_successive_hits(self, test_matrix_1: np.ndarray):
+        test_matrix_1[0, 0] = 0.2
+        test_matrix_1[0, 2] = 0.0
+        res = RandomWalk(test_matrix_1, stop_ixs=[0]).simulate_one(
+            max_iter=1000, seed=42, successive_hits=1
         )
 
-        assert len(res) <= 1000
-        assert res[-1] == 1
-        assert res[-2] == 1
-        assert res[-3] == 1
-        assert res[-4] != 1
+        assert len(res) <= 1001
+        assert res[-1] == 0
+        assert res[-2] == 0
+        assert res[-3] != 0
 
     @pytest.mark.parametrize("n_jobs", [1, 2])
     @pytest.mark.parametrize("backend", ["threading", "loky"])
diff --git a/tests/test_read.py b/tests/test_read.py
index c1c65db42..8a3debb34 100644
--- a/tests/test_read.py
+++ b/tests/test_read.py
@@ -1,15 +1,14 @@
 from typing import Tuple, Callable
 
-import pytest
 from pathlib import Path
 
 import scanpy as sc
 import cellrank as cr
 from anndata import AnnData
+from cellrank._key import Key
 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
+from cellrank.tl.estimators import CFLARE
 
 import numpy as np
 
@@ -22,13 +21,14 @@ def decorator(self, adata_cflare_fwd: Tuple[AnnData, CFLARE], tmpdir):
 
             dirname = func.__name__
             path = tmpdir.mkdir(dirname).join("tmp.h5ad")
+            key = Key.obs.term_states(False)
 
             return func(
                 self,
                 adata,
                 path,
-                str(AbsProbKey.FORWARD),
-                adata.obsm[str(AbsProbKey.FORWARD)].shape[1],
+                Key.obsm.abs_probs(False),
+                len(adata.obs[key].cat.categories),
             )
 
         return decorator
@@ -47,24 +47,9 @@ def test_no_lineage(self, adata: AnnData, path: Path, lin_key: str, _: int):
         assert adata_new is not adata  # sanity check
         assert lin_key not in adata_new.obsm.keys()
 
-    @test_fwd()
-    def test_no_names(self, adata: AnnData, path: Path, lin_key: str, n_lins: int):
-        names_key = _lin_names(lin_key)
-        del adata.uns[names_key]
-
-        sc.write(path, adata)
-        adata_new = cr.read(path)
-        lins = adata_new.obsm[lin_key]
-
-        assert isinstance(lins, Lineage)
-        np.testing.assert_array_equal(
-            lins.names, [f"Lineage {i}" for i in range(n_lins)]
-        )
-        np.testing.assert_array_equal(lins.names, adata_new.uns[names_key])
-
     @test_fwd()
     def test_no_colors(self, adata: AnnData, path: Path, lin_key: str, n_lins: int):
-        colors_key = _colors(lin_key)
+        colors_key = Key.uns.colors(Key.obs.term_states(False))
         del adata.uns[colors_key]
 
         sc.write(path, adata)
@@ -75,38 +60,11 @@ def test_no_colors(self, adata: AnnData, path: Path, lin_key: str, n_lins: int):
         np.testing.assert_array_equal(lins.colors, _create_categorical_colors(n_lins))
         np.testing.assert_array_equal(lins.colors, adata_new.uns[colors_key])
 
-    @test_fwd()
-    def test_wrong_names_length(
-        self, adata: AnnData, path: Path, lin_key: str, n_lins: int
-    ):
-        names_key = _lin_names(lin_key)
-        adata.uns[names_key] = list(adata.uns[names_key])
-        adata.uns[names_key] += ["foo", "bar", "baz"]
-
-        sc.write(path, adata)
-        adata_new = cr.read(path)
-        lins = adata_new.obsm[lin_key]
-
-        assert isinstance(lins, Lineage)
-        np.testing.assert_array_equal(
-            lins.names, [f"Lineage {i}" for i in range(n_lins)]
-        )
-        np.testing.assert_array_equal(lins.names, adata_new.uns[names_key])
-
-    @test_fwd()
-    def test_non_unique_names(self, adata: AnnData, path: Path, lin_key: str, _: int):
-        names_key = _lin_names(lin_key)
-        adata.uns[names_key][0] = adata.uns[names_key][1]
-
-        sc.write(path, adata)
-        with pytest.raises(ValueError):
-            _ = cr.read(path)
-
     @test_fwd()
     def test_wrong_colors_length(
         self, adata: AnnData, path: Path, lin_key: str, n_lins: int
     ):
-        colors_key = _colors(lin_key)
+        colors_key = Key.uns.colors(Key.obs.term_states(False))
         adata.uns[colors_key] = list(adata.uns[colors_key])
         adata.uns[colors_key] += [adata.uns[colors_key][0]]
 
@@ -122,7 +80,7 @@ def test_wrong_colors_length(
     def test_colors_not_colorlike(
         self, adata: AnnData, path: Path, lin_key: str, n_lins: int
     ):
-        colors_key = _colors(lin_key)
+        colors_key = Key.uns.colors(Key.obs.term_states(False))
         adata.uns[colors_key][0] = "foo"
 
         sc.write(path, adata)
@@ -135,15 +93,14 @@ def test_colors_not_colorlike(
 
     @test_fwd()
     def test_normal_run(self, adata: AnnData, path: Path, lin_key: str, n_lins: int):
+        key = Key.obs.term_states(False)
         colors = _create_categorical_colors(10)[-n_lins:]
-        names = [f"foo {i}" for i in range(n_lins)]
-
-        adata.uns[_colors(lin_key)] = colors
-        adata.uns[_lin_names(lin_key)] = names
+        names = adata.obs[key].cat.categories
+        adata.uns[Key.uns.colors(key)] = colors
 
         sc.write(path, adata)
         adata_new = cr.read(path)
         lins_new = adata_new.obsm[lin_key]
 
-        np.testing.assert_array_equal(lins_new.colors, colors)
         np.testing.assert_array_equal(lins_new.names, names)
+        np.testing.assert_array_equal(lins_new.colors, colors)
diff --git a/tests/test_tools.py b/tests/test_tools.py
index 872ba95b9..27e789d9d 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -1,10 +1,9 @@
-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 import Lineage
+from cellrank._key import Key
 from cellrank.tl.kernels._base_kernel import KernelAdd
 
 import numpy as np
@@ -19,7 +18,8 @@ def test_no_root_cells(self, adata: AnnData):
     def test_normal_run(self, adata_cflare):
         cr.tl.lineages(adata_cflare)
 
-        assert str(AbsProbKey.FORWARD) in adata_cflare.obsm.keys()
+        key = Key.obsm.abs_probs(False)
+        assert isinstance(adata_cflare.obsm[key], Lineage)
 
     def test_normal_run_copy(self, adata_cflare):
         adata_cr2 = cr.tl.lineages(adata_cflare, copy=True)
@@ -30,82 +30,35 @@ def test_normal_run_copy(self, adata_cflare):
 
 class TestLineageDrivers:
     def test_no_abs_probs(self, adata: AnnData):
-        with pytest.raises(RuntimeError):
+        with pytest.raises(KeyError):
             cr.tl.lineage_drivers(adata)
 
-    def test_normal_run(self, adata_cflare):
+    @pytest.mark.parametrize("use_raw", [False, True])
+    def test_normal_run(self, adata_cflare: AnnData, use_raw: bool):
+        if use_raw:
+            adata_cflare.raw = adata_cflare.copy()
         cr.tl.lineages(adata_cflare)
-        cr.tl.lineage_drivers(adata_cflare, use_raw=False)
-
-        for name in adata_cflare.obsm[AbsProbKey.FORWARD.s].names:
-            assert np.all(adata_cflare.var[f"to {name} corr"] >= -1.0)
-            assert np.all(adata_cflare.var[f"to {name} corr"] <= 1.0)
+        cr.tl.lineage_drivers(adata_cflare, use_raw=use_raw)
 
-            assert np.all(adata_cflare.var[f"to {name} qval"] >= 0)
-            assert np.all(adata_cflare.var[f"to {name} qval"] <= 1.0)
+        bwd = False
+        key = Key.varm.lineage_drivers(bwd)
+        names = adata_cflare.obsm[Key.obsm.abs_probs(bwd)].names
 
-    def test_normal_run_raw(self, adata_cflare):
-        adata_cflare.raw = adata_cflare.copy()
-        cr.tl.lineages(adata_cflare)
-        cr.tl.lineage_drivers(adata_cflare, use_raw=True)
-
-        for name in adata_cflare.obsm[AbsProbKey.FORWARD.s].names:
-            assert np.all(adata_cflare.raw.var[f"to {name} corr"] >= -1.0)
-            assert np.all(adata_cflare.raw.var[f"to {name} corr"] <= 1.0)
+        if use_raw:
+            adata_cflare = adata_cflare.raw
+        assert isinstance(adata_cflare.varm[key], pd.DataFrame)
+        for name in names:
+            assert np.all(adata_cflare.varm[key][f"{name}_corr"] >= -1.0)
+            assert np.all(adata_cflare.varm[key][f"{name}_corr"] <= 1.0)
 
-            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)
-
-    @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, 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
-
-        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
-            np.testing.assert_allclose(
-                adata_cflare.var[f"to {name} corr"],
-                res[f"{name} corr"].loc[adata_cflare.var_names],
-            )
-
-            assert np.all(res[f"{name} pval"] >= 0)
-            assert np.all(res[f"{name} pval"] <= 1.0)
-
-            assert np.all(adata_cflare.var[f"to {name} qval"] >= 0)
-            assert np.all(adata_cflare.var[f"to {name} qval"] <= 1.0)
-            # res is sorted
-            np.testing.assert_allclose(
-                adata_cflare.var[f"to {name} qval"],
-                res[f"{name} qval"].loc[adata_cflare.var_names],
-            )
-
-            assert np.all(res[f"{name} corr"] >= res[f"{name} ci low"])
-            assert np.all(res[f"{name} corr"] <= res[f"{name} ci high"])
+            assert np.all(adata_cflare.varm[key][f"{name}_qval"] >= 0)
+            assert np.all(adata_cflare.varm[key][f"{name}_qval"] <= 1.0)
 
     def test_invalid_mode(self, adata_cflare: AnnData):
         cr.tl.lineages(adata_cflare)
         with pytest.raises(ValueError):
             cr.tl.lineage_drivers(adata_cflare, use_raw=False, method="foobar")
 
-    def test_invalid_n_perms_type(self, adata_cflare: AnnData):
-        cr.tl.lineages(adata_cflare)
-        with pytest.raises(TypeError):
-            cr.tl.lineage_drivers(
-                adata_cflare, use_raw=False, n_perms=None, method="perm_test"
-            )
-
     def test_invalid_n_perms_value(self, adata_cflare: AnnData):
         cr.tl.lineages(adata_cflare)
         with pytest.raises(ValueError):
@@ -118,7 +71,6 @@ def test_seed_reproducible(self, adata_cflare: AnnData):
         res_a = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             n_perms=10,
             n_jobs=1,
             seed=0,
@@ -127,7 +79,6 @@ def test_seed_reproducible(self, adata_cflare: AnnData):
         res_b = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             n_perms=10,
             n_jobs=1,
             seed=0,
@@ -136,7 +87,6 @@ def test_seed_reproducible(self, adata_cflare: AnnData):
         res_diff_seed = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             n_perms=10,
             n_jobs=1,
             seed=1,
@@ -155,7 +105,6 @@ def test_seed_reproducible_parallel(self, adata_cflare: AnnData):
         res_a = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             n_perms=10,
             n_jobs=2,
             backend="threading",
@@ -165,7 +114,6 @@ def test_seed_reproducible_parallel(self, adata_cflare: AnnData):
         res_b = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             n_perms=10,
             n_jobs=2,
             backend="threading",
@@ -183,33 +131,33 @@ def test_confidence_level(self, adata_cflare: AnnData):
         res_narrow = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             confidence_level=0.95,
         )
         res_wide = cr.tl.lineage_drivers(
             adata_cflare,
             use_raw=False,
-            return_drivers=True,
             confidence_level=0.99,
         )
 
         for name in ["0", "1"]:
-            assert np.all(res_narrow[f"{name} ci low"] >= res_wide[f"{name} ci low"])
-            assert np.all(res_narrow[f"{name} ci high"] <= res_wide[f"{name} ci high"])
+            assert np.all(res_narrow[f"{name}_ci_low"] >= res_wide[f"{name}_ci_low"])
+            assert np.all(res_narrow[f"{name}_ci_high"] <= res_wide[f"{name}_ci_high"])
 
 
 class TestRootFinal:
     def test_find_root(self, adata: AnnData):
         cr.tl.initial_states(adata)
 
-        assert str(TermStatesKey.BACKWARD) in adata.obs.keys()
-        assert _probs(TermStatesKey.BACKWARD) in adata.obs.keys()
+        key = Key.obs.term_states(True)
+        assert key in adata.obs
+        assert Key.obs.probs(key) in adata.obs
 
     def test_find_final(self, adata: AnnData):
         cr.tl.terminal_states(adata, n_states=5, fit_kwargs=dict(n_cells=5))
 
-        assert str(TermStatesKey.FORWARD) in adata.obs.keys()
-        assert _probs(TermStatesKey.FORWARD) in adata.obs.keys()
+        key = Key.obs.term_states(False)
+        assert key in adata.obs
+        assert Key.obs.probs(key) in adata.obs
 
     def test_invalid_cluster_key(self, adata: AnnData):
         with pytest.raises(KeyError):