pyphi/pyphi/visualize/connectivity.py at develop · wmayner/pyphi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
# visualize/connectivity.py
"""Visualize system connectivity information."""

import matplotlib.pyplot as plt
import networkx as nx
import numpy as np

from .distribution import all_states_str

NODE_COLORS = {
    # (in subsystem, state)
    (False, 0): "lightgrey",
    (False, 1): "darkgrey",
    (True, 0): "lightblue",
    (True, 1): "darkblue",
}


def plot_graph(g, **kwargs):
    kwargs = {
        **{
            "with_labels": True,
            "arrowsize": 20,
            "node_size": 600,
            "font_color": "white",
        },
        **kwargs,
    }
    nx.draw_networkx(
        g,
        **kwargs,
    )


def plot_subsystem(subsystem, **kwargs):
    g = nx.from_numpy_array(subsystem.cm, create_using=nx.DiGraph)
    nx.relabel_nodes(
        g,
        dict(zip(range(subsystem.network.size), subsystem.node_labels, strict=False)),
        copy=False,
    )
    if "node_color" not in kwargs:
        kwargs["node_color"] = [
            NODE_COLORS[(i in subsystem.node_indices, subsystem.state[i])]
            for i in range(subsystem.network.size)
        ]
    plot_graph(g, **kwargs)
    return g


def plot_tpm(
    tpm,
    figsize=(10, 12),
    clim=None,
    cmap="viridis",
    label_fontsize=8,
    show_label_threshold=64,
    xticks_top=True,
):
    fig = plt.figure(figsize=figsize)
    ax = plt.axes()
    im = ax.imshow(tpm, cmap=cmap)
    plt.grid(False)
    for spine in ax.spines.values():
        spine.set_visible(False)
    cax = fig.add_axes(
        [
            ax.get_position().x1 + 0.05,
            ax.get_position().y0,
            0.05,
            ax.get_position().height,
        ]
    )
    plt.colorbar(im, cax=cax)
    if clim is not None:
        im.set_clim(*clim)
    if tpm.shape[1] <= show_label_threshold:
        ax.set_xticks(
            list(range(tpm.shape[1])),
            labels=all_states_str(int(np.log2(tpm.shape[1]))),
            rotation=90,
            fontsize=label_fontsize,
        )
        ax.xaxis.set_ticks_position("top" if xticks_top else "bottom")
        ax.xaxis.set_label_position("top" if xticks_top else "bottom")
    if tpm.shape[0] <= show_label_threshold:
        ax.set_yticks(
            list(range(tpm.shape[0])),
            labels=all_states_str(int(np.log2(tpm.shape[0]))),
            fontsize=label_fontsize,
        )
    return fig, ax