Repository: theislab/cellrank (language: python)

Issue:
Change the way in which the `backward` mode works
<!-- Describe the feature request in the block below: -->
At the moment, our estimators have a `backward` mode which is indented to study the dynamics of the "backward" Markov process. I would change the interpretation of this and unify our treatment of initial, intermediate, and terminal states:

- the `backward` mode should not change the naming of any attribute or method, nor should it make any method available. Any estimator should have the same methods and attributes available, independent of its direction. 
- the `backward` characteristic of an estimator should thus only act as a description of the estimator; the user themselves should understand to interpret terminal states in the backward process as initial states of the forward process
- initial and terminal states should be treated equally; every estimator should have equivalent methods available for their identification and visualization, e.g. `compute_terminal_states` and `compute_initial_states`, `plot_terminal_states` and `plot_initial_states`, etc. They should also be written to equivalent attributes in the estimator and in `AnnData`. 
- the `predict()` method should compute both initial and terminal states
- we need to discuss the treatment of intermediate states - ideally, the `predict()` method would perform a full classification of states into either initial, terminal, or intermediate. 



Discussion / hints:
Changes:
- introduce a method `classify_states` that controls which states are computed via a `which` parameter (default: `terminal`)
- plotting: use the `plot_macrostates` method, introduce a `which` parameter (`None` by default)
- For manual setting: `set_states_from macrostates` and `set_states` with `which` argument
- The `predict` method is just an alias for `classify_states` (other way round)

Fix patch:
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d60b23530..e20a2160a 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -7,7 +7,7 @@ default_stages:
 minimum_pre_commit_version: 2.9.3
 repos:
 -   repo: https://github.com/psf/black
-    rev: 22.8.0
+    rev: 23.1.0
     hooks:
     -   id: black
         additional_dependencies: [toml]
@@ -23,12 +23,12 @@ repos:
     -   id: yesqa
         additional_dependencies: [flake8-tidy-imports, flake8-docstrings, flake8-comprehensions, flake8-bugbear, flake8-blind-except]
 -   repo: https://github.com/macisamuele/language-formatters-pre-commit-hooks
-    rev: v2.4.0
+    rev: v2.6.0
     hooks:
     -   id: pretty-format-yaml
         args: [--autofix, --indent, '4']
 -   repo: https://github.com/pre-commit/pre-commit-hooks
-    rev: v4.3.0
+    rev: v4.4.0
     hooks:
     -   id: detect-private-key
     -   id: check-merge-conflict
@@ -49,12 +49,12 @@ repos:
     -   id: check-toml
     -   id: requirements-txt-fixer
 -   repo: https://github.com/myint/autoflake
-    rev: v1.6.1
+    rev: v2.0.1
     hooks:
     -   id: autoflake
         args: [--in-place, --remove-unused-variable, --ignore-init-module-imports]
 -   repo: https://github.com/PyCQA/flake8
-    rev: 5.0.4
+    rev: 6.0.0
     hooks:
     -   id: flake8
         additional_dependencies: [flake8-tidy-imports, flake8-docstrings, flake8-comprehensions, flake8-bugbear, flake8-blind-except]
@@ -65,17 +65,17 @@ repos:
         name: Check executable files use .sh extension
         types: [shell, executable]
 -   repo: https://github.com/asottile/blacken-docs
-    rev: v1.12.1
+    rev: 1.13.0
     hooks:
     -   id: blacken-docs
-        additional_dependencies: [black==20.8b1]
+        additional_dependencies: [black==23.1.0]
 -   repo: https://github.com/asottile/pyupgrade
-    rev: v3.0.0
+    rev: v3.3.1
     hooks:
     -   id: pyupgrade
         args: [--py3-plus, --py37-plus]
 -   repo: https://github.com/pre-commit/pygrep-hooks
-    rev: v1.9.0
+    rev: v1.10.0
     hooks:
     -   id: python-no-eval
     -   id: python-use-type-annotations
@@ -83,6 +83,6 @@ repos:
     -   id: rst-directive-colons
     -   id: rst-inline-touching-normal
 -   repo: https://github.com/PyCQA/doc8
-    rev: v1.0.0
+    rev: v1.1.1
     hooks:
     -   id: doc8
diff --git a/README.rst b/README.rst
index be668ffdb..fbe0d635d 100644
--- a/README.rst
+++ b/README.rst
@@ -4,8 +4,8 @@ CellRank for directed single-cell fate mapping
 ==============================================
 
 .. image:: https://raw.githubusercontent.com/theislab/cellrank/master/resources/images/cellrank_fate_map.png
-   :width: 600px
-   :align: center
+    :width: 600px
+    :align: center
 
 .. sidebar:: Key Contributors
 
diff --git a/cellrank/_utils/_docs.py b/cellrank/_utils/_docs.py
index 02812efd6..13cc8cc49 100644
--- a/cellrank/_utils/_docs.py
+++ b/cellrank/_utils/_docs.py
@@ -189,6 +189,13 @@
     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."""
+which = """\
+which
+    Whether to compute initial or terminal states."""
+allow_overlap = """\
+allow_overlap
+    Whether to allow overlapping names between initial and terminal states.
+"""
 
 
 def inject_docs(**kwargs: Any):  # noqa
@@ -245,4 +252,6 @@ def decorator2(obj):
     rw_ixs=_rw_ixs,
     gene_symbols=_gene_symbols,
     absorption_utils=_absorption_utils,
+    which=which,
+    allow_overlap=allow_overlap,
 )
diff --git a/cellrank/_utils/_key.py b/cellrank/_utils/_key.py
index 53fb968d2..5c15af8bd 100644
--- a/cellrank/_utils/_key.py
+++ b/cellrank/_utils/_key.py
@@ -40,8 +40,9 @@ def macrostates(cls, bwd: Optional[bool]) -> str:
             return f"macrostates_{Key.backward(bwd)}"
 
         @classmethod
-        def term_states(cls, bwd: Optional[bool]) -> str:
-            return f"{Key.initial(bwd)}_states"
+        def term_states(cls, estim_bwd: Optional[bool], *, bwd: bool = False) -> str:
+            states = "init_states" if bwd else "term_states"
+            return f"{states}_{Key.backward(estim_bwd)}"
 
         @classmethod
         def priming_degree(cls, bwd: Optional[bool]) -> str:
@@ -62,7 +63,7 @@ def macrostates(cls, bwd: Optional[bool]) -> str:
 
         @classmethod
         def abs_probs(cls, bwd: Optional[bool]) -> str:
-            return Key.where(bwd) + "_" + Key.obs.term_states(bwd)
+            return f"lineages_{Key.backward(bwd)}"
 
         @classmethod
         def abs_times(cls, bwd: Optional[bool]) -> str:
diff --git a/cellrank/_utils/_lineage.py b/cellrank/_utils/_lineage.py
index 2266c2ef9..bd8089e84 100644
--- a/cellrank/_utils/_lineage.py
+++ b/cellrank/_utils/_lineage.py
@@ -858,6 +858,7 @@ def from_adata(
         cls,
         adata: AnnData,
         backward: bool = False,
+        estimator_backward: Optional[bool] = None,
         kind: Literal["macrostates", "term_states", "abs_probs"] = LinKind.ABS_PROBS,
         copy: bool = False,
     ) -> "Lineage":
@@ -886,10 +887,10 @@ def from_adata(
             nkey = Key.obs.macrostates(backward)
             key = Key.obsm.memberships(nkey)
         elif kind == LinKind.TERM_STATES:
-            nkey = Key.obs.term_states(backward)
+            nkey = Key.obs.term_states(estim_bwd=estimator_backward, bwd=backward)
             key = Key.obsm.memberships(nkey)
         elif kind == LinKind.ABS_PROBS:
-            nkey = Key.obs.term_states(backward)
+            nkey = Key.obs.term_states(estim_bwd=estimator_backward, bwd=backward)
             key = Key.obsm.abs_probs(backward)
         else:
             raise NotImplementedError(f"Lineage kind `{kind}` is not yet implemented.")
diff --git a/cellrank/estimators/_base_estimator.py b/cellrank/estimators/_base_estimator.py
index 7f1dbc997..a3be45363 100644
--- a/cellrank/estimators/_base_estimator.py
+++ b/cellrank/estimators/_base_estimator.py
@@ -129,7 +129,7 @@ def _set(
         shadow_only: bool = False,
     ) -> None:
         """
-        Set an attribute and optionally update ``obj[{key}]``.
+        Set an attribute and optionally update ``obj['{key}']``.
 
         Parameters
         ----------
