def set_rcParams_scvelo(fontsize=12, color_map=None, frameon=None):
def set_rcParams_scanpy(fontsize=12, color_map=None, frameon=None):
def set_figure_params(
    style="scvelo",
    dpi=100,
    dpi_save=150,
    frameon=None,
    vector_friendly=True,
    transparent=True,
    fontsize=12,
    figsize=None,
    color_map=None,
    facecolor=None,
    format="pdf",
    ipython_format="png2x",
):
def set_rcParams_defaults():
def _set_ipython(ipython_format="png2x"):
def _set_start_time(): def get_neighs(adata, mode="distances"):
    if hasattr(adata, "obsp") and mode in adata.obsp.keys():
        return adata.obsp[mode]
    elif "neighbors" in adata.uns.keys() and mode in adata.uns["neighbors"]:
        return adata.uns["neighbors"][mode]
    else:
        raise ValueError("The selected mode is not valid.") def transition_matrix(
    adata,
    vkey="velocity",
    basis=None,
    backward=False,
    self_transitions=True,
    scale=10,
    perc=None,
    threshold=None,
    use_negative_cosines=False,
    weight_diffusion=0,
    scale_diffusion=1,
    weight_indirect_neighbors=None,
    n_neighbors=None,
    vgraph=None,
    basis_constraint=None,
):
    """Computes cell-to-cell transition probabilities

    .. math::
        \\tilde \\pi_{ij} = \\frac1{z_i} \\exp( \\pi_{ij} / \\sigma),

    from the velocity graph :math:`\\pi_{ij}`, with row-normalization :math:`z_i` and
    kernel width :math:`\\sigma` (scale parameter :math:`\\lambda = \\sigma^{-1}`).

    Alternatively, use :func:`cellrank.tl.transition_matrix` to account for uncertainty
    in the velocity estimates.

    Arguments
    ---------
    adata: :class:`~anndata.AnnData`
        Annotated data matrix.
    vkey: `str` (default: `'velocity'`)
        Name of velocity estimates to be used.
    basis: `str` or `None` (default: `None`)
        Restrict transition to embedding if specified
    backward: `bool` (default: `False`)
        Whether to use the transition matrix to
        push forward (`False`) or to pull backward (`True`)
    self_transitions: `bool` (default: `True`)
        Allow transitions from one node to itself.
    scale: `float` (default: 10)
        Scale parameter of gaussian kernel.
    perc: `float` between `0` and `100` or `None` (default: `None`)
        Determines threshold of transitions to include.
    use_negative_cosines: `bool` (default: `False`)
        If True, negatively similar transitions are taken into account.
    weight_diffusion: `float` (default: 0)
        Relative weight to be given to diffusion kernel (Brownian motion)
    scale_diffusion: `float` (default: 1)
        Scale of diffusion kernel.
    weight_indirect_neighbors: `float` between `0` and `1` or `None` (default: `None`)
        Weight to be assigned to indirect neighbors (i.e. neighbors of higher degrees).
    n_neighbors:`int` (default: None)
        Number of nearest neighbors to consider around each cell.
    vgraph: csr matrix or `None` (default: `None`)
        Velocity graph representation to use instead of adata.uns[f'{vkey}_graph'].

    Returns
    -------
    Returns sparse matrix with transition probabilities.
    """

    if f"{vkey}_graph" not in adata.uns:
        raise ValueError(
            "You need to run `tl.velocity_graph` first to compute cosine correlations."
        )

    graph_neg = None
    if vgraph is not None:
        graph = vgraph.copy()
    else:
        if hasattr(adata, "obsp") and f"{vkey}_graph" in adata.obsp.keys():
            graph = csr_matrix(adata.obsp[f"{vkey}_graph"]).copy()
            if f"{vkey}_graph_neg" in adata.obsp.keys():
                graph_neg = adata.obsp[f"{vkey}_graph_neg"]
        else:
            graph = csr_matrix(adata.uns[f"{vkey}_graph"]).copy()
            if f"{vkey}_graph_neg" in adata.uns.keys():
                graph_neg = adata.uns[f"{vkey}_graph_neg"]

    if basis_constraint is not None and f"X_{basis_constraint}" in adata.obsm.keys():
        from sklearn.neighbors import NearestNeighbors

        neighs = NearestNeighbors(n_neighbors=100)
        neighs.fit(adata.obsm[f"X_{basis_constraint}"])
        basis_graph = neighs.kneighbors_graph(mode="connectivity") > 0
        graph = graph.multiply(basis_graph)

    if self_transitions:
        confidence = graph.max(1).A.flatten()
        ub = np.percentile(confidence, 98)
        self_prob = np.clip(ub - confidence, 0, 1)
        graph.setdiag(self_prob)

    T = np.expm1(graph * scale)  # equivalent to np.exp(graph.A * scale) - 1
    if graph_neg is not None:
        graph_neg = adata.uns[f"{vkey}_graph_neg"]
        if use_negative_cosines:
            T -= np.expm1(-graph_neg * scale)
        else:
            T += np.expm1(graph_neg * scale)
            T.data += 1

    # weight direct and indirect (recursed) neighbors
    if weight_indirect_neighbors is not None and weight_indirect_neighbors < 1:
        direct_neighbors = get_neighs(adata, "distances") > 0
        direct_neighbors.setdiag(1)
        w = weight_indirect_neighbors
        T = w * T + (1 - w) * direct_neighbors.multiply(T)

    if n_neighbors is not None:
        T = T.multiply(
            get_connectivities(
                adata, mode="distances", n_neighbors=n_neighbors, recurse_neighbors=True
            )
        )

    if perc is not None or threshold is not None:
        if threshold is None:
            threshold = np.percentile(T.data, perc)
        T.data[T.data < threshold] = 0
        T.eliminate_zeros()

    if backward:
        T = T.T
    T = normalize(T)

    if f"X_{basis}" in adata.obsm.keys():
        dists_emb = (T > 0).multiply(squareform(pdist(adata.obsm[f"X_{basis}"])))
        scale_diffusion *= dists_emb.data.mean()

        diffusion_kernel = dists_emb.copy()
        diffusion_kernel.data = np.exp(
            -0.5 * dists_emb.data ** 2 / scale_diffusion ** 2
        )
        T = T.multiply(diffusion_kernel)  # combine velocity kernel & diffusion kernel

        if 0 < weight_diffusion < 1:  # add diffusion kernel (Brownian motion - like)
            diffusion_kernel.data = np.exp(
                -0.5 * dists_emb.data ** 2 / (scale_diffusion / 2) ** 2
            )
            T = (1 - weight_diffusion) * T + weight_diffusion * diffusion_kernel

        T = normalize(T)

    return T def doc_params(**kwds):
    def dec(obj): @doc_params(scatter=doc_scatter)
