Simple graph animator

networkx_animation.py

import matplotlib.pyplot as plt
import numpy as np, networkx as nx
from matplotlib.collections import LineCollection


def setup(
    g: nx.Graph, layout: dict or callable, node_kwargs={}, edge_kwargs={}, **kwargs
) -> list:
    """
    Simple setup: creates scatter points and line segments (graph)

    Parameters
    ----------
    g : nx.Graph or nx.DiGraph, input graph

    layout: dict or callable, if dict it contains node -> position, else should be one of the
    layout generators for networkx or something similar that generates the dict.

    node_kwargs: dict, formatter for nodes

    edge_kwargs: dict, formatter for the edges

    Return
    ------

    list containing [plt.Figure, plt.Axes, [plt.PatchCollection, plt.LineCollection]]
    """
    if hasattr(layout, "__call__"):
        layout = layout(g)
    coordinates = np.array([i for i in layout.values()]).T

    segments = np.array([[layout[x], layout[y]] for x, y in g.edges()])

    fig, ax = plt.subplots()
    nodes = ax.scatter(*coordinates, **node_kwargs)
    edges = LineCollection(segments, **edge_kwargs)
    ax.add_artist(edges)
    fig.show()

    return [fig, ax, [nodes, edges]]


def update(frame_idx: int, artists: list, ax=plt.Axes) -> list:
    """
    Simple update function
    """

    nodes, edges = artists
    # emulate some change
    current_pos = nodes.get_offsets()
    new_pos = current_pos + np.random.randn(*current_pos.shape) * 1e-3
    n = len(new_pos)
    nodes.set_offsets(new_pos)
    new_paths = np.array(
        [[new_pos[idx], new_pos[np.random.randint(0, n)]] for idx in range(n)]
    )
    edges.set_paths(new_paths)

    # needed as artists are not considered in ax.relim
    ax.update_datalim(new_pos)
    ax.axis("equal")
    return [nodes, edges]


if __name__ == "__main__":
    g = nx.erdos_renyi_graph(100, 0.3)
    fig, ax, artists = setup(
        g,
        nx.circular_layout,
        edge_kwargs=dict(color="lightgray", alpha=0.8),
        node_kwargs=dict(s=10),
    )

    ax.axis("off")
    while True:
        update(1, artists, ax)
        ax.relim()
        plt.pause(1e-2)