@@ -137,16 +137,16 @@ def _set(
             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`.
+            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.
+            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.
+            Whether to update the ``obj`` if we are not in the shadow.
 
         Returns
         -------
@@ -179,37 +179,35 @@ def _set(
 
     def _get(
         self,
-        attr: str,
+        *,
         obj: Union[pd.DataFrame, Mapping[str, Any]],
         key: str,
-        where: Optional[Literal["obs", "obsm", "var", "varm", "uns"]] = None,
+        shadow_attr: Optional[Literal["obs", "obsm", "var", "varm", "uns"]] = None,
         dtype: Optional[Union[type, Tuple[type, ...]]] = None,
         copy: bool = True,
         allow_missing: bool = False,
-    ) -> None:
+    ) -> Any:
         """
         Get data from an object and set an attribute.
 
         Parameters
         ----------
-        attr
-            Attribute to set.
         obj
             Object from which to extract the data.
         key
             Key in ``obj`` where the data is stored.
-        where
+        shadow_attr
             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``.
+            Copy the data before setting the ``self_attr``.
         allow_missing
-            Whether or not to allow ``key`` to be missing in ``obj``.
+            Whether to allow ``key`` to be missing in ``obj``.
 
         Returns
         -------
-        Nothing, just updates ``attr`` with the extracted values and optionally :attr:`_shadow_adata`.
+        The extracted values and optionally updates :attr:`_shadow_adata` ``.{shadow_attr}``.
 
         Raises
         ------
@@ -220,23 +218,24 @@ def _get(
         KeyError
             If ``allow_missing = False`` and ``key`` was not found in ``obj``.
         """
-        if not hasattr(self, attr):
-            raise AttributeError(attr)
+        if shadow_attr is not None and not hasattr(self._shadow_adata, shadow_attr):
+            raise AttributeError(shadow_attr)
 
         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__}`."
+                    f"Expected object 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
+            if shadow_attr is not None:
+                getattr(self._shadow_adata, shadow_attr)[key] = data
+            return data
         except KeyError:
             if not allow_missing:
                 raise
+            return None
 
     @property
     @contextmanager
@@ -333,8 +332,7 @@ def to_adata(
         *,
         copy: Union[bool, Sequence[Attr_t]] = True,
     ) -> AnnData:
-        """
-        %(to_adata.full_desc)s
+        """%(to_adata.full_desc)s
 
         Parameters
         ----------
@@ -347,12 +345,12 @@ def to_adata(
                   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.
+            Whether to copy the data. Can be specified on per-attribute basis.
+            Useful for attributes that are array-like.
 
         Returns
         -------
-        %(adata)s
+        Annotated data object.
         """  # noqa: D400
 
         def handle_attribute(attr: Attr_t, keys: List[str], *, copy: bool) -> None:
diff --git a/cellrank/estimators/mixins/_absorption_probabilities.py b/cellrank/estimators/mixins/_absorption_probabilities.py
index 37e1d2ff8..2c5fbb898 100644
--- a/cellrank/estimators/mixins/_absorption_probabilities.py
+++ b/cellrank/estimators/mixins/_absorption_probabilities.py
@@ -26,11 +26,12 @@
 from cellrank._utils._lineage import Lineage
 from cellrank._utils._linear_solver import _solve_lin_system
 from cellrank.estimators.mixins._utils import (
+    PlotMode,
     SafeGetter,
     BaseProtocol,
+    StatesHolder,
     logger,
     shadow,
-    register_plotter,
 )
 
 import numpy as np
@@ -52,7 +53,7 @@ class RecTransStates(NamedTuple):
 
 
 class AbsProbsProtocol(BaseProtocol):
-    _term_states_colors: np.ndarray
+    _term_states: StatesHolder
 
     @property
     def transition_matrix(self) -> Union[np.ndarray, spmatrix]:
@@ -101,7 +102,14 @@ def _rec_trans_states(
     ) -> RecTransStates:
         ...
 
-    def _ensure_lineage_object(self, attr: str, **kwargs: Any) -> None:
+    def _ensure_lineage_object(
+        self,
+        obj: Union[str, np.ndarray, Lineage],
+        *,
+        kind: Literal["macrostates", "term_states", "abs_probs"],
+        backward: bool,
+        **kwargs: Any,
+    ) -> Lineage:
         ...
 
     def _write_absorption_probabilities(
@@ -123,6 +131,22 @@ def _write_lineage_priming(
     ) -> str:
         ...
 
+    def _plot_continuous(
+        self,
+        _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:
+        ...
+
 
 def _normalize_abs_times(
     keys: Sequence[str], time_to_absorption: Any = None
@@ -220,8 +244,7 @@ def compute_absorption_probabilities(
         tol: float = 1e-6,
         preconditioner: Optional[str] = None,
     ) -> None:
-        """
-        Compute absorption probabilities.
+        """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`
@@ -270,6 +293,61 @@ def compute_absorption_probabilities(
             time=start,
         )
 
+    @d.dedent
+    def plot_absorption_probabilities(
+        self: AbsProbsProtocol,
+        states: Optional[Union[str, Sequence[str]]] = None,
+        color: Optional[str] = None,
+        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 absorption probabilities.
+
+        Parameters
+        ----------
+        states
+            Subset of the macrostates to show. If ``None``, plot all macrostates.
+        color
+            Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var` used to color the observations.
+        time_key
+            Key in :attr:`anndata.AnnData.obs` where pseudotime is stored. Only used when ``mode = 'time'``.
+        title
+            Title of the plot.
+        same_plot
+            Whether to plot the data on the same plot or not. Only use when ``mode = 'embedding'``.
+            If `True` and ``discrete = False``, ``color`` is ignored.
+        cmap
+            Colormap for continuous annotations.
+        kwargs
+            Keyword arguments for :func:`scvelo.pl.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+        if self.absorption_probabilities is None:
+            raise RuntimeError(
+                "Compute absorption probabilities first as `.compute_absorption_probabilities()`."
+            )
+
+        return self._plot_continuous(
+            _data=self.absorption_probabilities,
+            _colors=self.absorption_probabilities.colors,
+            _title="absorption probabilities",
+            states=states,
+            color=color,
+            mode=mode,
+            time_key=time_key,
+            same_plot=same_plot,
+            title=title,
+            cmap=cmap,
+            **kwargs,
+        )
+
     @d.dedent
     def compute_absorption_times(
         self: AbsProbsProtocol,
@@ -409,7 +487,7 @@ def _rec_trans_states(
 
         # process the current annotations according to `keys`
         term_states, colors = _process_series(
-            series=self.terminal_states, keys=keys, colors=self._term_states_colors
+            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)
@@ -556,43 +634,66 @@ def _write_lineage_priming(
         # fmt: on
 
     def _read_absorption_probabilities(self: AbsProbsProtocol, adata: 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")
+            abs_probs = self._get(
+                obj=adata.obsm,
+                key=key1,
+                shadow_attr="obsm",
+                dtype=(np.ndarray, Lineage),
+            )
+            self._absorption_probabilities = self._ensure_lineage_object(
+                abs_probs, backward=self.backward, kind="abs_probs"
+            )
             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._priming_degree = self._get(
+                obj=adata.obs,
+                key=key,
+                shadow_attr="obs",
+                dtype=pd.Series,
+                allow_missing=True,
+            )
             self.params[key1] = self._read_params(key1)
-        # fmt: on
 
         return sg.ok
 
     def _read_absorption_times(self: AbsProbsProtocol, adata: AnnData) -> bool:
-        # fmt: off
         with SafeGetter(self, allowed=KeyError) as sg:
             key = Key.obsm.abs_times(self.backward)
-            self._get("_absorption_times", self.adata.obsm, key=key, where="obsm", dtype=pd.DataFrame,
-                      allow_missing=True)
+            self._absorption_times = self._get(
+                obj=adata.obsm,
+                key=key,
+                shadow_attr="obsm",
+                dtype=pd.DataFrame,
+                allow_missing=True,
+            )
             self.params[key] = self._read_params(key)
-        # fmt: on
 
         return sg.ok
 
     def _ensure_lineage_object(
-        self: AbsProbsProtocol, attr: str, **kwargs: Any
-    ) -> None:
-        if not isinstance(getattr(self, attr), Lineage):
-            try:
-                lineage = Lineage.from_adata(
-                    self.adata, backward=self.backward, copy=True, **kwargs
-                )
-                setattr(self, attr, lineage)
-            except Exception as e:  # noqa: B902
-                raise RuntimeError(
-                    f"Unable to reconstruct `.absorption_probabilities`. Reason: `{e}`."
-                ) from None
-
-    plot_absorption_probabilities = register_plotter(
-        continuous="absorption_probabilities"
-    )
+        self: AbsProbsProtocol,
+        obj: Union[str, np.ndarray, Lineage],
+        *,
+        kind: Literal["macrostates", "term_states", "abs_probs"],
+        backward: bool,
+        **kwargs: Any,
+    ) -> Lineage:
+        if isinstance(obj, str):
+            obj = getattr(self, obj)
+        if isinstance(obj, Lineage):
+            return obj
+
+        try:
+            return Lineage.from_adata(
+                self.adata,
+                backward=backward,
+                kind=kind,
+                estimator_backward=self.backward,
+                copy=True,
+                **kwargs,
+            )
+        except Exception as e:  # noqa: B902
+            raise RuntimeError(
+                f"Unable to reconstruct `.absorption_probabilities`. Reason: `{e}`."
+            ) from None
diff --git a/cellrank/estimators/mixins/_lineage_drivers.py b/cellrank/estimators/mixins/_lineage_drivers.py
index 1f245ac2c..32c35283d 100644
--- a/cellrank/estimators/mixins/_lineage_drivers.py
+++ b/cellrank/estimators/mixins/_lineage_drivers.py
@@ -584,15 +584,20 @@ def _read_absorption_probabilities(
         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)
+                self._lineage_drivers = self._get(
+                    obj=self.adata.raw.varm,
+                    key=key,
+                    shadow_attr="varm",
+                    dtype=pd.DataFrame,
+                )
             else:
-                self._get("_lineage_drivers", self.adata.varm, key=key, where="varm", dtype=pd.DataFrame)
-        # fmt: on
+                self._lineage_drivers = self._get(
+                    obj=self.adata.varm, key=key, shadow_attr="varm", dtype=pd.DataFrame
+                )
 
         return sg.ok
 
diff --git a/cellrank/estimators/mixins/_utils.py b/cellrank/estimators/mixins/_utils.py
index 5f02d5d66..ba42a9fb0 100644
--- a/cellrank/estimators/mixins/_utils.py
+++ b/cellrank/estimators/mixins/_utils.py
@@ -11,36 +11,40 @@
     Optional,
     Protocol,
     Sequence,
+    NamedTuple,
 )
 
 from copy import copy, deepcopy
 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._utils._docs import d, inject_docs
 from cellrank._utils._enum import ModeEnum
-from cellrank._utils._utils import RandomKeys, _unique_order_preserving
-from cellrank._utils._colors import _create_categorical_colors
 from cellrank._utils._lineage import Lineage
 
 import numpy as np
 import pandas as pd
-from pandas.api.types import infer_dtype, is_categorical_dtype
 
 __all__ = [
     "BaseProtocol",
-    "PlotterProtocol",
     "logger",
     "shadow",
-    "register_plotter",
     "SafeGetter",
+    "StatesHolder",
 ]
 
 
+class StatesHolder(NamedTuple):  # noqa: D101
+    assignment: Optional[pd.Series] = None
+    probs: Optional[pd.Series] = None
+    colors: Optional[np.ndarray] = None
+    memberships: Optional[Lineage] = None
+
+    def set(self, **kwargs: Any) -> "StatesHolder":  # noqa: D102
+        return self._replace(**kwargs)
+
+
 class PlotMode(ModeEnum):
     EMBEDDING = auto()
     TIME = auto()
@@ -96,12 +100,6 @@ 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]
@@ -138,317 +136,6 @@ def shadow(
     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 = True,
-    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.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.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 `discrete`
-        disc = discrete is not 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)
-
-
 class SafeGetter:
     """
     Context manager that does a simple rollback on :attr:`object.__dict__`.
diff --git a/cellrank/estimators/mixins/decomposition/_eigen.py b/cellrank/estimators/mixins/decomposition/_eigen.py
index 91d33c2a6..16f2286fe 100644
--- a/cellrank/estimators/mixins/decomposition/_eigen.py
+++ b/cellrank/estimators/mixins/decomposition/_eigen.py
@@ -374,11 +374,15 @@ def _write_eigendecomposition(
     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._eigendecomposition = self._get(
+                obj=adata.uns,
+                key=key,
+                shadow_attr="uns",
+                dtype=dict,
+                allow_missing=allow_missing,
+            )
             self.params[key] = self._read_params(key)
-        # fmt: on
 
         return sg.ok
diff --git a/cellrank/estimators/mixins/decomposition/_schur.py b/cellrank/estimators/mixins/decomposition/_schur.py
index 5fbb6c296..26f21de74 100644
--- a/cellrank/estimators/mixins/decomposition/_schur.py
+++ b/cellrank/estimators/mixins/decomposition/_schur.py
@@ -63,6 +63,7 @@ def schur_vectors(self) -> Optional[np.ndarray]:
         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
@@ -82,8 +83,8 @@ def schur_matrix(self) -> Optional[np.ndarray]:
     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.
+        For non-symmetric real matrices, left and right eigenvectors will in general be different and complex.
+        We compute both left and right eigenvectors.
 
         Returns
         -------
@@ -308,18 +309,32 @@ def _write_schur_decomposition(
     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)
+            self._eigendecomposition = self._get(
+                obj=adata.uns,
+                key=key,
+                shadow_attr="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)
+            self._schur_vectors = self._get(
+                obj=self.adata.obsm,
+                key=key,
+                shadow_attr="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)
+            self._schur_matrix = self._get(
+                obj=self.adata.uns,
+                key=key,
+                shadow_attr="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/estimators/terminal_states/_cflare.py b/cellrank/estimators/terminal_states/_cflare.py
index 668884d82..732054fc2 100644
--- a/cellrank/estimators/terminal_states/_cflare.py
+++ b/cellrank/estimators/terminal_states/_cflare.py
@@ -52,7 +52,7 @@ def fit(self, k: int = 20, **kwargs: Any) -> "TermStatesEstimator":
 
         Returns
         -------
-        Self and modifies the following field:
+        Returns self and modifies the following field:
 
             - :attr:`eigendecomposition` - %(eigen.summary)s
         """
@@ -77,8 +77,8 @@ def predict(
         """
         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.
+        Filter to obtain recurrent states from left eigenvectors.
+        Cluster to obtain approximate recurrent classes from right eigenvectors.
 
         Parameters
         ----------
@@ -98,7 +98,7 @@ def predict(
         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,
+            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.
@@ -118,7 +118,7 @@ def predict(
 
         Returns
         -------
-        Self and just updates the following fields:
+        Returns self and just updates the following fields:
 
             - :attr:`terminal_states` - %(tse_term_states.summary)s
             - :attr:`terminal_states_probabilities` - %(tse_term_states_probs.summary)s
@@ -218,23 +218,21 @@ def convert_use(
         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
+        # filtering to get rid of some of the leftover 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,
+        return self.set_terminal_states(
+            states=labels,
             cluster_key=cluster_key,
             probs=self._compute_term_states_probs(eig, use),
             params=self._create_params(),
             time=start,
         )
 
-        return self
-
     def _compute_term_states_probs(
         self, eig: Dict[str, Any], use: List[int]
     ) -> pd.Series:
@@ -261,9 +259,9 @@ def _compute_term_states_probs(
 
     def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
         ok = super()._read_from_adata(adata, **kwargs)
+        ok = ok and self._read_eigendecomposition(adata, allow_missing=False)
         return (
             ok
-            and self._read_eigendecomposition(adata, allow_missing=False)
             and self._read_absorption_probabilities(adata)
             and self._read_absorption_times(adata)
         )
diff --git a/cellrank/estimators/terminal_states/_gpcca.py b/cellrank/estimators/terminal_states/_gpcca.py
index 6e1f7a226..2538f674a 100644
--- a/cellrank/estimators/terminal_states/_gpcca.py
+++ b/cellrank/estimators/terminal_states/_gpcca.py
@@ -20,12 +20,7 @@
 from cellrank._utils._lineage import Lineage
 from cellrank.estimators.mixins import EigenMixin, SchurMixin, LinDriversMixin
 from cellrank.kernels._base_kernel import KernelExpression
-from cellrank.estimators.mixins._utils import (
-    SafeGetter,
-    logger,
-    shadow,
-    register_plotter,
-)
+from cellrank.estimators.mixins._utils import SafeGetter, StatesHolder, logger, shadow
 from cellrank.estimators.terminal_states._term_states_estimator import (
     TermStatesEstimator,
 )
@@ -86,40 +81,43 @@ def __init__(
         **kwargs: Any,
     ):
         super().__init__(object=object, **kwargs)
-
+        self._macrostates = StatesHolder()
         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
+        return self._macrostates.assignment
 
     @property
     @d.get_summary(base="gpcca_macro_memberships")
     def macrostates_memberships(self) -> Optional[Lineage]:
-        """Macrostate membership matrix.
+        """Macrostate memberships.
 
         Soft assignment of microstates (cells) to macrostates.
         """
-        return self._macrostates_memberships
+        return self._macrostates.memberships
+
+    @property
+    @d.get_summary(base="gpcca_init_states_memberships")
+    def initial_states_memberships(self) -> Optional[Lineage]:
+        """Initial states memberships.
+
+        Soft assignment of cells to initial states.
+        """
+        return self._init_states.memberships
 
     @property
     @d.get_summary(base="gpcca_term_states_memberships")
     def terminal_states_memberships(self) -> Optional[Lineage]:
-        """Terminal state membership matrix.
+        """Terminal states memberships.
 
         Soft assignment of cells to terminal states.
         """
-        return self._term_states_memberships
+        return self._term_states.memberships
 
     @property
     @d.get_summary(base="gpcca_coarse_init")
@@ -154,8 +152,8 @@ def compute_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.
+            Number of macrostates to compute. 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.
@@ -164,7 +162,7 @@ def compute_macrostates(
 
         Returns
         -------
-        Self and updates the following fields:
+        Returns self and updates the following fields:
 
             - :attr:`macrostates` - %(gpcca_macro.summary)s
             - :attr:`macrostates_memberships` - %(gpcca_macro_memberships.summary)s
@@ -220,7 +218,12 @@ def compute_macrostates(
         return self
 
     @d.dedent
-    def predict(
+    def predict(self, *args: Any, **kwargs: Any) -> "GPCCA":
+        """Alias for :meth:`predict_terminal_states`."""
+        return self.predict_terminal_states(*args, **kwargs)
+
+    @d.dedent
+    def predict_terminal_states(
         self,
         method: Literal[
             "stability", "top_n", "eigengap", "eigengap_coarse"
@@ -229,6 +232,7 @@ def predict(
         alpha: Optional[float] = 1,
         stability_threshold: float = 0.96,
         n_states: Optional[int] = None,
+        allow_overlap: bool = False,
     ) -> "GPCCA":
         """
         Automatically select terminal states from macrostates.
@@ -252,27 +256,32 @@ def predict(
             Threshold used when ``method = 'stability'``.
         n_states
             Number of states used when ``method = 'top_n'``.
+        %(allow_overlap)s
 
         Returns
         -------
-        Self and updates the following fields:
+        Returns self and 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
+            - :attr:`terminal_states_memberships` - %(gpcca_term_states_memberships.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")
-            return self.set_terminal_states_from_macrostates(None, n_cells=n_cells, params=self._create_params())
+        if len(self.macrostates.cat.categories) == 1:
+            logg.warning(
+                "Found only one macrostate, making it the singular terminal state"
+            )
+            return self.set_terminal_states(
+                states=None, n_cells=n_cells, params=self._create_states()
+            )
 
         method = TermStatesMethod(method)
         eig = self.eigendecomposition
         coarse_T = self.coarse_T
 
+        # fmt: off
         if method == TermStatesMethod.EIGENGAP:
             if eig is None:
                 raise RuntimeError("Compute eigendecomposition first as `.compute_eigendecomposition()`.")
@@ -291,43 +300,114 @@ def predict(
                 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
-            return self.set_terminal_states_from_macrostates(names, n_cells=n_cells, params=self._create_params())
+            return self.set_terminal_states(names, n_cells=n_cells, params=self._create_params())
         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 self.set_terminal_states(
+            names,
+            n_cells=n_cells,
+            allow_overlap=allow_overlap,
+            params=self._create_params(),
+        )
+
+    @d.dedent
+    def predict_initial_states(
+        self, n_states: int = 1, n_cells: int = 30, allow_overlap: bool = False
+    ) -> "GPCCA":
+        """
+        Compute initial states from macrostates using :attr:`coarse_stationary_distribution`.
+
+        Parameters
+        ----------
+        n_states
+            Number of initial states.
+        %(n_cells)s
+        %(allow_overlap)s
+
+        Returns
+        -------
+        Returns self and updates the following fields:
+
+            - :attr:`initial_states` - %(tse_init_states.summary)s
+            - :attr:`initial_states_probabilities` - %(tse_init_states_probs.summary)s
+            - :attr:`initial_states_memberships` - %(gpcca_init_states_memberships.summary)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:
+            return self.set_initial_states(n_cells=n_cells, allow_overlap=allow_overlap)
+
+        stat_dist = self.coarse_stationary_distribution
+        if stat_dist is None:
+            raise RuntimeError("No coarse-grained stationary distribution found.")
+
+        states = list(stat_dist[np.argsort(stat_dist)][:n_states].index)
+        return self.set_initial_states(
+            states, n_cells=n_cells, allow_overlap=allow_overlap
         )
-        return self
 
     @d.dedent
-    def set_terminal_states_from_macrostates(
+    def set_terminal_states(
         self,
-        names: Optional[Union[str, Sequence[str], Mapping[str, str]]] = None,
+        states: Optional[
+            Union[str, Sequence[str], Dict[str, Sequence[str]], pd.Series]
+        ] = None,
         n_cells: int = 30,
+        allow_overlap: bool = False,
+        cluster_key: Optional[str] = None,
         **kwargs: Any,
     ) -> "GPCCA":
-        """
-        Manually select terminal states from macrostates.
+        """Set :attr:`terminal_states`.
 
         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.
+        states
+            Which states to select. Valid options are:
+
+                - :class:`str`, :class:`~typing.Sequence` - subset of :attr:`macrostates`. Multiple states can be
+                  combined using ``','``, such as ``['Alpha, Beta', 'Epsilon']``.
+                - :class:`dict` - keys correspond to terminal states and values to cell IDs in
+                  :attr:`~anndata.AnnData.obs_names`.
+                - :class:`~pandas.Series` - categorical series where each category corresponds to a macrostate.
+                  `NaN` values mark cells that should not be marked as :attr:`terminal_states`.
+                - :obj:`None` - select all :attr:`macrostates`.
         %(n_cells)s
+        %(allow_overlap)s
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` to associate names and colors with :attr:`terminal_states`.
+            Each state will be given the name and color corresponding to the cluster it mostly overlaps with.
+            Only used when ``states`` is a :class:`dict` or :class:`~pandas.Series`.
+        kwargs
+            Additional keyword arguments.
 
         Returns
         -------
-        Self and updates the following fields:
+        Returns self and 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
+            - :attr:`terminal_states_memberships` - %(gpcca_term_states_memberships.summary)s
         """
+        if isinstance(states, (dict, pd.Series)):
+            return super().set_terminal_states(
+                states, cluster_key=cluster_key, allow_overlap=allow_overlap, **kwargs
+            )
+
         if n_cells <= 0:
             raise ValueError(f"Expected `n_cells` to be positive, found `{n_cells}`.")
 
@@ -335,78 +415,118 @@ def set_terminal_states_from_macrostates(
         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):
+        if states is None:
+            states = memberships.names
+        elif isinstance(states, str):
+            states = [states]
+        if not len(states):  # unset the states
             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()
+        memberships = memberships[list(states)].copy()
 
         states = self._create_states(memberships, n_cells=n_cells, check_row_sums=False)
         if is_singleton:
-            colors = self._macrostates_colors.copy()
+            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 = (memberships.X / memberships.X.max(axis=0)).max(axis=1)
         probs = pd.Series(probs, index=self.adata.obs_names)
 
-        self._write_terminal_states(
-            states, colors, probs, memberships, params=kwargs.pop("params", {})
+        self._write_states(
+            "terminal",
+            states=states,
+            colors=colors,
+            probs=probs,
+            memberships=memberships,
+            params=kwargs.pop("params", {}),
+            allow_overlap=allow_overlap,
         )
-        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()))
-            )
         return self
 
     @d.dedent
-    def rename_terminal_states(self, new_names: Mapping[str, str]) -> "GPCCA":
-        """
-        %(tse_rename_term_states.full_desc)s
+    def set_initial_states(
+        self,
+        states: Optional[
+            Union[str, Sequence[str], Dict[str, Sequence[str]], pd.Series]
+        ] = None,
+        n_cells: int = 30,
+        allow_overlap: bool = False,
+        cluster_key: Optional[str] = None,
+        **kwargs: Any,
+    ) -> "GPCCA":
+        """Set :attr:`initial_states`.
 
         Parameters
         ----------
-        %(tse_rename_term_states.parameters)s
+        states
+            Which states to select. Valid options are:
+
+                - :class:`str`, :class:`~typing.Sequence` - subset of :attr:`macrostates`. Multiple states can be
+                  combined using ``','``, such as ``['Alpha, Beta', 'Epsilon']``.
+                - :class:`dict` - keys correspond to initial states and values to cell IDs in
+                  :attr:`~anndata.AnnData.obs_names`.
+                - :class:`~pandas.Series` - categorical series where each category corresponds to a macrostate.
+                  `NaN` values mark cells that should not be marked as :attr:`initial_states`.
+                - :obj:`None` - select all :attr:`macrostates`.
+        %(n_cells)s
+        %(allow_overlap)s
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` to associate names and colors with :attr:`initial_states`.
+            Each state will be given the name and color corresponding to the cluster it mostly overlaps with.
+            Only used when ``states`` is a :class:`dict` or :class:`~pandas.Series`.
+        kwargs
+            Additional keyword arguments.
 
         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)
+        Returns self and updates the following fields:
 
-        # 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
+            - :attr:`initial_states` - %(tse_init_states.summary)s
+            - :attr:`initial_states_probabilities` - %(tse_init_states_probs.summary)s
+            - :attr:`initial_states_memberships` - %(gpcca_init_states_memberships.summary)s
+        """
+        if isinstance(states, (pd.Series, dict)):
+            return super().set_initial_states(
+                states, cluster_key=cluster_key, allow_overlap=allow_overlap, **kwargs
+            )
+
+        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()`.")
+
+        if states is None:
+            states = memberships.names
+        elif isinstance(states, str):
+            states = [states]
+        if not len(states):  # unset the states
+            raise ValueError("No macrostates have been selected.")
+
+        is_singleton = memberships.shape[1] == 1
+        memberships = memberships[list(states)].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(axis=0)).max(axis=1)
+        probs = pd.Series(probs, index=self.adata.obs_names)
 
-        with self._shadow:
-            key = Key.obsm.memberships(Key.obs.macrostates(self.backward))
-            self._set(obj=self.adata.obsm, key=key, value=term_states_memberships)
+        self._write_states(
+            "initial",
+            states=states,
+            colors=colors,
+            probs=probs,
+            memberships=memberships,
+            params=kwargs.pop("params", {}),
+            allow_overlap=allow_overlap,
+        )
         return self
 
     @d.dedent
@@ -437,7 +557,7 @@ def fit(
         n = n_states if isinstance(n_states, int) else max(n_states)
         # call explicitly since `compute_macrostates` doesn't handle the case
         # when `minChi` is used for `n_states` and `self._gpcca` is uninitialized
-        self.compute_schur(n, **kwargs)
+        _ = self.compute_schur(n, **kwargs)
         return self.compute_macrostates(
             n_states=n_states, cluster_key=cluster_key, n_cells=n_cells
         )
@@ -753,7 +873,7 @@ def plot_macrostate_composition(
             Position of the legend. If `None`, don't show legend.
         %(plotting)s
         show
-            If `False`, return :class:`matplotlib.pyplot.Axes`.
+            If `False`, return :class:`~matplotlib.axes.Axes`.
 
         Returns
         -------
@@ -968,7 +1088,7 @@ def _set_macrostates(
 
         Returns
         -------
-        Nothing, just updates the field as described in :meth:`compute_macrostates`.
+        Nothing, just updates the fields as described in :meth:`compute_macrostates`.
         """
 
         if n_cells is None:
@@ -977,7 +1097,6 @@ def _set_macrostates(
             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")
@@ -991,9 +1110,10 @@ def _set_macrostates(
             )
 
         # remove previous fields
-        self._write_terminal_states(None, None, None, None, log=False)
-
         # fmt: off
+        self._write_states("initial", states=None, colors=None, probs=None, memberships=None, log=False)
+        self._write_states("terminal", states=None, colors=None, probs=None, memberships=None, log=False)
+
         assignment, colors = self._set_categorical_labels(assignment, cluster_key=cluster_key)
         memberships = Lineage(memberships, names=list(assignment.cat.categories), colors=colors)
         # fmt: on
@@ -1007,7 +1127,11 @@ def _set_macrostates(
             )
 
         self._write_macrostates(
-            assignment, colors, memberships, time=time, params=params
+            macrostates=assignment,
+            colors=colors,
+            memberships=memberships,
+            time=time,
+            params=params,
         )
 
     @logger
@@ -1020,16 +1144,16 @@ def _write_macrostates(
         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)
+        self._set(obj=self.adata.obs, key=key, value=macrostates)
         ckey = Key.uns.colors(key)
-        self._set("_macrostates_colors", obj=self.adata.uns, key=ckey, value=colors, shadow_only=True)
+        self._set(obj=self.adata.uns, key=ckey, value=colors)
         mkey = Key.obsm.memberships(key)
-        self._set("_macrostates_memberships", obj=self.adata.obsm, key=mkey, value=memberships, shadow_only=True)
+        self._set(obj=self.adata.obsm, key=mkey, value=memberships)
+        self._macrostates = self._macrostates.set(assignment=macrostates, colors=colors, memberships=memberships)
         self.params[key] = dict(params)
 
+        names = list(macrostates.cat.categories)
         if len(names) > 1:
             # not using stationary distribution
             g = self._gpcca
@@ -1050,7 +1174,10 @@ def _write_macrostates(
             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))
+            self._set(
+                obj=self.adata.uns, key=Key.uns.coarse(self.backward),
+                value=AnnData(tmat, obs=dists, dtype=float)
+            )
         else:
             for attr in ["_schur_vectors", "_schur_matrix", "_coarse_tmat", "_coarse_init_dist", "_coarse_stat_dist"]:
                 self._set(attr, value=None, shadow_only=True)
@@ -1071,31 +1198,47 @@ def _write_macrostates(
 
     @logger
     @shadow
-    def _write_terminal_states(
+    def _write_states(
         self,
+        which: Literal["initial", "terminal"],
         states: Optional[pd.Series],
         colors: Optional[np.ndarray],
         probs: Optional[pd.Series] = None,
         memberships: Optional[Lineage] = None,
         params: Dict[str, Any] = MappingProxyType({}),
+        allow_overlap: bool = False,
     ) -> str:
+        msg = super()._write_states(
+            which,
+            states=states,
+            colors=colors,
+            probs=probs,
+            params=params,
+            allow_overlap=allow_overlap,
+            log=False,
+        )
         # 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`"
+        msg += f"\n       `.{which}_states_memberships\n    Finish`"
 
-        # TODO(michalk8): CFLARE doesn't remove the downstream properties
         self._write_absorption_probabilities(None, log=False)
+        # TODO(michalk8): CFLARE doesn't remove the downstream properties
         self._write_absorption_times(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)
+
+        backward = which == "initial"
+        key = Key.obsm.memberships(Key.obs.term_states(self.backward, bwd=backward))
+        self._set(obj=self.adata.obsm, key=key, value=memberships)
+        if backward:
+            self._init_states = self._init_states.set(memberships=memberships)
+        else:
+            self._term_states = self._term_states.set(memberships=memberships)
         # fmt: on
 
         return msg
 
     def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
+        # design choice: no need to eigen/Schur-decomposition
         _ = 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
@@ -1104,17 +1247,16 @@ def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
         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)
+            assignment = self._get(obj=adata.obs, key=key, shadow_attr="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))
+            colors = self._get(obj=adata.uns, key=ckey, shadow_attr="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()
+            memberships = self._get(obj=adata.obsm, key=mkey, shadow_attr="obsm", dtype=(Lineage, np.ndarray))
+            memberships = self._ensure_lineage_object(memberships, backward=self.backward, kind="macrostates")
+            self._macrostates = StatesHolder(assignment=assignment, colors=colors, memberships=memberships)
             self.params[key] = self._read_params(key)
 
-            tmat = self.adata.uns[Key.uns.coarse(self.backward)].copy()
+            tmat = adata.uns[Key.uns.coarse(self.backward)].copy()
             if not isinstance(tmat, AnnData):
                 raise TypeError(
                     f"Expected coarse-grained transition matrix to be stored "
@@ -1127,146 +1269,25 @@ def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
 
             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 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")
+        # status is based on `backward=False` by design
+        for backward in [True, False]:
+            with SafeGetter(self, allowed=KeyError) as sg:
+                key = Key.obsm.memberships(Key.obs.term_states(self.backward, bwd=backward))
+                memberships = self._get(obj=adata.obsm, key=key, shadow_attr="obsm", dtype=(np.ndarray, Lineage))
+                memberships = self._ensure_lineage_object(memberships, backward=backward, kind="term_states")
+                if backward:
+                    self._init_states = self._init_states.set(memberships=memberships)
+                else:
+                    self._term_states = self._term_states.set(memberships=memberships)
+                self._set(obj=self._shadow_adata.obsm, key=key, value=memberships)
         # fmt: on
 
         abs_prob_ok = self._read_absorption_probabilities(adata)
         abs_time_ok = self._read_absorption_times(adata)
         return sg.ok and abs_prob_ok and abs_time_ok
-
-    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,
-        )
-
-    def _format_params(self) -> str:
-        fmt = super()._format_params()
-        macro = (
-            None
-            if self.macrostates is None
-            else sorted(self.macrostates.cat.categories)
-        )
-        return fmt + f", macrostates={macro}"
diff --git a/cellrank/estimators/terminal_states/_term_states_estimator.py b/cellrank/estimators/terminal_states/_term_states_estimator.py
index 59243001c..946cb39e7 100644
--- a/cellrank/estimators/terminal_states/_term_states_estimator.py
+++ b/cellrank/estimators/terminal_states/_term_states_estimator.py
@@ -1,12 +1,16 @@
-from typing import Any, Dict, Tuple, Union, Mapping, Optional, Sequence
+from typing import Any, Dict, Tuple, Union, Literal, Optional, Sequence
 
 from abc import ABC
 from types import MappingProxyType
 
+import scvelo as scv
 from anndata import AnnData
+from cellrank import logging as logg
 from cellrank._utils._key import Key
-from cellrank._utils._docs import d
+from cellrank._utils._docs import d, inject_docs
 from cellrank._utils._utils import (
+    RandomKeys,
+    _unique_order_preserving,
     _merge_categorical_series,
     _convert_to_categorical_series,
 )
@@ -15,12 +19,14 @@
     _convert_to_hex_colors,
     _create_categorical_colors,
 )
+from cellrank._utils._lineage import Lineage
 from cellrank.kernels._base_kernel import KernelExpression
 from cellrank.estimators.mixins._utils import (
+    PlotMode,
     SafeGetter,
+    StatesHolder,
     logger,
     shadow,
-    register_plotter,
 )
 from cellrank.estimators._base_estimator import BaseEstimator
 
@@ -50,84 +56,133 @@ def __init__(
         **kwargs: Any,
     ):
         super().__init__(object=object, **kwargs)
-
-        self._term_states: Optional[pd.Series] = None
-        self._term_states_probs: Optional[pd.Series] = None
-        self._term_states_colors: Optional[np.ndarray] = None
+        self._init_states = StatesHolder()
+        self._term_states = StatesHolder()
 
     @property
     @d.get_summary(base="tse_term_states")
     def terminal_states(self) -> Optional[pd.Series]:
-        """
-        Categorical annotation of terminal states.
+        """Categorical annotation of terminal states.
 
-        By default, all cells in transient cells will be labeled as `NaN`.
+        By default, all transient cells will be labeled as `NaN`.
         """
-        return self._term_states
+        return self._term_states.assignment
 
     @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."""
-        return self._term_states_probs
+        """Probability to be a terminal state."""
+        return self._term_states.probs
+
+    @property
+    @d.get_summary(base="tse_init_states")
+    def initial_states(self) -> Optional[pd.Series]:
+        """Categorical annotation of initial states.
+
+        By default, all transient cells will be labeled as `NaN`.
+        """
+        return self._init_states.assignment
+
+    @property
+    @d.get_summary(base="tse_init_states_probs")
+    def initial_states_probabilities(self) -> Optional[pd.Series]:
+        """Probability to be an initial state."""
+        return self._init_states.probs
 
     @d.dedent
     def set_terminal_states(
         self,
-        labels: Union[pd.Series, Dict[str, Sequence[Any]]],
+        states: Union[pd.Series, Dict[str, Sequence[Any]]],
         cluster_key: Optional[str] = None,
-        add_to_existing: bool = False,
+        allow_overlap: bool = False,
         **kwargs: Any,
     ) -> "TermStatesEstimator":
-        """
-        Manually define terminal states.
+        """Set :attr:`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`.
-
+        states
+            States to select. Valid options are:
+
+                - categorical :class:`~pandas.Series` where each category corresponds to an individual state.
+                  `NaN` entries denote cells that do not belong to any state, i.e., transient cells.
+                - :class:`dict` where keys are states and values are lists of cell barcodes corresponding to
+                  annotations in :attr:`anndata.AnnData.obs_names`.
+                  If only 1 key is provided, values should correspond to 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 the 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.
+            Key in :attr:`anndata.AnnData.obs` to associate names and colors with :attr:`terminal_states`.
+            Each state will be given the name and color corresponding to the cluster it mostly overlaps with.
+        %(allow_overlap)s
+        kwargs
+            Additional keyword arguments.
 
         Returns
         -------
-        Self and updates the following fields:
+        Returns self and 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=states,
+            cluster_key=cluster_key,
+            existing=None,
+        )
+        self._write_states(
+            "terminal",
+            states=states,
+            colors=colors,
+            allow_overlap=allow_overlap,
+            **kwargs,
+        )
+        return self
 
+    @d.dedent
+    def set_initial_states(
+        self,
+        states: Union[pd.Series, Dict[str, Sequence[Any]]],
+        cluster_key: Optional[str] = None,
+        allow_overlap: bool = False,
+        **kwargs: Any,
+    ) -> "TermStatesEstimator":
+        """Set :attr:`initial_states`.
+
+        Parameters
+        ----------
+        states
+            Which states to select. Valid options are:
+
+                - categorical :class:`~pandas.Series` where each category corresponds to an individual state.
+                  `NaN` entries denote cells that do not belong to any state, i.e., transient cells.
+                - :class:`dict` where keys are states and values are lists of cell barcodes corresponding to
+                  annotations in :attr:`anndata.AnnData.obs_names`.
+                  If only 1 key is provided, values should correspond to clusters if a categorical
+                  :class:`~pandas.Series` can be found in :attr:`anndata.AnnData.obs`.
+        cluster_key
+            Key in :attr:`anndata.AnnData.obs` to associate names and colors :attr:`initial_states`.
+            Each state will be given the name and color corresponding to the cluster it mostly overlaps with.
+        %(allow_overlap)s
+        kwargs
+            Additional keyword arguments.
+
+        Returns
+        -------
+        Returns self and updates the following fields:
+
+            - :attr:`initial_states` - %(tse_init_states.summary)s
+            - :attr:`initial_states_probabilities` - %(tse_init_states_probs.summary)s
+        """
         states, colors = self._set_categorical_labels(
-            categories=labels,
+            categories=states,
             cluster_key=cluster_key,
-            existing=existing,
+            existing=None,
         )
-        self._write_terminal_states(
-            states,
-            colors,
+        self._write_states(
+            "initial",
+            states=states,
+            colors=colors,
+            allow_overlap=allow_overlap,
             **kwargs,
         )
         return self
@@ -135,64 +190,381 @@ def set_terminal_states(
     @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]
-    ) -> "TermStatesEstimator":
-        """
-        Rename categories in :attr:`terminal_states`.
+    def rename_terminal_states(self, old_new: Dict[str, str]) -> "TermStatesEstimator":
+        """Rename :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.
+        old_new
+            Dictionary that maps old names to unique new names.
 
         Returns
         -------
-        Self and updates the following field:
+        Returns self and updates the following field:
 
             - :attr:`terminal_states` - %(tse_term_states.summary)s
         """
-
-        term_states = self.terminal_states
-        if term_states is None:
+        states = self.terminal_states
+        if states is None:
             raise RuntimeError(
-                "Compute terminal states first as `.compute_terminal_states()` or "
+                "Compute terminal states first as `.predict_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 isinstance(old_new, dict):
+            raise TypeError(f"Expected new names to be a `dict`, found `{type(old_new)}`.")
+        if not len(old_new):
+            return self
+
+        old_names = states.cat.categories
+        mask = np.isin(list(old_new.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)}`")
+            invalid = sorted(np.array(list(old_new.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]
+        names_after_renaming = [old_new.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)
+        assignment = states.cat.rename_categories(old_new)
+        memberships = self._term_states.memberships  # save before overwriting
+        self._write_states(
+            "terminal",
+            states=assignment,
+            colors=self._term_states.colors,
+            probs=self.terminal_states_probabilities,
+            log=False,
+        )
         if memberships is not None:
-            memberships.names = [new_names.get(n, n) for n in memberships.names]
+            memberships.names = [old_new.get(n, n) for n in memberships.names]
+
+        self._term_states = self._term_states.set(
+            assignment=assignment, memberships=memberships
+        )
+        return self
 
-        self._write_terminal_states(
-            self.terminal_states,
-            self._term_states_colors,
-            self.terminal_states_probabilities,
+    @d.dedent
+    def rename_initial_states(self, old_new: Dict[str, str]) -> "TermStatesEstimator":
+        """Rename :attr:`initial_states`.
+
+        Parameters
+        ----------
+        old_new
+            Dictionary that maps old names to unique new names.
+
+        Returns
+        -------
+        Returns self and updates the following field:
+
+            - :attr:`initial_states` - %(tse_init_states.summary)s
+        """
+        states = self.initial_states
+        if states is None:
+            raise RuntimeError(
+                "Compute initial states first as `.predict_initial_states()` or "
+                "set them manually as `.set_initial_states()`."
+            )
+
+        # fmt: off
+        if not isinstance(old_new, dict):
+            raise TypeError(f"Expected new names to be a `dict`, found `{type(old_new)}`.")
+        if not len(old_new):
+            return self
+
+        old_names = states.cat.categories
+        mask = np.isin(list(old_new.keys()), old_names)
+        if not np.all(mask):
+            invalid = sorted(np.array(list(old_new.keys()))[~mask])
+            raise ValueError(f"Invalid terminal states names: `{invalid}`. Valid names are: `{sorted(old_names)}`.")
+
+        names_after_renaming = [old_new.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
+
+        assignment = states.cat.rename_categories(old_new)
+        memberships = self._init_states.memberships  # save overwriting
+        self._write_states(
+            "initial",
+            states=assignment,
+            colors=self._init_states.colors,
+            probs=self.initial_states_probabilities,
             log=False,
         )
+        if memberships is not None:
+            memberships.names = [old_new.get(n, n) for n in memberships.names]
+
+        self._init_states = self._init_states.set(
+            assignment=assignment, memberships=memberships
+        )
         return self
 
+    @d.dedent
+    @inject_docs(m=PlotMode)
+    def plot_macrostates(
+        self,
+        which: Literal["all", "initial", "terminal"],
+        states: Optional[Union[str, Sequence[str]]] = None,
+        color: Optional[str] = None,
+        discrete: bool = True,
+        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 macrostates on an embedding or along pseudotime.
+
+        Parameters
+        ----------
+        which
+            Which macrostates to plot. Valid options are:
+
+                - ``'all'`` - plot all macrostates.
+                - ``'initial'`` - plot macrostates marked as :attr:`initial_states`.
+                - ``'terminal'`` - plot macrostates marked as :attr:`terminal_states`.
+        states
+            Subset of the macrostates to show. If obj:`None`, plot all macrostates.
+        color
+            Key in :attr:`anndata.AnnData.obs` or :attr:`anndata.AnnData.var` used to color the observations.
+        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_key
+            Key in :attr:`anndata.AnnData.obs` where pseudotime is stored. Only used when ``mode = {m.TIME!r}``.
+        title
+            Title of the plot.
+        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 annotations.
+        kwargs
+            Keyword arguments for :func:`scvelo.pl.scatter`.
+
+        Returns
+        -------
+        %(just_plots)s
+        """
+        if which == "all":
+            obj: Optional[StatesHolder] = getattr(self, "_macrostates", None)
+            if obj is None:
+                raise RuntimeError(f"`{type(self).__name__}` cannot plot macrostates.")
+        elif which == "initial":
+            obj = self._init_states
+        elif which == "terminal":
+            obj = self._term_states
+        else:
+            raise ValueError(
+                f"Unable to plot `{which!r}` states. "
+                f"Valid options are: `{['macro', 'initial', 'terminal']}`."
+            )
+
+        name = "macrostates" if which == "macro" else f"{which} states"
+        if obj.assignment is None and obj.memberships is None:
+            raise RuntimeError(f"Compute {name} first.")
+
+        if not discrete and obj.memberships is None:
+            logg.warning(f"Unable to plot {name} in continuous mode, using discrete")
+            discrete = True
+
+        data = obj.assignment if discrete else obj.memberships
+        colors = obj.colors
+
+        if discrete:
+            return self._plot_discrete(
+                _data=data,
+                _colors=colors,
+                _title=name,
+                states=states,
+                color=color,
+                same_plot=same_plot,
+                title=title,
+                cmap=cmap,
+                **kwargs,
+            )
+        return self._plot_continuous(
+            _data=data,
+            _colors=colors,
+            _title=name,
+            states=states,
+            color=color,
+            mode=mode,
+            time_key=time_key,
+            same_plot=same_plot,
+            title=title,
+            cmap=cmap,
+            **kwargs,
+        )
+
+    def _plot_discrete(
+        self,
+        _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)}`."
+            )
+        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,
+        _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)}`."
+            )
+
+        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,
+        )
+
     def _set_categorical_labels(
         self,
         categories: Union[pd.Series, Dict[str, Any]],
@@ -248,26 +620,43 @@ def _set_categorical_labels(
 
     @logger
     @shadow
-    def _write_terminal_states(
+    def _write_states(
         self,
+        which: Literal["initial", "terminal"],
         states: Optional[pd.Series],
         colors: Optional[np.ndarray],
         probs: Optional[pd.Series] = None,
         params: Dict[str, Any] = MappingProxyType({}),
+        allow_overlap: bool = False,
     ) -> 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)
+        backward = which == "initial"
+        if not allow_overlap:
+            fwd_states, bwd_states = (self.terminal_states, states) if backward else (states, self.initial_states)
+            if fwd_states is not None and bwd_states is not None:
+                overlap = np.sum(~pd.isnull(fwd_states) & ~pd.isnull(bwd_states))
+                if overlap:
+                    raise ValueError(
+                        f"Found `{overlap}` overlapping cells between initial and terminal states. "
+                        f"If this is intended, use `allow_overlap=True`."
+                    )
+
+        key = Key.obs.term_states(self.backward, bwd=backward)
+        self._set(obj=self.adata.obs, key=key, value=states)
+        self._set(obj=self.adata.obs, key=Key.obs.probs(key), value=probs)
+        self._set(obj=self.adata.uns, key=Key.uns.colors(key), value=colors)
+        if backward:
+            self._init_states = self._init_states.set(assignment=states, probs=probs, colors=colors)
+        else:
+            self._term_states = self._term_states.set(assignment=states, probs=probs, colors=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"
+            f"       `.{which}_states`\n"
+            f"       `.{which}_states_probabilities`\n"
             "    Finish"
         )
 
@@ -277,52 +666,28 @@ def _read_from_adata(self, adata: AnnData, **kwargs: Any) -> bool:
             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)
+        for backward in [True, False]:
+            key = Key.obs.term_states(self.backward, bwd=backward)
+            with SafeGetter(self, allowed=KeyError) as sg:
+                assignment = self._get(obj=adata.obs, key=key, shadow_attr="obs", dtype=pd.Series)
+                probs = self._get(obj=adata.obs, key=Key.obs.probs(key), shadow_attr="obs", dtype=pd.Series)
+                colors = self._get(obj=adata.uns, key=Key.uns.colors(key), shadow_attr="uns",
+                                   dtype=(list, tuple, np.ndarray))
+                colors = np.asarray([to_hex(c) for c in colors])
+                if backward:
+                    self._init_states = StatesHolder(assignment=assignment, probs=probs, colors=colors)
+                else:
+                    self._term_states = StatesHolder(assignment=assignment, probs=probs, colors=colors)
+                self.params[key] = self._read_params(key)
         # fmt: on
 
+        # status is based on `backward=False` by design
         return sg.ok
 
-    @d.dedent
-    def compute_terminal_states(
-        self, *args: Any, **kwargs: Any
-    ) -> "TermStatesEstimator":
-        """
-        Compute terminal states of the process.
-
-        This is an alias for :meth:`predict`.
-
-        Parameters
-        ----------
-        args
-            Positional arguments.
-        kwargs
-            Keyword arguments arguments.
-
-        Return
-        ------
-        Self and 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"
-    )
-
     def _format_params(self) -> str:
         fmt = super()._format_params()
-        ts = (
-            None
-            if self.terminal_states is None
-            else sorted(self.terminal_states.cat.categories)
-        )
-        return fmt + f", terminal_states={ts}"
+        # fmt: off
+        init_states = None if self.initial_states is None else sorted(self.initial_states.cat.categories)
+        term_states = None if self.terminal_states is None else sorted(self.terminal_states.cat.categories)
+        return fmt + f", initial_states={init_states}, terminal_states={term_states}"
+        # fmt: on
diff --git a/cellrank/external/kernels/_statot_kernel.py b/cellrank/external/kernels/_statot_kernel.py
index 094584090..ab4f6eed0 100644
--- a/cellrank/external/kernels/_statot_kernel.py
+++ b/cellrank/external/kernels/_statot_kernel.py
@@ -57,9 +57,9 @@ def __init__(
         **kwargs: Any,
     ):
         if isinstance(terminal_states, str):
-            adata.obs[Key.obs.term_states(bwd=False)] = adata.obs[terminal_states]
+            adata.obs[Key.obs.term_states(False)] = adata.obs[terminal_states]
         elif isinstance(terminal_states, pd.Series):
-            adata.obs[Key.obs.term_states(bwd=False)] = terminal_states
+            adata.obs[Key.obs.term_states(False)] = terminal_states
 
         try:
             super().__init__(adata, g=g, **kwargs)
diff --git a/cellrank/kernels/utils/_pseudotime_scheme.py b/cellrank/kernels/utils/_pseudotime_scheme.py
index 7da2faaec..d7f728662 100644
--- a/cellrank/kernels/utils/_pseudotime_scheme.py
+++ b/cellrank/kernels/utils/_pseudotime_scheme.py
@@ -49,7 +49,6 @@ def _bias_knn_helper(
         queue=None,
         **kwargs: Any,
     ) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
-
         indices, indptr, data = [], [], []
 
         for i in ixs:
diff --git a/cellrank/pl/_circular_projection.py b/cellrank/pl/_circular_projection.py
index a4c4cabe6..894233efd 100644
--- a/cellrank/pl/_circular_projection.py
+++ b/cellrank/pl/_circular_projection.py
@@ -254,7 +254,6 @@ def circular_projection(
 
     _i, lineage_key = 0, Key.obsm.abs_probs(backward)
     for _i, (k, ax) in enumerate(zip(keys, axes)):
-
         set_lognorm, colorbar = False, kwargs.pop("colorbar", True)
         try:
             _ = PrimingDegree(k)
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 161a888df..567e9216e 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -63,7 +63,7 @@
     "anndata": ("https://anndata.readthedocs.io/en/stable/", None),
     "scanpy": ("https://scanpy.readthedocs.io/en/stable/", None),
     "squidpy": ("https://squidpy.readthedocs.io/en/latest/", None),
-    "scvelo": ("https://scvelo.readthedocs.io/", None),
+    "scvelo": ("https://scvelo.readthedocs.io/en/latest/", None),
     "python": ("https://docs.python.org/3", None),
     "numpy": ("https://docs.scipy.org/doc/numpy/", None),
     "scipy": ("https://docs.scipy.org/doc/scipy/reference/", None),
diff --git a/docs/source/index.rst b/docs/source/index.rst
index 7afedd606..d4d8812a8 100644
--- a/docs/source/index.rst
+++ b/docs/source/index.rst
@@ -4,8 +4,8 @@ CellRank: dynamics from multi-view single-cell data
 ===================================================
 
 .. image:: https://raw.githubusercontent.com/theislab/cellrank/master/docs/source/_static/img/cellrank_overview.png
-   :width: 600px
-   :align: center
+    :width: 600px
+    :align: center
 
 **CellRank** :cite:`lange:22` is a modular framework to study cellular dynamics based on Markov state modeling of
 multi-view single-cell data. It estimates differentiation direction based on a varied (and growing!) number of


Test patch:
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png
deleted file mode 100644
index 93b946bdd..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_not_same.png and /dev/null differ
diff --git a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png b/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png
deleted file mode 100644
index 16ff03e7a..000000000
Binary files a/tests/_ground_truth_figures/scvelo_gpcca_abs_probs_disc_same.png and /dev/null differ
diff --git a/tests/_helpers.py b/tests/_helpers.py
index 7ccda1583..cfa1db6fa 100644
--- a/tests/_helpers.py
+++ b/tests/_helpers.py
@@ -255,13 +255,25 @@ def assert_estimators_equal(
     from_adata: bool = False,
 ) -> None:
     def check_arrays(x, y):
-        if isinstance(x, (np.ndarray, list, tuple)):
+        if isinstance(x, cr.Lineage):
+            check_arrays(x.X, y.X)
+            check_arrays(x.names, y.names)
+            check_arrays(x.colors, y.colors)
+        elif isinstance(x, tuple) and hasattr(x, "_fields") and hasattr(x, "_asdict"):
+            # namedtuple
+            x, y = x._asdict(), y._asdict()
+            assert x.keys() == y.keys()
+            for xx, yy in zip(x.values(), y.values()):
+                check_arrays(xx, yy)
+        elif isinstance(x, pd.Series):
+            assert_series_equal(x, y, check_names=False)
+        elif isinstance(x, (np.ndarray, list, tuple)):
             try:
                 np.testing.assert_array_compare(np.array_equal, x, y, equal_nan=True)
-            except:
+            except AssertionError:
+                raise
+            except Exception:
                 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)
 
@@ -295,7 +307,7 @@ def check_arrays(x, y):
         if attr == "_invalid_n_states" and from_adata:
             continue
         actual_val, expected_val = getattr(actual, attr), getattr(expected, attr)
-        if isinstance(actual_val, cr._utils.Lineage):
+        if isinstance(actual_val, cr.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)):
diff --git a/tests/conftest.py b/tests/conftest.py
index 8aacbd5d2..8913509c2 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -60,7 +60,7 @@ def _create_cflare(*, backward: bool = False) -> Tuple[AnnData, CFLARE]:
     final_kernel.write_to_adata()
 
     mc.compute_eigendecomposition()
-    mc.compute_terminal_states(use=2, method="kmeans")
+    mc.predict(use=2, method="kmeans")
     mc.compute_absorption_probabilities(use_petsc=False)
     mc.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
@@ -87,7 +87,7 @@ def _create_gpcca(*, backward: bool = False) -> Tuple[AnnData, GPCCA]:
     mc.compute_eigendecomposition()
     mc.compute_schur(method="krylov")
     mc.compute_macrostates(n_states=2)
-    mc.set_terminal_states_from_macrostates()
+    mc.set_terminal_states()
     mc.compute_absorption_probabilities(use_petsc=False)
     mc.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
 
diff --git a/tests/test_cflare.py b/tests/test_cflare.py
index d233535e2..620e0032b 100644
--- a/tests/test_cflare.py
+++ b/tests/test_cflare.py
@@ -8,6 +8,7 @@
 from anndata import AnnData
 from cellrank.kernels import VelocityKernel, ConnectivityKernel
 from cellrank._utils._key import Key
+from cellrank.estimators.mixins._utils import StatesHolder
 
 import numpy as np
 import pandas as pd
@@ -43,7 +44,7 @@ def test_compute_terminal_states_no_eig(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         with pytest.raises(RuntimeError):
-            mc.compute_terminal_states(use=2)
+            mc.predict(use=2)
 
     def test_compute_terminal_states_too_large_use(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -53,7 +54,7 @@ def test_compute_terminal_states_too_large_use(self, adata_large: AnnData):
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=2)
         with pytest.raises(ValueError):
-            mc.compute_terminal_states(use=1000)
+            mc.predict(use=1000)
 
     def test_compute_approx_normal_run(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -62,7 +63,7 @@ def test_compute_approx_normal_run(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
 
         assert is_categorical_dtype(mc.terminal_states)
         assert mc.terminal_states_probabilities is not None
@@ -89,7 +90,7 @@ def test_rename_terminal_states_invalid_old_name(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         with pytest.raises(ValueError):
             mc.rename_terminal_states({"foo": "bar"})
 
@@ -100,7 +101,7 @@ def test_rename_terminal_states_invalid_new_name(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2, method="kmeans")
+        mc.predict(use=2, method="kmeans")
         with pytest.raises(ValueError):
             mc.rename_terminal_states({"0": "1"})
 
@@ -111,7 +112,7 @@ def test_rename_terminal_states_try_joining_states(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         with pytest.raises(ValueError):
             mc.rename_terminal_states({"0, 1": "foo"})
 
@@ -122,10 +123,10 @@ def test_rename_terminal_states_empty_mapping(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
 
         orig_cats = list(mc.terminal_states.cat.categories)
-        mc.rename_terminal_states({})
+        mc = mc.rename_terminal_states({})
 
         np.testing.assert_array_equal(mc.terminal_states.cat.categories, orig_cats)
 
@@ -136,7 +137,7 @@ def test_rename_terminal_states_normal_run(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2, method="kmeans")
+        mc.predict(use=2, method="kmeans")
         mc.rename_terminal_states({"0": "foo", "1": "bar"})
 
         np.testing.assert_array_equal(mc.terminal_states.cat.categories, ["foo", "bar"])
@@ -161,7 +162,7 @@ def test_compute_absorption_probabilities_normal_run(self, adata_large: AnnData)
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2, method="kmeans")
+        mc.predict(use=2, method="kmeans")
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
@@ -187,7 +188,7 @@ def test_compute_absorption_probabilities_normal_run(self, adata_large: AnnData)
         np.testing.assert_array_equal(
             mc.absorption_probabilities.colors, mc.adata.uns[key]
         )
-        np.testing.assert_array_equal(mc._term_states_colors, mc.adata.uns[key])
+        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):
@@ -198,7 +199,7 @@ def test_compute_absorption_probabilities_solver(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
 
         # compute lin probs using direct solver
         mc.compute_absorption_probabilities(solver="direct")
@@ -219,7 +220,7 @@ def test_compute_absorption_probabilities_solver_petsc(self, adata_large: AnnDat
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2, method="kmeans")
+        mc.predict(use=2, method="kmeans")
 
         # compute lin probs using direct solver
         mc.compute_absorption_probabilities(solver="gmres", use_petsc=False, tol=tol)
@@ -269,7 +270,7 @@ def test_compute_lineage_drivers_no_lineages(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         with pytest.raises(RuntimeError):
             mc.compute_lineage_drivers()
 
@@ -280,7 +281,7 @@ def test_compute_lineage_drivers_invalid_lineages(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
         with pytest.raises(KeyError):
             mc.compute_lineage_drivers(use_raw=False, lineages=["foo"])
@@ -292,7 +293,7 @@ def test_compute_lineage_drivers_invalid_clusters(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
         with pytest.raises(KeyError):
             mc.compute_lineage_drivers(
@@ -306,7 +307,7 @@ def test_compute_lineage_drivers_normal_run(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2, method="kmeans")
+        mc.predict(use=2, method="kmeans")
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -325,7 +326,7 @@ def test_plot_lineage_drivers_not_computed(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
 
         with pytest.raises(RuntimeError):
@@ -338,7 +339,7 @@ def test_plot_lineage_drivers_invalid_name(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -352,7 +353,7 @@ def test_plot_lineage_drivers_invalid_n_genes(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -366,7 +367,7 @@ def test_plot_lineage_drivers_normal_run(self, adata_large: AnnData):
 
         mc = cr.estimators.CFLARE(terminal_kernel)
         mc.compute_eigendecomposition(k=5)
-        mc.compute_terminal_states(use=2)
+        mc.predict(use=2)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
         mc.plot_lineage_drivers("0", use_raw=False)
@@ -379,13 +380,13 @@ def test_compute_absorption_probabilities_keys_colors(self, adata_large: AnnData
 
         mc_fwd = cr.estimators.CFLARE(terminal_kernel)
         mc_fwd.compute_eigendecomposition()
-        mc_fwd.compute_terminal_states(use=3, method="kmeans")
+        mc_fwd.predict(use=3, method="kmeans")
 
         arcs = ["0", "2"]
         arc_colors = [
             c
             for arc, c in zip(
-                mc_fwd.terminal_states.cat.categories, mc_fwd._term_states_colors
+                mc_fwd.terminal_states.cat.categories, mc_fwd._term_states.colors
             )
             if arc in arcs
         ]
@@ -436,8 +437,9 @@ def test_compare_absorption_probabilities_with_reference(self):
         state_annotation[7] = "terminal_1"
         state_annotation[10] = "terminal_2"
         state_annotation = state_annotation.astype("category")
-        c._term_states = state_annotation
-        c._term_states_colors = np.array(["#000000", "#ffffff"])
+        c._term_states = StatesHolder(
+            assignment=state_annotation, colors=np.array(["#000000", "#ffffff"])
+        )
 
         # compute absorption probabilities
         c.compute_absorption_probabilities()
@@ -457,7 +459,7 @@ def test_manual_approx_rc_set(self, adata_large):
         mc_fwd.compute_eigendecomposition()
         key = Key.obs.term_states(mc_fwd.backward)
 
-        mc_fwd.compute_terminal_states(use=3)
+        mc_fwd.predict(use=3)
         original = np.array(adata.obs[key].copy())
         zero_mask = original == "0"
 
diff --git a/tests/test_gpcca.py b/tests/test_gpcca.py
index 09556853b..c6f142bee 100644
--- a/tests/test_gpcca.py
+++ b/tests/test_gpcca.py
@@ -62,7 +62,7 @@ def key(self) -> Optional[str]:
         if self.value == "macro":
             return Key.obs.macrostates(bwd)
         if self.value == "term":
-            return Key.obs.term_states(bwd)
+            return Key.obs.term_states(False, bwd=bwd)
         if self.value == "abs":
             return Key.obsm.abs_probs(bwd)
         if self.value == "lin":
@@ -83,7 +83,7 @@ def attr_keys(self) -> Optional[Sequence[Tuple[str, str]]]:
             key3 = Key.obsm.memberships(key1)
             return ("obs", key1), ("uns", key2), ("obsm", key3)
         if self.value == "term":
-            key1 = Key.obs.term_states(bwd)
+            key1 = Key.obs.term_states(False, bwd=bwd)
             key2 = Key.obs.probs(key1)
             key3 = Key.uns.colors(key1)
             return ("obs", key1), ("obs", key2), ("uns", key3)
@@ -103,12 +103,11 @@ def attrs(self) -> Optional[Sequence[Tuple[str, type]]]:
             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),
+                ("macrostates", pd.Series), ("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)
+            return ("terminal_states", pd.Series), ("terminal_states_probabilities", pd.Series)
         if self.value == "abs":
             return (
                 ("_absorption_probabilities", Lineage), ("_absorption_times", pd.DataFrame),
@@ -145,7 +144,7 @@ def _check_eigdecomposition(mc: cr.estimators.GPCCA) -> None:
 
 def _check_compute_macro(mc: cr.estimators.GPCCA) -> None:
     assert isinstance(mc.macrostates, pd.Series)
-    assert len(mc._macrostates_colors) == len(mc.macrostates.cat.categories)
+    assert len(mc._macrostates.colors) == len(mc.macrostates.cat.categories)
 
     if "stationary_dist" in mc.eigendecomposition:  # one state
         assert isinstance(mc.macrostates_memberships, cr.Lineage)
@@ -193,7 +192,7 @@ def _check_abs_probs(mc: cr.estimators.GPCCA) -> None:
     # macrostates
     assert isinstance(mc.macrostates, pd.Series)
     assert isinstance(mc.macrostates_memberships, cr.Lineage)
-    np.testing.assert_array_equal(mc._macrostates_colors, mc.macrostates_memberships.colors)
+    np.testing.assert_array_equal(mc._macrostates.colors, mc.macrostates_memberships.colors)
 
     # term states
     key = Key.obs.term_states(mc.backward)
@@ -201,7 +200,7 @@ def _check_abs_probs(mc: cr.estimators.GPCCA) -> None:
     # 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)
+    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)
 
@@ -234,7 +233,7 @@ def _fit_gpcca(adata, state: str, backward: bool = False) -> cr.estimators.GPCCA
     mc.compute_macrostates(n_states=2)
     if state == State.MACRO:
         return mc
-    mc.set_terminal_states_from_macrostates()
+    mc.set_terminal_states()
     if state == State.TERM:
         return mc
     mc.compute_absorption_probabilities()
@@ -284,12 +283,12 @@ def _assert_adata(
         else:
             for attr, key in state.attr_keys:
                 obj = getattr(adata, attr)
-                assert key not in obj, sorted(obj.keys())
+                assert key not in obj, (state, attr, key)
         _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 key in obj, (state, attr, key)
         _assert_adata(adata, state.prev, fwd=False)
 
 
@@ -513,20 +512,15 @@ def test_set_initial_states_from_forward(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2, n_cells=5)
-        obsm_keys = set(mc.adata.obsm.keys())
-        # TODO
-        mc._set_initial_states_from_macrostates("0")
+        mc.set_initial_states("0")
 
-        key = Key.obs.term_states(not mc.backward)
+        key = Key.obs.term_states(mc.backward, bwd=True)
 
         assert key in mc.adata.obs
         np.testing.assert_array_equal(mc.adata.obs[key].cat.categories, ["0"])
         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
-
     def test_compute_initial_states_from_forward_no_macro(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
             softmax_scale=4
@@ -537,13 +531,10 @@ def test_compute_initial_states_from_forward_no_macro(self, adata_large: AnnData
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
 
-        with pytest.raises(RuntimeError):
-            # TODO
-            mc._compute_initial_states(1)
+        with pytest.raises(RuntimeError, match=r"Compute macro"):
+            mc.predict_initial_states(n_states=1)
 
-    def test_compute_initial_states_from_forward_too_many_states(
-        self, adata_large: AnnData
-    ):
+    def test_compute_initial_states_from_forward(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
@@ -551,43 +542,12 @@ def test_compute_initial_states_from_forward_too_many_states(
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
 
-        mc.compute_macrostates(n_states=2, n_cells=5)
-        with pytest.raises(ValueError):
-            mc._compute_initial_states(42)
-
-    def test_compute_initial_states_from_forward_too_few_states(
-        self, adata_large: AnnData
-    ):
-        vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
-            softmax_scale=4
-        )
-        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
-        terminal_kernel = 0.8 * vk + 0.2 * ck
-
-        mc = cr.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=2, n_cells=5)
-        with pytest.raises(ValueError):
-            mc._compute_initial_states(0)
+        mc.compute_macrostates(n_states=4, n_cells=5)
+        mc.predict_initial_states(n_states=3)
 
-    def test_compute_initial_states_from_forward_no_stat_dist(
-        self, adata_large: AnnData
-    ):
-        vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
-            softmax_scale=4
-        )
-        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
-        terminal_kernel = 0.8 * vk + 0.2 * ck
-
-        mc = cr.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=2, n_cells=5)
-        mc._coarse_stat_dist = None
-
-        with pytest.raises(ValueError):
-            mc._compute_initial_states(0)
+        assert mc.terminal_states is None
+        assert len(mc.initial_states.cat.categories) == 3
+        assert mc.initial_states_memberships.shape == (adata_large.n_obs, 3)
 
     def test_compute_initial_states_from_forward_normal_run(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large, backward=False).compute_transition_matrix(
@@ -600,22 +560,16 @@ def test_compute_initial_states_from_forward_normal_run(self, adata_large: AnnDa
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2, n_cells=5)
-        obsm_keys = set(mc.adata.obsm.keys())
-        csd = mc.coarse_stationary_distribution
-        expected = csd.index[np.argmin(csd)]
 
-        mc._compute_initial_states(1)
+        mc.predict_initial_states(n_states=1)
 
-        key = Key.obs.term_states(not mc.backward)
+        key = Key.obs.term_states(mc.backward, bwd=True)
 
         assert key in mc.adata.obs
-        np.testing.assert_array_equal(mc.adata.obs[key].cat.categories, [expected])
+        np.testing.assert_array_equal(mc.adata.obs[key].cat.categories, ["0"])
         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
-
     def test_set_terminal_states_from_macrostates(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
@@ -625,111 +579,12 @@ def test_set_terminal_states_from_macrostates(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2, n_cells=5)
-        mc.set_terminal_states_from_macrostates()
-        mc.compute_absorption_probabilities()
-        mc.compute_lineage_priming()
-
-        _check_abs_probs(mc)
-
-    def test_set_terminal_states_from_macrostates_rename_states_invalid(
-        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.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=2, n_cells=5)
-        with pytest.raises(ValueError):
-            mc.set_terminal_states_from_macrostates({"0": "1"})
-        _check_renaming_no_write_terminal(mc)
-
-    def test_set_terminal_states_from_macrostates_rename_not_unique_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.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=3, n_cells=5)
-        with pytest.raises(ValueError):
-            mc.set_terminal_states_from_macrostates({"0": "1", "1": "1"})
-        _check_renaming_no_write_terminal(mc)
-
-    def test_set_terminal_states_from_macrostates_rename_overlapping_old_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.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=3, n_cells=5)
-        with pytest.raises(ValueError):
-            mc.set_terminal_states_from_macrostates({"0, 1": "foo", "1, 2": "bar"})
-        _check_renaming_no_write_terminal(mc)
-
-    def test_set_terminal_states_from_macrostates_rename_states(
-        self, adata_large: AnnData
-    ):
-        vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
-        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
-        terminal_kernel = 0.8 * vk + 0.2 * ck
-
-        mc = cr.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=2, n_cells=5)
-        mc.set_terminal_states_from_macrostates({"0": "foo"})
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
         _check_abs_probs(mc)
 
-        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
-    ):
-        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.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=3, n_cells=5)
-        mc.set_terminal_states_from_macrostates({"0, 1": "foo", "2": "bar"})
-        mc.compute_absorption_probabilities()
-        mc.compute_lineage_priming()
-
-        _check_abs_probs(mc)
-
-        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)
-        ck = ConnectivityKernel(adata_large).compute_transition_matrix()
-        terminal_kernel = 0.8 * vk + 0.2 * ck
-
-        mc = cr.estimators.GPCCA(terminal_kernel)
-        mc.compute_schur(n_components=10, method="krylov")
-
-        mc.compute_macrostates(n_states=2, n_cells=None)
-        with pytest.raises(TypeError):
-            mc.set_terminal_states_from_macrostates(n_cells=None)
-
     def test_set_terminal_states_from_macrostates_non_positive_cells(
         self, adata_large: AnnData
     ):
@@ -742,7 +597,7 @@ def test_set_terminal_states_from_macrostates_non_positive_cells(
 
         mc.compute_macrostates(n_states=2, n_cells=None)
         with pytest.raises(ValueError):
-            mc.set_terminal_states_from_macrostates(n_cells=0)
+            mc.set_terminal_states(n_cells=0)
 
     def test_set_terminal_states_from_macrostates_invalid_name(
         self, adata_large: AnnData
@@ -756,7 +611,7 @@ def test_set_terminal_states_from_macrostates_invalid_name(
 
         mc.compute_macrostates(n_states=2)
         with pytest.raises(KeyError):
-            mc.set_terminal_states_from_macrostates(names=["foobar"])
+            mc.set_terminal_states(states=["foobar"])
 
     @pytest.mark.parametrize("values", ["Astrocytes", ["Astrocytes", "OPC"]])
     def test_set_terminal_states_clusters(
@@ -789,7 +644,7 @@ def test_compute_terminal_states_invalid_method(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2)
         with pytest.raises(ValueError):
-            mc.compute_terminal_states(method="foobar")
+            mc.predict(method="foobar")
 
     def test_compute_terminal_states_no_cells(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -801,7 +656,7 @@ def test_compute_terminal_states_no_cells(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2)
         with pytest.raises(TypeError):
-            mc.compute_terminal_states(n_cells=None)
+            mc.predict(n_cells=None)
 
     def test_compute_terminal_states_non_positive_cells(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -813,7 +668,7 @@ def test_compute_terminal_states_non_positive_cells(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2)
         with pytest.raises(ValueError):
-            mc.compute_terminal_states(n_cells=0)
+            mc.predict(n_cells=0)
 
     def test_compute_terminal_states_eigengap(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -824,7 +679,7 @@ def test_compute_terminal_states_eigengap(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2)
-        mc.compute_terminal_states(n_cells=5, method="eigengap")
+        mc.predict(n_cells=5, method="eigengap")
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
@@ -839,7 +694,7 @@ def test_compute_terminal_states_n_main_states(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2)
-        mc.compute_terminal_states(n_cells=5, method="top_n", n_states=1)
+        mc.predict(n_cells=5, method="top_n", n_states=1)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
@@ -855,9 +710,7 @@ def test_compute_terminal_states_stability(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=5)
-        mc.compute_terminal_states(
-            n_cells=5, method="stability", stability_threshold=thresh
-        )
+        mc.predict(n_cells=5, method="stability", stability_threshold=thresh)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
@@ -879,10 +732,8 @@ def test_compute_terminal_states_no_selected(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2)
-        with pytest.raises(ValueError):
-            mc.compute_terminal_states(
-                n_cells=5, method="stability", stability_threshold=42
-            )
+        with pytest.raises(ValueError, match=r"No macrostates"):
+            mc.predict(n_cells=5, method="stability", stability_threshold=42)
 
     def test_compute_terminal_states_too_many_cells(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -894,7 +745,7 @@ def test_compute_terminal_states_too_many_cells(self, adata_large: AnnData):
 
         mc.compute_macrostates(n_states=2)
         with pytest.raises(ValueError):
-            mc.compute_terminal_states(n_cells=4200)
+            mc.predict(n_cells=4200)
 
     def test_compute_terminal_states_default(self, adata_large: AnnData):
         vk = VelocityKernel(adata_large).compute_transition_matrix(softmax_scale=4)
@@ -905,7 +756,7 @@ def test_compute_terminal_states_default(self, adata_large: AnnData):
         mc.compute_schur(n_components=10, method="krylov")
 
         mc.compute_macrostates(n_states=2)
-        mc.compute_terminal_states(n_cells=5)
+        mc.predict(n_cells=5)
         mc.compute_absorption_probabilities()
         mc.compute_lineage_priming()
 
@@ -919,7 +770,7 @@ def test_compute_lineage_drivers_invalid_lineages(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
 
         with pytest.raises(KeyError):
@@ -933,7 +784,7 @@ def test_compute_lineage_drivers_invalid_clusters(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
 
         with pytest.raises(KeyError):
@@ -949,7 +800,7 @@ def test_compute_lineage_drivers_normal_run(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -969,7 +820,7 @@ def test_plot_lineage_drivers_not_computed(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
 
         with pytest.raises(RuntimeError):
@@ -983,7 +834,7 @@ def test_plot_lineage_drivers_invalid_name(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -998,7 +849,7 @@ def test_plot_lineage_drivers_invalid_n_genes(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -1013,7 +864,7 @@ def test_plot_lineage_drivers_normal_run(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
         mc.compute_lineage_drivers(use_raw=False, cluster_key="clusters")
 
@@ -1027,7 +878,7 @@ def test_compute_priming_clusters(self, adata_large: AnnData):
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
         mc.compute_absorption_probabilities()
 
         cat = adata_large.obs["clusters"].cat.categories[0]
@@ -1053,9 +904,9 @@ def test_recompute_terminal_states_different_n_states(self, adata_large: AnnData
         mc = cr.estimators.GPCCA(terminal_kernel)
         mc.compute_schur(n_components=10, method="krylov")
         mc.compute_macrostates(n_states=2)
-        mc.set_terminal_states_from_macrostates()
+        mc.set_terminal_states()
 
-        key = Key.obsm.memberships(Key.obs.term_states(mc.backward))
+        key = Key.obsm.memberships(Key.obs.term_states(estim_bwd=mc.backward))
         assert len(mc.terminal_states.cat.categories) == 2
         assert adata_large.obsm[key].shape == (adata_large.n_obs, 2)
 
@@ -1112,7 +963,7 @@ def test_compute_time_to_absorption(
         ck = ConnectivityKernel(adata_large).compute_transition_matrix()
         g = cr.estimators.GPCCA(ck)
         g.fit(n_states=n_states)
-        g.set_terminal_states_from_macrostates()
+        g.set_terminal_states()
         g.compute_absorption_probabilities()
         g.compute_absorption_times(keys=keys)
 
@@ -1128,7 +979,7 @@ def test_to_adata(self, adata_large: AnnData, state: State):
         _assert_adata(adata, state, fwd=True)
 
     @pytest.mark.parametrize("copy", [False, True])
-    @pytest.mark.parametrize("keep", ["X", ("obs", "obsm"), ("layers")])
+    @pytest.mark.parametrize("keep", ["X", ("obs", "obsm"), ("layers",)])
     def test_to_adata_keep(
         self, adata_large: AnnData, keep: Union[str, Sequence[str]], copy: bool
     ):
@@ -1138,12 +989,10 @@ def test_to_adata_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])
+            columns = adata.obs.columns
+            np.testing.assert_array_equal(adata.obs.shape, g.adata.obs.shape)
+            # reorder columns
+            assert_frame_equal(adata.obs, g.adata.obs[columns])
         if "obsm" in keep:
             for key in g.adata.obsm.keys():
                 res = shares_mem(adata.obsm[key], g.adata.obsm[key])
@@ -1174,10 +1023,11 @@ def test_from_adata(self, adata_large: AnnData, state: State):
         _assert_adata(g2.adata, state, fwd=True)
         assert_estimators_equal(g1, g2, from_adata=True)
 
+    # TODO(michalk8): parametrize by bwd, needs attr_keys modification
     @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")
+        if state == State.SCHUR:
+            pytest.xfail("Schur decomposition is not needed.")
         g_orig = _fit_gpcca(adata_large, State.LIN)
         adata = g_orig.to_adata()
 
@@ -1209,19 +1059,75 @@ def test_reconstruct_lineage(self, adata_large: AnnData, bwd: bool):
 
         for key in keys:
             assert isinstance(adata.obsm[key], Lineage)
-            adata.obsm[key] = np.array(adata.obsm[key])
+            adata.obsm[key] = np.array(adata.obsm[key], copy=True)
             assert isinstance(adata.obsm[key], np.ndarray)
 
         g = cr.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)
+            assert isinstance(obj, Lineage), attr
             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)
 
+    def test_overlapping_states_predict(self, adata_large: AnnData):
+        g = _fit_gpcca(adata_large, State.MACRO, backward=False)
+        n_states = g.macrostates_memberships.shape[1]
+
+        g.predict_initial_states(n_states=n_states)
+        with pytest.raises(ValueError, match=r"Found \`\d+\` overlapping"):
+            g.predict_terminal_states(method="top_n", n_states=n_states)
+
+        g.predict_terminal_states(method="top_n", n_states=n_states, allow_overlap=True)
+
+        pd.testing.assert_series_equal(g.initial_states, g.terminal_states)
+
+    def test_overlapping_states_from_macro(self, adata_large: AnnData):
+        g = _fit_gpcca(adata_large, State.TERM, backward=False)
+        initial_states = g.terminal_states.cat.categories
+
+        with pytest.raises(ValueError, match=r"Found \`\d+\` overlapping"):
+            g.set_initial_states(initial_states)
+        assert g.initial_states is None
+
+        g.set_initial_states(initial_states, allow_overlap=True)
+
+        pd.testing.assert_series_equal(g.initial_states, g.terminal_states)
+
+    def test_overlapping_states_explicit(self, adata_large: AnnData):
+        g = _fit_gpcca(adata_large, State.TERM, backward=False)
+        initial_states = {"overlapping_initial_state": g.adata.obs_names[:20]}
+
+        with pytest.raises(ValueError, match=r"Found \`\d+\` overlapping"):
+            g.set_initial_states(initial_states)
+        assert g.initial_states is None
+
+        g.set_initial_states(initial_states, allow_overlap=True)
+        assert g.initial_states is not None
+
+    @pytest.mark.parametrize("initial", [False, True])
+    def test_rename_states_normal_run(self, adata_large: AnnData, initial: bool):
+        g = _fit_gpcca(adata_large, State.MACRO)
+        expected = ["foo", "bar"]
+
+        if initial:
+            g = g.set_initial_states(states=None)
+            g = g.rename_initial_states({"0": expected[0], "1": expected[1]})
+            np.testing.assert_array_equal(g.initial_states.cat.categories, expected)
+            np.testing.assert_array_equal(g.initial_states_memberships.names, expected)
+        else:
+            g = g.set_terminal_states(states=None)
+            g = g.rename_terminal_states({"0": expected[0], "1": expected[1]})
+            np.testing.assert_array_equal(g.terminal_states.cat.categories, expected)
+            np.testing.assert_array_equal(g.terminal_states_memberships.names, expected)
+
+        np.testing.assert_array_equal(
+            g.adata.obs[Key.obs.term_states(g.backward, bwd=initial)].cat.categories,
+            expected,
+        )
+
 
 class TestGPCCAIO:
     @pytest.mark.parametrize("deep", [False, True])
diff --git a/tests/test_pipeline.py b/tests/test_pipeline.py
index 44df6d801..232820110 100644
--- a/tests/test_pipeline.py
+++ b/tests/test_pipeline.py
@@ -18,7 +18,7 @@ def _assert_has_all_keys(adata: AnnData, bwd: bool = False) -> None:
 
     # lineages
     abs_probs = adata.obsm[Key.obsm.abs_probs(bwd)]
-    assert isinstance(abs_probs, cr._utils.Lineage)
+    assert isinstance(abs_probs, cr.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)
@@ -40,8 +40,8 @@ def test_fwd_pipeline_cflare(self, adata: AnnData):
         estimator_fwd.plot_spectrum()
         estimator_fwd.plot_spectrum(real_only=True)
 
-        estimator_fwd.compute_terminal_states(use=1, method="leiden")
-        estimator_fwd.plot_terminal_states()
+        estimator_fwd.predict(use=1, method="leiden")
+        estimator_fwd.plot_macrostates(which="terminal")
 
         estimator_fwd.compute_absorption_probabilities()
         estimator_fwd.plot_absorption_probabilities()
@@ -63,8 +63,8 @@ def test_bwd_pipeline_cflare(self, adata: AnnData):
         estimator_bwd.plot_spectrum()
         estimator_bwd.plot_spectrum(real_only=True)
 
-        estimator_bwd.compute_terminal_states(use=1, method="kmeans")
-        estimator_bwd.plot_terminal_states()
+        estimator_bwd.predict(use=1, method="kmeans")
+        estimator_bwd.plot_macrostates(which="terminal")
 
         estimator_bwd.compute_absorption_probabilities()
         estimator_bwd.plot_absorption_probabilities()
@@ -87,13 +87,13 @@ def test_fwd_pipeline_gpcca(self, adata: AnnData):
         estimator_fwd.compute_schur(5, method="brandts")
 
         estimator_fwd.compute_macrostates(3, n_cells=10)
-        estimator_fwd.plot_macrostates()
+        estimator_fwd.plot_macrostates(which="all")
         estimator_fwd.plot_coarse_T(show_initial_dist=True, show_stationary_dist=True)
         estimator_fwd.plot_schur_matrix()
 
         # select all states
-        estimator_fwd.set_terminal_states_from_macrostates(n_cells=10)
-        estimator_fwd.plot_terminal_states()
+        estimator_fwd.set_terminal_states(n_cells=10)
+        estimator_fwd.plot_macrostates(which="terminal")
 
         estimator_fwd.compute_absorption_probabilities()
         estimator_fwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
@@ -101,11 +101,11 @@ def test_fwd_pipeline_gpcca(self, adata: AnnData):
         _assert_has_all_keys(adata)
 
         # select a subset of states
-        estimator_fwd.set_terminal_states_from_macrostates(
+        estimator_fwd.set_terminal_states(
             n_cells=16,
-            names=estimator_fwd.macrostates.cat.categories[:2],
+            states=estimator_fwd.macrostates.cat.categories[:2],
         )
-        estimator_fwd.plot_terminal_states()
+        estimator_fwd.plot_macrostates(which="terminal")
 
         estimator_fwd.compute_absorption_probabilities()
         estimator_fwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
@@ -128,13 +128,13 @@ def test_bwd_pipeline_gpcca(self, adata: AnnData):
         estimator_bwd.compute_schur(5, method="brandts")
 
         estimator_bwd.compute_macrostates(3, n_cells=16)
-        estimator_bwd.plot_macrostates()
+        estimator_bwd.plot_macrostates(which="all")
         estimator_bwd.plot_coarse_T(show_initial_dist=True, show_stationary_dist=True)
         estimator_bwd.plot_schur_matrix()
 
         # select all cells
-        estimator_bwd.set_terminal_states_from_macrostates(n_cells=16)
-        estimator_bwd.plot_terminal_states()
+        estimator_bwd.set_terminal_states(n_cells=16)
+        estimator_bwd.plot_macrostates(which="terminal")
 
         estimator_bwd.compute_absorption_probabilities()
         estimator_bwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
@@ -142,11 +142,11 @@ def test_bwd_pipeline_gpcca(self, adata: AnnData):
         _assert_has_all_keys(adata, bwd=True)
 
         # select a subset of states
-        estimator_bwd.set_terminal_states_from_macrostates(
+        estimator_bwd.set_terminal_states(
             n_cells=16,
-            names=estimator_bwd.macrostates.cat.categories[:2],
+            states=estimator_bwd.macrostates.cat.categories[:2],
         )
-        estimator_bwd.plot_terminal_states()
+        estimator_bwd.plot_macrostates(which="terminal")
 
         estimator_bwd.compute_absorption_probabilities()
         estimator_bwd.compute_lineage_drivers(cluster_key="clusters", use_raw=False)
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
index 99588de05..0fbf42f19 100644
--- a/tests/test_plotting.py
+++ b/tests/test_plotting.py
@@ -103,7 +103,9 @@ def _assert_equal(fpath: str) -> None:
 
     def compare_cflare_fwd(
         func: Callable,
-    ) -> Callable:  # mustn't use functools.wraps - it think's the fact that `adata` is fixture
+    ) -> (
+        Callable
+    ):  # mustn't use functools.wraps - it think's the fact that `adata` is fixture
         def decorator(self, adata_cflare_fwd) -> None:
             adata, mc = adata_cflare_fwd
             fpath, path = _prepare_fname(func)
@@ -2073,11 +2075,11 @@ def test_mc_spectrum_evals_complex(self, mc: CFLARE, fpath: str):
 
     @compare(kind="cflare")
     def test_scvelo_final_states(self, mc: CFLARE, fpath: str):
-        mc.plot_terminal_states(dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_final_states_clusters(self, mc: CFLARE, fpath: str):
-        mc.plot_terminal_states(color="clusters", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="cflare")
     def test_scvelo_lin_probs(self, mc: CFLARE, fpath: str):
@@ -2183,96 +2185,86 @@ def test_gpcca_coarse_T_no_order(self, mc: GPCCA, fpath: str):
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_lineages(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(states=["0"], dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", states=["0"], dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_discrete(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(discrete=True, dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", discrete=True, dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_cluster_key(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(color="clusters", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_no_same_plot(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(same_plot=False, dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", same_plot=False, dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_cmap(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(cmap=cm.inferno, same_plot=False, dpi=DPI, save=fpath)
+        mc.plot_macrostates(
+            which="all", cmap=cm.inferno, same_plot=False, dpi=DPI, save=fpath
+        )
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_title(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(title="foobar", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", title="foobar", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_meta_states_time(self, mc: GPCCA, fpath: str):
-        mc.plot_macrostates(mode="time", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="all", mode="time", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_lineages(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(states=["0"], dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", states=["0"], dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_discrete(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(discrete=True, dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", discrete=True, dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_cluster_key(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(color="clusters", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", color="clusters", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_no_same_plot(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(same_plot=False, dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", same_plot=False, dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_cmap(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(cmap=cm.inferno, same_plot=False, dpi=DPI, save=fpath)
+        mc.plot_macrostates(
+            which="terminal", cmap=cm.inferno, same_plot=False, dpi=DPI, save=fpath
+        )
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_title(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(title="foobar", dpi=DPI, save=fpath)
+        mc.plot_macrostates(which="terminal", title="foobar", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_final_states_time(self, mc: GPCCA, fpath: str):
-        mc.plot_terminal_states(mode="time", dpi=DPI, save=fpath)
-
-    @compare(kind="gpcca")
-    def test_scvelo_gpcca_abs_probs_disc_same(self, mc: GPCCA, fpath: str):
-        mc.plot_absorption_probabilities(
-            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(
-            color="clusters", discrete=True, same_plot=False, dpi=DPI, save=fpath
-        )
+        mc.plot_macrostates(which="terminal", mode="time", dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_cont_same_no_clusters(self, mc: GPCCA, fpath: str):
-        mc.plot_absorption_probabilities(
-            discrete=False, same_plot=True, dpi=DPI, save=fpath
-        )
+        mc.plot_absorption_probabilities(same_plot=True, dpi=DPI, save=fpath)
 
     @compare(kind="gpcca")
     def test_scvelo_gpcca_abs_probs_cont_same_clusters(self, mc: GPCCA, fpath: str):
         mc.plot_absorption_probabilities(
-            color="clusters", discrete=False, same_plot=True, dpi=DPI, save=fpath
+            color="clusters", 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(
-            color="clusters", discrete=False, same_plot=False, dpi=DPI, save=fpath
+            color="clusters", same_plot=False, dpi=DPI, save=fpath
         )
 
     @compare(kind="gpcca")
diff --git a/tests/test_utils.py b/tests/test_utils.py
index 671cc2ac1..0eea4b7bd 100644
--- a/tests/test_utils.py
+++ b/tests/test_utils.py
@@ -163,14 +163,12 @@ def test_matrix_irreducibility(
     def test_matrix_connectivity(
         self, test_matrix_1: np.ndarray, test_matrix_3: np.ndarray
     ):
-
         assert _connected(test_matrix_1)
         assert not _connected(test_matrix_3)
 
     def test_matrix_symmetry(
         self, test_matrix_1: np.ndarray, test_matrix_4: np.ndarray
     ):
-
         assert not _symmetric(test_matrix_1)
         assert _symmetric(test_matrix_4)
 
@@ -180,7 +178,6 @@ def test_matrix_partition(
         test_matrix_2: np.ndarray,
         test_matrix_3: np.ndarray,
     ):
-
         np.testing.assert_array_equal(
             _partition(test_matrix_1)[0][0], [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11]
         )