def scatter(
    adata=None,
    basis=None,
    x=None,
    y=None,
    vkey=None,
    color=None,
    use_raw=None,
    layer=None,
    color_map=None,
    colorbar=None,
    palette=None,
    size=None,
    alpha=None,
    linewidth=None,
    linecolor=None,
    perc=None,
    groups=None,
    sort_order=True,
    components=None,
    projection=None,
    legend_loc=None,
    legend_loc_lines=None,
    legend_fontsize=None,
    legend_fontweight=None,
    legend_fontoutline=None,
    legend_align_text=None,
    xlabel=None,
    ylabel=None,
    title=None,
    fontsize=None,
    figsize=None,
    xlim=None,
    ylim=None,
    add_density=None,
    add_assignments=None,
    add_linfit=None,
    add_polyfit=None,
    add_rug=None,
    add_text=None,
    add_text_pos=None,
    add_margin=None,
    add_outline=None,
    outline_width=None,
    outline_color=None,
    n_convolve=None,
    smooth=None,
    normalize_data=None,
    rescale_color=None,
    color_gradients=None,
    dpi=None,
    frameon=None,
    zorder=None,
    ncols=None,
    nrows=None,
    wspace=None,
    hspace=None,
    show=None,
    save=None,
    ax=None,
    **kwargs,
):
    """\
    Scatter plot along observations or variables axes.

    Arguments
    ---------
    adata: :class:`~anndata.AnnData`
        Annotated data matrix.
    x: `str`, `np.ndarray` or `None` (default: `None`)
        x coordinate
    y: `str`, `np.ndarray` or `None` (default: `None`)
        y coordinate
    {scatter}

    Returns
    -------
    If `show==False` a `matplotlib.Axis`
    """

    if adata is None and (x is not None and y is not None):
        adata = AnnData(np.stack([x, y]).T)

    # restore old conventions
    add_assignments = kwargs.pop("show_assignments", add_assignments)
    add_linfit = kwargs.pop("show_linear_fit", add_linfit)
    add_polyfit = kwargs.pop("show_polyfit", add_polyfit)
    add_density = kwargs.pop("show_density", add_density)
    add_rug = kwargs.pop("rug", add_rug)
    basis = kwargs.pop("var_names", basis)

    # keys for figures (fkeys) and multiple plots (mkeys)
    fkeys = ["adata", "show", "save", "groups", "ncols", "nrows", "wspace", "hspace"]
    fkeys += ["add_margin", "ax", "kwargs"]
    mkeys = ["color", "layer", "basis", "components", "x", "y", "xlabel", "ylabel"]
    mkeys += ["title", "color_map", "add_text"]
    scatter_kwargs = {"show": False, "save": False}
    for key in signature(scatter).parameters:
        if key not in mkeys + fkeys:
            scatter_kwargs[key] = eval(key)
    mkwargs = {}
    for key in mkeys:  # mkwargs[key] = key for key in mkeys
        mkwargs[key] = eval("{0}[0] if is_list({0}) else {0}".format(key))

    # use c & color and cmap & color_map interchangeably,
    # and plot each group separately if groups is 'all'
    if "c" in kwargs:
        color = kwargs.pop("c")
    if "cmap" in kwargs:
        color_map = kwargs.pop("cmap")
    if "rasterized" not in kwargs:
        kwargs["rasterized"] = settings._vector_friendly
    if isinstance(color_map, (list, tuple)) and all(
        [is_color_like(c) or c == "transparent" for c in color_map]
    ):
        color_map = rgb_custom_colormap(colors=color_map)
    if isinstance(groups, str) and groups == "all":
        if color is None:
            color = default_color(adata)
        if is_categorical(adata, color):
            vc = adata.obs[color].value_counts()
            groups = [[c] for c in vc[vc > 0].index]
    if isinstance(add_text, (list, tuple, np.ndarray, np.record)):
        add_text = list(np.array(add_text, dtype=str))

    # create list of each mkey and check if all bases are valid.
    color = to_list(color, max_len=None)
    layer, components = to_list(layer), to_list(components)
    x, y, basis = to_list(x), to_list(y), to_valid_bases_list(adata, basis)

    # get multikey (with more than one element)
    multikeys = eval(f"[{','.join(mkeys)}]")
    if is_list_of_list(groups):
        multikeys.append(groups)
    key_lengths = np.array([len(key) if is_list(key) else 1 for key in multikeys])
    multikey = (
        multikeys[np.where(key_lengths > 1)[0][0]] if np.max(key_lengths) > 1 else None
    )

    # gridspec frame for plotting multiple keys (mkeys: list or tuple)
    if multikey is not None:
        if np.sum(key_lengths > 1) == 1 and is_list_of_str(multikey):
            multikey = unique(multikey)  # take unique set if no more than one multikey
        if len(multikey) > 20:
            raise ValueError("Please restrict the passed list to max 20 elements.")
        if ax is not None:
            logg.warn("Cannot specify `ax` when plotting multiple panels.")
        if is_list(title):
            title *= int(np.ceil(len(multikey) / len(title)))
        if nrows is None:
            ncols = len(multikey) if ncols is None else min(len(multikey), ncols)
            nrows = int(np.ceil(len(multikey) / ncols))
        else:
            ncols = int(np.ceil(len(multikey) / nrows))
        if not frameon or frameon == "artist":
            lloc, llines = "legend_loc", "legend_loc_lines"
            if lloc in scatter_kwargs and scatter_kwargs[lloc] is None:
                scatter_kwargs[lloc] = "none"
            if llines in scatter_kwargs and scatter_kwargs[llines] is None:
                scatter_kwargs[llines] = "none"

        grid_figsize, dpi = get_figure_params(figsize, dpi, ncols)
        grid_figsize = (grid_figsize[0] * ncols, grid_figsize[1] * nrows)
        fig = pl.figure(None, grid_figsize, dpi=dpi)
        hspace = 0.3 if hspace is None else hspace
        gspec = pl.GridSpec(nrows, ncols, fig, hspace=hspace, wspace=wspace)

        ax = []
        for i, gs in enumerate(gspec):
            if i < len(multikey):
                g = groups[i * (len(groups) > i)] if is_list_of_list(groups) else groups
                multi_kwargs = {"groups": g}
                for key in mkeys:  # multi_kwargs[key] = key[i] if is multikey else key
                    multi_kwargs[key] = eval(
                        "{0}[i * (len({0}) > i)] if is_list({0}) else {0}".format(key)
                    )
                ax.append(
                    scatter(
                        adata,
                        ax=pl.subplot(gs),
                        **multi_kwargs,
                        **scatter_kwargs,
                        **kwargs,
                    )
                )

        if not frameon and isinstance(ylabel, str):
            set_label(xlabel, ylabel, fontsize, ax=ax[0], fontweight="bold")
        savefig_or_show(dpi=dpi, save=save, show=show)
        if show is False:
            return ax

    else:
        # make sure that there are no more lists, e.g. ['clusters'] becomes 'clusters'
        color_map = to_val(color_map)
        color, layer, basis = to_val(color), to_val(layer), to_val(basis)
        x, y, components = to_val(x), to_val(y), to_val(components)
        xlabel, ylabel, title = to_val(xlabel), to_val(ylabel), to_val(title)

        # multiple plots within one ax for comma-separated y or layers (string).

        if any([isinstance(key, str) and "," in key for key in [y, layer]]):
            # comma split
            y, layer, color = [
                [k.strip() for k in key.split(",")]
                if isinstance(key, str) and "," in key
                else to_list(key)
                for key in [y, layer, color]
            ]
            multikey = y if len(y) > 1 else layer if len(layer) > 1 else None

            if multikey is not None:
                for i, mi in enumerate(multikey):
                    ax = scatter(
                        adata,
                        x=x,
                        y=y[i * (len(y) > i)],
                        color=color[i * (len(color) > i)],
                        layer=layer[i * (len(layer) > i)],
                        basis=basis,
                        components=components,
                        groups=groups,
                        xlabel=xlabel,
                        ylabel="expression" if ylabel is None else ylabel,
                        color_map=color_map,
                        title=y[i * (len(y) > i)] if title is None else title,
                        ax=ax,
                        **scatter_kwargs,
                    )
                if legend_loc is None:
                    legend_loc = "best"
                if legend_loc and legend_loc != "none":
                    multikey = [key.replace("Ms", "spliced") for key in multikey]
                    multikey = [key.replace("Mu", "unspliced") for key in multikey]
                    ax.legend(multikey, fontsize=legend_fontsize, loc=legend_loc)

                savefig_or_show(dpi=dpi, save=save, show=show)
                if show is False:
                    return ax

        elif color_gradients is not None and color_gradients is not False:
            vals, names, color, scatter_kwargs = gets_vals_from_color_gradients(
                adata, color, **scatter_kwargs
            )
            cols = zip(adata.obs[color].cat.categories, adata.uns[f"{color}_colors"])
            c_colors = {cat: col for (cat, col) in cols}
            mkwargs.pop("color")
            ax = scatter(
                adata,
                color="grey",
                ax=ax,
                **mkwargs,
                **get_kwargs(scatter_kwargs, {"alpha": 0.05}),
            )  # background
            ax = scatter(
                adata,
                color=color,
                ax=ax,
                **mkwargs,
                **get_kwargs(scatter_kwargs, {"s": 0}),
            )  # set legend
            sorted_idx = np.argsort(vals, 1)[:, ::-1][:, :2]
            for id0 in range(len(names)):
                for id1 in range(id0 + 1, len(names)):
                    cmap = rgb_custom_colormap(
                        [c_colors[names[id0]], "white", c_colors[names[id1]]],
                        alpha=[1, 0, 1],
                    )
                    mkwargs.update({"color_map": cmap})
                    c_vals = np.array(vals[:, id1] - vals[:, id0]).flatten()
                    c_bool = np.array([id0 in c and id1 in c for c in sorted_idx])
                    if np.sum(c_bool) > 1:
                        _adata = adata[c_bool] if np.sum(~c_bool) > 0 else adata
                        mkwargs["color"] = c_vals[c_bool]
                        ax = scatter(
                            _adata, ax=ax, **mkwargs, **scatter_kwargs, **kwargs
                        )
            savefig_or_show(dpi=dpi, save=save, show=show)
            if show is False:
                return ax

        # actual scatter plot
        else:
            # set color, color_map, edgecolor, basis, linewidth, frameon, use_raw
            if color is None:
                color = default_color(adata, add_outline)
            if "cmap" not in kwargs:
                kwargs["cmap"] = (
                    default_color_map(adata, color) if color_map is None else color_map
                )
            if "s" not in kwargs:
                kwargs["s"] = default_size(adata) if size is None else size
            if "edgecolor" not in kwargs:
                kwargs["edgecolor"] = "none"
            is_embedding = ((x is None) | (y is None)) and basis not in adata.var_names
            if basis is None and is_embedding:
                basis = default_basis(adata)
            if linewidth is None:
                linewidth = 1
            if frameon is None:
                frameon = True if not is_embedding else settings._frameon
            if isinstance(groups, str):
                groups = [groups]
            if use_raw is None and basis not in adata.var_names:
                use_raw = layer is None and adata.raw is not None

            ax, show = get_ax(ax, show, figsize, dpi, projection)

            # phase portrait: get x and y from .layers (e.g. spliced vs. unspliced)
            if basis in adata.var_names:
                if title is None:
                    title = basis
                if x is None and y is None:
                    x = default_xkey(adata, use_raw=use_raw)
                    y = default_ykey(adata, use_raw=use_raw)
                elif x is None or y is None:
                    raise ValueError("Both x and y have to specified.")
                if isinstance(x, str) and isinstance(y, str):
                    layers_keys = list(adata.layers.keys()) + ["X"]
                    if any([key not in layers_keys for key in [x, y]]):
                        raise ValueError("Could not find x or y in layers.")

                    if xlabel is None:
                        xlabel = x
                    if ylabel is None:
                        ylabel = y

                    x = get_obs_vector(adata, basis, layer=x, use_raw=use_raw)
                    y = get_obs_vector(adata, basis, layer=y, use_raw=use_raw)

                if legend_loc is None:
                    legend_loc = "none"

                if use_raw and perc is not None:
                    ub = np.percentile(x, 99.9 if not isinstance(perc, int) else perc)
                    ax.set_xlim(right=ub * 1.05)
                    ub = np.percentile(y, 99.9 if not isinstance(perc, int) else perc)
                    ax.set_ylim(top=ub * 1.05)

                # velocity model fits (full dynamics and steady-state ratios)
                if any(["gamma" in key or "alpha" in key for key in adata.var.keys()]):
                    plot_velocity_fits(
                        adata,
                        basis,
                        vkey,
                        use_raw,
                        linewidth,
                        linecolor,
                        legend_loc_lines,
                        legend_fontsize,
                        add_assignments,
                        ax=ax,
                    )

            # embedding: set x and y to embedding coordinates
            elif is_embedding:
                X_emb = adata.obsm[f"X_{basis}"][:, get_components(components, basis)]
                x, y = X_emb[:, 0], X_emb[:, 1]
                # todo: 3d plotting
                # z = X_emb[:, 2] if projection == "3d" and X_emb.shape[1] > 2 else None

            elif isinstance(x, str) and isinstance(y, str):
                var_names = (
                    adata.raw.var_names
                    if use_raw and adata.raw is not None
                    else adata.var_names
                )
                if layer is None:
                    layer = default_xkey(adata, use_raw=use_raw)
                x_keys = list(adata.obs.keys()) + list(adata.layers.keys())
                is_timeseries = y in var_names and x in x_keys
                if xlabel is None:
                    xlabel = x
                if ylabel is None:
                    ylabel = layer if is_timeseries else y
                if title is None:
                    title = y if is_timeseries else color
                if legend_loc is None:
                    legend_loc = "none"

                # gene trend: x and y as gene along obs/layers (e.g. pseudotime)
                if is_timeseries:
                    x = (
                        adata.obs[x]
                        if x in adata.obs.keys()
                        else adata.obs_vector(y, layer=x)
                    )
                    y = get_obs_vector(adata, basis=y, layer=layer, use_raw=use_raw)
                # get x and y from var_names, var or obs
                else:
                    if x in var_names and y in var_names:
                        if layer in adata.layers.keys():
                            x = adata.obs_vector(x, layer=layer)
                            y = adata.obs_vector(y, layer=layer)
                        else:
                            data = adata.raw if use_raw else adata
                            x, y = data.obs_vector(x), data.obs_vector(y)
                    elif x in adata.var.keys() and y in adata.var.keys():
                        x, y = adata.var[x], adata.var[y]
                    elif x in adata.obs.keys() and y in adata.obs.keys():
                        x, y = adata.obs[x], adata.obs[y]
                    elif np.any(
                        [var_key in x or var_key in y for var_key in adata.var.keys()]
                    ):
                        var_keys = [
                            k
                            for k in adata.var.keys()
                            if not isinstance(adata.var[k][0], str)
                        ]
                        var = adata.var[var_keys]
                        x = var.astype(np.float32).eval(x)
                        y = var.astype(np.float32).eval(y)
                    elif np.any(
                        [obs_key in x or obs_key in y for obs_key in adata.obs.keys()]
                    ):
                        obs_keys = [
                            k
                            for k in adata.obs.keys()
                            if not isinstance(adata.obs[k][0], str)
                        ]
                        obs = adata.obs[obs_keys]
                        x = obs.astype(np.float32).eval(x)
                        y = obs.astype(np.float32).eval(y)
                    else:
                        raise ValueError(
                            "x or y is invalid! pass valid observation or a gene name"
                        )

            x, y = make_dense(x).flatten(), make_dense(y).flatten()

            # convolve along x axes (e.g. pseudotime)
            if n_convolve is not None:
                vec_conv = np.ones(n_convolve) / n_convolve
                y[np.argsort(x)] = np.convolve(y[np.argsort(x)], vec_conv, mode="same")

            # if color is set to a cell index, plot that cell on top
            if is_int(color) or is_list_of_int(color) and len(color) != len(x):
                color = np.array(np.isin(np.arange(len(x)), color), dtype=bool)
                size = kwargs["s"] * 2 if np.sum(color) == 1 else kwargs["s"]
                if zorder is None:
                    zorder = 10
                ax.scatter(
                    np.ravel(x[color]),
                    np.ravel(y[color]),
                    s=size,
                    zorder=zorder,
                    color=palette[-1] if palette is not None else "darkblue",
                )
                color = (
                    palette[0] if palette is not None and len(palette) > 1 else "gold"
                )
                zorder -= 1

            # if color is in {'ascending', 'descending'}
            elif isinstance(color, str):
                if color == "ascending":
                    color = np.linspace(0, 1, len(x))
                elif color == "descending":
                    color = np.linspace(1, 0, len(x))

            # set palette if categorical color vals
            if is_categorical(adata, color):
                set_colors_for_categorical_obs(adata, color, palette)

            # set color
            if (
                basis in adata.var_names
                and isinstance(color, str)
                and color in adata.layers.keys()
            ):
                # phase portrait: color=basis, layer=color
                c = interpret_colorkey(adata, basis, color, perc, use_raw)
            else:
                # embedding, gene trend etc.
                c = interpret_colorkey(adata, color, layer, perc, use_raw)

            if c is not None and not isinstance(c, str) and not isinstance(c[0], str):
                # smooth color values across neighbors and rescale
                if smooth and len(c) == adata.n_obs:
                    n_neighbors = None if isinstance(smooth, bool) else smooth
                    c = get_connectivities(adata, n_neighbors=n_neighbors).dot(c)
                # rescale color values to min and max acc. to rescale_color tuple
                if rescale_color is not None:
                    try:
                        c += rescale_color[0] - np.nanmin(c)
                        c *= rescale_color[1] / np.nanmax(c)
                    except Exception:
                        logg.warn("Could not rescale colors. Pass a tuple, e.g. [0,1].")

            # set vmid to 0 if color values obtained from velocity expression
            if not np.any([v in kwargs for v in ["vmin", "vmid", "vmax"]]) and np.any(
                [
                    isinstance(v, str)
                    and "time" not in v
                    and (v.endswith("velocity") or v.endswith("transition"))
                    for v in [color, layer]
                ]
            ):
                kwargs["vmid"] = 0

            # introduce vmid by setting vmin and vmax accordingly
            if "vmid" in kwargs:
                vmid = kwargs.pop("vmid")
                if vmid is not None:
                    if not (isinstance(c, str) or isinstance(c[0], str)):
                        lb, ub = np.min(c), np.max(c)
                        crange = max(np.abs(vmid - lb), np.abs(ub - vmid))
                        kwargs.update({"vmin": vmid - crange, "vmax": vmid + crange})

            x, y = np.ravel(x), np.ravel(y)
            if len(x) != len(y):
                raise ValueError("x or y do not share the same dimension.")

            if normalize_data:
                x = (x - np.nanmin(x)) / (np.nanmax(x) - np.nanmin(x))
                y = (y - np.nanmin(x)) / (np.nanmax(y) - np.nanmin(y))

            if not isinstance(c, str):
                c = np.ravel(c) if len(np.ravel(c)) == len(x) else c

            # store original order of color values
            color_array, scatter_array = c, np.stack([x, y]).T

            # set color to grey for NAN values and for cells that are not in groups
            if (
                groups is not None
                or is_categorical(adata, color)
                and np.any(pd.isnull(adata.obs[color]))
            ):
                if isinstance(groups, (list, tuple, np.record)):
                    groups = unique(groups)
                zorder = 0 if zorder is None else zorder
                pop_keys = ["groups", "add_linfit", "add_polyfit", "add_density"]
                _ = [scatter_kwargs.pop(key, None) for key in pop_keys]
                ax = scatter(
                    adata,
                    x=x,
                    y=y,
                    basis=basis,
                    layer=layer,
                    color="lightgrey",
                    ax=ax,
                    **scatter_kwargs,
                )
                if groups is not None and len(groups) == 1:
                    if (
                        isinstance(groups[0], str)
                        and groups[0] in adata.var.keys()
                        and basis in adata.var_names
                    ):
                        groups = f"{adata[:, basis].var[groups[0]][0]}"
                idx = groups_to_bool(adata, groups, color)
                if idx is not None:
                    if np.sum(idx) > 0:  # if any group to be highlighted
                        x, y = x[idx], y[idx]
                        if not isinstance(c, str) and len(c) == adata.n_obs:
                            c = c[idx]
                        if isinstance(kwargs["s"], np.ndarray):
                            kwargs["s"] = np.array(kwargs["s"])[idx]
                        if (
                            title is None
                            and groups is not None
                            and len(groups) == 1
                            and isinstance(groups[0], str)
                        ):
                            title = groups[0]
                    else:  # if nothing to be highlighted
                        add_linfit, add_polyfit, add_density = None, None, None
            else:
                idx = None

            if not isinstance(c, str) and len(c) != len(x):
                c = "grey"
                if not isinstance(color, str) or color != default_color(adata):
                    logg.warn("Invalid color key. Using grey instead.")

            # check if higher value points should be plotted on top
            if not isinstance(c, str) and len(c) == len(x):
                order = None
                if sort_order and not is_categorical(adata, color):
                    order = np.argsort(c)
                elif not sort_order and is_categorical(adata, color):
                    counts = get_value_counts(
                        adata[idx] if idx is not None else adata, color
                    )
                    np.random.seed(0)
                    nums, p = np.arange(0, len(x)), counts / np.sum(counts)
                    order = np.random.choice(nums, len(x), replace=False, p=p)
                if order is not None:
                    x, y, c = x[order], y[order], c[order]
                    if isinstance(kwargs["s"], np.ndarray):  # sort sizes if array-type
                        kwargs["s"] = np.array(kwargs["s"])[order]

            marker = kwargs.pop("marker", ".")
            smp = ax.scatter(
                x, y, c=c, alpha=alpha, marker=marker, zorder=zorder, **kwargs
            )

            outline_dtypes = (list, tuple, np.ndarray, int, np.int_, str)
            if isinstance(add_outline, outline_dtypes) or add_outline:
                if isinstance(add_outline, (list, tuple, np.record)):
                    add_outline = unique(add_outline)
                if (
                    add_outline is not True
                    and isinstance(add_outline, (int, np.int_))
                    or is_list_of_int(add_outline)
                    and len(add_outline) != len(x)
                ):
                    add_outline = np.isin(np.arange(len(x)), add_outline)
                    add_outline = np.array(add_outline, dtype=bool)
                    if outline_width is None:
                        outline_width = (0.6, 0.3)
                if isinstance(add_outline, str):
                    if add_outline in adata.var.keys() and basis in adata.var_names:
                        add_outline = f"{adata[:, basis].var[add_outline][0]}"
                idx = groups_to_bool(adata, add_outline, color)
                if idx is not None and np.sum(idx) > 0:  # if anything to be outlined
                    zorder = 2 if zorder is None else zorder + 2
                    if kwargs["s"] is not None:
                        kwargs["s"] *= 1.2
                    # restore order of values
                    x, y = scatter_array[:, 0][idx], scatter_array[:, 1][idx]
                    c = color_array
                    if not isinstance(c, str) and len(c) == adata.n_obs:
                        c = c[idx]
                    if isinstance(kwargs["s"], np.ndarray):
                        kwargs["s"] = np.array(kwargs["s"])[idx]
                    if isinstance(c, np.ndarray) and not isinstance(c[0], str):
                        if "vmid" not in kwargs and "vmin" not in kwargs:
                            kwargs["vmin"] = np.min(color_array)
                        if "vmid" not in kwargs and "vmax" not in kwargs:
                            kwargs["vmax"] = np.max(color_array)
                    ax.scatter(
                        x, y, c=c, alpha=alpha, marker=".", zorder=zorder, **kwargs
                    )
                if idx is None or np.sum(idx) > 0:  # if all or anything to be outlined
                    plot_outline(
                        x, y, kwargs, outline_width, outline_color, zorder, ax=ax
                    )
                if idx is not None and np.sum(idx) == 0:  # if nothing to be outlined
                    add_linfit, add_polyfit, add_density = None, None, None

            # set legend if categorical categorical color vals
            if is_categorical(adata, color) and len(scatter_array) == adata.n_obs:
                legend_loc = default_legend_loc(adata, color, legend_loc)
                g_bool = groups_to_bool(adata, add_outline, color)
                if not (add_outline is None or g_bool is None):
                    groups = add_outline
                set_legend(
                    adata,
                    ax,
                    color,
                    legend_loc,
                    scatter_array,
                    legend_fontweight,
                    legend_fontsize,
                    legend_fontoutline,
                    legend_align_text,
                    groups,
                )
            if add_density:
                plot_density(x, y, add_density, ax=ax)

            if add_linfit:
                if add_linfit is True and basis in adata.var_names:
                    add_linfit = "no_intercept"  # without intercept
                plot_linfit(
                    x,
                    y,
                    add_linfit,
                    legend_loc != "none",
                    linecolor,
                    linewidth,
                    fontsize,
                    ax=ax,
                )

            if add_polyfit:
                if add_polyfit is True and basis in adata.var_names:
                    add_polyfit = "no_intercept"  # without intercept
                plot_polyfit(
                    x,
                    y,
                    add_polyfit,
                    legend_loc != "none",
                    linecolor,
                    linewidth,
                    fontsize,
                    ax=ax,
                )

            if add_rug:
                rug_color = add_rug if isinstance(add_rug, str) else color
                rug_color = np.ravel(interpret_colorkey(adata, rug_color))
                plot_rug(np.ravel(x), color=rug_color, ax=ax)

            if add_text:
                if add_text_pos is None:
                    add_text_pos = [0.05, 0.95]
                ax.text(
                    add_text_pos[0],
                    add_text_pos[1],
                    f"{add_text}",
                    ha="left",
                    va="top",
                    fontsize=fontsize,
                    transform=ax.transAxes,
                    bbox=dict(boxstyle="round", facecolor="wheat", alpha=0.2),
                )

            set_label(xlabel, ylabel, fontsize, basis, ax=ax)
            set_title(title, layer, color, fontsize, ax=ax)
            update_axes(ax, xlim, ylim, fontsize, is_embedding, frameon, figsize)
            if add_margin:
                set_margin(ax, x, y, add_margin)
            if colorbar is not False:
                if not isinstance(c, str) and not is_categorical(adata, color):
                    labelsize = fontsize * 0.75 if fontsize is not None else None
                    set_colorbar(smp, ax=ax, labelsize=labelsize)

            savefig_or_show(dpi=dpi, save=save, show=show)
            if show is False:
                return ax def default_basis(adata, **kwargs):
    if "x" in kwargs and "y" in kwargs:
        keys, x, y = ["embedding"], kwargs.pop("x"), kwargs.pop("y")
        adata.obsm["X_embedding"] = np.stack([x, y]).T
        if "velocity_embedding" in adata.obsm.keys():
            del adata.obsm["velocity_embedding"]
    else:
        keys = [
            key for key in ["pca", "tsne", "umap"] if f"X_{key}" in adata.obsm.keys()
        ]
    if not keys:
        raise ValueError("No basis specified.")
    return keys[-1] if len(keys) > 0 else None def default_size(adata):
    adjusted, classic = 1.2e5 / adata.n_obs, 20
    return (
        np.mean([adjusted, classic])
        if settings._rcParams_style == "scvelo"
        else adjusted
    ) def get_basis(adata, basis):
    if isinstance(basis, str) and basis.startswith("X_"):
        basis = basis[2:]
    check_basis(adata, basis)
    return basis def get_components(components=None, basis=None, projection=None):
    if components is None:
        components = "1,2,3" if projection == "3d" else "1,2"
    if isinstance(components, str):
        components = components.split(",")
    components = np.array(components).astype(int) - 1
    if "diffmap" in basis or "vmap" in basis:
        components += 1
    return components def groups_to_bool(adata, groups, groupby=None):
    groups, groupby = get_groups(adata, groups, groupby)
    if isinstance(groups, (list, tuple, np.ndarray, np.record)):
        if groupby is not None and isinstance(groups[0], str):
            groups = np.array([key in groups for key in adata.obs[groupby]])

    if groupby is not None and groupby in adata.obs.keys():
        c = adata.obs[groupby]
        if np.any(pd.isnull(c)):
            valid = np.array(~pd.isnull(c))
            groups = (
                valid if groups is None or len(groups) != len(c) else groups & valid
            )

    groups = (
        np.ravel(groups)
        if isinstance(groups, (list, tuple, np.ndarray, np.record))
        else None
    )
    return groups def savefig_or_show(writekey=None, show=None, dpi=None, ext=None, save=None):
    if isinstance(save, str):
        # check whether `save` contains a figure extension
        if ext is None:
            for try_ext in [".svg", ".pdf", ".png"]:
                if save.endswith(try_ext):
                    ext = try_ext[1:]
                    save = save.replace(try_ext, "")
                    break
        # append it
        if "/" in save:
            writekey = None
        writekey = (
            f"{writekey}_{save}" if writekey is not None and len(writekey) > 0 else save
        )
        save = True
    save = settings.autosave if save is None else save
    show = settings.autoshow if show is None else show

    if save:
        if dpi is None:
            # needed in nb b/c internal figures are also influenced by 'savefig.dpi'.
            if (
                not isinstance(rcParams["savefig.dpi"], str)
                and rcParams["savefig.dpi"] < 150
            ):
                if settings._low_resolution_warning:
                    logg.warn(
                        "You are using a low resolution (dpi<150) for saving figures.\n"
                        "Consider running `set_figure_params(dpi_save=...)`, which "
                        "will adjust `matplotlib.rcParams['savefig.dpi']`"
                    )
                    settings._low_resolution_warning = False
            else:
                dpi = rcParams["savefig.dpi"]
        if len(settings.figdir) > 0:
            if settings.figdir[-1] != "/":
                settings.figdir += "/"
            if not os.path.exists(settings.figdir):
                os.makedirs(settings.figdir)
        if ext is None:
            ext = settings.file_format_figs
        filepath = f"{settings.figdir}{settings.plot_prefix}{writekey}"
        if "/" in writekey:
            filepath = f"{writekey}"
        try:
            filename = filepath + f"{settings.plot_suffix}.{ext}"
            pl.savefig(filename, dpi=dpi, bbox_inches="tight")
        except Exception:
            # save as .png if .pdf is not feasible (e.g. specific streamplots)
            filename = filepath + f"{settings.plot_suffix}.png"
            pl.savefig(filename, dpi=dpi, bbox_inches="tight")
            logg.msg(f"figure cannot be saved as {ext}, using png instead.", v=1)
        logg.msg("saving figure to file", filename, v=1)
    if show:
        pl.show()
    if save:
        pl.close()  # clear figure
