Mahbodez/knee_report_checklist

1import networkx as nx2import json3import matplotlib.pyplot as plt4 5 6class Node:7    def __init__(self, name: str, value=None, parent=None, children: list = []):8        self.name = name9        self.children = set(children)10        self.parent = parent11        self.value = value12 13    def __repr__(self):14        return self.name15 16    def __str__(self):17        return self.name18 19    def __eq__(self, other):20        return self.name == other.name21 22    def __hash__(self) -> int:23        return hash(self.name)24 25    # make serializable for json26    def __getstate__(self):27        return self.__dict__28 29    def __dict__(self):30        # return a dict of the node's attributes31        return {32            "name": self.name,33            "children": self.children,34            "parent": self.parent,35            "value": self.value,36        }37 38    def to_json(self):39        """40        Returns a JSON string representation of the node.41        """42        return json.dumps(self.__dict__)43 44    def add_child(self, child):45        self.children.add(child)46 47    def has_children(self):48        return len(self.children) > 049 50    def set_parent(self, new_parent):51        self.parent = new_parent52 53    def set_value(self, new_value):54        self.value = new_value55 56 57def read_json(fname: str) -> dict:58    assert fname.endswith(".json"), "File must be a json file"59    with open(fname, "r") as f:60        data = json.load(f)61    return dict(data)62 63 64def build_tree_from_dict(data: dict, connect_children: bool = True):65    # every dict key is a node's name66    # dict value is a dict with keys "value", "parent", "children"67    # "value" is the node's value68    # "parent" is the node's parent's name69    # "children" is a list of the node's children's names70    # create a networkx graph71    G = nx.Graph()72    nodes_dict = dict()73    # build the nodes74    for name, info in data.items():75        value = info["value"]76        parent = info["parent"]77        children: list = info["children"]78        nodes_dict[name] = Node(79            name=name, parent=parent, children=children, value=value80        )81        G.add_node(nodes_dict[name], value=value)82    # build the edges83    for _, node in nodes_dict.items():84        for child in node.children:85            G.add_edge(node, nodes_dict[child])86            # connect children to each other if connect_children is True87            if connect_children:88                for child2 in node.children:89                    if child != child2:90                        G.add_edge(nodes_dict[child], nodes_dict[child2])91    return G, nodes_dict92 93 94def build_tree_from_file(fname: str):95    data = read_json(fname)96    return build_tree_from_dict(data)97 98 99# calculate the number of edges between two nodes100def num_edges_between_nodes(G, node1, node2):101    return len(nx.shortest_path(G, node1, node2)) - 1102 103 104def explore_bfs(G: nx.Graph, source: Node, nodes_dict: dict[str, Node]):105    # start from a source node and explore the graph in a breadth-first manner106    # prioritize nodes with non-empty values107    # explore the graph and return a list of nodes in the order they were explored108    explored_nodes = []109    queue = [source]110    while queue:111        node = queue.pop(0)112        explored_nodes.append(node)113        for child in node.children:114            if nodes_dict[child].value:115                queue.insert(0, nodes_dict[child])116            else:117                queue.append(nodes_dict[child])118    return explored_nodes119 120 121def from_list(node_list: list[Node], directional=True):122    # create a tree from a list of nodes123    # and label the edges from the first node to the last node from 1 to n124    if directional:125        G = nx.DiGraph()126    else:127        G = nx.Graph()128    G.add_nodes_from(node_list)129    for i in range(len(node_list) - 1):130        G.add_edge(node_list[i], node_list[i + 1], label=i + 1)131    return G132 133 134def visualize_graph(135    graph: nx.Graph,136    layout_graph: nx.Graph,137    title="BFS Tree",138    fig_size=(30, 20),139    title_fontsize=20,140    edge_width=1,141    font_size=9,142    node_size=500,143    node_shape="o",144    prog="dot",145):146    graphviz_args = "-Goverlap=false -Gsplines=true -Gsep=0.1 -Gnodesep=0.1 -Gmaxiter=1000 -Gepsilon=0.0001 -Gstart=0"147    _, ax = plt.subplots(figsize=fig_size)148    ax.set_title(title, fontsize=title_fontsize)149    # also draw edge labels150    nx.draw(151        graph,152        ax=ax,153        with_labels=True,154        # color every node lightblue except the root which is colored red155        node_color=(["lightgreen"] + ["lightblue"] * (len(graph.nodes) - 2) + ["red"])156        if len(graph.nodes) > 2157        else ["lightgreen", "red"]158        if len(graph.nodes) == 2159        else ["lightgreen"],160        edge_color="gray",161        width=edge_width,162        font_size=font_size,163        # node size to be proportional to the node's value164        node_size=node_size,165        # shape set to rectangle166        node_shape=node_shape,167        pos=nx.nx_agraph.graphviz_layout(168            layout_graph, prog=prog, root="root", args=graphviz_args169        ),170    )171    nx.draw_networkx_edge_labels(172        graph,173        pos=nx.nx_agraph.graphviz_layout(174            layout_graph, prog=prog, root="root", args=graphviz_args175        ),176        edge_labels=nx.get_edge_attributes(graph, "label"),177        font_size=font_size,178    )179    plt.show()180 181 182def get_graph(183    graph: nx.Graph,184    layout_graph: nx.Graph,185    title="BFS Tree",186    fig_size=(30, 20),187    title_fontsize=20,188    edge_width=1,189    font_size=9,190    node_size=500,191    node_shape="o",192    prog="dot",193):194    graphviz_args = "-Goverlap=false -Gsplines=true -Gsep=0.1 -Gnodesep=0.1 -Gmaxiter=1000 -Gepsilon=0.0001 -Gstart=0"195    fig, ax = plt.subplots(figsize=fig_size)196    ax.set_title(title, fontsize=title_fontsize)197    nx.draw(198        graph,199        ax=ax,200        with_labels=True,201        # color every node lightblue except the root which is colored red202        node_color=(["lightgreen"] + ["lightblue"] * (len(graph.nodes) - 2) + ["red"])203        if len(graph.nodes) > 2204        else ["lightgreen", "red"]205        if len(graph.nodes) == 2206        else ["lightgreen"],207        edge_color="gray",208        width=edge_width,209        font_size=font_size,210        # node size to be proportional to the node's value211        node_size=node_size,212        # shape set to rectangle213        node_shape=node_shape,214        pos=nx.nx_agraph.graphviz_layout(215            layout_graph, prog=prog, root="root", args=graphviz_args216        ),217    )218    nx.draw_networkx_edge_labels(219        graph,220        pos=nx.nx_agraph.graphviz_layout(221            layout_graph, prog=prog, root="root", args=graphviz_args222        ),223        edge_labels=nx.get_edge_attributes(graph, "label"),224        font_size=font_size,225    )226    return fig, ax227