persistence

def persistence_diagram(pers, ax=None, **kwargs):
    """Plot a persistence diagram.

    Parameters
    ----------
    pers :      pd.DataFrame
                Persistent points from [`navis.persistence_points`][].
    ax :        matplotlib ax, optional
                Ax to plot on.
    **kwargs
                Keyword arguments are passed to `LineCollection`.

    Returns
    -------
    ax :        matplotlib ax

    """
    if not isinstance(pers, pd.DataFrame):
        raise TypeError(f'Expected DataFrame, got "{type(pers)}"')

    if not ax:
        fig, ax = plt.subplots()

    segs = []
    for i, (b, d) in enumerate(zip(pers.birth.values, pers.death.values)):
        segs.append([[b, i], [d, i]])
    lc = LineCollection(segs, **kwargs)
    ax.add_collection(lc)

    ax.set_xlim(-5, pers.death.max())
    ax.set_ylim(-5, pers.shape[0])

    ax.set_ylabel('segments')
    ax.set_xlabel('time')

    return ax

def persistence_vector_plot(x, normalize=True, ax=None,
                            persistence_kwargs={}, vector_kwargs={}):
    """Plot persistence vectors.

    Parameters
    ----------
    x :         TreeNeuron | MeshNeuron | NeuronList
                Neuron(s) to calculate persistence points for. For MeshNeurons,
                we will use the skeleton produced by/associated with its
                `.skeleton` property.

    Returns
    -------
    ax

    """
    if not isinstance(x, core.NeuronList):
        x = core.NeuronList(x)

    # Get persistence points for each skeleton
    pers = persistence_points(x, **persistence_kwargs)

    # Get the vectors
    vectors, samples = persistence_vectors(pers, **vector_kwargs)

    # Normalizing the vectors will produce more useful distances
    if normalize:
        vectors = vectors / vectors.max(axis=1).reshape(-1, 1)
    else:
        vectors = vectors / vectors.max()

    if not ax:
        fig, ax = plt.subplots()

    for n, v in zip(x, vectors):
        ax.plot(samples, v, label=n.label)

    return ax