documentation/tools.py at main · CausalRivers/documentation · 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
93
94
95
96
97
98
99
100
101
102
103
104
105
import requests
import numpy as np
import requests
import matplotlib.pyplot as plt
import geopandas as gpd
import networkx as nx
import matplotlib as mpl


def plot_current_state_of_graph(
    G,
    dpi=100,
    lim=(50, 52),
    limx=(9.8, 13.4),
    node_size=50,
    arrowsize=20,
    fs=(10, 10),
    font_size=1,
    save=False,
    river_map=0,
    ger_map=0,
    emphasize=[],
    label=True,
    arrowstyle="fancy",
    autozoom=None,
    width=1,
    show_edge_origin=False,
    hardcode_colors = [],
    ger_path = "visualization/geomaps/vg2500_bld.shp",
    river_path = 'visualization/german_rivers_bg.shp',
    extra_points = []
):
#

    pos = {x: np.flip(np.array(G.nodes[x]["p"][:2]).astype(float)) for x in G.nodes}

    fig, ax = plt.subplots(1, 1, figsize=fs)

    if ger_map:
        fp = ger_path
        map_df2 = gpd.read_file(fp)
        map_df2.plot(color="green", ax=ax, alpha=0.3, linewidth=5, edgecolor="black")

    if river_map:
        fp = river_path
        map_df = gpd.read_file(fp)
        map_df.plot(color="blue", alpha=0.3, ax=ax, linewidth=0.5, edgecolor='blue')


    if hardcode_colors:
        colors = hardcode_colors
    else:
        colors = []
        for x in G.nodes:
            if x in emphasize:
                colors.append("black")
            else:
                colors.append(G.nodes[x]["c"])
    if show_edge_origin:
        cmap = mpl.colormaps['Set1']
        ege_base_colors = cmap(np.linspace(0, 1, 8))
        edge_colors = []
        for x in G.edges:
            edge_colors.append(tuple(ege_base_colors[G.edges[x]["origin"]]))
    nx.draw_networkx(
        G,
        pos,
        with_labels=label,
        font_size=font_size,
        node_size=node_size,
        arrows=True,
        node_color=colors,
        arrowsize=arrowsize,
        edge_color= edge_colors if show_edge_origin else "black",
        #arrowstyle=arrowstyle,
        width=width,
        ax=ax,
    )
    if autozoom:
        ax.set_xlim(
            min([pos[x][0] for x in pos.keys()]) - autozoom,
            max([pos[x][0] for x in pos.keys()]) + autozoom,
        )
        ax.set_ylim(
            min([pos[x][1] for x in pos.keys()]) - autozoom,
            max([pos[x][1] for x in pos.keys()]) + autozoom,
        )
    else:
        ax.set_ylim(lim[0], lim[1])
        ax.set_xlim(limx[0], limx[1])

    if save:
        plt.savefig("saves/" + save + "_G.svg", dpi=dpi)  # , dpi= 500
        plt.close()

    ax.set_title("River Causal Benchmark")
    ax.set_frame_on(True)

    if len(extra_points):
        for ex in extra_points:
            ax.scatter(ex[0],ex[1],color="black")
            ax.annotate(ex[2], (ex[0],ex[1]))

    plt.show()