contact-tracing/function.py at main · akwaed/contact-tracing · 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
import networkx as nx
import plotly.graph_objects as go


def print_map(network, node_trace, edge_trace, infected_nodes):
    node_adjacencies = []
    node_text = []
    node_sizes = []

    # Go through each nodes
    for node, adjacencies in enumerate(network.adjacency()):
        # Modification if you want to print node size based on number of edges
        # node_sizes.append(len(adjacencies[1]))

        # print infected node color
        if node == infected_nodes:
            node_adjacencies.append("violet")
            node_text.append("(infected) Node number " + str(node) + ' # of connections: ' + str(len(adjacencies[1])))

        # visit other nodes
        else:
            try:
                # get the distance between contaminated node and target node
                node_adjacencies.append(nx.astar_path_length(network, infected_nodes, node))
            except nx.NetworkXNoPath:
                # if path doesn't exist print it white
                node_adjacencies.append("white")
            node_text.append('Node number ' + str(node))

    # get the node colors and text
    node_trace.marker.color = node_adjacencies
    node_trace.text = node_text

    ## modification for size of nodes pased on number of edges
    # node_trace.marker.size = node_sizes

    # graph parameters
    fig = go.Figure(data=[edge_trace, node_trace],
                    layout=go.Layout(
                        title='<br><b>NETWORK GRAPH DEMONSTRATING INFECTION SPREAD<b>',
                        titlefont_size=16,
                        titlefont=dict(family="sans serif", size=18, color="Black"),
                        showlegend=False,
                        plot_bgcolor="#202A44",
                        hovermode='closest',
                        margin=dict(b=20, l=5, r=5, t=40),
                        annotations=[dict(
                            text=" by Elnoel Akwa CS315 Spring 2022 University of Kentucky",
                            align="center",
                            font=dict(family="sans serif", size=18, color="white"),
                            showarrow=False,
                            xref="paper", yref="paper",
                            x=0.005, y=-0.002)],
                        xaxis=dict(showgrid=False, zeroline=False, showticklabels=False),
                        yaxis=dict(showgrid=False, zeroline=False, showticklabels=False))
                    )

    fig.update_traces(textposition='top center')

    # print graph
    fig.show()