GEMLA/visualize_simulation_tree.py
2025-09-05 09:24:53 -07:00

104 lines
No EOL
4.1 KiB
Python

# Re-importing necessary libraries
import json
import matplotlib.pyplot as plt
import networkx as nx
import random
def hierarchy_pos(G, root=None, width=1., vert_gap=0.2, vert_loc=0, xcenter=0.5):
if not nx.is_tree(G):
raise TypeError('cannot use hierarchy_pos on a graph that is not a tree')
if root is None:
if isinstance(G, nx.DiGraph):
root = next(iter(nx.topological_sort(G)))
else:
root = random.choice(list(G.nodes))
def _hierarchy_pos(G, root, width=2., vert_gap=0.2, vert_loc=0, xcenter=0.5, pos=None, parent=None):
if pos is None:
pos = {root: (xcenter, vert_loc)}
else:
pos[root] = (xcenter, vert_loc)
children = list(G.successors(root)) # Use successors to get children for DiGraph
if not isinstance(G, nx.DiGraph):
if parent is not None:
children.remove(parent)
if len(children) != 0:
dx = width / len(children)
nextx = xcenter - width / 2 - dx / 2
for child in children:
nextx += dx
pos = _hierarchy_pos(G, child, width=dx*2.0, vert_gap=vert_gap,
vert_loc=vert_loc - vert_gap, xcenter=nextx,
pos=pos, parent=root)
return pos
return _hierarchy_pos(G, root, width, vert_gap, vert_loc, xcenter)
# Simplified JSON data for demonstration
with open('gemla/round4.json', 'r') as file:
simplified_json_data = json.load(file)
# Function to traverse the tree and create a graph
def traverse(node, graph, parent=None):
if node is None:
return
node_id = node["val"]["id"]
if "node" in node["val"] and node["val"]["node"]:
scores = node["val"]["node"]["scores"]
generations = node["val"]["node"]["generation"]
population_size = node["val"]["node"]["population_size"]
# Prepare to track the highest score across all generations and the corresponding individual
overall_max_score = float('-inf')
overall_max_score_individual = None
overall_max_score_gen = None
for gen, gen_scores in enumerate(scores):
if gen_scores: # Ensure the dictionary is not empty
# Find the max score and the individual for this generation
max_score_for_gen = max(gen_scores.values())
individual_with_max_score_for_gen = max(gen_scores, key=gen_scores.get)
# if max_score_for_gen > overall_max_score:
overall_max_score = max_score_for_gen
overall_max_score_individual = individual_with_max_score_for_gen
overall_max_score_gen = gen
# print debug statement
# print(f"Node {node_id}: Max score: {overall_max_score:.6f} (Individual {overall_max_score_individual} in Gen {overall_max_score_gen})")
# print(f"Left: {node.get('left')}, Right: {node.get('right')}")
label = f"{node_id}\nGenerations: {generations}, Population: {population_size}\nMax score: {overall_max_score:.6f} (Individual {overall_max_score_individual} in Gen {overall_max_score_gen + 1 if overall_max_score_gen is not None else 'N/A'})"
else:
label = node_id
graph.add_node(node_id, label=label)
if parent:
graph.add_edge(parent, node_id)
traverse(node.get("left"), graph, parent=node_id)
traverse(node.get("right"), graph, parent=node_id)
# Create a directed graph
G = nx.DiGraph()
# Populate the graph
traverse(simplified_json_data[0], G)
# Find the root node (a node with no incoming edges)
root_candidates = [node for node, indeg in G.in_degree() if indeg == 0]
if root_candidates:
root_node = root_candidates[0] # Assuming there's only one root candidate
else:
root_node = None # This should ideally never happen in a properly structured tree
# Use the determined root node for hierarchy_pos
if root_node is not None:
pos = hierarchy_pos(G, root=root_node)
labels = nx.get_node_attributes(G, 'label')
nx.draw(G, pos, labels=labels, with_labels=True, arrows=True)
plt.show()
else:
print("No root node found. Cannot draw the tree.")