<commit_before>from .. import settings
from ..preprocessing.neighbors import get_neighs
from ..tools.transition_matrix import transition_matrix
from .utils import savefig_or_show, default_basis, get_components, get_basis, groups_to_bool
from .scatter import scatter
from .docs import doc_scatter, doc_params

import warnings
import numpy as np
from scipy.sparse import issparse, csr_matrix


@doc_params(scatter=doc_scatter)
def velocity_graph(adata, basis=None, vkey='velocity', which_graph=None, n_neighbors=10, arrows=None, arrowsize=3,
                   alpha=.8, perc=None, threshold=None, edge_width=.2, edge_color='grey', edges_on_top=None, color=None,
                   layer=None, size=None, groups=None, components=None, title=None, dpi=None, show=True, 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 connectivities from neighbors graph.
    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)

    from networkx import Graph, DiGraph
    if which_graph in {'neighbors', 'connectivities'}:
        T = get_neighs(adata, 'connectivities').copy()
        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()
    elif which_graph in adata.uns.keys():
        T = adata.uns[which_graph].copy()
        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()
    elif hasattr(adata, 'obsp') and which_graph in adata.obsp.keys():
        T = adata.obsp[which_graph].copy()
        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()
    else:
        if threshold is None: threshold = .05
        T = transition_matrix(adata, vkey=vkey, weight_indirect_neighbors=0,
                              n_neighbors=n_neighbors, perc=perc, threshold=threshold)

    if groups is not None:
        if issparse(T): T = T.A
        T[~groups_to_bool(adata, groups, color)] = 0
        T = csr_matrix(T)
        T.eliminate_zeros()

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        X_emb = adata.obsm['X_' + basis][:, get_components(components, basis)]
        edges = draw_networkx_edges(DiGraph(T) if arrows else Graph(T), X_emb,
                                    width=edge_width, edge_color=edge_color, arrowsize=arrowsize, ax=ax)
        if not arrows and not edges_on_top:
            edges.set_zorder(-2)
            edges.set_rasterized(settings._vector_friendly)

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


def draw_networkx_edges(G, pos, edgelist=None, width=1.0, edge_color='k', style='solid', alpha=None, arrowstyle='-|>',
                        arrowsize=3, edge_cmap=None, edge_vmin=None, edge_vmax=None, ax=None, arrows=True, label=None,
                        node_size=300, nodelist=None, node_shape="o", connectionstyle=None, min_source_margin=0,
                        min_target_margin=0):<commit_msg>add node_size to pl.velocity_graph (relevant when arrows=True)<commit_after>from .. import settings
from ..preprocessing.neighbors import get_neighs
from ..tools.transition_matrix import transition_matrix
from .utils import savefig_or_show, default_basis, get_components, get_basis, groups_to_bool, default_size
from .scatter import scatter
from .docs import doc_scatter, doc_params

import warnings
import numpy as np
from scipy.sparse import issparse, csr_matrix


@doc_params(scatter=doc_scatter)
def velocity_graph(adata, basis=None, vkey='velocity', which_graph=None, n_neighbors=10, arrows=None, arrowsize=3,
                   alpha=.8, perc=None, threshold=None, edge_width=.2, edge_color='grey', edges_on_top=None, color=None,
                   layer=None, size=None, groups=None, components=None, title=None, dpi=None, show=True, 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 connectivities from neighbors graph.
    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)

    from networkx import Graph, DiGraph
    if which_graph in {'neighbors', 'connectivities'}:
        T = get_neighs(adata, 'connectivities').copy()
        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()
    elif which_graph in adata.uns.keys():
        T = adata.uns[which_graph].copy()
        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()
    elif hasattr(adata, 'obsp') and which_graph in adata.obsp.keys():
        T = adata.obsp[which_graph].copy()
        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()
    else:
        if threshold is None: threshold = .05
        T = transition_matrix(adata, vkey=vkey, weight_indirect_neighbors=0,
                              n_neighbors=n_neighbors, perc=perc, threshold=threshold)

    if groups is not None:
        if issparse(T): T = T.A
        T[~groups_to_bool(adata, groups, color)] = 0
        T = csr_matrix(T)
        T.eliminate_zeros()

    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        X_emb = adata.obsm['X_' + basis][:, get_components(components, basis)]
        node_size = (kwargs['size'] if 'size' in kwargs else default_size(adata)) / 4
        edges = draw_networkx_edges(DiGraph(T) if arrows else Graph(T), X_emb, node_size=node_size,
                                    width=edge_width, edge_color=edge_color, arrowsize=arrowsize, ax=ax)
        if not arrows and not edges_on_top:
            edges.set_zorder(-2)
            edges.set_rasterized(settings._vector_friendly)

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


def draw_networkx_edges(G, pos, edgelist=None, width=1.0, edge_color='k', style='solid', alpha=None, arrowstyle='-|>',
                        arrowsize=3, edge_cmap=None, edge_vmin=None, edge_vmax=None, ax=None, arrows=True, label=None,
                        node_size=300, nodelist=None, node_shape="o", connectionstyle=None, min_source_margin=0,
                        min_target_margin=0):
{'fraction_non_alphanumeric': 0.10982060717571297, 'fraction_numerical': 0.005864765409383624, 'mean_word_length': 2.886058981233244, 'pattern_counts': {'":': 16, '<': 9, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1987952', 'n_tokens_mistral': 2771, 'n_tokens_neox': 2626, 'n_words': 744}