import warnings
import numpy as np
        import matplotlib.pyplot as plt
from scipy.sparse import csr_matrix, issparse
from scvelo import settings
from scvelo.preprocessing.neighbors import get_neighs
from scvelo.tools.transition_matrix import transition_matrix
from .docs import doc_params, doc_scatter
from .scatter import scatter
from .utils import (
    default_basis,
    default_size,
    get_basis,
    get_components,
    groups_to_bool,
    savefig_or_show,
)
    from networkx import DiGraph, Graph
        from numbers import Number
        from matplotlib.collections import LineCollection
        from matplotlib.colors import colorConverter, Colormap, Normalize
        from matplotlib.patches import FancyArrowPatch




@doc_params(scatter=doc_scatter)
def velocity_graph(
    adata,
    basis=None,
    vkey="velocity",
    which_graph=None,
    n_neighbors=10,
    arrows=None,
    arrowsize=3,
    alpha=0.8,
    perc=None,
    threshold=None,
    edge_width=0.2,
    edge_color="grey",
    edges_on_top=None,
    color=None,
    layer=None,
    size=None,
    groups=None,
    components=None,
    title=None,
    dpi=None,
    show=None,
    save=None,
    ax=None,
    **kwargs,
):
    """\
    Plot of the velocity graph.

    Arguments
    ---------
    adata: :class:`~anndata.AnnData`
        Annotated data matrix.
    which_graph: `'velocity_graph'` or `'connectivities'`  (default: `None`)
        Whether to show transitions from velocity graph or neighbor connectivities.
    n_neighbors: `int` (default: 10)
        Number of neighbors to be included for generating connectivity / velocity graph.
    arrows: `bool` (default: `None`)
        Whether to display arrows instead of edges.
        Recommended to be used only on a cluster by setting groups parameter.
    arrowsize: `int` (default: 3)
        Size of the arrow heads.

    {scatter}

    Returns
    -------
    `matplotlib.Axis` if `show==False`
    """

    basis = default_basis(adata, **kwargs) if basis is None else get_basis(adata, basis)
    kwargs.update(
        {
            "basis": basis,
            "title": which_graph if title is None else title,
            "alpha": alpha,
            "components": components,
            "groups": groups,
            "dpi": dpi,
            "show": False,
            "save": None,
        }
    )

    ax = scatter(adata, layer=layer, color=color, size=size, ax=ax, zorder=0, **kwargs)